SenSu/services/web_panel/routes/status.py

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))
    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