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):
        """生成一个随机密钥"""
        self.key = key if key else get_random_bytes(16)

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

    def decrypt(self, dekey, encrypted_data: str, mode: str = 'ECB') -> str:
        """解密数据"""
        cipher = self._get_decipher(dekey, mode)
        encrypted_data = base64.b64decode(encrypted_data)  # 解码密文
        plaintext = unpad(cipher.decrypt(encrypted_data), AES.block_size)
        return plaintext.decode()

    def _get_encipher(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)

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

if __name__ == '__main__':
    key = b'1234567890abcdef'  # AES 密钥必须是 16 字节
    aes = AESUtils(key)

    # 加密
    plaintext = "Hello, AES!"
    enData, akey = aes.encrypt(plaintext, mode='ECB')
    print("Encrypted Data:", enData)
    print("Key:", akey)

    # 解密
    decrypted = aes.decrypt(aes.key, enData, mode='ECB')
    print("Decrypted:", decrypted)