regression-tests.recursor-dnssec/extendederrors.py

   1 #!/usr/bin/env python
   2 import struct
   3
   4 import dns
   5 import dns.edns
   6 import dns.flags
   7 import dns.message
   8 import dns.query
   9
  10 class ExtendedErrorOption(dns.edns.Option):
  11     """Implementation of rfc8914
  12     """
  13
  14     def __init__(self, code, extra):
  15         super(ExtendedErrorOption, self).__init__(15)
  16
  17         self.code = code
  18         self.extra = extra
  19
  20     def to_wire(self, file=None):
  21         """Create EDNS packet."""
  22
  23         data = struct.pack('!H', self.code)
  24         data = data + self.extra
  25         if file:
  26             file.write(data)
  27         else:
  28             return data
  29
  30     def from_wire(cls, otype, wire, current, olen):
  31         """Read EDNS packet.
  32
  33         Returns:
  34             An instance of ExtendedErrorOption based on the EDNS packet
  35         """
  36
  37         if olen < 2:
  38             raise Exception('Invalid EDNS Extended Error option')
  39
  40         (code,) = struct.unpack('!H', wire[current:current+2])
  41         if olen > 2:
  42             extra = wire[current + 2:current + olen]
  43         else:
  44             extra = b''
  45
  46         return cls(code, extra)
  47
  48     from_wire = classmethod(from_wire)
  49
  50     # needed in 2.0.0
  51     @classmethod
  52     def from_wire_parser(cls, otype, parser):
  53         data = parser.get_remaining()
  54
  55         if len(data) < 2:
  56             raise Exception('Invalid EDNS Extended Error option')
  57
  58         (code,) = struct.unpack('!H', data[0:2])
  59         if len(data) > 2:
  60             extra = data[2:]
  61         else:
  62             extra = b''
  63
  64         return cls(code, extra)
  65
  66     def __repr__(self):
  67         return '%s(%d, %s)' % (
  68             self.__class__.__name__,
  69             self.code,
  70             self.extra
  71         )
  72
  73     def to_text(self):
  74         return self.__repr__()
  75
  76     def __eq__(self, other):
  77         if not isinstance(other, ExtendedErrorOption):
  78             return False
  79         if self.code != other.code:
  80             return False
  81         if self.extra != other.extra:
  82             return False
  83         return True
  84
  85     def __ne__(self, other):
  86         return not self.__eq__(other)
  87
  88
  89 dns.edns._type_to_class[0x000F] = ExtendedErrorOption