Use aiodns instead of dnspython to query DNS records
This commit is contained in:
mathieui
parent
5b41fb98de
commit
711f8dc6af
@ -8,6 +8,7 @@
|
||||
:license: MIT, see LICENSE for more details
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import socket
|
||||
import logging
|
||||
import random
|
||||
@ -16,51 +17,42 @@ import random
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
#: Global flag indicating the availability of the ``dnspython`` package.
|
||||
#: Installing ``dnspython`` can be done via:
|
||||
#: Global flag indicating the availability of the ``aiodns`` package.
|
||||
#: Installing ``aiodns`` can be done via:
|
||||
#:
|
||||
#: .. code-block:: sh
|
||||
#:
|
||||
#: pip install dnspython
|
||||
#:
|
||||
#: 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
|
||||
#: pip install aiodns
|
||||
AIODNS_AVAILABLE = False
|
||||
try:
|
||||
import dns.resolver
|
||||
DNSPYTHON_AVAILABLE = True
|
||||
import aiodns
|
||||
AIODNS_AVAILABLE = True
|
||||
except ImportError as e:
|
||||
log.debug("Could not find dnspython package. " + \
|
||||
log.debug("Could not find aiodns package. " + \
|
||||
"Not all features will be available")
|
||||
|
||||
|
||||
def default_resolver():
|
||||
"""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``.
|
||||
"""
|
||||
if DNSPYTHON_AVAILABLE:
|
||||
return dns.resolver.get_default_resolver()
|
||||
if AIODNS_AVAILABLE:
|
||||
return aiodns.DNSResolver(loop=asyncio.get_event_loop())
|
||||
return None
|
||||
|
||||
|
||||
@asyncio.coroutine
|
||||
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.
|
||||
|
||||
Resolution may perform SRV record lookups if a service and protocol
|
||||
are specified. The returned addresses will be sorted according to
|
||||
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,
|
||||
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
|
||||
where it is either not available, or performance
|
||||
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
|
||||
library.
|
||||
|
||||
@ -85,25 +77,25 @@ def resolve(host, port=None, service=None, proto='tcp',
|
||||
:type port: int
|
||||
:type service: string
|
||||
:type proto: string
|
||||
:type resolver: :class:`dns.resolver.Resolver`
|
||||
:type resolver: :class:`aiodns.DNSResolver`
|
||||
:type use_ipv6: bool
|
||||
:type use_dnspython: bool
|
||||
:type use_aiodns: bool
|
||||
|
||||
:return: An iterable of IP address, port pairs in the order
|
||||
dictated by SRV priorities and weights, if applicable.
|
||||
"""
|
||||
|
||||
if not use_dnspython:
|
||||
if DNSPYTHON_AVAILABLE:
|
||||
log.debug("DNS: Not using dnspython, but dnspython is installed.")
|
||||
if not use_aiodns:
|
||||
if AIODNS_AVAILABLE:
|
||||
log.debug("DNS: Not using aiodns, but aiodns is installed.")
|
||||
else:
|
||||
log.debug("DNS: Not using dnspython.")
|
||||
log.debug("DNS: Not using aiodns.")
|
||||
|
||||
if not use_ipv6:
|
||||
log.debug("DNS: Use of IPv6 has been disabled.")
|
||||
|
||||
if resolver is None and DNSPYTHON_AVAILABLE and use_dnspython:
|
||||
resolver = dns.resolver.get_default_resolver()
|
||||
if resolver is None and AIODNS_AVAILABLE and use_aiodns:
|
||||
resolver = aiodns.DNSResolver(loop=asyncio.get_event_loop())
|
||||
|
||||
# An IPv6 literal is allowed to be enclosed in square brackets, but
|
||||
# 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:
|
||||
# If `host` is an IPv4 literal, we can return it immediately.
|
||||
ipv4 = socket.inet_aton(host)
|
||||
yield (host, host, port)
|
||||
return [(host, host, port)]
|
||||
except socket.error:
|
||||
pass
|
||||
|
||||
@ -123,7 +115,7 @@ def resolve(host, port=None, service=None, proto='tcp',
|
||||
# it immediately.
|
||||
if hasattr(socket, 'inet_pton'):
|
||||
ipv6 = socket.inet_pton(socket.AF_INET6, host)
|
||||
yield (host, host, port)
|
||||
return [(host, host, port)]
|
||||
except (socket.error, ValueError):
|
||||
pass
|
||||
|
||||
@ -133,29 +125,31 @@ def resolve(host, port=None, service=None, proto='tcp',
|
||||
if not service:
|
||||
hosts = [(host, port)]
|
||||
else:
|
||||
hosts = get_SRV(host, port, service, proto,
|
||||
resolver=resolver,
|
||||
use_dnspython=use_dnspython)
|
||||
|
||||
hosts = yield from get_SRV(host, port, service, proto,
|
||||
resolver=resolver,
|
||||
use_aiodns=use_aiodns)
|
||||
results = []
|
||||
for host, port in hosts:
|
||||
results = []
|
||||
if host == 'localhost':
|
||||
if use_ipv6:
|
||||
results.append((host, '::1', port))
|
||||
results.append((host, '127.0.0.1', port))
|
||||
|
||||
if use_ipv6:
|
||||
for address in get_AAAA(host, resolver=resolver,
|
||||
use_dnspython=use_dnspython):
|
||||
aaaa = yield from get_AAAA(host, resolver=resolver,
|
||||
use_aiodns=use_aiodns)
|
||||
for address in aaaa:
|
||||
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))
|
||||
|
||||
for host, address, port in results:
|
||||
yield host, address, port
|
||||
return results
|
||||
|
||||
|
||||
def get_A(host, resolver=None, use_dnspython=True):
|
||||
@asyncio.coroutine
|
||||
def get_A(host, resolver=None, use_aiodns=True):
|
||||
"""Lookup DNS A records for a given host.
|
||||
|
||||
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 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
|
||||
library.
|
||||
|
||||
:type host: string
|
||||
:type resolver: :class:`dns.resolver.Resolver` or ``None``
|
||||
:type use_dnspython: bool
|
||||
:type resolver: :class:`aiodns.DNSResolver` or ``None``
|
||||
:type use_aiodns: bool
|
||||
|
||||
:return: A list of IPv4 literals.
|
||||
"""
|
||||
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.
|
||||
if resolver is None or not use_dnspython:
|
||||
if resolver is None or not use_aiodns:
|
||||
try:
|
||||
recs = socket.getaddrinfo(host, None, socket.AF_INET,
|
||||
socket.SOCK_STREAM)
|
||||
return [rec[4][0] for rec in recs]
|
||||
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 []
|
||||
|
||||
# Using dnspython:
|
||||
# Using aiodns:
|
||||
future = resolver.query(host, 'A')
|
||||
try:
|
||||
recs = resolver.query(host, dns.rdatatype.A)
|
||||
return [rec.to_text() for rec in recs]
|
||||
except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
|
||||
log.debug("DNS: No A records for %s" % host)
|
||||
return []
|
||||
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 []
|
||||
recs = yield from future
|
||||
except Exception as e:
|
||||
log.debug('DNS: Exception while querying for %s A records: %s', host, e)
|
||||
recs = []
|
||||
return recs
|
||||
|
||||
|
||||
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.
|
||||
|
||||
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 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
|
||||
library.
|
||||
|
||||
:type host: string
|
||||
:type resolver: :class:`dns.resolver.Resolver` or ``None``
|
||||
:type use_dnspython: bool
|
||||
:type resolver: :class:`aiodns.DNSResolver` or ``None``
|
||||
:type use_aiodns: bool
|
||||
|
||||
:return: A list of IPv6 literals.
|
||||
"""
|
||||
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.
|
||||
if resolver is None or not use_dnspython:
|
||||
if resolver is None or not use_aiodns:
|
||||
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 []
|
||||
try:
|
||||
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)
|
||||
return []
|
||||
|
||||
# Using dnspython:
|
||||
# Using aiodns:
|
||||
future = resolver.query(host, 'AAAA')
|
||||
try:
|
||||
recs = resolver.query(host, dns.rdatatype.AAAA)
|
||||
return [rec.to_text() for rec in recs]
|
||||
except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
|
||||
log.debug("DNS: No AAAA records for %s" % host)
|
||||
return []
|
||||
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 []
|
||||
recs = yield from future
|
||||
except Exception as e:
|
||||
log.debug('DNS: Exception while querying for %s AAAA records: %s', host, e)
|
||||
recs = []
|
||||
return recs
|
||||
|
||||
|
||||
def get_SRV(host, port, service, proto='tcp', resolver=None, use_dnspython=True):
|
||||
@asyncio.coroutine
|
||||
def get_SRV(host, port, service, proto='tcp', resolver=None, use_aiodns=True):
|
||||
"""Perform SRV record resolution for a given host.
|
||||
|
||||
.. note::
|
||||
|
||||
This function requires the use of the ``dnspython`` package. Calling
|
||||
:func:`get_SRV` without ``dnspython`` will return the provided host
|
||||
This function requires the use of the ``aiodns`` package. Calling
|
||||
:func:`get_SRV` without ``aiodns`` will return the provided host
|
||||
and port without performing any DNS queries.
|
||||
|
||||
: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 service: 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
|
||||
by SRV priorities and weights.
|
||||
"""
|
||||
if resolver is None or not use_dnspython:
|
||||
log.warning("DNS: dnspython not found. Can not use SRV lookup.")
|
||||
if resolver is None or not use_aiodns:
|
||||
log.warning("DNS: aiodns not found. Can not use SRV lookup.")
|
||||
return [(host, port)]
|
||||
|
||||
log.debug("DNS: Querying SRV records for %s" % host)
|
||||
try:
|
||||
recs = resolver.query('_%s._%s.%s' % (service, proto, host),
|
||||
dns.rdatatype.SRV)
|
||||
except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
|
||||
log.debug("DNS: No SRV records for %s." % host)
|
||||
return [(host, port)]
|
||||
except dns.exception.Timeout:
|
||||
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)]
|
||||
future = resolver.query('_%s._%s.%s' % (service, proto, host),
|
||||
'SRV')
|
||||
recs = yield from future
|
||||
except Exception as e:
|
||||
log.debug('DNS: Exception while querying for %s SRV records: %s', host, e)
|
||||
return []
|
||||
|
||||
answers = {}
|
||||
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:
|
||||
if running_sum >= selected:
|
||||
rec = sums[running_sum]
|
||||
host = rec.target.to_text()
|
||||
host = rec.host
|
||||
if host.endswith('.'):
|
||||
host = host[:-1]
|
||||
sorted_recs.append((host, rec.port))
|
||||
|
||||
@ -162,7 +162,7 @@ class XMLStream(object):
|
||||
#: If set to ``True``, allow using the ``dnspython`` DNS library
|
||||
#: if available. If set to ``False``, the builtin DNS resolver
|
||||
#: 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
|
||||
#: to ``False``.
|
||||
@ -287,13 +287,32 @@ class XMLStream(object):
|
||||
def _connect_routine(self):
|
||||
loop = asyncio.get_event_loop()
|
||||
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:
|
||||
yield from loop.create_connection(lambda: self,
|
||||
self.address[0],
|
||||
self.address[1],
|
||||
address,
|
||||
port,
|
||||
ssl=self.use_ssl)
|
||||
except OSError as e:
|
||||
log.debug('Connection failed: %s', e)
|
||||
self.event("connection_failed", e)
|
||||
asyncio.async(self._connect_routine())
|
||||
|
||||
def process(self, timeout=None):
|
||||
"""Process all the available XMPP events (receiving or sending data on the
|
||||
@ -578,6 +597,7 @@ class XMLStream(object):
|
||||
idx += 1
|
||||
return False
|
||||
|
||||
@asyncio.coroutine
|
||||
def get_dns_records(self, domain, port=None):
|
||||
"""Get the DNS records for a domain.
|
||||
|
||||
@ -590,11 +610,14 @@ class XMLStream(object):
|
||||
resolver = default_resolver()
|
||||
self.configure_dns(resolver, domain=domain, port=port)
|
||||
|
||||
return resolve(domain, port, service=self.dns_service,
|
||||
resolver=resolver,
|
||||
use_ipv6=self.use_ipv6,
|
||||
use_dnspython=self.use_dnspython)
|
||||
result = yield from resolve(domain, port,
|
||||
service=self.dns_service,
|
||||
resolver=resolver,
|
||||
use_ipv6=self.use_ipv6,
|
||||
use_aiodns=self.use_aiodns)
|
||||
return result
|
||||
|
||||
@asyncio.coroutine
|
||||
def pick_dns_answer(self, domain, port=None):
|
||||
"""Pick a server and port from DNS answers.
|
||||
|
||||
@ -604,8 +627,9 @@ class XMLStream(object):
|
||||
:param domain: The domain in question.
|
||||
:param port: If the results don't include a port, use this one.
|
||||
"""
|
||||
if not self.dns_answers:
|
||||
self.dns_answers = self.get_dns_records(domain, port)
|
||||
if self.dns_answers is None:
|
||||
dns_records = yield from self.get_dns_records(domain, port)
|
||||
self.dns_answers = iter(dns_records)
|
||||
|
||||
return next(self.dns_answers)
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user