123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149 |
- import asyncio
- import socket
- from typing import Any, Dict, List, Optional
- from .abc import AbstractResolver
- from .helpers import get_running_loop
- __all__ = ("ThreadedResolver", "AsyncResolver", "DefaultResolver")
- try:
- import aiodns
- # aiodns_default = hasattr(aiodns.DNSResolver, 'gethostbyname')
- except ImportError: # pragma: no cover
- aiodns = None
- aiodns_default = False
- class ThreadedResolver(AbstractResolver):
- """Use Executor for synchronous getaddrinfo() calls, which defaults to
- concurrent.futures.ThreadPoolExecutor.
- """
- def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
- self._loop = get_running_loop(loop)
- async def resolve(
- self, hostname: str, port: int = 0, family: int = socket.AF_INET
- ) -> List[Dict[str, Any]]:
- infos = await self._loop.getaddrinfo(
- hostname,
- port,
- type=socket.SOCK_STREAM,
- family=family,
- flags=socket.AI_ADDRCONFIG,
- )
- hosts = []
- for family, _, proto, _, address in infos:
- if family == socket.AF_INET6 and address[3]: # type: ignore
- # This is essential for link-local IPv6 addresses.
- # LL IPv6 is a VERY rare case. Strictly speaking, we should use
- # getnameinfo() unconditionally, but performance makes sense.
- host, _port = socket.getnameinfo(
- address, socket.NI_NUMERICHOST | socket.NI_NUMERICSERV
- )
- port = int(_port)
- else:
- host, port = address[:2]
- hosts.append(
- {
- "hostname": hostname,
- "host": host,
- "port": port,
- "family": family,
- "proto": proto,
- "flags": socket.AI_NUMERICHOST | socket.AI_NUMERICSERV,
- }
- )
- return hosts
- async def close(self) -> None:
- pass
- class AsyncResolver(AbstractResolver):
- """Use the `aiodns` package to make asynchronous DNS lookups"""
- def __init__(
- self,
- loop: Optional[asyncio.AbstractEventLoop] = None,
- *args: Any,
- **kwargs: Any
- ) -> None:
- if aiodns is None:
- raise RuntimeError("Resolver requires aiodns library")
- self._loop = get_running_loop(loop)
- self._resolver = aiodns.DNSResolver(*args, loop=loop, **kwargs)
- if not hasattr(self._resolver, "gethostbyname"):
- # aiodns 1.1 is not available, fallback to DNSResolver.query
- self.resolve = self._resolve_with_query # type: ignore
- async def resolve(
- self, host: str, port: int = 0, family: int = socket.AF_INET
- ) -> List[Dict[str, Any]]:
- try:
- resp = await self._resolver.gethostbyname(host, family)
- except aiodns.error.DNSError as exc:
- msg = exc.args[1] if len(exc.args) >= 1 else "DNS lookup failed"
- raise OSError(msg) from exc
- hosts = []
- for address in resp.addresses:
- hosts.append(
- {
- "hostname": host,
- "host": address,
- "port": port,
- "family": family,
- "proto": 0,
- "flags": socket.AI_NUMERICHOST | socket.AI_NUMERICSERV,
- }
- )
- if not hosts:
- raise OSError("DNS lookup failed")
- return hosts
- async def _resolve_with_query(
- self, host: str, port: int = 0, family: int = socket.AF_INET
- ) -> List[Dict[str, Any]]:
- if family == socket.AF_INET6:
- qtype = "AAAA"
- else:
- qtype = "A"
- try:
- resp = await self._resolver.query(host, qtype)
- except aiodns.error.DNSError as exc:
- msg = exc.args[1] if len(exc.args) >= 1 else "DNS lookup failed"
- raise OSError(msg) from exc
- hosts = []
- for rr in resp:
- hosts.append(
- {
- "hostname": host,
- "host": rr.host,
- "port": port,
- "family": family,
- "proto": 0,
- "flags": socket.AI_NUMERICHOST,
- }
- )
- if not hosts:
- raise OSError("DNS lookup failed")
- return hosts
- async def close(self) -> None:
- return self._resolver.cancel()
- DefaultResolver = AsyncResolver if aiodns_default else ThreadedResolver
|