"""敏感字段落盘加密工具。""" from __future__ import annotations import base64 import hashlib import os import warnings from dataclasses import dataclass from functools import lru_cache from Crypto.Cipher import AES from sqlalchemy import text from sqlalchemy.engine import Engine from sqlalchemy.types import Text, TypeDecorator from loguru import logger _PREFIX = "enc:v1:" _FALLBACK_SECRET = "douyu-login-py-dev-storage-key-change-me" @dataclass(frozen=True) class _KeyCandidate: """可用于解密的候选密钥。""" name: str secret: str primary: bool = False @lru_cache(maxsize=1) def _key_candidates() -> tuple[_KeyCandidate, ...]: """返回主加密密钥与历史兼容解密密钥。""" app_key = os.getenv("APP_ENCRYPTION_KEY", "").strip() jwt_key = os.getenv("JWT_SECRET_KEY", "").strip() if app_key: primary = _KeyCandidate("APP_ENCRYPTION_KEY", app_key, True) elif jwt_key: primary = _KeyCandidate("JWT_SECRET_KEY", jwt_key, True) else: warnings.warn( "未设置 APP_ENCRYPTION_KEY,敏感字段将使用开发兜底密钥加密。" "生产环境请设置稳定且随机的 APP_ENCRYPTION_KEY。", stacklevel=2, ) primary = _KeyCandidate("fallback", _FALLBACK_SECRET, True) candidates = [primary] for name, secret in ( ("APP_ENCRYPTION_KEY", app_key), ("JWT_SECRET_KEY", jwt_key), ("fallback", _FALLBACK_SECRET), ): if secret and all(item.secret != secret for item in candidates): candidates.append(_KeyCandidate(name, secret, False)) return tuple(candidates) def _aes_key(secret: str) -> bytes: """从配置密钥派生 AES-256 key。""" return hashlib.sha256(secret.encode("utf-8")).digest() def _b64e(data: bytes) -> str: return base64.urlsafe_b64encode(data).decode("ascii") def _b64d(data: str) -> bytes: return base64.urlsafe_b64decode(data.encode("ascii")) def is_encrypted(value: str | None) -> bool: """判断值是否已经是当前加密格式。""" return bool(value) and value.startswith(_PREFIX) def encrypt_value(value: str | None) -> str | None: """加密字符串;空值保持空值,避免影响空值筛选。""" if value is None or value == "": return value if is_encrypted(value): return value candidate = _key_candidates()[0] nonce = os.urandom(12) cipher = AES.new(_aes_key(candidate.secret), AES.MODE_GCM, nonce=nonce) ciphertext, tag = cipher.encrypt_and_digest(value.encode("utf-8")) return _PREFIX + _b64e(nonce + tag + ciphertext) def _decrypt_with_candidate(value: str, candidate: _KeyCandidate) -> str: raw = _b64d(value[len(_PREFIX):]) if len(raw) < 28: raise ValueError("密文字段长度无效") nonce = raw[:12] tag = raw[12:28] ciphertext = raw[28:] cipher = AES.new(_aes_key(candidate.secret), AES.MODE_GCM, nonce=nonce) return cipher.decrypt_and_verify(ciphertext, tag).decode("utf-8") def decrypt_value(value: str | None) -> str | None: """解密字符串;旧明文会原样返回。""" if value is None or value == "" or not is_encrypted(value): return value last_error: Exception | None = None for candidate in _key_candidates(): try: return _decrypt_with_candidate(value, candidate) except Exception as exc: last_error = exc raise ValueError("敏感字段解密失败,请确认 APP_ENCRYPTION_KEY 是否正确") from last_error def decrypt_value_with_key_name(value: str) -> tuple[str, str]: """解密并返回实际使用的密钥名称,用于历史密钥迁移。""" if not is_encrypted(value): return value, "plain" last_error: Exception | None = None for candidate in _key_candidates(): try: return _decrypt_with_candidate(value, candidate), candidate.name except Exception as exc: last_error = exc raise ValueError("敏感字段解密失败,请确认 APP_ENCRYPTION_KEY 是否正确") from last_error class EncryptedText(TypeDecorator): """SQLAlchemy 字段类型:数据库密文,Python 属性明文。""" impl = Text cache_ok = True def process_bind_param(self, value, dialect): if value is None: return None return encrypt_value(str(value)) def process_result_value(self, value, dialect): return decrypt_value(value) _SENSITIVE_COLUMNS: tuple[tuple[str, str], ...] = ( ("accounts", "password"), ("accounts", "email"), ("accounts", "email_password"), ("login_tasks", "cookie"), ("huya_accounts", "account_password"), ("huya_accounts", "cookie"), ("proxy_config", "api_url"), ("proxy_config", "http"), ("proxy_config", "https"), ("proxy_config", "whitelist_uid"), ("proxy_config", "whitelist_ukey"), ) def _table_exists(engine: Engine, table_name: str) -> bool: query = text("SELECT name FROM sqlite_master WHERE type='table' AND name=:name") if not engine.url.get_backend_name().startswith("sqlite"): return True with engine.connect() as conn: return conn.execute(query, {"name": table_name}).first() is not None def encrypt_existing_sensitive_data(engine: Engine) -> int: """把历史明文敏感字段迁移成密文;已加密但非主密钥的值会重加密。 解密失败的值(密钥丢失/变更)会被跳过并记录警告。 """ primary_name = _key_candidates()[0].name changed = 0 with engine.begin() as conn: for table_name, column_name in _SENSITIVE_COLUMNS: if not _table_exists(engine, table_name): continue rows = conn.execute( text(f"SELECT id, {column_name} FROM {table_name} WHERE {column_name} IS NOT NULL") ).mappings().all() for row in rows: raw_value = row[column_name] if raw_value is None or raw_value == "": continue try: plain_value, key_name = decrypt_value_with_key_name(str(raw_value)) except ValueError: # 解密失败(密钥变更/丢失),跳过该值,避免启动阻塞 logger.warning( f"跳过无法解密的字段 {table_name}.{column_name}[id={row['id']}]," f"可能是加密密钥已变更" ) continue if is_encrypted(str(raw_value)) and key_name == primary_name: continue encrypted = encrypt_value(plain_value) conn.execute( text(f"UPDATE {table_name} SET {column_name} = :value WHERE id = :id"), {"value": encrypted, "id": row["id"]}, ) changed += 1 return changed