AstrBot 3.5.6__py3-none-any.whl → 4.7.0__py3-none-any.whl

astrbot/core/platform/sources/wecom_ai_bot/wecomai_queue_mgr.py

@@ -0,0 +1,153 @@
+"""企业微信智能机器人队列管理器
+参考 webchat_queue_mgr.py，为企业微信智能机器人实现队列机制
+支持异步消息处理和流式响应
+"""
+
+import asyncio
+from typing import Any
+
+from astrbot.api import logger
+
+
+class WecomAIQueueMgr:
+    """企业微信智能机器人队列管理器"""
+
+    def __init__(self) -> None:
+        self.queues: dict[str, asyncio.Queue] = {}
+        """StreamID 到输入队列的映射 - 用于接收用户消息"""
+
+        self.back_queues: dict[str, asyncio.Queue] = {}
+        """StreamID 到输出队列的映射 - 用于发送机器人响应"""
+
+        self.pending_responses: dict[str, dict[str, Any]] = {}
+        """待处理的响应缓存，用于流式响应"""
+
+    def get_or_create_queue(self, session_id: str) -> asyncio.Queue:
+        """获取或创建指定会话的输入队列
+
+        Args:
+            session_id: 会话ID
+
+        Returns:
+            输入队列实例
+
+        """
+        if session_id not in self.queues:
+            self.queues[session_id] = asyncio.Queue()
+            logger.debug(f"[WecomAI] 创建输入队列: {session_id}")
+        return self.queues[session_id]
+
+    def get_or_create_back_queue(self, session_id: str) -> asyncio.Queue:
+        """获取或创建指定会话的输出队列
+
+        Args:
+            session_id: 会话ID
+
+        Returns:
+            输出队列实例
+
+        """
+        if session_id not in self.back_queues:
+            self.back_queues[session_id] = asyncio.Queue()
+            logger.debug(f"[WecomAI] 创建输出队列: {session_id}")
+        return self.back_queues[session_id]
+
+    def remove_queues(self, session_id: str):
+        """移除指定会话的所有队列
+
+        Args:
+            session_id: 会话ID
+
+        """
+        if session_id in self.queues:
+            del self.queues[session_id]
+            logger.debug(f"[WecomAI] 移除输入队列: {session_id}")
+
+        if session_id in self.back_queues:
+            del self.back_queues[session_id]
+            logger.debug(f"[WecomAI] 移除输出队列: {session_id}")
+
+        if session_id in self.pending_responses:
+            del self.pending_responses[session_id]
+            logger.debug(f"[WecomAI] 移除待处理响应: {session_id}")
+
+    def has_queue(self, session_id: str) -> bool:
+        """检查是否存在指定会话的队列
+
+        Args:
+            session_id: 会话ID
+
+        Returns:
+            是否存在队列
+
+        """
+        return session_id in self.queues
+
+    def has_back_queue(self, session_id: str) -> bool:
+        """检查是否存在指定会话的输出队列
+
+        Args:
+            session_id: 会话ID
+
+        Returns:
+            是否存在输出队列
+
+        """
+        return session_id in self.back_queues
+
+    def set_pending_response(self, session_id: str, callback_params: dict[str, str]):
+        """设置待处理的响应参数
+
+        Args:
+            session_id: 会话ID
+            callback_params: 回调参数（nonce, timestamp等）
+
+        """
+        self.pending_responses[session_id] = {
+            "callback_params": callback_params,
+            "timestamp": asyncio.get_event_loop().time(),
+        }
+        logger.debug(f"[WecomAI] 设置待处理响应: {session_id}")
+
+    def get_pending_response(self, session_id: str) -> dict[str, Any] | None:
+        """获取待处理的响应参数
+
+        Args:
+            session_id: 会话ID
+
+        Returns:
+            响应参数，如果不存在则返回None
+
+        """
+        return self.pending_responses.get(session_id)
+
+    def cleanup_expired_responses(self, max_age_seconds: int = 300):
+        """清理过期的待处理响应
+
+        Args:
+            max_age_seconds: 最大存活时间（秒）
+
+        """
+        current_time = asyncio.get_event_loop().time()
+        expired_sessions = []
+
+        for session_id, response_data in self.pending_responses.items():
+            if current_time - response_data["timestamp"] > max_age_seconds:
+                expired_sessions.append(session_id)
+
+        for session_id in expired_sessions:
+            del self.pending_responses[session_id]
+            logger.debug(f"[WecomAI] 清理过期响应: {session_id}")
+
+    def get_stats(self) -> dict[str, int]:
+        """获取队列统计信息
+
+        Returns:
+            统计信息字典
+
+        """
+        return {
+            "input_queues": len(self.queues),
+            "output_queues": len(self.back_queues),
+            "pending_responses": len(self.pending_responses),
+        }

astrbot/core/platform/sources/wecom_ai_bot/wecomai_server.py

@@ -0,0 +1,168 @@
+"""企业微信智能机器人 HTTP 服务器
+处理企业微信智能机器人的 HTTP 回调请求
+"""
+
+import asyncio
+from collections.abc import Callable
+from typing import Any
+
+import quart
+
+from astrbot.api import logger
+
+from .wecomai_api import WecomAIBotAPIClient
+from .wecomai_utils import WecomAIBotConstants
+
+
+class WecomAIBotServer:
+    """企业微信智能机器人 HTTP 服务器"""
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        api_client: WecomAIBotAPIClient,
+        message_handler: Callable[[dict[str, Any], dict[str, str]], Any] | None = None,
+    ):
+        """初始化服务器
+
+        Args:
+            host: 监听地址
+            port: 监听端口
+            api_client: API客户端实例
+            message_handler: 消息处理回调函数
+
+        """
+        self.host = host
+        self.port = port
+        self.api_client = api_client
+        self.message_handler = message_handler
+
+        self.app = quart.Quart(__name__)
+        self._setup_routes()
+
+        self.shutdown_event = asyncio.Event()
+
+    def _setup_routes(self):
+        """设置 Quart 路由"""
+        # 使用 Quart 的 add_url_rule 方法添加路由
+        self.app.add_url_rule(
+            "/webhook/wecom-ai-bot",
+            view_func=self.verify_url,
+            methods=["GET"],
+        )
+
+        self.app.add_url_rule(
+            "/webhook/wecom-ai-bot",
+            view_func=self.handle_message,
+            methods=["POST"],
+        )
+
+    async def verify_url(self):
+        """验证回调 URL"""
+        args = quart.request.args
+        msg_signature = args.get("msg_signature")
+        timestamp = args.get("timestamp")
+        nonce = args.get("nonce")
+        echostr = args.get("echostr")
+
+        if not all([msg_signature, timestamp, nonce, echostr]):
+            logger.error("URL 验证参数缺失")
+            return "verify fail", 400
+
+        # 类型检查确保不为 None
+        assert msg_signature is not None
+        assert timestamp is not None
+        assert nonce is not None
+        assert echostr is not None
+
+        logger.info("收到企业微信智能机器人 WebHook URL 验证请求。")
+        result = self.api_client.verify_url(msg_signature, timestamp, nonce, echostr)
+        return result, 200, {"Content-Type": "text/plain"}
+
+    async def handle_message(self):
+        """处理消息回调"""
+        args = quart.request.args
+        msg_signature = args.get("msg_signature")
+        timestamp = args.get("timestamp")
+        nonce = args.get("nonce")
+
+        if not all([msg_signature, timestamp, nonce]):
+            logger.error("消息回调参数缺失")
+            return "缺少必要参数", 400
+
+        # 类型检查确保不为 None
+        assert msg_signature is not None
+        assert timestamp is not None
+        assert nonce is not None
+
+        logger.debug(
+            f"收到消息回调，msg_signature={msg_signature}, timestamp={timestamp}, nonce={nonce}",
+        )
+
+        try:
+            # 获取请求体
+            post_data = await quart.request.get_data()
+
+            # 确保 post_data 是 bytes 类型
+            if isinstance(post_data, str):
+                post_data = post_data.encode("utf-8")
+
+            # 解密消息
+            ret_code, message_data = await self.api_client.decrypt_message(
+                post_data,
+                msg_signature,
+                timestamp,
+                nonce,
+            )
+
+            if ret_code != WecomAIBotConstants.SUCCESS or not message_data:
+                logger.error("消息解密失败，错误码: %d", ret_code)
+                return "消息解密失败", 400
+
+            # 调用消息处理器
+            response = None
+            if self.message_handler:
+                try:
+                    response = await self.message_handler(
+                        message_data,
+                        {"nonce": nonce, "timestamp": timestamp},
+                    )
+                except Exception as e:
+                    logger.error("消息处理器执行异常: %s", e)
+                    return "消息处理异常", 500
+
+            if response:
+                return response, 200, {"Content-Type": "text/plain"}
+            return "success", 200, {"Content-Type": "text/plain"}
+
+        except Exception as e:
+            logger.error("处理消息时发生异常: %s", e)
+            return "内部服务器错误", 500
+
+    async def start_server(self):
+        """启动服务器"""
+        logger.info("启动企业微信智能机器人服务器，监听 %s:%d", self.host, self.port)
+
+        try:
+            await self.app.run_task(
+                host=self.host,
+                port=self.port,
+                shutdown_trigger=self.shutdown_trigger,
+            )
+        except Exception as e:
+            logger.error("服务器运行异常: %s", e)
+            raise
+
+    async def shutdown_trigger(self):
+        """关闭触发器"""
+        await self.shutdown_event.wait()
+
+    async def shutdown(self):
+        """关闭服务器"""
+        logger.info("企业微信智能机器人服务器正在关闭...")
+        self.shutdown_event.set()
+
+    def get_app(self):
+        """获取 Quart 应用实例"""
+        return self.app

astrbot/core/platform/sources/wecom_ai_bot/wecomai_utils.py

@@ -0,0 +1,209 @@
+"""企业微信智能机器人工具模块
+提供常量定义、工具函数和辅助方法
+"""
+
+import asyncio
+import base64
+import hashlib
+import secrets
+import string
+from typing import Any
+
+import aiohttp
+from Crypto.Cipher import AES
+
+from astrbot.api import logger
+
+
+# 常量定义
+class WecomAIBotConstants:
+    """企业微信智能机器人常量"""
+
+    # 消息类型
+    MSG_TYPE_TEXT = "text"
+    MSG_TYPE_IMAGE = "image"
+    MSG_TYPE_MIXED = "mixed"
+    MSG_TYPE_STREAM = "stream"
+    MSG_TYPE_EVENT = "event"
+
+    # 流消息状态
+    STREAM_CONTINUE = False
+    STREAM_FINISH = True
+
+    # 错误码
+    SUCCESS = 0
+    DECRYPT_ERROR = -40001
+    VALIDATE_SIGNATURE_ERROR = -40002
+    PARSE_XML_ERROR = -40003
+    COMPUTE_SIGNATURE_ERROR = -40004
+    ILLEGAL_AES_KEY = -40005
+    VALIDATE_APPID_ERROR = -40006
+    ENCRYPT_AES_ERROR = -40007
+    ILLEGAL_BUFFER = -40008
+
+
+def generate_random_string(length: int = 10) -> str:
+    """生成随机字符串
+
+    Args:
+        length: 字符串长度，默认为 10
+
+    Returns:
+        随机字符串
+
+    """
+    letters = string.ascii_letters + string.digits
+    return "".join(secrets.choice(letters) for _ in range(length))
+
+
+def calculate_image_md5(image_data: bytes) -> str:
+    """计算图片数据的 MD5 值
+
+    Args:
+        image_data: 图片二进制数据
+
+    Returns:
+        MD5 哈希值（十六进制字符串）
+
+    """
+    return hashlib.md5(image_data).hexdigest()
+
+
+def encode_image_base64(image_data: bytes) -> str:
+    """将图片数据编码为 Base64
+
+    Args:
+        image_data: 图片二进制数据
+
+    Returns:
+        Base64 编码的字符串
+
+    """
+    return base64.b64encode(image_data).decode("utf-8")
+
+
+def format_session_id(session_type: str, session_id: str) -> str:
+    """格式化会话 ID
+
+    Args:
+        session_type: 会话类型 ("user", "group")
+        session_id: 原始会话 ID
+
+    Returns:
+        格式化后的会话 ID
+
+    """
+    return f"wecom_ai_bot_{session_type}_{session_id}"
+
+
+def parse_session_id(formatted_session_id: str) -> tuple[str, str]:
+    """解析格式化的会话 ID
+
+    Args:
+        formatted_session_id: 格式化的会话 ID
+
+    Returns:
+        (会话类型, 原始会话ID)
+
+    """
+    parts = formatted_session_id.split("_", 3)
+    if (
+        len(parts) >= 4
+        and parts[0] == "wecom"
+        and parts[1] == "ai"
+        and parts[2] == "bot"
+    ):
+        return parts[3], "_".join(parts[4:]) if len(parts) > 4 else ""
+    return "user", formatted_session_id
+
+
+def safe_json_loads(json_str: str, default: Any = None) -> Any:
+    """安全地解析 JSON 字符串
+
+    Args:
+        json_str: JSON 字符串
+        default: 解析失败时的默认值
+
+    Returns:
+        解析结果或默认值
+
+    """
+    import json
+
+    try:
+        return json.loads(json_str)
+    except (json.JSONDecodeError, TypeError) as e:
+        logger.warning(f"JSON 解析失败: {e}, 原始字符串: {json_str}")
+        return default
+
+
+def format_error_response(error_code: int, error_msg: str) -> str:
+    """格式化错误响应
+
+    Args:
+        error_code: 错误码
+        error_msg: 错误信息
+
+    Returns:
+        格式化的错误响应字符串
+
+    """
+    return f"Error {error_code}: {error_msg}"
+
+
+async def process_encrypted_image(
+    image_url: str,
+    aes_key_base64: str,
+) -> tuple[bool, str]:
+    """下载并解密加密图片
+
+    Args:
+        image_url: 加密图片的URL
+        aes_key_base64: Base64编码的AES密钥(与回调加解密相同)
+
+    Returns:
+        Tuple[bool, str]: status 为 True 时 data 是解密后的图片数据的 base64 编码，
+            status 为 False 时 data 是错误信息
+
+    """
+    # 1. 下载加密图片
+    logger.info("开始下载加密图片: %s", image_url)
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(image_url, timeout=15) as response:
+                response.raise_for_status()
+                encrypted_data = await response.read()
+        logger.info("图片下载成功，大小: %d 字节", len(encrypted_data))
+    except (aiohttp.ClientError, asyncio.TimeoutError) as e:
+        error_msg = f"下载图片失败: {e!s}"
+        logger.error(error_msg)
+        return False, error_msg
+
+    # 2. 准备AES密钥和IV
+    if not aes_key_base64:
+        raise ValueError("AES密钥不能为空")
+
+    # Base64解码密钥 (自动处理填充)
+    aes_key = base64.b64decode(aes_key_base64 + "=" * (-len(aes_key_base64) % 4))
+    if len(aes_key) != 32:
+        raise ValueError("无效的AES密钥长度: 应为32字节")
+
+    iv = aes_key[:16]  # 初始向量为密钥前16字节
+
+    # 3. 解密图片数据
+    cipher = AES.new(aes_key, AES.MODE_CBC, iv)
+    decrypted_data = cipher.decrypt(encrypted_data)
+
+    # 4. 去除PKCS#7填充 (Python 3兼容写法)
+    pad_len = decrypted_data[-1]  # 直接获取最后一个字节的整数值
+    if pad_len > 32:  # AES-256块大小为32字节
+        raise ValueError("无效的填充长度 (大于32字节)")
+
+    decrypted_data = decrypted_data[:-pad_len]
+    logger.info("图片解密成功，解密后大小: %d 字节", len(decrypted_data))
+
+    # 5. 转换为base64编码
+    base64_data = base64.b64encode(decrypted_data).decode("utf-8")
+    logger.info("图片已转换为base64编码，编码后长度: %d", len(base64_data))
+
+    return True, base64_data