vps-manager/tests/test_escrow.py

"""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  # 旧钥匙已作废