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

from __future__ import annotations

import asyncio
import json
import logging
import os
from typing import TYPE_CHECKING

from fastapi import APIRouter, Header, Request
from fastapi.responses import JSONResponse

from .crypto import DecryptionError
from .dedupe import get_pay_deduplicator
from .types import KNOWN_EVENT_TYPES, CallbackPayload

if TYPE_CHECKING:
    from . import WechatPayNotifyPlugin

logger = logging.getLogger(__name__)

MAX_BODY_BYTES = 64 * 1024
BODY_TIMEOUT_S = 3.0

router = APIRouter(prefix="/webhook/wechatpay", tags=["wechatpay-notify"])

_plugin: WechatPayNotifyPlugin | None = None


def set_plugin(plugin: WechatPayNotifyPlugin) -> None:
    global _plugin
    _plugin = plugin


def clear_plugin() -> None:
    global _plugin
    _plugin = None


@router.get("/callback")
@router.head("/callback")
async def wechatpay_callback_health():
    plugin = _plugin
    if plugin is None or plugin._account is None or plugin._crypto is None:
        return JSONResponse(
            content={"code": "FAIL", "message": "服务未就绪"},
            status_code=503,
        )
    return JSONResponse(
        content={"code": "SUCCESS", "message": "ok"},
        status_code=200,
    )


@router.post("/callback")
async def wechatpay_callback(
    request: Request,
    wechatpay_signature: str = Header("", alias="Wechatpay-Signature"),
    wechatpay_serial: str = Header("", alias="Wechatpay-Serial"),
    wechatpay_timestamp: str = Header("", alias="Wechatpay-Timestamp"),
    wechatpay_nonce: str = Header("", alias="Wechatpay-Nonce"),
    wechatpay_signature_type: str = Header("", alias="Wechatpay-Signature-Type"),
):
    plugin = _plugin
    if plugin is None or plugin._account is None or plugin._crypto is None:
        return JSONResponse(
            content={"code": "FAIL", "message": "服务未就绪"},
            status_code=503,
        )

    account = plugin._account
    crypto = plugin._crypto
    cert_mgr = plugin._cert_mgr

    if wechatpay_signature_type and wechatpay_signature_type != "WECHATPAY2-SHA256-RSA2048":
        logger.warning("WechatPay callback: unsupported signature type=%s", wechatpay_signature_type)
        return JSONResponse(
            content={"code": "FAIL", "message": "不支持的签名类型"},
            status_code=400,
        )

    try:
        body = await asyncio.wait_for(request.body(), timeout=BODY_TIMEOUT_S)
    except TimeoutError:
        logger.warning("WechatPay callback: body read timeout")
        return JSONResponse(
            content={"code": "FAIL", "message": "读取超时"},
            status_code=408,
        )

    if len(body) > MAX_BODY_BYTES:
        logger.warning("WechatPay callback: body too large (%d bytes)", len(body))
        return JSONResponse(
            content={"code": "FAIL", "message": "请求体过大"},
            status_code=413,
        )

    if account.cert_mode == "public_key":
        if wechatpay_serial != account.public_key_id:
            logger.warning("WechatPay callback: public key id mismatch, serial=%s", wechatpay_serial)
            return JSONResponse(
                content={"code": "FAIL", "message": "公钥未找到"},
                status_code=500,
            )
        sig_valid = crypto.verify_signature_with_raw_key(
            body=body,
            signature=wechatpay_signature,
            timestamp=wechatpay_timestamp,
            nonce=wechatpay_nonce,
            public_key_pem=account.public_key_pem,
        )
    else:
        if not cert_mgr:
            return JSONResponse(
                content={"code": "FAIL", "message": "证书管理器未就绪"},
                status_code=503,
            )
        platform_cert_pem = await cert_mgr.get_public_key(wechatpay_serial)
        if not platform_cert_pem:
            logger.warning("WechatPay callback: platform cert not found for serial=%s", wechatpay_serial)
            asyncio.create_task(cert_mgr.refresh(), name=f"cert-refresh-{wechatpay_serial}")
            return JSONResponse(
                content={"code": "FAIL", "message": "平台证书暂不可用，请重试"},
                status_code=500,
            )
        sig_valid = crypto.verify_signature(
            body=body,
            signature=wechatpay_signature,
            timestamp=wechatpay_timestamp,
            nonce=wechatpay_nonce,
            serial_no=wechatpay_serial,
            platform_cert_pem=platform_cert_pem,
        )

    if not sig_valid:
        logger.warning("WechatPay callback: signature verification failed for serial=%s", wechatpay_serial)
        return JSONResponse(
            content={"code": "FAIL", "message": "签名验证失败"},
            status_code=401,
        )

    body_str = body.decode("utf-8")
    payload = CallbackPayload.from_json(body_str)
    if payload is None:
        logger.warning("WechatPay callback: invalid JSON body")
        return JSONResponse(
            content={"code": "FAIL", "message": "请求格式错误"},
            status_code=400,
        )

    if payload.event_type not in KNOWN_EVENT_TYPES:
        logger.info("WechatPay callback: unknown event_type=%s, acknowledged", payload.event_type)
        return JSONResponse(
            content={"code": "SUCCESS", "message": "已确认(未知事件类型)"},
            status_code=200,
        )

    dedup = get_pay_deduplicator()
    is_first = await dedup.try_claim(payload.notification_id)
    if not is_first:
        logger.debug("WechatPay callback: duplicate notification id=%s", payload.notification_id)
        return JSONResponse(
            content={"code": "SUCCESS", "message": "已确认(重复通知)"},
            status_code=200,
        )

    try:
        decrypted_str = crypto.decrypt_resource(
            ciphertext=payload.resource.ciphertext,
            nonce=payload.resource.nonce,
            associated_data=payload.resource.associated_data,
        )
        decrypted_data = json.loads(decrypted_str)
    except DecryptionError:
        logger.warning("WechatPay callback: decryption failed for notification id=%s", payload.notification_id)
        return JSONResponse(
            content={"code": "FAIL", "message": "解密失败"},
            status_code=500,
        )

    asyncio.create_task(
        _dispatch_event(
            mch_id=account.mch_id,
            event_type=payload.event_type,
            notification_id=payload.notification_id,
            resource_data=decrypted_data,
        ),
        name=f"wechatpay-dispatch-{payload.notification_id}",
    )

    return JSONResponse(
        content={"code": "SUCCESS", "message": "成功"},
        status_code=200,
    )


async def _dispatch_event(
    mch_id: str,
    event_type: str,
    notification_id: str,
    resource_data: dict,
) -> None:
    try:
        arq_pool = _get_arq_pool()
        if arq_pool:
            await arq_pool.enqueue_job(
                "process_wechatpay_notify",
                mch_id=mch_id,
                event_type=event_type,
                notification_id=notification_id,
                resource_data=resource_data,
            )
        else:
            from .event_router import handle_pay_event

            await handle_pay_event(mch_id, event_type, notification_id, resource_data)
    except Exception:
        logger.exception("WechatPay dispatch error for id=%s", notification_id)


_ARQ_POOL = None


def _get_arq_pool():
    global _ARQ_POOL
    if _ARQ_POOL is not None:
        return _ARQ_POOL

    redis_url = os.environ.get("YUXI_REDIS_URL", os.environ.get("REDIS_URL", ""))
    if not redis_url:
        return None

    try:
        import arq
        from arq.connections import RedisSettings

        settings = RedisSettings.from_dsn(redis_url)
        _ARQ_POOL = arq.ArqRedis(settings)
        logger.info("WechatPay ARQ pool created: redis=%s", redis_url)
        return _ARQ_POOL
    except ImportError:
        logger.warning("WechatPay: arq not installed, notifications will process synchronously")
        return None
    except Exception:
        logger.exception("WechatPay: failed to create ARQ pool")
        return None