from base64 import b64encode, b64decode
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad


class ShanHuZhuFuAes:
    def __init__(self):
        self.key = 'xlc2ze7qnqg8xi1d'.encode('utf-8')  # 需要一个bytes类型的key
        self.iv = self.key  # 在这个例子中，key和iv是相同的

    def encrypt(self, data):
        cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
        ct_bytes = cipher.encrypt(pad(data.encode('utf-8'), AES.block_size))
        ct = b64encode(ct_bytes).decode('utf-8')
        return ct

    def decrypt(self, data):
        try:
            ct = b64decode(data.encode('utf-8'))
            cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
            pt = unpad(cipher.decrypt(ct), AES.block_size)
            return pt.decode('utf-8')
        except Exception as e:
            print("Incorrect decryption")
            return None