136 lines
4.0 KiB
Python
136 lines
4.0 KiB
Python
"""mDNS peer discovery using zeroconf (optional dependency).
|
|
|
|
Falls back gracefully if zeroconf is not installed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import socket
|
|
from typing import Callable
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
MDNS_SERVICE_TYPE = "_meshnet._tcp.local."
|
|
|
|
try:
|
|
from zeroconf import ServiceInfo, ServiceBrowser, Zeroconf # type: ignore[import]
|
|
|
|
_HAS_ZEROCONF = True
|
|
except ImportError:
|
|
_HAS_ZEROCONF = False
|
|
|
|
|
|
def _local_ip() -> str:
|
|
try:
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
s.connect(("8.8.8.8", 80))
|
|
ip = s.getsockname()[0]
|
|
s.close()
|
|
return ip
|
|
except OSError:
|
|
return "127.0.0.1"
|
|
|
|
|
|
class MdnsDiscovery:
|
|
"""Announce this node on mDNS and discover peers on the same LAN.
|
|
|
|
If `zeroconf` is not installed, all methods are no-ops.
|
|
|
|
Usage::
|
|
|
|
disc = MdnsDiscovery(
|
|
peer_id="abc123",
|
|
port=8001,
|
|
on_peer_found=lambda peer_id, addr: print("found", peer_id, addr),
|
|
)
|
|
disc.start()
|
|
...
|
|
disc.stop()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
peer_id: str,
|
|
port: int,
|
|
on_peer_found: Callable[[str, str], None] | None = None,
|
|
on_peer_lost: Callable[[str], None] | None = None,
|
|
):
|
|
self.peer_id = peer_id
|
|
self.port = port
|
|
self.on_peer_found = on_peer_found
|
|
self.on_peer_lost = on_peer_lost
|
|
self._zc: "Zeroconf | None" = None # type: ignore[name-defined]
|
|
self._info: "ServiceInfo | None" = None # type: ignore[name-defined]
|
|
self._browser = None
|
|
|
|
def is_available(self) -> bool:
|
|
return _HAS_ZEROCONF
|
|
|
|
def start(self) -> None:
|
|
if not _HAS_ZEROCONF:
|
|
log.info("zeroconf not installed — mDNS discovery disabled")
|
|
return
|
|
try:
|
|
self._zc = Zeroconf()
|
|
local_ip = _local_ip()
|
|
self._info = ServiceInfo(
|
|
MDNS_SERVICE_TYPE,
|
|
f"{self.peer_id}.{MDNS_SERVICE_TYPE}",
|
|
addresses=[socket.inet_aton(local_ip)],
|
|
port=self.port,
|
|
properties={"peer_id": self.peer_id, "version": "1"},
|
|
)
|
|
self._zc.register_service(self._info)
|
|
if self.on_peer_found or self.on_peer_lost:
|
|
self._browser = ServiceBrowser(
|
|
self._zc, MDNS_SERVICE_TYPE, listener=_Listener(self)
|
|
)
|
|
log.info("mDNS announced: %s on %s:%d", self.peer_id, local_ip, self.port)
|
|
except Exception as exc:
|
|
log.warning("mDNS start failed: %s", exc)
|
|
|
|
def stop(self) -> None:
|
|
if not _HAS_ZEROCONF or self._zc is None:
|
|
return
|
|
try:
|
|
if self._info:
|
|
self._zc.unregister_service(self._info)
|
|
self._zc.close()
|
|
except Exception as exc:
|
|
log.debug("mDNS stop error: %s", exc)
|
|
self._zc = None
|
|
|
|
|
|
class _Listener:
|
|
"""Internal zeroconf service listener."""
|
|
|
|
def __init__(self, disc: MdnsDiscovery):
|
|
self._disc = disc
|
|
|
|
def add_service(self, zc, type_, name):
|
|
try:
|
|
info = zc.get_service_info(type_, name)
|
|
if info is None:
|
|
return
|
|
remote_peer_id = (info.properties or {}).get(b"peer_id", b"").decode()
|
|
if remote_peer_id == self._disc.peer_id:
|
|
return # ignore self
|
|
addr = f"http://{socket.inet_ntoa(info.addresses[0])}:{info.port}"
|
|
if self._disc.on_peer_found:
|
|
self._disc.on_peer_found(remote_peer_id, addr)
|
|
except Exception as exc:
|
|
log.debug("mDNS add_service error: %s", exc)
|
|
|
|
def remove_service(self, zc, type_, name):
|
|
try:
|
|
# name is like "peer_id._meshnet._tcp.local."
|
|
peer_id = name.split(".")[0]
|
|
if self._disc.on_peer_lost:
|
|
self._disc.on_peer_lost(peer_id)
|
|
except Exception as exc:
|
|
log.debug("mDNS remove_service error: %s", exc)
|
|
|
|
def update_service(self, zc, type_, name):
|
|
pass
|