|
|
|
|
# "ecb", "cbc", "cfb", "ofb"
|
|
|
|
|
|
|
|
|
|
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):
|
|
|
|
|
"""初始化 AES 加密工具，如果没有提供密钥，则生成一个随机密钥"""
|
|
|
|
|
if key is None:
|
|
|
|
|
self.key = get_random_bytes(32) # AES-256
|
|
|
|
|
else:
|
|
|
|
|
self.key = key
|
|
|
|
|
|
|
|
|
|
def encrypt(self, data: str, mode: str = 'ECB') -> (str,str):
|
|
|
|
|
"""加密数据"""
|
|
|
|
|
cipher = self._get_cipher(mode)
|
|
|
|
|
data = pad(data.encode(), AES.block_size) # 填充数据
|
|
|
|
|
ciphertext = cipher.encrypt(data)
|
|
|
|
|
return base64.b64encode(ciphertext).decode(),self.key
|
|
|
|
|
|
|
|
|
|
def _get_cipher(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)
|
|
|
|
|
|
