diff --git a/snip/sync/discovery.py b/snip/sync/discovery.py index b92863d..ff4e03e 100644 --- a/snip/sync/discovery.py +++ b/snip/sync/discovery.py @@ -1,87 +1,132 @@ -"""mDNS/Bonjour peer discovery for local network.""" - import json import socket +import uuid from pathlib import Path -from typing import Any +from typing import Any, Callable -from zeroconf import ServiceInfo, Zeroconf +try: + from zeroconf import ServiceInfo, Zeroconf + ZEROCONF_AVAILABLE = True +except ImportError: + ZEROCONF_AVAILABLE = False + Zeroconf = None + ServiceInfo = None class DiscoveryService: SERVICE_TYPE = "_snippets._tcp.local." SERVICE_NAME = "snip" + PEER_CACHE_FILE = Path.home() / ".snip" / "peers.json" - def __init__(self, port: int = 8765, peer_id: str | None = None): + def __init__(self, port: int = 8765, peer_name: str | None = None): self.port = port - self.peer_id = peer_id or socket.gethostname() - self.zeroconf = None - self.service_info = None - self._peer_cache_file = Path("~/.snip/peers.json").expanduser() + self.peer_id = str(uuid.uuid4())[:8] + self.peer_name = peer_name or socket.gethostname() + self._zeroconf: Zeroconf | None = None + self._service_info: ServiceInfo | None = None + self._listening = False + self._listeners: list[Callable[[dict[str, Any]], None]] = [] + self._peer_cache: list[dict[str, Any]] = [] + self._load_peer_cache() - def register(self, peer_id: str | None = None, host: str | None = None): - """Register this peer on the network.""" - if peer_id is None: - peer_id = self.peer_id - if host is None: + def _load_peer_cache(self): + if self.PEER_CACHE_FILE.exists(): try: - host = socket.gethostbyname(socket.gethostname()) + with open(self.PEER_CACHE_FILE) as f: + self._peer_cache = json.load(f) except Exception: - host = "127.0.0.1" - - 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.encode()}, - ) - 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() + self._peer_cache = [] + def _save_peer_cache(self): + self.PEER_CACHE_FILE.parent.mkdir(parents=True, exist_ok=True) 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() - peer_name = item.name.replace(f".{self.SERVICE_TYPE}", "") - peers.append({ - "peer_id": peer_id, - "peer_name": peer_name, - "host": peer_host, - "addresses": [peer_host], - "port": item.port, - }) + with open(self.PEER_CACHE_FILE, "w") as f: + json.dump(self._peer_cache, f) except Exception: pass - finally: - zeroconf.close() + def register(self) -> bool: + if not ZEROCONF_AVAILABLE: + return False + try: + self._zeroconf = Zeroconf() + self._service_info = ServiceInfo( + self.SERVICE_TYPE, + f"{self.SERVICE_NAME}_{self.peer_id}.{self.SERVICE_TYPE}", + addresses=[socket.inet_aton(socket.gethostbyname(socket.gethostname()))], + port=self.port, + properties={ + "peer_id": self.peer_id, + "peer_name": self.peer_name, + }, + ) + self._zeroconf.register_service(self._service_info) + return True + except Exception: + return False + + def unregister(self): + if self._zeroconf and self._service_info: + try: + self._zeroconf.unregister_service(self._service_info) + except Exception: + pass + self._zeroconf.close() + self._zeroconf = None + self._service_info = None + + def discover_peers(self, timeout: float = 5.0) -> list[dict[str, Any]]: + if not ZEROCONF_AVAILABLE: + return self._peer_cache + peers = [] + seen_ids = set() + try: + zeroconf = Zeroconf() + listener = PeerListener(seen_ids, peers) + browser = zeroconf.ServiceBrowser(self.SERVICE_TYPE, listener) + browser.cancel() + for peer in self._peer_cache: + if peer["peer_id"] not in seen_ids: + seen_ids.add(peer["peer_id"]) + zeroconf.close() + except Exception: + pass + self._peer_cache = peers + self._save_peer_cache() return peers - def save_peer_cache(self, peers: list[dict[str, Any]]): - """Save discovered peers to cache.""" - self._peer_cache_file.parent.mkdir(parents=True, exist_ok=True) - with open(self._peer_cache_file, "w") as f: - json.dump(peers, f) + def add_listener(self, callback: Callable[[dict[str, Any]], None]): + self._listeners.append(callback) - def load_peer_cache(self) -> list[dict[str, Any]]: - """Load cached peers.""" - if self._peer_cache_file.exists(): - with open(self._peer_cache_file, "r") as f: - return json.load(f) - return [] \ No newline at end of file + def remove_listener(self, callback: Callable[[dict[str, Any]], None]): + if callback in self._listeners: + self._listeners.remove(callback) + + +class PeerListener: + def __init__(self, seen_ids: set, peers: list): + self.seen_ids = seen_ids + self.peers = peers + + def add_service(self, zeroconf, service_type: str, name: str): + try: + info = zeroconf.get_service_info(service_type, name) + if info and info.properties: + peer_id = info.properties.get(b"peer_id", b"").decode() + peer_name = info.properties.get(b"peer_name", b"").decode() + if peer_id and peer_id not in self.seen_ids: + self.seen_ids.add(peer_id) + addresses = [ + socket.inet_ntoa(addr) for addr in info.addresses + ] if info.addresses else [] + self.peers.append({ + "peer_id": peer_id, + "peer_name": peer_name, + "addresses": addresses, + "port": info.port, + }) + except Exception: + pass + + def remove_service(self, zeroconf, service_type: str, name: str): + pass