""" API密钥加密工具 使用AES加密确保数据库中API密钥的安全性 """ import os import base64 import secrets from typing import Optional, Tuple from cryptography.fernet import Fernet from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC from src.utils.logger import setup_logger from src.utils.env_config import env_config logger = setup_logger("encryption") class APIKeyEncryption: """API密钥加密管理器""" def __init__(self): self._fernet = None self._init_encryption_key() def _init_encryption_key(self): """初始化加密密钥""" # 尝试从环境变量获取加密密钥 encryption_key = os.getenv('ENCRYPTION_KEY') if not encryption_key: # 首先尝试从数据库获取已存储的密钥 encryption_key = self._get_stored_encryption_key() if not encryption_key: # 检查数据库中是否已存在加密数据 has_encrypted_data = self._check_existing_encrypted_data() if has_encrypted_data: logger.error("Found encrypted API keys in database but no ENCRYPTION_KEY available!") logger.error("Option 1: Set ENCRYPTION_KEY in your .env file if you have the key") logger.error("Option 2: Delete encrypted channels and restart to generate new key") raise ValueError("Missing ENCRYPTION_KEY - cannot decrypt existing encrypted data") else: # 没有加密数据，生成新密钥并自动保存到数据库配置 encryption_key = self._generate_encryption_key() self._store_encryption_key(encryption_key) logger.info("Generated new encryption key and stored in database") logger.info("For better security, consider moving this to .env file:") logger.info(f"ENCRYPTION_KEY={encryption_key}") else: logger.info("Using stored encryption key from database") try: # 验证密钥格式 self._fernet = Fernet(encryption_key.encode()) logger.info("Encryption system initialized successfully") except Exception as e: logger.error(f"Failed to initialize encryption: {e}") raise ValueError("Invalid encryption key format") def _check_existing_encrypted_data(self) -> bool: """检查数据库中是否存在加密数据（避免循环导入）""" try: from src.utils.env_config import env_config import sqlite3 db_path = env_config.database_path if not os.path.exists(db_path): return False conn = sqlite3.connect(db_path) cursor = conn.execute( "SELECT COUNT(*) FROM channels WHERE api_key LIKE 'encrypted:%'" ) count = cursor.fetchone()[0] conn.close() return count > 0 except Exception: # 如果查询失败（表不存在等），假设没有加密数据 return False def _generate_encryption_key(self) -> str: """生成新的加密密钥""" # 生成32字节的随机密钥 key = Fernet.generate_key() return key.decode() def _store_encryption_key(self, encryption_key: str): """将加密密钥存储到数据库配置中""" try: from src.utils.env_config import env_config import sqlite3 db_path = env_config.database_path os.makedirs(os.path.dirname(db_path), exist_ok=True) conn = sqlite3.connect(db_path) # 创建配置表（如果不存在） conn.execute(''' CREATE TABLE IF NOT EXISTS config ( key TEXT PRIMARY KEY, value TEXT NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) ''') # 存储加密密钥 conn.execute( 'INSERT OR REPLACE INTO config (key, value, updated_at) VALUES (?, ?, CURRENT_TIMESTAMP)', ('encryption_key', encryption_key) ) conn.commit() conn.close() except Exception as e: logger.warning(f"Failed to store encryption key in database: {e}") def _get_stored_encryption_key(self) -> Optional[str]: """从数据库配置中获取加密密钥""" try: from src.utils.env_config import env_config import sqlite3 db_path = env_config.database_path if not os.path.exists(db_path): return None conn = sqlite3.connect(db_path) cursor = conn.execute('SELECT value FROM config WHERE key = ?', ('encryption_key',)) result = cursor.fetchone() conn.close() return result[0] if result else None except Exception: return None def encrypt_api_key(self, api_key: str) -> str: """加密API密钥""" if not api_key: return "" try: encrypted_data = self._fernet.encrypt(api_key.encode()) # 返回base64编码的加密数据，添加前缀标识 return f"encrypted:{base64.b64encode(encrypted_data).decode()}" except Exception as e: logger.error(f"Failed to encrypt API key: {e}") raise ValueError("Encryption failed") def decrypt_api_key(self, encrypted_api_key: str) -> str: """解密API密钥""" if not encrypted_api_key: return "" # 检查是否是加密格式 if not encrypted_api_key.startswith("encrypted:"): # 兼容未加密的旧数据 logger.warning("Found unencrypted API key, consider re-saving to encrypt it") return encrypted_api_key try: # 移除前缀并解码 encrypted_data = encrypted_api_key[10:] # 移除 "encrypted:" 前缀 encrypted_bytes = base64.b64decode(encrypted_data.encode()) # 解密 decrypted_data = self._fernet.decrypt(encrypted_bytes) return decrypted_data.decode() except Exception as e: logger.error(f"Failed to decrypt API key: {e}") raise ValueError("Decryption failed - possibly wrong encryption key") def is_encrypted(self, data: str) -> bool: """检查数据是否已加密""" return data.startswith("encrypted:") if data else False def rotate_encryption_key(self, new_key: str, old_encrypted_data: list) -> list: """ 轮换加密密钥（高级功能） 重新加密所有数据使用新密钥 """ # 保存当前密钥 old_fernet = self._fernet try: # 设置新密钥 self._fernet = Fernet(new_key.encode()) # 重新加密所有数据 reencrypted_data = [] for encrypted_item in old_encrypted_data: if self.is_encrypted(encrypted_item): # 使用旧密钥解密 self._fernet = old_fernet decrypted = self.decrypt_api_key(encrypted_item) # 使用新密钥加密 self._fernet = Fernet(new_key.encode()) reencrypted = self.encrypt_api_key(decrypted) reencrypted_data.append(reencrypted) else: reencrypted_data.append(encrypted_item) logger.info(f"Successfully rotated encryption key for {len(reencrypted_data)} items") return reencrypted_data except Exception as e: # 恢复旧密钥 self._fernet = old_fernet logger.error(f"Key rotation failed: {e}") raise # 全局加密管理器实例 encryption_manager = APIKeyEncryption()