wehub-resource-sync
53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
"""
|
|
加密工具
|
|
用于加密存储敏感信息如cookies
|
|
"""
|
|
|
|
import os
|
|
import base64
|
|
from cryptography.fernet import Fernet
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 获取加密密钥
|
|
def get_encryption_key():
|
|
"""获取加密密钥"""
|
|
# 从环境变量获取密钥，如果没有则生成一个
|
|
key = os.getenv('ENCRYPTION_KEY')
|
|
if not key:
|
|
# 生成新密钥
|
|
key = Fernet.generate_key()
|
|
logger.warning("ENCRYPTION_KEY环境变量未设置，使用临时密钥。请设置环境变量以保持数据安全。")
|
|
|
|
if isinstance(key, str):
|
|
key = key.encode()
|
|
|
|
return key
|
|
|
|
def encrypt_data(data: str) -> str:
|
|
"""加密数据"""
|
|
try:
|
|
key = get_encryption_key()
|
|
f = Fernet(key)
|
|
encrypted_data = f.encrypt(data.encode())
|
|
return base64.b64encode(encrypted_data).decode()
|
|
except Exception as e:
|
|
logger.error(f"加密数据失败: {str(e)}")
|
|
raise
|
|
|
|
def decrypt_data(encrypted_data: str) -> str:
|
|
"""解密数据"""
|
|
try:
|
|
key = get_encryption_key()
|
|
f = Fernet(key)
|
|
decoded_data = base64.b64decode(encrypted_data.encode())
|
|
decrypted_data = f.decrypt(decoded_data)
|
|
return decrypted_data.decode()
|
|
except Exception as e:
|
|
logger.error(f"解密数据失败: {str(e)}")
|
|
raise
|
|
|