wecom-aibot-python-sdk 1.0.0__py3-none-any.whl

aibot/__init__.py

@@ -0,0 +1,50 @@
+"""
+企业微信智能机器人 Python SDK
+
+基于 WebSocket 长连接通道，提供消息收发、流式回复、模板卡片、事件回调、文件下载解密等核心能力。
+"""
+
+__version__ = "1.0.0"
+
+from .client import WSClient
+from .api import WeComApiClient
+from .ws import WsConnectionManager
+from .message_handler import MessageHandler
+from .crypto_utils import decrypt_file
+from .logger import DefaultLogger
+from .utils import generate_req_id, generate_random_string
+from .types import (
+    MessageType,
+    EventType,
+    TemplateCardType,
+    WsCmd,
+    WSClientOptions,
+    WsFrame,
+    WsFrameHeaders,
+    Logger,
+)
+
+__all__ = [
+    # 版本
+    "__version__",
+    # 类
+    "WSClient",
+    "WeComApiClient",
+    "WsConnectionManager",
+    "MessageHandler",
+    "DefaultLogger",
+    # 函数
+    "decrypt_file",
+    "generate_req_id",
+    "generate_random_string",
+    # 枚举/常量
+    "MessageType",
+    "EventType",
+    "TemplateCardType",
+    "WsCmd",
+    # 类型
+    "WSClientOptions",
+    "WsFrame",
+    "WsFrameHeaders",
+    "Logger",
+]

aibot/api.py

@@ -0,0 +1,74 @@
+"""
+企业微信 API 客户端
+
+对标 Node.js SDK src/api.ts
+仅负责文件下载等 HTTP 辅助功能，消息收发均走 WebSocket 通道。
+"""
+
+import re
+import ssl
+from typing import Any, Optional, Tuple
+from urllib.parse import unquote
+
+import aiohttp
+
+try:
+    import certifi
+    _SSL_CONTEXT = ssl.create_default_context(cafile=certifi.where())
+except ImportError:
+    # 未安装 certifi 时回退到系统默认证书
+    _SSL_CONTEXT = ssl.create_default_context()
+
+
+class WeComApiClient:
+    """企业微信 API 客户端"""
+
+    def __init__(self, logger: Any, timeout: int = 10000):
+        self._logger = logger
+        self._timeout = aiohttp.ClientTimeout(total=timeout / 1000)
+
+    async def download_file_raw(self, url: str) -> Tuple[bytes, Optional[str]]:
+        """
+        下载文件（返回原始 bytes 及文件名）
+
+        :param url: 文件下载地址
+        :return: (文件数据, 文件名)
+        """
+        self._logger.info("Downloading file...")
+
+        try:
+            connector = aiohttp.TCPConnector(ssl=_SSL_CONTEXT)
+            async with aiohttp.ClientSession(timeout=self._timeout, connector=connector) as session:
+                async with session.get(url) as response:
+                    response.raise_for_status()
+                    data = await response.read()
+
+                    # 从 Content-Disposition 头中解析文件名
+                    content_disposition = response.headers.get("Content-Disposition", "")
+                    filename: Optional[str] = None
+
+                    if content_disposition:
+                        # 优先匹配 filename*=UTF-8''xxx 格式（RFC 5987）
+                        utf8_match = re.search(
+                            r"filename\*=UTF-8''([^;\s]+)",
+                            content_disposition,
+                            re.IGNORECASE,
+                        )
+                        if utf8_match:
+                            filename = unquote(utf8_match.group(1))
+                        else:
+                            # 匹配 filename="xxx" 或 filename=xxx 格式
+                            match = re.search(
+                                r'filename="?([^";\s]+)"?',
+                                content_disposition,
+                                re.IGNORECASE,
+                            )
+                            if match:
+                                filename = unquote(match.group(1))
+
+                    self._logger.info("File downloaded successfully")
+                    return data, filename
+
+        except Exception as e:
+            self._logger.error("File download failed:", str(e))
+            raise

aibot/client.py

@@ -0,0 +1,362 @@
+"""
+WSClient 核心客户端
+
+对标 Node.js SDK src/client.ts
+继承自 pyee.AsyncIOEventEmitter，组合 WsConnectionManager + MessageHandler + WeComApiClient。
+"""
+
+import asyncio
+from typing import Any, Dict, List, Optional, Tuple
+
+from pyee.asyncio import AsyncIOEventEmitter
+
+from .api import WeComApiClient
+from .crypto_utils import decrypt_file
+from .logger import DefaultLogger
+from .message_handler import MessageHandler
+from .types import WsCmd, WsFrame, WsFrameHeaders, WSClientOptions
+from .utils import generate_req_id
+from .ws import WsConnectionManager
+
+
+class WSClient(AsyncIOEventEmitter):
+    """
+    企业微信智能机器人 Python SDK 核心客户端
+
+    基于 asyncio + pyee 的事件驱动架构，提供 WebSocket 长连接消息收发能力。
+    """
+
+    def __init__(self, options: WSClientOptions) -> None:
+        super().__init__()
+
+        self._options = options
+        self._logger = options.logger or DefaultLogger()
+        self._started = False
+
+        # 初始化 API 客户端（仅用于文件下载）
+        self._api_client = WeComApiClient(
+            self._logger,
+            self._options.request_timeout,
+        )
+
+        # 初始化 WebSocket 管理器
+        self._ws_manager = WsConnectionManager(
+            self._logger,
+            heartbeat_interval=self._options.heartbeat_interval,
+            reconnect_base_delay=self._options.reconnect_interval,
+            max_reconnect_attempts=self._options.max_reconnect_attempts,
+            ws_url=self._options.ws_url or None,
+        )
+
+        # 设置认证凭证
+        self._ws_manager.set_credentials(self._options.bot_id, self._options.secret)
+
+        # 初始化消息处理器
+        self._message_handler = MessageHandler(self._logger)
+
+        # 绑定 WebSocket 事件
+        self._setup_ws_events()
+
+    def _setup_ws_events(self) -> None:
+        """设置 WebSocket 事件处理"""
+        self._ws_manager.on_connected = lambda: self.emit("connected")
+
+        def _on_authenticated() -> None:
+            self._logger.info("Authenticated")
+            self.emit("authenticated")
+
+        self._ws_manager.on_authenticated = _on_authenticated
+
+        self._ws_manager.on_disconnected = lambda reason: self.emit(
+            "disconnected", reason
+        )
+        self._ws_manager.on_reconnecting = lambda attempt: self.emit(
+            "reconnecting", attempt
+        )
+        self._ws_manager.on_error = lambda error: self.emit("error", error)
+        self._ws_manager.on_message = lambda frame: self._message_handler.handle_frame(
+            frame, self
+        )
+
+    async def connect(self) -> "WSClient":
+        """
+        建立 WebSocket 长连接
+
+        SDK 使用内置默认地址建立连接，连接成功后自动发送认证帧（bot_id + secret）。
+
+        :return: 返回 self，支持链式调用
+        """
+        if self._started:
+            self._logger.warn("Client already connected")
+            return self
+
+        self._logger.info("Establishing WebSocket connection...")
+        self._started = True
+
+        await self._ws_manager.connect()
+
+        return self
+
+    def disconnect(self) -> None:
+        """断开 WebSocket 连接"""
+        if not self._started:
+            self._logger.warn("Client not connected")
+            return
+
+        self._logger.info("Disconnecting...")
+        self._started = False
+        self._ws_manager.disconnect()
+        self._logger.info("Disconnected")
+
+    async def reply(
+        self,
+        frame: WsFrameHeaders,
+        body: Dict[str, Any],
+        cmd: Optional[str] = None,
+    ) -> WsFrame:
+        """
+        通过 WebSocket 通道发送回复消息（通用方法）
+
+        :param frame: 收到的原始 WebSocket 帧，透传 headers.req_id
+        :param body: 回复消息体
+        :param cmd: 发送的命令类型
+        :return: 回执帧
+        """
+        headers = frame.get("headers", {})
+        req_id = headers.get("req_id", "")
+        return await self._ws_manager.send_reply(req_id, body, cmd or WsCmd.RESPONSE)
+
+    async def reply_stream(
+        self,
+        frame: WsFrameHeaders,
+        stream_id: str,
+        content: str,
+        finish: bool = False,
+        msg_item: Optional[List[Dict[str, Any]]] = None,
+        feedback: Optional[Dict[str, Any]] = None,
+    ) -> WsFrame:
+        """
+        发送流式文本回复（便捷方法）
+
+        :param frame: 收到的原始 WebSocket 帧，透传 headers.req_id
+        :param stream_id: 流式消息 ID
+        :param content: 回复内容（支持 Markdown）
+        :param finish: 是否结束流式消息，默认 False
+        :param msg_item: 图文混排项（仅在 finish=True 时有效）
+        :param feedback: 反馈信息（仅在首次回复时设置）
+        :return: 回执帧
+        """
+        stream: Dict[str, Any] = {
+            "id": stream_id,
+            "finish": finish,
+            "content": content,
+        }
+
+        # msg_item 仅在 finish=True 时支持
+        if finish and msg_item and len(msg_item) > 0:
+            stream["msg_item"] = msg_item
+
+        # feedback 仅在首次回复时设置
+        if feedback:
+            stream["feedback"] = feedback
+
+        return await self.reply(
+            frame,
+            {
+                "msgtype": "stream",
+                "stream": stream,
+            },
+        )
+
+    async def reply_welcome(
+        self,
+        frame: WsFrameHeaders,
+        body: Dict[str, Any],
+    ) -> WsFrame:
+        """
+        发送欢迎语回复
+
+        注意：此方法需要使用对应事件（如 enter_chat）的 req_id 才能调用。
+        收到事件回调后需在 5 秒内发送回复，超时将无法发送欢迎语。
+
+        :param frame: 对应事件的 WebSocket 帧
+        :param body: 欢迎语消息体（支持文本或模板卡片格式）
+        :return: 回执帧
+        """
+        return await self.reply(frame, body, WsCmd.RESPONSE_WELCOME)
+
+    async def reply_template_card(
+        self,
+        frame: WsFrameHeaders,
+        template_card: Dict[str, Any],
+        feedback: Optional[Dict[str, Any]] = None,
+    ) -> WsFrame:
+        """
+        回复模板卡片消息
+
+        :param frame: 收到的原始 WebSocket 帧
+        :param template_card: 模板卡片内容
+        :param feedback: 反馈信息
+        :return: 回执帧
+        """
+        card = {**template_card, "feedback": feedback} if feedback else template_card
+        body = {
+            "msgtype": "template_card",
+            "template_card": card,
+        }
+        return await self.reply(frame, body)
+
+    async def reply_stream_with_card(
+        self,
+        frame: WsFrameHeaders,
+        stream_id: str,
+        content: str,
+        finish: bool = False,
+        msg_item: Optional[List[Dict[str, Any]]] = None,
+        stream_feedback: Optional[Dict[str, Any]] = None,
+        template_card: Optional[Dict[str, Any]] = None,
+        card_feedback: Optional[Dict[str, Any]] = None,
+    ) -> WsFrame:
+        """
+        发送流式消息 + 模板卡片组合回复
+
+        :param frame: 收到的原始 WebSocket 帧
+        :param stream_id: 流式消息 ID
+        :param content: 回复内容（支持 Markdown）
+        :param finish: 是否结束流式消息，默认 False
+        :param msg_item: 图文混排项（仅在 finish=True 时有效）
+        :param stream_feedback: 流式消息反馈信息（首次回复时设置）
+        :param template_card: 模板卡片内容（同一消息只能回复一次）
+        :param card_feedback: 模板卡片反馈信息
+        :return: 回执帧
+        """
+        stream: Dict[str, Any] = {
+            "id": stream_id,
+            "finish": finish,
+            "content": content,
+        }
+
+        if finish and msg_item and len(msg_item) > 0:
+            stream["msg_item"] = msg_item
+
+        if stream_feedback:
+            stream["feedback"] = stream_feedback
+
+        body: Dict[str, Any] = {
+            "msgtype": "stream_with_template_card",
+            "stream": stream,
+        }
+
+        if template_card:
+            card = (
+                {**template_card, "feedback": card_feedback}
+                if card_feedback
+                else template_card
+            )
+            body["template_card"] = card
+
+        return await self.reply(frame, body)
+
+    async def update_template_card(
+        self,
+        frame: WsFrameHeaders,
+        template_card: Dict[str, Any],
+        userids: Optional[List[str]] = None,
+    ) -> WsFrame:
+        """
+        更新模板卡片
+
+        注意：此方法需要使用对应事件（template_card_event）的 req_id 才能调用。
+        收到事件回调后需在 5 秒内发送回复，超时将无法更新卡片。
+
+        :param frame: 对应事件的 WebSocket 帧
+        :param template_card: 模板卡片内容（task_id 需跟回调收到的 task_id 一致）
+        :param userids: 要替换模版卡片消息的 userid 列表
+        :return: 回执帧
+        """
+        body: Dict[str, Any] = {
+            "response_type": "update_template_card",
+            "template_card": template_card,
+        }
+        if userids and len(userids) > 0:
+            body["userids"] = userids
+
+        return await self.reply(frame, body, WsCmd.RESPONSE_UPDATE)
+
+    async def send_message(
+        self,
+        chatid: str,
+        body: Dict[str, Any],
+    ) -> WsFrame:
+        """
+        主动发送消息
+
+        向指定会话（单聊或群聊）主动推送消息，无需依赖收到的回调帧。
+
+        :param chatid: 会话 ID，单聊填用户的 userid，群聊填对应群聊的 chatid
+        :param body: 消息体（支持 markdown 或 template_card 格式）
+        :return: 回执帧
+        """
+        req_id = generate_req_id(WsCmd.SEND_MSG)
+        full_body = {"chatid": chatid, **body}
+        return await self._ws_manager.send_reply(req_id, full_body, WsCmd.SEND_MSG)
+
+    async def download_file(
+        self, url: str, aes_key: Optional[str] = None
+    ) -> Tuple[bytes, Optional[str]]:
+        """
+        下载文件并使用 AES 密钥解密
+
+        :param url: 文件下载地址
+        :param aes_key: AES 解密密钥（Base64 编码），取自消息中 image.aeskey 或 file.aeskey
+        :return: (解密后的文件数据, 文件名)
+        """
+        self._logger.info("Downloading and decrypting file...")
+
+        try:
+            # 下载加密的文件数据
+            encrypted_data, filename = await self._api_client.download_file_raw(url)
+
+            # 如果没有提供 aes_key，直接返回原始数据
+            if not aes_key:
+                self._logger.warn("No aes_key provided, returning raw file data")
+                return encrypted_data, filename
+
+            # 使用独立的解密模块进行 AES-256-CBC 解密
+            decrypted_data = decrypt_file(encrypted_data, aes_key)
+
+            self._logger.info("File downloaded and decrypted successfully")
+            return decrypted_data, filename
+
+        except Exception as e:
+            self._logger.error(f"File download/decrypt failed: {e}")
+            raise
+
+    @property
+    def is_connected(self) -> bool:
+        """获取当前连接状态"""
+        return self._ws_manager.is_connected
+
+    @property
+    def api(self) -> WeComApiClient:
+        """获取 API 客户端实例（供高级用途使用）"""
+        return self._api_client
+
+    def run(self) -> None:
+        """
+        便捷方法：启动事件循环并连接
+
+        等价于:
+            asyncio.get_event_loop().run_until_complete(client.connect())
+            asyncio.get_event_loop().run_forever()
+        """
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        try:
+            loop.run_until_complete(self.connect())
+            loop.run_forever()
+        except KeyboardInterrupt:
+            self.disconnect()
+        finally:
+            loop.close()

aibot/crypto_utils.py

@@ -0,0 +1,73 @@
+"""
+加解密工具模块
+
+对标 Node.js SDK src/crypto.ts
+提供文件加解密相关的功能函数，使用 AES-256-CBC 解密。
+"""
+
+import base64
+
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+
+def decrypt_file(encrypted_data: bytes, aes_key: str) -> bytes:
+    """
+    使用 AES-256-CBC 解密文件
+
+    :param encrypted_data: 加密的文件数据
+    :param aes_key: Base64 编码的 AES-256 密钥
+    :return: 解密后的文件数据
+    :raises ValueError: 参数无效时
+    :raises RuntimeError: 解密失败时
+    """
+    if not encrypted_data:
+        raise ValueError("decrypt_file: encrypted_data is empty or not provided")
+
+    if not aes_key or not isinstance(aes_key, str):
+        raise ValueError("decrypt_file: aes_key must be a non-empty string")
+
+    # 将 Base64 编码的 aesKey 解码为 bytes
+    # Node.js 的 Buffer.from(str, 'base64') 会自动容错处理缺少的 '=' padding，
+    # 但 Python 的 base64.b64decode 严格要求长度是 4 的倍数，需要手动补齐。
+    padded_aes_key = aes_key + '=' * (4 - len(aes_key) % 4) if len(aes_key) % 4 != 0 else aes_key
+    key = base64.b64decode(padded_aes_key)
+
+    # IV 取 aesKey 解码后的前 16 字节
+    iv = key[:16]
+
+    try:
+        cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
+        decryptor = cipher.decryptor()
+
+        # 确保加密数据长度是 AES block size (16字节) 的倍数
+        # Node.js 的 setAutoPadding(false) 不会对不对齐的数据报错，
+        # 但 Python 的 cryptography 库会抛出 "Incorrect padding"。
+        # 这里手动补零对齐，后续通过 PKCS#7 去除 padding 来获得正确数据。
+        block_size = 16
+        remainder = len(encrypted_data) % block_size
+        if remainder != 0:
+            encrypted_data = encrypted_data + b'\x00' * (block_size - remainder)
+
+        # 解密（不自动处理 padding）
+        decrypted = decryptor.update(encrypted_data) + decryptor.finalize()
+
+        # 手动去除 PKCS#7 填充（支持 32 字节 block）
+        if len(decrypted) == 0:
+            raise ValueError("Decrypted data is empty")
+
+        pad_len = decrypted[-1]
+        if pad_len < 1 or pad_len > 32 or pad_len > len(decrypted):
+            raise ValueError(f"Invalid PKCS#7 padding value: {pad_len}")
+
+        # 验证所有 padding 字节是否一致
+        for i in range(len(decrypted) - pad_len, len(decrypted)):
+            if decrypted[i] != pad_len:
+                raise ValueError("Invalid PKCS#7 padding: padding bytes mismatch")
+
+        return decrypted[: len(decrypted) - pad_len]
+
+    except Exception as e:
+        raise RuntimeError(
+            f"decrypt_file: Decryption failed - {e}. "
+            "This may indicate corrupted data or an incorrect aesKey."
+        ) from e

aibot/logger.py

@@ -0,0 +1,47 @@
+"""
+默认日志实现
+
+对标 Node.js SDK src/logger.ts
+带有日志级别和时间戳的控制台日志
+"""
+
+import sys
+from datetime import datetime, timezone
+
+
+class DefaultLogger:
+    """默认日志实现，带有日志级别和时间戳的控制台日志"""
+
+    def __init__(self, prefix: str = "AiBotSDK"):
+        self._prefix = prefix
+
+    def _format_time(self) -> str:
+        return datetime.now(timezone.utc).isoformat()
+
+    def debug(self, message: str, *args: object) -> None:
+        print(
+            f"[{self._format_time()}] [{self._prefix}] [DEBUG] {message}",
+            *args,
+            file=sys.stderr,
+        )
+
+    def info(self, message: str, *args: object) -> None:
+        print(
+            f"[{self._format_time()}] [{self._prefix}] [INFO] {message}",
+            *args,
+            file=sys.stderr,
+        )
+
+    def warn(self, message: str, *args: object) -> None:
+        print(
+            f"[{self._format_time()}] [{self._prefix}] [WARN] {message}",
+            *args,
+            file=sys.stderr,
+        )
+
+    def error(self, message: str, *args: object) -> None:
+        print(
+            f"[{self._format_time()}] [{self._prefix}] [ERROR] {message}",
+            *args,
+            file=sys.stderr,
+        )

aibot/message_handler.py

@@ -0,0 +1,89 @@
+"""
+消息处理器
+
+对标 Node.js SDK src/message-handler.ts
+负责解析 WebSocket 帧并分发为具体的消息事件和事件回调。
+"""
+
+import json
+from typing import Any
+
+from .types import MessageType, WsCmd, WsFrame
+
+
+class MessageHandler:
+    """
+    消息处理器
+
+    负责解析 WebSocket 帧并分发为具体的消息事件和事件回调。
+    """
+
+    def __init__(self, logger: Any):
+        self._logger = logger
+
+    def handle_frame(self, frame: WsFrame, emitter: Any) -> None:
+        """
+        处理收到的 WebSocket 帧，解析并触发对应的消息/事件
+
+        :param frame: WebSocket 接收帧
+        :param emitter: WSClient 实例，用于触发事件
+        """
+        try:
+            body = frame.get("body")
+
+            if not body or not body.get("msgtype"):
+                self._logger.warn(
+                    f"Received invalid message format: {json.dumps(frame)[:200]}"
+                )
+                return
+
+            # 事件推送回调处理
+            if frame.get("cmd") == WsCmd.EVENT_CALLBACK:
+                self._handle_event_callback(frame, emitter)
+                return
+
+            # 消息推送回调处理
+            self._handle_message_callback(frame, emitter)
+        except Exception as e:
+            self._logger.error(f"Failed to handle message: {e}")
+
+    def _handle_message_callback(self, frame: WsFrame, emitter: Any) -> None:
+        """处理消息推送回调 (aibot_msg_callback)"""
+        body = frame.get("body", {})
+
+        # 触发通用消息事件
+        emitter.emit("message", frame)
+
+        # 根据 body 中的消息类型触发特定事件
+        msgtype = body.get("msgtype", "")
+
+        if msgtype == MessageType.Text:
+            emitter.emit("message.text", frame)
+        elif msgtype == MessageType.Image:
+            emitter.emit("message.image", frame)
+        elif msgtype == MessageType.Mixed:
+            emitter.emit("message.mixed", frame)
+        elif msgtype == MessageType.Voice:
+            emitter.emit("message.voice", frame)
+        elif msgtype == MessageType.File:
+            emitter.emit("message.file", frame)
+        else:
+            self._logger.debug(f"Received unhandled message type: {msgtype}")
+
+    def _handle_event_callback(self, frame: WsFrame, emitter: Any) -> None:
+        """处理事件推送回调 (aibot_event_callback)"""
+        body = frame.get("body", {})
+
+        # 触发通用事件
+        emitter.emit("event", frame)
+
+        # 根据事件类型触发特定事件
+        event = body.get("event", {})
+        event_type = event.get("eventtype") if isinstance(event, dict) else None
+
+        if event_type:
+            emitter.emit(f"event.{event_type}", frame)
+        else:
+            self._logger.debug(
+                f"Received event callback without eventtype: {json.dumps(body)[:200]}"
+            )