ForcePilot/backend/package/yuxi/channel/extensions/wecom/webhook.py

import asyncio
import logging
import secrets
import time
from datetime import datetime, UTC

from fastapi import APIRouter, Query, Request
from fastapi.responses import PlainTextResponse, Response

from yuxi.channel.extensions.wecom.__init__ import _wecom_plugin
from yuxi.channel.extensions.wecom.crypto import WeChatCrypto
from yuxi.channel.extensions.wecom.config import WeComConfig
from yuxi.channel.extensions.wecom.events import handle_event
from yuxi.channel.extensions.wecom.message import extract_content, parse_xml_to_message
from yuxi.channel.extensions.wecom.mentions import is_bot_mentioned
from yuxi.channel.message.models import MessageType, PeerInfo, UnifiedMessage
from yuxi.channel.routing.models import PeerKind
from yuxi.channel.runtime.manager import gateway

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/webhook/wecom", tags=["wecom"])

_config = WeComConfig()


def _build_crypto() -> WeChatCrypto | None:
    account = _config.resolve_account()
    if not account.is_configured():
        return None
    return WeChatCrypto(
        token=account.token,
        encoding_aes_key=account.encoding_aes_key,
        app_id=account.corp_id,
    )


def _build_outbound():
    from yuxi.channel.extensions.wecom.gateway import _get_gateway

    gw = _get_gateway()
    if gw is None:
        return None
    from yuxi.channel.extensions.wecom.outbound import WeComOutbound

    return WeComOutbound(gw)


def _build_passive_reply_xml(to_user: str, from_user: str, content: str, crypto: WeChatCrypto) -> str:
    reply_xml = f"""<xml>
<ToUserName><![CDATA[{to_user}]]></ToUserName>
<FromUserName><![CDATA[{from_user}]]></FromUserName>
<CreateTime>{int(time.time())}</CreateTime>
<MsgType><![CDATA[text]]></MsgType>
<Content><![CDATA[{content}]]></Content>
</xml>"""

    timestamp = str(int(time.time()))
    nonce = secrets.token_hex(8)
    encrypted = crypto.encrypt(reply_xml)
    signature = crypto.compute_signature(timestamp, nonce, encrypted)

    return f"""<xml>
<Encrypt><![CDATA[{encrypted}]]></Encrypt>
<MsgSignature><![CDATA[{signature}]]></MsgSignature>
<TimeStamp>{timestamp}</TimeStamp>
<Nonce><![CDATA[{nonce}]]></Nonce>
</xml>"""


@router.get("/callback")
async def wecom_url_verify(
    msg_signature: str = Query(...),
    timestamp: str = Query(...),
    nonce: str = Query(...),
    echostr: str = Query(...),
):
    crypto = _build_crypto()
    if crypto is None:
        return PlainTextResponse("", status_code=503)

    if not crypto.verify_signature(msg_signature, timestamp, nonce, echostr):
        logger.warning("WeCom URL verification: invalid signature")
        return PlainTextResponse("", status_code=403)

    try:
        decrypted, _corp_id = crypto.decrypt(echostr)
        return PlainTextResponse(decrypted)
    except Exception:
        logger.exception("WeCom echostr decryption failed")
        return PlainTextResponse("", status_code=403)


@router.post("/callback")
async def wecom_message_callback(
    request: Request,
    msg_signature: str = Query(default=""),
    timestamp: str = Query(default=""),
    nonce: str = Query(default=""),
):
    crypto = _build_crypto()
    if crypto is None:
        return PlainTextResponse("", status_code=503)

    xml_body = await request.body()
    xml_text = xml_body.decode("utf-8")

    import xml.etree.ElementTree as ET

    root = ET.fromstring(xml_text)
    encrypt_el = root.find("Encrypt")
    encrypt = encrypt_el.text or "" if encrypt_el is not None else ""

    if not crypto.verify_signature(msg_signature, timestamp, nonce, encrypt):
        logger.warning("WeCom callback: invalid msg_signature")
        return PlainTextResponse("", status_code=403)

    try:
        decrypted_xml, _corp_id = crypto.decrypt(encrypt)
    except Exception:
        logger.exception("WeCom message decryption failed")
        return PlainTextResponse("", status_code=403)

    try:
        raw_msg = parse_xml_to_message(decrypted_xml)
    except Exception:
        logger.exception("WeCom XML parsing failed")
        return PlainTextResponse("", status_code=400)

    if raw_msg.from_user == raw_msg.to_user:
        return PlainTextResponse("success")

    account = _config.resolve_account()
    plugin = _wecom_plugin

    if raw_msg.msg_type == "event":
        reply_content = await handle_event(raw_msg)
        if reply_content:
            passive_xml = _build_passive_reply_xml(
                raw_msg.from_user,
                raw_msg.to_user,
                reply_content,
                crypto,
            )
            return Response(content=passive_xml, media_type="application/xml")
        return PlainTextResponse("success")

    if not raw_msg.msg_id or plugin.is_duplicate(raw_msg.msg_id):
        return PlainTextResponse("success")

    content = extract_content(raw_msg)
    if not content:
        return PlainTextResponse("success")

    if not raw_msg.is_group_chat:
        policy = account.dm_policy

        if policy == "disabled":
            return PlainTextResponse("success")

        if policy == "pairing":
            if not plugin.check_allowlist(raw_msg.from_user, "direct"):
                if content.strip().startswith("配对 "):
                    code_input = content.strip()[3:].strip()
                    if await plugin.verify_code(raw_msg.from_user, code_input):
                        plugin.add_to_allowlist(raw_msg.from_user)
                        passive_xml = _build_passive_reply_xml(
                            raw_msg.from_user,
                            raw_msg.to_user,
                            "配对成功！现在可以开始对话了。",
                            crypto,
                        )
                        return Response(content=passive_xml, media_type="application/xml")
                    else:
                        passive_xml = _build_passive_reply_xml(
                            raw_msg.from_user,
                            raw_msg.to_user,
                            "配对码无效或已过期。",
                            crypto,
                        )
                        return Response(content=passive_xml, media_type="application/xml")
                else:
                    code = await plugin.generate_code(raw_msg.from_user)
                    if code:
                        passive_xml = _build_passive_reply_xml(
                            raw_msg.from_user,
                            raw_msg.to_user,
                            f"首次对话需要验证身份：\n\n配对 {code}\n（有效期 10 分钟）",
                            crypto,
                        )
                        return Response(content=passive_xml, media_type="application/xml")
                return PlainTextResponse("success")

        if policy == "allowlist":
            if not plugin.check_allowlist(raw_msg.from_user, "direct"):
                return PlainTextResponse("success")

    else:
        group_policy = account.group_policy

        if group_policy == "disabled":
            return PlainTextResponse("success")

        if group_policy == "activate_on_mention":
            if not is_bot_mentioned(content):
                return PlainTextResponse("success")

    processor = gateway._processor
    if processor is None:
        return PlainTextResponse("success")

    peer_id = raw_msg.peer_id
    peer_kind = PeerKind.GROUP if raw_msg.is_group_chat else PeerKind.DIRECT

    unified = UnifiedMessage(
        msg_id=raw_msg.msg_id,
        channel_type="wecom",
        account_id="default",
        content=content,
        message_type=MessageType.TEXT,
        sender=PeerInfo(
            id=raw_msg.from_user,
            kind=peer_kind,
            display_name=raw_msg.from_user,
        ),
        recipient=PeerInfo(
            id=peer_id,
            kind=peer_kind,
            display_name=raw_msg.chat_id if raw_msg.is_group_chat else raw_msg.from_user,
        ),
        timestamp=datetime.fromtimestamp(raw_msg.create_time, tz=UTC) if raw_msg.create_time else None,
        raw_payload=raw_msg.raw_xml,
        body_for_agent=content,
        metadata={
            "FromUserName": raw_msg.from_user,
            "ToUserName": raw_msg.to_user,
            "MsgType": raw_msg.msg_type,
            "ChatType": raw_msg.chat_type,
            "ChatId": raw_msg.chat_id,
            "AgentID": raw_msg.agent_id,
            "MediaId": raw_msg.media_id,
            "PicUrl": raw_msg.pic_url,
        },
    )

    asyncio.create_task(
        _dispatch_to_agent(processor, unified),
        name=f"wecom-dispatch-{raw_msg.from_user[:20]}",
    )

    return PlainTextResponse("success")


async def _dispatch_to_agent(processor, msg: UnifiedMessage) -> None:
    try:
        await asyncio.wait_for(processor.process(msg), timeout=120.0)
    except TimeoutError:
        logger.error("Agent response timeout for wecom user %s", msg.sender.id)
    except Exception:
        logger.exception("Failed to process WeCom message for user %s", msg.sender.id)