zhifu-platform/scripts/redis_recovery.py

import argparse
import hashlib
import json
import os
import re
import subprocess
import time
from contextlib import contextmanager
from datetime import UTC, datetime
from pathlib import Path
from uuid import uuid4

ROOT = Path(__file__).resolve().parents[1]
BACKUPS = ROOT / ".work" / "redis-backups"
TRANSIENT = ("agent:active:global*", "sse:global", "sse:actor:*", "monitor:worker:*")


def docker(*args, **kwargs):
    return subprocess.run(["docker", *args], check=True, timeout=180, **kwargs)


def cli(container, *args):
    result = docker("exec", container, "redis-cli", "-e", "--json", *args, capture_output=True)
    return json.loads(result.stdout)


def digest(path):
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def backup():
    """从本机 Compose 的 redis-state 获取包含模块数据和绝对到期时间的 RDB。"""
    source = docker("compose", "ps", "-q", "redis-state", cwd=ROOT, capture_output=True)
    source_id = source.stdout.decode().strip()
    if not source_id:
        raise RuntimeError("redis-state 容器未启动")
    metadata = json.loads(docker("inspect", source_id, capture_output=True).stdout)[0]
    directory = BACKUPS / (datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") + "-" + uuid4().hex[:8])
    directory.mkdir(parents=True)
    entry = {
        "format": 1,
        "started_at": datetime.now(UTC).isoformat(),
        "image_id": metadata["Image"],
        "indexes": cli(source_id, "FT._LIST"),
    }
    remote = "/data/recovery-" + uuid4().hex + ".rdb"
    path = directory / "state.rdb"
    try:
        # 使用 Redis 自身的快照传输，包含 RedisJSON 和 Search 模块持久数据。
        docker("exec", source_id, "redis-cli", "--rdb", remote, capture_output=True)
        docker("cp", source_id + ":" + remote, str(path), capture_output=True)
    finally:
        docker("exec", source_id, "rm", "-f", remote, capture_output=True)
    entry.update(
        finished_at=datetime.now(UTC).isoformat(), sha256=digest(path), bytes=path.stat().st_size
    )
    (directory / "manifest.json").write_text(json.dumps(entry, indent=2), encoding="utf-8")
    return directory


def validate(directory):
    directory = directory.resolve(strict=True)
    if not directory.is_relative_to(BACKUPS.resolve()):
        raise ValueError("恢复文件必须位于项目 .work/redis-backups")
    entry = json.loads((directory / "manifest.json").read_text(encoding="utf-8"))
    path = (directory / "state.rdb").resolve(strict=True)
    if entry["format"] != 1 or not re.fullmatch(r"sha256:[0-9a-f]{64}", entry["image_id"]):
        raise ValueError("备份清单格式不符合要求")
    if not path.is_relative_to(directory) or digest(path) != entry["sha256"]:
        raise ValueError("备份文件 SHA256 校验失败")
    if path.stat().st_size != entry["bytes"]:
        raise ValueError("备份文件大小不符合要求")
    return directory, entry


@contextmanager
def recovered(directory):
    """仅创建无网络、无公开端口的独立 Redis；离开上下文时删除本次容器。"""
    directory, entry = validate(directory)
    name = "zhifu-redis-restore-" + uuid4().hex[:12]
    docker(
        "run",
        "-d",
        "--name",
        name,
        "--network",
        "none",
        "--read-only",
        "--memory",
        "1g",
        "--label",
        "zhifu.purpose=redis-recovery-verification",
        "--label",
        "zhifu.recovery-run=" + os.environ.get("ZHIFU_RECOVERY_RUN_ID", "standalone"),
        "--mount",
        f"type=bind,source={directory},target=/backup,readonly",
        entry["image_id"],
        "redis-server",
        "--dir",
        "/backup",
        "--dbfilename",
        "state.rdb",
        "--appendonly",
        "no",
        "--save",
        "",
        capture_output=True,
    )
    try:
        for _ in range(60):
            state = json.loads(docker("inspect", name, capture_output=True).stdout)[0]["State"]
            if not state["Running"]:
                raise RuntimeError("独立 Redis 恢复容器启动失败")
            ready = subprocess.run(
                ["docker", "exec", name, "redis-cli", "-e", "--json", "PING"],
                capture_output=True,
                timeout=5,
            )
            if ready.returncode == 0 and json.loads(ready.stdout) == "PONG":
                break
            time.sleep(0.5)
        else:
            raise TimeoutError("独立 Redis 恢复容器启动超时")
        yield name, entry
    finally:
        docker("rm", "-f", name, capture_output=True)


def clear_transient(container):
    """仅允许清除本工具创建的隔离容器中的旧运行额度、SSE 连接和进程心跳。"""
    metadata = json.loads(docker("inspect", container, capture_output=True).stdout)[0]
    if (
        metadata["HostConfig"]["NetworkMode"] != "none"
        or metadata["Config"]["Labels"].get("zhifu.purpose") != "redis-recovery-verification"
    ):
        raise ValueError("仅允许处理独立恢复容器")
    removed = 0
    for pattern in TRANSIENT:
        cursor = "0"
        while True:
            cursor, keys = cli(container, "SCAN", cursor, "MATCH", pattern, "COUNT", "100")
            if keys:
                removed += cli(container, "UNLINK", *keys)
            if str(cursor) == "0":
                break
            cursor = str(cursor)
    return removed


def verify(directory):
    started = time.monotonic()
    with recovered(directory) as (container, entry):
        indexes = cli(container, "FT._LIST")
        if not set(entry["indexes"]).issubset(indexes):
            raise RuntimeError("恢复后的 Search 索引不完整")
        report = {
            "started_at": datetime.now(UTC).isoformat(),
            "sha256_verified": True,
            "keys_before_cleanup": cli(container, "DBSIZE"),
            "indexes": indexes,
            "transient_removed": clear_transient(container),
            "keys_after_cleanup": cli(container, "DBSIZE"),
        }
    report["seconds"] = round(time.monotonic() - started, 2)
    (directory / "restore-report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    return report


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Redis state 备份与独立容器恢复验证")
    parser.add_argument("action", choices=["backup", "verify", "exercise"])
    parser.add_argument("--directory", type=Path)
    args = parser.parse_args()
    if args.action == "verify":
        if args.directory is None:
            parser.error("verify 需要 --directory")
        print(json.dumps(verify(args.directory), indent=2))
    else:
        destination = backup()
        print(destination)
        if args.action == "exercise":
            print(json.dumps(verify(destination), indent=2))