553 lines
21 KiB
Python
553 lines
21 KiB
Python
"""WeCom (Enterprise WeChat) channel implementation using wecom_aibot_sdk."""
|
||
|
||
import asyncio
|
||
import base64
|
||
import hashlib
|
||
import importlib.util
|
||
import os
|
||
import re
|
||
from collections import OrderedDict
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from loguru import logger
|
||
|
||
from nanobot.bus.events import OutboundMessage
|
||
from nanobot.bus.queue import MessageBus
|
||
from nanobot.channels.base import BaseChannel
|
||
from nanobot.config.paths import get_media_dir, get_workspace_path
|
||
from nanobot.config.schema import Base
|
||
from pydantic import Field
|
||
|
||
WECOM_AVAILABLE = importlib.util.find_spec("wecom_aibot_sdk") is not None
|
||
|
||
# Upload safety limits (matching QQ channel defaults)
|
||
WECOM_UPLOAD_MAX_BYTES = 1024 * 1024 * 200 # 200MB
|
||
|
||
# Replace unsafe characters with "_", keep Chinese and common safe punctuation.
|
||
_SAFE_NAME_RE = re.compile(r"[^\w.\-()\[\]（）【】\u4e00-\u9fff]+", re.UNICODE)
|
||
|
||
|
||
def _sanitize_filename(name: str) -> str:
|
||
"""Sanitize filename to avoid traversal and problematic chars."""
|
||
name = (name or "").strip()
|
||
name = Path(name).name
|
||
name = _SAFE_NAME_RE.sub("_", name).strip("._ ")
|
||
return name
|
||
|
||
|
||
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}
|
||
_VIDEO_EXTS = {".mp4", ".avi", ".mov"}
|
||
_AUDIO_EXTS = {".amr", ".mp3", ".wav", ".ogg"}
|
||
|
||
|
||
def _guess_wecom_media_type(filename: str) -> str:
|
||
"""Classify file extension as WeCom media_type string."""
|
||
ext = Path(filename).suffix.lower()
|
||
if ext in _IMAGE_EXTS:
|
||
return "image"
|
||
if ext in _VIDEO_EXTS:
|
||
return "video"
|
||
if ext in _AUDIO_EXTS:
|
||
return "voice"
|
||
return "file"
|
||
|
||
class WecomConfig(Base):
|
||
"""WeCom (Enterprise WeChat) AI Bot channel configuration."""
|
||
|
||
enabled: bool = False
|
||
bot_id: str = ""
|
||
secret: str = ""
|
||
allow_from: list[str] = Field(default_factory=list)
|
||
welcome_message: str = ""
|
||
|
||
|
||
# Message type display mapping
|
||
MSG_TYPE_MAP = {
|
||
"image": "[image]",
|
||
"voice": "[voice]",
|
||
"file": "[file]",
|
||
"mixed": "[mixed content]",
|
||
}
|
||
|
||
|
||
class WecomChannel(BaseChannel):
|
||
"""
|
||
WeCom (Enterprise WeChat) channel using WebSocket long connection.
|
||
|
||
Uses WebSocket to receive events - no public IP or webhook required.
|
||
|
||
Requires:
|
||
- Bot ID and Secret from WeCom AI Bot platform
|
||
"""
|
||
|
||
name = "wecom"
|
||
display_name = "WeCom"
|
||
|
||
@classmethod
|
||
def default_config(cls) -> dict[str, Any]:
|
||
return WecomConfig().model_dump(by_alias=True)
|
||
|
||
def __init__(self, config: Any, bus: MessageBus):
|
||
if isinstance(config, dict):
|
||
config = WecomConfig.model_validate(config)
|
||
super().__init__(config, bus)
|
||
self.config: WecomConfig = config
|
||
self._client: Any = None
|
||
self._processed_message_ids: OrderedDict[str, None] = OrderedDict()
|
||
self._loop: asyncio.AbstractEventLoop | None = None
|
||
self._generate_req_id = None
|
||
# Store frame headers for each chat to enable replies
|
||
self._chat_frames: dict[str, Any] = {}
|
||
|
||
async def start(self) -> None:
|
||
"""Start the WeCom bot with WebSocket long connection."""
|
||
if not WECOM_AVAILABLE:
|
||
logger.error("WeCom SDK not installed. Run: pip install nanobot-ai[wecom]")
|
||
return
|
||
|
||
if not self.config.bot_id or not self.config.secret:
|
||
logger.error("WeCom bot_id and secret not configured")
|
||
return
|
||
|
||
from wecom_aibot_sdk import WSClient, generate_req_id
|
||
|
||
self._running = True
|
||
self._loop = asyncio.get_running_loop()
|
||
self._generate_req_id = generate_req_id
|
||
|
||
# Create WebSocket client
|
||
self._client = WSClient({
|
||
"bot_id": self.config.bot_id,
|
||
"secret": self.config.secret,
|
||
"reconnect_interval": 1000,
|
||
"max_reconnect_attempts": -1, # Infinite reconnect
|
||
"heartbeat_interval": 30000,
|
||
})
|
||
|
||
# Register event handlers
|
||
self._client.on("connected", self._on_connected)
|
||
self._client.on("authenticated", self._on_authenticated)
|
||
self._client.on("disconnected", self._on_disconnected)
|
||
self._client.on("error", self._on_error)
|
||
self._client.on("message.text", self._on_text_message)
|
||
self._client.on("message.image", self._on_image_message)
|
||
self._client.on("message.voice", self._on_voice_message)
|
||
self._client.on("message.file", self._on_file_message)
|
||
self._client.on("message.mixed", self._on_mixed_message)
|
||
self._client.on("event.enter_chat", self._on_enter_chat)
|
||
|
||
logger.info("WeCom bot starting with WebSocket long connection")
|
||
logger.info("No public IP required - using WebSocket to receive events")
|
||
|
||
# Connect
|
||
await self._client.connect_async()
|
||
|
||
# Keep running until stopped
|
||
while self._running:
|
||
await asyncio.sleep(1)
|
||
|
||
async def stop(self) -> None:
|
||
"""Stop the WeCom bot."""
|
||
self._running = False
|
||
if self._client:
|
||
await self._client.disconnect()
|
||
logger.info("WeCom bot stopped")
|
||
|
||
async def _on_connected(self, frame: Any) -> None:
|
||
"""Handle WebSocket connected event."""
|
||
logger.info("WeCom WebSocket connected")
|
||
|
||
async def _on_authenticated(self, frame: Any) -> None:
|
||
"""Handle authentication success event."""
|
||
logger.info("WeCom authenticated successfully")
|
||
|
||
async def _on_disconnected(self, frame: Any) -> None:
|
||
"""Handle WebSocket disconnected event."""
|
||
reason = frame.body if hasattr(frame, 'body') else str(frame)
|
||
logger.warning("WeCom WebSocket disconnected: {}", reason)
|
||
|
||
async def _on_error(self, frame: Any) -> None:
|
||
"""Handle error event."""
|
||
logger.error("WeCom error: {}", frame)
|
||
|
||
async def _on_text_message(self, frame: Any) -> None:
|
||
"""Handle text message."""
|
||
await self._process_message(frame, "text")
|
||
|
||
async def _on_image_message(self, frame: Any) -> None:
|
||
"""Handle image message."""
|
||
await self._process_message(frame, "image")
|
||
|
||
async def _on_voice_message(self, frame: Any) -> None:
|
||
"""Handle voice message."""
|
||
await self._process_message(frame, "voice")
|
||
|
||
async def _on_file_message(self, frame: Any) -> None:
|
||
"""Handle file message."""
|
||
await self._process_message(frame, "file")
|
||
|
||
async def _on_mixed_message(self, frame: Any) -> None:
|
||
"""Handle mixed content message."""
|
||
await self._process_message(frame, "mixed")
|
||
|
||
async def _on_enter_chat(self, frame: Any) -> None:
|
||
"""Handle enter_chat event (user opens chat with bot)."""
|
||
try:
|
||
# Extract body from WsFrame dataclass or dict
|
||
if hasattr(frame, 'body'):
|
||
body = frame.body or {}
|
||
elif isinstance(frame, dict):
|
||
body = frame.get("body", frame)
|
||
else:
|
||
body = {}
|
||
|
||
chat_id = body.get("chatid", "") if isinstance(body, dict) else ""
|
||
|
||
if chat_id and self.config.welcome_message:
|
||
await self._client.reply_welcome(frame, {
|
||
"msgtype": "text",
|
||
"text": {"content": self.config.welcome_message},
|
||
})
|
||
except Exception as e:
|
||
logger.error("Error handling enter_chat: {}", e)
|
||
|
||
async def _process_message(self, frame: Any, msg_type: str) -> None:
|
||
"""Process incoming message and forward to bus."""
|
||
try:
|
||
# Extract body from WsFrame dataclass or dict
|
||
if hasattr(frame, 'body'):
|
||
body = frame.body or {}
|
||
elif isinstance(frame, dict):
|
||
body = frame.get("body", frame)
|
||
else:
|
||
body = {}
|
||
|
||
# Ensure body is a dict
|
||
if not isinstance(body, dict):
|
||
logger.warning("Invalid body type: {}", type(body))
|
||
return
|
||
|
||
# Extract message info
|
||
msg_id = body.get("msgid", "")
|
||
if not msg_id:
|
||
msg_id = f"{body.get('chatid', '')}_{body.get('sendertime', '')}"
|
||
|
||
# Deduplication check
|
||
if msg_id in self._processed_message_ids:
|
||
return
|
||
self._processed_message_ids[msg_id] = None
|
||
|
||
# Trim cache
|
||
while len(self._processed_message_ids) > 1000:
|
||
self._processed_message_ids.popitem(last=False)
|
||
|
||
# Extract sender info from "from" field (SDK format)
|
||
from_info = body.get("from", {})
|
||
sender_id = from_info.get("userid", "unknown") if isinstance(from_info, dict) else "unknown"
|
||
|
||
# For single chat, chatid is the sender's userid
|
||
# For group chat, chatid is provided in body
|
||
chat_type = body.get("chattype", "single")
|
||
chat_id = body.get("chatid", sender_id)
|
||
|
||
content_parts = []
|
||
media_paths: list[str] = []
|
||
|
||
if msg_type == "text":
|
||
text = body.get("text", {}).get("content", "")
|
||
if text:
|
||
content_parts.append(text)
|
||
|
||
elif msg_type == "image":
|
||
image_info = body.get("image", {})
|
||
file_url = image_info.get("url", "")
|
||
aes_key = image_info.get("aeskey", "")
|
||
|
||
if file_url and aes_key:
|
||
file_path = await self._download_and_save_media(file_url, aes_key, "image")
|
||
if file_path:
|
||
filename = os.path.basename(file_path)
|
||
content_parts.append(f"[image: {filename}]")
|
||
media_paths.append(file_path)
|
||
else:
|
||
content_parts.append("[image: download failed]")
|
||
else:
|
||
content_parts.append("[image: download failed]")
|
||
|
||
elif msg_type == "voice":
|
||
voice_info = body.get("voice", {})
|
||
# Voice message already contains transcribed content from WeCom
|
||
voice_content = voice_info.get("content", "")
|
||
if voice_content:
|
||
content_parts.append(f"[voice] {voice_content}")
|
||
else:
|
||
content_parts.append("[voice]")
|
||
|
||
elif msg_type == "file":
|
||
file_info = body.get("file", {})
|
||
file_url = file_info.get("url", "")
|
||
aes_key = file_info.get("aeskey", "")
|
||
file_name = file_info.get("name", "unknown")
|
||
|
||
if file_url and aes_key:
|
||
file_path = await self._download_and_save_media(file_url, aes_key, "file", file_name)
|
||
if file_path:
|
||
content_parts.append(f"[file: {file_name}]")
|
||
media_paths.append(file_path)
|
||
else:
|
||
content_parts.append(f"[file: {file_name}: download failed]")
|
||
else:
|
||
content_parts.append(f"[file: {file_name}: download failed]")
|
||
|
||
elif msg_type == "mixed":
|
||
# Mixed content contains multiple message items
|
||
msg_items = body.get("mixed", {}).get("msg_item", [])
|
||
for item in msg_items:
|
||
item_type = item.get("msgtype", "")
|
||
if item_type == "text":
|
||
text = item.get("text", {}).get("content", "")
|
||
if text:
|
||
content_parts.append(text)
|
||
elif item_type == "image":
|
||
file_url = item.get("image", {}).get("url", "")
|
||
aes_key = item.get("image", {}).get("aeskey", "")
|
||
if file_url and aes_key:
|
||
file_path = await self._download_and_save_media(file_url, aes_key, "image")
|
||
if file_path:
|
||
filename = os.path.basename(file_path)
|
||
content_parts.append(f"[image: {filename}]")
|
||
media_paths.append(file_path)
|
||
else:
|
||
content_parts.append(MSG_TYPE_MAP.get(item_type, f"[{item_type}]"))
|
||
|
||
else:
|
||
content_parts.append(MSG_TYPE_MAP.get(msg_type, f"[{msg_type}]"))
|
||
|
||
content = "\n".join(content_parts) if content_parts else ""
|
||
|
||
if not content:
|
||
return
|
||
|
||
# Store frame for this chat to enable replies
|
||
self._chat_frames[chat_id] = frame
|
||
|
||
# Forward to message bus
|
||
await self._handle_message(
|
||
sender_id=sender_id,
|
||
chat_id=chat_id,
|
||
content=content,
|
||
media=media_paths or None,
|
||
metadata={
|
||
"message_id": msg_id,
|
||
"msg_type": msg_type,
|
||
"chat_type": chat_type,
|
||
}
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error("Error processing WeCom message: {}", e)
|
||
|
||
async def _download_and_save_media(
|
||
self,
|
||
file_url: str,
|
||
aes_key: str,
|
||
media_type: str,
|
||
filename: str | None = None,
|
||
) -> str | None:
|
||
"""
|
||
Download and decrypt media from WeCom.
|
||
|
||
Returns:
|
||
file_path or None if download failed
|
||
"""
|
||
try:
|
||
data, fname = await self._client.download_file(file_url, aes_key)
|
||
|
||
if not data:
|
||
logger.warning("Failed to download media from WeCom")
|
||
return None
|
||
|
||
if len(data) > WECOM_UPLOAD_MAX_BYTES:
|
||
logger.warning(
|
||
"WeCom inbound media too large: {} bytes (max {})",
|
||
len(data),
|
||
WECOM_UPLOAD_MAX_BYTES,
|
||
)
|
||
return None
|
||
|
||
media_dir = get_media_dir("wecom")
|
||
if not filename:
|
||
filename = fname or f"{media_type}_{hash(file_url) % 100000}"
|
||
filename = _sanitize_filename(filename)
|
||
|
||
file_path = media_dir / filename
|
||
await asyncio.to_thread(file_path.write_bytes, data)
|
||
logger.debug("Downloaded {} to {}", media_type, file_path)
|
||
return str(file_path)
|
||
|
||
except Exception as e:
|
||
logger.error("Error downloading media: {}", e)
|
||
return None
|
||
|
||
async def _upload_media_ws(
|
||
self, client: Any, file_path: str,
|
||
) -> "tuple[str, str] | tuple[None, None]":
|
||
"""Upload a local file to WeCom via WebSocket 3-step protocol (base64).
|
||
|
||
Uses the WeCom WebSocket upload commands directly via
|
||
``client._ws_manager.send_reply()``:
|
||
|
||
``aibot_upload_media_init`` → upload_id
|
||
``aibot_upload_media_chunk`` × N (≤512 KB raw per chunk, base64)
|
||
``aibot_upload_media_finish`` → media_id
|
||
|
||
Returns (media_id, media_type) on success, (None, None) on failure.
|
||
"""
|
||
from wecom_aibot_sdk.utils import generate_req_id as _gen_req_id
|
||
|
||
try:
|
||
fname = os.path.basename(file_path)
|
||
media_type = _guess_wecom_media_type(fname)
|
||
|
||
# Read file size and data in a thread to avoid blocking the event loop
|
||
def _read_file():
|
||
file_size = os.path.getsize(file_path)
|
||
if file_size > WECOM_UPLOAD_MAX_BYTES:
|
||
raise ValueError(
|
||
f"File too large: {file_size} bytes (max {WECOM_UPLOAD_MAX_BYTES})"
|
||
)
|
||
with open(file_path, "rb") as f:
|
||
return file_size, f.read()
|
||
|
||
file_size, data = await asyncio.to_thread(_read_file)
|
||
# MD5 is used for file integrity only, not cryptographic security
|
||
md5_hash = hashlib.md5(data).hexdigest()
|
||
|
||
CHUNK_SIZE = 512 * 1024 # 512 KB raw (before base64)
|
||
mv = memoryview(data)
|
||
chunk_list = [bytes(mv[i : i + CHUNK_SIZE]) for i in range(0, file_size, CHUNK_SIZE)]
|
||
n_chunks = len(chunk_list)
|
||
del mv, data
|
||
|
||
# Step 1: init
|
||
req_id = _gen_req_id("upload_init")
|
||
resp = await client._ws_manager.send_reply(req_id, {
|
||
"type": media_type,
|
||
"filename": fname,
|
||
"total_size": file_size,
|
||
"total_chunks": n_chunks,
|
||
"md5": md5_hash,
|
||
}, "aibot_upload_media_init")
|
||
if resp.errcode != 0:
|
||
logger.warning("WeCom upload init failed ({}): {}", resp.errcode, resp.errmsg)
|
||
return None, None
|
||
upload_id = resp.body.get("upload_id") if resp.body else None
|
||
if not upload_id:
|
||
logger.warning("WeCom upload init: no upload_id in response")
|
||
return None, None
|
||
|
||
# Step 2: send chunks
|
||
for i, chunk in enumerate(chunk_list):
|
||
req_id = _gen_req_id("upload_chunk")
|
||
resp = await client._ws_manager.send_reply(req_id, {
|
||
"upload_id": upload_id,
|
||
"chunk_index": i,
|
||
"base64_data": base64.b64encode(chunk).decode(),
|
||
}, "aibot_upload_media_chunk")
|
||
if resp.errcode != 0:
|
||
logger.warning("WeCom upload chunk {} failed ({}): {}", i, resp.errcode, resp.errmsg)
|
||
return None, None
|
||
|
||
# Step 3: finish
|
||
req_id = _gen_req_id("upload_finish")
|
||
resp = await client._ws_manager.send_reply(req_id, {
|
||
"upload_id": upload_id,
|
||
}, "aibot_upload_media_finish")
|
||
if resp.errcode != 0:
|
||
logger.warning("WeCom upload finish failed ({}): {}", resp.errcode, resp.errmsg)
|
||
return None, None
|
||
|
||
media_id = resp.body.get("media_id") if resp.body else None
|
||
if not media_id:
|
||
logger.warning("WeCom upload finish: no media_id in response body={}", resp.body)
|
||
return None, None
|
||
|
||
suffix = "..." if len(media_id) > 16 else ""
|
||
logger.debug("WeCom uploaded {} ({}) → media_id={}", fname, media_type, media_id[:16] + suffix)
|
||
return media_id, media_type
|
||
|
||
except ValueError as e:
|
||
logger.warning("WeCom upload skipped for {}: {}", file_path, e)
|
||
return None, None
|
||
except Exception as e:
|
||
logger.error("WeCom _upload_media_ws error for {}: {}", file_path, e)
|
||
return None, None
|
||
|
||
async def send(self, msg: OutboundMessage) -> None:
|
||
"""Send a message through WeCom."""
|
||
if not self._client:
|
||
logger.warning("WeCom client not initialized")
|
||
return
|
||
|
||
try:
|
||
content = (msg.content or "").strip()
|
||
is_progress = bool(msg.metadata.get("_progress"))
|
||
|
||
# Get the stored frame for this chat
|
||
frame = self._chat_frames.get(msg.chat_id)
|
||
|
||
# Send media files via WebSocket upload
|
||
for file_path in msg.media or []:
|
||
upload_path = file_path
|
||
if not os.path.isfile(upload_path) and not os.path.isabs(file_path):
|
||
upload_path = str(get_workspace_path() / file_path)
|
||
if not os.path.isfile(upload_path):
|
||
logger.warning("WeCom media file not found: {}", file_path)
|
||
continue
|
||
media_id, media_type = await self._upload_media_ws(self._client, upload_path)
|
||
if media_id:
|
||
if frame:
|
||
await self._client.reply(frame, {
|
||
"msgtype": media_type,
|
||
media_type: {"media_id": media_id},
|
||
})
|
||
else:
|
||
await self._client.send_message(msg.chat_id, {
|
||
"msgtype": media_type,
|
||
media_type: {"media_id": media_id},
|
||
})
|
||
logger.debug("WeCom sent {} → {}", media_type, msg.chat_id)
|
||
else:
|
||
content += f"\n[file upload failed: {os.path.basename(file_path)}]"
|
||
|
||
if not content:
|
||
return
|
||
|
||
if frame:
|
||
# Both progress and final messages must use reply_stream (cmd="aibot_respond_msg").
|
||
# The plain reply() uses cmd="reply" which does not support "text" msgtype
|
||
# and causes errcode=40008 from WeCom API.
|
||
stream_id = self._generate_req_id("stream")
|
||
await self._client.reply_stream(
|
||
frame,
|
||
stream_id,
|
||
content,
|
||
finish=not is_progress,
|
||
)
|
||
logger.debug(
|
||
"WeCom {} sent to {}",
|
||
"progress" if is_progress else "message",
|
||
msg.chat_id,
|
||
)
|
||
else:
|
||
# No frame (e.g. cron push): proactive send only supports markdown
|
||
await self._client.send_message(msg.chat_id, {
|
||
"msgtype": "markdown",
|
||
"markdown": {"content": content},
|
||
})
|
||
logger.info("WeCom proactive send to {}", msg.chat_id)
|
||
|
||
except Exception:
|
||
logger.exception("Error sending WeCom message to chat_id={}", msg.chat_id)
|