From 91d329a748c4d57c94ab2fef977d75d64b6bb66f Mon Sep 17 00:00:00 2001
From: Brian Wellington <bwelling@xbill.org>
Date: Fri, 11 Oct 2024 13:34:57 -0700
Subject: [PATCH] Add dns.edns.Option.to_generic() (#1145)

* Add dns.edns.Option.to_generic()

Converts an EDNS option represented by a custom class into an equivalent
option represented by the generic option class.  This is similar to the
existing dns.rdata.Rdata.to_generic() method.

Also, adds a specialization to the existing dns.rdata.Rdata.to_generic()
method for GenericRdata, to avoid extra work for applications that want
to convert all rdata to generic form.

* Fix typing issue.

* Fix typing again.
---
 dns/edns.py         | 12 ++++++++++++
 dns/rdata.py        |  5 +++++
 tests/test_edns.py  | 10 ++++++++++
 tests/test_rdata.py |  3 +++
 4 files changed, 30 insertions(+)

diff --git a/dns/edns.py b/dns/edns.py
index f7d9ff99..c3603686 100644
--- a/dns/edns.py
+++ b/dns/edns.py
@@ -81,6 +81,15 @@ class Option:
     def to_text(self) -> str:
         raise NotImplementedError  # pragma: no cover
 
+    def to_generic(self) -> "dns.edns.GenericOption":
+        """Creates a dns.edns.GenericOption equivalent of this rdata.
+
+        Returns a ``dns.edns.GenericOption``.
+        """
+        wire = self.to_wire()
+        assert wire is not None  # for mypy
+        return dns.edns.GenericOption(self.otype, wire)
+
     @classmethod
     def from_wire_parser(cls, otype: OptionType, parser: "dns.wire.Parser") -> "Option":
         """Build an EDNS option object from wire format.
@@ -166,6 +175,9 @@ class GenericOption(Option):  # lgtm[py/missing-equals]
     def to_text(self) -> str:
         return "Generic %d" % self.otype
 
+    def to_generic(self) -> "dns.edns.GenericOption":
+        return self
+
     @classmethod
     def from_wire_parser(
         cls, otype: Union[OptionType, str], parser: "dns.wire.Parser"
diff --git a/dns/rdata.py b/dns/rdata.py
index 8099c26a..0189f240 100644
--- a/dns/rdata.py
+++ b/dns/rdata.py
@@ -639,6 +639,11 @@ class GenericRdata(Rdata):
     def _to_wire(self, file, compress=None, origin=None, canonicalize=False):
         file.write(self.data)
 
+    def to_generic(
+        self, origin: Optional[dns.name.Name] = None
+    ) -> "dns.rdata.GenericRdata":
+        return self
+
     @classmethod
     def from_wire_parser(cls, rdclass, rdtype, parser, origin=None):
         return cls(rdclass, rdtype, parser.get_remaining())
diff --git a/tests/test_edns.py b/tests/test_edns.py
index 952ca05e..1229c1e9 100644
--- a/tests/test_edns.py
+++ b/tests/test_edns.py
@@ -302,3 +302,13 @@ class OptionTestCase(unittest.TestCase):
 
         opt = dns.edns.option_from_wire_parser(9999, dns.wire.Parser(wire1))
         self.assertEqual(opt, generic)
+
+    def test_to_generic(self):
+        nsid = dns.edns.NSIDOption(b"testing")
+        assert nsid.to_generic().data == b"testing"
+
+        ecs = dns.edns.ECSOption("1.2.3.0", 24)
+        assert ecs.to_generic().data == b"\x00\x01\x18\x00\x01\x02\x03"
+
+        generic = dns.edns.GenericOption(12345, "foo")
+        assert generic.to_generic() is generic
diff --git a/tests/test_rdata.py b/tests/test_rdata.py
index c63dc18c..4c62aa1d 100644
--- a/tests/test_rdata.py
+++ b/tests/test_rdata.py
@@ -130,6 +130,9 @@ class RdataTestCase(unittest.TestCase):
             str(ns.to_generic(origin=origin)), r"\# 13 03666f6f076578616d706c6500"
         )
 
+        generic = dns.rdata.from_text("in", "type45678", "\\# 4 00010203")
+        assert generic.to_generic() is generic
+
     def test_txt_unicode(self):
         # TXT records are not defined for Unicode, but if we get
         # Unicode we should convert it to UTF-8 to preserve meaning as
-- 
2.47.3