import os
import base64
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from fastapi import HTTPException

# 硬编码AES密钥（32字节，AES-256）
AES_SECRET_KEY = b"jr1vA6tfWMHOYi6UXw67UuO6fdak2rMa"
AES_BLOCK_SIZE = 16  # AES固定块大小

# 校验密钥长度（确保符合AES规范）
valid_key_lengths = [16, 24, 32]
if len(AES_SECRET_KEY) not in valid_key_lengths:
    raise ValueError(
        f"AES密钥长度必须为{valid_key_lengths}字节，当前为{len(AES_SECRET_KEY)}字节"
    )


def aes_encrypt(plaintext: str) -> dict:
    """AES-CBC模式加密（返回密文+IV，均为Base64编码）"""
    try:
        # 生成随机IV（16字节）
        iv = os.urandom(AES_BLOCK_SIZE)

        # 创建加密器
        cipher = AES.new(AES_SECRET_KEY, AES.MODE_CBC, iv)

        # 明文填充并加密
        padded_plaintext = pad(plaintext.encode("utf-8"), AES_BLOCK_SIZE)
        ciphertext = base64.b64encode(cipher.encrypt(padded_plaintext)).decode("utf-8")
        iv_base64 = base64.b64encode(iv).decode("utf-8")

        return {
            "ciphertext": ciphertext,
            "iv": iv_base64,
            "algorithm": "AES-CBC"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"AES加密失败：{str(e)}") from e


def aes_decrypt(ciphertext: str, iv: str) -> str:
    """AES-CBC模式解密"""
    try:
        # 解码Base64
        ciphertext_bytes = base64.b64decode(ciphertext)
        iv_bytes = base64.b64decode(iv)

        # 创建解密器
        cipher = AES.new(AES_SECRET_KEY, AES.MODE_CBC, iv_bytes)

        # 解密并去填充
        decrypted_bytes = unpad(cipher.decrypt(ciphertext_bytes), AES_BLOCK_SIZE)
        return decrypted_bytes.decode("utf-8")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"AES解密失败：{str(e)}") from e