AstrBot 4.5.1__py3-none-any.whl → 4.5.2__py3-none-any.whl

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

@@ -1,19 +1,20 @@
-"""
-企业微信智能机器人 API 客户端
+"""企业微信智能机器人 API 客户端
 处理消息加密解密、API 调用等
 """
 
-import json
 import base64
 import hashlib
-from typing import Dict, Any, Optional, Tuple, Union
-from Crypto.Cipher import AES
+import json
+from typing import Any
+
 import aiohttp
+from Crypto.Cipher import AES
 
-from .WXBizJsonMsgCrypt import WXBizJsonMsgCrypt
-from .wecomai_utils import WecomAIBotConstants
 from astrbot import logger
 
+from .wecomai_utils import WecomAIBotConstants
+from .WXBizJsonMsgCrypt import WXBizJsonMsgCrypt
+
 
 class WecomAIBotAPIClient:
     """企业微信智能机器人 API 客户端"""
@@ -24,14 +25,19 @@ class WecomAIBotAPIClient:
         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]]]:
+        self,
+        encrypted_data: bytes,
+        msg_signature: str,
+        timestamp: str,
+        nonce: str,
+    ) -> tuple[int, dict[str, Any] | None]:
         """解密企业微信消息
 
         Args:
@@ -42,10 +48,14 @@ class WecomAIBotAPIClient:
 
         Returns:
             (错误码, 解密后的消息数据字典)
+
         """
         try:
             ret, decrypted_msg = self.wxcpt.DecryptMsg(
-                encrypted_data, msg_signature, timestamp, nonce
+                encrypted_data,
+                msg_signature,
+                timestamp,
+                nonce,
             )
 
             if ret != WecomAIBotConstants.SUCCESS:
@@ -70,8 +80,11 @@ class WecomAIBotAPIClient:
             return WecomAIBotConstants.DECRYPT_ERROR, None
 
     async def encrypt_message(
-        self, plain_message: str, nonce: str, timestamp: str
-    ) -> Optional[str]:
+        self,
+        plain_message: str,
+        nonce: str,
+        timestamp: str,
+    ) -> str | None:
         """加密消息
 
         Args:
@@ -81,6 +94,7 @@ class WecomAIBotAPIClient:
 
         Returns:
             加密后的消息，失败时返回 None
+
         """
         try:
             ret, encrypted_msg = self.wxcpt.EncryptMsg(plain_message, nonce, timestamp)
@@ -97,7 +111,11 @@ class WecomAIBotAPIClient:
             return None
 
     def verify_url(
-        self, msg_signature: str, timestamp: str, nonce: str, echostr: str
+        self,
+        msg_signature: str,
+        timestamp: str,
+        nonce: str,
+        echostr: str,
     ) -> str:
         """验证回调 URL
 
@@ -109,10 +127,14 @@ class WecomAIBotAPIClient:
 
         Returns:
             验证结果字符串
+
         """
         try:
             ret, echo_result = self.wxcpt.VerifyURL(
-                msg_signature, timestamp, nonce, echostr
+                msg_signature,
+                timestamp,
+                nonce,
+                echostr,
             )
 
             if ret != WecomAIBotConstants.SUCCESS:
@@ -127,8 +149,10 @@ class WecomAIBotAPIClient:
             return "verify fail"
 
     async def process_encrypted_image(
-        self, image_url: str, aes_key_base64: Optional[str] = None
-    ) -> Tuple[bool, Union[bytes, str]]:
+        self,
+        image_url: str,
+        aes_key_base64: str | None = None,
+    ) -> tuple[bool, bytes | str]:
         """下载并解密加密图片
 
         Args:
@@ -137,6 +161,7 @@ class WecomAIBotAPIClient:
 
         Returns:
             (是否成功, 图片数据或错误信息)
+
         """
         try:
             # 下载图片
@@ -161,7 +186,7 @@ class WecomAIBotAPIClient:
 
             # Base64 解码密钥
             aes_key = base64.b64decode(
-                aes_key_base64 + "=" * (-len(aes_key_base64) % 4)
+                aes_key_base64 + "=" * (-len(aes_key_base64) % 4),
             )
             if len(aes_key) != 32:
                 raise ValueError("无效的 AES 密钥长度: 应为 32 字节")
@@ -183,17 +208,17 @@ class WecomAIBotAPIClient:
             return True, decrypted_data
 
         except aiohttp.ClientError as e:
-            error_msg = f"图片下载失败: {str(e)}"
+            error_msg = f"图片下载失败: {e!s}"
             logger.error(error_msg)
             return False, error_msg
 
         except ValueError as e:
-            error_msg = f"参数错误: {str(e)}"
+            error_msg = f"参数错误: {e!s}"
             logger.error(error_msg)
             return False, error_msg
 
         except Exception as e:
-            error_msg = f"图片处理异常: {str(e)}"
+            error_msg = f"图片处理异常: {e!s}"
             logger.error(error_msg)
             return False, error_msg
 
@@ -212,6 +237,7 @@ class WecomAIBotStreamMessageBuilder:
 
         Returns:
             JSON 格式的流消息字符串
+
         """
         plain = {
             "msgtype": WecomAIBotConstants.MSG_TYPE_STREAM,
@@ -221,7 +247,9 @@ class WecomAIBotStreamMessageBuilder:
 
     @staticmethod
     def make_image_stream(
-        stream_id: str, image_data: bytes, finish: bool = False
+        stream_id: str,
+        image_data: bytes,
+        finish: bool = False,
     ) -> str:
         """构建图片流消息
 
@@ -232,6 +260,7 @@ class WecomAIBotStreamMessageBuilder:
 
         Returns:
             JSON 格式的流消息字符串
+
         """
         image_md5 = hashlib.md5(image_data).hexdigest()
         image_base64 = base64.b64encode(image_data).decode("utf-8")
@@ -245,7 +274,7 @@ class WecomAIBotStreamMessageBuilder:
                     {
                         "msgtype": WecomAIBotConstants.MSG_TYPE_IMAGE,
                         "image": {"base64": image_base64, "md5": image_md5},
-                    }
+                    },
                 ],
             },
         }
@@ -253,7 +282,10 @@ class WecomAIBotStreamMessageBuilder:
 
     @staticmethod
     def make_mixed_stream(
-        stream_id: str, content: str, msg_items: list, finish: bool = False
+        stream_id: str,
+        content: str,
+        msg_items: list,
+        finish: bool = False,
     ) -> str:
         """构建混合类型流消息
 
@@ -265,6 +297,7 @@ class WecomAIBotStreamMessageBuilder:
 
         Returns:
             JSON 格式的流消息字符串
+
         """
         plain = {
             "msgtype": WecomAIBotConstants.MSG_TYPE_STREAM,
@@ -283,6 +316,7 @@ class WecomAIBotStreamMessageBuilder:
 
         Returns:
             JSON 格式的文本消息字符串
+
         """
         plain = {"msgtype": "text", "text": {"content": content}}
         return json.dumps(plain, ensure_ascii=False)
@@ -292,7 +326,7 @@ class WecomAIBotMessageParser:
     """企业微信智能机器人消息解析器"""
 
     @staticmethod
-    def parse_text_message(data: Dict[str, Any]) -> Optional[str]:
+    def parse_text_message(data: dict[str, Any]) -> str | None:
         """解析文本消息
 
         Args:
@@ -300,6 +334,7 @@ class WecomAIBotMessageParser:
 
         Returns:
             文本内容，解析失败返回 None
+
         """
         try:
             return data.get("text", {}).get("content")
@@ -308,7 +343,7 @@ class WecomAIBotMessageParser:
             return None
 
     @staticmethod
-    def parse_image_message(data: Dict[str, Any]) -> Optional[str]:
+    def parse_image_message(data: dict[str, Any]) -> str | None:
         """解析图片消息
 
         Args:
@@ -316,6 +351,7 @@ class WecomAIBotMessageParser:
 
         Returns:
             图片 URL，解析失败返回 None
+
         """
         try:
             return data.get("image", {}).get("url")
@@ -324,7 +360,7 @@ class WecomAIBotMessageParser:
             return None
 
     @staticmethod
-    def parse_stream_message(data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    def parse_stream_message(data: dict[str, Any]) -> dict[str, Any] | None:
         """解析流消息
 
         Args:
@@ -332,6 +368,7 @@ class WecomAIBotMessageParser:
 
         Returns:
             流消息数据，解析失败返回 None
+
         """
         try:
             stream_data = data.get("stream", {})
@@ -346,7 +383,7 @@ class WecomAIBotMessageParser:
             return None
 
     @staticmethod
-    def parse_mixed_message(data: Dict[str, Any]) -> Optional[list]:
+    def parse_mixed_message(data: dict[str, Any]) -> list | None:
         """解析混合消息
 
         Args:
@@ -354,6 +391,7 @@ class WecomAIBotMessageParser:
 
         Returns:
             消息项列表，解析失败返回 None
+
         """
         try:
             return data.get("mixed", {}).get("msg_item", [])
@@ -362,7 +400,7 @@ class WecomAIBotMessageParser:
             return None
 
     @staticmethod
-    def parse_event_message(data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    def parse_event_message(data: dict[str, Any]) -> dict[str, Any] | None:
         """解析事件消息
 
         Args:
@@ -370,6 +408,7 @@ class WecomAIBotMessageParser:
 
         Returns:
             事件数据，解析失败返回 None
+
         """
         try:
             return data.get("event", {})

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

@@ -1,13 +1,11 @@
-"""
-企业微信智能机器人事件处理模块，处理消息事件的发送和接收
-"""
+"""企业微信智能机器人事件处理模块，处理消息事件的发送和接收"""
 
+from astrbot.api import logger
 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
@@ -32,6 +30,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
             platform_meta: 平台元数据
             session_id: 会话 ID
             api_client: API 客户端
+
         """
         super().__init__(message_str, message_obj, platform_meta, session_id)
         self.api_client = api_client
@@ -50,7 +49,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
                     "type": "end",
                     "data": "",
                     "streaming": False,
-                }
+                },
             )
             return ""
 
@@ -64,7 +63,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
                         "data": data,
                         "streaming": streaming,
                         "session_id": stream_id,
-                    }
+                    },
                 )
             elif isinstance(comp, Image):
                 # 处理图片消息
@@ -77,7 +76,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
                                 "image_data": image_base64,
                                 "streaming": streaming,
                                 "session_id": stream_id,
-                            }
+                            },
                         )
                     else:
                         logger.warning("图片数据为空，跳过")
@@ -127,7 +126,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
                         "data": final_data,
                         "streaming": True,
                         "session_id": self.session_id,
-                    }
+                    },
                 )
                 final_data = ""
                 continue
@@ -144,6 +143,6 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
                 "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

@@ -1,11 +1,11 @@
-"""
-企业微信智能机器人队列管理器
+"""企业微信智能机器人队列管理器
 参考 webchat_queue_mgr.py，为企业微信智能机器人实现队列机制
 支持异步消息处理和流式响应
 """
 
 import asyncio
-from typing import Dict, Any, Optional
+from typing import Any
+
 from astrbot.api import logger
 
 
@@ -13,13 +13,13 @@ class WecomAIQueueMgr:
     """企业微信智能机器人队列管理器"""
 
     def __init__(self) -> None:
-        self.queues: Dict[str, asyncio.Queue] = {}
+        self.queues: dict[str, asyncio.Queue] = {}
         """StreamID 到输入队列的映射 - 用于接收用户消息"""
 
-        self.back_queues: Dict[str, asyncio.Queue] = {}
+        self.back_queues: dict[str, asyncio.Queue] = {}
         """StreamID 到输出队列的映射 - 用于发送机器人响应"""
 
-        self.pending_responses: Dict[str, Dict[str, Any]] = {}
+        self.pending_responses: dict[str, dict[str, Any]] = {}
         """待处理的响应缓存，用于流式响应"""
 
     def get_or_create_queue(self, session_id: str) -> asyncio.Queue:
@@ -30,6 +30,7 @@ class WecomAIQueueMgr:
 
         Returns:
             输入队列实例
+
         """
         if session_id not in self.queues:
             self.queues[session_id] = asyncio.Queue()
@@ -44,6 +45,7 @@ class WecomAIQueueMgr:
 
         Returns:
             输出队列实例
+
         """
         if session_id not in self.back_queues:
             self.back_queues[session_id] = asyncio.Queue()
@@ -55,6 +57,7 @@ class WecomAIQueueMgr:
 
         Args:
             session_id: 会话ID
+
         """
         if session_id in self.queues:
             del self.queues[session_id]
@@ -76,6 +79,7 @@ class WecomAIQueueMgr:
 
         Returns:
             是否存在队列
+
         """
         return session_id in self.queues
 
@@ -87,15 +91,17 @@ class WecomAIQueueMgr:
 
         Returns:
             是否存在输出队列
+
         """
         return session_id in self.back_queues
 
-    def set_pending_response(self, session_id: str, callback_params: Dict[str, str]):
+    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,
@@ -103,7 +109,7 @@ class WecomAIQueueMgr:
         }
         logger.debug(f"[WecomAI] 设置待处理响应: {session_id}")
 
-    def get_pending_response(self, session_id: str) -> Optional[Dict[str, Any]]:
+    def get_pending_response(self, session_id: str) -> dict[str, Any] | None:
         """获取待处理的响应参数
 
         Args:
@@ -111,6 +117,7 @@ class WecomAIQueueMgr:
 
         Returns:
             响应参数，如果不存在则返回None
+
         """
         return self.pending_responses.get(session_id)
 
@@ -119,6 +126,7 @@ class WecomAIQueueMgr:
 
         Args:
             max_age_seconds: 最大存活时间（秒）
+
         """
         current_time = asyncio.get_event_loop().time()
         expired_sessions = []
@@ -131,11 +139,12 @@ class WecomAIQueueMgr:
             del self.pending_responses[session_id]
             logger.debug(f"[WecomAI] 清理过期响应: {session_id}")
 
-    def get_stats(self) -> Dict[str, int]:
+    def get_stats(self) -> dict[str, int]:
         """获取队列统计信息
 
         Returns:
             统计信息字典
+
         """
         return {
             "input_queues": len(self.queues),

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

@@ -1,12 +1,13 @@
-"""
-企业微信智能机器人 HTTP 服务器
+"""企业微信智能机器人 HTTP 服务器
 处理企业微信智能机器人的 HTTP 回调请求
 """
 
 import asyncio
-from typing import Dict, Any, Optional, Callable
+from collections.abc import Callable
+from typing import Any
 
 import quart
+
 from astrbot.api import logger
 
 from .wecomai_api import WecomAIBotAPIClient
@@ -21,9 +22,7 @@ class WecomAIBotServer:
         host: str,
         port: int,
         api_client: WecomAIBotAPIClient,
-        message_handler: Optional[
-            Callable[[Dict[str, Any], Dict[str, str]], Any]
-        ] = None,
+        message_handler: Callable[[dict[str, Any], dict[str, str]], Any] | None = None,
     ):
         """初始化服务器
 
@@ -32,6 +31,7 @@ class WecomAIBotServer:
             port: 监听端口
             api_client: API客户端实例
             message_handler: 消息处理回调函数
+
         """
         self.host = host
         self.port = port
@@ -45,7 +45,6 @@ class WecomAIBotServer:
 
     def _setup_routes(self):
         """设置 Quart 路由"""
-
         # 使用 Quart 的 add_url_rule 方法添加路由
         self.app.add_url_rule(
             "/webhook/wecom-ai-bot",
@@ -98,7 +97,7 @@ class WecomAIBotServer:
         assert nonce is not None
 
         logger.debug(
-            f"收到消息回调，msg_signature={msg_signature}, timestamp={timestamp}, nonce={nonce}"
+            f"收到消息回调，msg_signature={msg_signature}, timestamp={timestamp}, nonce={nonce}",
         )
 
         try:
@@ -111,7 +110,10 @@ class WecomAIBotServer:
 
             # 解密消息
             ret_code, message_data = await self.api_client.decrypt_message(
-                post_data, msg_signature, timestamp, nonce
+                post_data,
+                msg_signature,
+                timestamp,
+                nonce,
             )
 
             if ret_code != WecomAIBotConstants.SUCCESS or not message_data:
@@ -123,7 +125,8 @@ class WecomAIBotServer:
             if self.message_handler:
                 try:
                     response = await self.message_handler(
-                        message_data, {"nonce": nonce, "timestamp": timestamp}
+                        message_data,
+                        {"nonce": nonce, "timestamp": timestamp},
                     )
                 except Exception as e:
                     logger.error("消息处理器执行异常: %s", e)
@@ -131,8 +134,7 @@ class WecomAIBotServer:
 
             if response:
                 return response, 200, {"Content-Type": "text/plain"}
-            else:
-                return "success", 200, {"Content-Type": "text/plain"}
+            return "success", 200, {"Content-Type": "text/plain"}
 
         except Exception as e:
             logger.error("处理消息时发生异常: %s", e)

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

@@ -1,16 +1,17 @@
-"""
-企业微信智能机器人工具模块
+"""企业微信智能机器人工具模块
 提供常量定义、工具函数和辅助方法
 """
 
-import string
-import random
-import hashlib
+import asyncio
 import base64
+import hashlib
+import secrets
+import string
+from typing import Any
+
 import aiohttp
-import asyncio
 from Crypto.Cipher import AES
-from typing import Any, Tuple
+
 from astrbot.api import logger
 
 
@@ -49,9 +50,10 @@ def generate_random_string(length: int = 10) -> str:
 
     Returns:
         随机字符串
+
     """
     letters = string.ascii_letters + string.digits
-    return "".join(random.choice(letters) for _ in range(length))
+    return "".join(secrets.choice(letters) for _ in range(length))
 
 
 def calculate_image_md5(image_data: bytes) -> str:
@@ -62,6 +64,7 @@ def calculate_image_md5(image_data: bytes) -> str:
 
     Returns:
         MD5 哈希值（十六进制字符串）
+
     """
     return hashlib.md5(image_data).hexdigest()
 
@@ -74,6 +77,7 @@ def encode_image_base64(image_data: bytes) -> str:
 
     Returns:
         Base64 编码的字符串
+
     """
     return base64.b64encode(image_data).decode("utf-8")
 
@@ -87,11 +91,12 @@ def format_session_id(session_type: str, session_id: str) -> str:
 
     Returns:
         格式化后的会话 ID
+
     """
     return f"wecom_ai_bot_{session_type}_{session_id}"
 
 
-def parse_session_id(formatted_session_id: str) -> Tuple[str, str]:
+def parse_session_id(formatted_session_id: str) -> tuple[str, str]:
     """解析格式化的会话 ID
 
     Args:
@@ -99,6 +104,7 @@ def parse_session_id(formatted_session_id: str) -> Tuple[str, str]:
 
     Returns:
         (会话类型, 原始会话ID)
+
     """
     parts = formatted_session_id.split("_", 3)
     if (
@@ -120,6 +126,7 @@ def safe_json_loads(json_str: str, default: Any = None) -> Any:
 
     Returns:
         解析结果或默认值
+
     """
     import json
 
@@ -139,13 +146,15 @@ def format_error_response(error_code: int, error_msg: str) -> str:
 
     Returns:
         格式化的错误响应字符串
+
     """
     return f"Error {error_code}: {error_msg}"
 
 
 async def process_encrypted_image(
-    image_url: str, aes_key_base64: str
-) -> Tuple[bool, str]:
+    image_url: str,
+    aes_key_base64: str,
+) -> tuple[bool, str]:
     """下载并解密加密图片
 
     Args:
@@ -155,6 +164,7 @@ async def process_encrypted_image(
     Returns:
         Tuple[bool, str]: status 为 True 时 data 是解密后的图片数据的 base64 编码，
             status 为 False 时 data 是错误信息
+
     """
     # 1. 下载加密图片
     logger.info("开始下载加密图片: %s", image_url)
@@ -165,7 +175,7 @@ async def process_encrypted_image(
                 encrypted_data = await response.read()
         logger.info("图片下载成功，大小: %d 字节", len(encrypted_data))
     except (aiohttp.ClientError, asyncio.TimeoutError) as e:
-        error_msg = f"下载图片失败: {str(e)}"
+        error_msg = f"下载图片失败: {e!s}"
         logger.error(error_msg)
         return False, error_msg