diff options
author | Brian Wellington <bwelling@xbill.org> | 2020-07-17 17:01:32 -0700 |
---|---|---|
committer | Brian Wellington <bwelling@xbill.org> | 2020-07-17 17:01:32 -0700 |
commit | 8669f5740c336e58aaed6e5fa9bbfdcff97fd8be (patch) | |
tree | 7339da17ef9935b0b232898121b6ae910ca715bb | |
parent | 1836e45b701352ec86de5517b2c3e017e5189708 (diff) | |
download | dnspython-8669f5740c336e58aaed6e5fa9bbfdcff97fd8be.tar.gz |
Add dns.edns.register_type().
This allows an application to register a custom EDNS option type.
-rw-r--r-- | dns/edns.py | 10 | ||||
-rw-r--r-- | tests/test_edns.py | 37 |
2 files changed, 47 insertions, 0 deletions
diff --git a/dns/edns.py b/dns/edns.py index 28718d5..05f76e6 100644 --- a/dns/edns.py +++ b/dns/edns.py @@ -342,3 +342,13 @@ def option_from_wire(otype, wire, current, olen): parser = dns.wire.Parser(wire, current) with parser.restrict_to(olen): return option_from_wire_parser(otype, parser) + +def register_type(implementation, otype): + """Register the implementation of an option type. + + *implementation*, a ``class``, is a subclass of ``dns.edns.Option``. + + *otype*, an ``int``, is the option type. + """ + + _type_to_class[otype] = implementation diff --git a/tests/test_edns.py b/tests/test_edns.py index a640a74..e5dca06 100644 --- a/tests/test_edns.py +++ b/tests/test_edns.py @@ -17,6 +17,7 @@ # OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. import operator +import struct import unittest from io import BytesIO @@ -138,3 +139,39 @@ class OptionTestCase(unittest.TestCase): self.assertTrue(o1 != o2) self.assertFalse(o1 == 123) self.assertTrue(o1 != 123) + + def test_option_registration(self): + U32OptionType = 9999 + + class U32Option(dns.edns.Option): + def __init__(self, value=None): + super().__init__(U32OptionType) + self.value = value + + def to_wire(self, file=None): + data = struct.pack('!I', self.value) + if file: + file.write(data) + else: + return data + + @classmethod + def from_wire_parser(cls, otype, parser): + (value,) = parser.get_struct('!I') + return cls(value) + + try: + dns.edns.register_type(U32Option, U32OptionType) + generic = dns.edns.GenericOption(U32OptionType, b'\x00\x00\x00\x01') + wire1 = generic.to_wire() + u32 = dns.edns.option_from_wire_parser(U32OptionType, + dns.wire.Parser(wire1)) + self.assertEqual(u32.value, 1) + wire2 = u32.to_wire() + self.assertEqual(wire1, wire2) + self.assertEqual(u32, generic) + finally: + dns.edns._type_to_class.pop(U32OptionType, None) + + opt = dns.edns.option_from_wire_parser(9999, dns.wire.Parser(wire1)) + self.assertEqual(opt, generic) |