jshiffer/slixmpp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use aiodns instead of dnspython to query DNS records

Browse Source
This commit is contained in:
mathieui 2014-11-02 17:26:29 +01:00
parent 5b41fb98de
commit 711f8dc6af
2 changed files with 115 additions and 117 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

190
slixmpp/xmlstream/resolver.py
View File

@ -8,6 +8,7 @@
:license: MIT, see LICENSE for more details :license: MIT, see LICENSE for more details
""" """
import asyncio
import socket import socket
import logging import logging
import random import random
@ -16,51 +17,42 @@ import random
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
#: Global flag indicating the availability of the ``dnspython`` package. #: Global flag indicating the availability of the ``aiodns`` package.
#: Installing ``dnspython`` can be done via: #: Installing ``aiodns`` can be done via:
#: #:
#: .. code-block:: sh #: .. code-block:: sh
#: #:
#: pip install dnspython #: pip install aiodns
#: AIODNS_AVAILABLE = False
#: For Python3, installation may require installing from source using
#: the ``python3`` branch:
#:
#: .. code-block:: sh
#:
#: git clone http://github.com/rthalley/dnspython
#: cd dnspython
#: git checkout python3
#: python3 setup.py install
DNSPYTHON_AVAILABLE = False
try: try:
import dns.resolver import aiodns
DNSPYTHON_AVAILABLE = True AIODNS_AVAILABLE = True
except ImportError as e: except ImportError as e:
log.debug("Could not find dnspython package. " + \ log.debug("Could not find aiodns package. " + \
"Not all features will be available") "Not all features will be available")
def default_resolver(): def default_resolver():
"""Return a basic DNS resolver object. """Return a basic DNS resolver object.
:returns: A :class:`dns.resolver.Resolver` object if dnspython :returns: A :class:`aiodns.DNSResolver` object if aiodns
is available. Otherwise, ``None``. is available. Otherwise, ``None``.
""" """
if DNSPYTHON_AVAILABLE: if AIODNS_AVAILABLE:
return dns.resolver.get_default_resolver() return aiodns.DNSResolver(loop=asyncio.get_event_loop())
return None return None
@asyncio.coroutine
def resolve(host, port=None, service=None, proto='tcp', def resolve(host, port=None, service=None, proto='tcp',
resolver=None, use_ipv6=True, use_dnspython=True): resolver=None, use_ipv6=True, use_aiodns=True):
"""Peform DNS resolution for a given hostname. """Peform DNS resolution for a given hostname.
Resolution may perform SRV record lookups if a service and protocol Resolution may perform SRV record lookups if a service and protocol
are specified. The returned addresses will be sorted according to are specified. The returned addresses will be sorted according to
the SRV priorities and weights. the SRV priorities and weights.
If no resolver is provided, the dnspython resolver will be used if If no resolver is provided, the aiodns resolver will be used if
available. Otherwise the built-in socket facilities will be used, available. Otherwise the built-in socket facilities will be used,
but those do not provide SRV support. but those do not provide SRV support.
@ -77,7 +69,7 @@ def resolve(host, port=None, service=None, proto='tcp',
:param use_ipv6: Optionally control the use of IPv6 in situations :param use_ipv6: Optionally control the use of IPv6 in situations
where it is either not available, or performance where it is either not available, or performance
is degraded. Defaults to ``True``. is degraded. Defaults to ``True``.
:param use_dnspython: Optionally control if dnspython is used to make :param use_aiodns: Optionally control if aiodns is used to make
the DNS queries instead of the built-in DNS the DNS queries instead of the built-in DNS
library. library.
@ -85,25 +77,25 @@ def resolve(host, port=None, service=None, proto='tcp',
:type port: int :type port: int
:type service: string :type service: string
:type proto: string :type proto: string
:type resolver: :class:`dns.resolver.Resolver` :type resolver: :class:`aiodns.DNSResolver`
:type use_ipv6: bool :type use_ipv6: bool
:type use_dnspython: bool :type use_aiodns: bool
:return: An iterable of IP address, port pairs in the order :return: An iterable of IP address, port pairs in the order
dictated by SRV priorities and weights, if applicable. dictated by SRV priorities and weights, if applicable.
""" """
if not use_dnspython: if not use_aiodns:
if DNSPYTHON_AVAILABLE: if AIODNS_AVAILABLE:
log.debug("DNS: Not using dnspython, but dnspython is installed.") log.debug("DNS: Not using aiodns, but aiodns is installed.")
else: else:
log.debug("DNS: Not using dnspython.") log.debug("DNS: Not using aiodns.")
if not use_ipv6: if not use_ipv6:
log.debug("DNS: Use of IPv6 has been disabled.") log.debug("DNS: Use of IPv6 has been disabled.")
if resolver is None and DNSPYTHON_AVAILABLE and use_dnspython: if resolver is None and AIODNS_AVAILABLE and use_aiodns:
resolver = dns.resolver.get_default_resolver() resolver = aiodns.DNSResolver(loop=asyncio.get_event_loop())
# An IPv6 literal is allowed to be enclosed in square brackets, but # An IPv6 literal is allowed to be enclosed in square brackets, but
# the brackets must be stripped in order to process the literal; # the brackets must be stripped in order to process the literal;
@ -113,7 +105,7 @@ def resolve(host, port=None, service=None, proto='tcp',
try: try:
# If `host` is an IPv4 literal, we can return it immediately. # If `host` is an IPv4 literal, we can return it immediately.
ipv4 = socket.inet_aton(host) ipv4 = socket.inet_aton(host)
yield (host, host, port) return [(host, host, port)]
except socket.error: except socket.error:
pass pass
@ -123,7 +115,7 @@ def resolve(host, port=None, service=None, proto='tcp',
# it immediately. # it immediately.
if hasattr(socket, 'inet_pton'): if hasattr(socket, 'inet_pton'):
ipv6 = socket.inet_pton(socket.AF_INET6, host) ipv6 = socket.inet_pton(socket.AF_INET6, host)
yield (host, host, port) return [(host, host, port)]
except (socket.error, ValueError): except (socket.error, ValueError):
pass pass
@ -133,29 +125,31 @@ def resolve(host, port=None, service=None, proto='tcp',
if not service: if not service:
hosts = [(host, port)] hosts = [(host, port)]
else: else:
hosts = get_SRV(host, port, service, proto, hosts = yield from get_SRV(host, port, service, proto,
resolver=resolver, resolver=resolver,
use_dnspython=use_dnspython) use_aiodns=use_aiodns)
results = []
for host, port in hosts: for host, port in hosts:
results = []
if host == 'localhost': if host == 'localhost':
if use_ipv6: if use_ipv6:
results.append((host, '::1', port)) results.append((host, '::1', port))
results.append((host, '127.0.0.1', port)) results.append((host, '127.0.0.1', port))
if use_ipv6: if use_ipv6:
for address in get_AAAA(host, resolver=resolver, aaaa = yield from get_AAAA(host, resolver=resolver,
use_dnspython=use_dnspython): use_aiodns=use_aiodns)
for address in aaaa:
results.append((host, address, port)) results.append((host, address, port))
for address in get_A(host, resolver=resolver,
use_dnspython=use_dnspython): a = yield from get_A(host, resolver=resolver,
use_aiodns=use_aiodns)
for address in a:
results.append((host, address, port)) results.append((host, address, port))
for host, address, port in results: return results
yield host, address, port
@asyncio.coroutine
def get_A(host, resolver=None, use_dnspython=True): def get_A(host, resolver=None, use_aiodns=True):
"""Lookup DNS A records for a given host. """Lookup DNS A records for a given host.
If ``resolver`` is not provided, or is ``None``, then resolution will If ``resolver`` is not provided, or is ``None``, then resolution will
@ -163,46 +157,41 @@ def get_A(host, resolver=None, use_dnspython=True):
:param host: The hostname to resolve for A record IPv4 addresses. :param host: The hostname to resolve for A record IPv4 addresses.
:param resolver: Optional DNS resolver object to use for the query. :param resolver: Optional DNS resolver object to use for the query.
:param use_dnspython: Optionally control if dnspython is used to make :param use_aiodns: Optionally control if aiodns is used to make
the DNS queries instead of the built-in DNS the DNS queries instead of the built-in DNS
library. library.
:type host: string :type host: string
:type resolver: :class:`dns.resolver.Resolver` or ``None`` :type resolver: :class:`aiodns.DNSResolver` or ``None``
:type use_dnspython: bool :type use_aiodns: bool
:return: A list of IPv4 literals. :return: A list of IPv4 literals.
""" """
log.debug("DNS: Querying %s for A records." % host) log.debug("DNS: Querying %s for A records." % host)
# If not using dnspython, attempt lookup using the OS level # If not using aiodns, attempt lookup using the OS level
# getaddrinfo() method. # getaddrinfo() method.
if resolver is None or not use_dnspython: if resolver is None or not use_aiodns:
try: try:
recs = socket.getaddrinfo(host, None, socket.AF_INET, recs = socket.getaddrinfo(host, None, socket.AF_INET,
socket.SOCK_STREAM) socket.SOCK_STREAM)
return [rec[4][0] for rec in recs] return [rec[4][0] for rec in recs]
except socket.gaierror: except socket.gaierror:
log.debug("DNS: Error retreiving A address info for %s." % host) log.debug("DNS: Error retrieving A address info for %s." % host)
return [] return []
# Using dnspython: # Using aiodns:
future = resolver.query(host, 'A')
try: try:
recs = resolver.query(host, dns.rdatatype.A) recs = yield from future
return [rec.to_text() for rec in recs] except Exception as e:
except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer): log.debug('DNS: Exception while querying for %s A records: %s', host, e)
log.debug("DNS: No A records for %s" % host) recs = []
return [] return recs
except dns.exception.Timeout:
log.debug("DNS: A record resolution timed out for %s" % host)
return []
except dns.exception.DNSException as e:
log.debug("DNS: Error querying A records for %s" % host)
log.exception(e)
return []
def get_AAAA(host, resolver=None, use_dnspython=True): @asyncio.coroutine
def get_AAAA(host, resolver=None, use_aiodns=True):
"""Lookup DNS AAAA records for a given host. """Lookup DNS AAAA records for a given host.
If ``resolver`` is not provided, or is ``None``, then resolution will If ``resolver`` is not provided, or is ``None``, then resolution will
@ -210,23 +199,23 @@ def get_AAAA(host, resolver=None, use_dnspython=True):
:param host: The hostname to resolve for AAAA record IPv6 addresses. :param host: The hostname to resolve for AAAA record IPv6 addresses.
:param resolver: Optional DNS resolver object to use for the query. :param resolver: Optional DNS resolver object to use for the query.
:param use_dnspython: Optionally control if dnspython is used to make :param use_aiodns: Optionally control if aiodns is used to make
the DNS queries instead of the built-in DNS the DNS queries instead of the built-in DNS
library. library.
:type host: string :type host: string
:type resolver: :class:`dns.resolver.Resolver` or ``None`` :type resolver: :class:`aiodns.DNSResolver` or ``None``
:type use_dnspython: bool :type use_aiodns: bool
:return: A list of IPv6 literals. :return: A list of IPv6 literals.
""" """
log.debug("DNS: Querying %s for AAAA records." % host) log.debug("DNS: Querying %s for AAAA records." % host)
# If not using dnspython, attempt lookup using the OS level # If not using aiodns, attempt lookup using the OS level
# getaddrinfo() method. # getaddrinfo() method.
if resolver is None or not use_dnspython: if resolver is None or not use_aiodns:
if not socket.has_ipv6: if not socket.has_ipv6:
log.debug("Unable to query %s for AAAA records: IPv6 is not supported", host) log.debug("DNS: Unable to query %s for AAAA records: IPv6 is not supported", host)
return [] return []
try: try:
recs = socket.getaddrinfo(host, None, socket.AF_INET6, recs = socket.getaddrinfo(host, None, socket.AF_INET6,
@ -237,29 +226,23 @@ def get_AAAA(host, resolver=None, use_dnspython=True):
"info for %s." % host) "info for %s." % host)
return [] return []
# Using dnspython: # Using aiodns:
future = resolver.query(host, 'AAAA')
try: try:
recs = resolver.query(host, dns.rdatatype.AAAA) recs = yield from future
return [rec.to_text() for rec in recs] except Exception as e:
except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer): log.debug('DNS: Exception while querying for %s AAAA records: %s', host, e)
log.debug("DNS: No AAAA records for %s" % host) recs = []
return [] return recs
except dns.exception.Timeout:
log.debug("DNS: AAAA record resolution timed out for %s" % host)
return []
except dns.exception.DNSException as e:
log.debug("DNS: Error querying AAAA records for %s" % host)
log.exception(e)
return []
@asyncio.coroutine
def get_SRV(host, port, service, proto='tcp', resolver=None, use_dnspython=True): def get_SRV(host, port, service, proto='tcp', resolver=None, use_aiodns=True):
"""Perform SRV record resolution for a given host. """Perform SRV record resolution for a given host.
.. note:: .. note::
This function requires the use of the ``dnspython`` package. Calling This function requires the use of the ``aiodns`` package. Calling
:func:`get_SRV` without ``dnspython`` will return the provided host :func:`get_SRV` without ``aiodns`` will return the provided host
and port without performing any DNS queries. and port without performing any DNS queries.
:param host: The hostname to resolve. :param host: The hostname to resolve.
@ -274,32 +257,23 @@ def get_SRV(host, port, service, proto='tcp', resolver=None, use_dnspython=True)
:type port: int :type port: int
:type service: string :type service: string
:type proto: string :type proto: string
:type resolver: :class:`dns.resolver.Resolver` :type resolver: :class:`aiodns.DNSResolver`
:return: A list of hostname, port pairs in the order dictacted :return: A list of hostname, port pairs in the order dictacted
by SRV priorities and weights. by SRV priorities and weights.
""" """
if resolver is None or not use_dnspython: if resolver is None or not use_aiodns:
log.warning("DNS: dnspython not found. Can not use SRV lookup.") log.warning("DNS: aiodns not found. Can not use SRV lookup.")
return [(host, port)] return [(host, port)]
log.debug("DNS: Querying SRV records for %s" % host) log.debug("DNS: Querying SRV records for %s" % host)
try: try:
recs = resolver.query('_%s._%s.%s' % (service, proto, host), future = resolver.query('_%s._%s.%s' % (service, proto, host),
dns.rdatatype.SRV) 'SRV')
except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer): recs = yield from future
log.debug("DNS: No SRV records for %s." % host) except Exception as e:
return [(host, port)] log.debug('DNS: Exception while querying for %s SRV records: %s', host, e)
except dns.exception.Timeout: return []
log.debug("DNS: SRV record resolution timed out for %s." % host)
return [(host, port)]
except dns.exception.DNSException as e:
log.debug("DNS: Error querying SRV records for %s." % host)
log.exception(e)
return [(host, port)]
if len(recs) == 1 and recs[0].target == '.':
return [(host, port)]
answers = {} answers = {}
for rec in recs: for rec in recs:
@ -323,7 +297,7 @@ def get_SRV(host, port, service, proto='tcp', resolver=None, use_dnspython=True)
for running_sum in sums: for running_sum in sums:
if running_sum >= selected: if running_sum >= selected:
rec = sums[running_sum] rec = sums[running_sum]
host = rec.target.to_text() host = rec.host
if host.endswith('.'): if host.endswith('.'):
host = host[:-1] host = host[:-1]
sorted_recs.append((host, rec.port)) sorted_recs.append((host, rec.port))

42
slixmpp/xmlstream/xmlstream.py
View File

@ -162,7 +162,7 @@ class XMLStream(object):
#: If set to ``True``, allow using the ``dnspython`` DNS library #: If set to ``True``, allow using the ``dnspython`` DNS library
#: if available. If set to ``False``, the builtin DNS resolver #: if available. If set to ``False``, the builtin DNS resolver
#: will be used, even if ``dnspython`` is installed. #: will be used, even if ``dnspython`` is installed.
self.use_dnspython = True self.use_aiodns = True
#: Use CDATA for escaping instead of XML entities. Defaults #: Use CDATA for escaping instead of XML entities. Defaults
#: to ``False``. #: to ``False``.
@ -287,13 +287,32 @@ class XMLStream(object):
def _connect_routine(self): def _connect_routine(self):
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
self.event_when_connected = "connected" self.event_when_connected = "connected"
try:
record = yield from self.pick_dns_answer(self.default_domain)
except StopIteration:
# No more DNS records to try
self.dns_answers = None
return
else:
if record:
host, address, port = record
self._service_name = host
else:
self.event('connection_failed',
'No DNS record available for %s' % self.default_domain)
self.dns_answers = None
return
try: try:
yield from loop.create_connection(lambda: self, yield from loop.create_connection(lambda: self,
self.address[0], address,
self.address[1], port,
ssl=self.use_ssl) ssl=self.use_ssl)
except OSError as e: except OSError as e:
log.debug('Connection failed: %s', e)
self.event("connection_failed", e) self.event("connection_failed", e)
asyncio.async(self._connect_routine())
def process(self, timeout=None): def process(self, timeout=None):
"""Process all the available XMPP events (receiving or sending data on the """Process all the available XMPP events (receiving or sending data on the
@ -578,6 +597,7 @@ class XMLStream(object):
idx += 1 idx += 1
return False return False
@asyncio.coroutine
def get_dns_records(self, domain, port=None): def get_dns_records(self, domain, port=None):
"""Get the DNS records for a domain. """Get the DNS records for a domain.
@ -590,11 +610,14 @@ class XMLStream(object):
resolver = default_resolver() resolver = default_resolver()
self.configure_dns(resolver, domain=domain, port=port) self.configure_dns(resolver, domain=domain, port=port)
return resolve(domain, port, service=self.dns_service, result = yield from resolve(domain, port,
resolver=resolver, service=self.dns_service,
use_ipv6=self.use_ipv6, resolver=resolver,
use_dnspython=self.use_dnspython) use_ipv6=self.use_ipv6,
use_aiodns=self.use_aiodns)
return result
@asyncio.coroutine
def pick_dns_answer(self, domain, port=None): def pick_dns_answer(self, domain, port=None):
"""Pick a server and port from DNS answers. """Pick a server and port from DNS answers.
@ -604,8 +627,9 @@ class XMLStream(object):
:param domain: The domain in question. :param domain: The domain in question.
:param port: If the results don't include a port, use this one. :param port: If the results don't include a port, use this one.
""" """
if not self.dns_answers: if self.dns_answers is None:
self.dns_answers = self.get_dns_records(domain, port) dns_records = yield from self.get_dns_records(domain, port)
self.dns_answers = iter(dns_records)
return next(self.dns_answers) return next(self.dns_answers)