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


def generate_key(): # 自动生成密钥
    key = os.urandom(16) 
    return base64.b64encode(key).decode('utf-8') 


def encrypt(plaintext, key): # 加密函数
    cipher = AES.new(base64.b64decode(key), AES.MODE_CFB, bytes([0] * 16))
    ciphertext = cipher.encrypt(pad(plaintext.encode('utf-8'), AES.block_size))
    return base64.b64encode(ciphertext)


def decrypt(ciphertext, key): # 解密函数
    cipher = AES.new(base64.b64decode(key), AES.MODE_CFB, bytes([0] * 16))
    decrypted = unpad(cipher.decrypt(base64.b64decode(ciphertext)), AES.block_size)
    return decrypted.decode('utf-8')


def test_encrypt_decrypt(): # 测试加解密
    testplaintext = plaintext
    key = generate_key()

    testencrypted = encrypt(testplaintext, key)
    print("加密:", testencrypted)

    testdecrypted = decrypt(testencrypted, key)
    testdecrypted = testdecrypted.encode('utf-8').decode('unicode_escape')
    print("解密:", testdecrypted)


# 示例
key = generate_key()
plaintext = "wo shi yi ge test"

test_encrypt_decrypt()

pip3 install pycryptodome