import time import json import asyncio import logging from aiohttp import web from ..utils.system_info import SystemInfoCollector from ..utils.auth import panel_auth collector = SystemInfoCollector() logger = logging.getLogger(__name__) # Track active system-status WS clients _sys_ws_clients: set = set() def _ws_auth_wrapper(handler): """WebSocket 鉴权 — Cookie 优先(同源自动发送), query token 备用""" async def wrapper(request): # 1. Cookie (浏览器同源自动发送, 支持 HttpOnly) token = request.cookies.get("panel_token", "") # 2. Query string token 备用 (跨域/非浏览器客户端) if not token: token = request.query.get("token", "") # 3. API Key 验证 session_store = request.app.get("panel_session_store", {}) is_valid = token and token in session_store if not is_valid and token: from services.web_panel.routes.apikeys import validate_api_key is_valid = validate_api_key(token) is not None if not is_valid: ws = web.WebSocketResponse() await ws.prepare(request) await ws.send_str(json.dumps({"error": "Unauthorized"})) await ws.close(code=4001, message="Unauthorized") return ws return await handler(request) return wrapper def setup_routes(app, prefix=''): app.router.add_get(f'{prefix}/api/framework', panel_auth(get_framework)) app.router.add_get(f'{prefix}/api/system', panel_auth(get_system)) app.router.add_get(f'{prefix}/api/system/ws', _ws_auth_wrapper(sys_ws_handler)) app.router.add_get(f'{prefix}/api/updates', panel_auth(get_updates)) app.router.add_post(f'{prefix}/api/updates/check', panel_auth(check_now)) app.router.add_post(f'{prefix}/api/updates/apply', panel_auth(apply_update)) logger.info(f"📡 系统状态WS端点已注册: {prefix}/api/system/ws (已加认证)") def _read_version(): """Read version from config file (shared helper)""" ver = "v0.6.0" try: import yaml, os config_path = os.path.join(os.path.dirname(__file__), "..", "..", "..", "config", "framework", "base_config.yaml") with open(config_path) as f: cfg = yaml.safe_load(f) ver = cfg.get("framework", {}).get("version", ver) except: pass return ver def _get_framework_data(sm): """Collect framework status (reused by HTTP and WS handlers)""" ver = _read_version() ps = sm.get_service("plugin") if sm else None uptime = int(time.time() - getattr(sm, 'start_time', time.time())) return { "version": ver, "uptime": uptime, "plugins": len(ps.plugins) if ps else 0 } async def get_framework(req): sm = req.app.get('service_manager') if not sm: return web.json_response({"error": "Missing"}, status=500) return web.json_response(_get_framework_data(sm)) def _get_collector(req): """获取共享的 SystemInfoCollector(优先从 service_manager),fallback 到模块级实例""" sm = req.app.get('service_manager') if sm: try: shared = sm.get_service("sys_collector") if shared: return shared except Exception: pass return collector async def get_system(req): c = _get_collector(req) return web.json_response(c.get_all()) async def sys_ws_handler(req): """WebSocket push endpoint: system+framwork stats every 2s (no logging per message)""" ws = web.WebSocketResponse(heartbeat=30.0) await ws.prepare(req) _sys_ws_clients.add(ws) async def push(): """Send one snapshot to this client (silent on error)""" try: sm = req.app.get('service_manager') c = _get_collector(req) payload = json.dumps({ "type": "sys", "system": c.get_all(), "framework": _get_framework_data(sm) }, ensure_ascii=False) if not ws.closed: await ws.send_str(payload) except Exception: pass # Background push loop — runs until client disconnects async def push_loop(): while not ws.closed: await push() await asyncio.sleep(2) task = asyncio.ensure_future(push_loop()) try: async for msg in ws: if msg.type == web.WSMsgType.ERROR: break finally: task.cancel() try: await task except asyncio.CancelledError: pass _sys_ws_clients.discard(ws) return ws async def get_updates(req): """获取更新状态""" sm = req.app.get("service_manager") us = sm.get_service("update") if sm else None if not us: return web.json_response({"error": "更新服务未就绪"}, status=503) return web.json_response(us.status) async def check_now(req): """手动触发更新检查""" sm = req.app.get("service_manager") us = sm.get_service("update") if sm else None if not us: return web.json_response({"error": "更新服务未就绪"}, status=503) await us.check_now() return web.json_response(us.status) async def apply_update(req): """执行更新""" sm = req.app.get("service_manager") us = sm.get_service("update") if sm else None if not us: return web.json_response({"error": "更新服务未就绪"}, status=503) result = await us.apply_update() return web.json_response(result)