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

from __future__ import annotations

import base64
import logging

from cryptography import x509
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

logger = logging.getLogger(__name__)

MAX_TIMESTAMP_DRIFT_SECONDS = 300


class SignatureVerificationError(Exception):
    pass


class DecryptionError(Exception):
    pass


class WechatPayCrypto:
    def __init__(self, api_v3_key: str):
        if len(api_v3_key) != 64:
            raise ValueError("APIv3 key must be 32 bytes (64 hex characters)")
        self._api_v3_key = api_v3_key

    @staticmethod
    def build_authorization_signature(
        method: str,
        url: str,
        body: str,
        timestamp: str,
        nonce_str: str,
        private_key_pem: str,
    ) -> str:
        from cryptography.hazmat.primitives.serialization import load_pem_private_key

        private_key = load_pem_private_key(private_key_pem.encode("utf-8"), password=None, backend=default_backend())
        sign_message = f"{method}\n{url}\n{timestamp}\n{nonce_str}\n{body}\n"
        signature = private_key.sign(sign_message.encode("utf-8"), padding.PKCS1v15(), hashes.SHA256())
        return base64.b64encode(signature).decode("utf-8")

    def verify_signature(
        self,
        body: bytes,
        signature: str,
        timestamp: str,
        nonce: str,
        serial_no: str,
        platform_cert_pem: str,
    ) -> bool:
        try:
            dr = _check_timestamp_drift(timestamp)
            if dr:
                logger.warning("WechatPay signature timestamp drift: %s", dr)
                return False
        except (ValueError, TypeError):
            return False

        sign_message = f"{timestamp}\n{nonce}\n{body.decode('utf-8')}\n"

        try:
            cert = x509.load_pem_x509_certificate(platform_cert_pem.encode("utf-8"), default_backend())
            public_key = cert.public_key()
            public_key.verify(
                base64.b64decode(signature),
                sign_message.encode("utf-8"),
                padding.PKCS1v15(),
                hashes.SHA256(),
            )
            return True
        except Exception:
            logger.warning("WechatPay signature verification failed for serial=%s", serial_no)
            return False

    def decrypt_resource(
        self,
        ciphertext: str,
        nonce: str,
        associated_data: str,
    ) -> str:
        try:
            aesgcm = AESGCM(bytes.fromhex(self._api_v3_key))
            plaintext = aesgcm.decrypt(
                nonce.encode("utf-8"),
                base64.b64decode(ciphertext),
                associated_data.encode("utf-8") if associated_data else None,
            )
            return plaintext.decode("utf-8")
        except Exception as e:
            raise DecryptionError(f"AES-GCM decryption failed: {e}") from e

    def verify_signature_with_raw_key(
        self,
        body: bytes,
        signature: str,
        timestamp: str,
        nonce: str,
        public_key_pem: str,
    ) -> bool:
        from cryptography.hazmat.primitives.serialization import load_pem_public_key

        dr = _check_timestamp_drift(timestamp)
        if dr:
            logger.warning("WechatPay signature timestamp drift: %s", dr)
            return False

        sign_message = f"{timestamp}\n{nonce}\n{body.decode('utf-8')}\n"
        try:
            public_key = load_pem_public_key(public_key_pem.encode("utf-8"), default_backend())
            public_key.verify(
                base64.b64decode(signature),
                sign_message.encode("utf-8"),
                padding.PKCS1v15(),
                hashes.SHA256(),
            )
            return True
        except Exception:
            logger.warning("WechatPay signature verification failed with raw public key")
            return False


def _check_timestamp_drift(timestamp: str) -> str | None:
    import time

    ts = int(timestamp)
    now = int(time.time())
    diff = abs(now - ts)
    if diff > MAX_TIMESTAMP_DRIFT_SECONDS:
        return f"drift={diff}s > {MAX_TIMESTAMP_DRIFT_SECONDS}s"
    return None