#!/usr/bin/env python3
"""Volatility 3 comprehensive memory analysis wrapper.

Runs key Volatility 3 plugins against a memory dump, auto-detects the OS,
and outputs a consolidated JSON report.

Usage:
    python vol3_analyze.py --dump memory.raw --output report.json
    python vol3_analyze.py --dump memory.raw --plugins pslist,netscan,malfind
    python vol3_analyze.py --dump memory.raw --detect-only
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import platform
import shutil
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

# Plugin definitions organized by OS
WINDOWS_PLUGINS = {
    "pslist": "windows.pslist",
    "psscan": "windows.psscan",
    "pstree": "windows.pstree",
    "netscan": "windows.netscan",
    "malfind": "windows.malfind",
    "dlllist": "windows.dlllist",
    "handles": "windows.handles",
    "svcscan": "windows.svcscan",
    "modules": "windows.modules",
    "driverscan": "windows.driverscan",
    "ssdt": "windows.ssdt",
    "callbacks": "windows.callbacks",
    "cmdline": "windows.cmdline",
    "envars": "windows.envars",
    "filescan": "windows.filescan",
    "registry_hivelist": "windows.registry.hivelist",
    "ldrmodules": "windows.ldrmodules",
    "vadinfo": "windows.vadinfo",
}

LINUX_PLUGINS = {
    "pslist": "linux.pslist",
    "pstree": "linux.pstree",
    "lsmod": "linux.lsmod",
    "sockstat": "linux.sockstat",
    "bash": "linux.bash",
    "elfs": "linux.elfs",
    "tty_check": "linux.tty_check",
    "check_syscall": "linux.check_syscall",
    "check_modules": "linux.check_modules",
    "proc_maps": "linux.proc.Maps",
    "mountinfo": "linux.mountinfo",
}

MAC_PLUGINS = {
    "pslist": "mac.pslist",
    "pstree": "mac.pstree",
    "lsmod": "mac.lsmod",
    "netstat": "mac.netstat",
    "bash": "mac.bash",
    "check_syscall": "mac.check_syscall",
    "mount": "mac.mount",
}

# Default plugins for quick triage (per OS)
DEFAULT_PLUGINS = {
    "windows": ["pslist", "psscan", "netscan", "malfind", "dlllist", "handles", "svcscan", "cmdline"],
    "linux": ["pslist", "pstree", "lsmod", "sockstat", "bash", "check_syscall"],
    "mac": ["pslist", "pstree", "lsmod", "netstat", "bash"],
}

OS_PLUGIN_MAP = {
    "windows": WINDOWS_PLUGINS,
    "linux": LINUX_PLUGINS,
    "mac": MAC_PLUGINS,
}


def find_volatility() -> str:
    """Locate the Volatility 3 executable."""
    # Check for 'vol' or 'vol3' or 'volatility3' in PATH
    for name in ("vol", "vol3", "volatility3"):
        path = shutil.which(name)
        if path:
            return path

    # Check if volatility3 is importable as a Python module
    try:
        result = subprocess.run(
            [sys.executable, "-m", "volatility3", "--help"],
            capture_output=True, text=True, timeout=15,
        )
        if result.returncode == 0:
            return f"{sys.executable} -m volatility3"
    except (subprocess.TimeoutExpired, FileNotFoundError):
        pass

    # Check common installation paths
    common_paths = [
        Path.home() / "volatility3" / "vol.py",
        Path("/opt/volatility3/vol.py"),
        Path("/usr/local/bin/vol.py"),
    ]
    for p in common_paths:
        if p.exists():
            return f"{sys.executable} {p}"

    logger.error(
        "Volatility 3 not found. Install with: pip install volatility3"
    )
    sys.exit(1)


def detect_os(vol_cmd: str, dump_path: str) -> Optional[str]:
    """Auto-detect the OS of the memory dump."""
    info_plugins = {
        "windows": "windows.info",
        "linux": "linux.info",
        "mac": "mac.info",
    }

    for os_name, plugin in info_plugins.items():
        logger.info(f"Trying {os_name} detection with {plugin}...")
        try:
            cmd = f"{vol_cmd} -f {dump_path} {plugin}"
            result = subprocess.run(
                cmd.split() if " -m " not in vol_cmd else cmd,
                shell=" -m " in vol_cmd,
                capture_output=True, text=True, timeout=120,
            )
            if result.returncode == 0 and result.stdout.strip():
                logger.info(f"Detected OS: {os_name}")
                return os_name
        except subprocess.TimeoutExpired:
            logger.warning(f"Timeout detecting {os_name}")
        except Exception as e:
            logger.debug(f"Error detecting {os_name}: {e}")

    logger.warning("Could not auto-detect OS. Defaulting to 'windows'.")
    return "windows"


def run_plugin(
    vol_cmd: str,
    dump_path: str,
    plugin_name: str,
    extra_args: Optional[list] = None,
    timeout: int = 300,
) -> dict[str, Any]:
    """Run a single Volatility 3 plugin and return structured results."""
    cmd_parts = f"{vol_cmd} -f {dump_path} -r json {plugin_name}"
    if extra_args:
        cmd_parts += " " + " ".join(extra_args)

    logger.info(f"Running plugin: {plugin_name}")
    start_time = time.time()

    try:
        result = subprocess.run(
            cmd_parts if " -m " in vol_cmd else cmd_parts.split(),
            shell=" -m " in vol_cmd,
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        elapsed = round(time.time() - start_time, 2)

        if result.returncode == 0:
            try:
                data = json.loads(result.stdout)
            except json.JSONDecodeError:
                data = {"raw_output": result.stdout}

            return {
                "plugin": plugin_name,
                "status": "success",
                "elapsed_seconds": elapsed,
                "data": data,
                "errors": result.stderr.strip() if result.stderr.strip() else None,
            }
        else:
            return {
                "plugin": plugin_name,
                "status": "error",
                "elapsed_seconds": elapsed,
                "data": None,
                "errors": result.stderr.strip() or result.stdout.strip(),
            }

    except subprocess.TimeoutExpired:
        elapsed = round(time.time() - start_time, 2)
        logger.warning(f"Plugin {plugin_name} timed out after {timeout}s")
        return {
            "plugin": plugin_name,
            "status": "timeout",
            "elapsed_seconds": elapsed,
            "data": None,
            "errors": f"Timed out after {timeout} seconds",
        }
    except Exception as e:
        elapsed = round(time.time() - start_time, 2)
        logger.error(f"Plugin {plugin_name} failed: {e}")
        return {
            "plugin": plugin_name,
            "status": "error",
            "elapsed_seconds": elapsed,
            "data": None,
            "errors": str(e),
        }


def compare_process_lists(results: dict) -> list[dict]:
    """Compare pslist and psscan to find hidden processes."""
    hidden = []
    pslist_data = results.get("pslist", {}).get("data")
    psscan_data = results.get("psscan", {}).get("data")

    if not pslist_data or not psscan_data:
        return hidden

    # Extract PIDs from pslist
    pslist_pids = set()
    if isinstance(pslist_data, list):
        for entry in pslist_data:
            pid = entry.get("PID") or entry.get("pid")
            if pid is not None:
                pslist_pids.add(int(pid))

    # Find PIDs in psscan but not in pslist
    if isinstance(psscan_data, list):
        for entry in psscan_data:
            pid = entry.get("PID") or entry.get("pid")
            if pid is not None and int(pid) not in pslist_pids:
                exit_time = entry.get("ExitTime") or entry.get("exit_time")
                # Processes with exit times are terminated, not hidden
                if not exit_time or exit_time == "N/A":
                    hidden.append({
                        "pid": int(pid),
                        "name": entry.get("ImageFileName") or entry.get("name", "unknown"),
                        "reason": "Found in psscan but not pslist (potential DKOM hiding)",
                    })

    return hidden


def analyze_malfind(results: dict) -> list[dict]:
    """Extract and summarize malfind results for suspicious injections."""
    suspicious = []
    malfind_data = results.get("malfind", {}).get("data")

    if not malfind_data or not isinstance(malfind_data, list):
        return suspicious

    for entry in malfind_data:
        pid = entry.get("PID") or entry.get("pid")
        process = entry.get("Process") or entry.get("process", "unknown")
        protection = entry.get("Protection") or entry.get("protection", "unknown")
        hexdump = entry.get("Hexdump") or entry.get("hexdump", "")

        # Check for MZ header (PE injection)
        has_mz = False
        if isinstance(hexdump, str):
            has_mz = "4d 5a" in hexdump.lower() or "MZ" in hexdump

        suspicious.append({
            "pid": pid,
            "process": process,
            "protection": protection,
            "has_pe_header": has_mz,
            "severity": "high" if has_mz else "medium",
        })

    return suspicious


def generate_summary(results: dict, detected_os: str) -> dict:
    """Generate an analysis summary with key findings."""
    summary = {
        "os_detected": detected_os,
        "plugins_run": len(results),
        "plugins_succeeded": sum(1 for r in results.values() if r.get("status") == "success"),
        "plugins_failed": sum(1 for r in results.values() if r.get("status") != "success"),
        "hidden_processes": [],
        "suspicious_injections": [],
        "key_findings": [],
    }

    # Hidden process detection
    hidden = compare_process_lists(results)
    if hidden:
        summary["hidden_processes"] = hidden
        summary["key_findings"].append(
            f"Found {len(hidden)} potentially hidden process(es)"
        )

    # Malfind analysis
    injections = analyze_malfind(results)
    if injections:
        summary["suspicious_injections"] = injections
        pe_injections = sum(1 for i in injections if i.get("has_pe_header"))
        summary["key_findings"].append(
            f"Found {len(injections)} suspicious memory regions "
            f"({pe_injections} with PE headers)"
        )

    # Network connection count
    netscan_data = results.get("netscan", {}).get("data")
    if netscan_data and isinstance(netscan_data, list):
        summary["key_findings"].append(
            f"Found {len(netscan_data)} network connection(s)"
        )

    # Service count
    svcscan_data = results.get("svcscan", {}).get("data")
    if svcscan_data and isinstance(svcscan_data, list):
        summary["key_findings"].append(
            f"Found {len(svcscan_data)} service(s)"
        )

    return summary


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Volatility 3 comprehensive memory analysis wrapper",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s --dump memory.raw --output report.json
  %(prog)s --dump memory.raw --plugins pslist,netscan,malfind
  %(prog)s --dump memory.raw --detect-only
  %(prog)s --dump memory.raw --os linux --output report.json
        """,
    )
    parser.add_argument(
        "--input", "--dump", "-f", dest="dump", required=True,
        help="Path to memory dump file",
    )
    parser.add_argument(
        "--output", "-o",
        help="Output JSON report path (default: stdout)",
    )
    parser.add_argument(
        "--plugins", "-p",
        help="Comma-separated list of plugins to run (default: OS-specific defaults)",
    )
    parser.add_argument(
        "--os",
        choices=["windows", "linux", "mac"],
        help="Force OS type (skip auto-detection)",
    )
    parser.add_argument(
        "--detect-only",
        action="store_true",
        help="Only detect the OS and exit",
    )
    parser.add_argument(
        "--dump-dir",
        help="Directory to dump extracted artifacts",
    )
    parser.add_argument(
        "--timeout",
        type=int, default=300,
        help="Per-plugin timeout in seconds (default: 300)",
    )
    parser.add_argument(
        "--vol-path",
        help="Path to Volatility 3 executable",
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Enable verbose logging",
    )
    parser.add_argument(
        "--format", default="json", choices=["json", "text", "csv"],
        help="Output format (default: json)",
    )

    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # Validate dump file
    dump_path = os.path.abspath(args.dump)
    if not os.path.isfile(dump_path):
        logger.error(f"Memory dump not found: {dump_path}")
        sys.exit(1)

    dump_size = os.path.getsize(dump_path)
    logger.info(f"Memory dump: {dump_path} ({dump_size / (1024**3):.2f} GB)")

    # Find Volatility 3
    vol_cmd = args.vol_path or find_volatility()
    logger.info(f"Using Volatility 3: {vol_cmd}")

    # Detect OS
    detected_os = args.os or detect_os(vol_cmd, dump_path)
    if not detected_os:
        logger.error("Could not detect OS. Specify with --os flag.")
        sys.exit(1)

    logger.info(f"OS: {detected_os}")

    if args.detect_only:
        print(json.dumps({"os": detected_os}, indent=2))
        sys.exit(0)

    # Determine plugins to run
    plugin_map = OS_PLUGIN_MAP.get(detected_os, WINDOWS_PLUGINS)

    if args.plugins:
        requested = [p.strip() for p in args.plugins.split(",")]
        plugins_to_run = {}
        for name in requested:
            if name in plugin_map:
                plugins_to_run[name] = plugin_map[name]
            else:
                # Allow raw plugin names
                plugins_to_run[name] = name
    else:
        default_names = DEFAULT_PLUGINS.get(detected_os, [])
        plugins_to_run = {n: plugin_map[n] for n in default_names if n in plugin_map}

    logger.info(f"Running {len(plugins_to_run)} plugin(s): {', '.join(plugins_to_run.keys())}")

    # Create dump directory if requested
    if args.dump_dir:
        os.makedirs(args.dump_dir, exist_ok=True)

    # Run plugins
    results = {}
    for short_name, full_name in plugins_to_run.items():
        extra_args = []
        if args.dump_dir and short_name == "malfind":
            extra_args = ["--dump", "--dump-dir", args.dump_dir]

        result = run_plugin(vol_cmd, dump_path, full_name, extra_args, args.timeout)
        results[short_name] = result

    # Generate summary
    summary = generate_summary(results, detected_os)

    # Build final report
    report = {
        "metadata": {
            "tool": "vol3_analyze.py",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "dump_file": dump_path,
            "dump_size_bytes": dump_size,
            "detected_os": detected_os,
            "volatility_cmd": vol_cmd,
            "platform": platform.system(),
        },
        "summary": summary,
        "plugin_results": results,
    }

    # Output report
    report_json = json.dumps(report, indent=2, default=str)

    if args.output:
        output_path = os.path.abspath(args.output)
        with open(output_path, "w") as f:
            f.write(report_json)
        logger.info(f"Report written to: {output_path}")
    else:
        print(report_json)

    # Print summary to stderr
    print("\n=== Analysis Summary ===", file=sys.stderr)
    print(f"OS: {detected_os}", file=sys.stderr)
    print(
        f"Plugins: {summary['plugins_succeeded']}/{summary['plugins_run']} succeeded",
        file=sys.stderr,
    )
    for finding in summary["key_findings"]:
        print(f"  - {finding}", file=sys.stderr)
    if not summary["key_findings"]:
        print("  No significant findings.", file=sys.stderr)


if __name__ == "__main__":
    main()