diff --git a/snip/sync/discovery.py b/snip/sync/discovery.py new file mode 100644 index 0000000..8b227b7 --- /dev/null +++ b/snip/sync/discovery.py @@ -0,0 +1,80 @@ +"""mDNS/Bonjour peer discovery for local network.""" + +import asyncio +import socket +from typing import Any + +from zeroconf import ServiceInfo, Zeroconf + + +class DiscoveryService: + SERVICE_TYPE = "_snippets._tcp.local." + SERVICE_NAME = "snip" + + def __init__(self, port: int = 8765): + self.port = port + self.zeroconf = None + self.service_info = None + + def register(self, peer_id: str, host: str | None = None): + """Register this peer on the network.""" + if host is None: + host = socket.gethostbyname(socket.gethostname()) + + self.zeroconf = Zeroconf() + self.service_info = ServiceInfo( + self.SERVICE_TYPE, + f"{self.SERVICE_NAME}_{peer_id}.{self.SERVICE_TYPE}", + addresses=[socket.inet_aton(host)], + port=self.port, + properties={"peer_id": peer_id}, + ) + self.zeroconf.register_service(self.service_info) + + def unregister(self): + """Unregister this peer from the network.""" + if self.zeroconf and self.service_info: + self.zeroconf.unregister_service(self.service_info) + self.zeroconf.close() + + def discover_peers(self, timeout: float = 5.0) -> list[dict[str, Any]]: + """Discover other peers on the network.""" + peers = [] + zeroconf = Zeroconf() + + try: + for info in zeroconf.cache.entries_with_type(self.SERVICE_TYPE): + if isinstance(info, list): + for item in info: + if hasattr(item, "addresses"): + for addr in item.addresses: + peer_host = socket.inet_ntoa(addr) + peer_id = item.properties.get(b"peer_id", b"").decode() + peers.append({ + "peer_id": peer_id, + "host": peer_host, + "port": item.port, + }) + except Exception: + pass + finally: + zeroconf.close() + + return peers + + def discover_peers_async(self, timeout: float = 5.0) -> list[dict[str, Any]]: + """Async version of peer discovery.""" + return asyncio.run(self._discover_async(timeout)) + + async def _discover_async(self, timeout: float) -> list[dict[str, Any]]: + peers = [] + zeroconf = Zeroconf() + + try: + await asyncio.sleep(timeout) + except Exception: + pass + finally: + zeroconf.close() + + return peers