#!/usr/bin/env python3
"""Generate ATT&CK Navigator layer JSON files for visualization.

Creates ATT&CK Navigator-compatible layer files from technique mappings with
confidence levels, enabling visual representation of observed TTPs in the
MITRE ATT&CK Navigator web application.
"""

from __future__ import annotations

import argparse
import json
import sys
import uuid
from pathlib import Path
from typing import Any

# ATT&CK Navigator layer schema version
LAYER_VERSION = "4.5"
NAV_VERSION = "4.9.1"
ATTACK_VERSION = "14"
DOMAIN = "enterprise-attack"

# Color schemes for confidence levels
CONFIDENCE_COLORS: dict[str, str] = {
    "high": "#ff6666",      # Red
    "medium": "#ffaf66",    # Orange
    "low": "#ffe766",       # Yellow
    "unknown": "#a1d99b",   # Light green
}

CONFIDENCE_SCORES: dict[str, int] = {
    "high": 100,
    "medium": 60,
    "low": 30,
    "unknown": 10,
}


def parse_techniques_string(techniques_str: str) -> list[dict[str, Any]]:
    """Parse a comma-separated technique:confidence string.

    Format: T1055:high,T1071.001:medium,T1082:low
    """
    techniques: list[dict[str, Any]] = []
    for entry in techniques_str.split(","):
        entry = entry.strip()
        if not entry:
            continue
        parts = entry.split(":")
        technique_id = parts[0].strip().upper()
        confidence = parts[1].strip().lower() if len(parts) > 1 else "unknown"
        if confidence not in CONFIDENCE_COLORS:
            confidence = "unknown"
        techniques.append({
            "techniqueID": technique_id,
            "confidence": confidence,
        })
    return techniques


def load_techniques_from_file(file_path: Path) -> list[dict[str, Any]]:
    """Load technique mappings from a JSON file."""
    with open(file_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    if isinstance(data, list):
        return data
    if isinstance(data, dict) and "techniques" in data:
        return data["techniques"]
    return [data]


def build_navigator_layer(
    name: str,
    techniques: list[dict[str, Any]],
    description: str = "",
    color_scheme: str = "confidence",
    comments: dict[str, str] | None = None,
    metadata: dict[str, list[dict[str, str]]] | None = None,
) -> dict[str, Any]:
    """Build an ATT&CK Navigator layer JSON structure."""
    layer_techniques: list[dict[str, Any]] = []

    for tech in techniques:
        technique_id = tech.get("techniqueID", tech.get("technique_id", ""))
        confidence = tech.get("confidence", "unknown").lower()

        layer_tech: dict[str, Any] = {
            "techniqueID": technique_id,
            "tactic": "",
            "color": CONFIDENCE_COLORS.get(confidence, CONFIDENCE_COLORS["unknown"]),
            "comment": "",
            "enabled": True,
            "metadata": [],
            "links": [],
            "showSubtechniques": False,
            "score": CONFIDENCE_SCORES.get(confidence, 10),
        }

        # Add comment if provided
        if comments and technique_id in comments:
            layer_tech["comment"] = comments[technique_id]
        elif "comment" in tech:
            layer_tech["comment"] = tech["comment"]

        # Add confidence as metadata
        layer_tech["metadata"].append({
            "name": "confidence",
            "value": confidence,
        })

        # Add custom metadata if provided
        if metadata and technique_id in metadata:
            layer_tech["metadata"].extend(metadata[technique_id])

        # Handle subtechnique display
        if "." in technique_id:
            layer_tech["showSubtechniques"] = False
        else:
            # Check if there are subtechniques for this technique
            has_subtechniques = any(
                t.get("techniqueID", t.get("technique_id", "")).startswith(
                    f"{technique_id}."
                )
                for t in techniques
            )
            layer_tech["showSubtechniques"] = has_subtechniques

        layer_techniques.append(layer_tech)

    layer: dict[str, Any] = {
        "name": name,
        "versions": {
            "attack": ATTACK_VERSION,
            "navigator": NAV_VERSION,
            "layer": LAYER_VERSION,
        },
        "domain": DOMAIN,
        "description": description,
        "filters": {
            "platforms": [
                "Linux",
                "macOS",
                "Windows",
                "Network",
                "PRE",
                "Containers",
                "Office 365",
                "SaaS",
                "Google Workspace",
                "IaaS",
                "Azure AD",
            ]
        },
        "sorting": 0,
        "layout": {
            "layout": "side",
            "aggregateFunction": "average",
            "showID": True,
            "showName": True,
            "showAggregateScores": False,
            "countUnscored": False,
        },
        "hideDisabled": False,
        "techniques": layer_techniques,
        "gradient": {
            "colors": [
                CONFIDENCE_COLORS["low"],
                CONFIDENCE_COLORS["medium"],
                CONFIDENCE_COLORS["high"],
            ],
            "minValue": 0,
            "maxValue": 100,
        },
        "legendItems": [
            {"label": "High confidence", "color": CONFIDENCE_COLORS["high"]},
            {"label": "Medium confidence", "color": CONFIDENCE_COLORS["medium"]},
            {"label": "Low confidence", "color": CONFIDENCE_COLORS["low"]},
        ],
        "metadata": [],
        "links": [],
        "showTacticRowBackground": True,
        "tacticRowBackground": "#dddddd",
        "selectTechniquesAcrossTactics": True,
        "selectSubtechniquesWithParent": False,
        "selectVisibleTechniques": False,
    }

    return layer


def analyze(input_path: Path, output_format: str = "json") -> dict[str, Any]:
    """Load techniques from input and generate a Navigator layer."""
    techniques = load_techniques_from_file(input_path)
    return build_navigator_layer(
        name="Analysis Layer",
        techniques=techniques,
    )


def main() -> None:
    """Entry point for the ATT&CK Navigator layer generator CLI."""
    parser = argparse.ArgumentParser(
        description="Generate MITRE ATT&CK Navigator layer JSON files."
    )
    parser.add_argument(
        "--input",
        type=Path,
        help="Path to technique mappings JSON file (alternative to --techniques)",
    )
    parser.add_argument(
        "--output", type=Path, help="Path to output file (stdout if omitted)"
    )
    parser.add_argument(
        "--format",
        default="json",
        choices=["json", "text"],
        help="Output format (default: json)",
    )
    parser.add_argument(
        "--name",
        type=str,
        default="Analysis Layer",
        help="Layer name displayed in ATT&CK Navigator",
    )
    parser.add_argument(
        "--techniques",
        type=str,
        help="Comma-separated technique:confidence pairs (e.g., T1055:high,T1071.001:medium)",
    )
    parser.add_argument(
        "--description",
        type=str,
        default="",
        help="Layer description",
    )
    parser.add_argument(
        "--color-scheme",
        choices=["confidence", "tactic", "custom"],
        default="confidence",
        help="Color scheme for the layer (default: confidence)",
    )
    parser.add_argument(
        "--comments",
        type=Path,
        help="JSON file with technique-level comments (format: {\"T1055\": \"comment text\"})",
    )
    parser.add_argument(
        "--metadata",
        type=Path,
        dest="metadata_file",
        help="JSON file with custom metadata per technique",
    )
    args = parser.parse_args()

    # Load techniques from either --techniques string or --input file
    techniques: list[dict[str, Any]] = []
    if args.techniques:
        techniques = parse_techniques_string(args.techniques)
    elif args.input and args.input.exists():
        techniques = load_techniques_from_file(args.input)
    else:
        print(
            "Error: Provide either --techniques or --input with technique mappings.",
            file=sys.stderr,
        )
        sys.exit(1)

    # Load optional comments
    comments: dict[str, str] | None = None
    if args.comments and args.comments.exists():
        with open(args.comments, "r", encoding="utf-8") as f:
            comments = json.load(f)

    # Load optional metadata
    custom_metadata: dict[str, list[dict[str, str]]] | None = None
    if args.metadata_file and args.metadata_file.exists():
        with open(args.metadata_file, "r", encoding="utf-8") as f:
            custom_metadata = json.load(f)

    layer = build_navigator_layer(
        name=args.name,
        techniques=techniques,
        description=args.description,
        color_scheme=args.color_scheme,
        comments=comments,
        metadata=custom_metadata,
    )

    if args.format == "text":
        output_text = _format_text(layer)
    else:
        output_text = json.dumps(layer, indent=2)

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(output_text, encoding="utf-8")
        print(f"Navigator layer written to {args.output}", file=sys.stderr)
    else:
        print(output_text)


def _format_text(layer: dict[str, Any]) -> str:
    """Format layer as human-readable text summary."""
    lines: list[str] = [
        f"=== ATT&CK Navigator Layer: {layer['name']} ===",
        f"Domain: {layer['domain']}",
        f"Description: {layer.get('description', 'N/A')}",
        f"Techniques mapped: {len(layer['techniques'])}",
        "",
        "Techniques:",
    ]
    for tech in layer["techniques"]:
        confidence = "unknown"
        for md in tech.get("metadata", []):
            if md.get("name") == "confidence":
                confidence = md["value"]
                break
        lines.append(f"  {tech['techniqueID']} [{confidence}] score={tech.get('score', 'N/A')}")
        if tech.get("comment"):
            lines.append(f"    Comment: {tech['comment']}")
    return "\n".join(lines)


if __name__ == "__main__":
    main()