Merge pull request #14195 from rgacogne/ddist-no-assertions

[thirdparty/pdns.git] / regression-tests.dnsdist / clientsubnetoption.py

diff --git a/regression-tests.dnsdist/clientsubnetoption.py b/regression-tests.dnsdist/clientsubnetoption.py
index 5306985288fa3f82d8a4c3914bf57a77bd8f618f..5da748ab0b0901758a0a5bb4349f75955599d229 100644 (file)
--- a/regression-tests.dnsdist/clientsubnetoption.py
+++ b/regression-tests.dnsdist/clientsubnetoption.py
@@ -35,6 +35,7 @@ Requirements:
 from __future__ import print_function
 from __future__ import division
 
+import math
 import socket
 import struct
 import dns
@@ -58,7 +59,7 @@ class ClientSubnetOption(dns.edns.Option):
     """Implementation of draft-vandergaast-edns-client-subnet-01.
 
     Attributes:
-        family: An integer inidicating which address family is being sent
+        family: An integer indicating which address family is being sent
         ip: IP address in integer notation
         mask: An integer representing the number of relevant bits being sent
         scope: An integer representing the number of significant bits used by
@@ -125,8 +126,8 @@ class ClientSubnetOption(dns.edns.Option):
         """" Determines whether this instance is using the draft option code """
         return self.option == DRAFT_OPTION_CODE
 
-    def to_wire(self, file):
-        """Create EDNS packet as definied in draft-vandergaast-edns-client-subnet-01."""
+    def to_wire(self, file=None):
+        """Create EDNS packet as defined in draft-vandergaast-edns-client-subnet-01."""
 
         ip = self.calculate_ip()
 
@@ -142,7 +143,10 @@ class ClientSubnetOption(dns.edns.Option):
 
         format = "!HBB%ds" % (mask_bits // 8)
         data = struct.pack(format, self.family, self.mask, self.scope, test)
-        file.write(data)
+        if file:
+            file.write(data)
+        else:
+            return data
 
     def from_wire(cls, otype, wire, current, olen):
         """Read EDNS packet as defined in draft-vandergaast-edns-client-subnet-01.
@@ -173,6 +177,23 @@ class ClientSubnetOption(dns.edns.Option):
 
     from_wire = classmethod(from_wire)
 
+    # needed in 2.0.0..
+    @classmethod
+    def from_wire_parser(cls, otype, parser):
+        family, src, scope = parser.get_struct('!HBB')
+        addrlen = int(math.ceil(src / 8.0))
+        prefix = parser.get_bytes(addrlen)
+        if family == 1:
+            pad = 4 - addrlen
+            addr = dns.ipv4.inet_ntoa(prefix + b'\x00' * pad)
+        elif family == 2:
+            pad = 16 - addrlen
+            addr = dns.ipv6.inet_ntoa(prefix + b'\x00' * pad)
+        else:
+            raise ValueError('unsupported family')
+
+        return cls(addr, src, scope, otype)
+
     def __repr__(self):
         if self.family == FAMILY_IPV4:
             ip = socket.inet_ntop(socket.AF_INET, struct.pack('!L', self.ip))
@@ -189,6 +210,9 @@ class ClientSubnetOption(dns.edns.Option):
             self.scope
         )
 
+    def to_text(self):
+        return self.__repr__()
+
     def __eq__(self, other):
         """Rich comparison method for equality.