#!/usr/bin/env python3
"""Fileless malware detection agent using Windows event logs and Volatility."""

import json
import os
import re
import subprocess
import sys
from datetime import datetime

try:
    import Evtx.Evtx as evtx
    HAS_EVTX = True
except ImportError:
    HAS_EVTX = False


LOLBINS = {
    "mshta.exe": {"risk": "HIGH", "usage": "Execute HTA with embedded VBScript/JScript"},
    "regsvr32.exe": {"risk": "HIGH", "usage": "Proxy execution via COM scriptlets"},
    "rundll32.exe": {"risk": "HIGH", "usage": "Execute DLL exports or JavaScript"},
    "certutil.exe": {"risk": "HIGH", "usage": "Download files, decode base64 payloads"},
    "bitsadmin.exe": {"risk": "MEDIUM", "usage": "Download files via BITS service"},
    "wmic.exe": {"risk": "HIGH", "usage": "Remote execution, XSL script processing"},
    "cmstp.exe": {"risk": "HIGH", "usage": "UAC bypass, COM object registration"},
    "msbuild.exe": {"risk": "HIGH", "usage": "Execute inline C# tasks from XML"},
    "installutil.exe": {"risk": "MEDIUM", "usage": "Execute .NET assemblies"},
    "regasm.exe": {"risk": "MEDIUM", "usage": "Execute .NET COM assemblies"},
    "powershell.exe": {"risk": "CONTEXT", "usage": "Script execution, download cradle"},
    "cmd.exe": {"risk": "CONTEXT", "usage": "Command execution, script chaining"},
    "wscript.exe": {"risk": "MEDIUM", "usage": "Execute VBScript/JScript files"},
    "cscript.exe": {"risk": "MEDIUM", "usage": "Execute VBScript/JScript files"},
}

SUSPICIOUS_PS_PATTERNS = [
    (r'-enc\s', "Encoded command execution"),
    (r'IEX\s*\(', "Invoke-Expression (download cradle)"),
    (r'Invoke-Expression', "Invoke-Expression"),
    (r'Net\.WebClient', "WebClient download"),
    (r'DownloadString\(', "Remote script download"),
    (r'DownloadFile\(', "File download"),
    (r'FromBase64String', "Base64 decoding"),
    (r'Reflection\.Assembly', ".NET reflection loading"),
    (r'\[System\.Convert\]', "Type conversion (possible decode)"),
    (r'New-Object\s+IO\.MemoryStream', "In-memory stream (reflective load)"),
    (r'VirtualAlloc', "Memory allocation (shellcode)"),
    (r'CreateThread', "Thread creation (injection)"),
    (r'Add-MpPreference.*ExclusionPath', "Defender exclusion modification"),
    (r'Set-MpPreference.*DisableRealtimeMonitoring', "Defender disablement"),
]


def scan_powershell_logs(log_dir=None):
    """Scan PowerShell script block logs for suspicious patterns."""
    if not log_dir:
        log_dir = r"C:\Windows\System32\winevt\Logs"

    ps_log = os.path.join(log_dir, "Microsoft-Windows-PowerShell%4Operational.evtx")
    if not os.path.exists(ps_log) or not HAS_EVTX:
        return {"error": "PowerShell log not found or python-evtx not installed"}

    alerts = []
    with evtx.Evtx(ps_log) as log:
        for record in log.records():
            try:
                xml = record.xml()
                if "<EventID>4104</EventID>" not in xml:
                    continue
                for pattern, desc in SUSPICIOUS_PS_PATTERNS:
                    if re.search(pattern, xml, re.IGNORECASE):
                        alerts.append({
                            "event_id": 4104,
                            "timestamp": record.timestamp().isoformat(),
                            "detection": desc,
                            "snippet": xml[:500],
                        })
                        break
            except Exception:
                continue

    return {"log_file": ps_log, "suspicious_events": len(alerts), "alerts": alerts[:50]}


def scan_sysmon_for_lolbins(log_dir=None):
    """Scan Sysmon logs for LOLBin process creation events."""
    if not log_dir:
        log_dir = r"C:\Windows\System32\winevt\Logs"

    sysmon_log = os.path.join(log_dir, "Microsoft-Windows-Sysmon%4Operational.evtx")
    if not os.path.exists(sysmon_log) or not HAS_EVTX:
        return {"error": "Sysmon log not found or python-evtx not installed"}

    detections = []
    with evtx.Evtx(sysmon_log) as log:
        for record in log.records():
            try:
                xml = record.xml()
                if "<EventID>1</EventID>" not in xml:
                    continue
                for lolbin, info in LOLBINS.items():
                    if lolbin.lower() in xml.lower():
                        detections.append({
                            "timestamp": record.timestamp().isoformat(),
                            "lolbin": lolbin,
                            "risk": info["risk"],
                            "known_abuse": info["usage"],
                            "snippet": xml[:500],
                        })
                        break
            except Exception:
                continue

    return {"log_file": sysmon_log, "lolbin_detections": len(detections), "detections": detections[:50]}


def scan_wmi_persistence():
    """Detect WMI event subscription persistence mechanisms."""
    cmd = [
        "powershell", "-Command",
        "Get-WMIObject -Namespace root\\Subscription -Class __EventFilter | "
        "Select-Object Name, Query | ConvertTo-Json"
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
        if result.returncode == 0 and result.stdout.strip():
            filters = json.loads(result.stdout)
            if not isinstance(filters, list):
                filters = [filters]
            return {"wmi_event_filters": filters, "count": len(filters)}
        return {"wmi_event_filters": [], "count": 0}
    except Exception as e:
        return {"error": str(e)}


def scan_registry_run_keys():
    """Check registry Run keys for suspicious persistence entries."""
    keys_to_check = [
        r"HKLM\SOFTWARE\Microsoft\Windows\CurrentVersion\Run",
        r"HKCU\SOFTWARE\Microsoft\Windows\CurrentVersion\Run",
        r"HKLM\SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce",
    ]
    results = []
    for key in keys_to_check:
        cmd = ["reg", "query", key]
        try:
            r = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
            if r.returncode == 0:
                for line in r.stdout.strip().splitlines():
                    line = line.strip()
                    if line and not line.startswith("HKEY"):
                        for lolbin in LOLBINS:
                            if lolbin.lower() in line.lower():
                                results.append({
                                    "key": key,
                                    "entry": line,
                                    "lolbin_detected": lolbin,
                                    "risk": "HIGH",
                                })
        except Exception:
            continue
    return {"registry_persistence": results, "count": len(results)}


def run_volatility_malfind(memory_dump):
    """Run Volatility malfind to detect injected code in memory."""
    if not os.path.exists(memory_dump):
        return {"error": f"Memory dump not found: {memory_dump}"}
    cmd = ["vol3", "-f", memory_dump, "windows.malfind"]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        return {"output": result.stdout.strip(), "exit_code": result.returncode}
    except FileNotFoundError:
        return {"error": "Volatility 3 (vol3) not installed"}
    except subprocess.TimeoutExpired:
        return {"error": "Volatility analysis timed out"}


def generate_report():
    """Generate fileless malware detection report."""
    return {
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "powershell_scan": scan_powershell_logs(),
        "lolbin_scan": scan_sysmon_for_lolbins(),
        "wmi_persistence": scan_wmi_persistence(),
        "registry_persistence": scan_registry_run_keys(),
    }


if __name__ == "__main__":
    action = sys.argv[1] if len(sys.argv) > 1 else "report"
    if action == "report":
        print(json.dumps(generate_report(), indent=2, default=str))
    elif action == "powershell":
        log_dir = sys.argv[2] if len(sys.argv) > 2 else None
        print(json.dumps(scan_powershell_logs(log_dir), indent=2, default=str))
    elif action == "lolbins":
        print(json.dumps(scan_sysmon_for_lolbins(), indent=2, default=str))
    elif action == "wmi":
        print(json.dumps(scan_wmi_persistence(), indent=2))
    elif action == "registry":
        print(json.dumps(scan_registry_run_keys(), indent=2))
    elif action == "malfind" and len(sys.argv) > 2:
        print(json.dumps(run_volatility_malfind(sys.argv[2]), indent=2))
    else:
        print("Usage: agent.py [report|powershell [log_dir]|lolbins|wmi|registry|malfind <memory.dmp>]")