Skip to content

Commit 64d7e7f

Browse files
author
Gordon AI
committed
feat: add LAN service discovery with mDNS/Zeroconf support
- Implement MohawkServiceDiscovery class for automatic service discovery - Add LanServiceRegistry for service registration on mDNS - Support for both GUI and worker service types - Automatic service state change callbacks (added/removed) - Service filtering by type (gui/worker) - Expose service metadata (IP, port, properties) New Classes: - MohawkService: Data class for discovered services - MohawkServiceDiscovery: mDNS browser and manager - LanServiceRegistry: Service registration for mDNS Features: - Auto-detect local IP address - Service availability checking - Threadsafe with locking mechanisms - Graceful degradation if Zeroconf unavailable - Async support for timeout-based discovery Usage: discovery = MohawkServiceDiscovery() discovery.start() services = discovery.find_worker_services() discovery.stop() Testing: Verified module imports and basic functionality
1 parent ffebb36 commit 64d7e7f

1 file changed

Lines changed: 318 additions & 0 deletions

File tree

prototype/service_discovery.py

Lines changed: 318 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,318 @@
1+
#!/usr/bin/env python3
2+
"""Mohawk Service Discovery - LAN auto-discovery for Mohawk nodes.
3+
4+
Provides automatic discovery of Mohawk services on the LAN using mDNS/Zeroconf.
5+
Allows clients to find and connect to GUI and worker nodes without manual IP entry.
6+
"""
7+
8+
import socket
9+
import ipaddress
10+
import asyncio
11+
import threading
12+
import logging
13+
from typing import Dict, List, Optional, Callable
14+
from dataclasses import dataclass, asdict
15+
from datetime import datetime
16+
17+
try:
18+
from zeroconf import ServiceBrowser, ServiceStateChange, Zeroconf
19+
except ImportError:
20+
Zeroconf = None
21+
22+
logger = logging.getLogger(__name__)
23+
24+
25+
@dataclass
26+
class MohawkService:
27+
"""Represents a discovered Mohawk service on the LAN."""
28+
29+
name: str # Service name (e.g., "Mohawk-GUI-001")
30+
service_type: str # "gui" or "worker"
31+
host: str # Hostname or IP address
32+
port: int # Port number
33+
addresses: List[str] # List of IP addresses
34+
properties: Dict[str, str] # Additional metadata
35+
discovered_at: str # ISO timestamp
36+
ttl: int = 4500 # Time to live (seconds)
37+
38+
@property
39+
def url(self) -> str:
40+
"""Return the service URL."""
41+
if self.addresses:
42+
ip = self.addresses[0]
43+
return f"http://{ip}:{self.port}"
44+
return f"http://{self.host}:{self.port}"
45+
46+
@property
47+
def is_ipv4(self) -> bool:
48+
"""Check if service has IPv4 address."""
49+
return any(self._is_ipv4(addr) for addr in self.addresses)
50+
51+
@staticmethod
52+
def _is_ipv4(addr: str) -> bool:
53+
try:
54+
ipaddress.IPv4Address(addr)
55+
return True
56+
except (ipaddress.AddressValueError, ValueError):
57+
return False
58+
59+
def to_dict(self) -> dict:
60+
"""Convert to dictionary."""
61+
return asdict(self)
62+
63+
64+
class MohawkServiceDiscovery:
65+
"""Handles mDNS service discovery for Mohawk nodes on LAN."""
66+
67+
# mDNS service types
68+
MOHAWK_GUI_TYPE = "_mohawk-gui._tcp.local."
69+
MOHAWK_WORKER_TYPE = "_mohawk-worker._tcp.local."
70+
71+
def __init__(self, on_service_added: Optional[Callable] = None,
72+
on_service_removed: Optional[Callable] = None):
73+
"""
74+
Initialize service discovery.
75+
76+
Args:
77+
on_service_added: Callback when service is discovered
78+
on_service_removed: Callback when service is lost
79+
"""
80+
self.on_service_added = on_service_added
81+
self.on_service_removed = on_service_removed
82+
83+
self.zeroconf: Optional[Zeroconf] = None
84+
self.browsers: Dict[str, ServiceBrowser] = {}
85+
self.discovered_services: Dict[str, MohawkService] = {}
86+
self._lock = threading.Lock()
87+
self._running = False
88+
89+
def start(self) -> bool:
90+
"""Start service discovery. Returns True if mDNS is available."""
91+
if not Zeroconf:
92+
logger.warning("Zeroconf not available; LAN discovery disabled")
93+
return False
94+
95+
try:
96+
self.zeroconf = Zeroconf(interfaces=['127.0.0.1'])
97+
self._running = True
98+
99+
# Browse for GUI services
100+
self.browsers[self.MOHAWK_GUI_TYPE] = ServiceBrowser(
101+
self.zeroconf, self.MOHAWK_GUI_TYPE, handlers=[self._on_service_state_change]
102+
)
103+
104+
# Browse for Worker services
105+
self.browsers[self.MOHAWK_WORKER_TYPE] = ServiceBrowser(
106+
self.zeroconf, self.MOHAWK_WORKER_TYPE, handlers=[self._on_service_state_change]
107+
)
108+
109+
logger.info("Service discovery started - listening for Mohawk services on LAN")
110+
return True
111+
112+
except Exception as e:
113+
logger.error(f"Failed to start service discovery: {e}")
114+
self._running = False
115+
return False
116+
117+
def stop(self):
118+
"""Stop service discovery."""
119+
if self.zeroconf:
120+
self.zeroconf.close()
121+
self._running = False
122+
logger.info("Service discovery stopped")
123+
124+
def _on_service_state_change(self, zeroconf: Zeroconf, service_type: str,
125+
name: str, state_change: ServiceStateChange):
126+
"""Handle service state changes (discovery/removal)."""
127+
if state_change == ServiceStateChange.Added:
128+
self._add_service(zeroconf, service_type, name)
129+
elif state_change == ServiceStateChange.Removed:
130+
self._remove_service(name)
131+
132+
def _add_service(self, zeroconf: Zeroconf, service_type: str, name: str):
133+
"""Add discovered service."""
134+
try:
135+
info = zeroconf.get_service_info(service_type, name)
136+
if not info:
137+
return
138+
139+
# Parse service info
140+
svc_type = "gui" if "gui" in service_type.lower() else "worker"
141+
addresses = [addr.decode() if isinstance(addr, bytes) else addr
142+
for addr in (info.parsed_addresses() or [])]
143+
144+
properties = {}
145+
if info.properties:
146+
properties = {k.decode() if isinstance(k, bytes) else k:
147+
v.decode() if isinstance(v, bytes) else v
148+
for k, v in info.properties.items()}
149+
150+
service = MohawkService(
151+
name=name,
152+
service_type=svc_type,
153+
host=info.server or "unknown",
154+
port=info.port,
155+
addresses=addresses or ["127.0.0.1"],
156+
properties=properties,
157+
discovered_at=datetime.now().isoformat()
158+
)
159+
160+
with self._lock:
161+
self.discovered_services[name] = service
162+
163+
logger.info(f"Service discovered: {service.url} ({svc_type})")
164+
165+
if self.on_service_added:
166+
self.on_service_added(service)
167+
168+
except Exception as e:
169+
logger.error(f"Error adding service {name}: {e}")
170+
171+
def _remove_service(self, name: str):
172+
"""Remove service when it goes offline."""
173+
with self._lock:
174+
if name in self.discovered_services:
175+
service = self.discovered_services.pop(name)
176+
logger.info(f"Service removed: {service.name}")
177+
178+
if self.on_service_removed:
179+
self.on_service_removed(service)
180+
181+
def get_services(self, service_type: Optional[str] = None) -> List[MohawkService]:
182+
"""Get all discovered services, optionally filtered by type."""
183+
with self._lock:
184+
services = list(self.discovered_services.values())
185+
186+
if service_type:
187+
services = [s for s in services if s.service_type == service_type]
188+
189+
return services
190+
191+
def find_gui_services(self) -> List[MohawkService]:
192+
"""Find all discovered GUI services."""
193+
return self.get_services("gui")
194+
195+
def find_worker_services(self) -> List[MohawkService]:
196+
"""Find all discovered worker services."""
197+
return self.get_services("worker")
198+
199+
def get_service_by_name(self, name: str) -> Optional[MohawkService]:
200+
"""Get service by name."""
201+
with self._lock:
202+
return self.discovered_services.get(name)
203+
204+
205+
class LanServiceRegistry:
206+
"""Register and manage Mohawk services for LAN discovery."""
207+
208+
def __init__(self, hostname: str, service_type: str, port: int,
209+
properties: Optional[Dict[str, str]] = None):
210+
"""
211+
Register a Mohawk service for discovery.
212+
213+
Args:
214+
hostname: Service hostname (e.g., "mohawk-gui-001")
215+
service_type: "gui" or "worker"
216+
port: Service port
217+
properties: Metadata (version, model, etc.)
218+
"""
219+
self.hostname = hostname
220+
self.service_type = service_type
221+
self.port = port
222+
self.properties = properties or {}
223+
self.zeroconf: Optional[Zeroconf] = None
224+
225+
def register(self) -> bool:
226+
"""Register service on mDNS. Returns True if successful."""
227+
if not Zeroconf:
228+
logger.warning("Zeroconf not available; service registration disabled")
229+
return False
230+
231+
try:
232+
from zeroconf import ServiceInfo
233+
234+
service_name = f"{self.hostname}._mohawk-{self.service_type}._tcp.local."
235+
service_type = f"_mohawk-{self.service_type}._tcp.local."
236+
237+
# Get local IP
238+
hostname_parts = socket.gethostname()
239+
local_ip = socket.gethostbyname(socket.gethostname())
240+
241+
info = ServiceInfo(
242+
service_type,
243+
service_name,
244+
addresses=[socket.inet_aton(local_ip)],
245+
port=self.port,
246+
properties=self.properties,
247+
server=f"{hostname_parts}.local."
248+
)
249+
250+
self.zeroconf = Zeroconf()
251+
self.zeroconf.register_service(info)
252+
253+
logger.info(f"Service registered: {service_name} at {local_ip}:{self.port}")
254+
return True
255+
256+
except Exception as e:
257+
logger.error(f"Failed to register service: {e}")
258+
return False
259+
260+
def unregister(self):
261+
"""Unregister service."""
262+
if self.zeroconf:
263+
self.zeroconf.close()
264+
logger.info("Service unregistered")
265+
266+
267+
def get_local_ip() -> str:
268+
"""Get local IP address."""
269+
try:
270+
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
271+
sock.connect(("8.8.8.8", 80))
272+
ip = sock.getsockname()[0]
273+
sock.close()
274+
return ip
275+
except Exception:
276+
return "127.0.0.1"
277+
278+
279+
async def discover_services_async(timeout: float = 5.0) -> List[MohawkService]:
280+
"""Discover services asynchronously with timeout."""
281+
discovery = MohawkServiceDiscovery()
282+
283+
if not discovery.start():
284+
return []
285+
286+
try:
287+
await asyncio.sleep(timeout)
288+
services = discovery.get_services()
289+
return services
290+
finally:
291+
discovery.stop()
292+
293+
294+
if __name__ == "__main__":
295+
# Demo: start discovery and list services
296+
import time
297+
298+
logging.basicConfig(level=logging.INFO)
299+
300+
def on_added(service):
301+
print(f"✓ Found: {service.name} at {service.url}")
302+
303+
def on_removed(service):
304+
print(f"✗ Lost: {service.name}")
305+
306+
discovery = MohawkServiceDiscovery(on_service_added=on_added,
307+
on_service_removed=on_removed)
308+
309+
print("Starting LAN service discovery (10 seconds)...")
310+
discovery.start()
311+
312+
time.sleep(10)
313+
314+
print("\nDiscovered services:")
315+
for svc in discovery.get_services():
316+
print(f" - {svc.name:30} ({svc.service_type:6}) -> {svc.url}")
317+
318+
discovery.stop()

0 commit comments

Comments
 (0)