muse-agent-example/.claude/skills/runtime/scripts/test_raw_vault.py

#!/usr/bin/env python3
"""raw vault 的 lease 顺序、权限、路径边界与恢复测试。"""

from __future__ import annotations

import json
import os
import pathlib
import stat
import sys
import tempfile
import unittest
from datetime import datetime, timedelta, timezone

SCRIPT_DIR = pathlib.Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR))

from raw_vault import RawVaultError, RawVaultManager  # noqa: E402


def future_time(hours: int = 1) -> str:
    """生成不超过 24 小时的 UTC 保留截止时间。"""

    return (datetime.now(timezone.utc) + timedelta(hours=hours)).isoformat()


class RawVaultTest(unittest.TestCase):
    """验证敏感字节只在有效 lease 覆盖下存在。"""

    def setUp(self) -> None:
        """为每个测试创建独立 vault 根和非敏感 journal。"""

        self.temporary = tempfile.TemporaryDirectory()
        root = pathlib.Path(self.temporary.name)
        self.vault_root = root / "private-tmp"
        self.journal_root = root / "journal"
        self.vault_root.mkdir(mode=0o700)
        self.manager = RawVaultManager(self.journal_root, vault_root=self.vault_root)

    def tearDown(self) -> None:
        """删除测试产生的独立目录。"""

        self.temporary.cleanup()

    def _create(self):
        """创建字段完整、审批角色合法的单次 lease。"""

        return self.manager.create_vault(
            authorization_id="auth-1",
            approved_by="创始人",
            run_id="run-1",
            source_version="source-v1",
            content_hashes=("sha256:" + "a" * 64,),
            purpose="offline_evaluation",
            retain_until=future_time(),
        )

    def test_lease_precedes_vault_and_permissions_are_fixed(self):
        """journal 必须先存在，随后 vault=0700、raw 文件=0600。"""

        lease = self._create()
        lease_file = self.journal_root / "leases" / f"{lease.vault_id}.json"
        self.assertTrue(lease_file.exists())
        self.assertEqual(stat.S_IMODE(lease_file.stat().st_mode), 0o600)
        lease_json = json.loads(lease_file.read_text(encoding="utf-8"))
        self.assertEqual(lease_json["status"], "open")
        self.assertNotIn("vaultPath", lease_json)

        vaults = list(self.vault_root.iterdir())
        self.assertEqual(len(vaults), 1)
        self.assertEqual(stat.S_IMODE(vaults[0].stat().st_mode), 0o700)
        self.manager.write_bytes(lease, "responses/model.json", b"raw-secret")
        raw_file = vaults[0] / "responses" / "model.json"
        self.assertEqual(raw_file.read_bytes(), b"raw-secret")
        self.assertEqual(stat.S_IMODE(raw_file.stat().st_mode), 0o600)

    def test_path_escape_absolute_path_and_symlink_are_rejected(self):
        """调用方不能用绝对路径、.. 或软链接把 raw 写到 vault 外。"""

        lease = self._create()
        for candidate in ("../escape", "/tmp/escape", "nested/../../escape"):
            with self.subTest(candidate=candidate), self.assertRaises(RawVaultError) as raised:
                self.manager.write_bytes(lease, candidate, b"secret")
            self.assertEqual(raised.exception.code, "RAW_PATH_INVALID")

        vault = next(self.vault_root.iterdir())
        outside = pathlib.Path(self.temporary.name) / "outside"
        outside.write_bytes(b"unchanged")
        os.symlink(outside, vault / "link")
        with self.assertRaises(RawVaultError) as raised:
            self.manager.write_bytes(lease, "link", b"changed")
        self.assertEqual(raised.exception.code, "RAW_PATH_INVALID")
        self.assertEqual(outside.read_bytes(), b"unchanged")

    def test_cleanup_closes_lease_without_disclosing_raw_path(self):
        """清理成功后 vault 消失、lease 关闭，回执不含本机 raw 路径。"""

        lease = self._create()
        self.manager.write_bytes(lease, "input.json", b"secret")
        receipt = self.manager.cleanup(lease)

        self.assertEqual(list(self.vault_root.iterdir()), [])
        self.assertEqual(receipt["status"], "closed")
        self.assertNotIn(str(self.vault_root), json.dumps(receipt, ensure_ascii=False))
        lease_json = json.loads(
            (self.journal_root / "leases" / f"{lease.vault_id}.json").read_text(encoding="utf-8")
        )
        self.assertEqual(lease_json["status"], "closed")

    def test_recovery_cleans_open_vault_and_closes_missing_lease(self):
        """恢复扫描必须清理未关闭 vault，并补记已经缺失的 lease。"""

        existing = self._create()
        self.manager.write_bytes(existing, "response.json", b"secret")

        missing = self.manager.create_vault(
            authorization_id="auth-2",
            approved_by="用户",
            run_id="run-2",
            source_version="source-v2",
            content_hashes=("sha256:" + "b" * 64,),
            purpose="offline_evaluation",
            retain_until=future_time(),
        )
        missing_dir = [path for path in self.vault_root.iterdir() if missing.vault_id in path.name][0]
        missing_dir.rmdir()

        recovery = self.manager.recover()

        self.assertIn(existing.vault_id, recovery["cleanedVaultIds"])
        self.assertIn(missing.vault_id, recovery["closedMissingVaultIds"])
        self.assertEqual(list(self.vault_root.iterdir()), [])

    def test_recovery_removes_orphan_in_manager_namespace(self):
        """同一管理器命名空间内没有 lease 的 raw 目录不能遗留。"""

        orphan = self.vault_root / f"{self.manager.vault_prefix}orphan"
        orphan.mkdir(mode=0o700)
        (orphan / "raw").write_bytes(b"secret")
        os.chmod(orphan / "raw", 0o600)

        recovery = self.manager.recover()

        self.assertIn("orphan", recovery["cleanedOrphanIds"])
        self.assertFalse(orphan.exists())

    def test_recovery_removes_raw_directory_left_after_closed_lease(self):
        """lease 已关闭后若同名 raw 目录再次出现，恢复扫描仍必须删除。"""

        lease = self._create()
        vault_name = next(self.vault_root.iterdir()).name
        self.manager.cleanup(lease)
        stale = self.vault_root / vault_name
        stale.mkdir(mode=0o700)
        (stale / "raw").write_bytes(b"secret")
        os.chmod(stale / "raw", 0o600)

        recovery = self.manager.recover()

        self.assertIn(lease.vault_id, recovery["cleanedVaultIds"])
        self.assertFalse(stale.exists())


if __name__ == "__main__":
    unittest.main()