regression-tests.dnsdist/clientsubnetoption.py

   1 #!/usr/bin/env python
   2 #
   3 # Copyright (c) 2012 OpenDNS, Inc.
   4 # All rights reserved.
   5 #
   6 # Redistribution and use in source and binary forms, with or without
   7 # modification, are permitted provided that the following conditions are met:
   8 #    * Redistributions of source code must retain the above copyright
   9 #      notice, this list of conditions and the following disclaimer.
  10 #    * Redistributions in binary form must reproduce the above copyright
  11 #      notice, this list of conditions and the following disclaimer in the
  12 #      documentation and/or other materials provided with the distribution.
  13 #    * Neither the name of the OpenDNS nor the names of its contributors may be
  14 #      used to endorse or promote products derived from this software without
  15 #      specific prior written permission.
  16 #
  17 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  18 # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  19 # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  20 # DISCLAIMED. IN NO EVENT SHALL OPENDNS BE LIABLE FOR ANY DIRECT, INDIRECT,
  21 # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  22 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
  23 # OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
  24 # LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
  25 # NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
  26 # EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  27
  28 """ Class to implement draft-ietf-dnsop-edns-client-subnet (previously known as
  29 draft-vandergaast-edns-client-subnet.
  30
  31 The contained class supports both IPv4 and IPv6 addresses.
  32 Requirements:
  33   dnspython (http://www.dnspython.org/)
  34 """
  35 from __future__ import print_function
  36 from __future__ import division
  37
  38 import socket
  39 import struct
  40 import dns
  41 import dns.edns
  42 import dns.flags
  43 import dns.message
  44 import dns.query
  45
  46 __author__ = "bhartvigsen@opendns.com (Brian Hartvigsen)"
  47 __version__ = "2.0.0"
  48
  49 ASSIGNED_OPTION_CODE = 0x0008
  50 DRAFT_OPTION_CODE = 0x50FA
  51
  52 FAMILY_IPV4 = 1
  53 FAMILY_IPV6 = 2
  54 SUPPORTED_FAMILIES = (FAMILY_IPV4, FAMILY_IPV6)
  55
  56
  57 class ClientSubnetOption(dns.edns.Option):
  58     """Implementation of draft-vandergaast-edns-client-subnet-01.
  59
  60     Attributes:
  61         family: An integer indicating which address family is being sent
  62         ip: IP address in integer notation
  63         mask: An integer representing the number of relevant bits being sent
  64         scope: An integer representing the number of significant bits used by
  65             the authoritative server.
  66     """
  67
  68     def __init__(self, ip, bits=24, scope=0, option=ASSIGNED_OPTION_CODE):
  69         super(ClientSubnetOption, self).__init__(option)
  70
  71         n = None
  72         f = None
  73
  74         for family in (socket.AF_INET, socket.AF_INET6):
  75             try:
  76                 n = socket.inet_pton(family, ip)
  77                 if family == socket.AF_INET6:
  78                     f = FAMILY_IPV6
  79                     hi, lo = struct.unpack('!QQ', n)
  80                     ip = hi << 64 | lo
  81                 elif family == socket.AF_INET:
  82                     f = FAMILY_IPV4
  83                     ip = struct.unpack('!L', n)[0]
  84             except Exception:
  85                 pass
  86
  87         if n is None:
  88             raise Exception("%s is an invalid ip" % ip)
  89
  90         self.family = f
  91         self.ip = ip
  92         self.mask = bits
  93         self.scope = scope
  94         self.option = option
  95
  96         if self.family == FAMILY_IPV4 and self.mask > 32:
  97             raise Exception("32 bits is the max for IPv4 (%d)" % bits)
  98         if self.family == FAMILY_IPV6 and self.mask > 128:
  99             raise Exception("128 bits is the max for IPv6 (%d)" % bits)
 100
 101     def calculate_ip(self):
 102         """Calculates the relevant ip address based on the network mask.
 103
 104         Calculates the relevant bits of the IP address based on network mask.
 105         Sizes up to the nearest octet for use with wire format.
 106
 107         Returns:
 108             An integer of only the significant bits sized up to the nearest
 109             octect.
 110         """
 111
 112         if self.family == FAMILY_IPV4:
 113             bits = 32
 114         elif self.family == FAMILY_IPV6:
 115             bits = 128
 116
 117         ip = self.ip >> bits - self.mask
 118
 119         if (self.mask % 8 != 0):
 120             ip = ip << 8 - (self.mask % 8)
 121
 122         return ip
 123
 124     def is_draft(self):
 125         """" Determines whether this instance is using the draft option code """
 126         return self.option == DRAFT_OPTION_CODE
 127
 128     def to_wire(self, file):
 129         """Create EDNS packet as defined in draft-vandergaast-edns-client-subnet-01."""
 130
 131         ip = self.calculate_ip()
 132
 133         mask_bits = self.mask
 134         if mask_bits % 8 != 0:
 135                 mask_bits += 8 - (self.mask % 8)
 136
 137         if self.family == FAMILY_IPV4:
 138             test = struct.pack("!L", ip)
 139         elif self.family == FAMILY_IPV6:
 140             test = struct.pack("!QQ", ip >> 64, ip & (2 ** 64 - 1))
 141         test = test[-(mask_bits // 8):]
 142
 143         format = "!HBB%ds" % (mask_bits // 8)
 144         data = struct.pack(format, self.family, self.mask, self.scope, test)
 145         file.write(data)
 146
 147     def from_wire(cls, otype, wire, current, olen):
 148         """Read EDNS packet as defined in draft-vandergaast-edns-client-subnet-01.
 149
 150         Returns:
 151             An instance of ClientSubnetOption based on the ENDS packet
 152         """
 153
 154         data = wire[current:current + olen]
 155         (family, mask, scope) = struct.unpack("!HBB", data[:4])
 156
 157         c_mask = mask
 158         if mask % 8 != 0:
 159             c_mask += 8 - (mask % 8)
 160
 161         ip = struct.unpack_from("!%ds" % (c_mask // 8), data, 4)[0]
 162
 163         if (family == FAMILY_IPV4):
 164             ip = ip + b'\0' * ((32 - c_mask) // 8)
 165             ip = socket.inet_ntop(socket.AF_INET, ip)
 166         elif (family == FAMILY_IPV6):
 167             ip = ip + b'\0' * ((128 - c_mask) // 8)
 168             ip = socket.inet_ntop(socket.AF_INET6, ip)
 169         else:
 170             raise Exception("Returned a family other then IPv4 or IPv6")
 171
 172         return cls(ip, mask, scope, otype)
 173
 174     from_wire = classmethod(from_wire)
 175
 176     def __repr__(self):
 177         if self.family == FAMILY_IPV4:
 178             ip = socket.inet_ntop(socket.AF_INET, struct.pack('!L', self.ip))
 179         elif self.family == FAMILY_IPV6:
 180             ip = socket.inet_ntop(socket.AF_INET6,
 181                                   struct.pack('!QQ',
 182                                               self.ip >> 64,
 183                                               self.ip & (2 ** 64 - 1)))
 184
 185         return "%s(%s, %s, %s)" % (
 186             self.__class__.__name__,
 187             ip,
 188             self.mask,
 189             self.scope
 190         )
 191
 192     def __eq__(self, other):
 193         """Rich comparison method for equality.
 194
 195         Two ClientSubnetOptions are equal if their relevant ip bits, mask, and
 196         family are identical. We ignore scope since generally we want to
 197         compare questions to responses and that bit is only relevant when
 198         determining caching behavior.
 199
 200         Returns:
 201             boolean
 202         """
 203
 204         if not isinstance(other, ClientSubnetOption):
 205             return False
 206         if self.calculate_ip() != other.calculate_ip():
 207             return False
 208         if self.mask != other.mask:
 209             return False
 210         if self.family != other.family:
 211             return False
 212         return True
 213
 214     def __ne__(self, other):
 215         """Rich comparison method for inequality.
 216
 217         See notes for __eq__()
 218
 219         Returns:
 220             boolean
 221         """
 222         return not self.__eq__(other)
 223
 224
 225 dns.edns._type_to_class[DRAFT_OPTION_CODE] = ClientSubnetOption
 226 dns.edns._type_to_class[ASSIGNED_OPTION_CODE] = ClientSubnetOption
 227
 228 if __name__ == "__main__":
 229     import argparse
 230     import sys
 231
 232     def CheckForClientSubnetOption(addr, args, option_code=ASSIGNED_OPTION_CODE):
 233         print("Testing for edns-clientsubnet using option code", hex(option_code), file=sys.stderr)
 234         cso = ClientSubnetOption(args.subnet, args.mask, option=option_code)
 235         message = dns.message.make_query(args.rr, args.type)
 236         # Tested authoritative servers seem to use the last code in cases
 237         # where they support both. We make the official code last to allow
 238         # us to check for support of both draft and official
 239         message.use_edns(options=[cso])
 240
 241         try:
 242             r = dns.query.udp(message, addr, timeout=args.timeout)
 243             if r.flags & dns.flags.TC:
 244                 r = dns.query.tcp(message, addr, timeout=args.timeout)
 245         except dns.exception.Timeout:
 246             print("Timeout: No answer received from %s\n" % args.nameserver, file=sys.stderr)
 247             sys.exit(3)
 248
 249         error = False
 250         found = False
 251         for options in r.options:
 252             # Have not run into anyone who passes back both codes yet
 253             # but just in case, we want to check all possible options
 254             if isinstance(options, ClientSubnetOption):
 255                 found = True
 256                 print("Found ClientSubnetOption...", end=None, file=sys.stderr)
 257                 if not cso.family == options.family:
 258                     error = True
 259                     print("\nFailed: returned family (%d) is different from the passed family (%d)" % (options.family, cso.family), file=sys.stderr)
 260                 if not cso.calculate_ip() == options.calculate_ip():
 261                     error = True
 262                     print("\nFailed: returned ip (%s) is different from the passed ip (%s)." % (options.calculate_ip(), cso.calculate_ip()), file=sys.stderr)
 263                 if not options.mask == cso.mask:
 264                     error = True
 265                     print("\nFailed: returned mask bits (%d) is different from the passed mask bits (%d)" % (options.mask, cso.mask), file=sys.stderr)
 266                 if not options.scope != 0:
 267                     print("\nWarning: scope indicates edns-clientsubnet data is not used", file=sys.stderr)
 268                 if options.is_draft():
 269                     print("\nWarning: detected support for edns-clientsubnet draft code", file=sys.stderr)
 270
 271         if found and not error:
 272             print("Success", file=sys.stderr)
 273         elif found:
 274             print("Failed: See error messages above", file=sys.stderr)
 275         else:
 276             print("Failed: No ClientSubnetOption returned", file=sys.stderr)
 277
 278     parser = argparse.ArgumentParser(description='draft-vandergaast-edns-client-subnet-01 tester')
 279     parser.add_argument('nameserver', help='The nameserver to test')
 280     parser.add_argument('rr', help='DNS record that should return an EDNS enabled response')
 281     parser.add_argument('-s', '--subnet', help='Specifies an IP to pass as the client subnet.', default='192.0.2.0')
 282     parser.add_argument('-m', '--mask', type=int, help='CIDR mask to use for subnet')
 283     parser.add_argument('--timeout', type=int, help='Set the timeout for query to TIMEOUT seconds, default=10', default=10)
 284     parser.add_argument('-t', '--type', help='DNS query type, default=A', default='A')
 285     args = parser.parse_args()
 286
 287     if not args.mask:
 288         if ':' in args.subnet:
 289             args.mask = 48
 290         else:
 291             args.mask = 24
 292
 293     try:
 294         addr = socket.gethostbyname(args.nameserver)
 295     except socket.gaierror:
 296         print("Unable to resolve %s\n" % args.nameserver, file=sys.stderr)
 297         sys.exit(3)
 298
 299     CheckForClientSubnetOption(addr, args, DRAFT_OPTION_CODE)
 300     print("", file=sys.stderr)
 301     CheckForClientSubnetOption(addr, args, ASSIGNED_OPTION_CODE)