Spaces:
Build error
Build error
| import asyncio | |
| import socket | |
| from typing import Any, Dict, List, Optional, Type, Union | |
| 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): | |
| """Threaded resolver. | |
| Uses an Executor for synchronous getaddrinfo() calls. | |
| concurrent.futures.ThreadPoolExecutor is used by default. | |
| """ | |
| 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: | |
| if len(address) < 3: | |
| # IPv6 is not supported by Python build, | |
| # or IPv6 is not enabled in the host | |
| continue | |
| if address[3]: | |
| # 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] | |
| else: # IPv4 | |
| assert family == socket.AF_INET | |
| host, port = address # type: ignore[misc] | |
| 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: | |
| self._resolver.cancel() | |
| _DefaultType = Type[Union[AsyncResolver, ThreadedResolver]] | |
| DefaultResolver: _DefaultType = AsyncResolver if aiodns_default else ThreadedResolver | |