"""Key Escrow(MASTER_KEY 密钥托管)测试 覆盖两层: - crypto 层:build_escrow/recover_master_key 往返、错误钥匙拒绝、格式校验 - CLI 层:setup → 模拟 .env 丢失 key → recover --write 找回回写(§3.3 丢失演练) 运行:.venv/bin/pytest tests/test_escrow.py -v """ import os import re import subprocess import sys from pathlib import Path import pytest from cryptography.fernet import Fernet from app.core import crypto from app.core.config import settings BASE_DIR = Path(__file__).resolve().parent.parent SETUP_SCRIPT = BASE_DIR / "scripts" / "setup_key_escrow.py" RECOVER_SCRIPT = BASE_DIR / "scripts" / "recover_master_key.py" # ----------------------------- crypto 层 ----------------------------- # def test_build_recover_roundtrip(monkeypatch): """托管往返:build_escrow 后用同一把钥匙能还原 MASTER_KEY""" monkeypatch.setattr(settings, "MASTER_KEY", Fernet.generate_key().decode()) restore_key = Fernet.generate_key().decode() escrow = crypto.build_escrow(restore_key) assert escrow.startswith(crypto.ESCROW_PREFIX) assert crypto.recover_master_key(restore_key, escrow) == settings.MASTER_KEY def test_recover_wrong_key_rejected(monkeypatch): """错误恢复钥匙必须抛 ValueError(不能静默返回错误 key)""" monkeypatch.setattr(settings, "MASTER_KEY", Fernet.generate_key().decode()) escrow = crypto.build_escrow(Fernet.generate_key().decode()) with pytest.raises(ValueError): crypto.recover_master_key(Fernet.generate_key().decode(), escrow) def test_recover_invalid_prefix(): """缺 v1: 前缀视为格式无效""" with pytest.raises(ValueError): crypto.recover_master_key(Fernet.generate_key().decode(), "garbage-no-prefix") def test_build_escrow_requires_master_key(monkeypatch): """MASTER_KEY 未配置时拒绝建立托管(避免生成解不开真实数据的废档)""" monkeypatch.setattr(settings, "MASTER_KEY", "") with pytest.raises(RuntimeError): crypto.build_escrow(Fernet.generate_key().decode()) # ----------------------------- CLI 端到端 ----------------------------- # def _run_script(script: Path, env_extra: dict, args: list) -> subprocess.CompletedProcess: """以子进程运行脚本(环境变量覆盖 .env:python-dotenv 不覆盖已存在变量)""" env = {**os.environ, **env_extra} return subprocess.run( [sys.executable, str(script), *args], capture_output=True, text=True, env=env ) # Fernet key = base64url(32 字节) = 43 个 [A-Za-z0-9_-] 字符 + '=' _FERNET_KEY_RE = re.compile(r"^[A-Za-z0-9_-]{43}=$") def _extract_restore_key(stdout: str) -> str: """从 setup 输出提取 RESTORE_KEY(44 字符 Fernet key,独立一行)""" for line in stdout.splitlines(): s = line.strip() if _FERNET_KEY_RE.match(s): return s raise AssertionError("setup 输出未包含 RESTORE_KEY:\n" + stdout) def test_cli_loss_drill(tmp_path): """丢失演练:托管建立 → .env 中 key 变为错误值 → recover --write 找回真实 key""" master_key = Fernet.generate_key().decode() escrow_path = tmp_path / "master_key.escrow" env_path = tmp_path / ".env" # 模拟丢失现场:.env 里是错误的 key,其它配置需原样保留 env_path.write_text("API_KEY=abc\nMASTER_KEY=lost-wrong-key\n", encoding="utf-8") # 1. 建立托管(MASTER_KEY 环境变量 = 正确 key) r = _run_script(SETUP_SCRIPT, {"MASTER_KEY": master_key}, ["--escrow", str(escrow_path)]) assert r.returncode == 0, r.stderr assert escrow_path.exists() assert (escrow_path.stat().st_mode & 0o777) == 0o600 restore_key = _extract_restore_key(r.stdout) # 2. 回验模式:--key 校验托管可解且与当前 MASTER_KEY 一致 r = _run_script( SETUP_SCRIPT, {"MASTER_KEY": master_key}, ["--key", restore_key, "--escrow", str(escrow_path)], ) assert r.returncode == 0 and "通过" in r.stdout # 3. 恢复:即使运行时环境里的 MASTER_KEY 是错的,也能凭 RESTORE_KEY 找回并回写 .env r = _run_script( RECOVER_SCRIPT, {"MASTER_KEY": "irrelevant", "RESTORE_KEY": restore_key}, ["--escrow", str(escrow_path), "--env", str(env_path), "--write"], ) assert r.returncode == 0, r.stderr text = env_path.read_text(encoding="utf-8") assert f"MASTER_KEY={master_key}" in text assert "API_KEY=abc" in text # 其它配置行不动 assert "lost-wrong-key" not in text assert (tmp_path / ".env.bak-pre-recover").exists() # 回写前已备份 # 4. 错误钥匙恢复必须失败(退出码 1,.env 不被污染) r = _run_script( RECOVER_SCRIPT, {"RESTORE_KEY": Fernet.generate_key().decode()}, ["--escrow", str(escrow_path), "--env", str(env_path), "--write"], ) assert r.returncode == 1 and "失败" in (r.stdout + r.stderr) assert f"MASTER_KEY={master_key}" in env_path.read_text(encoding="utf-8") # 5. 幂等:托管已存在时默认不重建 r = _run_script(SETUP_SCRIPT, {"MASTER_KEY": master_key}, ["--escrow", str(escrow_path)]) assert r.returncode == 2 and "已存在" in r.stderr # 6. --force 重建:新钥匙生效,旧钥匙作废 r = _run_script( SETUP_SCRIPT, {"MASTER_KEY": master_key}, ["--escrow", str(escrow_path), "--force"] ) assert r.returncode == 0 new_key = _extract_restore_key(r.stdout) assert new_key != restore_key r = _run_script( RECOVER_SCRIPT, {"RESTORE_KEY": new_key}, ["--escrow", str(escrow_path)] ) assert r.returncode == 0 and f"MASTER_KEY={master_key}" in r.stdout r = _run_script( RECOVER_SCRIPT, {"RESTORE_KEY": restore_key}, ["--escrow", str(escrow_path)] ) assert r.returncode == 1 # 旧钥匙已作废