515 lines
17 KiB
Python
515 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import functools
|
|
import logging
|
|
import socket
|
|
import sys
|
|
import warnings
|
|
import weakref
|
|
from collections.abc import Callable, Iterable, Sequence
|
|
from types import TracebackType
|
|
from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload
|
|
|
|
import pycares
|
|
|
|
from . import error
|
|
from .compat import (
|
|
AresHostResult,
|
|
AresQueryAAAAResult,
|
|
AresQueryAResult,
|
|
AresQueryCAAResult,
|
|
AresQueryCNAMEResult,
|
|
AresQueryMXResult,
|
|
AresQueryNAPTRResult,
|
|
AresQueryNSResult,
|
|
AresQueryPTRResult,
|
|
AresQuerySOAResult,
|
|
AresQuerySRVResult,
|
|
AresQueryTXTResult,
|
|
QueryResult,
|
|
convert_result,
|
|
)
|
|
|
|
__version__ = '4.0.0'
|
|
|
|
__all__ = (
|
|
'DNSResolver',
|
|
'error',
|
|
)
|
|
|
|
_T = TypeVar('_T')
|
|
|
|
WINDOWS_SELECTOR_ERR_MSG = (
|
|
'aiodns needs a SelectorEventLoop on Windows. See more: '
|
|
'https://github.com/aio-libs/aiodns#note-for-windows-users'
|
|
)
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
|
|
query_type_map = {
|
|
'A': pycares.QUERY_TYPE_A,
|
|
'AAAA': pycares.QUERY_TYPE_AAAA,
|
|
'ANY': pycares.QUERY_TYPE_ANY,
|
|
'CAA': pycares.QUERY_TYPE_CAA,
|
|
'CNAME': pycares.QUERY_TYPE_CNAME,
|
|
'MX': pycares.QUERY_TYPE_MX,
|
|
'NAPTR': pycares.QUERY_TYPE_NAPTR,
|
|
'NS': pycares.QUERY_TYPE_NS,
|
|
'PTR': pycares.QUERY_TYPE_PTR,
|
|
'SOA': pycares.QUERY_TYPE_SOA,
|
|
'SRV': pycares.QUERY_TYPE_SRV,
|
|
'TXT': pycares.QUERY_TYPE_TXT,
|
|
}
|
|
|
|
query_class_map = {
|
|
'IN': pycares.QUERY_CLASS_IN,
|
|
'CHAOS': pycares.QUERY_CLASS_CHAOS,
|
|
'HS': pycares.QUERY_CLASS_HS,
|
|
'NONE': pycares.QUERY_CLASS_NONE,
|
|
'ANY': pycares.QUERY_CLASS_ANY,
|
|
}
|
|
|
|
|
|
class DNSResolver:
|
|
def __init__(
|
|
self,
|
|
nameservers: Sequence[str] | None = None,
|
|
loop: asyncio.AbstractEventLoop | None = None,
|
|
**kwargs: Any,
|
|
) -> None: # TODO(PY311): Use Unpack for kwargs.
|
|
self._closed = True
|
|
self.loop = loop or asyncio.get_event_loop()
|
|
if TYPE_CHECKING:
|
|
assert self.loop is not None
|
|
kwargs.pop('sock_state_cb', None)
|
|
timeout = kwargs.pop('timeout', None)
|
|
self._timeout = timeout
|
|
self._event_thread, self._channel = self._make_channel(**kwargs)
|
|
if nameservers:
|
|
self.nameservers = nameservers
|
|
self._read_fds: set[int] = set()
|
|
self._write_fds: set[int] = set()
|
|
self._timer: asyncio.TimerHandle | None = None
|
|
self._closed = False
|
|
|
|
def _make_channel(self, **kwargs: Any) -> tuple[bool, pycares.Channel]:
|
|
# pycares 5+ uses event_thread by default when sock_state_cb
|
|
# is not provided
|
|
try:
|
|
return True, pycares.Channel(timeout=self._timeout, **kwargs)
|
|
except pycares.AresError as e:
|
|
if sys.platform == 'linux':
|
|
_LOGGER.warning(
|
|
'Failed to create DNS resolver channel with automatic '
|
|
'monitoring of resolver configuration changes. This '
|
|
'usually means the system ran out of inotify watches. '
|
|
'Falling back to socket state callback. Consider '
|
|
'increasing the system inotify watch limit: %s',
|
|
e,
|
|
)
|
|
else:
|
|
_LOGGER.warning(
|
|
'Failed to create DNS resolver channel with automatic '
|
|
'monitoring of resolver configuration changes. '
|
|
'Falling back to socket state callback: %s',
|
|
e,
|
|
)
|
|
# Fall back to sock_state_cb (needs SelectorEventLoop on Windows)
|
|
if sys.platform == 'win32' and not isinstance(
|
|
self.loop, asyncio.SelectorEventLoop
|
|
):
|
|
try:
|
|
import winloop
|
|
|
|
if not isinstance(self.loop, winloop.Loop):
|
|
raise RuntimeError(WINDOWS_SELECTOR_ERR_MSG)
|
|
except ModuleNotFoundError as ex:
|
|
raise RuntimeError(WINDOWS_SELECTOR_ERR_MSG) from ex
|
|
# Use weak reference for deterministic cleanup. Without it there's a
|
|
# reference cycle (DNSResolver -> _channel -> callback -> DNSResolver).
|
|
# Python 3.4+ can handle cycles with __del__, but weak ref ensures
|
|
# cleanup happens immediately when last reference is dropped.
|
|
weak_self = weakref.ref(self)
|
|
|
|
def sock_state_cb_wrapper(
|
|
fd: int, readable: bool, writable: bool
|
|
) -> None:
|
|
this = weak_self()
|
|
if this is not None:
|
|
this._sock_state_cb(fd, readable, writable)
|
|
|
|
return False, pycares.Channel(
|
|
sock_state_cb=sock_state_cb_wrapper,
|
|
timeout=self._timeout,
|
|
**kwargs,
|
|
)
|
|
|
|
@property
|
|
def nameservers(self) -> Sequence[str]:
|
|
# pycares 5.x returns servers with port (e.g., '8.8.8.8:53')
|
|
# Strip port for backward compatibility with pycares 4.x
|
|
return [s.rsplit(':', 1)[0] for s in self._channel.servers]
|
|
|
|
@nameservers.setter
|
|
def nameservers(self, value: Iterable[str | bytes]) -> None:
|
|
self._channel.servers = value
|
|
|
|
def _callback(
|
|
self, fut: asyncio.Future[_T], result: _T, errorno: int | None
|
|
) -> None:
|
|
if fut.cancelled():
|
|
return
|
|
if errorno is not None:
|
|
fut.set_exception(
|
|
error.DNSError(errorno, pycares.errno.strerror(errorno))
|
|
)
|
|
else:
|
|
fut.set_result(result)
|
|
|
|
def _get_future_callback(
|
|
self,
|
|
) -> tuple[asyncio.Future[_T], Callable[[_T, int | None], None]]:
|
|
"""Return a future and a callback to set the result of the future."""
|
|
cb: Callable[[_T, int | None], None]
|
|
future: asyncio.Future[_T] = self.loop.create_future()
|
|
if self._event_thread:
|
|
cb = functools.partial( # type: ignore[assignment]
|
|
self.loop.call_soon_threadsafe,
|
|
self._callback, # type: ignore[arg-type]
|
|
future,
|
|
)
|
|
else:
|
|
cb = functools.partial(self._callback, future)
|
|
return future, cb
|
|
|
|
def _query_callback(
|
|
self,
|
|
fut: asyncio.Future[QueryResult],
|
|
qtype: int,
|
|
result: pycares.DNSResult,
|
|
errorno: int | None,
|
|
) -> None:
|
|
"""Callback for query that converts results to compatible format."""
|
|
if fut.cancelled():
|
|
return
|
|
if errorno is not None:
|
|
fut.set_exception(
|
|
error.DNSError(errorno, pycares.errno.strerror(errorno))
|
|
)
|
|
else:
|
|
fut.set_result(convert_result(result, qtype))
|
|
|
|
def _get_query_future_callback(
|
|
self, qtype: int
|
|
) -> tuple[asyncio.Future[QueryResult], Callable[..., None]]:
|
|
"""Return a future and callback for query with result conversion."""
|
|
future: asyncio.Future[QueryResult] = self.loop.create_future()
|
|
cb: Callable[..., None]
|
|
if self._event_thread:
|
|
cb = functools.partial( # type: ignore[assignment]
|
|
self.loop.call_soon_threadsafe,
|
|
self._query_callback, # type: ignore[arg-type]
|
|
future,
|
|
qtype,
|
|
)
|
|
else:
|
|
cb = functools.partial(self._query_callback, future, qtype)
|
|
return future, cb
|
|
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['A'], qclass: str | None = ...
|
|
) -> asyncio.Future[list[AresQueryAResult]]: ...
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['AAAA'], qclass: str | None = ...
|
|
) -> asyncio.Future[list[AresQueryAAAAResult]]: ...
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['CAA'], qclass: str | None = ...
|
|
) -> asyncio.Future[list[AresQueryCAAResult]]: ...
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['CNAME'], qclass: str | None = ...
|
|
) -> asyncio.Future[AresQueryCNAMEResult]: ...
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['MX'], qclass: str | None = ...
|
|
) -> asyncio.Future[list[AresQueryMXResult]]: ...
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['NAPTR'], qclass: str | None = ...
|
|
) -> asyncio.Future[list[AresQueryNAPTRResult]]: ...
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['NS'], qclass: str | None = ...
|
|
) -> asyncio.Future[list[AresQueryNSResult]]: ...
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['PTR'], qclass: str | None = ...
|
|
) -> asyncio.Future[AresQueryPTRResult]: ...
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['SOA'], qclass: str | None = ...
|
|
) -> asyncio.Future[AresQuerySOAResult]: ...
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['SRV'], qclass: str | None = ...
|
|
) -> asyncio.Future[list[AresQuerySRVResult]]: ...
|
|
@overload
|
|
def query(
|
|
self, host: str, qtype: Literal['TXT'], qclass: str | None = ...
|
|
) -> asyncio.Future[list[AresQueryTXTResult]]: ...
|
|
|
|
def query(
|
|
self, host: str, qtype: str, qclass: str | None = None
|
|
) -> asyncio.Future[list[Any]] | asyncio.Future[Any]:
|
|
"""Query DNS records (deprecated, use query_dns instead)."""
|
|
warnings.warn(
|
|
'query() is deprecated, use query_dns() instead',
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
try:
|
|
qtype_int = query_type_map[qtype]
|
|
except KeyError as e:
|
|
raise ValueError(f'invalid query type: {qtype}') from e
|
|
qclass_int: int | None = None
|
|
if qclass is not None:
|
|
try:
|
|
qclass_int = query_class_map[qclass]
|
|
except KeyError as e:
|
|
raise ValueError(f'invalid query class: {qclass}') from e
|
|
|
|
fut, cb = self._get_query_future_callback(qtype_int)
|
|
if qclass_int is not None:
|
|
self._channel.query(
|
|
host, qtype_int, query_class=qclass_int, callback=cb
|
|
)
|
|
else:
|
|
self._channel.query(host, qtype_int, callback=cb)
|
|
return fut
|
|
|
|
def query_dns(
|
|
self, host: str, qtype: str, qclass: str | None = None
|
|
) -> asyncio.Future[pycares.DNSResult]:
|
|
"""Query DNS records, returning native pycares 5.x DNSResult."""
|
|
try:
|
|
qtype_int = query_type_map[qtype]
|
|
except KeyError as e:
|
|
raise ValueError(f'invalid query type: {qtype}') from e
|
|
qclass_int: int | None = None
|
|
if qclass is not None:
|
|
try:
|
|
qclass_int = query_class_map[qclass]
|
|
except KeyError as e:
|
|
raise ValueError(f'invalid query class: {qclass}') from e
|
|
|
|
fut: asyncio.Future[pycares.DNSResult]
|
|
fut, cb = self._get_future_callback()
|
|
if qclass_int is not None:
|
|
self._channel.query(
|
|
host, qtype_int, query_class=qclass_int, callback=cb
|
|
)
|
|
else:
|
|
self._channel.query(host, qtype_int, callback=cb)
|
|
return fut
|
|
|
|
def _gethostbyname_callback(
|
|
self,
|
|
fut: asyncio.Future[AresHostResult],
|
|
host: str,
|
|
result: pycares.AddrInfoResult | None,
|
|
errorno: int | None,
|
|
) -> None:
|
|
"""Callback for gethostbyname that converts AddrInfoResult."""
|
|
if fut.cancelled():
|
|
return
|
|
if errorno is not None:
|
|
fut.set_exception(
|
|
error.DNSError(errorno, pycares.errno.strerror(errorno))
|
|
)
|
|
else:
|
|
assert result is not None # noqa: S101
|
|
# node.addr is (address_bytes, port) - extract and decode
|
|
addresses = [node.addr[0].decode() for node in result.nodes]
|
|
# Get canonical name from cnames if available
|
|
name = result.cnames[0].name if result.cnames else host
|
|
fut.set_result(
|
|
AresHostResult(name=name, aliases=[], addresses=addresses)
|
|
)
|
|
|
|
def gethostbyname(
|
|
self, host: str, family: socket.AddressFamily
|
|
) -> asyncio.Future[AresHostResult]:
|
|
"""
|
|
Resolve hostname to addresses.
|
|
|
|
Deprecated: Use getaddrinfo() instead. This is implemented using
|
|
getaddrinfo as pycares 5.x removed the gethostbyname method.
|
|
"""
|
|
warnings.warn(
|
|
'gethostbyname() is deprecated, use getaddrinfo() instead',
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
fut: asyncio.Future[AresHostResult] = self.loop.create_future()
|
|
cb: Callable[..., None]
|
|
if self._event_thread:
|
|
cb = functools.partial( # type: ignore[assignment]
|
|
self.loop.call_soon_threadsafe,
|
|
self._gethostbyname_callback, # type: ignore[arg-type]
|
|
fut,
|
|
host,
|
|
)
|
|
else:
|
|
cb = functools.partial(self._gethostbyname_callback, fut, host)
|
|
self._channel.getaddrinfo(host, None, family=family, callback=cb)
|
|
return fut
|
|
|
|
def getaddrinfo(
|
|
self,
|
|
host: str,
|
|
family: socket.AddressFamily = socket.AF_UNSPEC,
|
|
port: int | None = None,
|
|
proto: int = 0,
|
|
type: int = 0,
|
|
flags: int = 0,
|
|
) -> asyncio.Future[pycares.AddrInfoResult]:
|
|
fut: asyncio.Future[pycares.AddrInfoResult]
|
|
fut, cb = self._get_future_callback()
|
|
self._channel.getaddrinfo(
|
|
host,
|
|
port,
|
|
family=family,
|
|
type=type,
|
|
proto=proto,
|
|
flags=flags,
|
|
callback=cb,
|
|
)
|
|
return fut
|
|
|
|
def getnameinfo(
|
|
self,
|
|
sockaddr: tuple[str, int] | tuple[str, int, int, int],
|
|
flags: int = 0,
|
|
) -> asyncio.Future[pycares.NameInfoResult]:
|
|
fut: asyncio.Future[pycares.NameInfoResult]
|
|
fut, cb = self._get_future_callback()
|
|
self._channel.getnameinfo(sockaddr, flags, callback=cb)
|
|
return fut
|
|
|
|
def gethostbyaddr(self, name: str) -> asyncio.Future[pycares.HostResult]:
|
|
fut: asyncio.Future[pycares.HostResult]
|
|
fut, cb = self._get_future_callback()
|
|
self._channel.gethostbyaddr(name, callback=cb)
|
|
return fut
|
|
|
|
def cancel(self) -> None:
|
|
self._channel.cancel()
|
|
|
|
def _sock_state_cb(self, fd: int, readable: bool, writable: bool) -> None:
|
|
if readable or writable:
|
|
if readable:
|
|
self.loop.add_reader(
|
|
fd, self._channel.process_fd, fd, pycares.ARES_SOCKET_BAD
|
|
)
|
|
self._read_fds.add(fd)
|
|
if writable:
|
|
self.loop.add_writer(
|
|
fd, self._channel.process_fd, pycares.ARES_SOCKET_BAD, fd
|
|
)
|
|
self._write_fds.add(fd)
|
|
if self._timer is None:
|
|
self._start_timer()
|
|
else:
|
|
# socket is now closed
|
|
if fd in self._read_fds:
|
|
self._read_fds.discard(fd)
|
|
self.loop.remove_reader(fd)
|
|
|
|
if fd in self._write_fds:
|
|
self._write_fds.discard(fd)
|
|
self.loop.remove_writer(fd)
|
|
|
|
if (
|
|
not self._read_fds
|
|
and not self._write_fds
|
|
and self._timer is not None
|
|
):
|
|
self._timer.cancel()
|
|
self._timer = None
|
|
|
|
def _timer_cb(self) -> None:
|
|
if self._read_fds or self._write_fds:
|
|
self._channel.process_fd(
|
|
pycares.ARES_SOCKET_BAD, pycares.ARES_SOCKET_BAD
|
|
)
|
|
self._start_timer()
|
|
else:
|
|
self._timer = None
|
|
|
|
def _start_timer(self) -> None:
|
|
timeout = self._timeout
|
|
if timeout is None or timeout < 0 or timeout > 1:
|
|
timeout = 1
|
|
elif timeout == 0:
|
|
timeout = 0.1
|
|
|
|
self._timer = self.loop.call_later(timeout, self._timer_cb)
|
|
|
|
def _cleanup(self) -> None:
|
|
"""Cleanup timers and file descriptors when closing resolver."""
|
|
if self._closed:
|
|
return
|
|
# Mark as closed first to prevent double cleanup
|
|
self._closed = True
|
|
# Cancel timer if running
|
|
if self._timer is not None:
|
|
self._timer.cancel()
|
|
self._timer = None
|
|
|
|
# Remove all file descriptors
|
|
for fd in self._read_fds:
|
|
self.loop.remove_reader(fd)
|
|
for fd in self._write_fds:
|
|
self.loop.remove_writer(fd)
|
|
|
|
self._read_fds.clear()
|
|
self._write_fds.clear()
|
|
self._channel.close()
|
|
|
|
async def close(self) -> None:
|
|
"""
|
|
Cleanly close the DNS resolver.
|
|
|
|
This should be called to ensure all resources are properly released.
|
|
After calling close(), the resolver should not be used again.
|
|
"""
|
|
if not self._closed:
|
|
self._channel.cancel()
|
|
self._cleanup()
|
|
|
|
async def __aenter__(self) -> DNSResolver:
|
|
"""Enter the async context manager."""
|
|
return self
|
|
|
|
async def __aexit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_val: BaseException | None,
|
|
exc_tb: TracebackType | None,
|
|
) -> None:
|
|
"""Exit the async context manager."""
|
|
await self.close()
|
|
|
|
def __del__(self) -> None:
|
|
"""Handle cleanup when the resolver is garbage collected."""
|
|
# Check if we have a channel to clean up
|
|
# This can happen if an exception occurs during __init__ before
|
|
# _channel is created (e.g., RuntimeError on Windows
|
|
# without proper loop)
|
|
if hasattr(self, '_channel'):
|
|
self._cleanup()
|