更新: 359 个文件 - 2026-03-16 23:30:01

scripts/lab/main.py

@@ -0,0 +1,380 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from typing import Any, Dict, List
+
+CURRENT_DIR = Path(__file__).resolve().parent
+SCRIPTS_DIR = CURRENT_DIR.parent
+if str(SCRIPTS_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPTS_DIR))
+
+from lab import attack, baseline, browser, catalog, evidence, provision, render, repro, seed, task_queue, validators  # noqa: E402
+from lab.config import ADVISORIES_DIR, CASE_RUNS_DIR, ENV_PROFILES_DIR, RUNS_DIR  # noqa: E402
+from lab.utils import command_available, ensure_dir, isoformat, load_json_dir, now_utc, read_json, read_yaml, write_json  # noqa: E402
+
+
+def _load_advisory(canonical_id: str) -> Dict[str, Any]:
+    advisory = read_json(ADVISORIES_DIR / f"{canonical_id}.json", default=None)
+    if not advisory:
+        raise ValueError(f"Unknown advisory: {canonical_id}")
+    return advisory
+
+
+def _run_dir(run_id: str) -> Path:
+    path = CASE_RUNS_DIR / run_id
+    ensure_dir(path)
+    ensure_dir(path / "logs")
+    ensure_dir(path / "assets")
+    return path
+
+
+def _compose_run_id(advisory: Dict[str, Any]) -> str:
+    return f"{advisory['system_id']}-{advisory['canonical_id']}-{now_utc().strftime('%Y%m%d%H%M%S')}"
+
+
+def _resolve_profile(advisory: Dict[str, Any]) -> Dict[str, Any]:
+    profile = repro.resolve_profile(advisory["canonical_id"], advisory)
+    current_profile = read_yaml(ENV_PROFILES_DIR / "core" / advisory["system_id"] / "current.yaml", default={}) or {}
+    merged = dict(current_profile)
+    merged.update(profile)
+    if current_profile.get("services") and not merged.get("services"):
+        merged["services"] = current_profile["services"]
+    if current_profile.get("baseline_urls") and not merged.get("baseline_urls"):
+        merged["baseline_urls"] = current_profile["baseline_urls"]
+    if current_profile.get("artifact_mode") and not merged.get("artifact_mode"):
+        merged["artifact_mode"] = current_profile["artifact_mode"]
+    if current_profile.get("verification_mode") and not merged.get("verification_mode"):
+        merged["verification_mode"] = current_profile["verification_mode"]
+    if current_profile.get("browser_required"):
+        merged.setdefault("browser_assertions", {})
+        merged["browser_assertions"].setdefault("required", current_profile["browser_required"])
+    if not profile.get("system_id"):
+        merged["system_id"] = advisory["system_id"]
+    if not profile.get("profile_id"):
+        merged["profile_id"] = advisory["canonical_id"]
+    return merged
+
+
+def _build_run_bundle(
+    advisory: Dict[str, Any],
+    profile: Dict[str, Any],
+    run_id: str,
+    verification_status: str,
+    verification_mode: str,
+    artifact_mode: str,
+    baseline_refs: List[str],
+    attack_steps: List[Dict[str, Any]],
+    browser_refs: List[str],
+    container_log_refs: List[str],
+    request_log_refs: List[str],
+    blocked_reason: str | None,
+) -> Dict[str, Any]:
+    return {
+        "run_id": run_id,
+        "system_id": advisory["system_id"],
+        "advisory_id": advisory["canonical_id"],
+        "repro_profile_id": profile["profile_id"],
+        "verification_status": verification_status,
+        "verification_mode": verification_mode,
+        "artifact_mode": artifact_mode,
+        "target_env": "local-docker",
+        "compose_services": sorted(profile.get("services", {}).keys()),
+        "baseline_refs": baseline_refs,
+        "attack_steps": attack_steps,
+        "browser_refs": browser_refs,
+        "container_log_refs": container_log_refs,
+        "request_log_refs": request_log_refs,
+        "timeline": [],
+        "started_at": isoformat(now_utc()),
+        "finished_at": isoformat(now_utc()),
+        "blocked_reason": blocked_reason,
+    }
+
+
+def cmd_catalog_sync(args) -> int:
+    summary = catalog.sync_catalog(write_profiles=True, write_repro_map=True)
+    print(summary)
+    return 0
+
+
+def cmd_compose_generate(args) -> int:
+    advisory = _load_advisory(args.case)
+    profile = _resolve_profile(advisory)
+    run_dir = _run_dir(args.run_id or f"compose-{advisory['canonical_id']}")
+    compose_result = provision.prepare(profile, run_dir, dry_run=True)
+    print(compose_result)
+    return 0
+
+
+def cmd_provision(args) -> int:
+    advisory = _load_advisory(args.case)
+    profile = _resolve_profile(advisory)
+    run_dir = _run_dir(args.run_id or _compose_run_id(advisory))
+    result = provision.prepare(profile, run_dir, dry_run=args.dry_run)
+    print(result)
+    return 0
+
+
+def cmd_seed(args) -> int:
+    advisory = _load_advisory(args.case)
+    profile = _resolve_profile(advisory)
+    print({"steps": seed.run_seed(profile)})
+    return 0
+
+
+def cmd_baseline(args) -> int:
+    advisory = _load_advisory(args.case)
+    profile = _resolve_profile(advisory)
+    run_dir = _run_dir(args.run_id or _compose_run_id(advisory))
+    result = baseline.collect(profile, run_dir)
+    print(result)
+    return 0
+
+
+def cmd_attack(args) -> int:
+    advisory = _load_advisory(args.case)
+    profile = _resolve_profile(advisory)
+    run_dir = _run_dir(args.run_id or _compose_run_id(advisory))
+    result = attack.run_attack(profile, advisory, run_dir, dry_run=args.dry_run)
+    print(result)
+    return 0
+
+
+def cmd_verify(args) -> int:
+    advisory = _load_advisory(args.case)
+    profile = _resolve_profile(advisory)
+    browser_required = bool(profile.get("browser_assertions", {}).get("required"))
+    payload = {
+        "advisory": advisory["canonical_id"],
+        "profile_id": profile["profile_id"],
+        "browser_required": browser_required,
+        "result": "ready-for-run",
+    }
+    print(payload)
+    return 0
+
+
+def cmd_run_case(args) -> int:
+    advisory = _load_advisory(args.case)
+    profile = _resolve_profile(advisory)
+    run_id = args.run_id or _compose_run_id(advisory)
+    run_dir = _run_dir(run_id)
+
+    provision_result = provision.prepare(profile, run_dir, dry_run=args.dry_run)
+    baseline_payload = baseline.collect(profile, run_dir) if profile.get("baseline_urls") else {"observations": []}
+    attack_payload = attack.run_attack(profile, advisory, run_dir, dry_run=args.dry_run)
+
+    browser_payload = {"required": bool(profile.get("browser_assertions", {}).get("required")), "present": False, "refs": []}
+    blocked_reason = provision_result.get("blocked_reason")
+    if browser_payload["required"] and not args.dry_run and profile.get("baseline_urls"):
+        browser_payload = browser.capture(profile["baseline_urls"][0], run_dir, prefix="proof")
+        if not browser_payload.get("present"):
+            blocked_reason = blocked_reason or browser_payload.get("reason")
+
+    compose_path = Path(provision_result["compose_path"])
+    container_logs = evidence.collect_container_logs(run_dir, compose_path) if compose_path.exists() else []
+
+    verification_status = "triage-manual"
+    verification_mode = profile.get("verification_mode", "synthetic")
+    artifact_mode = profile.get("artifact_mode", profile.get("provisioning_mode", "synthetic"))
+    if args.dry_run:
+        verification_status = "triage-manual"
+        blocked_reason = blocked_reason or "dry-run only"
+    elif provision_result.get("status") == "blocked-artifact":
+        verification_status = "blocked-artifact"
+    elif browser_payload.get("required") and not browser_payload.get("present"):
+        verification_status = "triage-manual"
+    elif artifact_mode == "synthetic":
+        verification_status = "verified-synthetic"
+    else:
+        verification_status = "verified-real"
+
+    bundle = _build_run_bundle(
+        advisory=advisory,
+        profile=profile,
+        run_id=run_id,
+        verification_status=verification_status,
+        verification_mode=verification_mode,
+        artifact_mode=artifact_mode,
+        baseline_refs=[str(run_dir / "logs" / "baseline.json")] if baseline_payload.get("observations") else [],
+        attack_steps=attack_payload.get("steps", []),
+        browser_refs=browser_payload.get("refs", []),
+        container_log_refs=container_logs,
+        request_log_refs=[str(run_dir / "logs" / "attack.json"), str(run_dir / "logs" / "baseline.json")],
+        blocked_reason=blocked_reason,
+    )
+    report_refs = render.render_run(bundle)
+    bundle["report_refs"] = report_refs
+    evidence.write_run_bundle(run_dir, bundle)
+    ensure_dir(RUNS_DIR)
+    write_json(RUNS_DIR / f"{run_id}.json", bundle)
+    render.render_dashboard()
+    print(bundle)
+    return 0
+
+
+def cmd_run_system(args) -> int:
+    advisories = [item for item in load_json_dir(ADVISORIES_DIR) if item.get("system_id") == args.system]
+    selected = advisories[: args.limit]
+    for advisory in selected:
+        cmd_run_case(argparse.Namespace(case=advisory["canonical_id"], run_id=None, dry_run=args.dry_run))
+    print({"system": args.system, "count": len(selected)})
+    return 0
+
+
+def cmd_run_batch(args) -> int:
+    if args.from_queue:
+        items = task_queue.dequeue(limit=args.limit)
+    else:
+        task_queue.enqueue_from_registry(only_hotlane=args.only_hotlane, limit=args.limit)
+        items = task_queue.dequeue(limit=args.limit)
+    for item in items:
+        cmd_run_case(argparse.Namespace(case=item["advisory_id"], run_id=None, dry_run=args.dry_run))
+    print({"processed": len(items)})
+    return 0
+
+
+def cmd_render_run(args) -> int:
+    run = read_json(RUNS_DIR / f"{args.run_id}.json", default=None)
+    if not run:
+        raise ValueError(f"Unknown run: {args.run_id}")
+    print(render.render_run(run))
+    return 0
+
+
+def cmd_serve_dashboard(args) -> int:
+    render.render_dashboard()
+    import http.server
+    import socketserver
+
+    os_dir = str(render.DASHBOARD_DIR if hasattr(render, "DASHBOARD_DIR") else "")
+    if not os_dir:
+        from lab.config import DASHBOARD_DIR
+
+        os_dir = str(DASHBOARD_DIR)
+    handler = http.server.SimpleHTTPRequestHandler
+    with socketserver.TCPServer(("127.0.0.1", args.port), handler) as httpd:
+        print(f"serving dashboard at http://127.0.0.1:{args.port}/")
+        import os
+
+        os.chdir(os_dir)
+        httpd.serve_forever()
+
+
+def cmd_cleanup(args) -> int:
+    run = read_json(RUNS_DIR / f"{args.run_id}.json", default=None)
+    if not run:
+        raise ValueError(f"Unknown run: {args.run_id}")
+    compose_path = Path(run["report_refs"]["bundle_dir"]) / "compose" / "compose.yaml"
+    if command_available("docker") and compose_path.exists():
+        from lab.utils import run as shell_run
+
+        shell_run(["docker", "compose", "-f", str(compose_path), "down", "-v"], cwd=compose_path.parent.parent)
+    print({"cleaned": args.run_id})
+    return 0
+
+
+def cmd_retry_failures(args) -> int:
+    failed = [
+        item
+        for item in load_json_dir(RUNS_DIR)
+        if item.get("verification_status") in {"blocked-artifact", "triage-manual"}
+    ]
+    task_queue.enqueue_items(
+        [{"advisory_id": item["advisory_id"], "system_id": item["system_id"], "priority": "retry"} for item in failed[: args.limit]]
+    )
+    print({"requeued": min(len(failed), args.limit)})
+    return 0
+
+
+def cmd_validate(args) -> int:
+    errors = validators.validate_assets()
+    if errors:
+        print("Validation failed:")
+        for error in errors:
+            print(f"- {error}")
+        return 1
+    print("Validation passed.")
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Websafe local lab orchestrator")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    catalog_sync = subparsers.add_parser("catalog", help="catalog operations")
+    catalog_sub = catalog_sync.add_subparsers(dest="catalog_command", required=True)
+    catalog_sync_cmd = catalog_sub.add_parser("sync", help="sync environment catalog and repro map")
+    catalog_sync_cmd.set_defaults(func=cmd_catalog_sync)
+
+    compose_generate = subparsers.add_parser("compose", help="compose operations")
+    compose_sub = compose_generate.add_subparsers(dest="compose_command", required=True)
+    compose_generate_cmd = compose_sub.add_parser("generate", help="generate compose file for a case")
+    compose_generate_cmd.add_argument("--case", required=True)
+    compose_generate_cmd.add_argument("--run-id")
+    compose_generate_cmd.set_defaults(func=cmd_compose_generate)
+
+    for name, func in [
+        ("provision", cmd_provision),
+        ("seed", cmd_seed),
+        ("baseline", cmd_baseline),
+        ("attack", cmd_attack),
+        ("verify", cmd_verify),
+    ]:
+        sub = subparsers.add_parser(name)
+        sub.add_argument("--case", required=True)
+        sub.add_argument("--run-id")
+        sub.add_argument("--dry-run", action="store_true")
+        sub.set_defaults(func=func)
+
+    run_case = subparsers.add_parser("run-case", help="run a single advisory through the lab pipeline")
+    run_case.add_argument("--case", required=True)
+    run_case.add_argument("--run-id")
+    run_case.add_argument("--dry-run", action="store_true")
+    run_case.set_defaults(func=cmd_run_case)
+
+    run_system = subparsers.add_parser("run-system", help="run the first N advisories for a system")
+    run_system.add_argument("--system", required=True)
+    run_system.add_argument("--limit", type=int, default=5)
+    run_system.add_argument("--dry-run", action="store_true")
+    run_system.set_defaults(func=cmd_run_system)
+
+    run_batch = subparsers.add_parser("run-batch", help="process repro queue or enqueue from registry")
+    run_batch.add_argument("--limit", type=int, default=10)
+    run_batch.add_argument("--only-hotlane", action="store_true")
+    run_batch.add_argument("--from-queue", action="store_true")
+    run_batch.add_argument("--dry-run", action="store_true")
+    run_batch.set_defaults(func=cmd_run_batch)
+
+    render_run = subparsers.add_parser("render-run", help="re-render a stored run")
+    render_run.add_argument("--run-id", required=True)
+    render_run.set_defaults(func=cmd_render_run)
+
+    serve = subparsers.add_parser("serve-dashboard", help="serve the static dashboard locally")
+    serve.add_argument("--port", type=int, default=8734)
+    serve.set_defaults(func=cmd_serve_dashboard)
+
+    cleanup = subparsers.add_parser("cleanup", help="tear down a stored run compose environment")
+    cleanup.add_argument("--run-id", required=True)
+    cleanup.set_defaults(func=cmd_cleanup)
+
+    retry = subparsers.add_parser("retry-failures", help="requeue blocked or manual runs")
+    retry.add_argument("--limit", type=int, default=50)
+    retry.set_defaults(func=cmd_retry_failures)
+
+    validate = subparsers.add_parser("validate", help="validate lab assets")
+    validate.set_defaults(func=cmd_validate)
+    return parser
+
+
+def main() -> int:
+    parser = build_parser()
+    args = parser.parse_args()
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())