Coverage for src\core\discovery.py: 0%

1""" 

2设备发现服务 

3使用 Zeroconf 实现局域网内的设备自动发现 

4""" 

5 

6import asyncio 

7import socket 

8import logging 

9from typing import Dict, Any, Optional, Callable 

10from dataclasses import dataclass 

11from contextlib import asynccontextmanager 

12 

13import zeroconf 

14from zeroconf.asyncio import AsyncZeroconf, AsyncServiceInfo, AsyncServiceBrowser 

15 

16from .device import DeviceInfo, DeviceType, DeviceCapability 

17from .config import PAOConfig 

18 

19logger = logging.getLogger(__name__) 

20 

21 

22@dataclass 

23class DiscoveryConfig: 

24 """发现服务配置""" 

25 service_type: str = "_pao._tcp.local." 

26 service_name: str = "pao-device" 

27 domain: str = "local." 

28 port: int = 8765 

29 scan_interval: int = 30 # 扫描间隔（秒） 

30 timeout: int = 10 # 发现超时（秒） 

31 retry_count: int = 3 # 重试次数 

32 

33 

34class DeviceDiscoveryService: 

35 """设备发现服务""" 

36 

37 def __init__( 

38 self, 

39 config: PAOConfig, 

40 device_registry, 

41 on_device_discovered: Optional[Callable[[DeviceInfo], None]] = None, 

42 on_device_lost: Optional[Callable[[str], None]] = None 

43 ): 

44 self.config = config 

45 self.device_registry = device_registry 

46 self.on_device_discovered = on_device_discovered 

47 self.on_device_lost = on_device_lost 

48 

49 self.discovery_config = DiscoveryConfig() 

50 self.zeroconf: Optional[AsyncZeroconf] = None 

51 self.service_browser: Optional[AsyncServiceBrowser] = None 

52 self.service_info: Optional[AsyncServiceInfo] = None 

53 

54 # 运行状态 

55 self.is_running = False 

56 self.discovery_task: Optional[asyncio.Task] = None 

57 

58 # 本地设备信息 

59 self.local_device = self._create_local_device_info() 

60 

61 def _create_local_device_info(self) -> DeviceInfo: 

62 """创建本地设备信息""" 

63 device = DeviceInfo() 

64 device.name = self.config.device_name 

65 device.device_type = self._detect_device_type() 

66 device.capabilities = self._detect_capabilities() 

67 

68 # 获取系统信息 

69 import platform 

70 import psutil 

71 

72 device.os_name = platform.system() 

73 device.os_version = platform.version() 

74 device.cpu_cores = psutil.cpu_count(logical=False) or 1 

75 device.memory_gb = psutil.virtual_memory().total // (1024**3) 

76 

77 # 存储信息 

78 try: 

79 device.storage_gb = psutil.disk_usage('/').total // (1024**3) 

80 except: 

81 device.storage_gb = 0 

82 

83 # MAC地址 

84 try: 

85 import uuid 

86 device.mac_address = ':'.join(['{:02x}'.format((uuid.getnode() >> elements) & 0xff) 

87 for elements in range(0, 8*6, 8)][::-1]) 

88 except: 

89 device.mac_address = "unknown" 

90 

91 return device 

92 

93 def _detect_device_type(self) -> DeviceType: 

94 """检测设备类型""" 

95 import platform 

96 

97 system = platform.system().lower() 

98 machine = platform.machine().lower() 

99 

100 if "linux" in system: 

101 # 可能是服务器或IoT 

102 if "arm" in machine or "aarch" in machine: 

103 return DeviceType.IOT 

104 else: 

105 return DeviceType.SERVER 

106 elif "windows" in system or "darwin" in system: 

107 # 桌面设备 

108 if "laptop" in platform.node().lower(): 

109 return DeviceType.LAPTOP 

110 return DeviceType.DESKTOP 

111 elif "android" in system: 

112 return DeviceType.PHONE 

113 elif "ios" in system: 

114 return DeviceType.TABLET 

115 else: 

116 return DeviceType.UNKNOWN 

117 

118 def _detect_capabilities(self) -> list: 

119 """检测设备能力""" 

120 capabilities = [] 

121 

122 # 所有设备都有计算和网络能力 

123 capabilities.append(DeviceCapability.COMPUTE) 

124 capabilities.append(DeviceCapability.NETWORK) 

125 

126 import platform 

127 import psutil 

128 

129 # 存储能力 

130 if psutil.disk_usage('/').total > 0: 

131 capabilities.append(DeviceCapability.STORAGE) 

132 

133 # 内存能力 

134 if psutil.virtual_memory().total > 0: 

135 capabilities.append(DeviceCapability.MEMORY) 

136 

137 # 显示能力（根据平台判断） 

138 system = platform.system().lower() 

139 if system in ['windows', 'darwin', 'linux']: 

140 # 尝试检测是否有显示 

141 try: 

142 import os 

143 if 'DISPLAY' in os.environ or system != 'linux': 

144 capabilities.append(DeviceCapability.DISPLAY) 

145 capabilities.append(DeviceCapability.INPUT) 

146 except: 

147 pass 

148 

149 return capabilities 

150 

151 async def start(self) -> None: 

152 """启动发现服务""" 

153 if self.is_running: 

154 logger.warning("发现服务已经在运行") 

155 return 

156 

157 logger.info(f"启动设备发现服务: {self.local_device.name}") 

158 

159 # 创建 Zeroconf 实例 

160 self.zeroconf = AsyncZeroconf() 

161 

162 # 注册本地服务 

163 await self._register_service() 

164 

165 # 启动服务浏览器 

166 await self._start_browser() 

167 

168 # 设置运行状态 

169 self.is_running = True 

170 

171 # 启动定期扫描任务 

172 self.discovery_task = asyncio.create_task(self._periodic_discovery()) 

173 

174 logger.info("设备发现服务启动完成") 

175 

176 async def stop(self) -> None: 

177 """停止发现服务""" 

178 if not self.is_running: 

179 return 

180 

181 logger.info("停止设备发现服务") 

182 

183 # 取消定期任务 

184 if self.discovery_task and not self.discovery_task.done(): 

185 self.discovery_task.cancel() 

186 try: 

187 await self.discovery_task 

188 except asyncio.CancelledError: 

189 pass 

190 

191 # 取消服务注册 

192 if self.service_info and self.zeroconf: 

193 await self.zeroconf.async_unregister_service(self.service_info) 

194 

195 # 关闭浏览器 

196 if self.service_browser: 

197 self.service_browser.cancel() 

198 

199 # 关闭 Zeroconf 

200 if self.zeroconf: 

201 await self.zeroconf.async_close() 

202 

203 self.is_running = False 

204 logger.info("设备发现服务已停止") 

205 

206 async def _register_service(self) -> None: 

207 """注册本地设备服务""" 

208 service_type = self.discovery_config.service_type 

209 service_name = f"{self.local_device.name}.{service_type}" 

210 

211 # 服务属性 

212 properties = self.local_device.to_dict() 

213 

214 # 创建服务信息 

215 self.service_info = AsyncServiceInfo( 

216 service_type, 

217 service_name, 

218 addresses=[socket.inet_aton(self.local_device.ip_address)], 

219 port=self.discovery_config.port, 

220 properties=properties, 

221 server=f"{self.local_device.hostname}.local." 

222 ) 

223 

224 # 注册服务 

225 if self.zeroconf: 

226 await self.zeroconf.async_register_service(self.service_info) 

227 logger.info(f"已注册服务: {service_name}") 

228 

229 async def _start_browser(self) -> None: 

230 """启动服务浏览器""" 

231 if not self.zeroconf: 

232 return 

233 

234 # 创建浏览器回调 

235 listener = self._create_service_listener() 

236 

237 # 启动浏览器 

238 self.service_browser = AsyncServiceBrowser( 

239 self.zeroconf.zeroconf, 

240 self.discovery_config.service_type, 

241 listener 

242 ) 

243 

244 logger.info("服务浏览器已启动") 

245 

246 def _create_service_listener(self): 

247 """创建服务监听器""" 

248 class ServiceListener: 

249 def __init__(self, discovery_service): 

250 self.discovery_service = discovery_service 

251 

252 def add_service(self, zeroconf, service_type, name): 

253 asyncio.create_task( 

254 self.discovery_service._on_service_added(zeroconf, service_type, name) 

255 ) 

256 

257 def remove_service(self, zeroconf, service_type, name): 

258 asyncio.create_task( 

259 self.discovery_service._on_service_removed(zeroconf, service_type, name) 

260 ) 

261 

262 def update_service(self, zeroconf, service_type, name): 

263 asyncio.create_task( 

264 self.discovery_service._on_service_updated(zeroconf, service_type, name) 

265 ) 

266 

267 return ServiceListener(self) 

268 

269 async def _on_service_added(self, zeroconf, service_type, name) -> None: 

270 """处理服务添加事件""" 

271 logger.info(f"发现新服务: {name}") 

272 

273 # 获取服务信息 

274 info = await AsyncServiceInfo.async_from_zeroconf(zeroconf, service_type, name) 

275 if info: 

276 await self._process_service_info(info) 

277 

278 async def _on_service_removed(self, zeroconf, service_type, name) -> None: 

279 """处理服务移除事件""" 

280 logger.info(f"服务移除: {name}") 

281 

282 # 从设备ID中提取设备ID 

283 device_id = name.split('.')[0] 

284 

285 # 通知设备丢失 

286 if self.on_device_lost: 

287 self.on_device_lost(device_id) 

288 

289 async def _on_service_updated(self, zeroconf, service_type, name) -> None: 

290 """处理服务更新事件""" 

291 logger.debug(f"服务更新: {name}") 

292 

293 # 重新获取服务信息 

294 info = await AsyncServiceInfo.async_from_zeroconf(zeroconf, service_type, name) 

295 if info: 

296 await self._process_service_info(info) 

297 

298 async def _process_service_info(self, info: AsyncServiceInfo) -> None: 

299 """处理服务信息""" 

300 try: 

301 # 解析设备信息 

302 properties = info.properties or {} 

303 

304 # 转换为设备信息 

305 device_info = DeviceInfo.from_dict({ 

306 **properties, 

307 "ip_address": socket.inet_ntoa(info.addresses[0]) if info.addresses else "0.0.0.0", 

308 "port": info.port, 

309 "hostname": info.server.rstrip('.local.') 

310 }) 

311 

312 # 更新最后出现时间 

313 device_info.update_last_seen() 

314 

315 # 注册设备 

316 self.device_registry.register_device(device_info) 

317 

318 logger.info(f"发现设备: {device_info.name} ({device_info.device_type.value})") 

319 

320 # 通知回调 

321 if self.on_device_discovered: 

322 self.on_device_discovered(device_info) 

323 

324 except Exception as e: 

325 logger.error(f"处理服务信息失败: {e}") 

326 

327 async def _periodic_discovery(self) -> None: 

328 """定期发现任务""" 

329 while self.is_running: 

330 try: 

331 await asyncio.sleep(self.discovery_config.scan_interval) 

332 

333 # 主动扫描 

334 await self._active_scan() 

335 

336 # 清理离线设备 

337 self.device_registry.cleanup_offline_devices() 

338 

339 except asyncio.CancelledError: 

340 break 

341 except Exception as e: 

342 logger.error(f"定期发现任务出错: {e}") 

343 

344 async def _active_scan(self) -> None: 

345 """主动扫描网络中的设备""" 

346 if not self.zeroconf: 

347 return 

348 

349 # 使用 Zeroconf 主动查询服务 

350 try: 

351 services = await self.zeroconf.async_get_service_info( 

352 self.discovery_config.service_type, 

353 f"{self.local_device.name}.{self.discovery_config.service_type}" 

354 ) 

355 

356 if services: 

357 await self._process_service_info(services) 

358 

359 except Exception as e: 

360 logger.debug(f"主动扫描出错: {e}") 

361 

362 async def discover_devices(self, timeout: Optional[int] = None) -> list: 

363 """主动发现设备（阻塞方式）""" 

364 if timeout is None: 

365 timeout = self.discovery_config.timeout 

366 

367 # 记录当前设备数量 

368 initial_count = len(self.device_registry.list_devices()) 

369 

370 # 等待一段时间 

371 await asyncio.sleep(timeout) 

372 

373 # 返回新发现的设备 

374 current_count = len(self.device_registry.list_devices()) 

375 new_devices = current_count - initial_count 

376 

377 logger.info(f"发现 {new_devices} 个新设备") 

378 return self.device_registry.list_devices() 

379 

380 

381@asynccontextmanager 

382async def discovery_service(config: PAOConfig, device_registry): 

383 """发现服务的上下文管理器""" 

384 service = DeviceDiscoveryService(config, device_registry) 

385 try: 

386 await service.start() 

387 yield service 

388 finally: 

389 await service.stop() 

390 

391 

392async def demo_discovery(): 

393 """演示设备发现功能""" 

394 from .config import load_config 

395 

396 config = load_config() 

397 from .device import DeviceRegistry 

398 

399 device_registry = DeviceRegistry() 

400 

401 # 创建设备发现服务 

402 async with discovery_service(config, device_registry) as service: 

403 print(f"本地设备: {service.local_device.name}") 

404 print("开始发现设备...") 

405 

406 # 等待10秒发现设备 

407 devices = await service.discover_devices(timeout=10) 

408 

409 print(f"\n发现 {len(devices)} 个设备:") 

410 for device in devices: 

411 print(f" - {device.name} ({device.device_type.value}) " 

412 f"[{device.ip_address}:{device.port}]") 

413 

414 # 保持运行 

415 print("\n按 Ctrl+C 停止...") 

416 try: 

417 while True: 

418 await asyncio.sleep(1) 

419 except KeyboardInterrupt: 

420 print("停止发现服务") 

421 

422 

423if __name__ == "__main__": 

424 import asyncio 

425 asyncio.run(demo_discovery())