| 1 | """WeCom (Enterprise WeChat) channel implementation using wecom_aibot_sdk.""" |
| 2 | |
| 3 | import asyncio |
| 4 | import base64 |
| 5 | import hashlib |
| 6 | import importlib.util |
| 7 | import os |
| 8 | import re |
| 9 | from collections import OrderedDict |
| 10 | from pathlib import Path |
| 11 | from typing import Any |
| 12 | |
| 13 | from loguru import logger |
| 14 | |
| 15 | from nanobot.bus.events import OutboundMessage |
| 16 | from nanobot.bus.queue import MessageBus |
| 17 | from nanobot.channels.base import BaseChannel |
| 18 | from nanobot.config.paths import get_media_dir |
| 19 | from nanobot.config.schema import Base |
| 20 | from pydantic import Field |
| 21 | |
| 22 | WECOM_AVAILABLE = importlib.util.find_spec("wecom_aibot_sdk") is not None |
| 23 | |
| 24 | # Upload safety limits (matching QQ channel defaults) |
| 25 | WECOM_UPLOAD_MAX_BYTES = 1024 * 1024 * 200 # 200MB |
| 26 | |
| 27 | # Replace unsafe characters with "_", keep Chinese and common safe punctuation. |
| 28 | _SAFE_NAME_RE = re.compile(r"[^\w.\-()\[\]（）【】\u4e00-\u9fff]+", re.UNICODE) |
| 29 | |
| 30 | |
| 31 | def _sanitize_filename(name: str) -> str: |
| 32 | """Sanitize filename to avoid traversal and problematic chars.""" |
| 33 | name = (name or "").strip() |
| 34 | name = Path(name).name |
| 35 | name = _SAFE_NAME_RE.sub("_", name).strip("._ ") |
| 36 | return name |
| 37 | |
| 38 | |
| 39 | _IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"} |
| 40 | _VIDEO_EXTS = {".mp4", ".avi", ".mov"} |
| 41 | _AUDIO_EXTS = {".amr", ".mp3", ".wav", ".ogg"} |
| 42 | |
| 43 | |
| 44 | def _guess_wecom_media_type(filename: str) -> str: |
| 45 | """Classify file extension as WeCom media_type string.""" |
| 46 | ext = Path(filename).suffix.lower() |
| 47 | if ext in _IMAGE_EXTS: |
| 48 | return "image" |
| 49 | if ext in _VIDEO_EXTS: |
| 50 | return "video" |
| 51 | if ext in _AUDIO_EXTS: |
| 52 | return "voice" |
| 53 | return "file" |
| 54 | |
| 55 | class WecomConfig(Base): |
| 56 | """WeCom (Enterprise WeChat) AI Bot channel configuration.""" |
| 57 | |
| 58 | enabled: bool = False |
| 59 | bot_id: str = "" |
| 60 | secret: str = "" |
| 61 | allow_from: list[str] = Field(default_factory=list) |
| 62 | welcome_message: str = "" |
| 63 | |
| 64 | |
| 65 | # Message type display mapping |
| 66 | MSG_TYPE_MAP = { |
| 67 | "image": "[image]", |
| 68 | "voice": "[voice]", |
| 69 | "file": "[file]", |
| 70 | "mixed": "[mixed content]", |
| 71 | } |
| 72 | |
| 73 | |
| 74 | class WecomChannel(BaseChannel): |
| 75 | """ |
| 76 | WeCom (Enterprise WeChat) channel using WebSocket long connection. |
| 77 | |
| 78 | Uses WebSocket to receive events - no public IP or webhook required. |
| 79 | |
| 80 | Requires: |
| 81 | - Bot ID and Secret from WeCom AI Bot platform |
| 82 | """ |
| 83 | |
| 84 | name = "wecom" |
| 85 | display_name = "WeCom" |
| 86 | |
| 87 | @classmethod |
| 88 | def default_config(cls) -> dict[str, Any]: |
| 89 | return WecomConfig().model_dump(by_alias=True) |
| 90 | |
| 91 | def __init__(self, config: Any, bus: MessageBus): |
| 92 | if isinstance(config, dict): |
| 93 | config = WecomConfig.model_validate(config) |
| 94 | super().__init__(config, bus) |
| 95 | self.config: WecomConfig = config |
| 96 | self._client: Any = None |
| 97 | self._processed_message_ids: OrderedDict[str, None] = OrderedDict() |
| 98 | self._loop: asyncio.AbstractEventLoop | None = None |
| 99 | self._generate_req_id = None |
| 100 | # Store frame headers for each chat to enable replies |
| 101 | self._chat_frames: dict[str, Any] = {} |
| 102 | |
| 103 | async def start(self) -> None: |
| 104 | """Start the WeCom bot with WebSocket long connection.""" |
| 105 | if not WECOM_AVAILABLE: |
| 106 | logger.error("WeCom SDK not installed. Run: pip install echo-director-agent[wecom]") |
| 107 | return |
| 108 | |
| 109 | if not self.config.bot_id or not self.config.secret: |
| 110 | logger.error("WeCom bot_id and secret not configured") |
| 111 | return |
| 112 | |
| 113 | from wecom_aibot_sdk import WSClient, generate_req_id |
| 114 | |
| 115 | self._running = True |
| 116 | self._loop = asyncio.get_running_loop() |
| 117 | self._generate_req_id = generate_req_id |
| 118 | |
| 119 | # Create WebSocket client |
| 120 | self._client = WSClient({ |
| 121 | "bot_id": self.config.bot_id, |
| 122 | "secret": self.config.secret, |
| 123 | "reconnect_interval": 1000, |
| 124 | "max_reconnect_attempts": -1, # Infinite reconnect |
| 125 | "heartbeat_interval": 30000, |
| 126 | }) |
| 127 | |
| 128 | # Register event handlers |
| 129 | self._client.on("connected", self._on_connected) |
| 130 | self._client.on("authenticated", self._on_authenticated) |
| 131 | self._client.on("disconnected", self._on_disconnected) |
| 132 | self._client.on("error", self._on_error) |
| 133 | self._client.on("message.text", self._on_text_message) |
| 134 | self._client.on("message.image", self._on_image_message) |
| 135 | self._client.on("message.voice", self._on_voice_message) |
| 136 | self._client.on("message.file", self._on_file_message) |
| 137 | self._client.on("message.mixed", self._on_mixed_message) |
| 138 | self._client.on("event.enter_chat", self._on_enter_chat) |
| 139 | |
| 140 | logger.info("WeCom bot starting with WebSocket long connection") |
| 141 | logger.info("No public IP required - using WebSocket to receive events") |
| 142 | |
| 143 | # Connect |
| 144 | await self._client.connect_async() |
| 145 | |
| 146 | # Keep running until stopped |
| 147 | while self._running: |
| 148 | await asyncio.sleep(1) |
| 149 | |
| 150 | async def stop(self) -> None: |
| 151 | """Stop the WeCom bot.""" |
| 152 | self._running = False |
| 153 | if self._client: |
| 154 | await self._client.disconnect() |
| 155 | logger.info("WeCom bot stopped") |
| 156 | |
| 157 | async def _on_connected(self, frame: Any) -> None: |
| 158 | """Handle WebSocket connected event.""" |
| 159 | logger.info("WeCom WebSocket connected") |
| 160 | |
| 161 | async def _on_authenticated(self, frame: Any) -> None: |
| 162 | """Handle authentication success event.""" |
| 163 | logger.info("WeCom authenticated successfully") |
| 164 | |
| 165 | async def _on_disconnected(self, frame: Any) -> None: |
| 166 | """Handle WebSocket disconnected event.""" |
| 167 | reason = frame.body if hasattr(frame, 'body') else str(frame) |
| 168 | logger.warning("WeCom WebSocket disconnected: {}", reason) |
| 169 | |
| 170 | async def _on_error(self, frame: Any) -> None: |
| 171 | """Handle error event.""" |
| 172 | logger.error("WeCom error: {}", frame) |
| 173 | |
| 174 | async def _on_text_message(self, frame: Any) -> None: |
| 175 | """Handle text message.""" |
| 176 | await self._process_message(frame, "text") |
| 177 | |
| 178 | async def _on_image_message(self, frame: Any) -> None: |
| 179 | """Handle image message.""" |
| 180 | await self._process_message(frame, "image") |
| 181 | |
| 182 | async def _on_voice_message(self, frame: Any) -> None: |
| 183 | """Handle voice message.""" |
| 184 | await self._process_message(frame, "voice") |
| 185 | |
| 186 | async def _on_file_message(self, frame: Any) -> None: |
| 187 | """Handle file message.""" |
| 188 | await self._process_message(frame, "file") |
| 189 | |
| 190 | async def _on_mixed_message(self, frame: Any) -> None: |
| 191 | """Handle mixed content message.""" |
| 192 | await self._process_message(frame, "mixed") |
| 193 | |
| 194 | async def _on_enter_chat(self, frame: Any) -> None: |
| 195 | """Handle enter_chat event (user opens chat with bot).""" |
| 196 | try: |
| 197 | # Extract body from WsFrame dataclass or dict |
| 198 | if hasattr(frame, 'body'): |
| 199 | body = frame.body or {} |
| 200 | elif isinstance(frame, dict): |
| 201 | body = frame.get("body", frame) |
| 202 | else: |
| 203 | body = {} |
| 204 | |
| 205 | chat_id = body.get("chatid", "") if isinstance(body, dict) else "" |
| 206 | |
| 207 | if chat_id and self.config.welcome_message: |
| 208 | await self._client.reply_welcome(frame, { |
| 209 | "msgtype": "text", |
| 210 | "text": {"content": self.config.welcome_message}, |
| 211 | }) |
| 212 | except Exception as e: |
| 213 | logger.error("Error handling enter_chat: {}", e) |
| 214 | |
| 215 | async def _process_message(self, frame: Any, msg_type: str) -> None: |
| 216 | """Process incoming message and forward to bus.""" |
| 217 | try: |
| 218 | # Extract body from WsFrame dataclass or dict |
| 219 | if hasattr(frame, 'body'): |
| 220 | body = frame.body or {} |
| 221 | elif isinstance(frame, dict): |
| 222 | body = frame.get("body", frame) |
| 223 | else: |
| 224 | body = {} |
| 225 | |
| 226 | # Ensure body is a dict |
| 227 | if not isinstance(body, dict): |
| 228 | logger.warning("Invalid body type: {}", type(body)) |
| 229 | return |
| 230 | |
| 231 | # Extract message info |
| 232 | msg_id = body.get("msgid", "") |
| 233 | if not msg_id: |
| 234 | msg_id = f"{body.get('chatid', '')}_{body.get('sendertime', '')}" |
| 235 | |
| 236 | # Deduplication check |
| 237 | if msg_id in self._processed_message_ids: |
| 238 | return |
| 239 | self._processed_message_ids[msg_id] = None |
| 240 | |
| 241 | # Trim cache |
| 242 | while len(self._processed_message_ids) > 1000: |
| 243 | self._processed_message_ids.popitem(last=False) |
| 244 | |
| 245 | # Extract sender info from "from" field (SDK format) |
| 246 | from_info = body.get("from", {}) |
| 247 | sender_id = from_info.get("userid", "unknown") if isinstance(from_info, dict) else "unknown" |
| 248 | |
| 249 | # For single chat, chatid is the sender's userid |
| 250 | # For group chat, chatid is provided in body |
| 251 | chat_type = body.get("chattype", "single") |
| 252 | chat_id = body.get("chatid", sender_id) |
| 253 | |
| 254 | content_parts = [] |
| 255 | media_paths: list[str] = [] |
| 256 | |
| 257 | if msg_type == "text": |
| 258 | text = body.get("text", {}).get("content", "") |
| 259 | if text: |
| 260 | content_parts.append(text) |
| 261 | |
| 262 | elif msg_type == "image": |
| 263 | image_info = body.get("image", {}) |
| 264 | file_url = image_info.get("url", "") |
| 265 | aes_key = image_info.get("aeskey", "") |
| 266 | |
| 267 | if file_url and aes_key: |
| 268 | file_path = await self._download_and_save_media(file_url, aes_key, "image") |
| 269 | if file_path: |
| 270 | filename = os.path.basename(file_path) |
| 271 | content_parts.append(f"[image: {filename}]") |
| 272 | media_paths.append(file_path) |
| 273 | else: |
| 274 | content_parts.append("[image: download failed]") |
| 275 | else: |
| 276 | content_parts.append("[image: download failed]") |
| 277 | |
| 278 | elif msg_type == "voice": |
| 279 | voice_info = body.get("voice", {}) |
| 280 | # Voice message already contains transcribed content from WeCom |
| 281 | voice_content = voice_info.get("content", "") |
| 282 | if voice_content: |
| 283 | content_parts.append(f"[voice] {voice_content}") |
| 284 | else: |
| 285 | content_parts.append("[voice]") |
| 286 | |
| 287 | elif msg_type == "file": |
| 288 | file_info = body.get("file", {}) |
| 289 | file_url = file_info.get("url", "") |
| 290 | aes_key = file_info.get("aeskey", "") |
| 291 | file_name = file_info.get("name", "unknown") |
| 292 | |
| 293 | if file_url and aes_key: |
| 294 | file_path = await self._download_and_save_media(file_url, aes_key, "file", file_name) |
| 295 | if file_path: |
| 296 | content_parts.append(f"[file: {file_name}]") |
| 297 | media_paths.append(file_path) |
| 298 | else: |
| 299 | content_parts.append(f"[file: {file_name}: download failed]") |
| 300 | else: |
| 301 | content_parts.append(f"[file: {file_name}: download failed]") |
| 302 | |
| 303 | elif msg_type == "mixed": |
| 304 | # Mixed content contains multiple message items |
| 305 | msg_items = body.get("mixed", {}).get("msg_item", []) |
| 306 | for item in msg_items: |
| 307 | item_type = item.get("msgtype", "") |
| 308 | if item_type == "text": |
| 309 | text = item.get("text", {}).get("content", "") |
| 310 | if text: |
| 311 | content_parts.append(text) |
| 312 | elif item_type == "image": |
| 313 | file_url = item.get("image", {}).get("url", "") |
| 314 | aes_key = item.get("image", {}).get("aeskey", "") |
| 315 | if file_url and aes_key: |
| 316 | file_path = await self._download_and_save_media(file_url, aes_key, "image") |
| 317 | if file_path: |
| 318 | filename = os.path.basename(file_path) |
| 319 | content_parts.append(f"[image: {filename}]") |
| 320 | media_paths.append(file_path) |
| 321 | else: |
| 322 | content_parts.append(MSG_TYPE_MAP.get(item_type, f"[{item_type}]")) |
| 323 | |
| 324 | else: |
| 325 | content_parts.append(MSG_TYPE_MAP.get(msg_type, f"[{msg_type}]")) |
| 326 | |
| 327 | content = "\n".join(content_parts) if content_parts else "" |
| 328 | |
| 329 | if not content: |
| 330 | return |
| 331 | |
| 332 | # Store frame for this chat to enable replies |
| 333 | self._chat_frames[chat_id] = frame |
| 334 | |
| 335 | # Forward to message bus |
| 336 | await self._handle_message( |
| 337 | sender_id=sender_id, |
| 338 | chat_id=chat_id, |
| 339 | content=content, |
| 340 | media=media_paths or None, |
| 341 | metadata={ |
| 342 | "message_id": msg_id, |
| 343 | "msg_type": msg_type, |
| 344 | "chat_type": chat_type, |
| 345 | } |
| 346 | ) |
| 347 | |
| 348 | except Exception as e: |
| 349 | logger.error("Error processing WeCom message: {}", e) |
| 350 | |
| 351 | async def _download_and_save_media( |
| 352 | self, |
| 353 | file_url: str, |
| 354 | aes_key: str, |
| 355 | media_type: str, |
| 356 | filename: str | None = None, |
| 357 | ) -> str | None: |
| 358 | """ |
| 359 | Download and decrypt media from WeCom. |
| 360 | |
| 361 | Returns: |
| 362 | file_path or None if download failed |
| 363 | """ |
| 364 | try: |
| 365 | data, fname = await self._client.download_file(file_url, aes_key) |
| 366 | |
| 367 | if not data: |
| 368 | logger.warning("Failed to download media from WeCom") |
| 369 | return None |
| 370 | |
| 371 | if len(data) > WECOM_UPLOAD_MAX_BYTES: |
| 372 | logger.warning( |
| 373 | "WeCom inbound media too large: {} bytes (max {})", |
| 374 | len(data), |
| 375 | WECOM_UPLOAD_MAX_BYTES, |
| 376 | ) |
| 377 | return None |
| 378 | |
| 379 | media_dir = get_media_dir("wecom") |
| 380 | if not filename: |
| 381 | filename = fname or f"{media_type}_{hash(file_url) % 100000}" |
| 382 | filename = _sanitize_filename(filename) |
| 383 | |
| 384 | file_path = media_dir / filename |
| 385 | await asyncio.to_thread(file_path.write_bytes, data) |
| 386 | logger.debug("Downloaded {} to {}", media_type, file_path) |
| 387 | return str(file_path) |
| 388 | |
| 389 | except Exception as e: |
| 390 | logger.error("Error downloading media: {}", e) |
| 391 | return None |
| 392 | |
| 393 | async def _upload_media_ws( |
| 394 | self, client: Any, file_path: str, |
| 395 | ) -> "tuple[str, str] | tuple[None, None]": |
| 396 | """Upload a local file to WeCom via WebSocket 3-step protocol (base64). |
| 397 | |
| 398 | Uses the WeCom WebSocket upload commands directly via |
| 399 | ``client._ws_manager.send_reply()``: |
| 400 | |
| 401 | ``aibot_upload_media_init`` → upload_id |
| 402 | ``aibot_upload_media_chunk`` × N (≤512 KB raw per chunk, base64) |
| 403 | ``aibot_upload_media_finish`` → media_id |
| 404 | |
| 405 | Returns (media_id, media_type) on success, (None, None) on failure. |
| 406 | """ |
| 407 | from wecom_aibot_sdk.utils import generate_req_id as _gen_req_id |
| 408 | |
| 409 | try: |
| 410 | fname = os.path.basename(file_path) |
| 411 | media_type = _guess_wecom_media_type(fname) |
| 412 | |
| 413 | # Read file size and data in a thread to avoid blocking the event loop |
| 414 | def _read_file(): |
| 415 | file_size = os.path.getsize(file_path) |
| 416 | if file_size > WECOM_UPLOAD_MAX_BYTES: |
| 417 | raise ValueError( |
| 418 | f"File too large: {file_size} bytes (max {WECOM_UPLOAD_MAX_BYTES})" |
| 419 | ) |
| 420 | with open(file_path, "rb") as f: |
| 421 | return file_size, f.read() |
| 422 | |
| 423 | file_size, data = await asyncio.to_thread(_read_file) |
| 424 | # MD5 is used for file integrity only, not cryptographic security |
| 425 | md5_hash = hashlib.md5(data).hexdigest() |
| 426 | |
| 427 | CHUNK_SIZE = 512 * 1024 # 512 KB raw (before base64) |
| 428 | mv = memoryview(data) |
| 429 | chunk_list = [bytes(mv[i : i + CHUNK_SIZE]) for i in range(0, file_size, CHUNK_SIZE)] |
| 430 | n_chunks = len(chunk_list) |
| 431 | del mv, data |
| 432 | |
| 433 | # Step 1: init |
| 434 | req_id = _gen_req_id("upload_init") |
| 435 | resp = await client._ws_manager.send_reply(req_id, { |
| 436 | "type": media_type, |
| 437 | "filename": fname, |
| 438 | "total_size": file_size, |
| 439 | "total_chunks": n_chunks, |
| 440 | "md5": md5_hash, |
| 441 | }, "aibot_upload_media_init") |
| 442 | if resp.errcode != 0: |
| 443 | logger.warning("WeCom upload init failed ({}): {}", resp.errcode, resp.errmsg) |
| 444 | return None, None |
| 445 | upload_id = resp.body.get("upload_id") if resp.body else None |
| 446 | if not upload_id: |
| 447 | logger.warning("WeCom upload init: no upload_id in response") |
| 448 | return None, None |
| 449 | |
| 450 | # Step 2: send chunks |
| 451 | for i, chunk in enumerate(chunk_list): |
| 452 | req_id = _gen_req_id("upload_chunk") |
| 453 | resp = await client._ws_manager.send_reply(req_id, { |
| 454 | "upload_id": upload_id, |
| 455 | "chunk_index": i, |
| 456 | "base64_data": base64.b64encode(chunk).decode(), |
| 457 | }, "aibot_upload_media_chunk") |
| 458 | if resp.errcode != 0: |
| 459 | logger.warning("WeCom upload chunk {} failed ({}): {}", i, resp.errcode, resp.errmsg) |
| 460 | return None, None |
| 461 | |
| 462 | # Step 3: finish |
| 463 | req_id = _gen_req_id("upload_finish") |
| 464 | resp = await client._ws_manager.send_reply(req_id, { |
| 465 | "upload_id": upload_id, |
| 466 | }, "aibot_upload_media_finish") |
| 467 | if resp.errcode != 0: |
| 468 | logger.warning("WeCom upload finish failed ({}): {}", resp.errcode, resp.errmsg) |
| 469 | return None, None |
| 470 | |
| 471 | media_id = resp.body.get("media_id") if resp.body else None |
| 472 | if not media_id: |
| 473 | logger.warning("WeCom upload finish: no media_id in response body={}", resp.body) |
| 474 | return None, None |
| 475 | |
| 476 | suffix = "..." if len(media_id) > 16 else "" |
| 477 | logger.debug("WeCom uploaded {} ({}) → media_id={}", fname, media_type, media_id[:16] + suffix) |
| 478 | return media_id, media_type |
| 479 | |
| 480 | except ValueError as e: |
| 481 | logger.warning("WeCom upload skipped for {}: {}", file_path, e) |
| 482 | return None, None |
| 483 | except Exception as e: |
| 484 | logger.error("WeCom _upload_media_ws error for {}: {}", file_path, e) |
| 485 | return None, None |
| 486 | |
| 487 | async def send(self, msg: OutboundMessage) -> None: |
| 488 | """Send a message through WeCom.""" |
| 489 | if not self._client: |
| 490 | logger.warning("WeCom client not initialized") |
| 491 | return |
| 492 | |
| 493 | try: |
| 494 | content = (msg.content or "").strip() |
| 495 | is_progress = bool(msg.metadata.get("_progress")) |
| 496 | |
| 497 | # Get the stored frame for this chat |
| 498 | frame = self._chat_frames.get(msg.chat_id) |
| 499 | |
| 500 | # Send media files via WebSocket upload |
| 501 | for file_path in msg.media or []: |
| 502 | if not os.path.isfile(file_path): |
| 503 | logger.warning("WeCom media file not found: {}", file_path) |
| 504 | continue |
| 505 | media_id, media_type = await self._upload_media_ws(self._client, file_path) |
| 506 | if media_id: |
| 507 | if frame: |
| 508 | await self._client.reply(frame, { |
| 509 | "msgtype": media_type, |
| 510 | media_type: {"media_id": media_id}, |
| 511 | }) |
| 512 | else: |
| 513 | await self._client.send_message(msg.chat_id, { |
| 514 | "msgtype": media_type, |
| 515 | media_type: {"media_id": media_id}, |
| 516 | }) |
| 517 | logger.debug("WeCom sent {} → {}", media_type, msg.chat_id) |
| 518 | else: |
| 519 | content += f"\n[file upload failed: {os.path.basename(file_path)}]" |
| 520 | |
| 521 | if not content: |
| 522 | return |
| 523 | |
| 524 | if frame: |
| 525 | # Both progress and final messages must use reply_stream (cmd="aibot_respond_msg"). |
| 526 | # The plain reply() uses cmd="reply" which does not support "text" msgtype |
| 527 | # and causes errcode=40008 from WeCom API. |
| 528 | stream_id = self._generate_req_id("stream") |
| 529 | await self._client.reply_stream( |
| 530 | frame, |
| 531 | stream_id, |
| 532 | content, |
| 533 | finish=not is_progress, |
| 534 | ) |
| 535 | logger.debug( |
| 536 | "WeCom {} sent to {}", |
| 537 | "progress" if is_progress else "message", |
| 538 | msg.chat_id, |
| 539 | ) |
| 540 | else: |
| 541 | # No frame (e.g. cron push): proactive send only supports markdown |
| 542 | await self._client.send_message(msg.chat_id, { |
| 543 | "msgtype": "markdown", |
| 544 | "markdown": {"content": content}, |
| 545 | }) |
| 546 | logger.info("WeCom proactive send to {}", msg.chat_id) |
| 547 | |
| 548 | except Exception: |
| 549 | logger.exception("Error sending WeCom message to chat_id={}", msg.chat_id) |
| 550 |