ForcomeBot 2.2.4__py3-none-any.whl

src/api/websocket.py

@@ -0,0 +1,280 @@
+"""WebSocket推送模块 - 实时状态更新
+
+功能：
+- WebSocketManager 连接管理
+- 日志实时推送
+- 状态变更推送
+"""
+import asyncio
+import logging
+from typing import Set, Optional, Dict, Any, TYPE_CHECKING
+
+from fastapi import WebSocket, WebSocketDisconnect
+
+if TYPE_CHECKING:
+    from ..core.log_collector import LogCollector
+    from ..clients.langbot import LangBotClient
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocketManager:
+    """WebSocket连接管理器"""
+    
+    def __init__(self, max_connections: int = 100):
+        """初始化WebSocket管理器
+        
+        Args:
+            max_connections: 最大连接数
+        """
+        self._connections: Set[WebSocket] = set()
+        self._max_connections = max_connections
+        self._lock = asyncio.Lock()
+    
+    async def connect(self, websocket: WebSocket) -> bool:
+        """接受WebSocket连接
+        
+        Args:
+            websocket: WebSocket连接
+            
+        Returns:
+            是否成功连接
+        """
+        async with self._lock:
+            if len(self._connections) >= self._max_connections:
+                logger.warning(f"WebSocket连接数已达上限: {self._max_connections}")
+                await websocket.close(code=1013, reason="连接数已达上限")
+                return False
+            
+            await websocket.accept()
+            self._connections.add(websocket)
+            logger.info(f"WebSocket连接已建立，当前连接数: {len(self._connections)}")
+            return True
+    
+    def disconnect(self, websocket: WebSocket):
+        """断开WebSocket连接
+        
+        Args:
+            websocket: WebSocket连接
+        """
+        self._connections.discard(websocket)
+        logger.info(f"WebSocket连接已断开，当前连接数: {len(self._connections)}")
+    
+    async def broadcast(self, message: Dict[str, Any]):
+        """广播消息给所有连接
+        
+        Args:
+            message: 要广播的消息
+        """
+        if not self._connections:
+            return
+        
+        dead_connections = set()
+        
+        for ws in self._connections.copy():
+            try:
+                await ws.send_json(message)
+            except Exception as e:
+                logger.debug(f"发送WebSocket消息失败: {e}")
+                dead_connections.add(ws)
+        
+        # 清理失效连接
+        for ws in dead_connections:
+            self._connections.discard(ws)
+    
+    async def send_to(self, websocket: WebSocket, message: Dict[str, Any]) -> bool:
+        """发送消息给指定连接
+        
+        Args:
+            websocket: 目标WebSocket连接
+            message: 要发送的消息
+            
+        Returns:
+            是否发送成功
+        """
+        try:
+            await websocket.send_json(message)
+            return True
+        except Exception as e:
+            logger.debug(f"发送WebSocket消息失败: {e}")
+            self._connections.discard(websocket)
+            return False
+    
+    @property
+    def connection_count(self) -> int:
+        """获取当前连接数"""
+        return len(self._connections)
+
+
+# 全局WebSocket管理器实例
+ws_manager = WebSocketManager()
+
+# 全局引用（由main.py设置）
+_log_collector: Optional["LogCollector"] = None
+_langbot_client: Optional["LangBotClient"] = None
+
+
+def set_websocket_dependencies(
+    log_collector: "LogCollector",
+    langbot_client: "LangBotClient"
+):
+    """设置WebSocket依赖
+    
+    Args:
+        log_collector: 日志收集器
+        langbot_client: LangBot客户端
+    """
+    global _log_collector, _langbot_client
+    _log_collector = log_collector
+    _langbot_client = langbot_client
+    logger.info("WebSocket依赖已设置")
+
+
+async def websocket_endpoint(websocket: WebSocket):
+    """WebSocket端点处理函数
+    
+    处理WebSocket连接，推送日志更新和状态变更
+    """
+    # 尝试连接
+    if not await ws_manager.connect(websocket):
+        return
+    
+    # 订阅日志更新
+    log_queue = None
+    if _log_collector:
+        log_queue = _log_collector.subscribe()
+    
+    # 启动状态推送任务
+    status_task = asyncio.create_task(_push_status_updates(websocket))
+    
+    try:
+        # 发送初始状态
+        await _send_initial_status(websocket)
+        
+        # 主循环：推送日志更新
+        while True:
+            if log_queue:
+                try:
+                    # 等待日志更新（带超时，以便检查连接状态）
+                    log_entry = await asyncio.wait_for(
+                        log_queue.get(), 
+                        timeout=30.0
+                    )
+                    await ws_manager.send_to(websocket, {
+                        "type": "log",
+                        "data": log_entry
+                    })
+                except asyncio.TimeoutError:
+                    # 发送心跳
+                    try:
+                        await websocket.send_json({"type": "ping"})
+                    except:
+                        break
+            else:
+                # 没有日志收集器，只发送心跳
+                await asyncio.sleep(30)
+                try:
+                    await websocket.send_json({"type": "ping"})
+                except:
+                    break
+                    
+    except WebSocketDisconnect:
+        logger.debug("WebSocket客户端断开连接")
+    except Exception as e:
+        logger.error(f"WebSocket处理异常: {e}")
+    finally:
+        # 取消状态推送任务
+        status_task.cancel()
+        try:
+            await status_task
+        except asyncio.CancelledError:
+            pass
+        
+        # 取消日志订阅
+        if log_queue and _log_collector:
+            _log_collector.unsubscribe(log_queue)
+        
+        # 断开连接
+        ws_manager.disconnect(websocket)
+
+
+async def _send_initial_status(websocket: WebSocket):
+    """发送初始状态
+    
+    Args:
+        websocket: WebSocket连接
+    """
+    status = _get_current_status()
+    await ws_manager.send_to(websocket, {
+        "type": "status",
+        "data": status
+    })
+
+
+async def _push_status_updates(websocket: WebSocket):
+    """定期推送状态更新
+    
+    Args:
+        websocket: WebSocket连接
+    """
+    last_status = None
+    
+    try:
+        while True:
+            await asyncio.sleep(5)  # 每5秒检查一次状态变化
+            
+            current_status = _get_current_status()
+            
+            # 只在状态变化时推送
+            if current_status != last_status:
+                success = await ws_manager.send_to(websocket, {
+                    "type": "status",
+                    "data": current_status
+                })
+                if not success:
+                    break
+                last_status = current_status
+                
+    except asyncio.CancelledError:
+        pass
+    except Exception as e:
+        logger.debug(f"状态推送任务异常: {e}")
+
+
+def _get_current_status() -> Dict[str, Any]:
+    """获取当前状态
+    
+    Returns:
+        状态字典
+    """
+    return {
+        "langbot_connected": _langbot_client.is_connected if _langbot_client else False,
+        "langbot_reconnecting": _langbot_client.is_reconnecting if _langbot_client else False,
+        "websocket_connections": ws_manager.connection_count
+    }
+
+
+async def broadcast_status_change(status_type: str, data: Dict[str, Any]):
+    """广播状态变更
+    
+    Args:
+        status_type: 状态类型（如 langbot_connected, config_updated）
+        data: 状态数据
+    """
+    await ws_manager.broadcast({
+        "type": "status_change",
+        "status_type": status_type,
+        "data": data
+    })
+
+
+async def broadcast_log(log_entry: Dict[str, Any]):
+    """广播日志条目
+    
+    Args:
+        log_entry: 日志条目
+    """
+    await ws_manager.broadcast({
+        "type": "log",
+        "data": log_entry
+    })

src/auth/__init__.py

@@ -0,0 +1,33 @@
+"""认证模块"""
+from .database import Database, get_db, init_database
+from .models import User, OperationLog
+from .jwt_handler import JWTHandler
+from .dingtalk import DingTalkClient, init_dingtalk_client, get_dingtalk_client
+from .middleware import (
+    AuthMiddleware,
+    get_current_user,
+    get_optional_user,
+    init_auth,
+    is_auth_enabled,
+    log_operation
+)
+from .routes import router as auth_router
+
+__all__ = [
+    "Database",
+    "get_db",
+    "init_database",
+    "User",
+    "OperationLog",
+    "JWTHandler",
+    "DingTalkClient",
+    "init_dingtalk_client",
+    "get_dingtalk_client",
+    "AuthMiddleware",
+    "get_current_user",
+    "get_optional_user",
+    "init_auth",
+    "is_auth_enabled",
+    "log_operation",
+    "auth_router",
+]

src/auth/database.py

@@ -0,0 +1,87 @@
+"""数据库连接管理"""
+import logging
+from pathlib import Path
+from typing import AsyncGenerator
+
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
+from sqlalchemy.orm import DeclarativeBase
+
+logger = logging.getLogger(__name__)
+
+
+class Base(DeclarativeBase):
+    """SQLAlchemy 基类"""
+    pass
+
+
+class Database:
+    """数据库管理器"""
+
+    def __init__(self, db_path: str = "data/app.db"):
+        self.db_path = db_path
+        self.engine = None
+        self.session_factory = None
+
+    async def init(self):
+        """初始化数据库连接"""
+        # 确保目录存在
+        Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
+
+        # 创建异步引擎
+        self.engine = create_async_engine(
+            f"sqlite+aiosqlite:///{self.db_path}",
+            echo=False,
+            future=True
+        )
+
+        # 创建会话工厂
+        self.session_factory = async_sessionmaker(
+            self.engine,
+            class_=AsyncSession,
+            expire_on_commit=False
+        )
+
+        # 创建表
+        async with self.engine.begin() as conn:
+            from .models import User, OperationLog  # 导入模型以注册
+            await conn.run_sync(Base.metadata.create_all)
+
+        logger.info(f"数据库初始化完成: {self.db_path}")
+
+    async def close(self):
+        """关闭数据库连接"""
+        if self.engine:
+            await self.engine.dispose()
+            logger.info("数据库连接已关闭")
+
+    async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
+        """获取数据库会话"""
+        if not self.session_factory:
+            raise RuntimeError("数据库未初始化")
+
+        async with self.session_factory() as session:
+            try:
+                yield session
+                await session.commit()
+            except Exception:
+                await session.rollback()
+                raise
+
+
+# 全局数据库实例
+_database: Database = None
+
+
+async def init_database(db_path: str = "data/app.db") -> Database:
+    """初始化全局数据库实例"""
+    global _database
+    _database = Database(db_path)
+    await _database.init()
+    return _database
+
+
+def get_db() -> Database:
+    """获取全局数据库实例"""
+    if _database is None:
+        raise RuntimeError("数据库未初始化，请先调用 init_database()")
+    return _database