# "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
from fontTools.misc.eexec import encrypt


class AESUtils:
    def __init__(self, key: bytes = None):
        """生成一个随机密钥"""
        self.key = get_random_bytes(32)

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

    def decrypt(self, encrypted_data: str, mode: str = 'ECB') -> str:
        """解密数据"""
        cipher = self._get_cipher(mode)
        encrypted_data = base64.b64decode(encrypted_data)
        plaintext = unpad(cipher.decrypt(encrypted_data), AES.block_size)
        return plaintext.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)
        elif mode == 'CFB':
            return AES.new(self.key, AES.MODE_CFB)
        elif mode == 'OFB':
            return AES.new(self.key, AES.MODE_OFB)
        else:  # 默认是 ECB
            return AES.new(self.key, AES.MODE_ECB)

if __name__ == "__main__":
    # 示例数据和密钥
    data = "Hello, SM4!"
    key = "1234567890abcdef"

    enData , akey = AESUtils(key)

    print(enData)