veltro-cli 0.15.7__py3-none-any.whl

veltro_cli/lib/demo.py

@@ -0,0 +1,421 @@
+"""`veltro demo list|seed|reset` — the demo sandbox (§7.3).
+
+`seed` provisions and replays one or more datasets from `content/datasets/`:
+
+* the `brute-force` dataset is the golden path and is provisioned by
+  `scripts/seed-demo.sh`, which also builds the VectorFlow pipeline, the CHAD
+  push data source and the Warden bindings the other datasets replay through;
+* every other dataset is provisioned and replayed by
+  `veltro_cli.lib.demo_replay` against the running suite.
+
+`reset` undoes what is reversible per dataset and names what is retained (see
+`demo_replay.reset_dataset`), then clears the suite demo-mode banner.
+
+Both write `.demo-datasets.json` in the suite directory (mode 0600): the
+per-dataset run token, rule, alert ids and case keys the parameterised
+`e2e/demo-datasets.spec.ts` follows. `seed` claims every dataset it is about
+to touch — in that handoff and in identity demo mode — BEFORE its first write,
+and rewrites each entry with the real identifiers as the dataset completes, so
+a run that dies half-way still shows the synthetic-data banner and is still
+undone by `veltro demo reset`. That covers the golden path, whose shell seeder
+records its identifiers only in its last step: `reset` resolves its rule by
+the title `content/datasets/brute-force/dataset.json` carries, then reaches
+that rule's alerts and Warden cases through it.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shlex
+import subprocess
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable, Mapping, Sequence
+
+from veltro_cli.lib.datasets import (
+    GOLDEN_PATH_ENGINE,
+    Dataset,
+    load_catalog,
+    resolve_datasets,
+)
+from veltro_cli.lib.demo_replay import (
+    INGEST_SUFFIX_DEFAULT,
+    Clock,
+    clear_demo_mode,
+    read_demo_mode,
+    reset_dataset,
+    resolve_ingest,
+    reset_golden_path,
+    seed_dataset,
+    set_demo_mode,
+)
+from veltro_cli.lib.errors import VeltroCliError
+from veltro_cli.lib.paths import enrollment_token_dir, resolve_dir
+from veltro_cli.lib.suite_http import SuiteSession
+from veltro_cli.lib.toml_config import VELTRO_TOML, read_config
+
+ACTIONS = ("list", "seed", "reset")
+HANDOFF_NAME = ".demo-datasets.json"
+HANDOFF_SCHEMA = "veltro-demo-handoff/v1"
+DEFAULT_DATASET = "brute-force"
+REPLAY_SUBDIR = "demo"
+SEED_HANDOFF_NAME = ".seed-demo.env"
+GOLDEN_HANDOFF_FIELDS = {
+    "CHAD_RULE_ID",
+    "CHAD_ALERT_ID",
+    "CHAD_IOC_ALERT_ID",
+    "CHAD_INGEST_SUFFIX",
+    "CHAD_NOVEL_ALERT_ID",
+    "WARDEN_CASE_KEY",
+}
+
+
+def _run_token(now: datetime) -> str:
+    return f"{int(now.timestamp())}-{os.getpid()}"
+
+
+def public_url(root: Path, env: dict[str, str]) -> str:
+    """The suite origin: PUBLIC_BASE_URL, else veltro.toml, else localhost."""
+    from_env = env.get("PUBLIC_BASE_URL", "").strip()
+    if from_env:
+        return from_env
+    config = read_config(root / VELTRO_TOML)
+    if config and config.public_url:
+        return config.public_url
+    return "http://localhost"
+
+
+def replay_dir(root: Path, env: Mapping[str, str]) -> Path:
+    """Where corpus JSONL is queued for the VectorFlow agent's file source.
+
+    The agent reads `/run/vf-agent-enrollment/demo/*.jsonl`, and compose binds
+    `VF_ENROLLMENT_TOKEN_DIR` (default `<suite>/.suite/vf-agent`) onto that
+    target. `.env.example` says that variable is the ONLY name for the
+    directory, so this resolves exactly the way `scripts/seed-demo.sh` does:
+    `VELTRO_DEMO_REPLAY_DIR` overrides the whole path, else
+    `VF_ENROLLMENT_TOKEN_DIR`, else the in-tree default — with the same `demo/`
+    subdirectory the deployed source globs. Hardcoding the default instead put
+    the corpus somewhere the relocated agent could not see, and the seed then
+    failed 90s later blaming the rule and the field mappings (veltro#984).
+    """
+    explicit = env.get("VELTRO_DEMO_REPLAY_DIR", "").strip()
+    if explicit:
+        return Path(explicit).expanduser().resolve()
+    base = env.get("VF_ENROLLMENT_TOKEN_DIR", "").strip()
+    root_dir = Path(base).expanduser().resolve() if base else enrollment_token_dir(root)
+    return root_dir / REPLAY_SUBDIR
+
+
+def _admin_credentials(env: dict[str, str]) -> tuple[str, str]:
+    email = env.get("VF_ADMIN_EMAIL", "").strip()
+    password = env.get("VF_ADMIN_PASSWORD", "")
+    if not email or not password:
+        raise VeltroCliError(
+            "VF_ADMIN_EMAIL and VF_ADMIN_PASSWORD are required to seed a dataset",
+            remedy="export a suite admin login; the same pair scripts/seed-demo.sh uses",
+        )
+    return email, password
+
+
+def write_handoff(root: Path, payload: dict[str, Any]) -> Path:
+    path = root / HANDOFF_NAME
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+    path.chmod(0o600)
+    return path
+
+
+def read_handoff(root: Path) -> dict[str, Any]:
+    path = root / HANDOFF_NAME
+    if not path.is_file():
+        return {}
+    try:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+    except json.JSONDecodeError:
+        return {}
+    return payload if isinstance(payload, dict) else {}
+
+def read_golden_path_handoff(root: Path) -> dict[str, str]:
+    """Read only non-secret identifiers from the shell handoff."""
+    path = root / SEED_HANDOFF_NAME
+    if not path.is_file():
+        return {}
+    values: dict[str, str] = {}
+    for line in path.read_text(encoding="utf-8").splitlines():
+        key, separator, encoded = line.partition("=")
+        if separator != "=" or key not in GOLDEN_HANDOFF_FIELDS:
+            continue
+        parsed = shlex.split(encoded)
+        if len(parsed) == 1 and parsed[0]:
+            values[key] = parsed[0]
+    return values
+
+
+def _golden_path_argv(root: Path, extra_args: Sequence[str]) -> list[str]:
+    script = root / "scripts" / "seed-demo.sh"
+    if not script.is_file():
+        raise VeltroCliError(
+            f"{script} not found",
+            remedy="run from an installed suite directory or a repository checkout",
+        )
+    return ["bash", str(script), *extra_args]
+
+
+def _seed_golden_path(
+    root: Path,
+    extra_args: Sequence[str],
+    runner: Callable[..., Any] | None,
+    env: dict[str, str],
+) -> int:
+    argv = _golden_path_argv(root, extra_args)
+    if runner is not None:
+        result = runner(argv, env, root)
+        return int(getattr(result, "returncode", result))
+    completed = subprocess.run(argv, cwd=str(root), check=False)
+    return int(completed.returncode)
+
+
+def format_catalog(catalog: dict[str, Dataset]) -> str:
+    lines = ["datasets (content/datasets):"]
+    for dataset in sorted(catalog.values(), key=lambda d: (d.engine != GOLDEN_PATH_ENGINE, d.id)):
+        techniques = ", ".join(dataset.attack_techniques)
+        lines.append(f"  {dataset.id:<22} {dataset.title}")
+        lines.append(f"  {'':<22} engine={dataset.engine} attack={techniques}")
+    lines.append("")
+    lines.append("seed one:   veltro demo seed --dataset <id>")
+    lines.append("seed all:   veltro demo seed --dataset all")
+    lines.append("undo:       veltro demo reset [--dataset <id>]")
+    return "\n".join(lines) + "\n"
+
+
+def run_demo(
+    action: str,
+    *,
+    directory: str | Path | None = None,
+    datasets: Sequence[str] = (),
+    extra_args: Sequence[str] = (),
+    runner: Callable[..., Any] | None = None,
+    session: SuiteSession | None = None,
+    clock: Clock | None = None,
+    env: dict[str, str] | None = None,
+    out: Callable[[str], None] = print,
+) -> int:
+    if action not in ACTIONS:
+        raise VeltroCliError(
+            f"unknown demo action {action!r}",
+            remedy="use `veltro demo list`, `veltro demo seed` or `veltro demo reset`",
+        )
+    root = resolve_dir(directory)
+    environment = dict(os.environ if env is None else env)
+    clock = clock or Clock()
+
+    if action == "list":
+        out(format_catalog(load_catalog(root)).rstrip("\n"))
+        return 0
+    if action == "seed":
+        return _seed(root, datasets, extra_args, runner, session, clock, environment, out)
+    return _reset(root, datasets, session, environment, out)
+
+
+def _open_session(root: Path, environment: dict[str, str], session: SuiteSession | None) -> SuiteSession:
+    if session is not None:
+        return session
+    live = SuiteSession(base_url=public_url(root, environment))
+    live.login(*_admin_credentials(environment))
+    return live
+
+
+def _persist(
+    root: Path,
+    live: SuiteSession,
+    results: dict[str, Any],
+    seeded_at: datetime,
+) -> Path:
+    """Record what is (or is being) seeded, in both places `reset` reads.
+
+    Called before the first mutation and again after every dataset, because a
+    seed that dies half-way used to leave a deployed rule, an alert, a Warden
+    case and a queued JSONL file behind with demo mode OFF: no product banner
+    said the data was synthetic, and `veltro demo reset` answered "nothing to
+    reset" (veltro#984). The marker is what makes a partial seed visible and
+    undoable, so it is written first and kept current, not written last.
+    """
+    set_demo_mode(live, results.keys(), seeded_at=seeded_at)
+    return write_handoff(
+        root,
+        {
+            "schema": HANDOFF_SCHEMA,
+            "seeded_at": seeded_at.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"),
+            "datasets": results,
+        },
+    )
+
+
+def _seed(
+    root: Path,
+    names: Sequence[str],
+    extra_args: Sequence[str],
+    runner: Callable[..., Any] | None,
+    session: SuiteSession | None,
+    clock: Clock,
+    environment: dict[str, str],
+    out: Callable[[str], None],
+) -> int:
+    catalog = load_catalog(root)
+    selected = resolve_datasets(root, names or (DEFAULT_DATASET,))
+    corpus = [dataset for dataset in selected if dataset.is_corpus]
+    golden = [dataset for dataset in selected if not dataset.is_corpus]
+    if corpus and not golden:
+        golden = [catalog[DEFAULT_DATASET]]
+        out("== brute-force: preparing the VectorFlow → CHAD → Warden replay path ==")
+    queue = replay_dir(root, environment)
+
+    # Open the session BEFORE anything is mutated: it validates the admin
+    # credentials and it is what records the demo-mode marker the whole run
+    # depends on. Failing here costs nothing; failing later leaves state.
+    live = _open_session(root, environment, session)
+    seeded_at = clock.now()
+    results: dict[str, Any] = dict(read_handoff(root).get("datasets") or {})
+    for dataset in (*golden, *corpus):
+        results.setdefault(
+            dataset.id,
+            {
+                "dataset": dataset.id,
+                "engine": dataset.engine,
+                "status": "seeding",
+                "notes": [
+                    "claimed before the first write; `veltro demo reset --dataset "
+                    f"{dataset.id}` undoes whatever this run created"
+                ],
+            },
+        )
+    _persist(root, live, results, seeded_at)
+    out("   demo mode on: every product shows the demo banner while this seed runs")
+
+    for dataset in golden:
+        out(f"== {dataset.id}: golden-path seeder (scripts/seed-demo.sh) ==")
+        code = _seed_golden_path(root, extra_args, runner, environment)
+        if code != 0:
+            raise VeltroCliError(
+                f"the golden-path seeder failed (exit {code})",
+                remedy="read its FAIL line above; every step names what to fix",
+            )
+        golden_handoff = read_golden_path_handoff(root)
+        results[dataset.id] = {
+            "dataset": dataset.id,
+            "engine": dataset.engine,
+            "rule_id": golden_handoff.get("CHAD_RULE_ID", ""),
+            "ingest_suffix": golden_handoff.get("CHAD_INGEST_SUFFIX", INGEST_SUFFIX_DEFAULT),
+            "alert_ids": [
+                identifier
+                for key in ("CHAD_ALERT_ID", "CHAD_IOC_ALERT_ID", "CHAD_NOVEL_ALERT_ID")
+                if (identifier := golden_handoff.get(key))
+            ],
+            "case_keys": (
+                [golden_handoff["WARDEN_CASE_KEY"]]
+                if golden_handoff.get("WARDEN_CASE_KEY")
+                else []
+            ),
+            "notes": ["provisioned by scripts/seed-demo.sh; see .seed-demo.env for its private handoff"],
+        }
+        _persist(root, live, results, seeded_at)
+
+    if corpus:
+        ingest = resolve_ingest(live)
+        for dataset in corpus:
+            token = _run_token(clock.now())
+            out(f"== {dataset.id}: replaying {dataset.title} ==")
+            result = seed_dataset(
+                live,
+                dataset,
+                run_token=token,
+                replay_dir=queue,
+                clock=clock,
+                ingest=ingest,
+            )
+            results[dataset.id] = result.as_dict()
+            _persist(root, live, results, seeded_at)
+            out(
+                f"   {result.events_replayed} events queued through VectorFlow, "
+                f"{len(result.alert_ids)} CHAD alert(s), Warden case(s): "
+                f"{', '.join(result.case_keys) or 'none'}"
+            )
+
+    path = _persist(root, live, results, seeded_at)
+    out(f"handoff: {path}")
+    return 0
+
+
+def _reset(
+    root: Path,
+    names: Sequence[str],
+    session: SuiteSession | None,
+    environment: dict[str, str],
+    out: Callable[[str], None],
+) -> int:
+    catalog = load_catalog(root)
+    handoff = read_handoff(root)
+    recorded = list((handoff.get("datasets") or {}).keys())
+    live = _open_session(root, environment, session)
+    if not names:
+        state = read_demo_mode(live) or {}
+        from_identity = [str(identifier) for identifier in state.get("datasets") or []]
+        names = [name for name in dict.fromkeys([*recorded, *from_identity]) if name in catalog]
+        if not names:
+            out("nothing to reset: no seeded dataset recorded in the handoff or in identity demo mode")
+            clear_demo_mode(live)
+            return 0
+
+    selected = resolve_datasets(root, names)
+    queue = replay_dir(root, environment)
+    remaining = dict(handoff.get("datasets") or {})
+    for dataset in selected:
+        recorded_result = remaining.get(dataset.id) or {}
+        case_keys = [str(key) for key in recorded_result.get("case_keys") or []]
+        if dataset.is_corpus:
+            result = reset_dataset(
+                live,
+                dataset,
+                replay_dir=queue,
+                case_keys=case_keys,
+            )
+        else:
+            result = reset_golden_path(
+                live,
+                replay_dir=queue,
+                rule_title=dataset.rule_title,
+                rule_id=str(recorded_result.get("rule_id") or ""),
+                alert_ids=[str(identifier) for identifier in recorded_result.get("alert_ids") or []],
+                case_keys=case_keys,
+                ingest_suffix=str(recorded_result.get("ingest_suffix") or INGEST_SUFFIX_DEFAULT),
+            )
+        remaining.pop(dataset.id, None)
+        out(f"== {dataset.id} ==")
+        if result.replay_files_deleted:
+            out(f"   replay files deleted: {result.replay_files_deleted}")
+        if result.rules_undeployed:
+            out(f"   undeployed rule(s): {', '.join(result.rules_undeployed)}")
+        out(f"   alerts deleted: {result.alerts_deleted}")
+        if result.cases_closed:
+            out(f"   cases closed: {', '.join(result.cases_closed)}")
+        for note in result.notes:
+            out(f"   note: {note}")
+        for retained in result.retained:
+            out(f"   retained: {retained}")
+
+    if remaining:
+        write_handoff(
+            root,
+            {
+                "schema": HANDOFF_SCHEMA,
+                "seeded_at": handoff.get("seeded_at", ""),
+                "datasets": remaining,
+            },
+        )
+        set_demo_mode(live, remaining.keys(), seeded_at=datetime.now(timezone.utc))
+        out(f"still seeded: {', '.join(sorted(remaining))} (demo mode stays on)")
+    else:
+        (root / HANDOFF_NAME).unlink(missing_ok=True)
+        clear_demo_mode(live)
+        out("demo mode off: the demo banner is gone from every product")
+    return 0