ForcePilot/backend/package/yuxi/channels/plugins/wecom/signature.py

"""企业微信 Webhook 签名校验与消息解密模块。

提供以下能力：

- ``verify_webhook_signature``: 校验 Webhook 签名（SHA1 + 恒定时间比对）。
- ``is_timestamp_valid``: 校验时间戳是否在允许的时钟偏移窗口内（防重放）。
- ``decrypt_event``: AES-CBC 解密企业微信回调加密内容。
- ``extract_echostr``: 识别 URL 验证 echostr 请求。

企业微信回调加密算法（与飞书不同）：
- 签名：``SHA1(token + timestamp + nonce + encrypt_msg)``（企业微信用 SHA1）
- 密钥派生：EncodingAESKey（43 位）Base64 解码后补 ``=`` 得到 32 字节 AES key
- 加密结构：Base64(16 字节随机串 + 4 字节消息长度(大端) + 消息体 + CorpID)
- 填充：PKCS7，块大小 32 字节（企业微信自定义，非标准 16 字节）

依赖方向：仅依赖标准库、``cryptography``（可选）与 ``yuxi.channels.contract``
错误类型，不污染框架层。
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import struct
import time
import xml.etree.ElementTree as ET
from typing import Any

from yuxi.channels.contract.errors import DependencyError, ValidationError

try:
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

    _HAS_CRYPTOGRAPHY = True
except ImportError:
    _HAS_CRYPTOGRAPHY = False

# 企业微信 PKCS7 填充块大小（32 字节，非标准 AES 16 字节）
_WECOM_BLOCK_SIZE = 32


def verify_webhook_signature(
    token: str,
    timestamp: str,
    nonce: str,
    encrypt_msg: str,
    msg_signature: str,
) -> bool:
    """校验企业微信 Webhook 签名。

    签名算法：``SHA1(token + timestamp + nonce + encrypt_msg)``，
    使用 ``hmac.compare_digest`` 进行恒定时间比对以防止时序攻击。

    Args:
        token: 配置的回调签名 Token。
        timestamp: 时间戳字符串。
        nonce: 随机字符串。
        encrypt_msg: 加密消息体（Encrypt 字段值）。
        msg_signature: 请求查询参数中的签名（十六进制小写）。

    Returns:
        签名校验结果，``True`` 表示通过。**不抛异常**：输入类型异常或
        编码失败时返回 ``False``。
    """
    try:
        parts = sorted([token, timestamp, nonce, encrypt_msg])
        raw = "".join(parts).encode("utf-8")
        expected = hashlib.sha1(raw).hexdigest()
    except (TypeError, UnicodeEncodeError, AttributeError):
        return False
    return hmac.compare_digest(expected, msg_signature)


def is_timestamp_valid(timestamp: str, max_skew_seconds: int = 300) -> bool:
    """校验时间戳是否在允许的时钟偏移窗口内（防重放）。

    Args:
        timestamp: 时间戳字符串（秒级）。
        max_skew_seconds: 允许的最大时间偏差（秒），默认 300（5 分钟）。

    Returns:
        ``True`` 表示时间戳有效。解析失败或超出窗口返回 ``False``，
        **不抛异常**。
    """
    try:
        ts = int(timestamp)
    except (TypeError, ValueError):
        return False
    skew = abs(int(time.time()) - ts)
    return skew <= max_skew_seconds


def _decode_aes_key(encoding_aes_key: str) -> bytes:
    """将 43 位 EncodingAESKey Base64 解码为 32 字节 AES 密钥。

    企业微信 EncodingAESKey 为 43 位 Base64 字符串，需在末尾补 ``=``
    后 Base64 解码得到 32 字节密钥。
    """
    return base64.b64decode(encoding_aes_key + "=")


def decrypt_event(
    encrypt_payload: str,
    encoding_aes_key: str,
    corp_id: str,
) -> dict[str, Any]:
    """AES-CBC 解密企业微信回调加密内容。

    解密流程：
        1. EncodingAESKey（43 位）Base64 解码为 32 字节 AES key
        2. 加密内容 Base64 解码
        3. AES-CBC 解密（IV 取前 16 字节，PKCS7 块大小 32）
        4. 去除前 16 字节随机串 + 4 字节消息长度
        5. 提取消息体 XML 与 CorpID，校验 CorpID 一致
        6. XML 解析为 dict

    Args:
        encrypt_payload: Base64 编码的加密内容字符串。
        encoding_aes_key: 配置的 EncodingAESKey（43 位）。
        corp_id: 企业微信企业 ID，用于解密后校验。

    Returns:
        解密后的 XML dict。

    Raises:
        DependencyError: ``cryptography`` 依赖不可用。
        ValidationError: 解密失败（Base64 解码失败、CorpID 不匹配、
            XML 解析失败等），``field`` 为 ``encrypt_payload``。
    """
    if not _HAS_CRYPTOGRAPHY:
        raise DependencyError(
            dep="cryptography",
            cause=ImportError("cryptography package is required for event decryption"),
        )

    try:
        key = _decode_aes_key(encoding_aes_key)
        encrypted = base64.b64decode(encrypt_payload)
        iv = encrypted[:16]
        ciphertext = encrypted[16:]

        cipher = Cipher(
            algorithms.AES(key),
            modes.CBC(iv),
            backend=default_backend(),
        )
        decryptor = cipher.decryptor()
        padded = decryptor.update(ciphertext) + decryptor.finalize()

        # 企业微信使用 PKCS7 块大小 32（非标准 16），手动去填充
        pad_len = padded[-1]
        if pad_len < 1 or pad_len > _WECOM_BLOCK_SIZE:
            raise ValueError(f"invalid pad length: {pad_len}")
        content = padded[:-pad_len]

        # 去除 16 字节随机串 + 4 字节消息长度（大端）
        if len(content) < 20:
            raise ValueError("decrypted content too short")
        msg_len = struct.unpack(">I", content[16:20])[0]
        if 20 + msg_len + len(corp_id) != len(content):
            raise ValueError("message length mismatch")
        msg = content[20 : 20 + msg_len]
        from_corp_id = content[20 + msg_len :].decode("utf-8")

        if from_corp_id != corp_id:
            raise ValueError(f"CorpID mismatch: expected {corp_id}, got {from_corp_id}")

        xml_str = msg.decode("utf-8")
        return _xml_to_dict(xml_str)
    except Exception as exc:
        if isinstance(exc, ValidationError):
            raise
        raise ValidationError(
            field="encrypt_payload",
            message="decrypt_failed",
        ) from exc


def decrypt_echostr(
    echostr: str,
    encoding_aes_key: str,
    corp_id: str,
) -> str:
    """解密 URL 验证 echostr。

    企业微信注册回调 URL 时发送 GET 请求含 ``echostr`` 参数，需解密后
    原样返回完成握手。

    Args:
        echostr: 加密的 echostr 字符串。
        encoding_aes_key: 配置的 EncodingAESKey（43 位）。
        corp_id: 企业微信企业 ID。

    Returns:
        解密后的 echostr 字符串。

    Raises:
        DependencyError: ``cryptography`` 依赖不可用。
        ValidationError: 解密失败。
    """
    if not _HAS_CRYPTOGRAPHY:
        raise DependencyError(
            dep="cryptography",
            cause=ImportError("cryptography package is required for echostr decryption"),
        )

    try:
        key = _decode_aes_key(encoding_aes_key)
        encrypted = base64.b64decode(echostr)
        iv = encrypted[:16]
        ciphertext = encrypted[16:]

        cipher = Cipher(
            algorithms.AES(key),
            modes.CBC(iv),
            backend=default_backend(),
        )
        decryptor = cipher.decryptor()
        padded = decryptor.update(ciphertext) + decryptor.finalize()

        pad_len = padded[-1]
        if pad_len < 1 or pad_len > _WECOM_BLOCK_SIZE:
            raise ValueError(f"invalid pad length: {pad_len}")
        content = padded[:-pad_len]

        if len(content) < 20:
            raise ValueError("decrypted content too short")
        msg_len = struct.unpack(">I", content[16:20])[0]
        if 20 + msg_len + len(corp_id) != len(content):
            raise ValueError("message length mismatch")
        msg = content[20 : 20 + msg_len]
        from_corp_id = content[20 + msg_len :].decode("utf-8")

        if from_corp_id != corp_id:
            raise ValueError(f"CorpID mismatch: expected {corp_id}, got {from_corp_id}")

        return msg.decode("utf-8")
    except Exception as exc:
        raise ValidationError(
            field="echostr",
            message="decrypt_failed",
        ) from exc


def _xml_to_dict(xml_str: str) -> dict[str, Any]:
    """将企业微信 XML 消息解析为 dict。

    企业微信回调消息为 XML 格式，根元素为 ``<xml>``，子元素为消息字段。
    """
    try:
        root = ET.fromstring(xml_str)
    except ET.ParseError as exc:
        raise ValidationError(
            field="xml",
            message="parse_failed",
        ) from exc

    result: dict[str, Any] = {}
    for child in root:
        result[child.tag] = child.text or ""
    return result


def extract_echostr(payload: dict[str, Any]) -> str | None:
    """识别 URL 验证 echostr 请求。

    企业微信在配置回调 URL 时发送 GET 请求含 ``echostr`` 参数，本函数
    从查询参数中提取 echostr 值。

    Args:
        payload: Webhook 请求参数 dict（含 query params）。

    Returns:
        echostr 字符串；若不含则返回 ``None``。
    """
    echostr = payload.get("echostr")
    if isinstance(echostr, str) and echostr:
        return echostr
    return None