81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
"""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
|