import asyncio import contextlib import functools import socket from dataclasses import dataclass from ipaddress import IPv4Address, IPv6Address from typing import List, Optional, Tuple, Union, cast import zeroconf try: import zeroconf.asyncio ZC_ASYNCIO = True except ImportError: ZC_ASYNCIO = False from .core import APIConnectionError, ResolveAPIError ZeroconfInstanceType = Union[zeroconf.Zeroconf, "zeroconf.asyncio.AsyncZeroconf", None] @dataclass(frozen=True) class Sockaddr: pass @dataclass(frozen=True) class IPv4Sockaddr(Sockaddr): address: str port: int @dataclass(frozen=True) class IPv6Sockaddr(Sockaddr): address: str port: int flowinfo: int scope_id: int @dataclass(frozen=True) class AddrInfo: family: int type: int proto: int sockaddr: Sockaddr def _sync_zeroconf_get_service_info( zeroconf_instance: ZeroconfInstanceType, service_type: str, service_name: str, timeout: float, ) -> Optional["zeroconf.ServiceInfo"]: # Use or create zeroconf instance, ensure it's an AsyncZeroconf if zeroconf_instance is None: try: zc = zeroconf.Zeroconf() except Exception: raise ResolveAPIError( "Cannot start mDNS sockets, is this a docker container without " "host network mode?" ) do_close = True elif isinstance(zeroconf_instance, zeroconf.Zeroconf): zc = zeroconf_instance do_close = False else: raise ValueError( f"Invalid type passed for zeroconf_instance: {type(zeroconf_instance)}" ) try: info = zc.get_service_info(service_type, service_name, int(timeout * 1000)) except Exception as exc: raise ResolveAPIError( f"Error resolving mDNS {service_name} via mDNS: {exc}" ) from exc finally: if do_close: zc.close() return info async def _async_zeroconf_get_service_info( eventloop: asyncio.events.AbstractEventLoop, zeroconf_instance: ZeroconfInstanceType, service_type: str, service_name: str, timeout: float, ) -> Optional["zeroconf.ServiceInfo"]: if not ZC_ASYNCIO: return await eventloop.run_in_executor( None, functools.partial( _sync_zeroconf_get_service_info, zeroconf_instance, service_type, service_name, timeout, ), ) # Use or create zeroconf instance, ensure it's an AsyncZeroconf if zeroconf_instance is None: try: zc = zeroconf.asyncio.AsyncZeroconf() except Exception: raise ResolveAPIError( "Cannot start mDNS sockets, is this a docker container without " "host network mode?" ) do_close = True elif isinstance(zeroconf_instance, zeroconf.asyncio.AsyncZeroconf): zc = zeroconf_instance do_close = False elif isinstance(zeroconf_instance, zeroconf.Zeroconf): zc = zeroconf.asyncio.AsyncZeroconf(zc=zeroconf_instance) do_close = False else: raise ValueError( f"Invalid type passed for zeroconf_instance: {type(zeroconf_instance)}" ) try: info = await zc.async_get_service_info( service_type, service_name, int(timeout * 1000) ) except Exception as exc: raise ResolveAPIError( f"Error resolving mDNS {service_name} via mDNS: {exc}" ) from exc finally: if do_close: await zc.async_close() return info async def _async_resolve_host_zeroconf( eventloop: asyncio.events.AbstractEventLoop, host: str, port: int, *, timeout: float = 3.0, zeroconf_instance: ZeroconfInstanceType = None, ) -> List[AddrInfo]: service_type = "_esphomelib._tcp.local." service_name = f"{host}.{service_type}" info = await _async_zeroconf_get_service_info( eventloop, zeroconf_instance, service_type, service_name, timeout ) if info is None: return [] addrs: List[AddrInfo] = [] for raw in info.addresses_by_version(zeroconf.IPVersion.All): is_ipv6 = len(raw) == 16 sockaddr: Sockaddr if is_ipv6: sockaddr = IPv6Sockaddr( address=socket.inet_ntop(socket.AF_INET6, raw), port=port, flowinfo=0, scope_id=0, ) else: sockaddr = IPv4Sockaddr( address=socket.inet_ntop(socket.AF_INET, raw), port=port, ) addrs.append( AddrInfo( family=socket.AF_INET6 if is_ipv6 else socket.AF_INET, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP, sockaddr=sockaddr, ) ) return addrs async def _async_resolve_host_getaddrinfo( eventloop: asyncio.events.AbstractEventLoop, host: str, port: int ) -> List[AddrInfo]: try: # Limit to TCP IP protocol res = await eventloop.getaddrinfo(host, port, proto=socket.IPPROTO_TCP) except OSError as err: raise APIConnectionError("Error resolving IP address: {}".format(err)) addrs: List[AddrInfo] = [] for family, type_, proto, _, raw in res: sockaddr: Sockaddr if family == socket.AF_INET: raw = cast(Tuple[str, int], raw) address, port = raw sockaddr = IPv4Sockaddr(address=address, port=port) elif family == socket.AF_INET6: raw = cast(Tuple[str, int, int, int], raw) address, port, flowinfo, scope_id = raw sockaddr = IPv6Sockaddr( address=address, port=port, flowinfo=flowinfo, scope_id=scope_id ) else: # Unknown family continue addrs.append( AddrInfo(family=family, type=type_, proto=proto, sockaddr=sockaddr) ) return addrs def _async_ip_address_to_addrs(host: str, port: int) -> List[AddrInfo]: """Convert an ipaddress to AddrInfo.""" with contextlib.suppress(ValueError): return [ AddrInfo( family=socket.AF_INET6, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP, sockaddr=IPv6Sockaddr( address=str(IPv6Address(host)), port=port, flowinfo=0, scope_id=0 ), ) ] with contextlib.suppress(ValueError): return [ AddrInfo( family=socket.AF_INET, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP, sockaddr=IPv4Sockaddr( address=str(IPv4Address(host)), port=port, ), ) ] return [] async def async_resolve_host( eventloop: asyncio.events.AbstractEventLoop, host: str, port: int, zeroconf_instance: ZeroconfInstanceType = None, ) -> AddrInfo: addrs: List[AddrInfo] = [] zc_error = None if host.endswith(".local"): name = host[: -len(".local")] try: addrs.extend( await _async_resolve_host_zeroconf( eventloop, name, port, zeroconf_instance=zeroconf_instance ) ) except APIConnectionError as err: zc_error = err if not addrs: addrs.extend(_async_ip_address_to_addrs(host, port)) if not addrs: addrs.extend(await _async_resolve_host_getaddrinfo(eventloop, host, port)) if not addrs: if zc_error: # Only show ZC error if getaddrinfo also didn't work raise zc_error raise ResolveAPIError(f"Could not resolve host {host} - got no results from OS") # Use first matching result # Future: return all matches and use first working one return addrs[0]