pymobiledevice3 4.14.6__py3-none-any.whl → 7.0.6__py3-none-any.whl

pymobiledevice3/bonjour.py

@@ -1,135 +1,380 @@
+# Async, dependency-light mDNS browser returning dataclasses with per-address interface names.
+# Works for any DNS-SD type, e.g. "_remoted._tcp.local."
+# - Uses ifaddr (optional) to map IPs -> local interfaces; otherwise iface will be None.
+
 import asyncio
-import dataclasses
-from socket import AF_INET, AF_INET6, inet_ntop
+import contextlib
+import ipaddress
+import socket
+import struct
+import sys
+from collections import defaultdict
+from dataclasses import dataclass, field
 from typing import Optional
 
-from ifaddr import get_adapters
-from zeroconf import IPVersion, ServiceListener, ServiceStateChange, Zeroconf
-from zeroconf.asyncio import AsyncServiceBrowser, AsyncServiceInfo, AsyncZeroconf
+import ifaddr  # pip install ifaddr
 
 from pymobiledevice3.osu.os_utils import get_os_utils
 
-REMOTEPAIRING_SERVICE_NAMES = ['_remotepairing._tcp.local.']
-REMOTEPAIRING_MANUAL_PAIRING_SERVICE_NAMES = ['_remotepairing-manual-pairing._tcp.local.']
-MOBDEV2_SERVICE_NAMES = ['_apple-mobdev2._tcp.local.']
-REMOTED_SERVICE_NAMES = ['_remoted._tcp.local.']
+REMOTEPAIRING_SERVICE_NAME = "_remotepairing._tcp.local."
+REMOTEPAIRING_MANUAL_PAIRING_SERVICE_NAME = "_remotepairing-manual-pairing._tcp.local."
+MOBDEV2_SERVICE_NAME = "_apple-mobdev2._tcp.local."
+REMOTED_SERVICE_NAME = "_remoted._tcp.local."
 OSUTILS = get_os_utils()
 DEFAULT_BONJOUR_TIMEOUT = OSUTILS.bonjour_timeout
 
+MDNS_PORT = 5353
+MDNS_MCAST_V4 = "224.0.0.251"
+MDNS_MCAST_V6 = "ff02::fb"
 
-@dataclasses.dataclass
-class BonjourAnswer:
-    name: str
-    properties: dict[bytes, bytes]
-    ips: list[str]
-    port: int
-
-
-class BonjourListener(ServiceListener):
-    def __init__(self, ip: str):
-        super().__init__()
-        self.name: Optional[str] = None
-        self.properties: dict[bytes, bytes] = {}
-        self.ip = ip
-        self.port: Optional[int] = None
-        self.addresses: list[str] = []
-        self.queue: asyncio.Queue = asyncio.Queue()
-        self.querying_task: Optional[asyncio.Task] = asyncio.create_task(self.query_addresses())
-
-    def async_on_service_state_change(
-            self, zeroconf: Zeroconf, service_type: str, name: str, state_change: ServiceStateChange) -> None:
-        self.queue.put_nowait((zeroconf, service_type, name, state_change))
-
-    async def query_addresses(self) -> None:
-        zeroconf, service_type, name, state_change = await self.queue.get()
-        self.name = name
-        service_info = AsyncServiceInfo(service_type, name)
-        await service_info.async_request(zeroconf, 3000)
-        ipv4 = [inet_ntop(AF_INET, address.packed) for address in
-                service_info.ip_addresses_by_version(IPVersion.V4Only)]
-        ipv6 = []
-        if '%' in self.ip:
-            ipv6 = [inet_ntop(AF_INET6, address.packed) + '%' + self.ip.split('%')[1] for address in
-                    service_info.ip_addresses_by_version(IPVersion.V6Only)]
-        self.addresses = ipv4 + ipv6
-        self.properties = service_info.properties
-        self.port = service_info.port
-
-    async def close(self) -> None:
-        self.querying_task.cancel()
-        try:
-            await self.querying_task
-        except asyncio.CancelledError:
-            pass
-
-
-@dataclasses.dataclass
-class BonjourQuery:
-    zc: AsyncZeroconf
-    service_browser: AsyncServiceBrowser
-    listener: BonjourListener
-
-
-def query_bonjour(service_names: list[str], ip: str) -> BonjourQuery:
-    aiozc = AsyncZeroconf(interfaces=[ip])
-    listener = BonjourListener(ip)
-    service_browser = AsyncServiceBrowser(aiozc.zeroconf, service_names,
-                                          handlers=[listener.async_on_service_state_change])
-    return BonjourQuery(aiozc, service_browser, listener)
-
-
-async def browse(service_names: list[str], ips: list[str], timeout: float = DEFAULT_BONJOUR_TIMEOUT) \
-        -> list[BonjourAnswer]:
-    bonjour_queries = [query_bonjour(service_names, adapter) for adapter in ips]
-    answers = []
-    await asyncio.sleep(timeout)
-    for bonjour_query in bonjour_queries:
-        if bonjour_query.listener.addresses:
-            answer = BonjourAnswer(
-                bonjour_query.listener.name, bonjour_query.listener.properties, bonjour_query.listener.addresses,
-                bonjour_query.listener.port)
-            if answer not in answers:
-                answers.append(answer)
-        await bonjour_query.listener.close()
-        await bonjour_query.service_browser.async_cancel()
-        await bonjour_query.zc.async_close()
-    return answers
-
-
-async def browse_ipv6(service_names: list[str], timeout: float = DEFAULT_BONJOUR_TIMEOUT) -> list[BonjourAnswer]:
-    return await browse(service_names, OSUTILS.get_ipv6_ips(), timeout=timeout)
-
-
-def get_ipv4_addresses() -> list[str]:
-    ips = []
-    for adapter in get_adapters():
-        if adapter.nice_name.startswith('tun'):
-            # skip browsing on already established tunnels
+QTYPE_A = 1
+QTYPE_PTR = 12
+QTYPE_TXT = 16
+QTYPE_AAAA = 28
+QTYPE_SRV = 33
+
+CLASS_IN = 0x0001
+CLASS_QU = 0x8000  # unicast-response bit (we use multicast queries)
+
+
+# ---------------- Dataclasses ----------------
+
+
+# --- Dataclass decorator shim (adds slots only on 3.10+)
+def dataclass_compat(*d_args, **d_kwargs):
+    if sys.version_info < (3, 10):
+        d_kwargs.pop("slots", None)  # ignore on 3.9
+    return dataclass(*d_args, **d_kwargs)
+
+
+@dataclass_compat(slots=True)
+class Address:
+    ip: str
+    iface: str  # local interface name (e.g., "en0"), or None if unknown
+
+    @property
+    def full_ip(self) -> str:
+        if self.iface and self.ip.lower().startswith("fe80:"):
+            return f"{self.ip}%{self.iface}"
+        return self.ip
+
+
+@dataclass_compat(slots=True)
+class ServiceInstance:
+    instance: str  # "<Instance Name>._type._proto.local."
+    host: Optional[str]  # "host.local" (without trailing dot), or None if unresolved
+    port: Optional[int]  # SRV port
+    addresses: list[Address] = field(default_factory=list)  # IPs with interface names
+    properties: dict[str, str] = field(default_factory=dict)  # TXT key/values
+
+
+# ---------------- DNS helpers ----------------
+
+
+def encode_name(name: str) -> bytes:
+    name = name.rstrip(".")
+    out = bytearray()
+    for label in name.split(".") if name else []:
+        b = label.encode("utf-8")
+        if len(b) > 63:
+            raise ValueError("label too long")
+        out.append(len(b))
+        out += b
+    out.append(0)
+    return bytes(out)
+
+
+def decode_name(data: bytes, off: int) -> tuple[str, int]:
+    labels = []
+    jumped = False
+    orig_end = off
+    for _ in range(128):  # loop guard
+        if off >= len(data):
+            break
+        length = data[off]
+        if length == 0:
+            off += 1
+            break
+        if (length & 0xC0) == 0xC0:
+            if off + 1 >= len(data):
+                raise ValueError("truncated name pointer")
+            ptr = ((length & 0x3F) << 8) | data[off + 1]
+            if ptr >= len(data):
+                raise ValueError("bad name pointer")
+            if not jumped:
+                orig_end = off + 2
+            off = ptr
+            jumped = True
             continue
-        for ip in adapter.ips:
-            if ip.ip == '127.0.0.1':
-                continue
-            if not ip.is_IPv4:
+        off += 1
+        end = off + length
+        if end > len(data):
+            raise ValueError("truncated label")
+        labels.append(data[off:end].decode("utf-8", errors="replace"))
+        off = end
+    return ".".join(labels) + ".", (orig_end if jumped else off)
+
+
+def build_query(name: str, qtype: int, unicast: bool = False) -> bytes:
+    hdr = struct.pack("!HHHHHH", 0, 0, 1, 0, 0, 0)  # TXID=0, flags=0, 1 question
+    qclass = CLASS_IN | (CLASS_QU if unicast else 0)
+    return hdr + encode_name(name) + struct.pack("!HH", qtype, qclass)
+
+
+def parse_rr(data: bytes, off: int):
+    name, off = decode_name(data, off)
+    if off + 10 > len(data):
+        raise ValueError("truncated RR header")
+    rtype, rclass, ttl, rdlen = struct.unpack("!HHIH", data[off : off + 10])
+    off += 10
+    rdata = data[off : off + rdlen]
+    off += rdlen
+
+    rr = {"name": name, "type": rtype, "class": rclass & 0x7FFF, "ttl": ttl}
+    if rtype == QTYPE_PTR:
+        target, _ = decode_name(data, off - rdlen)
+        rr["ptrdname"] = target
+    elif rtype == QTYPE_SRV and rdlen >= 6:
+        priority, weight, port = struct.unpack("!HHH", rdata[:6])
+        target, _ = decode_name(data, off - rdlen + 6)
+        rr.update({"priority": priority, "weight": weight, "port": port, "target": target})
+    elif rtype == QTYPE_TXT:
+        kv = {}
+        i = 0
+        while i < rdlen:
+            b = rdata[i]
+            i += 1
+            seg = rdata[i : i + b]
+            i += b
+            if not seg:
                 continue
-            ips.append(ip.ip)
-    return ips
+            if b"=" in seg:
+                k, v = seg.split(b"=", 1)
+                kv[k.decode()] = v.decode(errors="replace")
+            else:
+                kv[seg.decode()] = ""
+        rr["txt"] = kv
+    elif rtype == QTYPE_A and rdlen == 4:
+        rr["address"] = socket.inet_ntop(socket.AF_INET, rdata)
+    elif rtype == QTYPE_AAAA and rdlen == 16:
+        rr["address"] = socket.inet_ntop(socket.AF_INET6, rdata)
+    else:
+        rr["raw"] = rdata
+    return rr, off
+
+
+def parse_mdns_message(data: bytes):
+    if len(data) < 12:
+        return []
+    _, _, qd, an, ns, ar = struct.unpack("!HHHHHH", data[:12])
+    off = 12
+    for _ in range(qd):
+        _, off = decode_name(data, off)
+        off += 4
+    rrs = []
+    for _ in range(an + ns + ar):
+        rr, off = parse_rr(data, off)
+        rrs.append(rr)
+    return rrs
+
+
+# ---------------- Interface mapping helpers ----------------
+
+
+class _Adapters:
+    def __init__(self):
+        self.adapters = ifaddr.get_adapters() if ifaddr is not None else []
+
+    def pick_iface_for_ip(self, ip_str: str, family: int, v6_scopeid: Optional[int]) -> Optional[str]:
+        # Prefer scope id for IPv6 link-local
+        if family == socket.AF_INET6 and ip_str.lower().startswith("fe80:") and v6_scopeid:
+            try:
+                return socket.if_indextoname(v6_scopeid)
+            except OSError:
+                pass
+
+        # Otherwise, try to match destination ip to local subnet via ifaddr
+        if not self.adapters:
+            return None
+        ip = ipaddress.ip_address(ip_str)
+        best = (None, -1)  # (name, prefix_len)
+        for ad in self.adapters:
+            for ipn in ad.ips:
+                if isinstance(ipn.ip, str):
+                    # IPv4
+                    fam = socket.AF_INET
+                    ipn_ip = ipn.ip
+                else:
+                    # IPv6
+                    fam = socket.AF_INET6
+                    ipn_ip = ipn.ip[0]
+                if fam != family:
+                    continue
+                net = ipaddress.ip_network(f"{ipn_ip}/{ipn.network_prefix}", strict=False)
+                if ip in net and ipn.network_prefix > best[1]:
+                    best = (ad.nice_name or ad.name, ipn.network_prefix)
+        return best[0]
+
+
+# ---------------- async sockets ----------------
+
+
+class _DatagramProtocol(asyncio.DatagramProtocol):
+    def __init__(self, queue: asyncio.Queue):
+        self.queue = queue
+
+    def datagram_received(self, data, addr):
+        # addr: IPv4 -> (host, port); IPv6 -> (host, port, flowinfo, scopeid)
+        self.queue.put_nowait((data, addr))
+
+
+async def _bind_ipv4(queue: asyncio.Queue):
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    if hasattr(socket, "SO_REUSEPORT"):
+        with contextlib.suppress(OSError):
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
+    s.bind(("0.0.0.0", MDNS_PORT))
+    try:
+        mreq = struct.pack("=4s4s", socket.inet_aton(MDNS_MCAST_V4), socket.inet_aton("0.0.0.0"))
+        s.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)
+    except OSError:
+        pass
+    transport, _ = await asyncio.get_running_loop().create_datagram_endpoint(lambda: _DatagramProtocol(queue), sock=s)
+    return transport, s
+
+
+async def _bind_ipv6_all_ifaces(queue: asyncio.Queue):
+    s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
+    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    if hasattr(socket, "SO_REUSEPORT"):
+        with contextlib.suppress(OSError):
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
+    s.bind(("::", MDNS_PORT))
+    grp = socket.inet_pton(socket.AF_INET6, MDNS_MCAST_V6)
+    for ifindex, _ in socket.if_nameindex():
+        mreq6 = grp + struct.pack("@I", ifindex)
+        try:
+            s.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_JOIN_GROUP, mreq6)
+        except OSError:
+            continue
+    transport, _ = await asyncio.get_running_loop().create_datagram_endpoint(lambda: _DatagramProtocol(queue), sock=s)
+    return transport, s
+
+
+async def _open_mdns_sockets():
+    queue = asyncio.Queue()
+    transports: list[tuple[asyncio.BaseTransport, socket.socket]] = []
+    t4, s4 = await _bind_ipv4(queue)
+    transports.append((t4, s4))
+    t6, s6 = await _bind_ipv6_all_ifaces(queue)
+    transports.append((t6, s6))
+    if not transports:
+        raise RuntimeError("Failed to open mDNS sockets (UDP/5353)")
+    return transports, queue
+
+
+async def _send_query_all(transports, pkt: bytes):
+    for transport, sock in transports:
+        if sock.family == socket.AF_INET:
+            transport.sendto(pkt, (MDNS_MCAST_V4, MDNS_PORT))
+        else:
+            # Send once per iface index for better reachability
+            for ifindex, _ in socket.if_nameindex():
+                transport.sendto(pkt, (MDNS_MCAST_V6, MDNS_PORT, 0, ifindex))
+
+
+# ---------------- Public API ----------------
+
+
+async def browse_service(service_type: str, timeout: float = 4.0) -> list[ServiceInstance]:
+    """
+    Discover a DNS-SD/mDNS service type (e.g. "_remoted._tcp.local.") on the local network.
+
+    Returns: List[ServiceInstance] with Address(ip, iface) entries.
+    """
+    if not service_type.endswith("."):
+        service_type += "."
+
+    transports, queue = await _open_mdns_sockets()
+    adapters = _Adapters()
+
+    ptr_targets: set[str] = set()
+    srv_map: dict[str, list[dict]] = defaultdict(list)  # instance_name -> list of {"target", "port"}
+    txt_map: dict[str, dict] = {}
+    host_addrs: dict[str, list[Address]] = defaultdict(list)  # host -> list[(ip, iface)]
+
+    def _record_addr(rr_name: str, ip_str: str, pkt_addr):
+        # Determine family and possible scopeid from the packet that delivered this RR
+        family = socket.AF_INET6 if ":" in ip_str else socket.AF_INET
+        scopeid = None
+        if isinstance(pkt_addr, tuple) and len(pkt_addr) == 4:  # IPv6 remote tuple
+            scopeid = pkt_addr[3]
+        iface = adapters.pick_iface_for_ip(ip_str, family, scopeid)
+        if iface is None:
+            return
+        # Avoid duplicates for the same host/ip
+        existing = host_addrs[rr_name]
+        if not any(a.ip == ip_str for a in existing):
+            existing.append(Address(ip=ip_str, iface=iface))
 
+    try:
+        await _send_query_all(transports, build_query(service_type, QTYPE_PTR, unicast=False))
+        loop = asyncio.get_running_loop()
+        end = loop.time() + timeout
+        while loop.time() < end:
+            try:
+                data, pkt_addr = await asyncio.wait_for(queue.get(), timeout=end - loop.time())
+            except asyncio.TimeoutError:
+                break
+            for rr in parse_mdns_message(data):
+                t = rr.get("type")
+                if t == QTYPE_PTR and rr.get("name") == service_type:
+                    ptr_targets.add(rr.get("ptrdname"))
+                elif t == QTYPE_SRV:
+                    srv_map[rr["name"]].append({"target": rr.get("target"), "port": rr.get("port")})
+                elif t == QTYPE_TXT:
+                    # TODO: This could possibly mix the properties of multiple TXT records for the same instance.
+                    #       However, it's currently unused.
+                    txt_map[rr["name"]] = rr.get("txt", {})
+                elif (t == QTYPE_A and rr.get("address")) or (t == QTYPE_AAAA and rr.get("address")):
+                    _record_addr(rr["name"], rr["address"], pkt_addr)
+    finally:
+        for transport, _ in transports:
+            transport.close()
 
-async def browse_ipv4(service_names: list[str], timeout: float = DEFAULT_BONJOUR_TIMEOUT) -> list[BonjourAnswer]:
-    return await browse(service_names, get_ipv4_addresses(), timeout=timeout)
+    # Assemble dataclasses
+    results: list[ServiceInstance] = []
+    for inst in sorted(ptr_targets):
+        srv_entries = srv_map.get(inst, [])
+        props = txt_map.get(inst, {})
+        for srv in srv_entries:
+            target = srv.get("target")
+            host = (target[:-1] if target and target.endswith(".") else target) or None
+            addrs = host_addrs.get(target, []) if target else []
+            results.append(
+                ServiceInstance(
+                    instance=inst,
+                    host=host,
+                    port=srv.get("port"),
+                    addresses=addrs,
+                    properties=props,
+                )
+            )
+    return results
 
 
-async def browse_remoted(timeout: float = DEFAULT_BONJOUR_TIMEOUT) -> list[BonjourAnswer]:
-    return await browse_ipv6(REMOTED_SERVICE_NAMES, timeout=timeout)
+async def browse_remoted(timeout: float = DEFAULT_BONJOUR_TIMEOUT) -> list[ServiceInstance]:
+    return await browse_service(REMOTED_SERVICE_NAME, timeout=timeout)
 
 
-async def browse_mobdev2(timeout: float = DEFAULT_BONJOUR_TIMEOUT) -> list[BonjourAnswer]:
-    return await browse(MOBDEV2_SERVICE_NAMES, get_ipv4_addresses() + OSUTILS.get_ipv6_ips(), timeout=timeout)
+async def browse_mobdev2(timeout: float = DEFAULT_BONJOUR_TIMEOUT) -> list[ServiceInstance]:
+    return await browse_service(MOBDEV2_SERVICE_NAME, timeout=timeout)
 
 
-async def browse_remotepairing(timeout: float = DEFAULT_BONJOUR_TIMEOUT) -> list[BonjourAnswer]:
-    return await browse_ipv4(REMOTEPAIRING_SERVICE_NAMES, timeout=timeout)
+async def browse_remotepairing(timeout: float = DEFAULT_BONJOUR_TIMEOUT) -> list[ServiceInstance]:
+    return await browse_service(REMOTEPAIRING_SERVICE_NAME, timeout=timeout)
 
 
-async def browse_remotepairing_manual_pairing(timeout: float = DEFAULT_BONJOUR_TIMEOUT) -> list[BonjourAnswer]:
-    return await browse_ipv4(REMOTEPAIRING_MANUAL_PAIRING_SERVICE_NAMES, timeout=timeout)
+async def browse_remotepairing_manual_pairing(timeout: float = DEFAULT_BONJOUR_TIMEOUT) -> list[ServiceInstance]:
+    return await browse_service(REMOTEPAIRING_MANUAL_PAIRING_SERVICE_NAME, timeout=timeout)