#!/usr/bin/env python3
"""Build a MITRE ATT&CK Navigator layer from Wazuh 5.0 detection content.

Two sources, because they answer two different questions:

  --from-rules DIR   What the rule catalogue CLAIMS to cover.
                     Reads the `mitre:` block of every Sigma YAML rule under DIR.
                     Score = number of rules mapped to that technique.

  --from-findings    What the deployment has ACTUALLY OBSERVED.
                     Runs a terms aggregation over wazuh-findings-v5-* on
                     wazuh.rule.mitre.technique.id and wazuh.rule.mitre.subtechnique.id.
                     Score = finding count in the selected window.

                     Verified on a Wazuh 5.0.0-beta5 lab (23 Sep 2026): a finding
                     carries mitre as an OBJECT, not a flat list:
                       wazuh.rule.mitre.technique.id / .name
                       wazuh.rule.mitre.tactic.id / .name
                     so the aggregation field must be the .id sub-field. Findings
                     produced by the shipped content carried technique and tactic
                     but no subtechnique; that aggregation may legitimately return
                     no buckets.

Output conforms to ATT&CK Navigator Layer File Format v4.5:
https://github.com/mitre-attack/attack-navigator/blob/master/layers/spec/v4.5/layerformat.md

WHY THERE IS NO "read the rules from the API" MODE
--------------------------------------------------
The Wazuh 5.0-beta Content Manager API specification defines POST, PUT and DELETE on
/_plugins/_content_manager/rules but no GET or list operation, so there is no documented
endpoint that returns the rule catalogue. Read your rules from source control (--from-rules)
or measure what fired (--from-findings).

Python 3.9+. Standard library only for --from-findings; --from-rules needs PyYAML.
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import re
import ssl
import sys
import urllib.error
import urllib.request
from collections import Counter

LAYER_SPEC_VERSION = "4.5"
MIN_NAVIGATOR_VERSION = "4.9.0"

# T1234 or T1234.001
TECHNIQUE_RE = re.compile(r"^T\d{4}(\.\d{3})?$")


# --------------------------------------------------------------------------- rules


def counts_from_rules(root: str) -> Counter:
    """Count Sigma rules per ATT&CK technique / subtechnique ID."""
    try:
        import yaml  # type: ignore
    except ImportError:
        sys.exit("--from-rules needs PyYAML: pip install pyyaml")

    counts: Counter = Counter()
    seen = 0

    for dirpath, _dirnames, filenames in os.walk(root):
        for filename in filenames:
            if not filename.endswith((".yml", ".yaml")):
                continue
            path = os.path.join(dirpath, filename)
            with open(path, encoding="utf-8") as handle:
                try:
                    documents = list(yaml.safe_load_all(handle))
                except yaml.YAMLError as exc:
                    print(f"skipping {path}: {exc}", file=sys.stderr)
                    continue

            for document in documents:
                if not isinstance(document, dict):
                    continue
                mitre = document.get("mitre")
                if not isinstance(mitre, dict):
                    continue
                seen += 1
                for key in ("technique", "subtechnique"):
                    for value in mitre.get(key) or []:
                        identifier = str(value).strip()
                        if TECHNIQUE_RE.match(identifier):
                            counts[identifier] += 1
                        else:
                            print(
                                f"{path}: ignoring non-technique value {identifier!r} "
                                f"under mitre.{key}",
                                file=sys.stderr,
                            )

    print(f"read {seen} rules with a mitre block from {root}", file=sys.stderr)
    return counts


# ------------------------------------------------------------------------ findings


def counts_from_findings(
    url: str,
    user: str,
    password: str,
    index: str,
    since: str,
    field_suffix: str,
    insecure: bool,
    bucket_size: int,
) -> Counter:
    """Terms aggregation over the findings data streams, per technique ID."""
    body = {
        "size": 0,
        "query": {"range": {"@timestamp": {"gte": since}}},
        "aggs": {
            "technique": {
                "terms": {
                    "field": f"wazuh.rule.mitre.technique.id{field_suffix}",
                    "size": bucket_size,
                }
            },
            "subtechnique": {
                "terms": {
                    "field": f"wazuh.rule.mitre.subtechnique.id{field_suffix}",
                    "size": bucket_size,
                }
            },
        },
    }

    endpoint = f"{url.rstrip('/')}/{index}/_search"
    request = urllib.request.Request(
        endpoint,
        data=json.dumps(body).encode("utf-8"),
        method="POST",
        headers={
            "Content-Type": "application/json",
            "Authorization": "Basic "
            + base64.b64encode(f"{user}:{password}".encode()).decode(),
        },
    )

    context = ssl.create_default_context()
    if insecure:
        context.check_hostname = False
        context.verify_mode = ssl.CERT_NONE

    try:
        with urllib.request.urlopen(request, context=context, timeout=60) as response:
            payload = json.load(response)
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")
        sys.exit(
            f"{endpoint} returned {exc.code}.\n{detail}\n\n"
            "A 400 here usually means the mitre fields are not aggregatable in your "
            "mapping. Inspect the mapping and retry with --field-suffix .keyword"
        )
    except urllib.error.URLError as exc:
        sys.exit(f"could not reach {endpoint}: {exc.reason}")

    counts: Counter = Counter()
    for aggregation in ("technique", "subtechnique"):
        buckets = payload.get("aggregations", {}).get(aggregation, {}).get("buckets", [])
        if not buckets:
            print(f"no buckets returned for {aggregation}", file=sys.stderr)
        for bucket in buckets:
            identifier = str(bucket["key"]).strip()
            if TECHNIQUE_RE.match(identifier):
                counts[identifier] += int(bucket["doc_count"])
    return counts


# --------------------------------------------------------------------------- layer


def build_layer(
    counts: Counter, name: str, description: str, unit: str, attack_version: str | None
) -> dict:
    if not counts:
        sys.exit("no techniques found, refusing to write an empty layer")

    highest = max(counts.values())

    techniques = [
        {
            "techniqueID": identifier,
            "score": score,
            "comment": f"{score} {unit if score != 1 else unit.rstrip('s')}",
            "enabled": True,
        }
        for identifier, score in sorted(counts.items())
    ]

    versions = {"navigator": MIN_NAVIGATOR_VERSION, "layer": LAYER_SPEC_VERSION}
    if attack_version:
        versions["attack"] = attack_version

    return {
        "name": name,
        "versions": versions,
        "domain": "enterprise-attack",
        "description": description,
        "sorting": 3,
        "hideDisabled": False,
        "techniques": techniques,
        "gradient": {
            "colors": ["#ffe766", "#8ec843"],
            "minValue": 0,
            "maxValue": highest,
        },
        "legendItems": [
            {"label": f"low ({unit})", "color": "#ffe766"},
            {"label": f"high ({unit})", "color": "#8ec843"},
        ],
        "layout": {"layout": "side", "showID": True, "showName": True},
        "showTacticRowBackground": True,
        "tacticRowBackground": "#dddddd",
    }


# ---------------------------------------------------------------------------- main


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Export a MITRE ATT&CK Navigator layer from Wazuh 5.0 content."
    )
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--from-rules", metavar="DIR", help="directory of Sigma YAML rules")
    source.add_argument(
        "--from-findings", action="store_true", help="aggregate over the findings streams"
    )

    parser.add_argument("--url", default="https://127.0.0.1:9200", help="Wazuh indexer URL")
    parser.add_argument("--user", default=os.environ.get("WAZUH_INDEXER_USER", ""))
    parser.add_argument("--password", default=os.environ.get("WAZUH_INDEXER_PASSWORD", ""))
    parser.add_argument("--index", default="wazuh-findings-v5-*")
    parser.add_argument("--since", default="now-30d", help="date-math lower bound on @timestamp")
    parser.add_argument(
        "--field-suffix",
        default="",
        help="append to the mitre field names, e.g. .keyword, if they are not aggregatable",
    )
    parser.add_argument("--bucket-size", type=int, default=1000)
    parser.add_argument(
        "--insecure",
        action="store_true",
        help="skip TLS verification (self-signed installer certificates)",
    )
    parser.add_argument(
        "--attack-version",
        default=None,
        help="ATT&CK version to pin the layer to; omitted means the Navigator's current version",
    )
    parser.add_argument("--name", default=None, help="layer name")
    parser.add_argument("-o", "--output", default="wazuh-attack-layer.json")

    args = parser.parse_args()

    if args.from_rules:
        counts = counts_from_rules(args.from_rules)
        unit = "rules"
        name = args.name or "Wazuh rule coverage"
        description = (
            "Techniques mapped in the mitre block of Sigma rules under "
            f"{args.from_rules}. Score is the number of rules per technique. "
            "This is claimed coverage, not observed coverage."
        )
    else:
        if not args.user or not args.password:
            sys.exit(
                "set --user/--password, or WAZUH_INDEXER_USER and WAZUH_INDEXER_PASSWORD"
            )
        counts = counts_from_findings(
            args.url,
            args.user,
            args.password,
            args.index,
            args.since,
            args.field_suffix,
            args.insecure,
            args.bucket_size,
        )
        unit = "findings"
        name = args.name or "Wazuh observed coverage"
        description = (
            f"Techniques present on findings in {args.index} since {args.since}. "
            "Score is the finding count per technique. This is observed activity, "
            "not rule coverage."
        )

    layer = build_layer(counts, name, description, unit, args.attack_version)

    with open(args.output, "w", encoding="utf-8") as handle:
        json.dump(layer, handle, indent=2)
        handle.write("\n")

    print(
        f"wrote {args.output}: {len(layer['techniques'])} techniques, "
        f"max score {layer['gradient']['maxValue']}",
        file=sys.stderr,
    )


if __name__ == "__main__":
    main()