sleekxmpp/xmlstream/resolver.py

   1 # -*- encoding: utf-8 -*-
   2
   3 """
   4     sleekxmpp.xmlstream.dns
   5     ~~~~~~~~~~~~~~~~~~~~~~~
   6
   7     :copyright: (c) 2012 Nathanael C. Fritz
   8     :license: MIT, see LICENSE for more details
   9 """
  10
  11 import socket
  12 import logging
  13 import random
  14
  15
  16 log = logging.getLogger(__name__)
  17
  18
  19 #: Global flag indicating the availability of the ``dnspython`` package.
  20 #: Installing ``dnspython`` can be done via:
  21 #:
  22 #: .. code-block:: sh
  23 #:
  24 #:     pip install dnspython
  25 #:
  26 #: For Python3, installation may require installing from source using
  27 #: the ``python3`` branch:
  28 #:
  29 #: .. code-block:: sh
  30 #:
  31 #:     git clone http://github.com/rthalley/dnspython
  32 #:     cd dnspython
  33 #:     git checkout python3
  34 #:     python3 setup.py install
  35 DNSPYTHON_AVAILABLE = False
  36 try:
  37     import dns.resolver
  38     DNSPYTHON_AVAILABLE = True
  39 except ImportError as e:
  40     log.debug("Could not find dnspython package. " + \
  41               "Not all features will be available")
  42
  43
  44 def default_resolver():
  45     """Return a basic DNS resolver object.
  46
  47     :returns: A :class:`dns.resolver.Resolver` object if dnspython
  48               is available. Otherwise, ``None``.
  49     """
  50     if DNSPYTHON_AVAILABLE:
  51         return dns.resolver.get_default_resolver()
  52     return None
  53
  54
  55 def resolve(host, port=None, service=None, proto='tcp',
  56             resolver=None, use_ipv6=True, use_dnspython=True):
  57     """Peform DNS resolution for a given hostname.
  58
  59     Resolution may perform SRV record lookups if a service and protocol
  60     are specified. The returned addresses will be sorted according to
  61     the SRV priorities and weights.
  62
  63     If no resolver is provided, the dnspython resolver will be used if
  64     available. Otherwise the built-in socket facilities will be used,
  65     but those do not provide SRV support.
  66
  67     If SRV records were used, queries to resolve alternative hosts will
  68     be made as needed instead of all at once.
  69
  70     :param     host: The hostname to resolve.
  71     :param     port: A default port to connect with. SRV records may
  72                      dictate use of a different port.
  73     :param  service: Optional SRV service name without leading underscore.
  74     :param    proto: Optional SRV protocol name without leading underscore.
  75     :param resolver: Optionally provide a DNS resolver object that has
  76                      been custom configured.
  77     :param use_ipv6: Optionally control the use of IPv6 in situations
  78                      where it is either not available, or performance
  79                      is degraded. Defaults to ``True``.
  80     :param use_dnspython: Optionally control if dnspython is used to make
  81                           the DNS queries instead of the built-in DNS
  82                           library.
  83
  84     :type     host: string
  85     :type     port: int
  86     :type  service: string
  87     :type    proto: string
  88     :type resolver: :class:`dns.resolver.Resolver`
  89     :type use_ipv6: bool
  90     :type use_dnspython: bool
  91
  92     :return: An iterable of IP address, port pairs in the order
  93              dictated by SRV priorities and weights, if applicable.
  94     """
  95
  96     if not use_dnspython:
  97         if DNSPYTHON_AVAILABLE:
  98             log.debug("DNS: Not using dnspython, but dnspython is installed.")
  99         else:
 100             log.debug("DNS: Not using dnspython.")
 101
 102     if not use_ipv6:
 103         log.debug("DNS: Use of IPv6 has been disabled.")
 104
 105     if resolver is None and DNSPYTHON_AVAILABLE and use_dnspython:
 106         resolver = dns.resolver.get_default_resolver()
 107
 108     # An IPv6 literal is allowed to be enclosed in square brackets, but
 109     # the brackets must be stripped in order to process the literal;
 110     # otherwise, things break.
 111     host = host.strip('[]')
 112
 113     try:
 114         # If `host` is an IPv4 literal, we can return it immediately.
 115         ipv4 = socket.inet_aton(host)
 116         yield (host, host, port)
 117     except socket.error:
 118         pass
 119
 120     if use_ipv6:
 121         try:
 122             # Likewise, If `host` is an IPv6 literal, we can return
 123             # it immediately.
 124             if hasattr(socket, 'inet_pton'):
 125                 ipv6 = socket.inet_pton(socket.AF_INET6, host)
 126                 yield (host, host, port)
 127         except (socket.error, ValueError):
 128             pass
 129
 130     # If no service was provided, then we can just do A/AAAA lookups on the
 131     # provided host. Otherwise we need to get an ordered list of hosts to
 132     # resolve based on SRV records.
 133     if not service:
 134         hosts = [(host, port)]
 135     else:
 136         hosts = get_SRV(host, port, service, proto,
 137                         resolver=resolver,
 138                         use_dnspython=use_dnspython)
 139
 140     for host, port in hosts:
 141         results = []
 142         if host == 'localhost':
 143             if use_ipv6:
 144                 results.append((host, '::1', port))
 145             results.append((host, '127.0.0.1', port))
 146         if use_ipv6:
 147             for address in get_AAAA(host, resolver=resolver,
 148                                           use_dnspython=use_dnspython):
 149                 results.append((host, address, port))
 150         for address in get_A(host, resolver=resolver,
 151                                    use_dnspython=use_dnspython):
 152             results.append((host, address, port))
 153
 154         for host, address, port in results:
 155             yield host, address, port
 156
 157
 158 def get_A(host, resolver=None, use_dnspython=True):
 159     """Lookup DNS A records for a given host.
 160
 161     If ``resolver`` is not provided, or is ``None``, then resolution will
 162     be performed using the built-in :mod:`socket` module.
 163
 164     :param     host: The hostname to resolve for A record IPv4 addresses.
 165     :param resolver: Optional DNS resolver object to use for the query.
 166     :param use_dnspython: Optionally control if dnspython is used to make
 167                           the DNS queries instead of the built-in DNS
 168                           library.
 169
 170     :type     host: string
 171     :type resolver: :class:`dns.resolver.Resolver` or ``None``
 172     :type use_dnspython: bool
 173
 174     :return: A list of IPv4 literals.
 175     """
 176     log.debug("DNS: Querying %s for A records." % host)
 177
 178     # If not using dnspython, attempt lookup using the OS level
 179     # getaddrinfo() method.
 180     if resolver is None or not use_dnspython:
 181         try:
 182             recs = socket.getaddrinfo(host, None, socket.AF_INET,
 183                                                   socket.SOCK_STREAM)
 184             return [rec[4][0] for rec in recs]
 185         except socket.gaierror:
 186             log.debug("DNS: Error retreiving A address info for %s." % host)
 187             return []
 188
 189     # Using dnspython:
 190     try:
 191         recs = resolver.query(host, dns.rdatatype.A)
 192         return [rec.to_text() for rec in recs]
 193     except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
 194         log.debug("DNS: No A records for %s" % host)
 195         return []
 196     except dns.exception.Timeout:
 197         log.debug("DNS: A record resolution timed out for %s" % host)
 198         return []
 199     except dns.exception.DNSException as e:
 200         log.debug("DNS: Error querying A records for %s" % host)
 201         log.exception(e)
 202         return []
 203
 204
 205 def get_AAAA(host, resolver=None, use_dnspython=True):
 206     """Lookup DNS AAAA records for a given host.
 207
 208     If ``resolver`` is not provided, or is ``None``, then resolution will
 209     be performed using the built-in :mod:`socket` module.
 210
 211     :param     host: The hostname to resolve for AAAA record IPv6 addresses.
 212     :param resolver: Optional DNS resolver object to use for the query.
 213     :param use_dnspython: Optionally control if dnspython is used to make
 214                           the DNS queries instead of the built-in DNS
 215                           library.
 216
 217     :type     host: string
 218     :type resolver: :class:`dns.resolver.Resolver` or ``None``
 219     :type use_dnspython: bool
 220
 221     :return: A list of IPv6 literals.
 222     """
 223     log.debug("DNS: Querying %s for AAAA records." % host)
 224
 225     # If not using dnspython, attempt lookup using the OS level
 226     # getaddrinfo() method.
 227     if resolver is None or not use_dnspython:
 228         if not socket.has_ipv6:
 229             log.debug("Unable to query %s for AAAA records: IPv6 is not supported", host)
 230             return []
 231         try:
 232             recs = socket.getaddrinfo(host, None, socket.AF_INET6,
 233                                                   socket.SOCK_STREAM)
 234             return [rec[4][0] for rec in recs]
 235         except (OSError, socket.gaierror):
 236             log.debug("DNS: Error retreiving AAAA address " + \
 237                       "info for %s." % host)
 238             return []
 239
 240     # Using dnspython:
 241     try:
 242         recs = resolver.query(host, dns.rdatatype.AAAA)
 243         return [rec.to_text() for rec in recs]
 244     except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
 245         log.debug("DNS: No AAAA records for %s" % host)
 246         return []
 247     except dns.exception.Timeout:
 248         log.debug("DNS: AAAA record resolution timed out for %s" % host)
 249         return []
 250     except dns.exception.DNSException as e:
 251         log.debug("DNS: Error querying AAAA records for %s" % host)
 252         log.exception(e)
 253         return []
 254
 255
 256 def get_SRV(host, port, service, proto='tcp', resolver=None, use_dnspython=True):
 257     """Perform SRV record resolution for a given host.
 258
 259     .. note::
 260
 261         This function requires the use of the ``dnspython`` package. Calling
 262         :func:`get_SRV` without ``dnspython`` will return the provided host
 263         and port without performing any DNS queries.
 264
 265     :param     host: The hostname to resolve.
 266     :param     port: A default port to connect with. SRV records may
 267                      dictate use of a different port.
 268     :param  service: Optional SRV service name without leading underscore.
 269     :param    proto: Optional SRV protocol name without leading underscore.
 270     :param resolver: Optionally provide a DNS resolver object that has
 271                      been custom configured.
 272
 273     :type     host: string
 274     :type     port: int
 275     :type  service: string
 276     :type    proto: string
 277     :type resolver: :class:`dns.resolver.Resolver`
 278
 279     :return: A list of hostname, port pairs in the order dictacted
 280              by SRV priorities and weights.
 281     """
 282     if resolver is None or not use_dnspython:
 283         log.warning("DNS: dnspython not found. Can not use SRV lookup.")
 284         return [(host, port)]
 285
 286     log.debug("DNS: Querying SRV records for %s" % host)
 287     try:
 288         recs = resolver.query('_%s._%s.%s' % (service, proto, host),
 289                               dns.rdatatype.SRV)
 290     except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
 291         log.debug("DNS: No SRV records for %s." % host)
 292         return [(host, port)]
 293     except dns.exception.Timeout:
 294         log.debug("DNS: SRV record resolution timed out for %s." % host)
 295         return [(host, port)]
 296     except dns.exception.DNSException as e:
 297         log.debug("DNS: Error querying SRV records for %s." % host)
 298         log.exception(e)
 299         return [(host, port)]
 300
 301     if len(recs) == 1 and recs[0].target == '.':
 302         return [(host, port)]
 303
 304     answers = {}
 305     for rec in recs:
 306         if rec.priority not in answers:
 307             answers[rec.priority] = []
 308         if rec.weight == 0:
 309             answers[rec.priority].insert(0, rec)
 310         else:
 311             answers[rec.priority].append(rec)
 312
 313     sorted_recs = []
 314     for priority in sorted(answers.keys()):
 315         while answers[priority]:
 316             running_sum = 0
 317             sums = {}
 318             for rec in answers[priority]:
 319                 running_sum += rec.weight
 320                 sums[running_sum] = rec
 321
 322             selected = random.randint(0, running_sum + 1)
 323             for running_sum in sums:
 324                 if running_sum >= selected:
 325                     rec = sums[running_sum]
 326                     host = rec.target.to_text()
 327                     if host.endswith('.'):
 328                         host = host[:-1]
 329                     sorted_recs.append((host, rec.port))
 330                     answers[priority].remove(rec)
 331                     break
 332
 333     return sorted_recs