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

"""微信 iLink AES 加解密与 X-WECHAT-UIN 工具模块。

提供 AES-128-ECB 加解密（PKCS7 填充）、三种 AES 密钥编码格式归一化、
X-WECHAT-UIN 防重放头生成与上传密文大小计算。

安全说明：
- AES-128-**ECB** 模式为 iLink CDN 协议强制要求（上传密钥由服务端
  ``getuploadurl`` 返回，下载密钥由消息体提供，客户端无法单方面改 GCM/CBC）。
  ECB 对相同明文块产生相同密文块，存在模式泄露风险；此处通过每次上传
  生成一次性随机密钥（``generate_random_aes_key``）缓解，且加密对象为
  高熵文件内容（图片/视频），风险有限。**禁止** 将本模块用于通用数据加密。
- ``generate_x_wechat_uin`` 仅 4 字节（32-bit）熵，为 iLink 协议头格式约定，
  不应用于高安全场景的防重放。

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

from __future__ import annotations

import base64
import os
import re
import struct

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

from yuxi.channels.contract.errors import ValidationError

# 32 个十六进制字符的正则，用于识别 Format C 与 Format B 的内层 hex 字符串
_HEX32_PATTERN = re.compile(r"[0-9a-fA-F]{32}")

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


def normalize_aes_key(aes_key: str) -> bytes:
    """将三种 AES 密钥编码格式归一化为 16 字节原始密钥。

    支持的格式（按检测优先级）：

    - **Format C**: 直接十六进制（32 字符），匹配 ``^[0-9a-fA-F]{32}$``，直接 hex 解码。
    - **Format A**: base64(原始 16 字节)，base64 解码后正好 16 字节。
    - **Format B**: base64(十六进制字符串)，base64 解码后得到 32 字符 hex 串，再 hex 解码。

    Args:
        aes_key: AES 密钥字符串，可能为上述任一格式。

    Returns:
        长度恒为 16 的原始密钥字节。

    Raises:
        ValidationError: 格式无法识别时抛出 ``invalid_aes_key_format: {aes_key}``；
            最终长度不为 16 字节时抛出 ``aes_key_must_be_16_bytes``。
    """
    # Format C: 32 个十六进制字符，直接 hex 解码
    if _HEX32_PATTERN.fullmatch(aes_key):
        key = bytes.fromhex(aes_key)
    else:
        # 尝试 base64 解码；validate=True 保证非 base64 字符直接报错
        try:
            decoded = base64.b64decode(aes_key, validate=True)
        except ValueError:
            raise ValidationError(
                field="aes_key",
                message=f"invalid_aes_key_format: {aes_key}",
            ) from None

        if len(decoded) == 16:
            # Format A: base64 解码后正好 16 字节
            key = decoded
        elif len(decoded) == 32:
            # Format B: base64 解码后应为 32 字符的十六进制字符串
            try:
                hex_str = decoded.decode("ascii")
            except UnicodeDecodeError:
                raise ValidationError(
                    field="aes_key",
                    message=f"invalid_aes_key_format: {aes_key}",
                ) from None
            if not _HEX32_PATTERN.fullmatch(hex_str):
                raise ValidationError(
                    field="aes_key",
                    message=f"invalid_aes_key_format: {aes_key}",
                )
            key = bytes.fromhex(hex_str)
        else:
            raise ValidationError(
                field="aes_key",
                message=f"invalid_aes_key_format: {aes_key}",
            )

    if len(key) != 16:
        raise ValidationError(
            field="aes_key",
            message="aes_key_must_be_16_bytes",
        )
    return key


def encrypt(plaintext: bytes, aes_key: str) -> bytes:
    """AES-128-ECB 加密（PKCS7 填充）。

    Args:
        plaintext: 待加密的明文字节。
        aes_key: AES 密钥字符串（任一支持格式，由 ``normalize_aes_key`` 归一化）。

    Returns:
        密文字节。
    """
    key = normalize_aes_key(aes_key)
    padder = PKCS7(_AES_BLOCK_BITS).padder()
    padded = padder.update(plaintext) + padder.finalize()
    cipher = Cipher(algorithms.AES(key), modes.ECB())
    encryptor = cipher.encryptor()
    return encryptor.update(padded) + encryptor.finalize()


def decrypt(ciphertext: bytes, aes_key: str) -> bytes:
    """AES-128-ECB 解密（PKCS7 去填充）。

    Args:
        ciphertext: 待解密的密文字节。
        aes_key: AES 密钥字符串（任一支持格式，由 ``normalize_aes_key`` 归一化）。

    Returns:
        明文字节。
    """
    key = normalize_aes_key(aes_key)
    cipher = Cipher(algorithms.AES(key), modes.ECB())
    decryptor = cipher.decryptor()
    padded = decryptor.update(ciphertext) + decryptor.finalize()
    unpadder = PKCS7(_AES_BLOCK_BITS).unpadder()
    return unpadder.update(padded) + unpadder.finalize()


def generate_x_wechat_uin() -> str:
    """生成 X-WECHAT-UIN 防重放头。

    流程：4 字节随机数 → uint32 大端 → 十进制字符串 → base64 编码。
    每次调用产生唯一值，用于请求防重放。

    Returns:
        base64 编码的 X-WECHAT-UIN 字符串。
    """
    raw = os.urandom(4)
    uint32_value = struct.unpack(">I", raw)[0]
    decimal_str = str(uint32_value)
    return base64.b64encode(decimal_str.encode("utf-8")).decode("ascii")


def calc_ciphertext_size(raw_size: int) -> int:
    """计算上传密文大小（含 PKCS7 填充）。

    公式：``ceil((raw_size + 1) / 16) * 16``，其中 ``+1`` 保证 ``raw_size``
    为 16 整数倍时仍追加一个完整填充块。使用整数算术实现避免浮点误差。

    Args:
        raw_size: 原始数据字节数。

    Returns:
        密文字节数（16 的整数倍）。
    """
    return ((raw_size + 1 + 15) // 16) * 16


def generate_random_aes_key() -> str:
    """生成随机 AES 密钥（上传一次性密钥，Format A）。

    生成 16 个随机字节，返回其 base64 编码字符串，可被 ``normalize_aes_key``
    直接识别为 Format A。

    Returns:
        base64 编码的 16 字节随机密钥字符串。
    """
    raw = os.urandom(16)
    return base64.b64encode(raw).decode("ascii")


__all__ = [
    "normalize_aes_key",
    "encrypt",
    "decrypt",
    "generate_x_wechat_uin",
    "calc_ciphertext_size",
    "generate_random_aes_key",
]