live-hub-py/web/backend/crypto_storage.py

"""敏感字段落盘加密工具。"""

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


_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"),
    ("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
                plain_value, key_name = decrypt_value_with_key_name(str(raw_value))
                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