SenSu/services/internet_service.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
import asyncio
from typing import Dict, List, Callable, Any, Optional
from pathlib import Path
import aiohttp
from aiohttp import web
import json
import ssl

logger = logging.getLogger(__name__)

class InternetService:
    """网络服务 - 管理HTTP和WebSocket服务器"""
    
    def __init__(self, config: Dict, service_manager):
        self.config = config
        self.service_manager = service_manager
        self.http_app = web.Application()
        self._setup_security_middleware()
        self.http_runner = None
        self.ws_connections: Dict[str, List] = {}
        self.plugin_routes: Dict[str, List] = {}
        self.is_running = False
        
        # 从配置获取端口
        internet_config = config.get('internet', {})
        
        ws_config = internet_config.get('websocket', {})
        self.ws_host = ws_config.get('host', '0.0.0.0')
        self.ws_port = ws_config.get('port', 8765)

        http_config = internet_config.get('http', {})        
        self.http_host = http_config.get('host', '0.0.0.0')
        self.http_port = http_config.get('port', 8000)
        
        logger.debug("InternetService初始化开始")
    
    async def start(self):
        """启动网络服务 - 精准错误处理与端口复用版"""
        try:
            logger.info("启动网络服务")
            self._setup_default_routes()
            
            # 初始化 AppRunner
            self.http_runner = web.AppRunner(self.http_app)
            await self.http_runner.setup()
            
            # 1. 启动 HTTP 站点
            try:
                self.site = web.TCPSite(
                    self.http_runner, self.http_host, self.http_port, 
                    reuse_address=True, reuse_port=True
                )
                await self.site.start()
                logger.info(f"✅ HTTP 服务已绑定: {self.http_host}:{self.http_port}")
            except OSError as e:
                logger.error(f"❌ HTTP 端口 {self.http_port} 绑定失败: {e}")
                await self.http_runner.cleanup()
                return False
            
            # 2. 启动 WebSocket 站点 (独立端口)
            try:
                self.ws_site = web.TCPSite(
                    self.http_runner, self.ws_host, self.ws_port, 
                    reuse_address=True, reuse_port=True
                )
                await self.ws_site.start()
                logger.info(f"✅ WebSocket 服务已绑定: {self.ws_host}:{self.ws_port}")
            except OSError as e:
                logger.error(f"❌ WebSocket 端口 {self.ws_port} 绑定失败: {e}")
                logger.warning("💡 WS端口可能处于 TIME_WAIT，请等待30秒或更换 config 中的 websocket.port")
                await self.http_runner.cleanup() # 回滚已启动的 HTTP
                return False
            
            await self.save_network_config()
            self.is_running = True
            logger.info("🌐 网络服务启动完成")
            return True
            
        except Exception as e:
            logger.error(f"❌ 启动网络服务异常: {str(e)}", exc_info=True)
            return False

    
    async def check_service_health(self) -> Dict[str, Any]:
        """检查服务健康状况"""
        try:
            import socket
            
            health_info = {
                "is_running": self.is_running,
                "http_port": self.http_port,
                "websocket_port": self.ws_port,
                "http_active": False,
                "dependencies_available": self._check_dependencies(),
                "error": None
            }
            
            # 检查端口是否在监听
            if self.is_running:
                try:
                    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                    sock.settimeout(2)
                    result = sock.connect_ex(('localhost', self.http_port))
                    sock.close()
                    health_info["http_active"] = (result == 0)
                except Exception as e:
                    health_info["error"] = f"端口检查失败: {str(e)}"
            
            return health_info
            
        except Exception as e:
            return {
                "is_running": False,
                "error": f"健康检查失败: {str(e)}"
            }
    
    def _check_dependencies(self) -> bool:
        """检查必要的依赖包"""
        try:
            import aiohttp
            import yaml
            return True
        except ImportError as e:
            logger.error(f"❌ 缺少依赖包: {str(e)}")
            return False    
        
    async def save_network_config(self):
        """保存网络配置"""
        try:
            config_path = Path("config") / "services" / "network_routes.yaml"
            config_path.parent.mkdir(parents=True, exist_ok=True)
            
            config_data = {
                "plugin_routes": self.get_plugin_routes(),
                "last_updated": asyncio.get_event_loop().time(),
                "http_port": self.http_port,
                "websocket_port": self.ws_port
            }
            
            import yaml
            with open(config_path, 'w', encoding='utf-8') as f:
                yaml.dump(config_data, f, default_flow_style=False, allow_unicode=True, indent=2)
            
            logger.info(f"网络配置已保存: {config_path}")
            
        except Exception as e:
            logger.error(f"保存网络配置时出错: {str(e)}")
        
    def _setup_security_middleware(self):
        """注入安全响应头 + 错误脱敏中间件"""

        @web.middleware
        async def security_headers(request, handler):
            try:
                resp = await handler(request)
            except web.HTTPException:
                raise
            except Exception as e:
                # 生产模式: 脱敏错误，仅返回通用消息，详细信息写日志
                logger.error(f"未捕获异常 {request.method} {request.path}: {e}", exc_info=True)
                resp = web.json_response(
                    {"error": "Internal server error"}, status=500
                )
            resp.headers.setdefault("X-Content-Type-Options", "nosniff")
            resp.headers.setdefault("X-Frame-Options", "DENY")
            resp.headers.setdefault("X-XSS-Protection", "1; mode=block")
            resp.headers.setdefault("Referrer-Policy", "strict-origin-when-cross-origin")
            return resp

        self.http_app.middlewares.append(security_headers)

    def _setup_default_routes(self):
        """设置默认路由"""
        from services.web_panel.utils.auth import panel_auth

        # 健康检查端点 (公开，不暴露内部信息)
        self.http_app.router.add_get('/health', self._handle_health_check)

        # 插件API端点 (加认证)
        self.http_app.router.add_get('/api/plugins', panel_auth(self._handle_get_plugins))
        self.http_app.router.add_get('/api/commands', panel_auth(self._handle_get_commands))

        # 数据接收端点 (加认证)
        self.http_app.router.add_post('/api/data', panel_auth(self._handle_data_receive))

        logger.debug("默认路由设置完成 (已加认证)")
    
    async def register_plugin_route(self, plugin_name: str, route_path: str, 
                                  handler: Callable, methods: List[str] = ["GET"],
                                  require_auth: bool = True):
        """为插件注册HTTP路由 - 修复冻结路由器问题"""
        try:
            # 规范化路径
            if not route_path.startswith('/'):
                route_path = '/' + route_path
            
            full_path = f"/{plugin_name}{route_path}"
            
            # 创建包装器处理权限验证
            async def wrapped_handler(request):
                try:
                    # 权限验证
                    if require_auth:
                        auth_result = await self._check_plugin_auth(plugin_name, request)
                        if not auth_result['allowed']:
                            return web.json_response(
                                {"error": "权限不足", "details": auth_result['reason']}, 
                                status=403
                            )
                    
                    # 调用插件处理器
                    return await handler(request)
                    
                except Exception as e:
                    logger.error(f"插件路由处理出错 {full_path}: {str(e)}")
                    return web.json_response(
                        {"error": "内部服务器错误", "details": str(e)}, 
                        status=500
                    )
            
            # 检查路由器是否已冻结
            if hasattr(self.http_app.router, '_frozen') and self.http_app.router._frozen:
                logger.warning(f"路由器已冻结，无法注册新路由: {full_path}")
                logger.info("💡 建议: 在启动网络服务前注册所有插件路由")
                return
            
            # 注册路由
            for method in methods:
                self.http_app.router.add_route(method.upper(), full_path, wrapped_handler)
            
            # 记录路由信息
            if plugin_name not in self.plugin_routes:
                self.plugin_routes[plugin_name] = []
            
            self.plugin_routes[plugin_name].append({
                'path': full_path,
                'methods': methods,
                'require_auth': require_auth
            })
            
            logger.info(f"注册插件路由: {plugin_name} -> {full_path} [{','.join(methods)}]")
            
        except RuntimeError as e:
            if "frozen router" in str(e):
                logger.error(f"❌ 无法注册路由 {full_path}: 路由器已冻结")
                logger.info("💡 解决方案: 在启动网络服务前注册插件路由")
            else:
                logger.error(f"注册插件路由时出错: {str(e)}", exc_info=True)
            raise
        except Exception as e:
            logger.error(f"注册插件路由时出错: {str(e)}", exc_info=True)
            raise
        
    async def register_plugin_websocket(self, plugin_name: str, ws_path: str,
                                      handler: Callable, require_auth: bool = True):
        """为插件注册WebSocket路由"""
        try:
            # 规范化路径
            if not ws_path.startswith('/'):
                ws_path = '/' + ws_path
            
            full_path = f"/plugin/{plugin_name}/ws{ws_path}"
            
            async def websocket_handler(request):
                try:
                    # 权限验证
                    if require_auth:
                        auth_result = await self._check_plugin_auth(plugin_name, request)
                        if not auth_result['allowed']:
                            return web.json_response(
                                {"error": "WebSocket连接权限不足"}, 
                                status=403
                            )
                    
                    # 建立WebSocket连接
                    ws = web.WebSocketResponse()
                    await ws.prepare(request)
                    
                    # 记录连接
                    connection_id = f"{plugin_name}_{id(ws)}"
                    if plugin_name not in self.ws_connections:
                        self.ws_connections[plugin_name] = []
                    self.ws_connections[plugin_name].append(ws)
                    
                    logger.debug(f"WebSocket连接建立: {connection_id}")
                    
                    # 调用插件处理器
                    await handler(ws, request)
                    
                    # 清理连接
                    self.ws_connections[plugin_name].remove(ws)
                    
                    return ws
                    
                except Exception as e:
                    logger.error(f"WebSocket处理出错 {full_path}: {str(e)}")
                    return web.json_response(
                        {"error": "WebSocket连接失败"}, 
                        status=500
                    )
            
            # 注册WebSocket路由
            self.http_app.router.add_route('GET', full_path, websocket_handler)
            
            # 记录路由信息
            if plugin_name not in self.plugin_routes:
                self.plugin_routes[plugin_name] = []
            
            self.plugin_routes[plugin_name].append({
                'path': full_path,
                'methods': ['WEBSOCKET'],
                'require_auth': require_auth
            })
            
            logger.info(f"注册插件WebSocket: {plugin_name} -> {full_path}")
            
        except Exception as e:
            logger.error(f"注册插件WebSocket时出错: {str(e)}", exc_info=True)
            raise
    
    async def _check_plugin_auth(self, plugin_name: str, request) -> Dict[str, Any]:
        """检查插件路由权限 — 先验证用户身份，再检查插件权限"""
        try:
            # 1. 验证用户身份 (panel token)
            token = request.cookies.get("panel_token")
            if not token:
                auth_hdr = request.headers.get("Authorization", "")
                if auth_hdr.startswith("Bearer "):
                    token = auth_hdr.split(" ", 1)[1]
            if not token:
                return {"allowed": False, "reason": "未认证"}

            session_store = request.app.get("panel_session_store", {})
            if token not in session_store:
                # 回退到 API Key 验证
                from services.web_panel.routes.apikeys import validate_api_key
                if not validate_api_key(token):
                    return {"allowed": False, "reason": "会话无效或已过期"}

            # 2. 检查插件是否有网络访问权限
            permission_service = self.service_manager.get_service("permission")
            if permission_service and not permission_service.has_permission(
                plugin_name, "plugin.network.access"
            ):
                return {"allowed": False, "reason": "插件没有网络访问权限"}

            return {"allowed": True, "reason": "权限验证通过"}

        except Exception as e:
            logger.error(f"权限检查时出错: {str(e)}")
            return {"allowed": False, "reason": "权限检查失败"}
    
    async def broadcast_to_websockets(self, plugin_name: str, message: Dict):
        """向插件的所有WebSocket连接广播消息"""
        try:
            if plugin_name not in self.ws_connections:
                return
            
            message_json = json.dumps(message, ensure_ascii=False)
            disconnected = []
            
            for ws in self.ws_connections[plugin_name]:
                try:
                    if not ws.closed:
                        await ws.send_str(message_json)
                    else:
                        disconnected.append(ws)
                except Exception as e:
                    logger.error(f"WebSocket广播消息失败: {str(e)}")
                    disconnected.append(ws)
            
            # 清理断开连接的WebSocket
            for ws in disconnected:
                self.ws_connections[plugin_name].remove(ws)
            
            logger.debug(f"WebSocket广播完成: {plugin_name} -> {len(self.ws_connections[plugin_name])} 个连接")
            
        except Exception as e:
            logger.error(f"WebSocket广播时出错: {str(e)}", exc_info=True)
    
    # 默认路由处理器
    async def _handle_health_check(self, request):
        """健康检查端点"""
        return web.json_response({
            "status": "healthy",
            "service": "internet",
            "timestamp": asyncio.get_event_loop().time()
        })
    
    async def _handle_get_plugins(self, request):
        """获取插件列表"""
        try:
            plugin_service = self.service_manager.get_service("plugin")
            if not plugin_service:
                return web.json_response({"error": "插件服务不可用"}, status=503)
            
            plugins_info = []
            for name, info in plugin_service.plugin_info.items():
                plugins_info.append({
                    "name": name,
                    "version": info.version,
                    "description": info.description,
                    "enabled": info.enabled,
                    "loaded": info.loaded
                })
            
            return web.json_response({
                "plugins": plugins_info,
                "count": len(plugins_info)
            })
            
        except Exception as e:
            logger.error(f"获取插件列表时出错: {str(e)}")
            return web.json_response({"error": "内部服务器错误"}, status=500)
    
    async def _handle_get_commands(self, request):
        """获取命令列表"""
        try:
            command_service = self.service_manager.get_service("command")
            if not command_service:
                return web.json_response({"error": "命令服务不可用"}, status=503)
            
            command_list = command_service.get_command_list()
            
            return web.json_response({
                "commands": command_list,
                "count": len(command_list)
            })
            
        except Exception as e:
            logger.error(f"获取命令列表时出错: {str(e)}")
            return web.json_response({"error": "内部服务器错误"}, status=500)
    
    async def _handle_data_receive(self, request):
        """处理跨端数据传输"""
        try:
            data = await request.json()
            
            # 获取插件桥接服务
            plugin_bridge = self.service_manager.get_service("plugin_bridge")
            if plugin_bridge:
                # 广播数据到所有插件
                await plugin_bridge.broadcast_to_plugins(
                    "network.data.receive",
                    {
                        "source": request.remote,
                        "data": data,
                        "timestamp": asyncio.get_event_loop().time()
                    }
                )
            
            return web.json_response({
                "status": "success",
                "message": "数据接收成功",
                "timestamp": asyncio.get_event_loop().time()
            })
            
        except Exception as e:
            logger.error(f"处理跨端数据时出错: {str(e)}")
            return web.json_response({
                "error": "数据接收失败",
                "details": str(e)
            }, status=400)
    
    def get_plugin_routes(self, plugin_name: str = None) -> Dict:
        """获取插件路由信息"""
        if plugin_name:
            return self.plugin_routes.get(plugin_name, [])
        else:
            return self.plugin_routes.copy()
    
    async def shutdown(self):
        """关闭网络服务"""
        try:
            logger.info("关闭网络服务")
            self.is_running = False
            
            # 关闭所有WebSocket连接
            for plugin_name, connections in self.ws_connections.items():
                for ws in connections:
                    if not ws.closed:
                        await ws.close()
                self.ws_connections[plugin_name].clear()
            
            # 关闭HTTP服务器
            if self.http_runner:
                await self.http_runner.cleanup()
            
            logger.debug("网络服务关闭完成")
            
        except Exception as e:
            logger.error(f"关闭网络服务时出错: {str(e)}", exc_info=True)