# "ecb", "cbc", "cfb", "ofb"

from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes
import base64

class AESUtils:
    def __init__(self, key: bytes = None):
        """初始化 AES 加密工具，如果没有提供密钥，则生成一个随机密钥"""
        if key is None:
            self.key = get_random_bytes(32)  # AES-256
        else:
            self.key = key

    def encrypt(self, data: str, mode: str = 'ECB') -> str:
        """加密数据"""
        cipher = self._get_cipher(mode)
        data = pad(data.encode(), AES.block_size)  # 填充数据
        ciphertext = cipher.encrypt(data)
        return base64.b64encode(ciphertext).decode()

    def _get_cipher(self, mode: str) -> AES:
        """根据模式返回相应的 AES cipher"""
        iv = None
        if mode == 'CBC':
            iv = get_random_bytes(AES.block_size)
            return AES.new(self.key, AES.MODE_CBC, iv), self.key
        elif mode == 'CFB':
            return AES.new(self.key, AES.MODE_CFB), self.key
        elif mode == 'OFB':
            return AES.new(self.key, AES.MODE_OFB), self.key
        else:  # 默认是 ECB
            return AES.new(self.key, AES.MODE_ECB), self.key