veltro-cli 0.15.7__py3-none-any.whl

veltro_cli/lib/datasets.py

@@ -0,0 +1,611 @@
+"""Demo dataset catalogue: discovery, validation, selection and replay shaping.
+
+The corpora live in ``content/datasets/<id>/`` (Apache-2.0) and are described by
+a ``veltro-demo-dataset/v1`` manifest. This module is the only reader: the CLI
+(`veltro demo seed --dataset …`), the replay engine
+(`veltro_cli.lib.demo_replay`) and the CI validator
+(`scripts/test_demo_datasets.py`) all go through it, so a corpus cannot ship a
+manifest that one consumer accepts and another rejects.
+
+Everything here is pure: no network, no clock except what the caller passes in.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from typing import Any, Iterable, Mapping, Sequence
+
+from veltro_cli.lib.errors import VeltroCliError
+
+MANIFEST_SCHEMA = "veltro-demo-dataset/v1"
+MANIFEST_NAME = "dataset.json"
+CONTENT_SUBDIR = ("content", "datasets")
+GOLDEN_PATH_ENGINE = "golden-path-seed"
+CORPUS_ENGINE = "corpus"
+ENGINES = (GOLDEN_PATH_ENGINE, CORPUS_ENGINE)
+REQUIRED_LICENSE = "Apache-2.0"
+APACHE_SPDX = "SPDX-License-Identifier: Apache-2.0"
+ID_PATTERN = re.compile(r"^[a-z][a-z0-9-]{1,40}$")
+# The label every replayed event carries, so `veltro demo reset` can find demo
+# data again and an operator can tell it apart from their own telemetry.
+DEMO_LABEL_FIELD = "veltro_demo"
+DEMO_TAG = "veltro-demo"
+ALL_KEYWORD = "all"
+
+
+class DatasetError(VeltroCliError):
+    """A dataset directory does not satisfy the manifest contract."""
+
+
+@dataclass(frozen=True)
+class Threshold:
+    count: int
+    window_minutes: int
+    group_by: str
+
+
+@dataclass(frozen=True)
+class Detection:
+    rule_title: str
+    description: str
+    severity: str
+    yaml_path: Path
+    field_mappings: tuple[tuple[str, str], ...]
+    threshold: Threshold | None
+    sample_event_index: int
+
+    def yaml_content(self) -> str:
+        return self.yaml_path.read_text(encoding="utf-8")
+
+
+@dataclass(frozen=True)
+class Dataset:
+    id: str
+    title: str
+    summary: str
+    engine: str
+    directory: Path
+    attack_techniques: tuple[str, ...]
+    events_path: Path | None
+    detection: Detection | None
+    # The CHAD rule title this dataset's detection carries. Taken from
+    # `detection.rule_title` for a corpus dataset and from the manifest's
+    # `golden_path.rule_title` for the golden path, whose rule the shell
+    # seeder owns: `reset` resolves the rule by title when a run died before
+    # it could record the id (veltro#984), so every dataset has one.
+    rule_title: str
+    min_alerts: int
+    expect_warden_case: bool
+    matching_event_ids: tuple[str, ...]
+
+    @property
+    def is_corpus(self) -> bool:
+        return self.engine == CORPUS_ENGINE
+
+    def events(self) -> list[dict[str, Any]]:
+        """The corpus, in authored order. Empty for the golden-path engine."""
+        if self.events_path is None:
+            return []
+        return [
+            json.loads(line)
+            for line in self.events_path.read_text(encoding="utf-8").splitlines()
+            if line.strip()
+        ]
+
+
+def _require(condition: bool, message: str, remedy: str) -> None:
+    if not condition:
+        raise DatasetError(message, remedy=remedy)
+
+
+def _string(mapping: Mapping[str, Any], key: str, where: str) -> str:
+    value = mapping.get(key)
+    _require(
+        isinstance(value, str) and value.strip() != "",
+        f"{where}: {key!r} must be a non-empty string",
+        remedy=f"set {key!r} in {where}",
+    )
+    return str(value).strip()
+
+
+def _event_id(event: Mapping[str, Any]) -> str:
+    event_block = event.get("event")
+    if isinstance(event_block, Mapping):
+        candidate = event_block.get("id")
+        if isinstance(candidate, str) and candidate.strip():
+            return candidate.strip()
+    return ""
+
+
+def field_value(event: Mapping[str, Any], dotted: str) -> Any:
+    """Read a dotted ECS path out of a nested event document."""
+    cursor: Any = event
+    for part in dotted.split("."):
+        if not isinstance(cursor, Mapping) or part not in cursor:
+            return None
+        cursor = cursor[part]
+    return cursor
+
+
+def parse_simple_detection(rule_yaml: str) -> dict[str, str]:
+    """Return the ``selection`` field/value pairs of a shipped demo rule.
+
+    Demo rules are deliberately restricted to one flat ``selection`` mapping
+    with ``condition: selection`` — the shape a stdlib reader can parse exactly,
+    with no YAML dependency and no partial interpretation. CHAD's pySigma engine
+    remains the real evaluator; this parser exists so CI can prove a shipped
+    rule matches the events its manifest claims, and it refuses anything richer
+    rather than guessing.
+    """
+    lines = rule_yaml.splitlines()
+    try:
+        start = next(i for i, line in enumerate(lines) if line.rstrip() == "detection:")
+    except StopIteration as exc:
+        raise DatasetError(
+            "rule has no `detection:` block",
+            remedy="a demo rule needs a detection block with one flat selection",
+        ) from exc
+    selection: dict[str, str] = {}
+    condition = ""
+    in_selection = False
+    for line in lines[start + 1 :]:
+        if line.strip() == "" or line.lstrip().startswith("#"):
+            continue
+        indent = len(line) - len(line.lstrip())
+        if indent == 0:
+            break
+        stripped = line.strip()
+        if indent == 2:
+            in_selection = stripped == "selection:"
+            if stripped.startswith("condition:"):
+                condition = stripped.split(":", 1)[1].strip()
+            elif not in_selection:
+                raise DatasetError(
+                    f"unsupported detection key {stripped!r}",
+                    remedy="demo rules support exactly `selection:` and `condition: selection`",
+                )
+            continue
+        if indent == 4 and in_selection:
+            _require(
+                ":" in stripped,
+                f"unsupported selection entry {stripped!r}",
+                remedy="each selection entry must be `field: value`",
+            )
+            key, _, value = stripped.partition(":")
+            _require(
+                "|" not in key,
+                f"selection modifier in {key!r} is not supported by the demo validator",
+                remedy="keep demo rule selections to plain equality so CI can verify them",
+            )
+            selection[key.strip()] = value.strip().strip("'\"")
+            continue
+        raise DatasetError(
+            f"unsupported detection line {line!r}",
+            remedy="demo rules support exactly one flat selection mapping",
+        )
+    _require(
+        condition == "selection",
+        "demo rules must use `condition: selection`",
+        remedy="simplify the rule condition, or move the rule out of content/datasets",
+    )
+    _require(
+        bool(selection),
+        "demo rule selection is empty",
+        remedy="add at least one `field: value` entry under selection",
+    )
+    return selection
+
+
+def selection_matches(selection: Mapping[str, str], event: Mapping[str, Any]) -> bool:
+    """True when every selection field equals its value on this event."""
+    return all(str(field_value(event, field)) == value for field, value in selection.items())
+
+
+def _load_detection(directory: Path, block: Mapping[str, Any], event_count: int) -> Detection:
+    where = f"{directory.name}/dataset.json detection"
+    yaml_name = _string(block, "yaml", where)
+    yaml_path = directory / yaml_name
+    _require(
+        yaml_path.is_file(),
+        f"{where}: rule file {yaml_name!r} is missing",
+        remedy=f"add {directory.name}/{yaml_name}",
+    )
+    head = "".join(yaml_path.read_text(encoding="utf-8").splitlines(keepends=True)[:5])
+    _require(
+        APACHE_SPDX in head,
+        f"{where}: {yaml_name} has no Apache-2.0 SPDX header",
+        remedy=f"add `# {APACHE_SPDX}` as the first line of {yaml_name}",
+    )
+    raw_mappings = block.get("field_mappings")
+    _require(
+        isinstance(raw_mappings, list) and len(raw_mappings) > 0,
+        f"{where}: field_mappings must be a non-empty list",
+        remedy="map every Sigma field the rule keys on to its indexed field",
+    )
+    mappings: list[tuple[str, str]] = []
+    for entry in raw_mappings or []:
+        _require(
+            isinstance(entry, Mapping),
+            f"{where}: each field mapping must be an object",
+            remedy="use {\"sigma_field\": …, \"target_field\": …}",
+        )
+        mappings.append(
+            (_string(entry, "sigma_field", where), _string(entry, "target_field", where))
+        )
+    threshold_block = block.get("threshold")
+    threshold: Threshold | None = None
+    if threshold_block is not None:
+        _require(
+            isinstance(threshold_block, Mapping),
+            f"{where}: threshold must be an object or null",
+            remedy="omit the threshold with null, or give count/window_minutes/group_by",
+        )
+        assert isinstance(threshold_block, Mapping)
+        count = threshold_block.get("count")
+        window = threshold_block.get("window_minutes")
+        _require(
+            isinstance(count, int) and count >= 1,
+            f"{where}: threshold.count must be a positive integer",
+            remedy="set threshold.count",
+        )
+        _require(
+            isinstance(window, int) and window >= 1,
+            f"{where}: threshold.window_minutes must be a positive integer",
+            remedy="set threshold.window_minutes",
+        )
+        threshold = Threshold(
+            count=int(count),  # type: ignore[arg-type]
+            window_minutes=int(window),  # type: ignore[arg-type]
+            group_by=_string(threshold_block, "group_by", where),
+        )
+    sample = block.get("sample")
+    _require(
+        isinstance(sample, Mapping) and isinstance(sample.get("event_index"), int),
+        f"{where}: sample.event_index must be an integer",
+        remedy="name the corpus event the rule's deploy-floor sample test uses",
+    )
+    assert isinstance(sample, Mapping)
+    index = int(sample["event_index"])
+    _require(
+        0 <= index < event_count,
+        f"{where}: sample.event_index {index} is outside the corpus (0..{event_count - 1})",
+        remedy="point sample.event_index at an event the rule matches",
+    )
+    return Detection(
+        rule_title=_string(block, "rule_title", where),
+        description=_string(block, "description", where),
+        severity=_string(block, "severity", where),
+        yaml_path=yaml_path,
+        field_mappings=tuple(mappings),
+        threshold=threshold,
+        sample_event_index=index,
+    )
+
+
+def load_dataset(directory: Path) -> Dataset:
+    """Read and fully validate one dataset directory."""
+    directory = Path(directory)
+    manifest_path = directory / MANIFEST_NAME
+    _require(
+        manifest_path.is_file(),
+        f"{directory} has no {MANIFEST_NAME}",
+        remedy=f"add {directory.name}/{MANIFEST_NAME} (schema {MANIFEST_SCHEMA})",
+    )
+    try:
+        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    except json.JSONDecodeError as exc:
+        raise DatasetError(
+            f"{manifest_path} is not valid JSON: {exc}",
+            remedy="fix the manifest syntax",
+        ) from exc
+    _require(
+        isinstance(manifest, Mapping),
+        f"{manifest_path} must contain a JSON object",
+        remedy="wrap the manifest in an object",
+    )
+    where = f"{directory.name}/{MANIFEST_NAME}"
+    _require(
+        manifest.get("schema") == MANIFEST_SCHEMA,
+        f"{where}: schema must be {MANIFEST_SCHEMA!r}",
+        remedy=f"set \"schema\": \"{MANIFEST_SCHEMA}\"",
+    )
+    dataset_id = _string(manifest, "id", where)
+    _require(
+        bool(ID_PATTERN.match(dataset_id)),
+        f"{where}: id {dataset_id!r} must be lowercase letters, digits and dashes",
+        remedy="rename the dataset id",
+    )
+    _require(
+        dataset_id == directory.name,
+        f"{where}: id {dataset_id!r} must equal the directory name {directory.name!r}",
+        remedy="rename the directory or the id so they agree",
+    )
+    _require(
+        manifest.get("license") == REQUIRED_LICENSE,
+        f"{where}: license must be {REQUIRED_LICENSE!r}",
+        remedy="demo corpora ship Apache-2.0; see content/datasets/LICENSE",
+    )
+    engine = _string(manifest, "engine", where)
+    _require(
+        engine in ENGINES,
+        f"{where}: engine must be one of {', '.join(ENGINES)}",
+        remedy="use \"corpus\" for a dataset that ships its own events",
+    )
+    techniques = manifest.get("attack_techniques")
+    _require(
+        isinstance(techniques, list) and all(isinstance(t, str) and t for t in techniques),
+        f"{where}: attack_techniques must be a list of technique ids",
+        remedy="list the ATT&CK technique ids the corpus depicts",
+    )
+    provenance = manifest.get("provenance")
+    _require(
+        isinstance(provenance, Mapping) and provenance.get("synthetic") is True,
+        f"{where}: provenance.synthetic must be true",
+        remedy="only authored, synthetic corpora may ship here — see content/datasets/README.md",
+    )
+    events_path: Path | None = None
+    detection: Detection | None = None
+    rule_title = ""
+    event_ids: list[str] = []
+    if engine == CORPUS_ENGINE:
+        corpus = manifest.get("corpus")
+        _require(
+            isinstance(corpus, Mapping),
+            f"{where}: engine \"corpus\" requires a corpus block",
+            remedy="add corpus.events and corpus.count",
+        )
+        assert isinstance(corpus, Mapping)
+        events_name = _string(corpus, "events", where + " corpus")
+        events_path = directory / events_name
+        _require(
+            events_path.is_file(),
+            f"{where}: corpus file {events_name!r} is missing",
+            remedy=f"add {directory.name}/{events_name}",
+        )
+        lines = [
+            line
+            for line in events_path.read_text(encoding="utf-8").splitlines()
+            if line.strip()
+        ]
+        parsed: list[dict[str, Any]] = []
+        for number, line in enumerate(lines, start=1):
+            try:
+                event = json.loads(line)
+            except json.JSONDecodeError as exc:
+                raise DatasetError(
+                    f"{events_name} line {number} is not valid JSON: {exc}",
+                    remedy="one JSON object per line, no trailing commas",
+                ) from exc
+            _require(
+                isinstance(event, dict),
+                f"{events_name} line {number} must be a JSON object",
+                remedy="one JSON event object per line",
+            )
+            identifier = _event_id(event)
+            _require(
+                identifier != "",
+                f"{events_name} line {number} has no event.id",
+                remedy="every corpus event needs a stable authored event.id",
+            )
+            _require(
+                isinstance(event.get("@timestamp"), str) and event["@timestamp"].endswith("Z"),
+                f"{events_name} line {number} needs a UTC @timestamp ending in Z",
+                remedy="use an ISO-8601 UTC timestamp such as 2026-03-11T09:00:12Z",
+            )
+            event_ids.append(identifier)
+            parsed.append(event)
+        _require(
+            len(set(event_ids)) == len(event_ids),
+            f"{events_name}: duplicate event.id values",
+            remedy="every authored event.id must be unique inside the corpus",
+        )
+        declared = corpus.get("count")
+        _require(
+            declared == len(parsed),
+            f"{where}: corpus.count is {declared!r} but the file holds {len(parsed)} events",
+            remedy="update corpus.count",
+        )
+        detection_block = manifest.get("detection")
+        _require(
+            isinstance(detection_block, Mapping),
+            f"{where}: engine \"corpus\" requires a detection block",
+            remedy="describe the rule this dataset deploys",
+        )
+        assert isinstance(detection_block, Mapping)
+        detection = _load_detection(directory, detection_block, len(parsed))
+        rule_title = detection.rule_title
+    else:
+        _require(
+            manifest.get("corpus") is None and manifest.get("detection") is None,
+            f"{where}: engine {engine!r} must declare corpus and detection as null",
+            remedy="the golden-path seeder owns this dataset's events and rule",
+        )
+        # The seeder owns the rule, but `reset` still has to find it when a run
+        # dies before recording its id, so the title is manifest data rather
+        # than a literal only the shell knows (veltro#984).
+        golden = manifest.get("golden_path")
+        _require(
+            isinstance(golden, Mapping),
+            f"{where}: engine {engine!r} requires a golden_path block",
+            remedy="add golden_path.rule_title — the title the seeder gives its CHAD rule",
+        )
+        assert isinstance(golden, Mapping)
+        rule_title = _string(golden, "rule_title", where + " golden_path")
+
+    expect = manifest.get("expect")
+    _require(
+        isinstance(expect, Mapping),
+        f"{where}: expect block is required",
+        remedy="declare min_alerts and warden_case",
+    )
+    assert isinstance(expect, Mapping)
+    min_alerts = expect.get("min_alerts")
+    _require(
+        isinstance(min_alerts, int) and min_alerts >= 1,
+        f"{where}: expect.min_alerts must be a positive integer",
+        remedy="a dataset that cannot raise an alert does not belong in the sandbox",
+    )
+    _require(
+        isinstance(expect.get("warden_case"), bool),
+        f"{where}: expect.warden_case must be a boolean",
+        remedy="state whether the alert is expected to open a Warden case",
+    )
+    matching = expect.get("matching_event_ids", [])
+    _require(
+        isinstance(matching, list) and all(isinstance(m, str) for m in matching),
+        f"{where}: expect.matching_event_ids must be a list of authored event ids",
+        remedy="name the corpus events the rule matches",
+    )
+    unknown = [m for m in matching if m not in event_ids]
+    _require(
+        not unknown,
+        f"{where}: expect.matching_event_ids not present in the corpus: {', '.join(unknown)}",
+        remedy="use authored event.id values from events.ndjson",
+    )
+
+
+    return Dataset(
+        id=dataset_id,
+        title=_string(manifest, "title", where),
+        summary=_string(manifest, "summary", where),
+        engine=engine,
+        directory=directory,
+        attack_techniques=tuple(str(t) for t in techniques or ()),
+        events_path=events_path,
+        detection=detection,
+        rule_title=rule_title,
+        min_alerts=int(min_alerts),  # type: ignore[arg-type]
+        expect_warden_case=bool(expect["warden_case"]),
+        matching_event_ids=tuple(str(m) for m in matching or ()),
+    )
+
+
+def datasets_root(root: Path) -> Path:
+    return Path(root).joinpath(*CONTENT_SUBDIR)
+
+
+def load_catalog(root: Path) -> dict[str, Dataset]:
+    """Every dataset under ``<root>/content/datasets``, keyed by id."""
+    base = datasets_root(root)
+    if not base.is_dir():
+        raise DatasetError(
+            f"{base} not found",
+            remedy="run from an installed suite directory or a repository checkout",
+        )
+    catalog: dict[str, Dataset] = {}
+    for directory in sorted(p for p in base.iterdir() if p.is_dir()):
+        if not (directory / MANIFEST_NAME).is_file():
+            continue
+        dataset = load_dataset(directory)
+        catalog[dataset.id] = dataset
+    _require(
+        bool(catalog),
+        f"{base} holds no datasets",
+        remedy="restore content/datasets from the release",
+    )
+    return catalog
+
+
+def split_dataset_names(values: Iterable[str]) -> list[str]:
+    """Accept ``--dataset a --dataset b`` and ``--dataset a,b`` alike."""
+    names: list[str] = []
+    for value in values:
+        for part in str(value).split(","):
+            part = part.strip()
+            if part:
+                names.append(part)
+    return names
+
+
+def resolve_datasets(root: Path, names: Sequence[str]) -> list[Dataset]:
+    """Selected datasets, de-duplicated, golden-path first.
+
+    The golden-path dataset is ordered first deliberately: it provisions the
+    VectorFlow pipeline, the CHAD push data source and the Warden bindings the
+    corpus datasets replay through.
+    """
+    catalog = load_catalog(root)
+    requested = split_dataset_names(names)
+    if not requested or requested == [ALL_KEYWORD]:
+        selected = list(catalog.values()) if requested == [ALL_KEYWORD] else []
+    else:
+        selected = []
+        for name in requested:
+            if name == ALL_KEYWORD:
+                selected.extend(catalog.values())
+                continue
+            dataset = catalog.get(name)
+            if dataset is None:
+                raise DatasetError(
+                    f"unknown dataset {name!r}",
+                    remedy=f"available datasets: {', '.join(sorted(catalog))}",
+                )
+            selected.append(dataset)
+    seen: set[str] = set()
+    unique: list[Dataset] = []
+    for dataset in selected:
+        if dataset.id in seen:
+            continue
+        seen.add(dataset.id)
+        unique.append(dataset)
+    unique.sort(key=lambda d: (d.engine != GOLDEN_PATH_ENGINE, d.id))
+    return unique
+
+
+def _parse_timestamp(value: str) -> datetime:
+    try:
+        return datetime.strptime(value, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)
+    except ValueError as exc:
+        raise DatasetError(
+            f"corpus timestamp {value!r} is not `YYYY-MM-DDTHH:MM:SSZ`",
+            remedy="use whole-second UTC timestamps in the corpus",
+        ) from exc
+
+
+def rebase_events(
+    dataset: Dataset,
+    *,
+    now: datetime,
+    run_token: str,
+    lead_seconds: int = 30,
+) -> list[dict[str, Any]]:
+    """Shift a corpus onto the wall clock and stamp it as demo data.
+
+    The newest authored event lands ``lead_seconds`` before ``now`` and every
+    earlier event keeps its exact offset, so a threshold rule sees the interval
+    the corpus author intended. Each ``event.id`` becomes
+    ``demo-<dataset>-<run token>-<authored id>``: a re-seed produces distinct
+    alerts, and any assertion can follow exactly its own run.
+    """
+    events = dataset.events()
+    if not events:
+        return []
+    stamps = [_parse_timestamp(str(event["@timestamp"])) for event in events]
+    newest = max(stamps)
+    target = now.astimezone(timezone.utc) - timedelta(seconds=lead_seconds)
+    shift = target - newest
+    prepared: list[dict[str, Any]] = []
+    for event, stamp in zip(events, stamps):
+        document = json.loads(json.dumps(event))  # deep copy, corpus stays untouched
+        document["@timestamp"] = (stamp + shift).strftime("%Y-%m-%dT%H:%M:%SZ")
+        block = document.setdefault("event", {})
+        authored = str(block.get("id", ""))
+        block["id"] = f"demo-{dataset.id}-{run_token}-{authored}"
+        block["original_id"] = authored
+        labels = document.setdefault("labels", {})
+        if isinstance(labels, dict):
+            labels[DEMO_LABEL_FIELD] = dataset.id
+            labels["veltro_demo_run"] = run_token
+        tags = document.get("tags")
+        document["tags"] = sorted({*(tags if isinstance(tags, list) else []), DEMO_TAG, f"{DEMO_TAG}:{dataset.id}"})
+        prepared.append(document)
+    return prepared
+
+
+def replay_event_prefix(dataset: Dataset, run_token: str) -> str:
+    """The ``event.id`` prefix every event of one replay carries."""
+    return f"demo-{dataset.id}-{run_token}-"