#!/usr/bin/env python3
"""Map observed malware behaviors to MITRE ATT&CK techniques.

Takes observed behaviors and maps them to ATT&CK technique IDs using
a local technique database.
"""

from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime
from pathlib import Path

# Local ATT&CK technique database (subset of most common techniques)
ATTACK_TECHNIQUES = {
    "T1566.001": {
        "name": "Phishing: Spearphishing Attachment",
        "tactic": "Initial Access",
        "description": "Adversaries send spearphishing emails with malicious attachments.",
        "detection": "Monitor for suspicious email attachments, file execution from email clients.",
        "mitigations": ["User training", "Email filtering", "Antivirus"],
    },
    "T1566.002": {
        "name": "Phishing: Spearphishing Link",
        "tactic": "Initial Access",
        "description": "Adversaries send spearphishing emails with malicious links.",
        "detection": "URL filtering, email link analysis, proxy logs.",
        "mitigations": ["URL filtering", "User training", "Web proxy"],
    },
    "T1059.001": {
        "name": "Command and Scripting Interpreter: PowerShell",
        "tactic": "Execution",
        "description": "Adversaries abuse PowerShell for execution and automation.",
        "detection": "PowerShell logging (ScriptBlock, Module), AMSI, process monitoring.",
        "mitigations": ["Constrained Language Mode", "Script signing", "AMSI"],
    },
    "T1059.003": {
        "name": "Command and Scripting Interpreter: Windows Command Shell",
        "tactic": "Execution",
        "description": "Adversaries abuse cmd.exe to execute commands.",
        "detection": "Process monitoring, command-line logging.",
        "mitigations": ["Execution prevention", "Process monitoring"],
    },
    "T1059.005": {
        "name": "Command and Scripting Interpreter: Visual Basic",
        "tactic": "Execution",
        "description": "Adversaries abuse VBA macros for execution.",
        "detection": "Office macro logging, process creation monitoring.",
        "mitigations": ["Disable macros", "Attack Surface Reduction rules"],
    },
    "T1204.002": {
        "name": "User Execution: Malicious File",
        "tactic": "Execution",
        "description": "Adversary relies on user executing a malicious file.",
        "detection": "Monitor file execution events, endpoint detection.",
        "mitigations": ["User training", "Application allowlisting"],
    },
    "T1547.001": {
        "name": "Boot or Logon Autostart Execution: Registry Run Keys",
        "tactic": "Persistence",
        "description": "Adversaries add programs to registry Run keys for persistence.",
        "detection": "Registry monitoring (Run, RunOnce keys), Sysmon Event ID 13.",
        "mitigations": ["Registry auditing", "Endpoint detection"],
    },
    "T1053.005": {
        "name": "Scheduled Task/Job: Scheduled Task",
        "tactic": "Persistence",
        "description": "Adversaries create scheduled tasks for persistence or execution.",
        "detection": "Task Scheduler logs, schtasks.exe monitoring.",
        "mitigations": ["Privileged account management", "Task auditing"],
    },
    "T1543.003": {
        "name": "Create or Modify System Process: Windows Service",
        "tactic": "Persistence",
        "description": "Adversaries create or modify services for persistence.",
        "detection": "Service creation events, sc.exe monitoring.",
        "mitigations": ["Least privilege", "Service auditing"],
    },
    "T1055.001": {
        "name": "Process Injection: DLL Injection",
        "tactic": "Defense Evasion",
        "description": "Adversaries inject DLLs into processes to evade detection.",
        "detection": "API monitoring (CreateRemoteThread, LoadLibrary), Sysmon.",
        "mitigations": ["Endpoint detection", "Process integrity"],
    },
    "T1055.012": {
        "name": "Process Injection: Process Hollowing",
        "tactic": "Defense Evasion",
        "description": "Adversaries hollow out legitimate processes to host malicious code.",
        "detection": "API monitoring, memory analysis, process integrity checks.",
        "mitigations": ["Endpoint detection", "Memory protection"],
    },
    "T1027": {
        "name": "Obfuscated Files or Information",
        "tactic": "Defense Evasion",
        "description": "Adversaries obfuscate payloads to evade detection.",
        "detection": "File analysis, deobfuscation tools, behavioral detection.",
        "mitigations": ["Antivirus", "Behavioral detection"],
    },
    "T1027.002": {
        "name": "Obfuscated Files or Information: Software Packing",
        "tactic": "Defense Evasion",
        "description": "Adversaries pack executables to evade signature detection.",
        "detection": "Packer detection (entropy analysis, signatures).",
        "mitigations": ["Behavioral detection", "Sandbox analysis"],
    },
    "T1071.001": {
        "name": "Application Layer Protocol: Web Protocols",
        "tactic": "Command and Control",
        "description": "Adversaries use HTTP/HTTPS for C2 communication.",
        "detection": "Network monitoring, proxy logs, TLS inspection.",
        "mitigations": ["Network filtering", "SSL inspection"],
    },
    "T1573.001": {
        "name": "Encrypted Channel: Symmetric Cryptography",
        "tactic": "Command and Control",
        "description": "Adversaries encrypt C2 with symmetric algorithms.",
        "detection": "Network traffic analysis, JA3 fingerprinting.",
        "mitigations": ["SSL inspection", "Network monitoring"],
    },
    "T1568.002": {
        "name": "Dynamic Resolution: Domain Generation Algorithms",
        "tactic": "Command and Control",
        "description": "Adversaries use DGAs to generate C2 domains.",
        "detection": "DNS monitoring, DGA detection algorithms, NXDomain analysis.",
        "mitigations": ["DNS filtering", "Sinkholing"],
    },
    "T1486": {
        "name": "Data Encrypted for Impact",
        "tactic": "Impact",
        "description": "Adversaries encrypt data for ransom or destruction.",
        "detection": "File modification monitoring, volume shadow copy deletion.",
        "mitigations": ["Backups", "Endpoint detection", "File integrity monitoring"],
    },
    "T1490": {
        "name": "Inhibit System Recovery",
        "tactic": "Impact",
        "description": "Adversaries delete backups and shadow copies.",
        "detection": "Monitor vssadmin, wbadmin, bcdedit commands.",
        "mitigations": ["Offline backups", "Access controls on backup tools"],
    },
    "T1003.001": {
        "name": "OS Credential Dumping: LSASS Memory",
        "tactic": "Credential Access",
        "description": "Adversaries dump LSASS process memory for credentials.",
        "detection": "LSASS access monitoring, Credential Guard alerts.",
        "mitigations": ["Credential Guard", "LSA protection", "Privileged access management"],
    },
    "T1082": {
        "name": "System Information Discovery",
        "tactic": "Discovery",
        "description": "Adversaries gather system information (hostname, OS, hardware).",
        "detection": "Monitor for systeminfo, hostname, ver commands.",
        "mitigations": ["N/A (legitimate behavior)"],
    },
    "T1083": {
        "name": "File and Directory Discovery",
        "tactic": "Discovery",
        "description": "Adversaries enumerate files and directories.",
        "detection": "Monitor for dir, find, ls commands from unusual processes.",
        "mitigations": ["N/A (legitimate behavior)"],
    },
    "T1041": {
        "name": "Exfiltration Over C2 Channel",
        "tactic": "Exfiltration",
        "description": "Adversaries exfiltrate data over existing C2 channel.",
        "detection": "Network monitoring, data volume analysis, DLP.",
        "mitigations": ["DLP", "Network segmentation", "Egress filtering"],
    },
    "T1105": {
        "name": "Ingress Tool Transfer",
        "tactic": "Command and Control",
        "description": "Adversaries transfer tools/files from external systems.",
        "detection": "Monitor for file downloads, certutil, bitsadmin, curl.",
        "mitigations": ["Network filtering", "Application allowlisting"],
    },
    "T1218.011": {
        "name": "System Binary Proxy Execution: Rundll32",
        "tactic": "Defense Evasion",
        "description": "Adversaries abuse rundll32.exe to execute malicious DLLs.",
        "detection": "Monitor rundll32 execution with unusual DLLs or arguments.",
        "mitigations": ["Application allowlisting", "Process monitoring"],
    },
    "T1014": {
        "name": "Rootkit",
        "tactic": "Defense Evasion",
        "description": "Adversaries use rootkits to hide malware presence.",
        "detection": "Memory forensics, integrity checking, behavioral detection.",
        "mitigations": ["Secure Boot", "Kernel integrity checking"],
    },
    "T1542.001": {
        "name": "Pre-OS Boot: System Firmware",
        "tactic": "Persistence",
        "description": "Adversaries modify system firmware for persistence.",
        "detection": "Firmware integrity verification, CHIPSEC, Secure Boot.",
        "mitigations": ["Secure Boot", "Firmware update verification"],
    },
    "T1505.003": {
        "name": "Server Software Component: Web Shell",
        "tactic": "Persistence",
        "description": "Adversaries install webshells on compromised web servers.",
        "detection": "File integrity monitoring, web traffic analysis, process monitoring.",
        "mitigations": ["File integrity monitoring", "Web application firewall"],
    },
    "T1195.002": {
        "name": "Supply Chain Compromise: Compromise Software Supply Chain",
        "tactic": "Initial Access",
        "description": "Adversaries compromise software supply chains.",
        "detection": "Software verification, code signing, dependency analysis.",
        "mitigations": ["Code signing verification", "Vendor security assessment"],
    },
}

# Behavior-to-technique mapping keywords
BEHAVIOR_KEYWORDS = {
    "registry run key": ["T1547.001"],
    "persistence": ["T1547.001", "T1053.005", "T1543.003"],
    "scheduled task": ["T1053.005"],
    "service creation": ["T1543.003"],
    "process injection": ["T1055.001", "T1055.012"],
    "dll injection": ["T1055.001"],
    "process hollowing": ["T1055.012"],
    "powershell": ["T1059.001"],
    "cmd.exe": ["T1059.003"],
    "vba macro": ["T1059.005"],
    "obfuscation": ["T1027"],
    "packing": ["T1027.002"],
    "http c2": ["T1071.001"],
    "encrypted c2": ["T1573.001"],
    "dga": ["T1568.002"],
    "ransomware": ["T1486", "T1490"],
    "encryption": ["T1486"],
    "shadow copy": ["T1490"],
    "credential dump": ["T1003.001"],
    "lsass": ["T1003.001"],
    "phishing": ["T1566.001", "T1566.002"],
    "rootkit": ["T1014"],
    "webshell": ["T1505.003"],
    "firmware": ["T1542.001"],
    "supply chain": ["T1195.002"],
    "download": ["T1105"],
    "rundll32": ["T1218.011"],
    "exfiltration": ["T1041"],
}


def map_behaviors_to_attack(behaviors) -> dict:
    """Map observed behaviors to ATT&CK techniques."""
    mappings = []
    matched_techniques = set()

    for behavior in behaviors:
        behavior_lower = behavior.lower()
        for keyword, technique_ids in BEHAVIOR_KEYWORDS.items():
            if keyword in behavior_lower:
                for tid in technique_ids:
                    if tid not in matched_techniques and tid in ATTACK_TECHNIQUES:
                        tech = ATTACK_TECHNIQUES[tid]
                        mappings.append({
                            "technique_id": tid,
                            "technique_name": tech["name"],
                            "tactic": tech["tactic"],
                            "matched_behavior": behavior,
                            "detection": tech["detection"],
                            "mitigations": tech["mitigations"],
                        })
                        matched_techniques.add(tid)

    return mappings


def lookup_techniques(technique_ids) -> dict:
    """Look up specific ATT&CK techniques by ID."""
    results = []
    for tid in technique_ids:
        tid = tid.strip().upper()
        if tid in ATTACK_TECHNIQUES:
            tech = ATTACK_TECHNIQUES[tid]
            results.append({
                "technique_id": tid,
                "technique_name": tech["name"],
                "tactic": tech["tactic"],
                "description": tech["description"],
                "detection": tech["detection"],
                "mitigations": tech["mitigations"],
            })
        else:
            results.append({
                "technique_id": tid,
                "error": "Technique not found in local database",
            })
    return results


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Map malware behaviors to MITRE ATT&CK techniques"
    )
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--input", "--behaviors", dest="behaviors", help="JSON file with list of observed behaviors")
    group.add_argument("--techniques", help="Comma-separated ATT&CK technique IDs to look up")
    group.add_argument("--list-all", action="store_true", help="List all techniques in database")

    parser.add_argument("--output", "-o", help="Output file (JSON)")
    parser.add_argument("--format", choices=["json", "text"], default="text")
    args = parser.parse_args()

    if args.list_all:
        result = {
            "total_techniques": len(ATTACK_TECHNIQUES),
            "techniques": [
                {"id": tid, "name": t["name"], "tactic": t["tactic"]}
                for tid, t in sorted(ATTACK_TECHNIQUES.items())
            ]
        }
    elif args.behaviors:
        path = Path(args.behaviors)
        if path.exists():
            behaviors = json.loads(path.read_text())
            if isinstance(behaviors, dict):
                behaviors = behaviors.get("behaviors", [])
        else:
            behaviors = [args.behaviors]

        mappings = map_behaviors_to_attack(behaviors)
        result = {
            "input_behaviors": len(behaviors),
            "techniques_mapped": len(mappings),
            "mappings": mappings,
        }
    else:
        technique_ids = [t.strip() for t in args.techniques.split(",")]
        results = lookup_techniques(technique_ids)
        result = {"techniques": results}

    result["timestamp"] = datetime.now().isoformat()
    result["tool"] = "attack_mapper"

    if args.format == "json":
        output = json.dumps(result, indent=2)
    else:
        output = "=== MITRE ATT&CK Mapping ===\n\n"
        if "mappings" in result:
            for m in result["mappings"]:
                output += f"[{m['technique_id']}] {m['technique_name']}\n"
                output += f"  Tactic: {m['tactic']}\n"
                output += f"  Evidence: {m['matched_behavior']}\n"
                output += f"  Detection: {m['detection']}\n\n"
        elif "techniques" in result:
            for t in result["techniques"]:
                if "error" in t:
                    output += f"[{t['technique_id']}] {t['error']}\n"
                else:
                    output += f"[{t['technique_id']}] {t['technique_name']}\n"
                    output += f"  Tactic: {t['tactic']}\n"
                    output += f"  Description: {t.get('description', 'N/A')}\n"
                    output += f"  Detection: {t.get('detection', 'N/A')}\n\n"

    if args.output:
        Path(args.output).write_text(output)
        print(f"[+] Output saved to {args.output}")
    else:
        print(output)


if __name__ == "__main__":
    main()