#!/usr/bin/env python3
"""Look up MITRE ATT&CK techniques by ID, keyword, or tactic.

Provides a command-line interface for searching the ATT&CK Enterprise technique
database. Supports keyword search, technique ID lookup, tactic filtering, and
detailed technique information including detection and mitigation guidance.
"""

from __future__ import annotations

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

# Bundled technique database path (relative to script location)
DEFAULT_DATA_PATH = Path(__file__).parent.parent / "assets" / "attack-data.json"

# Fallback stub data when the bundled database is not available
STUB_TECHNIQUES: list[dict[str, Any]] = [
    {
        "id": "T1055",
        "name": "Process Injection",
        "tactic": ["defense-evasion", "privilege-escalation"],
        "description": "Adversaries may inject code into processes to evade process-based defenses and elevate privileges.",
        "platforms": ["Windows", "Linux", "macOS"],
        "detection": "Monitor for suspicious process behavior such as unexpected memory writes or thread creation in remote processes.",
        "mitigation": "Use endpoint detection tools that monitor for injection techniques. Apply behavior-based detection.",
        "subtechniques": ["T1055.001", "T1055.002", "T1055.003", "T1055.012"],
    },
    {
        "id": "T1071",
        "name": "Application Layer Protocol",
        "tactic": ["command-and-control"],
        "description": "Adversaries may communicate using application layer protocols to avoid detection.",
        "platforms": ["Windows", "Linux", "macOS"],
        "detection": "Analyze network data for uncommon data flows. Monitor for unusual HTTP, DNS, or HTTPS traffic patterns.",
        "mitigation": "Use network intrusion detection systems to identify unusual traffic patterns.",
        "subtechniques": ["T1071.001", "T1071.002", "T1071.003", "T1071.004"],
    },
    {
        "id": "T1071.001",
        "name": "Web Protocols",
        "tactic": ["command-and-control"],
        "description": "Adversaries may communicate using HTTP/HTTPS to blend C2 traffic with normal web traffic.",
        "platforms": ["Windows", "Linux", "macOS"],
        "detection": "Monitor HTTP/HTTPS traffic for unusual patterns, beaconing behavior, or connections to known-bad domains.",
        "mitigation": "Use TLS inspection and web proxy filtering to detect anomalous traffic.",
        "subtechniques": [],
    },
    {
        "id": "T1082",
        "name": "System Information Discovery",
        "tactic": ["discovery"],
        "description": "An adversary may attempt to get detailed information about the operating system and hardware.",
        "platforms": ["Windows", "Linux", "macOS"],
        "detection": "Monitor for execution of system information gathering commands such as systeminfo, hostname, or uname.",
        "mitigation": "This technique is difficult to mitigate as it uses legitimate system functionality.",
        "subtechniques": [],
    },
    {
        "id": "T1547.001",
        "name": "Registry Run Keys / Startup Folder",
        "tactic": ["persistence", "privilege-escalation"],
        "description": "Adversaries may achieve persistence by adding programs to registry run keys or the startup folder.",
        "platforms": ["Windows"],
        "detection": "Monitor registry modifications to run keys and file creation in startup folders.",
        "mitigation": "Restrict registry permissions and monitor changes to startup locations.",
        "subtechniques": [],
    },
    {
        "id": "T1027",
        "name": "Obfuscated Files or Information",
        "tactic": ["defense-evasion"],
        "description": "Adversaries may use obfuscation to make files or information difficult to discover or analyze.",
        "platforms": ["Windows", "Linux", "macOS"],
        "detection": "Monitor for file obfuscation indicators such as high entropy sections or encoded content.",
        "mitigation": "Use anti-malware solutions that can detect obfuscated content and behavioral analysis.",
        "subtechniques": ["T1027.001", "T1027.002", "T1027.003"],
    },
    {
        "id": "T1059.001",
        "name": "PowerShell",
        "tactic": ["execution"],
        "description": "Adversaries may abuse PowerShell for execution of commands and scripts.",
        "platforms": ["Windows"],
        "detection": "Monitor PowerShell execution, enable script block logging, and module logging.",
        "mitigation": "Use Constrained Language mode and disable PowerShell v2. Enable logging.",
        "subtechniques": [],
    },
    {
        "id": "T1486",
        "name": "Data Encrypted for Impact",
        "tactic": ["impact"],
        "description": "Adversaries may encrypt data on target systems to interrupt availability.",
        "platforms": ["Windows", "Linux", "macOS"],
        "detection": "Monitor for mass file modification events and unusual cryptographic API usage.",
        "mitigation": "Maintain offline backups. Use anti-ransomware protections.",
        "subtechniques": [],
    },
    {
        "id": "T1053.005",
        "name": "Scheduled Task",
        "tactic": ["execution", "persistence", "privilege-escalation"],
        "description": "Adversaries may use Windows Task Scheduler to execute programs at system startup or on a scheduled basis.",
        "platforms": ["Windows"],
        "detection": "Monitor scheduled task creation via schtasks.exe, Task Scheduler COM objects, and WMI.",
        "mitigation": "Restrict task creation to authorized administrators.",
        "subtechniques": [],
    },
    {
        "id": "T1003.001",
        "name": "LSASS Memory",
        "tactic": ["credential-access"],
        "description": "Adversaries may access LSASS process memory to extract credential material.",
        "platforms": ["Windows"],
        "detection": "Monitor for processes accessing lsass.exe memory. Enable Credential Guard.",
        "mitigation": "Enable Windows Credential Guard. Restrict debug privileges.",
        "subtechniques": [],
    },
]


def load_attack_data(data_path: Path | None = None) -> list[dict[str, Any]]:
    """Load ATT&CK technique data from JSON file or use stub data."""
    if data_path and data_path.exists():
        with open(data_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]

    # Check default path
    if DEFAULT_DATA_PATH.exists():
        with open(DEFAULT_DATA_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 STUB_TECHNIQUES


def search_techniques(
    techniques: list[dict[str, Any]], keyword: str
) -> list[dict[str, Any]]:
    """Search techniques by keyword in name, description, and ID."""
    keyword_lower = keyword.lower()
    results: list[dict[str, Any]] = []
    for tech in techniques:
        searchable = (
            f"{tech.get('id', '')} {tech.get('name', '')} "
            f"{tech.get('description', '')}"
        ).lower()
        if keyword_lower in searchable:
            results.append(tech)
    return results


def lookup_by_id(
    techniques: list[dict[str, Any]], technique_id: str
) -> dict[str, Any] | None:
    """Look up a specific technique by its ATT&CK ID."""
    technique_id_upper = technique_id.upper()
    for tech in techniques:
        if tech.get("id", "").upper() == technique_id_upper:
            return tech
    return None


def filter_by_tactic(
    techniques: list[dict[str, Any]], tactic: str
) -> list[dict[str, Any]]:
    """Filter techniques by tactic name."""
    tactic_lower = tactic.lower()
    results: list[dict[str, Any]] = []
    for tech in techniques:
        tactics = tech.get("tactic", [])
        if isinstance(tactics, str):
            tactics = [tactics]
        if any(tactic_lower in t.lower() for t in tactics):
            results.append(tech)
    return results


def format_technique_brief(tech: dict[str, Any]) -> dict[str, Any]:
    """Format a technique as a brief summary."""
    return {
        "id": tech.get("id", ""),
        "name": tech.get("name", ""),
        "tactic": tech.get("tactic", []),
        "platforms": tech.get("platforms", []),
    }


def format_technique_full(tech: dict[str, Any]) -> dict[str, Any]:
    """Format a technique with full details."""
    return {
        "id": tech.get("id", ""),
        "name": tech.get("name", ""),
        "tactic": tech.get("tactic", []),
        "description": tech.get("description", ""),
        "platforms": tech.get("platforms", []),
        "detection": tech.get("detection", ""),
        "mitigation": tech.get("mitigation", ""),
        "subtechniques": tech.get("subtechniques", []),
    }


def analyze(input_path: Path, output_format: str = "json") -> dict[str, Any]:
    """Load ATT&CK data from the given path and return technique listing."""
    techniques = load_attack_data(input_path)
    return {
        "total_techniques": len(techniques),
        "techniques": [format_technique_brief(t) for t in techniques],
    }


def main() -> None:
    """Entry point for the ATT&CK technique lookup CLI."""
    parser = argparse.ArgumentParser(
        description="Look up MITRE ATT&CK techniques by ID, keyword, or tactic."
    )
    parser.add_argument(
        "--input",
        "--data",
        type=Path,
        dest="input",
        help="Path to ATT&CK data JSON file (uses bundled data if omitted)",
    )
    parser.add_argument(
        "--output", type=Path, help="Path to output file (stdout if omitted)"
    )
    parser.add_argument(
        "--format",
        default="json",
        choices=["json", "text", "csv"],
        help="Output format (default: json)",
    )
    parser.add_argument(
        "--search",
        type=str,
        help="Search techniques by keyword",
    )
    parser.add_argument(
        "--id",
        type=str,
        dest="technique_id",
        help="Look up a specific technique by ATT&CK ID (e.g., T1055)",
    )
    parser.add_argument(
        "--tactic",
        type=str,
        help="Filter techniques by tactic (e.g., defense-evasion)",
    )
    parser.add_argument(
        "--detail",
        choices=["brief", "full"],
        default="brief",
        help="Detail level for results (default: brief)",
    )
    args = parser.parse_args()

    techniques = load_attack_data(args.input)

    if args.technique_id:
        tech = lookup_by_id(techniques, args.technique_id)
        if tech is None:
            print(f"Technique {args.technique_id} not found.", file=sys.stderr)
            sys.exit(1)
        if args.detail == "full":
            result: Any = format_technique_full(tech)
        else:
            result = format_technique_brief(tech)
    elif args.search:
        matches = search_techniques(techniques, args.search)
        formatter = format_technique_full if args.detail == "full" else format_technique_brief
        result = {
            "query": args.search,
            "results_count": len(matches),
            "results": [formatter(t) for t in matches],
        }
    elif args.tactic:
        matches = filter_by_tactic(techniques, args.tactic)
        formatter = format_technique_full if args.detail == "full" else format_technique_brief
        result = {
            "tactic": args.tactic,
            "results_count": len(matches),
            "results": [formatter(t) for t in matches],
        }
    else:
        formatter = format_technique_full if args.detail == "full" else format_technique_brief
        result = {
            "total_techniques": len(techniques),
            "techniques": [formatter(t) for t in techniques],
        }

    if args.format == "text":
        output_text = _format_text(result)
    elif args.format == "csv":
        output_text = _format_csv(result)
    else:
        output_text = json.dumps(result, 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"Results written to {args.output}", file=sys.stderr)
    else:
        print(output_text)


def _format_text(result: Any) -> str:
    """Format results as human-readable text."""
    lines: list[str] = []
    if isinstance(result, dict) and "results" in result:
        lines.append(f"Query: {result.get('query', result.get('tactic', 'all'))}")
        lines.append(f"Results: {result['results_count']}")
        lines.append("")
        for tech in result["results"]:
            lines.append(f"  {tech['id']} - {tech['name']}")
            if "description" in tech:
                lines.append(f"    {tech['description'][:120]}...")
            lines.append(f"    Tactics: {', '.join(tech.get('tactic', []))}")
            lines.append("")
    elif isinstance(result, dict) and "id" in result:
        lines.append(f"{result['id']} - {result['name']}")
        lines.append(f"Tactics: {', '.join(result.get('tactic', []))}")
        if "description" in result:
            lines.append(f"\n{result['description']}")
        if "detection" in result:
            lines.append(f"\nDetection:\n  {result['detection']}")
        if "mitigation" in result:
            lines.append(f"\nMitigation:\n  {result['mitigation']}")
    else:
        lines.append(json.dumps(result, indent=2))
    return "\n".join(lines)


def _format_csv(result: Any) -> str:
    """Format results as CSV."""
    lines: list[str] = ["id,name,tactic,platforms"]
    techs = []
    if isinstance(result, dict) and "results" in result:
        techs = result["results"]
    elif isinstance(result, dict) and "techniques" in result:
        techs = result["techniques"]
    elif isinstance(result, dict) and "id" in result:
        techs = [result]
    for tech in techs:
        tactics = "|".join(tech.get("tactic", []))
        platforms = "|".join(tech.get("platforms", []))
        name = tech.get("name", "").replace(",", ";")
        lines.append(f"{tech.get('id', '')},{name},{tactics},{platforms}")
    return "\n".join(lines)


if __name__ == "__main__":
    main()