# -*- coding: utf-8 -*-

import base64

from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad

BLOCK_SIZE = 16
KEY = "base64 encoded aes key"


class AesFunc(object):

    def __init__(self):
        key_b64decode = base64.b64decode(KEY)
        # ECB模式
        self.cipher = AES.new(key_b64decode, AES.MODE_ECB)

    def encrypt(self, plaintext):
        plaintext_bytes = plaintext.encode('utf-8')
        # pkcs7填充
        plaintext_paded = pad(plaintext_bytes, BLOCK_SIZE)
        # 加密
        ciphertext_bytes = self.cipher.encrypt(plaintext_paded)
        # base64编码
        ciphertext = base64.b64encode(ciphertext_bytes)
        return ciphertext.decode('utf-8')

    def decrypt(self, ciphertext):
        # base64解码
        ciphertext_b64decode = base64.b64decode(ciphertext)
        # 解密
        plaintext_bytes = self.cipher.decrypt(ciphertext_b64decode)
        # 取消填充
        plaintext = unpad(plaintext_bytes, BLOCK_SIZE)
        return plaintext.decode('utf-8')

    def run(self, msg):
        ciphertext = self.encrypt(msg)
        print("密文: " + ciphertext)

        plaintext = self.decrypt(ciphertext)
        print("明文: " + plaintext)


if __name__ == '__main__':
    msg = '测试 AES/ECB/PKCS5Padding 加解密'
    AesFunc().run(msg)