import base64
import hashlib
import json
from datetime import datetime

from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad

AES_KEY = "djangocms888@"


class AuthEncryptor:

    def __init__(self, key):
        """初始化加密器，使用固定密钥"""
        self.key = self._prepare_key(key)

    def _prepare_key(self, key):
        """将字符串密钥转换为32字节的AES密钥"""
        # 使用SHA256哈希确保密钥长度为32字节（AES-256）
        return hashlib.sha256(key.encode('utf-8')).digest()

    def encrypt(self, data):
        """加密数据"""
        # 生成随机初始化向量
        iv = hashlib.md5(str(
            datetime.now().timestamp()).encode()).digest()[:16]

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

        # 对数据进行填充并加密
        if isinstance(data, dict):
            data_str = json.dumps(data, ensure_ascii=False)
        else:
            data_str = str(data)

        padded_data = pad(data_str.encode('utf-8'), AES.block_size)
        encrypted_data = cipher.encrypt(padded_data)

        # 组合IV和加密数据，并进行Base64编码
        combined = iv + encrypted_data
        return base64.b64encode(combined).decode('utf-8')

    def decrypt(self, encrypted_text):
        """解密数据"""
        try:
            # Base64解码
            combined = base64.b64decode(encrypted_text)

            # 分离IV和加密数据
            iv = combined[:16]
            encrypted_data = combined[16:]

            # 创建AES解密器
            cipher = AES.new(self.key, AES.MODE_CBC, iv)

            # 解密并去除填充
            decrypted_data = unpad(cipher.decrypt(encrypted_data),
                                   AES.block_size)
            decrypted_str = decrypted_data.decode('utf-8')

            # 尝试解析为JSON，否则返回字符串
            try:
                return json.loads(decrypted_str)
            except:
                return decrypted_str

        except Exception as e:
            print(f"解密错误: {e}")
            return None
