87 lines
3.1 KiB
Python
87 lines
3.1 KiB
Python
"""mDNS/Bonjour peer discovery for local network."""
|
|
|
|
import json
|
|
import socket
|
|
from pathlib import Path
|
|
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, peer_id: 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()
|
|
|
|
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:
|
|
try:
|
|
host = socket.gethostbyname(socket.gethostname())
|
|
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()
|
|
|
|
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,
|
|
})
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
zeroconf.close()
|
|
|
|
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 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 [] |