diff --git a/channel/wecom_bot/wecom_bot_channel.py b/channel/wecom_bot/wecom_bot_channel.py index ebc1104b..9ecbbefe 100644 --- a/channel/wecom_bot/wecom_bot_channel.py +++ b/channel/wecom_bot/wecom_bot_channel.py @@ -12,16 +12,19 @@ import hashlib import json import math import os +import re import threading import time import uuid import requests +import web import websocket from bridge.context import Context, ContextType from bridge.reply import Reply, ReplyType from channel.chat_channel import ChatChannel, check_prefix +from channel.wecom_bot.wecom_bot_crypt import WecomBotCrypt from channel.wecom_bot.wecom_bot_message import WecomBotMessage from common.expired_dict import ExpiredDict from common.log import logger @@ -32,6 +35,9 @@ from config import conf WECOM_WS_URL = "wss://openws.work.weixin.qq.com" HEARTBEAT_INTERVAL = 30 MEDIA_CHUNK_SIZE = 512 * 1024 # 512KB per chunk (before base64 encoding) +# Fixed URL path for the callback (webhook) HTTP server. The bot's +# receive-message URL must point at this path, e.g. http://host:9892/wecombot +CALLBACK_PATH = "/wecombot" def _escape_control_chars_inside_json_strings(s: str) -> str: @@ -97,6 +103,14 @@ class WecomBotChannel(ChatChannel): self._pending_lock = threading.Lock() self._stream_states = {} # req_id -> {"stream_id": str, "content": str} + # Transport mode: "websocket" (long connection) or "webhook" (HTTP callback) + self.mode = "websocket" + self._crypt = None + self._http_server = None + # stream_id -> {"committed", "current", "finished", "images", "last_access"} + self._callback_streams = ExpiredDict(60 * 10) # auto-expire after 10min (max poll window is 6min) + self._callback_lock = threading.Lock() + conf()["group_name_white_list"] = ["ALL_GROUP"] conf()["single_chat_prefix"] = [""] @@ -105,6 +119,11 @@ class WecomBotChannel(ChatChannel): # ------------------------------------------------------------------ def startup(self): + self.mode = conf().get("wecom_bot_mode", "websocket") + if self.mode == "webhook": + self._startup_callback() + return + self.bot_id = conf().get("wecom_bot_id", "") self.bot_secret = conf().get("wecom_bot_secret", "") @@ -127,6 +146,13 @@ class WecomBotChannel(ChatChannel): pass self._ws = None self._connected = False + if self._http_server: + try: + self._http_server.stop() + logger.info("[WecomBot] Callback HTTP server stopped") + except Exception as e: + logger.warning(f"[WecomBot] Error stopping HTTP server: {e}") + self._http_server = None # ------------------------------------------------------------------ # WebSocket connection @@ -183,6 +209,192 @@ class WecomBotChannel(ChatChannel): def _gen_req_id(self) -> str: return uuid.uuid4().hex[:16] + # ------------------------------------------------------------------ + # Callback (webhook) mode + # ------------------------------------------------------------------ + + def _startup_callback(self): + """Start an HTTP server that receives encrypted callbacks (webhook mode). + + The bot's "接收消息" URL in the WeCom admin console should point at this + server (any path is accepted). Verification (GET) and message delivery + (POST) are both handled by ``WecomBotCallbackController``. + """ + token = conf().get("wecom_bot_token", "") + aes_key = conf().get("wecom_bot_encoding_aes_key", "") + if not token or not aes_key: + err = "[WecomBot] callback mode requires wecom_bot_token and wecom_bot_encoding_aes_key" + logger.error(err) + self.report_startup_error(err) + return + + try: + # Enterprise-internal smart bot: receive_id is an empty string. + self._crypt = WecomBotCrypt(token, aes_key, "") + except Exception as e: + err = f"[WecomBot] invalid callback credentials: {e}" + logger.error(err) + self.report_startup_error(err) + return + + port = int(conf().get("wecom_bot_port", 9892)) + logger.info(f"[WecomBot] Starting callback (webhook) server on port {port}, path {CALLBACK_PATH} ...") + # Only serve the fixed callback path; everything else 404s instead of being + # treated as a (signature-failing) WeCom callback. + urls = (re.escape(CALLBACK_PATH), "channel.wecom_bot.wecom_bot_channel.WecomBotCallbackController") + app = web.application(urls, globals(), autoreload=False) + func = web.httpserver.StaticMiddleware(app.wsgifunc()) + func = web.httpserver.LogMiddleware(func) + server = web.httpserver.WSGIServer(("0.0.0.0", port), func) + self._http_server = server + self.report_startup_success() + try: + server.start() + except (KeyboardInterrupt, SystemExit): + server.stop() + + def _new_callback_stream(self, response_url: str = "") -> str: + """Create a new stream state and return its id.""" + stream_id = uuid.uuid4().hex[:16] + now = time.time() + with self._callback_lock: + self._callback_streams[stream_id] = { + "committed": "", + "current": "", + "finished": False, + "images": [], # list of (base64_str, md5_str), flushed only at finish + "image_urls": [], # public http(s) image urls (usable in response_url markdown) + "image_pending": False, # an image reply is being prepared; don't finish on text yet + "last_access": now, + "created_at": now, + "response_url": response_url or "", + "delivered": False, # final answer handed to WeCom via a poll + "url_sent": False, # final answer pushed via response_url (active reply) + } + return stream_id + + def _callback_handle_message(self, data: dict) -> dict: + """Handle a freshly-received user message in callback mode. + + Produces the context for async processing and returns the initial passive + reply (a stream packet with finish=false) so WeCom starts polling for the + agent's streamed answer. Returns ``None`` when there's nothing to reply + (e.g. an image/file silently cached for the next query). + """ + msg_id = data.get("msgid", "") + if msg_id and self.received_msgs.get(msg_id): + logger.debug(f"[WecomBot] Duplicate msg filtered: {msg_id}") + return None + if msg_id: + self.received_msgs[msg_id] = True + + chattype = data.get("chattype", "single") + is_group = chattype == "group" + + default_aeskey = conf().get("wecom_bot_encoding_aes_key", "") + result = self._build_context(data, is_group, default_aeskey=default_aeskey) + if not result: + return None + context, wecom_msg = result + + # response_url lets us actively reply once within 1h, used as a fallback + # when the agent finishes after WeCom stops polling (max ~6min window). + response_url = data.get("response_url", "") or "" + stream_id = self._new_callback_stream(response_url=response_url) + wecom_msg.stream_id = stream_id + context["wecom_stream_id"] = stream_id + context["on_event"] = self._make_callback_stream_callback(stream_id) + self.produce(context) + + # First passive reply: register the stream id, WeCom will poll for updates. + return { + "msgtype": "stream", + "stream": {"id": stream_id, "finish": False, "content": ""}, + } + + def _callback_handle_stream_poll(self, data: dict) -> dict: + """Handle a "流式消息刷新" poll: return the latest accumulated content.""" + stream_id = data.get("stream", {}).get("id", "") + with self._callback_lock: + state = self._callback_streams.get(stream_id) + if state is None: + # Unknown / expired stream: tell WeCom we're done to stop polling. + return {"msgtype": "stream", "stream": {"id": stream_id, "finish": True, "content": ""}} + state["last_access"] = time.time() + if state.get("url_sent"): + # Final answer already pushed via response_url; finish silently. + return {"msgtype": "stream", "stream": {"id": stream_id, "finish": True, "content": ""}} + # We never force-finish on a timer: while a task is still running the + # bubble should keep spinning until either the task finishes or the + # user cancels. If WeCom's 6min window closes before completion, the + # answer is delivered later via response_url instead. + finished = state["finished"] + content = state["committed"] + state["current"] + images = state["images"] if finished else [] + if finished: + state["delivered"] = True + logger.debug(f"[WecomBot] stream {stream_id} delivered via poll, len={len(content)}, images={len(images)}") + + stream = {"id": stream_id, "finish": finished, "content": content} + if images: + stream["msg_item"] = [ + {"msgtype": "image", "image": {"base64": b64, "md5": md5}} + for (b64, md5) in images + ] + return {"msgtype": "stream", "stream": stream} + + def _make_callback_stream_callback(self, stream_id: str): + """Build an on_event callback that accumulates agent output into stream state. + + Mirrors the websocket streaming behaviour: intermediate turns (text before + a tool call) are committed with a '---' separator; WeCom reads the full + accumulated content on each poll. + """ + def on_event(event: dict): + event_type = event.get("type") + edata = event.get("data", {}) + cancelled = False + with self._callback_lock: + state = self._callback_streams.get(stream_id) + if not state: + return + + if event_type == "turn_start": + state["current"] = "" + elif event_type == "message_update": + delta = edata.get("delta", "") + if delta: + state["current"] += delta + elif event_type == "message_end": + tool_calls = edata.get("tool_calls", []) + if tool_calls: + if state["current"].strip(): + state["committed"] += state["current"].strip() + "\n\n---\n\n" + state["current"] = "" + else: + state["committed"] += state["current"] + state["current"] = "" + elif event_type == "agent_cancelled": + # Mechanism 1: a cancelled run never reaches send(), so finalize + # its stream here to stop the "···" bubble immediately. + if state["current"]: + state["committed"] += state["current"] + state["current"] = "" + state["committed"] = state["committed"].rstrip() + if state["committed"].endswith("---"): + state["committed"] = state["committed"][:-3].rstrip() + if not state["committed"].strip(): + state["committed"] = "🛑 已中止" + state["finished"] = True + state["last_access"] = time.time() + cancelled = True + + if cancelled: + # Outside the lock: response_url fallback re-acquires it. + self._schedule_response_url_fallback(stream_id) + + return on_event + # ------------------------------------------------------------------ # Subscribe & heartbeat # ------------------------------------------------------------------ @@ -287,16 +499,31 @@ class WecomBotChannel(ChatChannel): chattype = body.get("chattype", "single") is_group = chattype == "group" + result = self._build_context(body, is_group) + if not result: + return + context, wecom_msg = result + wecom_msg.req_id = req_id + if req_id: + context["on_event"] = self._make_stream_callback(req_id) + self.produce(context) + + def _build_context(self, body: dict, is_group: bool, default_aeskey: str = ""): + """Parse a wecom message body into a Context, applying file-cache logic. + + Shared by both the websocket (long-connection) and callback (webhook) + receive paths. Returns ``(context, wecom_msg)`` when the message should be + handed to the agent, or ``None`` when it was consumed (cached image/file, + parse failure, etc.). + """ try: - wecom_msg = WecomBotMessage(body, is_group=is_group) + wecom_msg = WecomBotMessage(body, is_group=is_group, default_aeskey=default_aeskey) except NotImplementedError as e: logger.warning(f"[WecomBot] {e}") - return + return None except Exception as e: logger.error(f"[WecomBot] Failed to parse message: {e}", exc_info=True) - return - - wecom_msg.req_id = req_id + return None # File cache logic (same pattern as feishu) from channel.file_cache import get_file_cache @@ -314,13 +541,13 @@ class WecomBotChannel(ChatChannel): if hasattr(wecom_msg, "image_path") and wecom_msg.image_path: file_cache.add(session_id, wecom_msg.image_path, file_type="image") logger.info(f"[WecomBot] Image cached for session {session_id}") - return + return None if wecom_msg.ctype == ContextType.FILE: wecom_msg.prepare() file_cache.add(session_id, wecom_msg.content, file_type="file") logger.info(f"[WecomBot] File cached for session {session_id}: {wecom_msg.content}") - return + return None if wecom_msg.ctype == ContextType.TEXT: cached_files = file_cache.get(session_id) @@ -346,10 +573,9 @@ class WecomBotChannel(ChatChannel): msg=wecom_msg, no_need_at=True, ) - if context: - if req_id: - context["on_event"] = self._make_stream_callback(req_id) - self.produce(context) + if not context: + return None + return context, wecom_msg # ------------------------------------------------------------------ # Event callback @@ -490,11 +716,233 @@ class WecomBotChannel(ChatChannel): return context + # ------------------------------------------------------------------ + # Callback (webhook) send: write the final reply into the stream state + # so the next "流式消息刷新" poll returns it with finish=true. + # ------------------------------------------------------------------ + + def _callback_send(self, reply: Reply, context: Context): + msg = context.get("msg") + stream_id = getattr(msg, "stream_id", None) if msg else None + if not stream_id: + stream_id = context.get("wecom_stream_id") + if not stream_id: + logger.warning("[WecomBot] callback send without stream_id, dropping reply") + return + + if reply.type == ReplyType.TEXT: + self._callback_finalize_text(stream_id, reply.content) + elif reply.type in (ReplyType.IMAGE_URL, ReplyType.IMAGE): + self._callback_finalize_image(stream_id, reply.content) + elif reply.type == ReplyType.FILE: + # Passive callback replies only support text + image (base64); files + # are not supported by the protocol, so append a notice to whatever + # text the agent already streamed (do not drop it). + text = getattr(reply, "text_content", "") or "" + note = (text + "\n\n" if text else "") + "(文件无法在企微回调模式下直接发送)" + self._callback_finalize_text(stream_id, note, append=True) + elif reply.type in (ReplyType.VIDEO, ReplyType.VIDEO_URL, ReplyType.VOICE): + logger.warning(f"[WecomBot] reply type {reply.type} not supported in callback mode") + text = getattr(reply, "text_content", "") or "" + note = (text + "\n\n" if text else "") + "(该消息类型无法在企微回调模式下直接发送)" + self._callback_finalize_text(stream_id, note, append=True) + else: + self._callback_finalize_text(stream_id, str(reply.content)) + + def _callback_get_or_create_state(self, stream_id: str) -> dict: + state = self._callback_streams.get(stream_id) + if state is None: + now = time.time() + state = { + "committed": "", + "current": "", + "finished": False, + "images": [], + "image_urls": [], + "image_pending": False, + "last_access": now, + "created_at": now, + "response_url": "", + "delivered": False, + "url_sent": False, + } + self._callback_streams[stream_id] = state + return state + + def _callback_finalize_text(self, stream_id: str, content: str, append: bool = False): + with self._callback_lock: + state = self._callback_get_or_create_state(stream_id) + accumulated = (state["committed"] + state["current"]).strip() + if append and accumulated: + state["committed"] = (accumulated + "\n\n" + (content or "")).strip() + else: + state["committed"] = accumulated if accumulated else (content or "") + state["current"] = "" + state["last_access"] = time.time() + # Don't finish synchronously: chat_channel splits an image-with-caption + # reply into a TEXT send followed (0.3s later) by the IMAGE send. If the + # text finished the stream immediately, WeCom would close it before the + # image arrives. Defer the finish so a trailing image can merge in. + self._schedule_text_finish(stream_id) + + def _schedule_text_finish(self, stream_id: str, delay: float = 1.2): + def _run(): + time.sleep(delay) + with self._callback_lock: + state = self._callback_streams.get(stream_id) + if not state or state["finished"] or state.get("image_pending"): + return # already finished, or an image reply is on its way + state["finished"] = True + state["last_access"] = time.time() + self._schedule_response_url_fallback(stream_id) + + threading.Thread(target=_run, daemon=True, name=f"wecom-textfin-{stream_id}").start() + + def _callback_finalize_image(self, stream_id: str, img_path_or_url: str): + # Mark the image as pending up front (before the slow load/compress) so a + # preceding text finalize won't close the stream while we work. + with self._callback_lock: + self._callback_get_or_create_state(stream_id)["image_pending"] = True + b64md5 = self._load_image_base64(img_path_or_url) + with self._callback_lock: + state = self._callback_get_or_create_state(stream_id) + accumulated = (state["committed"] + state["current"]).strip() + state["current"] = "" + if b64md5: + state["images"].append(b64md5) + state["committed"] = accumulated + # Remember the public url (if any) so the response_url fallback + # can embed it as markdown when the poll window has closed. + if img_path_or_url.startswith(("http://", "https://")): + state["image_urls"].append(img_path_or_url) + else: + state["committed"] = accumulated or "[图片发送失败]" + state["finished"] = True + state["image_pending"] = False + state["last_access"] = time.time() + self._schedule_response_url_fallback(stream_id) + + # ------------------------------------------------------------------ + # Active reply fallback (response_url): rescue replies that finish after + # WeCom stops polling (the passive stream window is ~6 min from the user's + # message). A short delay lets an in-flight poll deliver first; only if no + # poll picks up the finished answer do we push it actively via response_url. + # ------------------------------------------------------------------ + + def _schedule_response_url_fallback(self, stream_id: str, delay: float = 3.0): + def _run(): + time.sleep(delay) + with self._callback_lock: + state = self._callback_streams.get(stream_id) + if not state: + return + if state.get("delivered") or state.get("url_sent"): + return # a poll already delivered (or fallback already ran) + response_url = state.get("response_url") or "" + if not response_url: + logger.warning( + f"[WecomBot] stream {stream_id} finished after poll window but no response_url; reply dropped" + ) + return + content = (state["committed"] + state["current"]).strip() + image_urls = list(state.get("image_urls") or []) + has_images = bool(state.get("images")) + state["url_sent"] = True + + self._send_via_response_url(stream_id, response_url, content, image_urls, has_images) + + threading.Thread(target=_run, daemon=True, name=f"wecom-respurl-{stream_id}").start() + + def _send_via_response_url(self, stream_id, response_url, content, image_urls, has_images): + """Push a one-shot active markdown reply to response_url (valid 1h, single use).""" + md = content or "" + if image_urls: + md += ("\n\n" if md else "") + "\n".join(f"![]({u})" for u in image_urls) + elif has_images: + md += ("\n\n" if md else "") + "(图片已生成,但因处理超时无法通过回调发送)" + if not md: + md = "(处理完成)" + payload = {"msgtype": "markdown", "markdown": {"content": md}} + try: + resp = requests.post(response_url, json=payload, timeout=15) + logger.info( + f"[WecomBot] response_url active reply sent for {stream_id}: " + f"status={resp.status_code}, body={resp.text[:200]}" + ) + except Exception as e: + logger.error(f"[WecomBot] response_url active reply failed for {stream_id}: {e}") + + def _load_image_base64(self, img_path_or_url: str): + """Load a local/remote image, ensure JPG/PNG within 10MB, return (base64, md5).""" + local_path = img_path_or_url + if local_path.startswith("file://"): + local_path = local_path[7:] + + # Temp files we create here (downloads/conversions/compressions) must be + # cleaned up afterwards; the caller's original local file must not be. + temp_files = [] + try: + if local_path.startswith(("http://", "https://")): + try: + resp = requests.get(local_path, timeout=30) + resp.raise_for_status() + tmp_path = f"/tmp/wecom_cb_img_{uuid.uuid4().hex[:8]}" + with open(tmp_path, "wb") as f: + f.write(resp.content) + temp_files.append(tmp_path) + local_path = tmp_path + except Exception as e: + logger.error(f"[WecomBot] Failed to download image for callback reply: {e}") + return None + + if not os.path.exists(local_path): + logger.error(f"[WecomBot] Image file not found: {local_path}") + return None + + formatted = self._ensure_image_format(local_path) + if not formatted: + return None + if formatted != local_path: + temp_files.append(formatted) + local_path = formatted + + # Unlike the long-connection path (which uploads and sends only a tiny + # media_id), the callback reply embeds the whole image as base64 inside + # an AES-encrypted body that is returned on EVERY poll. Empirically a + # ~1.5MB image (base64 ~2.1MB, encrypted ~2.8MB) makes WeCom reject the + # finish packet and poll forever, so cap well below that. + callback_max_size = 512 * 1024 + if os.path.getsize(local_path) > callback_max_size: + compressed = self._compress_image(local_path, callback_max_size) + if compressed: + temp_files.append(compressed) + local_path = compressed + else: + logger.warning("[WecomBot] callback image compress failed; sending original (may be rejected)") + + try: + with open(local_path, "rb") as f: + raw = f.read() + return base64.b64encode(raw).decode("utf-8"), hashlib.md5(raw).hexdigest() + except Exception as e: + logger.error(f"[WecomBot] Failed to read image for callback reply: {e}") + return None + finally: + for path in temp_files: + try: + os.remove(path) + except OSError: + pass + # ------------------------------------------------------------------ # Send reply # ------------------------------------------------------------------ def send(self, reply: Reply, context: Context): + if self.mode == "webhook": + self._callback_send(reply, context) + return + msg = context.get("msg") is_group = context.get("isgroup", False) receiver = context.get("receiver", "") @@ -906,3 +1354,85 @@ class WecomBotChannel(ChatChannel): else: logger.error("[WecomBot] Failed to get media_id from finish response") return media_id + + +class WecomBotCallbackController: + """HTTP controller for wecom bot callback (webhook) mode. + + - GET : URL verification (echo the decrypted echostr). + - POST : encrypted message / stream-refresh / event callbacks; returns an + encrypted passive reply (or "success" for an empty reply). + """ + + @staticmethod + def _channel() -> "WecomBotChannel": + return WecomBotChannel() + + def GET(self): + channel = self._channel() + params = web.input(msg_signature="", timestamp="", nonce="", echostr="") + if not channel._crypt: + return "wecom bot callback not ready" + ret, echo = channel._crypt.verify_url( + params.msg_signature, params.timestamp, params.nonce, params.echostr + ) + if ret != 0: + logger.error(f"[WecomBot] URL verify failed: ret={ret}") + return "verify fail" + if isinstance(echo, bytes): + echo = echo.decode("utf-8") + return echo + + def POST(self): + channel = self._channel() + if not channel._crypt: + return "success" + + params = web.input(msg_signature="", timestamp="", nonce="") + body = web.data() + ret, plain = channel._crypt.decrypt_msg( + body, params.msg_signature, params.timestamp, params.nonce + ) + if ret != 0: + logger.error(f"[WecomBot] callback decrypt failed: ret={ret}") + return "success" + + try: + data = json.loads(plain) + except Exception as e: + logger.error(f"[WecomBot] callback json parse failed: {e}") + return "success" + + msgtype = data.get("msgtype", "") + # Stream polls arrive ~1/s; logging each is noisy, so only log non-poll + # callbacks here (poll completion is logged in the stream-poll handler). + if msgtype != "stream": + logger.debug(f"[WecomBot] callback received msgtype={msgtype}") + + try: + if msgtype == "stream": + reply = channel._callback_handle_stream_poll(data) + elif msgtype == "event": + event_type = data.get("event", {}).get("eventtype", "") + logger.info(f"[WecomBot] callback event: {event_type}") + reply = None + elif msgtype in ("text", "image", "voice", "file", "video", "mixed"): + reply = channel._callback_handle_message(data) + else: + logger.warning(f"[WecomBot] unsupported callback msgtype: {msgtype}") + reply = None + except Exception as e: + logger.error(f"[WecomBot] callback handling error: {e}", exc_info=True) + reply = None + + if not reply: + # Empty reply package is acceptable. + return "success" + + plain_reply = json.dumps(reply, ensure_ascii=False) + ret, enc = channel._crypt.encrypt_msg(plain_reply, params.nonce, params.timestamp) + if ret != 0: + logger.error(f"[WecomBot] callback encrypt failed: ret={ret}") + return "success" + web.header("Content-Type", "application/json; charset=utf-8") + return json.dumps(enc, ensure_ascii=False) diff --git a/channel/wecom_bot/wecom_bot_crypt.py b/channel/wecom_bot/wecom_bot_crypt.py new file mode 100644 index 00000000..262ca3f1 --- /dev/null +++ b/channel/wecom_bot/wecom_bot_crypt.py @@ -0,0 +1,203 @@ +""" +WeCom (企业微信) smart-bot callback message encryption/decryption. + +Adapted from the official `WXBizJsonMsgCrypt` sample (JSON variant) used by the +AI bot callback (webhook) mode. The bot's receive-message callback delivers +AES-256-CBC encrypted JSON payloads, and passive replies must be encrypted the +same way before being returned in the HTTP response. + +For an enterprise-internal smart bot, ``receive_id`` is always an empty string. +""" + +import base64 +import hashlib +import random +import socket +import struct +import time + +from Crypto.Cipher import AES + +from common.log import logger + +# Error codes (mirrors the official ierror.py) +WXBizMsgCrypt_OK = 0 +WXBizMsgCrypt_ValidateSignature_Error = -40001 +WXBizMsgCrypt_ParseJson_Error = -40002 +WXBizMsgCrypt_ComputeSignature_Error = -40003 +WXBizMsgCrypt_IllegalAesKey = -40004 +WXBizMsgCrypt_ValidateCorpid_Error = -40005 +WXBizMsgCrypt_EncryptAES_Error = -40006 +WXBizMsgCrypt_DecryptAES_Error = -40007 +WXBizMsgCrypt_IllegalBuffer = -40008 +WXBizMsgCrypt_EncodeBase64_Error = -40009 +WXBizMsgCrypt_DecodeBase64_Error = -40010 +WXBizMsgCrypt_GenReturnJson_Error = -40011 + + +class FormatException(Exception): + pass + + +def _gen_sha1(token, timestamp, nonce, encrypt): + """Compute the WeCom message signature with SHA1 over the sorted parts.""" + try: + if isinstance(encrypt, bytes): + encrypt = encrypt.decode("utf-8") + sortlist = [str(token), str(timestamp), str(nonce), str(encrypt)] + sortlist.sort() + sha = hashlib.sha1() + sha.update("".join(sortlist).encode("utf-8")) + return WXBizMsgCrypt_OK, sha.hexdigest() + except Exception as e: + logger.error(f"[WecomBot] compute signature error: {e}") + return WXBizMsgCrypt_ComputeSignature_Error, None + + +class _PKCS7Encoder: + """PKCS#7 padding with a 32-byte block size (AES-256).""" + + block_size = 32 + + def encode(self, text: bytes) -> bytes: + text_length = len(text) + amount_to_pad = self.block_size - (text_length % self.block_size) + if amount_to_pad == 0: + amount_to_pad = self.block_size + pad = bytes([amount_to_pad]) + return text + pad * amount_to_pad + + def decode(self, decrypted: bytes) -> bytes: + pad = decrypted[-1] + if pad < 1 or pad > 32: + pad = 0 + return decrypted[:-pad] if pad else decrypted + + +class _Prpcrypt: + """AES-256-CBC encrypt/decrypt for WeCom callback messages.""" + + def __init__(self, key: bytes): + self.key = key + self.mode = AES.MODE_CBC + + def encrypt(self, text: str, receive_id: str): + text_bytes = text.encode() + # 16-byte random prefix + network-order length + body + receive_id + text_bytes = ( + self._get_random_str() + + struct.pack("I", socket.htonl(len(text_bytes))) + + text_bytes + + receive_id.encode() + ) + text_bytes = _PKCS7Encoder().encode(text_bytes) + try: + cryptor = AES.new(self.key, self.mode, self.key[:16]) + ciphertext = cryptor.encrypt(text_bytes) + return WXBizMsgCrypt_OK, base64.b64encode(ciphertext) + except Exception as e: + logger.error(f"[WecomBot] AES encrypt error: {e}") + return WXBizMsgCrypt_EncryptAES_Error, None + + def decrypt(self, text, receive_id: str): + try: + cryptor = AES.new(self.key, self.mode, self.key[:16]) + plain_text = cryptor.decrypt(base64.b64decode(text)) + except Exception as e: + logger.error(f"[WecomBot] AES decrypt error: {e}") + return WXBizMsgCrypt_DecryptAES_Error, None + try: + pad = plain_text[-1] + content = plain_text[16:-pad] + json_len = socket.ntohl(struct.unpack("I", content[:4])[0]) + json_content = content[4 : json_len + 4].decode("utf-8") + from_receive_id = content[json_len + 4 :].decode("utf-8") + except Exception as e: + logger.error(f"[WecomBot] illegal buffer when decrypting: {e}") + return WXBizMsgCrypt_IllegalBuffer, None + if from_receive_id != receive_id: + logger.error( + f"[WecomBot] receive_id not match: expect={receive_id}, got={from_receive_id}" + ) + return WXBizMsgCrypt_ValidateCorpid_Error, None + return WXBizMsgCrypt_OK, json_content + + @staticmethod + def _get_random_str() -> bytes: + return str(random.randint(1000000000000000, 9999999999999999)).encode() + + +class WecomBotCrypt: + """High-level helper for verifying URLs and (de)crypting callback messages.""" + + def __init__(self, token: str, encoding_aes_key: str, receive_id: str = ""): + try: + self.key = base64.b64decode(encoding_aes_key + "=") + assert len(self.key) == 32 + except Exception: + raise FormatException("[WecomBot] invalid EncodingAESKey") + self.token = token + self.receive_id = receive_id + + def verify_url(self, msg_signature, timestamp, nonce, echostr): + ret, signature = _gen_sha1(self.token, timestamp, nonce, echostr) + if ret != 0: + return ret, None + if signature != msg_signature: + return WXBizMsgCrypt_ValidateSignature_Error, None + pc = _Prpcrypt(self.key) + return pc.decrypt(echostr, self.receive_id) + + def encrypt_msg(self, reply_msg: str, nonce: str, timestamp: str = None): + """Encrypt a passive-reply JSON string and return the full response JSON. + + Returns (ret, response_dict). On success ret==0 and response_dict is a + dict with encrypt/msgsignature/timestamp/nonce fields. + """ + pc = _Prpcrypt(self.key) + ret, encrypt = pc.encrypt(reply_msg, self.receive_id) + if ret != 0: + return ret, None + encrypt = encrypt.decode("utf-8") + if timestamp is None: + timestamp = str(int(time.time())) + ret, signature = _gen_sha1(self.token, timestamp, nonce, encrypt) + if ret != 0: + return ret, None + return WXBizMsgCrypt_OK, { + "encrypt": encrypt, + "msgsignature": signature, + "timestamp": timestamp, + "nonce": nonce, + } + + def decrypt_msg(self, post_data, msg_signature, timestamp, nonce): + """Verify signature and decrypt the encrypted callback payload. + + ``post_data`` may be the raw request body (bytes/str) containing + ``{"encrypt": "..."}`` or the already-extracted encrypt string. + Returns (ret, plaintext_json_str). + """ + import json + + encrypt = None + if isinstance(post_data, (bytes, bytearray)): + post_data = post_data.decode("utf-8") + if isinstance(post_data, str): + try: + encrypt = json.loads(post_data).get("encrypt") + except Exception: + encrypt = post_data + elif isinstance(post_data, dict): + encrypt = post_data.get("encrypt") + if not encrypt: + return WXBizMsgCrypt_ParseJson_Error, None + + ret, signature = _gen_sha1(self.token, timestamp, nonce, encrypt) + if ret != 0: + return ret, None + if signature != msg_signature: + logger.error("[WecomBot] callback signature not match") + return WXBizMsgCrypt_ValidateSignature_Error, None + pc = _Prpcrypt(self.key) + return pc.decrypt(encrypt, self.receive_id) diff --git a/channel/wecom_bot/wecom_bot_message.py b/channel/wecom_bot/wecom_bot_message.py index 16b7dfdc..f021a24f 100644 --- a/channel/wecom_bot/wecom_bot_message.py +++ b/channel/wecom_bot/wecom_bot_message.py @@ -87,11 +87,14 @@ def _get_tmp_dir() -> str: class WecomBotMessage(ChatMessage): """Message wrapper for wecom bot (websocket long-connection mode).""" - def __init__(self, msg_body: dict, is_group: bool = False): + def __init__(self, msg_body: dict, is_group: bool = False, default_aeskey: str = ""): super().__init__(msg_body) self.msg_id = msg_body.get("msgid") self.create_time = msg_body.get("create_time") self.is_group = is_group + # In callback (webhook) mode the media bodies carry no per-message aeskey; + # the download url is encrypted with the bot's EncodingAESKey instead. + self._default_aeskey = default_aeskey msg_type = msg_body.get("msgtype") from_userid = msg_body.get("from", {}).get("userid", "") @@ -113,7 +116,7 @@ class WecomBotMessage(ChatMessage): self.ctype = ContextType.IMAGE image_info = msg_body.get("image", {}) image_url = image_info.get("url", "") - aeskey = image_info.get("aeskey", "") + aeskey = image_info.get("aeskey", "") or self._default_aeskey tmp_dir = _get_tmp_dir() image_path = os.path.join(tmp_dir, f"wecom_{self.msg_id}.png") @@ -147,7 +150,7 @@ class WecomBotMessage(ChatMessage): elif item_type == "image": img_info = item.get("image", {}) img_url = img_info.get("url", "") - img_aeskey = img_info.get("aeskey", "") + img_aeskey = img_info.get("aeskey", "") or self._default_aeskey img_path = os.path.join(tmp_dir, f"wecom_{self.msg_id}_{idx}.png") try: img_data = _decrypt_media(img_url, img_aeskey) @@ -166,7 +169,7 @@ class WecomBotMessage(ChatMessage): self.ctype = ContextType.FILE file_info = msg_body.get("file", {}) file_url = file_info.get("url", "") - aeskey = file_info.get("aeskey", "") + aeskey = file_info.get("aeskey", "") or self._default_aeskey tmp_dir = _get_tmp_dir() base_path = os.path.join(tmp_dir, f"wecom_{self.msg_id}") self.content = base_path @@ -188,7 +191,7 @@ class WecomBotMessage(ChatMessage): self.ctype = ContextType.FILE video_info = msg_body.get("video", {}) video_url = video_info.get("url", "") - aeskey = video_info.get("aeskey", "") + aeskey = video_info.get("aeskey", "") or self._default_aeskey tmp_dir = _get_tmp_dir() self.content = os.path.join(tmp_dir, f"wecom_{self.msg_id}.mp4") diff --git a/config.py b/config.py index 6d81f5ff..d5e98950 100644 --- a/config.py +++ b/config.py @@ -183,6 +183,11 @@ available_setting = { # WeCom smart bot config (long connection mode) "wecom_bot_id": "", # WeCom smart bot BotID "wecom_bot_secret": "", # WeCom smart bot long-connection secret + # WeCom smart bot transport mode: "websocket" (long connection) or "webhook" (HTTP callback) + "wecom_bot_mode": "websocket", + "wecom_bot_token": "", # webhook mode: Token configured on the bot's receive-message URL + "wecom_bot_encoding_aes_key": "", # webhook mode: EncodingAESKey configured on the bot's receive-message URL + "wecom_bot_port": 9892, # webhook mode: local HTTP server port for the receive-message URL # Telegram config "telegram_token": "", # Bot token from @BotFather "telegram_proxy": "", # Optional HTTP/SOCKS5 proxy, e.g. http://127.0.0.1:7890 or socks5://127.0.0.1:1080 (empty falls back to env vars)