#!/usr/bin/env python3
"""
Sandbox Monitor - Process, filesystem, and network monitoring for malware analysis.

Monitors system activity during malware execution and generates timelines
for behavioral analysis. Cross-platform support via psutil and watchdog.

Usage:
    # Monitor a specific process
    python sandbox_monitor.py --pid 1234 --output ./results --duration 300

    # Watch for new processes
    python sandbox_monitor.py --watch-new --output ./results --duration 300

    # Generate timeline from collected data
    python sandbox_monitor.py --timeline ./results --format json
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import platform
import signal
import sys
import time
import threading
import hashlib
from datetime import datetime, timezone
from pathlib import Path
from collections import defaultdict

try:
    import psutil
except ImportError:
    print("Error: psutil is required. Install with: pip install psutil", file=sys.stderr)
    sys.exit(1)

try:
    from watchdog.observers import Observer
    from watchdog.events import FileSystemEventHandler
    HAS_WATCHDOG = True
except ImportError:
    HAS_WATCHDOG = False
    print("Warning: watchdog not installed. Filesystem monitoring disabled.", file=sys.stderr)
    print("Install with: pip install watchdog", file=sys.stderr)


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%dT%H:%M:%S",
)
logger = logging.getLogger(__name__)


class EventRecorder:
    """Thread-safe event recording with categorization."""

    def __init__(self):
        self._events = []
        self._lock = threading.Lock()
        self._start_time = time.time()

    def record(self, category, event_type, details) -> dict:
        """Record a timestamped event."""
        event = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "elapsed_seconds": round(time.time() - self._start_time, 3),
            "category": category,
            "type": event_type,
            "details": details,
        }
        with self._lock:
            self._events.append(event)
        return event

    def get_events(self, category=None) -> dict:
        """Retrieve events, optionally filtered by category."""
        with self._lock:
            if category:
                return [e for e in self._events if e["category"] == category]
            return list(self._events)

    def get_summary(self) -> dict:
        """Get event count summary by category."""
        with self._lock:
            counts = defaultdict(int)
            for event in self._events:
                counts[event["category"]] += 1
            return dict(counts)


class ProcessMonitor:
    """Monitor process creation, termination, and resource usage."""

    def __init__(self, recorder, target_pid=None, watch_new=False):
        self.recorder = recorder
        self.target_pid = target_pid
        self.watch_new = watch_new
        self._running = False
        self._known_pids = set()
        self._process_info_cache = {}

    def _get_process_info(self, proc):
        """Safely extract process information."""
        info = {"pid": proc.pid}
        try:
            info["name"] = proc.name()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            info["name"] = "<unknown>"
        try:
            info["exe"] = proc.exe()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            info["exe"] = "<unknown>"
        try:
            info["cmdline"] = proc.cmdline()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            info["cmdline"] = []
        try:
            info["ppid"] = proc.ppid()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            info["ppid"] = None
        try:
            info["username"] = proc.username()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            info["username"] = "<unknown>"
        try:
            info["create_time"] = datetime.fromtimestamp(
                proc.create_time(), tz=timezone.utc
            ).isoformat()
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            info["create_time"] = None
        return info

    def _get_resource_usage(self, proc):
        """Get CPU and memory usage for a process."""
        usage = {"pid": proc.pid}
        try:
            usage["cpu_percent"] = proc.cpu_percent(interval=0.1)
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            usage["cpu_percent"] = 0.0
        try:
            mem = proc.memory_info()
            usage["memory_rss_bytes"] = mem.rss
            usage["memory_vms_bytes"] = mem.vms
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            usage["memory_rss_bytes"] = 0
            usage["memory_vms_bytes"] = 0
        return usage

    def _get_connections(self, proc):
        """Get network connections for a process."""
        connections = []
        try:
            for conn in proc.net_connections(kind="all"):
                conn_info = {
                    "pid": proc.pid,
                    "fd": conn.fd,
                    "family": str(conn.family),
                    "type": str(conn.type),
                    "status": conn.status,
                    "local_address": (
                        f"{conn.laddr.ip}:{conn.laddr.port}" if conn.laddr else None
                    ),
                    "remote_address": (
                        f"{conn.raddr.ip}:{conn.raddr.port}" if conn.raddr else None
                    ),
                }
                connections.append(conn_info)
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
        return connections

    def _monitor_target(self, pid):
        """Monitor a specific process and its children."""
        try:
            proc = psutil.Process(pid)
        except psutil.NoSuchProcess:
            logger.warning("Process %d not found", pid)
            return

        info = self._get_process_info(proc)
        self.recorder.record("process", "target_start", info)
        self._known_pids.add(pid)

        while self._running:
            try:
                if not proc.is_running():
                    self.recorder.record(
                        "process", "terminated", {"pid": pid, "name": info["name"]}
                    )
                    break

                # Record resource usage
                usage = self._get_resource_usage(proc)
                self.recorder.record("process", "resource_usage", usage)

                # Record network connections
                connections = self._get_connections(proc)
                for conn in connections:
                    self.recorder.record("network", "connection", conn)

                # Check for child processes
                try:
                    children = proc.children(recursive=True)
                    for child in children:
                        if child.pid not in self._known_pids:
                            self._known_pids.add(child.pid)
                            child_info = self._get_process_info(child)
                            child_info["parent_pid"] = pid
                            self.recorder.record("process", "child_created", child_info)
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    pass

                time.sleep(1)
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                self.recorder.record(
                    "process", "terminated", {"pid": pid, "name": info["name"]}
                )
                break

    def _monitor_new_processes(self):
        """Watch for newly created processes."""
        # Snapshot existing processes
        self._known_pids = set(psutil.pids())
        logger.info("Baseline: %d existing processes", len(self._known_pids))

        while self._running:
            current_pids = set(psutil.pids())
            new_pids = current_pids - self._known_pids
            gone_pids = self._known_pids - current_pids

            for pid in new_pids:
                try:
                    proc = psutil.Process(pid)
                    info = self._get_process_info(proc)
                    self.recorder.record("process", "new_process", info)
                    logger.info("New process: PID=%d Name=%s", pid, info["name"])
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    pass

            for pid in gone_pids:
                cached = self._process_info_cache.get(pid, {})
                self.recorder.record(
                    "process",
                    "process_exit",
                    {"pid": pid, "name": cached.get("name", "<unknown>")},
                )

            self._known_pids = current_pids

            # Cache process info for terminated process reporting
            for pid in new_pids:
                try:
                    proc = psutil.Process(pid)
                    self._process_info_cache[pid] = self._get_process_info(proc)
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    pass

            time.sleep(0.5)

    def start(self) -> None:
        """Start process monitoring."""
        self._running = True
        if self.target_pid:
            self._thread = threading.Thread(
                target=self._monitor_target, args=(self.target_pid,), daemon=True
            )
        elif self.watch_new:
            self._thread = threading.Thread(
                target=self._monitor_new_processes, daemon=True
            )
        else:
            raise ValueError("Must specify either target_pid or watch_new")
        self._thread.start()
        logger.info("Process monitor started")

    def stop(self) -> None:
        """Stop process monitoring."""
        self._running = False
        if hasattr(self, "_thread"):
            self._thread.join(timeout=5)
        logger.info("Process monitor stopped")


class FilesystemMonitor:
    """Monitor filesystem changes using watchdog."""

    class _Handler(FileSystemEventHandler):
        def __init__(self, recorder):
            self.recorder = recorder

        def on_created(self, event) -> None:
            details = {"path": event.src_path, "is_directory": event.is_directory}
            if not event.is_directory:
                try:
                    details["size"] = os.path.getsize(event.src_path)
                except OSError:
                    details["size"] = None
            self.recorder.record("filesystem", "created", details)

        def on_modified(self, event) -> None:
            details = {"path": event.src_path, "is_directory": event.is_directory}
            if not event.is_directory:
                try:
                    details["size"] = os.path.getsize(event.src_path)
                except OSError:
                    details["size"] = None
            self.recorder.record("filesystem", "modified", details)

        def on_deleted(self, event) -> None:
            self.recorder.record(
                "filesystem",
                "deleted",
                {"path": event.src_path, "is_directory": event.is_directory},
            )

        def on_moved(self, event) -> None:
            self.recorder.record(
                "filesystem",
                "moved",
                {
                    "src_path": event.src_path,
                    "dest_path": event.dest_path,
                    "is_directory": event.is_directory,
                },
            )

    def __init__(self, recorder, watch_paths=None):
        if not HAS_WATCHDOG:
            raise RuntimeError("watchdog library required for filesystem monitoring")
        self.recorder = recorder
        self.watch_paths = watch_paths or self._default_watch_paths()
        self._observer = Observer()

    @staticmethod
    def _default_watch_paths():
        """Return default paths to monitor based on platform."""
        system = platform.system()
        if system == "Windows":
            return [
                os.environ.get("TEMP", r"C:\Windows\Temp"),
                os.environ.get("APPDATA", ""),
                os.environ.get("LOCALAPPDATA", ""),
                os.path.join(os.environ.get("USERPROFILE", ""), "Desktop"),
                os.path.join(os.environ.get("USERPROFILE", ""), "Documents"),
            ]
        elif system == "Linux":
            return ["/tmp", "/var/tmp", os.path.expanduser("~")]
        elif system == "Darwin":
            return ["/tmp", "/var/tmp", os.path.expanduser("~")]
        return ["/tmp"]

    def start(self) -> None:
        """Start filesystem monitoring."""
        handler = self._Handler(self.recorder)
        for path in self.watch_paths:
            if path and os.path.isdir(path):
                try:
                    self._observer.schedule(handler, path, recursive=True)
                    logger.info("Watching filesystem: %s", path)
                except Exception as e:
                    logger.warning("Cannot watch %s: %s", path, e)
        self._observer.start()
        logger.info("Filesystem monitor started")

    def stop(self) -> None:
        """Stop filesystem monitoring."""
        self._observer.stop()
        self._observer.join(timeout=5)
        logger.info("Filesystem monitor stopped")


class NetworkMonitor:
    """Monitor network connections system-wide."""

    def __init__(self, recorder):
        self.recorder = recorder
        self._running = False
        self._seen_connections = set()

    def _connection_key(self, conn):
        """Create a unique key for a connection."""
        laddr = f"{conn.laddr.ip}:{conn.laddr.port}" if conn.laddr else "none"
        raddr = f"{conn.raddr.ip}:{conn.raddr.port}" if conn.raddr else "none"
        return (conn.pid, laddr, raddr, conn.status)

    def _monitor_loop(self):
        """Main monitoring loop."""
        while self._running:
            try:
                connections = psutil.net_connections(kind="inet")
                for conn in connections:
                    key = self._connection_key(conn)
                    if key not in self._seen_connections:
                        self._seen_connections.add(key)
                        details = {
                            "pid": conn.pid,
                            "family": str(conn.family),
                            "type": str(conn.type),
                            "status": conn.status,
                            "local_address": (
                                f"{conn.laddr.ip}:{conn.laddr.port}"
                                if conn.laddr
                                else None
                            ),
                            "remote_address": (
                                f"{conn.raddr.ip}:{conn.raddr.port}"
                                if conn.raddr
                                else None
                            ),
                        }
                        # Try to get process name
                        if conn.pid:
                            try:
                                proc = psutil.Process(conn.pid)
                                details["process_name"] = proc.name()
                            except (psutil.NoSuchProcess, psutil.AccessDenied):
                                details["process_name"] = "<unknown>"
                        self.recorder.record("network", "connection", details)
            except (psutil.AccessDenied, OSError) as e:
                logger.debug("Network monitoring access issue: %s", e)
            time.sleep(1)

    def start(self) -> None:
        """Start network monitoring."""
        self._running = True
        self._thread = threading.Thread(target=self._monitor_loop, daemon=True)
        self._thread.start()
        logger.info("Network monitor started")

    def stop(self) -> None:
        """Stop network monitoring."""
        self._running = False
        if hasattr(self, "_thread"):
            self._thread.join(timeout=5)
        logger.info("Network monitor stopped")


def generate_timeline(results_dir, output_format="json") -> dict:
    """Generate a consolidated timeline from collected monitoring data."""
    results_path = Path(results_dir)
    all_events = []

    # Load all event files
    for json_file in results_path.glob("events_*.json"):
        try:
            with open(json_file, "r") as f:
                data = json.load(f)
                if isinstance(data, list):
                    all_events.extend(data)
                elif isinstance(data, dict) and "events" in data:
                    all_events.extend(data["events"])
        except (json.JSONDecodeError, IOError) as e:
            logger.warning("Error reading %s: %s", json_file, e)

    if not all_events:
        logger.error("No events found in %s", results_dir)
        return None

    # Sort by timestamp
    all_events.sort(key=lambda e: e.get("timestamp", ""))

    # Build timeline
    timeline = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "total_events": len(all_events),
        "duration_seconds": (
            all_events[-1].get("elapsed_seconds", 0)
            if all_events
            else 0
        ),
        "summary": {},
        "events": all_events,
    }

    # Summarize by category
    categories = defaultdict(lambda: defaultdict(int))
    for event in all_events:
        categories[event["category"]][event["type"]] += 1
    timeline["summary"] = {
        cat: dict(types) for cat, types in categories.items()
    }

    # Output
    output_file = results_path / f"timeline.{output_format}"
    if output_format == "json":
        with open(output_file, "w") as f:
            json.dump(timeline, f, indent=2)
    elif output_format == "csv":
        import csv

        with open(output_file, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["timestamp", "elapsed_seconds", "category", "type", "details"])
            for event in all_events:
                writer.writerow([
                    event.get("timestamp", ""),
                    event.get("elapsed_seconds", ""),
                    event.get("category", ""),
                    event.get("type", ""),
                    json.dumps(event.get("details", {})),
                ])
    elif output_format == "text":
        with open(output_file, "w") as f:
            f.write(f"=== Malware Analysis Timeline ===\n")
            f.write(f"Generated: {timeline['generated_at']}\n")
            f.write(f"Total Events: {timeline['total_events']}\n")
            f.write(f"Duration: {timeline['duration_seconds']}s\n\n")
            for event in all_events:
                f.write(
                    f"[{event.get('elapsed_seconds', '?'):>8.3f}s] "
                    f"{event['category']:>12s} | {event['type']:>20s} | "
                    f"{json.dumps(event.get('details', {}))}\n"
                )

    logger.info("Timeline written to %s", output_file)
    return timeline


def run_monitor(args) -> dict:
    """Main monitoring execution."""
    recorder = EventRecorder()
    monitors = []
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Set up process monitor
    proc_monitor = ProcessMonitor(
        recorder, target_pid=args.pid, watch_new=args.watch_new
    )
    monitors.append(proc_monitor)

    # Set up filesystem monitor
    if HAS_WATCHDOG and not args.no_filesystem:
        watch_paths = args.watch_paths if args.watch_paths else None
        fs_monitor = FilesystemMonitor(recorder, watch_paths=watch_paths)
        monitors.append(fs_monitor)

    # Set up network monitor
    if not args.no_network:
        net_monitor = NetworkMonitor(recorder)
        monitors.append(net_monitor)

    # Handle graceful shutdown
    shutdown_event = threading.Event()

    def signal_handler(signum, frame) -> None:
        logger.info("Shutdown signal received")
        shutdown_event.set()

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    # Start all monitors
    for monitor in monitors:
        try:
            monitor.start()
        except Exception as e:
            logger.error("Failed to start %s: %s", monitor.__class__.__name__, e)

    logger.info(
        "Monitoring started. Duration: %s seconds. Press Ctrl+C to stop.",
        args.duration or "unlimited",
    )

    # Wait for duration or shutdown signal
    start_time = time.time()
    while not shutdown_event.is_set():
        if args.duration and (time.time() - start_time) >= args.duration:
            logger.info("Duration reached, stopping monitors")
            break
        shutdown_event.wait(timeout=1)

    # Stop all monitors
    for monitor in reversed(monitors):
        try:
            monitor.stop()
        except Exception as e:
            logger.error("Error stopping %s: %s", monitor.__class__.__name__, e)

    # Save events
    events = recorder.get_events()
    session_id = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
    events_file = output_dir / f"events_{session_id}.json"

    output_data = {
        "session_id": session_id,
        "platform": platform.system(),
        "start_time": datetime.fromtimestamp(
            start_time, tz=timezone.utc
        ).isoformat(),
        "duration_seconds": round(time.time() - start_time, 2),
        "summary": recorder.get_summary(),
        "events": events,
    }

    with open(events_file, "w") as f:
        json.dump(output_data, f, indent=2)

    logger.info("Saved %d events to %s", len(events), events_file)
    logger.info("Summary: %s", json.dumps(recorder.get_summary()))

    return output_data


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Sandbox Monitor - Process, filesystem, and network monitoring for malware analysis",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s --pid 1234 --output ./results --duration 300
  %(prog)s --watch-new --output ./results --duration 600
  %(prog)s --timeline ./results --format json
        """,
    )

    mode_group = parser.add_mutually_exclusive_group(required=True)
    mode_group.add_argument(
        "--input", "--pid", dest="pid", type=int, help="PID of the process to monitor"
    )
    mode_group.add_argument(
        "--watch-new",
        action="store_true",
        help="Watch for newly created processes",
    )
    mode_group.add_argument(
        "--timeline",
        metavar="DIR",
        help="Generate timeline from results directory",
    )

    parser.add_argument(
        "--output", "-o", default="./monitor_results", help="Output directory for results"
    )
    parser.add_argument(
        "--duration",
        "-d",
        type=int,
        default=None,
        help="Monitoring duration in seconds (default: unlimited)",
    )
    parser.add_argument(
        "--format",
        choices=["json", "csv", "text"],
        default="json",
        help="Timeline output format (default: json)",
    )
    parser.add_argument(
        "--watch-paths",
        nargs="+",
        help="Filesystem paths to monitor (default: platform-specific)",
    )
    parser.add_argument(
        "--no-filesystem",
        action="store_true",
        help="Disable filesystem monitoring",
    )
    parser.add_argument(
        "--no-network",
        action="store_true",
        help="Disable network monitoring",
    )
    parser.add_argument(
        "--verbose", "-v", action="store_true", help="Enable verbose logging"
    )

    args = parser.parse_args()

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

    if args.timeline:
        result = generate_timeline(args.timeline, args.format)
        if result:
            print(json.dumps(result["summary"], indent=2))
        else:
            sys.exit(1)
    else:
        run_monitor(args)


if __name__ == "__main__":
    main()