Add support for writing NSEC records (python-zeroconf#971)

bdraco · web-flow · commit 768a23c656e3 · 2021-08-29T17:40:12.000-05:00
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
@@ -54,6 +54,35 @@ def test_parse_own_packet_question(self):
         generated.add_question(r.DNSQuestion("testname.local.", const._TYPE_SRV, const._CLASS_IN))
         r.DNSIncoming(generated.packets()[0])
 
+    def test_parse_own_packet_nsec(self):
+        answer = r.DNSNsec(
+            'eufy HomeBase2-2464._hap._tcp.local.',
+            const._TYPE_NSEC,
+            const._CLASS_IN | const._CLASS_UNIQUE,
+            const._DNS_OTHER_TTL,
+            'eufy HomeBase2-2464._hap._tcp.local.',
+            [const._TYPE_TXT, const._TYPE_SRV],
+        )
+
+        generated = r.DNSOutgoing(const._FLAGS_QR_RESPONSE)
+        generated.add_answer_at_time(answer, 0)
+        parsed = r.DNSIncoming(generated.packets()[0])
+        assert answer in parsed.answers
+
+        # Types > 255 should be ignored
+        answer_invalid_types = r.DNSNsec(
+            'eufy HomeBase2-2464._hap._tcp.local.',
+            const._TYPE_NSEC,
+            const._CLASS_IN | const._CLASS_UNIQUE,
+            const._DNS_OTHER_TTL,
+            'eufy HomeBase2-2464._hap._tcp.local.',
+            [const._TYPE_TXT, const._TYPE_SRV, 1000],
+        )
+        generated = r.DNSOutgoing(const._FLAGS_QR_RESPONSE)
+        generated.add_answer_at_time(answer_invalid_types, 0)
+        parsed = r.DNSIncoming(generated.packets()[0])
+        assert answer in parsed.answers
+
     def test_parse_own_packet_response(self):
         generated = r.DNSOutgoing(const._FLAGS_QR_RESPONSE)
         generated.add_answer_at_time(
diff --git a/zeroconf/_dns.py b/zeroconf/_dns.py
@@ -458,9 +458,25 @@ def __init__(
     ) -> None:
         super().__init__(name, type_, class_, ttl, created)
         self.next_name = next_name
-        self.rdtypes = rdtypes
+        self.rdtypes = sorted(rdtypes)
         self._hash = hash((self.key, type_, self.class_, next_name, *self.rdtypes))
 
+    def write(self, out: 'DNSOutgoing') -> None:
+        """Used in constructing an outgoing packet."""
+        bitmap = bytearray(b'\0' * 32)
+        for rdtype in self.rdtypes:
+            if rdtype > 255:  # mDNS only supports window 0
+                continue
+            offset = rdtype % 256
+            byte = offset // 8
+            total_octets = byte + 1
+            bitmap[byte] |= 0x80 >> (offset % 8)
+        out_bytes = bytes(bitmap[0:total_octets])
+        out.write_name(self.next_name)
+        out.write_short(0)
+        out.write_short(len(out_bytes))
+        out.write_string(out_bytes)
+
     def __eq__(self, other: Any) -> bool:
         """Tests equality on cpu and os"""
         return (
