#!/usr/bin/python
# -*- coding=utf-8 -*-
# @Create Time:		2024-02-22 14:41:39
# @Last Modified time: 2024-02-22 17:07:53
from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex

class AES_Util(object):
    """docstring for AES_Util"""
    def __init__(self, mode, key, iv=None):
        super(AES_Util, self).__init__()
        self.set_key(key)
        if mode == "CBC":
            self.mode = AES.MODE_CBC
            self.set_iv(iv)
        else:
            self.mode = AES.MODE_ECB

    def set_key(self, key):
        self.key = key.encode("utf-8")

    def set_iv(self, iv):
        if isinstance(iv, bytes):
            self.iv = iv

    # 如果text不足16位的倍数就用空格补足为16位
    def add_to_16(self, text):
        if len(text.encode('utf-8')) % 16:
            add = 16 - (len(text.encode('utf-8')) % 16)
        else:
            add = 0
        text = text + ('\0' * add)
        return text.encode('utf-8')

    # 加密函数
    def encrypt(self, text):
        text = self.add_to_16(text)
        if self.mode == AES.MODE_CBC:
            cryptos = AES.new(self.key, self.mode, self.iv)
        else:
            cryptos = AES.new(self.key, self.mode)
        cipher_text = cryptos.encrypt(text)
        # 因为AES加密后的字符串不一定是ascii字符集的，输出保存可能存在问题，所以这里转为16进制字符串
        return b2a_hex(cipher_text)

    # 解密后，去掉补足的空格用strip() 去掉
    def decrypt(self, text):
        if self.mode == AES.MODE_CBC:
            cryptos = AES.new(self.key, self.mode, self.iv)
        else:
            cryptos = AES.new(self.key, self.mode)
        plain_text = cryptos.decrypt(a2b_hex(text))
        return bytes.decode(plain_text).rstrip('\0')

if __name__ == '__main__':
    aes = AES_Util("CBC", key='9999999999999999', iv=b'qqqqqqqqqqqqqqqq')
    e = aes.encrypt("hello world")  # 加密
    print(e)
    d = aes.decrypt(e)  # 解密
    print(d)