78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
import json
|
||
import base64
|
||
from Crypto.Cipher import AES
|
||
from Crypto.Util.Padding import pad, unpad
|
||
from Crypto.Random import get_random_bytes
|
||
|
||
|
||
# from .setting import ENCRYPTKEY
|
||
ENCRYPTKEY = b"lunminalinguaai_"
|
||
|
||
|
||
class AESCipher:
|
||
def __init__(self, key):
|
||
"""
|
||
初始化AES加密器
|
||
:param key: 加密密钥（16/24/32字节），可以是字符串或字节
|
||
"""
|
||
if isinstance(key, str):
|
||
key = key.encode("utf-8")
|
||
if len(key) not in [16, 24, 32]:
|
||
raise ValueError("密钥长度必须为16、24或32字节")
|
||
self.key = key
|
||
|
||
def encrypt(self, data):
|
||
"""
|
||
加密数据（支持字符串/列表/字典等JSON可序列化类型）
|
||
:param data: 要加密的数据
|
||
:return: 返回Base64编码的加密字符串
|
||
"""
|
||
# 生成随机初始化向量
|
||
iv = get_random_bytes(AES.block_size)
|
||
|
||
# 创建AES加密器
|
||
cipher = AES.new(self.key, AES.MODE_CBC, iv)
|
||
|
||
# 序列化数据为JSON字符串并编码为字节
|
||
json_data = json.dumps(data)
|
||
plain_bytes = json_data.encode("utf-8")
|
||
|
||
# 填充并加密数据
|
||
padded_bytes = pad(plain_bytes, AES.block_size)
|
||
cipher_bytes = cipher.encrypt(padded_bytes)
|
||
|
||
# 组合IV和密文并进行Base64编码
|
||
encrypted_data = iv + cipher_bytes
|
||
return base64.b64encode(encrypted_data).decode("utf-8")
|
||
|
||
def decrypt(self, enc_data):
|
||
"""
|
||
解密数据并恢复原始格式
|
||
:param enc_data: Base64编码的加密字符串
|
||
:return: 原始数据（保持原始格式）
|
||
"""
|
||
# Base64解码
|
||
encrypted_data = base64.b64decode(enc_data)
|
||
|
||
# 提取初始化向量
|
||
iv = encrypted_data[: AES.block_size]
|
||
cipher_bytes = encrypted_data[AES.block_size :]
|
||
|
||
# 创建AES解密器
|
||
cipher = AES.new(self.key, AES.MODE_CBC, iv)
|
||
|
||
# 解密并去除填充
|
||
decrypted_bytes = cipher.decrypt(cipher_bytes)
|
||
unpadded_bytes = unpad(decrypted_bytes, AES.block_size)
|
||
|
||
# 解码JSON并恢复原始数据结构
|
||
json_data = unpadded_bytes.decode("utf-8")
|
||
return json.loads(json_data)
|
||
|
||
encrypt = AESCipher(ENCRYPTKEY)
|
||
|
||
if __name__ == '__main__':
|
||
data = "sk-3e47a49bf60e49e8ab08bb1f1550aa86"
|
||
enc_data = encrypt.encrypt(data)
|
||
print(enc_data)
|