AstrBot 4.3.3__py3-none-any.whl → 4.3.5__py3-none-any.whl

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

@@ -0,0 +1,378 @@
+"""
+企业微信智能机器人 API 客户端
+处理消息加密解密、API 调用等
+"""
+
+import json
+import base64
+import hashlib
+from typing import Dict, Any, Optional, Tuple, Union
+from Crypto.Cipher import AES
+import aiohttp
+
+from .WXBizJsonMsgCrypt import WXBizJsonMsgCrypt
+from .wecomai_utils import WecomAIBotConstants
+from astrbot import logger
+
+
+class WecomAIBotAPIClient:
+    """企业微信智能机器人 API 客户端"""
+
+    def __init__(self, token: str, encoding_aes_key: str):
+        """初始化 API 客户端
+
+        Args:
+            token: 企业微信机器人 Token
+            encoding_aes_key: 消息加密密钥
+        """
+        self.token = token
+        self.encoding_aes_key = encoding_aes_key
+        self.wxcpt = WXBizJsonMsgCrypt(token, encoding_aes_key, "")  # receiveid 为空串
+
+    async def decrypt_message(
+        self, encrypted_data: bytes, msg_signature: str, timestamp: str, nonce: str
+    ) -> Tuple[int, Optional[Dict[str, Any]]]:
+        """解密企业微信消息
+
+        Args:
+            encrypted_data: 加密的消息数据
+            msg_signature: 消息签名
+            timestamp: 时间戳
+            nonce: 随机数
+
+        Returns:
+            (错误码, 解密后的消息数据字典)
+        """
+        try:
+            ret, decrypted_msg = self.wxcpt.DecryptMsg(
+                encrypted_data, msg_signature, timestamp, nonce
+            )
+
+            if ret != WecomAIBotConstants.SUCCESS:
+                logger.error(f"消息解密失败，错误码: {ret}")
+                return ret, None
+
+            # 解析 JSON
+            if decrypted_msg:
+                try:
+                    message_data = json.loads(decrypted_msg)
+                    logger.debug(f"解密成功，消息内容: {message_data}")
+                    return WecomAIBotConstants.SUCCESS, message_data
+                except json.JSONDecodeError as e:
+                    logger.error(f"JSON 解析失败: {e}, 原始消息: {decrypted_msg}")
+                    return WecomAIBotConstants.PARSE_XML_ERROR, None
+            else:
+                logger.error("解密消息为空")
+                return WecomAIBotConstants.DECRYPT_ERROR, None
+
+        except Exception as e:
+            logger.error(f"解密过程发生异常: {e}")
+            return WecomAIBotConstants.DECRYPT_ERROR, None
+
+    async def encrypt_message(
+        self, plain_message: str, nonce: str, timestamp: str
+    ) -> Optional[str]:
+        """加密消息
+
+        Args:
+            plain_message: 明文消息
+            nonce: 随机数
+            timestamp: 时间戳
+
+        Returns:
+            加密后的消息，失败时返回 None
+        """
+        try:
+            ret, encrypted_msg = self.wxcpt.EncryptMsg(plain_message, nonce, timestamp)
+
+            if ret != WecomAIBotConstants.SUCCESS:
+                logger.error(f"消息加密失败，错误码: {ret}")
+                return None
+
+            logger.debug("消息加密成功")
+            return encrypted_msg
+
+        except Exception as e:
+            logger.error(f"加密过程发生异常: {e}")
+            return None
+
+    def verify_url(
+        self, msg_signature: str, timestamp: str, nonce: str, echostr: str
+    ) -> str:
+        """验证回调 URL
+
+        Args:
+            msg_signature: 消息签名
+            timestamp: 时间戳
+            nonce: 随机数
+            echostr: 验证字符串
+
+        Returns:
+            验证结果字符串
+        """
+        try:
+            ret, echo_result = self.wxcpt.VerifyURL(
+                msg_signature, timestamp, nonce, echostr
+            )
+
+            if ret != WecomAIBotConstants.SUCCESS:
+                logger.error(f"URL 验证失败，错误码: {ret}")
+                return "verify fail"
+
+            logger.info("URL 验证成功")
+            return echo_result if echo_result else "verify fail"
+
+        except Exception as e:
+            logger.error(f"URL 验证发生异常: {e}")
+            return "verify fail"
+
+    async def process_encrypted_image(
+        self, image_url: str, aes_key_base64: Optional[str] = None
+    ) -> Tuple[bool, Union[bytes, str]]:
+        """下载并解密加密图片
+
+        Args:
+            image_url: 加密图片的 URL
+            aes_key_base64: Base64 编码的 AES 密钥，如果为 None 则使用实例的密钥
+
+        Returns:
+            (是否成功, 图片数据或错误信息)
+        """
+        try:
+            # 下载图片
+            logger.info(f"开始下载加密图片: {image_url}")
+
+            async with aiohttp.ClientSession() as session:
+                async with session.get(image_url, timeout=15) as response:
+                    if response.status != 200:
+                        error_msg = f"图片下载失败，状态码: {response.status}"
+                        logger.error(error_msg)
+                        return False, error_msg
+
+                    encrypted_data = await response.read()
+                    logger.info(f"图片下载成功，大小: {len(encrypted_data)} 字节")
+
+            # 准备解密密钥
+            if aes_key_base64 is None:
+                aes_key_base64 = self.encoding_aes_key
+
+            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 字节
+
+            # 解密图片数据
+            cipher = AES.new(aes_key, AES.MODE_CBC, iv)
+            decrypted_data = cipher.decrypt(encrypted_data)
+
+            # 去除 PKCS#7 填充
+            pad_len = decrypted_data[-1]
+            if pad_len > 32:  # AES-256 块大小为 32 字节
+                raise ValueError("无效的填充长度 (大于32字节)")
+
+            decrypted_data = decrypted_data[:-pad_len]
+            logger.info(f"图片解密成功，解密后大小: {len(decrypted_data)} 字节")
+
+            return True, decrypted_data
+
+        except aiohttp.ClientError as e:
+            error_msg = f"图片下载失败: {str(e)}"
+            logger.error(error_msg)
+            return False, error_msg
+
+        except ValueError as e:
+            error_msg = f"参数错误: {str(e)}"
+            logger.error(error_msg)
+            return False, error_msg
+
+        except Exception as e:
+            error_msg = f"图片处理异常: {str(e)}"
+            logger.error(error_msg)
+            return False, error_msg
+
+
+class WecomAIBotStreamMessageBuilder:
+    """企业微信智能机器人流消息构建器"""
+
+    @staticmethod
+    def make_text_stream(stream_id: str, content: str, finish: bool = False) -> str:
+        """构建文本流消息
+
+        Args:
+            stream_id: 流 ID
+            content: 文本内容
+            finish: 是否结束
+
+        Returns:
+            JSON 格式的流消息字符串
+        """
+        plain = {
+            "msgtype": WecomAIBotConstants.MSG_TYPE_STREAM,
+            "stream": {"id": stream_id, "finish": finish, "content": content},
+        }
+        return json.dumps(plain, ensure_ascii=False)
+
+    @staticmethod
+    def make_image_stream(
+        stream_id: str, image_data: bytes, finish: bool = False
+    ) -> str:
+        """构建图片流消息
+
+        Args:
+            stream_id: 流 ID
+            image_data: 图片二进制数据
+            finish: 是否结束
+
+        Returns:
+            JSON 格式的流消息字符串
+        """
+        image_md5 = hashlib.md5(image_data).hexdigest()
+        image_base64 = base64.b64encode(image_data).decode("utf-8")
+
+        plain = {
+            "msgtype": WecomAIBotConstants.MSG_TYPE_STREAM,
+            "stream": {
+                "id": stream_id,
+                "finish": finish,
+                "msg_item": [
+                    {
+                        "msgtype": WecomAIBotConstants.MSG_TYPE_IMAGE,
+                        "image": {"base64": image_base64, "md5": image_md5},
+                    }
+                ],
+            },
+        }
+        return json.dumps(plain, ensure_ascii=False)
+
+    @staticmethod
+    def make_mixed_stream(
+        stream_id: str, content: str, msg_items: list, finish: bool = False
+    ) -> str:
+        """构建混合类型流消息
+
+        Args:
+            stream_id: 流 ID
+            content: 文本内容
+            msg_items: 消息项列表
+            finish: 是否结束
+
+        Returns:
+            JSON 格式的流消息字符串
+        """
+        plain = {
+            "msgtype": WecomAIBotConstants.MSG_TYPE_STREAM,
+            "stream": {"id": stream_id, "finish": finish, "msg_item": msg_items},
+        }
+        if content:
+            plain["stream"]["content"] = content
+        return json.dumps(plain, ensure_ascii=False)
+
+    @staticmethod
+    def make_text(content: str) -> str:
+        """构建文本消息
+
+        Args:
+            content: 文本内容
+
+        Returns:
+            JSON 格式的文本消息字符串
+        """
+        plain = {"msgtype": "text", "text": {"content": content}}
+        return json.dumps(plain, ensure_ascii=False)
+
+
+class WecomAIBotMessageParser:
+    """企业微信智能机器人消息解析器"""
+
+    @staticmethod
+    def parse_text_message(data: Dict[str, Any]) -> Optional[str]:
+        """解析文本消息
+
+        Args:
+            data: 消息数据
+
+        Returns:
+            文本内容，解析失败返回 None
+        """
+        try:
+            return data.get("text", {}).get("content")
+        except (KeyError, TypeError):
+            logger.warning("文本消息解析失败")
+            return None
+
+    @staticmethod
+    def parse_image_message(data: Dict[str, Any]) -> Optional[str]:
+        """解析图片消息
+
+        Args:
+            data: 消息数据
+
+        Returns:
+            图片 URL，解析失败返回 None
+        """
+        try:
+            return data.get("image", {}).get("url")
+        except (KeyError, TypeError):
+            logger.warning("图片消息解析失败")
+            return None
+
+    @staticmethod
+    def parse_stream_message(data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """解析流消息
+
+        Args:
+            data: 消息数据
+
+        Returns:
+            流消息数据，解析失败返回 None
+        """
+        try:
+            stream_data = data.get("stream", {})
+            return {
+                "id": stream_data.get("id"),
+                "finish": stream_data.get("finish"),
+                "content": stream_data.get("content"),
+                "msg_item": stream_data.get("msg_item", []),
+            }
+        except (KeyError, TypeError):
+            logger.warning("流消息解析失败")
+            return None
+
+    @staticmethod
+    def parse_mixed_message(data: Dict[str, Any]) -> Optional[list]:
+        """解析混合消息
+
+        Args:
+            data: 消息数据
+
+        Returns:
+            消息项列表，解析失败返回 None
+        """
+        try:
+            return data.get("mixed", {}).get("msg_item", [])
+        except (KeyError, TypeError):
+            logger.warning("混合消息解析失败")
+            return None
+
+    @staticmethod
+    def parse_event_message(data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """解析事件消息
+
+        Args:
+            data: 消息数据
+
+        Returns:
+            事件数据，解析失败返回 None
+        """
+        try:
+            return data.get("event", {})
+        except (KeyError, TypeError):
+            logger.warning("事件消息解析失败")
+            return None

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

@@ -0,0 +1,149 @@
+"""
+企业微信智能机器人事件处理模块，处理消息事件的发送和接收
+"""
+
+from astrbot.api.event import AstrMessageEvent, MessageChain
+from astrbot.api.message_components import (
+    Image,
+    Plain,
+)
+from astrbot.api import logger
+
+from .wecomai_api import WecomAIBotAPIClient
+from .wecomai_queue_mgr import wecomai_queue_mgr
+
+
+class WecomAIBotMessageEvent(AstrMessageEvent):
+    """企业微信智能机器人消息事件"""
+
+    def __init__(
+        self,
+        message_str: str,
+        message_obj,
+        platform_meta,
+        session_id: str,
+        api_client: WecomAIBotAPIClient,
+    ):
+        """初始化消息事件
+
+        Args:
+            message_str: 消息字符串
+            message_obj: 消息对象
+            platform_meta: 平台元数据
+            session_id: 会话 ID
+            api_client: API 客户端
+        """
+        super().__init__(message_str, message_obj, platform_meta, session_id)
+        self.api_client = api_client
+
+    @staticmethod
+    async def _send(
+        message_chain: MessageChain,
+        stream_id: str,
+        streaming: bool = False,
+    ):
+        back_queue = wecomai_queue_mgr.get_or_create_back_queue(stream_id)
+
+        if not message_chain:
+            await back_queue.put(
+                {
+                    "type": "end",
+                    "data": "",
+                    "streaming": False,
+                }
+            )
+            return ""
+
+        data = ""
+        for comp in message_chain.chain:
+            if isinstance(comp, Plain):
+                data = comp.text
+                await back_queue.put(
+                    {
+                        "type": "plain",
+                        "data": data,
+                        "streaming": streaming,
+                        "session_id": stream_id,
+                    }
+                )
+            elif isinstance(comp, Image):
+                # 处理图片消息
+                try:
+                    image_base64 = await comp.convert_to_base64()
+                    if image_base64:
+                        await back_queue.put(
+                            {
+                                "type": "image",
+                                "image_data": image_base64,
+                                "streaming": streaming,
+                                "session_id": stream_id,
+                            }
+                        )
+                    else:
+                        logger.warning("图片数据为空，跳过")
+                except Exception as e:
+                    logger.error("处理图片消息失败: %s", e)
+            else:
+                logger.warning(f"[WecomAI] 不支持的消息组件类型: {type(comp)}, 跳过")
+
+        return data
+
+    async def send(self, message: MessageChain):
+        """发送消息"""
+        raw = self.message_obj.raw_message
+        assert isinstance(raw, dict), (
+            "wecom_ai_bot platform event raw_message should be a dict"
+        )
+        stream_id = raw.get("stream_id", self.session_id)
+        await WecomAIBotMessageEvent._send(message, stream_id)
+        await super().send(message)
+
+    async def send_streaming(self, generator, use_fallback=False):
+        """流式发送消息，参考webchat的send_streaming设计"""
+        final_data = ""
+        raw = self.message_obj.raw_message
+        assert isinstance(raw, dict), (
+            "wecom_ai_bot platform event raw_message should be a dict"
+        )
+        stream_id = raw.get("stream_id", self.session_id)
+        back_queue = wecomai_queue_mgr.get_or_create_back_queue(stream_id)
+
+        # 企业微信智能机器人不支持增量发送，因此我们需要在这里将增量内容累积起来，积累发送
+        increment_plain = ""
+        async for chain in generator:
+            # 累积增量内容，并改写 Plain 段
+            chain.squash_plain()
+            for comp in chain.chain:
+                if isinstance(comp, Plain):
+                    comp.text = increment_plain + comp.text
+                    increment_plain = comp.text
+                    break
+
+            if chain.type == "break" and final_data:
+                # 分割符
+                await back_queue.put(
+                    {
+                        "type": "break",  # break means a segment end
+                        "data": final_data,
+                        "streaming": True,
+                        "session_id": self.session_id,
+                    }
+                )
+                final_data = ""
+                continue
+
+            final_data += await WecomAIBotMessageEvent._send(
+                chain,
+                stream_id=stream_id,
+                streaming=True,
+            )
+
+        await back_queue.put(
+            {
+                "type": "complete",  # complete means we return the final result
+                "data": final_data,
+                "streaming": True,
+                "session_id": self.session_id,
+            }
+        )
+        await super().send_streaming(generator, use_fallback)

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

@@ -0,0 +1,148 @@
+"""
+企业微信智能机器人队列管理器
+参考 webchat_queue_mgr.py，为企业微信智能机器人实现队列机制
+支持异步消息处理和流式响应
+"""
+
+import asyncio
+from typing import Dict, Any, Optional
+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) -> Optional[Dict[str, Any]]:
+        """获取待处理的响应参数
+
+        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),
+        }
+
+
+# 全局队列管理器实例
+wecomai_queue_mgr = WecomAIQueueMgr()