ForcePilot/backend/package/yuxi/channels/plugins/wechat_mp/crypto.py

"""微信公众号 AES-128-CBC 加解密与 SHA1 签名工具模块。

提供微信公众号消息加解密（明文/兼容/安全三种模式）与 Webhook 签名校验。

密钥派生（功能点清单 §2.1）：
- ``EncodingAESKey``（43 位 base64）+ ``"="`` → base64 解码 → 32 字节
- AES 密钥 = 前 16 字节（AES-128）
- IV = 前 16 字节（与密钥相同，CBC 模式）

明文结构（解密后）：
- 16 字节随机字符串
- 4 字节消息长度（大端 uint32）
- XML 消息内容
- AppID

仅依赖标准库 + cryptography 库 + contract.errors，不 import 框架层。
"""

from __future__ import annotations

import base64
import hashlib
import os
import struct
import xml.etree.ElementTree as ET

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7

from yuxi.channels.contract.errors import ValidationError

from ._constants import ENCODING_AES_KEY_LENGTH

# AES-128 块大小（比特），PKCS7 填充参数
_AES_BLOCK_BITS = 128

# 微信消息体明文各段长度
_RANDOM_BYTES_LEN = 16
_MSG_LEN_BYTES = 4


def derive_aes_key(encoding_aes_key: str) -> bytes:
    """从 EncodingAESKey 派生 16 字节 AES 密钥。

    流程：43 位 base64 + ``"="`` → base64 解码 → 32 字节 → 前 16 字节作为 AES 密钥。

    Args:
        encoding_aes_key: 43 位 EncodingAESKey。

    Returns:
        16 字节 AES 密钥。

    Raises:
        ValidationError: EncodingAESKey 长度不为 43 或 base64 解码失败。
    """
    if not encoding_aes_key or len(encoding_aes_key) != ENCODING_AES_KEY_LENGTH:
        raise ValidationError(
            field="encoding_aes_key",
            message=f"encoding_aes_key_must_be_{ENCODING_AES_KEY_LENGTH}_chars",
        )
    try:
        decoded = base64.b64decode(encoding_aes_key + "=")
    except Exception as exc:
        raise ValidationError(
            field="encoding_aes_key",
            message="encoding_aes_key_base64_decode_failed",
        ) from exc
    if len(decoded) < 16:
        raise ValidationError(
            field="encoding_aes_key",
            message="encoding_aes_key_decoded_too_short",
        )
    return decoded[:16]


def encrypt_message(plaintext_xml: str, encoding_aes_key: str, app_id: str) -> str:
    """AES-128-CBC 加密消息，返回 base64 编码密文。

    明文结构：16 字节随机串 + 4 字节消息长度（大端）+ XML + AppID，PKCS7 填充。

    Args:
        plaintext_xml: 待加密的 XML 字符串。
        encoding_aes_key: 43 位 EncodingAESKey。
        app_id: 公众号 AppID（用于校验）。

    Returns:
        base64 编码的密文字符串。
    """
    key = derive_aes_key(encoding_aes_key)
    iv = key  # IV 与密钥相同（微信协议约定）
    random_bytes = os.urandom(_RANDOM_BYTES_LEN)
    msg_bytes = plaintext_xml.encode("utf-8")
    appid_bytes = app_id.encode("utf-8")
    msg_len = struct.pack(">I", len(msg_bytes))
    plaintext = random_bytes + msg_len + msg_bytes + appid_bytes

    padder = PKCS7(_AES_BLOCK_BITS).padder()
    padded = padder.update(plaintext) + padder.finalize()
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
    encryptor = cipher.encryptor()
    ciphertext = encryptor.update(padded) + encryptor.finalize()
    return base64.b64encode(ciphertext).decode("ascii")


def decrypt_message(ciphertext_b64: str, encoding_aes_key: str, expected_app_id: str) -> str:
    """AES-128-CBC 解密消息，返回 XML 字符串。

    流程：base64 解码 → AES-128-CBC 解密 → 去 PKCS7 填充 → 校验 AppID → 提取 XML。

    Args:
        ciphertext_b64: base64 编码的密文。
        encoding_aes_key: 43 位 EncodingAESKey。
        expected_app_id: 期望的 AppID（用于校验）。

    Returns:
        解密后的 XML 字符串。

    Raises:
        ValidationError: 解密失败、AppID 不匹配、消息长度异常。
    """
    key = derive_aes_key(encoding_aes_key)
    iv = key
    try:
        ciphertext = base64.b64decode(ciphertext_b64)
    except Exception as exc:
        raise ValidationError(
            field="Encrypt",
            message="decrypt_failed",
        ) from exc

    try:
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
        decryptor = cipher.decryptor()
        padded = decryptor.update(ciphertext) + decryptor.finalize()
        unpadder = PKCS7(_AES_BLOCK_BITS).unpadder()
        plaintext = unpadder.update(padded) + unpadder.finalize()
    except Exception as exc:
        raise ValidationError(
            field="Encrypt",
            message="decrypt_failed",
        ) from exc

    if len(plaintext) < _RANDOM_BYTES_LEN + _MSG_LEN_BYTES:
        raise ValidationError(
            field="Encrypt",
            message="decrypt_failed",
        )

    msg_len = struct.unpack(">I", plaintext[_RANDOM_BYTES_LEN : _RANDOM_BYTES_LEN + _MSG_LEN_BYTES])[0]
    msg_start = _RANDOM_BYTES_LEN + _MSG_LEN_BYTES
    msg_end = msg_start + msg_len
    if msg_end > len(plaintext):
        raise ValidationError(
            field="Encrypt",
            message="decrypt_failed",
        )

    msg_bytes = plaintext[msg_start:msg_end]
    appid_bytes = plaintext[msg_end:]
    try:
        app_id = appid_bytes.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise ValidationError(
            field="Encrypt",
            message="decrypt_failed",
        ) from exc

    if app_id != expected_app_id:
        raise ValidationError(
            field="AppID",
            message="appid_mismatch",
        )

    try:
        return msg_bytes.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise ValidationError(
            field="Encrypt",
            message="decrypt_failed",
        ) from exc


def compute_signature(token: str, timestamp: str, nonce: str, *extra: str) -> str:
    """计算微信 Webhook SHA1 签名。

    公式：``SHA1(sort([token, timestamp, nonce, *extra])).hexdigest()``

    Args:
        token: Webhook 签名校验令牌。
        timestamp: 时间戳。
        nonce: 随机数。
        *extra: 额外参与签名的字段（如加密消息体 ``encrypt``）。

    Returns:
        40 位十六进制 SHA1 签名。
    """
    parts = sorted([token, timestamp, nonce, *extra])
    raw = "".join(parts).encode("utf-8")
    return hashlib.sha1(raw).hexdigest()


def verify_signature(
    signature: str,
    token: str,
    timestamp: str,
    nonce: str,
    *extra: str,
) -> bool:
    """校验微信 Webhook 签名。

    Args:
        signature: 待校验的签名。
        token: Webhook 签名校验令牌。
        timestamp: 时间戳。
        nonce: 随机数。
        *extra: 额外参与签名的字段。

    Returns:
        签名匹配返回 True，否则 False。
    """
    expected = compute_signature(token, timestamp, nonce, *extra)
    return _constant_time_eq(signature, expected)


def extract_encrypt(xml_text: str) -> str | None:
    """从 XML 中提取 ``<Encrypt>`` 字段内容。

    兼容/安全模式下，入站 XML 含 ``<Encrypt>`` 字段（base64 密文）。

    Args:
        xml_text: 原始 XML 字符串。

    Returns:
        Encrypt 字段内容；不存在返回 None。
    """
    try:
        root = ET.fromstring(xml_text)
    except ET.ParseError:
        return None
    encrypt_node = root.find("Encrypt")
    if encrypt_node is None or encrypt_node.text is None:
        return None
    return encrypt_node.text


def _constant_time_eq(a: str, b: str) -> bool:
    """恒定时间字符串比较，避免时序攻击。"""
    if len(a) != len(b):
        return False
    result = 0
    for x, y in zip(a, b, strict=False):
        result |= ord(x) ^ ord(y)
    return result == 0


__all__ = [
    "derive_aes_key",
    "encrypt_message",
    "decrypt_message",
    "compute_signature",
    "verify_signature",
    "extract_encrypt",
]