204 lines
6.8 KiB
Python
204 lines
6.8 KiB
Python
"""敏感字段落盘加密工具。"""
|
||
|
||
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"),
|
||
("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
|