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 to_text(self):
 193         return self.__repr__()
 194
 195     def __eq__(self, other):
 196         """Rich comparison method for equality.
 197
 198         Two ClientSubnetOptions are equal if their relevant ip bits, mask, and
 199         family are identical. We ignore scope since generally we want to
 200         compare questions to responses and that bit is only relevant when
 201         determining caching behavior.
 202
 203         Returns:
 204             boolean
 205         """
 206
 207         if not isinstance(other, ClientSubnetOption):
 208             return False
 209         if self.calculate_ip() != other.calculate_ip():
 210             return False
 211         if self.mask != other.mask:
 212             return False
 213         if self.family != other.family:
 214             return False
 215         return True
 216
 217     def __ne__(self, other):
 218         """Rich comparison method for inequality.
 219
 220         See notes for __eq__()
 221
 222         Returns:
 223             boolean
 224         """
 225         return not self.__eq__(other)
 226
 227
 228 dns.edns._type_to_class[DRAFT_OPTION_CODE] = ClientSubnetOption
 229 dns.edns._type_to_class[ASSIGNED_OPTION_CODE] = ClientSubnetOption
 230
 231 if __name__ == "__main__":
 232     import argparse
 233     import sys
 234
 235     def CheckForClientSubnetOption(addr, args, option_code=ASSIGNED_OPTION_CODE):
 236         print("Testing for edns-clientsubnet using option code", hex(option_code), file=sys.stderr)
 237         cso = ClientSubnetOption(args.subnet, args.mask, option=option_code)
 238         message = dns.message.make_query(args.rr, args.type)
 239         # Tested authoritative servers seem to use the last code in cases
 240         # where they support both. We make the official code last to allow
 241         # us to check for support of both draft and official
 242         message.use_edns(options=[cso])
 243
 244         try:
 245             r = dns.query.udp(message, addr, timeout=args.timeout)
 246             if r.flags & dns.flags.TC:
 247                 r = dns.query.tcp(message, addr, timeout=args.timeout)
 248         except dns.exception.Timeout:
 249             print("Timeout: No answer received from %s\n" % args.nameserver, file=sys.stderr)
 250             sys.exit(3)
 251
 252         error = False
 253         found = False
 254         for options in r.options:
 255             # Have not run into anyone who passes back both codes yet
 256             # but just in case, we want to check all possible options
 257             if isinstance(options, ClientSubnetOption):
 258                 found = True
 259                 print("Found ClientSubnetOption...", end=None, file=sys.stderr)
 260                 if not cso.family == options.family:
 261                     error = True
 262                     print("\nFailed: returned family (%d) is different from the passed family (%d)" % (options.family, cso.family), file=sys.stderr)
 263                 if not cso.calculate_ip() == options.calculate_ip():
 264                     error = True
 265                     print("\nFailed: returned ip (%s) is different from the passed ip (%s)." % (options.calculate_ip(), cso.calculate_ip()), file=sys.stderr)
 266                 if not options.mask == cso.mask:
 267                     error = True
 268                     print("\nFailed: returned mask bits (%d) is different from the passed mask bits (%d)" % (options.mask, cso.mask), file=sys.stderr)
 269                 if not options.scope != 0:
 270                     print("\nWarning: scope indicates edns-clientsubnet data is not used", file=sys.stderr)
 271                 if options.is_draft():
 272                     print("\nWarning: detected support for edns-clientsubnet draft code", file=sys.stderr)
 273
 274         if found and not error:
 275             print("Success", file=sys.stderr)
 276         elif found:
 277             print("Failed: See error messages above", file=sys.stderr)
 278         else:
 279             print("Failed: No ClientSubnetOption returned", file=sys.stderr)
 280
 281     parser = argparse.ArgumentParser(description='draft-vandergaast-edns-client-subnet-01 tester')
 282     parser.add_argument('nameserver', help='The nameserver to test')
 283     parser.add_argument('rr', help='DNS record that should return an EDNS enabled response')
 284     parser.add_argument('-s', '--subnet', help='Specifies an IP to pass as the client subnet.', default='192.0.2.0')
 285     parser.add_argument('-m', '--mask', type=int, help='CIDR mask to use for subnet')
 286     parser.add_argument('--timeout', type=int, help='Set the timeout for query to TIMEOUT seconds, default=10', default=10)
 287     parser.add_argument('-t', '--type', help='DNS query type, default=A', default='A')
 288     args = parser.parse_args()
 289
 290     if not args.mask:
 291         if ':' in args.subnet:
 292             args.mask = 48
 293         else:
 294             args.mask = 24
 295
 296     try:
 297         addr = socket.gethostbyname(args.nameserver)
 298     except socket.gaierror:
 299         print("Unable to resolve %s\n" % args.nameserver, file=sys.stderr)
 300         sys.exit(3)
 301
 302     CheckForClientSubnetOption(addr, args, DRAFT_OPTION_CODE)
 303     print("", file=sys.stderr)
 304     CheckForClientSubnetOption(addr, args, ASSIGNED_OPTION_CODE)