watcher.py

"""Watch Security log for reg.exe hive exports and selected firewall events.

This is a merged watcher that handles registry hive exports (reg.exe save)
and the firewall policy events (4950, 4946, 4947, 4948). Detected incidents
are serialized to JSON (using `utils.json_formatter`) and uploaded via the
MongoDB uploader.
"""

from __future__ import annotations

import ctypes
import os
import shlex
import time
import sys
from ctypes import wintypes
from datetime import datetime, timezone, timedelta
from pathlib import Path
from typing import Optional, List, Dict, Any

import pywintypes
import win32evtlog
import win32com.client
import pythoncom
import argparse
import re

from utils.xml_parser import parse_event_xml
from utils.json_formatter import (
    build_registry_hive_json,
    build_firewall_json,
    build_priv_esc_json,
    build_generic_command_json,
    write_json_file,
)
from utils.firewall_parser import parse_firewall_event, format_firewall_summary
from utils.firewall_watcher import build_firewall_context
from utils.reg_parser import parse_reg_event
from utils.reg_parser import parse_reg_event
from utils.killchain import build_killchain
from utils.priv_esc_parser import parse_priv_esc_event
from uploaders.mongodb_uploader import upload_json_report
import subprocess
import hashlib
import http.client
import json

host_id = "67b3dba8bc6778101892eb77249db32e"
agent_id = "c4ca4238a0b923820dcc509a6f75849b"

# Event selection: include reg.exe 4688 events and firewall rule events
EVENT_FILTER = "Event[System[EventID=4688 or EventID=4950 or EventID=4946 or EventID=4947 or EventID=4948]]"
SECURITY_CHANNEL = "Security"
SYSMON_CHANNEL = "Microsoft-Windows-Sysmon/Operational"
POWER_SHELL_CHANNEL = "Microsoft-Windows-PowerShell/Operational"
TARGET_PROCESS = "reg.exe"
TRIGGER_KEYWORD = "save"
WATCH_INTERVAL_SECONDS = 1

FIREWALL_EVENT_IDS = {4950, 4946, 4947, 4948}

# Runtime flags
DEBUG = False

# Mapping and helpers adapted from network_attack.py
TARGET_EVENTS = {
    4950: "SETTING CHANGED",
    4946: "RULE ADDED",
    4947: "RULE MODIFIED",
    4948: "RULE DELETED",
}


def _should_upload_firewall_event(profile: Optional[str], new_state: Optional[bool]) -> bool:
    """Determine if a firewall state change should trigger an upload.
    
    Only upload if:
    1. This is a profile ON/OFF state change (not a rule modification)
    2. The state is actually different from the last recorded state
    3. Profile is one of the tracked profiles (Domain, Private, Public)
    
    Returns True if upload should proceed, False otherwise.
    """
    if not profile or profile not in g_firewall_state:
        return False
    
    if new_state is None:
        return False
    
    previous_state = g_firewall_state[profile]
    
    # If we haven't seen this profile before, record the state but don't upload
    if previous_state is None:
        g_firewall_state[profile] = new_state
        print(f"[DEBUG] Initialized firewall state for {profile}: {new_state}")
        return False
    
    # Check if state actually changed
    if previous_state == new_state:
        print(f"[DEBUG] Firewall state for {profile} unchanged: {new_state}")
        return False
    
    # State changed! Update and allow upload
    g_firewall_state[profile] = new_state
    print(f"[DEBUG] Firewall state change detected for {profile}: {previous_state} -> {new_state}")
    return True


def _extract_firewall_profile_state(context: dict) -> tuple[Optional[str], Optional[bool]]:
    """Extract the profile name and enabled/disabled state from firewall event context.
    
    Returns (profile_name, is_enabled) or (None, None) if not detectable.
    """
    automated_response = context.get("automated_response", {})
    profile = automated_response.get("profile")
    value = automated_response.get("value")
    
    if not profile:
        return None, None
    
    # Check if this is a profile enabled/disabled change
    if value is not None:
        value_lower = str(value).lower()
        if "disable" in value_lower or "false" in value_lower or value_lower == "0":
            return profile, False
        elif "enable" in value_lower or "true" in value_lower or value_lower == "1":
            return profile, True
    
    return None, None


def _is_profile_state_event(context: dict) -> bool:
    """Check if this firewall event is about profile state (on/off), not rule changes."""
    automated_response = context.get("automated_response", {})
    event_type = automated_response.get("type", "")
    
    # Only process "SETTING CHANGED" events for profile state changes
    if "SETTING CHANGED" not in event_type:
        return False
    
    setting = automated_response.get("setting", "")
    return "enabled" in str(setting).lower() or "state" in str(setting).lower()



    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except Exception:
        return False


def check_audit_policy() -> None:
    try:
        output = subprocess.check_output(
            'auditpol /get /subcategory:"MPSSVC Rule-Level Policy Change"',
            shell=True,
        ).decode("utf-8", errors="ignore")
        if "No Auditing" in output:
            print("[WARN] 'Audit MPSSVC Rule-Level Policy Change' appears disabled.")
            print(
                '    Run as Admin to enable: auditpol /set /subcategory:"MPSSVC Rule-Level Policy Change" /success:enable /failure:enable'
            )
        else:
            print("[INFO] Firewall audit policy appears active.")
    except Exception as e:
        print(f"[INFO] Could not verify audit policy automatically: {e}")


def _get_firewall_message(evt: dict) -> str:
    """Try to extract a concise human-readable message for firewall events from parsed XML."""
    system = evt.get("Event", {}).get("System", {})
    data = evt.get("Event", {}).get("EventData", {})
    try:
        eid = int(system.get("EventID"))
    except Exception:
        return "Unknown firewall event"
    # Build a lowercase-key lookup for fuzzy matching
    lower = {k.lower(): v for k, v in (data or {}).items()}

    def find_by_substr(subs: list[str]) -> Optional[str]:
        for s in subs:
            for k, v in lower.items():
                if s in k and v:
                    return v
        return None

    # 4950 typically signals a profile/setting change; look for keys containing
    # 'profile', 'setting', 'value', 'type', etc.
    if eid == 4950:
        profile = (
            find_by_substr(["profile"]) or find_by_substr(["profilechanged"]) or ""
        )
        # Setting field names vary: SettingType, SettingName, etc.
        setting = find_by_substr(["setting", "settingtype"]) or ""
        value = find_by_substr(["value", "settingvalue"]) or ""
        if profile or setting or value:
            return f"Profile: {profile} | Setting: {setting} | Value: {value}"

    # Rule-add/modify/delete events — try to find anything with 'rule' or 'name'
    if eid in (4946, 4947, 4948):
        rule = find_by_substr(["rule", "name"]) or ""
        if rule:
            return f"Rule Name: {rule}"

    # Generic fallback — include up to 3 non-empty values from EventData
    if data:
        vals = [v for v in data.values() if v]
        if vals:
            return f"Details: {' | '.join(vals[:3])}"
    return "No details available"


# Mitigation Constants
PROCESS_TERMINATE = 0x0001
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
kernel32.OpenProcess.restype = wintypes.HANDLE
kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
kernel32.TerminateProcess.restype = wintypes.BOOL
kernel32.TerminateProcess.argtypes = [wintypes.HANDLE, wintypes.UINT]
kernel32.CloseHandle.restype = wintypes.BOOL
kernel32.CloseHandle.argtypes = [wintypes.HANDLE]

_WMI_SERVICE = None
_WMI_FAILED = False

# Global state
g_subscription_handles: list = []
g_min_record_id: Optional[int] = None
# Simple cache of recent process metadata keyed by PID to capture short-lived CLI commands
g_process_cache: dict[int, dict] = {}
# Track processed event record IDs to prevent duplicate uploads
g_processed_event_ids: set = set()
# Track firewall state (profile -> on/off status) for 3 profiles: Domain, Private, Public
# Scenarios: each profile can be on or off (2^3 = 8 possible states, but we focus on transitions)
g_firewall_state: Dict[str, bool] = {
    "Domain": None,    # None = unknown, True = on, False = off
    "Private": None,
    "Public": None,
}


def _safe_evt_close(handle) -> None:
    close_fn = getattr(win32evtlog, "EvtClose", None)
    if not close_fn or not handle:
        return
    try:
        close_fn(handle)
    except pywintypes.error:
        pass


def _utc_now_iso() -> str:
    return (
        datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z")
    )


def _parse_pid(pid_hex: Optional[str]) -> Optional[int]:
    if not pid_hex:
        return None
    try:
        return int(pid_hex, 0)
    except (TypeError, ValueError):
        return None


def _ensure_wmi_service():
    global _WMI_SERVICE, _WMI_FAILED
    if _WMI_FAILED:
        return None
    if _WMI_SERVICE is not None:
        return _WMI_SERVICE
    try:
        try:
            pythoncom.CoInitialize()
        except pythoncom.error:
            pass
        locator = win32com.client.Dispatch("WbemScripting.SWbemLocator")
        _WMI_SERVICE = locator.ConnectServer(".", "root\\cimv2")
        return _WMI_SERVICE
    except pywintypes.com_error as exc:
        print(f"  [WARN] Unable to connect to WMI for process lookup: {exc}")
        _WMI_FAILED = True
        return None


def _lookup_process_metadata(pid: Optional[int]) -> dict[str, str]:
    empty = {"command_line": "", "creation_time": "", "owner": ""}
    if pid is None:
        return empty
    # Check recent-process cache first (helpful for short-lived processes like netsh)
    try:
        cached = g_process_cache.get(int(pid))
        if cached:
            # treat entries younger than 5 minutes as valid
            if (time.time() - cached.get("cached_at", 0)) < 300:
                return {
                    "command_line": cached.get("command_line", ""),
                    "creation_time": cached.get("creation_time", ""),
                    "owner": cached.get("owner", ""),
                    "parent_pid": cached.get("parent_pid", None),
                    "executable": cached.get("executable", ""),
                }
    except Exception:
        pass
    service = _ensure_wmi_service()
    if not service:
        return empty
    try:
        results = service.ExecQuery(
            f"SELECT CommandLine, CreationDate, ParentProcessId, ExecutablePath FROM Win32_Process WHERE ProcessId = {pid}"
        )
    except pywintypes.com_error:
        return empty
    for proc in results:
        cmd = getattr(proc, "CommandLine", None) or ""
        creation_raw = getattr(proc, "CreationDate", None)
        creation_iso = _wmi_datetime_to_iso(creation_raw) if creation_raw else ""
        owner = ""
        try:
            owner_info = proc.GetOwner()
            if owner_info:
                username = owner_info[0]
                domain = owner_info[1]
                if username:
                    owner = f"{domain}\\{username}" if domain else username
        except (pywintypes.com_error, AttributeError, TypeError):
            owner = ""
        parent_pid = getattr(proc, "ParentProcessId", None)
        exe = getattr(proc, "ExecutablePath", None) or ""
        return {
            "command_line": cmd,
            "creation_time": creation_iso,
            "owner": owner,
            "parent_pid": parent_pid,
            "executable": exe,
        }
    return empty


def _cache_process_creation(data: dict, evt: dict) -> None:
    """Cache short-lived process creation metadata from a 4688 event.

    This extracts command line, parent PID, creation time and executable
    and stores a small entry in the global `g_process_cache` keyed by PID.
    """
    try:
        new_proc_hex = data.get("NewProcessId")
        new_pid = _parse_pid(new_proc_hex)
        if not new_pid:
            return
        cmdline = data.get("CommandLine") or ""
        raw_parent = data.get("ParentProcessId") or None
        try:
            parent_pid_val = int(raw_parent, 0) if raw_parent not in (None, "", "-") else None
        except Exception:
            parent_pid_val = None
        g_process_cache[int(new_pid)] = {
            "command_line": cmdline,
            "creation_time": (evt.get("Event", {}).get("System", {}).get("TimeCreated") or ""),
            "owner": "",
            "parent_pid": parent_pid_val,
            "executable": (data.get("NewProcessName") or "").lower(),
            "cached_at": time.time(),
        }
    except Exception:
        # Swallow any errors to avoid impacting event processing
        pass


def _resolve_parent_meta(data: dict, child_meta: dict) -> Optional[dict]:
    """Resolve parent process metadata given event data and child metadata.

    Returns a metadata dict for the parent process or None if not found.
    """
    parent_meta = None
    try:
        parent_pid_hex = data.get("ProcessId") or data.get("ParentProcessId")
        parent_pid_int = _parse_pid(parent_pid_hex)
        if not parent_pid_int:
            parent_pid_int = child_meta.get("parent_pid")
            try:
                parent_pid_int = int(parent_pid_int) if parent_pid_int not in (None, "", "-") else None
            except Exception:
                parent_pid_int = None
        if parent_pid_int:
            parent_meta = _lookup_process_metadata(parent_pid_int)
    except Exception:
        parent_meta = None
    return parent_meta


# NOTE: historical Security log polling removed by request. Only subscribed 4688 events
# populate `g_process_cache` and are used for correlation with firewall events.


def _wmi_datetime_to_iso(raw: Optional[str]) -> str:
    if not raw or len(raw) < 22:
        return ""
    try:
        base = raw[:14]
        micros = raw[15:21]
        offset_sign = raw[21]
        offset_minutes = int(raw[22:]) if len(raw) > 22 else 0
        dt = datetime.strptime(base, "%Y%m%d%H%M%S")
        if micros.isdigit():
            dt = dt.replace(microsecond=int(micros))
        minutes = offset_minutes if offset_sign == "+" else -offset_minutes
        tz = timezone(timedelta(minutes=minutes))
        return (
            dt.replace(tzinfo=tz)
            .astimezone(timezone.utc)
            .isoformat()
            .replace("+00:00", "Z")
        )
    except Exception:
        return ""


def _query_log_channel(channel: str, query: str) -> List[Dict]:
    """
    Query a specific Event Log channel with an XPath query and return parsed events.
    Returns list of parsed event dicts (via `parse_event_xml`).
    """
    results: List[Dict] = []
    try:
        handle = win32evtlog.EvtQuery(
            channel,
            win32evtlog.EvtQueryChannelPath | win32evtlog.EvtQueryReverseDirection,
            query,
        )
    except pywintypes.error:
        return []

    try:
        while True:
            events = win32evtlog.EvtNext(handle, 10)
            if not events:
                break
            for event_h in events:
                try:
                    xml = win32evtlog.EvtRender(event_h, win32evtlog.EvtRenderEventXml)
                    parsed = parse_event_xml(xml)
                    if parsed:
                        results.append(parsed)
                except Exception:
                    continue
    finally:
        _safe_evt_close(handle)
    return results


def _build_process_ancestry(
    start_pid: Optional[int], max_depth: int = 5
) -> List[Dict[str, Any]]:
    """
    Trace backwards through Security 4688 events to build a parent chain for `start_pid`.

    Returns a list of process metadata dicts (parent-first, closest ancestor earliest),
    each containing: pid (int), executable (path), command_line, creation_time, parent_pid.
    """
    if not start_pid:
        return []

    chain: List[Dict[str, Any]] = []
    trace_pid = start_pid
    for depth in range(max_depth):
        if not trace_pid:
            break

        pid_hex = hex(int(trace_pid)).lower()
        # XPath: find 4688 events where NewProcessId equals our pid hex
        query = f"*[System[EventID=4688]] and *[EventData[Data[@Name='NewProcessId']='{pid_hex}']]"
        events = _query_log_channel(SECURITY_CHANNEL, query)
        if not events:
            break

        evt = events[0]
        evt_data = evt.get("Event", {}).get("EventData", {})
        system = evt.get("Event", {}).get("System", {})

        new_proc_hex = evt_data.get("NewProcessId")
        parent_pid_hex = evt_data.get("ProcessId")
        try:
            new_pid = int(new_proc_hex, 0) if new_proc_hex else None
        except Exception:
            new_pid = None

        try:
            parent_pid = int(parent_pid_hex, 0) if parent_pid_hex else None
        except Exception:
            parent_pid = None

        chain.append(
            {
                "pid": new_pid or trace_pid,
                "executable": evt_data.get("NewProcessName") or "",
                "command_line": evt_data.get("CommandLine") or "",
                "creation_time": system.get("TimeCreated") or "",
                "parent_pid": parent_pid,
            }
        )

        trace_pid = parent_pid

    # Reverse chain so parent-most appears first (optional, but consistent)
    try:
        chain.reverse()
    except Exception:
        pass

    return chain


def _terminate_process(pid_hex: Optional[str], label: str) -> bool:
    if not pid_hex:
        return False
    try:
        pid = int(pid_hex, 0)
    except (TypeError, ValueError):
        print(f"  [WARN] Invalid PID value for {label}: {pid_hex}")
        return False

    handle = kernel32.OpenProcess(PROCESS_TERMINATE, False, pid)
    if not handle:
        err = ctypes.get_last_error()
        if err == 87:
            print(
                f"  [INFO] OpenProcess indicates {label} PID {pid} already exited (error 87)"
            )
            return True
        print(f"  [WARN] OpenProcess failed for PID {pid} ({label}), error {err}")
        return False

    try:
        if not kernel32.TerminateProcess(handle, 1):
            err = ctypes.get_last_error()
            print(f"  ⚠️ TerminateProcess failed for PID {pid} ({label}), error {err}")
            return False
        return True
    finally:
        kernel32.CloseHandle(handle)


def _add_termination_action(
    pid_hex: Optional[str], label: str, actions: list, timeline: list, event_label: Optional[str] = None
) -> None:
    """Attempt to terminate a process and record the action and timeline entries.

    pid_hex: hex string like '0x141c'
    label: human label used in prints and action names
    event_label: optional explicit timeline event string; if None, a sensible default is used.
    """
    result = _terminate_process(pid_hex, label)
    action_time = _utc_now_iso()
    actions.append(
        {
            "name": f"Terminate {label}",
            "status": "Success" if result else "Failed",
            "timestamp": action_time,
            "details": f"PID {pid_hex}",
        }
    )
    timeline.append(
        {
            "timestamp": action_time,
            "event": (event_label or (f"Terminated {label}" if result else f"Failed to terminate {label}")),
            "details": f"PID {pid_hex}",
        }
    )
    print(f"  [{'OK' if result else 'WARN'}] {'Terminated' if result else 'Failed to terminate'} {label} (PID {pid_hex})")


def _attempt_corrupt_output(output_path: Optional[str], actions: list, timeline: list) -> tuple[str, str]:
    """Attempt to corrupt the output path and record action/timeline entries.

    Returns a tuple of (corruption_status, corruption_detail).
    """
    if not output_path:
        return "Skipped", "No output path detected"
    success, detail = _corrupt_file(output_path)
    action_time = _utc_now_iso()
    actions.append(
        {
            "name": "Corrupt hive dump",
            "status": "Success" if success else "Failed",
            "timestamp": action_time,
            "details": detail,
        }
    )
    timeline.append(
        {
            "timestamp": action_time,
            "event": ("Corrupted hive dump" if success else "Failed to corrupt hive dump"),
            "details": detail,
        }
    )
    print(f"  [{'CORRUPTED' if success else 'WARN'}] Corruption attempt: {detail}")
    return ("Success" if success else "Failed", detail)


def _extract_output_path(command_line: str) -> Optional[str]:
    if not command_line:
        return None
    try:
        tokens = shlex.split(command_line, posix=False)
    except ValueError:
        tokens = command_line.split()
    if not tokens:
        return None
    lowered = [token.lower() for token in tokens]
    try:
        idx = lowered.index(TRIGGER_KEYWORD)
    except ValueError:
        idx = -1
    if idx != -1:
        if len(tokens) > idx + 2:
            return tokens[idx + 2].strip('"')
        if len(tokens) > idx + 1:
            return tokens[idx + 1].strip('"')
    return tokens[-1].strip('"') if tokens else None


def _corrupt_file(path_str: str) -> tuple[bool, str]:
    """Overwrite a file with random data to render it unusable."""
    if not path_str:
        return False, "No output path supplied"
    if path_str.startswith("\\"):
        return False, "Remote path detected; skipping corruption"

    file_path = Path(path_str)
    try:
        if not file_path.exists():
            return False, "File not found"
        if file_path.is_dir():
            return False, "Target is a directory"

        size = file_path.stat().st_size
        bytes_to_write = size if size > 0 else 1024
        chunk_size = 1024 * 1024

        with file_path.open("r+b") as handle:
            handle.seek(0)
            remaining = bytes_to_write
            while remaining > 0:
                this_chunk = min(chunk_size, remaining)
                handle.write(os.urandom(this_chunk))
                remaining -= this_chunk
            if bytes_to_write < size:
                handle.truncate(bytes_to_write)

        return True, f"Overwrote {file_path} with {bytes_to_write} bytes of random data"
    except FileNotFoundError:
        return False, "File not found"
    except PermissionError as exc:
        return False, f"Permission denied: {exc}"
    except OSError as exc:
        return False, f"Corruption failed: {exc}"


# parsing logic for reg and firewall moved to utils.reg_parser and utils.firewall_parser


def _execute_mitigation_actions(context: dict) -> tuple[list[dict], list[dict]]:
    reg_pid = context["reg_pid"]
    parent_pid = context["parent_pid"]
    output_path = context["output_path"]
    raw_command = context["command_line"]

    timeline: list[dict[str, str]] = [
        {
            "timestamp": context["process_creation_time"],
            "event": "reg.exe save command detected",
            "details": raw_command or "-",
        }
    ]
    actions: list[dict[str, str]] = []

    # Terminate reg.exe
    _add_termination_action(reg_pid, "reg.exe", actions, timeline, event_label="Terminated reg.exe")

    # Terminate parent process
    if parent_pid and parent_pid != reg_pid:
        _add_termination_action(parent_pid, "parent process", actions, timeline, event_label="Terminated parent process")

    # Corrupt output file (if present)
    if output_path:
        corruption_status, corruption_detail = _attempt_corrupt_output(output_path, actions, timeline)
    else:
        print("  [WARN] Unable to determine dump output path for corruption")
        corruption_status, corruption_detail = "Skipped", "No output path detected"

    context["corruption_status"] = corruption_status
    context["corruption_detail"] = corruption_detail
    # Build a kill_chain object using available lookup/ancestry helpers.
    try:
        context["kill_chain"] = build_killchain(context, _lookup_process_metadata, _build_process_ancestry, primary_detection="T1003.002")
    except Exception:
        context.setdefault("kill_chain", {})
    return actions, timeline


def _report_incident(context: dict, event_type: str = "reg") -> None:
    """
    Serialize incident `context` to JSON using the appropriate builder and upload it.

    event_type: 'reg' or 'firewall'
    """
    ts = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S")
    json_dir = Path("json_data")
    json_dir.mkdir(parents=True, exist_ok=True)

    # Choose builder and filename based on event type
    if event_type == "reg":
        filename = f"registry_hive_attempt_{ts}.json"
        payload = build_registry_hive_json(context)
    elif event_type == "firewall":
        record_id = context.get("event_record_id") or ts
        filename = f"firewall_event_{record_id}.json"
        payload = build_firewall_json(context)
    elif event_type == "priv":
        # Privilege escalation report
        record_id = context.get("event_record_id") or ts
        filename = f"priv_esc_{record_id}.json"
        payload = build_priv_esc_json(context)
    elif event_type == "generic":
        record_id = context.get("event_record_id") or ts
        filename = f"incident_{record_id}.json"
        # Use generic builder; allow caller to pass desired attack type in context['attack_type']
        attack_type = context.get("attack_type") or "Generic Command"
        payload = build_generic_command_json(context, attack_type=attack_type)
    elif event_type == "priv":
        # Privilege escalation report
        record_id = context.get("event_record_id") or ts
        filename = f"priv_esc_{record_id}.json"
        payload = build_priv_esc_json(context)
    else:
        # Fallback to generic firewall builder if unknown
        record_id = context.get("event_record_id") or ts
        filename = f"incident_{record_id}.json"
        payload = build_firewall_json(context)

    json_path = json_dir / filename

    # Build agent_data block and prepend it to the payload so it appears
    # before `incident_summary` in the written JSON.
    # Use a randomly generated MD5 for event_id (unique per uploaded event).
    try:
        rnd = os.urandom(16)
    except Exception:
        # fallback to time-based randomness
        rnd = f"{ts}-{time.time()}".encode("utf-8")
    md5sum = hashlib.md5(rnd).hexdigest()
    agent_data = {
        "event_id": md5sum,
        "host_id": host_id,
        "agent_id": agent_id,
    }

    # Prepend agent_data and kill_chain to the payload (preserve insertion order)
    kc = context.get("kill_chain") or {}
    final_payload = {"agent_data": agent_data, "kill_chain": kc, **payload}

    # Extract command from incident_summary (or from context/process details) and ask classifier endpoint
    try:
        incident_summary = final_payload.get("incident_summary", {})
        cmd_to_send = ""
        if isinstance(incident_summary, dict):
            # try common keys
            cmd_to_send = incident_summary.get("Command") or incident_summary.get("Command Line") or ""

        # If no command in summary, try to extract from context (firewall events often don't include 'Command')
        if not cmd_to_send:
            # direct context keys
            cmd_to_send = context.get("command_line") or context.get("command") or ""
        if not cmd_to_send:
            # firewall initiator block
            fw = context.get("firewall") or {}
            if isinstance(fw, dict):
                initiator = fw.get("initiator")
                if isinstance(initiator, dict):
                    cmd_to_send = initiator.get("command_line") or initiator.get("command") or ""
        if not cmd_to_send:
            # inspect processes list from context
            procs = context.get("processes") or []
            if isinstance(procs, list):
                for p in procs:
                    if not isinstance(p, dict):
                        continue
                    cmd = p.get("command") or p.get("Command Line") or p.get("command_line")
                    if cmd:
                        cmd_to_send = cmd
                        break
        if not cmd_to_send:
            # inspect final_payload process_details (could be dict or list)
            pd = final_payload.get("process_details") or []
            if isinstance(pd, dict):
                for role, rec in pd.items():
                    if not isinstance(rec, dict):
                        continue
                    cmd = rec.get("command_line") or rec.get("Command Line") or rec.get("command")
                    if cmd:
                        cmd_to_send = cmd
                        break
            elif isinstance(pd, list):
                for rec in pd:
                    if not isinstance(rec, dict):
                        continue
                    cmd = rec.get("Command Line") or rec.get("Command") or rec.get("command_line")
                    if cmd:
                        cmd_to_send = cmd
                        break

        if cmd_to_send:
            try:
                conn = http.client.HTTPConnection("138.197.166.142", 80, timeout=5)
                # Send the command as the raw request body (plain text)
                conn.request("POST", "/predict", body=str(cmd_to_send), headers={"Content-Type": "text/plain"})
                resp = conn.getresponse()
                resp_body = resp.read().decode("utf-8", errors="ignore")
                try:
                    classifier_json = json.loads(resp_body)
                    # Response expected like: {"confidence score":0.999,"label":"Benign"}
                    label = classifier_json.get("label") or classifier_json.get("Label") or ""
                    confidence = (
                        classifier_json.get("confidence score")
                        if "confidence score" in classifier_json
                        else classifier_json.get("confidence")
                    )
                    # coerce confidence to a number when possible
                    try:
                        confidence = float(confidence)
                    except Exception:
                        # keep as-is (likely 0) if coercion fails
                        pass

                    # Inject results into the incident summary mapping
                    if isinstance(final_payload.get("incident_summary"), dict):
                        final_payload["incident_summary"]["MITRE ATT&CK"] = label
                        final_payload["incident_summary"]["Confidence"] = confidence
                except Exception:
                    print("  [WARN] Failed to parse classifier response")
            except Exception as exc:
                print(f"  [WARN] Failed to contact classifier endpoint: {exc}")
            finally:
                try:
                    conn.close()
                except Exception:
                    pass
    except Exception:
        # Ensure that any errors here do not prevent the report from being written/uploaded
        pass

    write_json_file(final_payload, json_path)
    print(f"  [INFO] JSON report generated: {json_path}")
    # If running in debug mode, skip the upload step so the user can inspect
    # generated JSON files without attempting to contact MongoDB.
    try:
        if DEBUG:
            print("  [DEBUG] Debug mode enabled - skipping upload")
            return
    except Exception:
        # If DEBUG isn't available for any reason, fall back to uploading.
        pass

    uploaded_id = upload_json_report(str(json_path))
    if uploaded_id:
        print("  [OK] JSON report uploaded to MongoDB (id=", uploaded_id, ")")
    else:
        print("  [WARN] JSON report upload failed")


def _handle_firewall_event(evt: dict) -> None:
    # Delegate building the firewall incident context to the dedicated parser
    try:
        context = build_firewall_context(
            evt, _lookup_process_metadata, g_process_cache, _build_process_ancestry
        )
    except Exception as exc:
        print(f"[ERROR] Failed to build firewall context: {exc}")
        return

    # Debug: show that we're about to report a firewall incident and a summary
    try:
        sys = evt.get("Event", {}).get("System", {})
        print(f"[DEBUG] Handling firewall event id={sys.get('EventID')} record={sys.get('EventRecordID')}")
    except Exception:
        pass

    # Check if this is a profile state change event and enforce state-change policy
    if _is_profile_state_event(context):
        profile, new_state = _extract_firewall_profile_state(context)
        if profile and new_state is not None:
            if not _should_upload_firewall_event(profile, new_state):
                print(f"[DEBUG] Firewall state event for {profile} does not meet upload criteria; skipping upload")
                return
    
    # Persist JSON and upload via centralized reporter
    _report_incident(context, event_type="firewall")


def _handle_event_callback(action, _context, event) -> None:
    if action == win32evtlog.EvtSubscribeActionError:
        print("[WARN] Subscription error while receiving events")
        return
    if action != win32evtlog.EvtSubscribeActionDeliver:
        return

    try:
        xml = win32evtlog.EvtRender(event, win32evtlog.EvtRenderEventXml)
    except pywintypes.error as exc:
        print("[ERROR] Failed to render event:", exc)
        return

    # Parse XML just to get EventRecordID for deduplication
    try:
        evt_temp = parse_event_xml(xml)
        system_temp = evt_temp.get("Event", {}).get("System", {})
        record_id = system_temp.get("EventRecordID")
        if record_id:
            if record_id in g_processed_event_ids:
                print(f"[DEBUG] Skipping duplicate event (EventRecordID={record_id})")
                return
            g_processed_event_ids.add(record_id)
    except Exception:
        pass

    # Defensive parsing: catch parsing errors and log the parsed System/EventData
    try:
        evt = parse_event_xml(xml)
    except Exception as exc:
        # Log a short excerpt of the XML to help debugging
        try:
            snippet = (xml[:200] + ""
            "...") if xml and len(xml) > 200 else xml
        except Exception:
            snippet = "<unprintable xml>"
        print(f"[ERROR] parse_event_xml failed: {exc} -- xml snippet: {snippet}")
        return

    if not evt:
        print("[WARN] parse_event_xml returned no event")
        return

    system = evt.get("Event", {}).get("System", {})
    eventdata = evt.get("Event", {}).get("EventData", {})
    event_id = None
    try:
        event_id = int(system.get("EventID"))
    except Exception:
        # Try to extract from EventData if missing
        pass
    
    channel = system.get("Channel") or ""
    is_firewall_event = event_id in FIREWALL_EVENT_IDS or (
        isinstance(channel, str) and "firewall" in channel.lower()
    )
    
    # Only print debug info for firewall events
    if is_firewall_event:
        try:
            print(f"[DEBUG] Parsed System fields: EventID={system.get('EventID')} EventRecordID={system.get('EventRecordID')} Channel={system.get('Channel')}")
            # Also show a short summary of EventData for troubleshooting
            if eventdata:
                sample = ", ".join([f"{k}={v}" for k, v in list(eventdata.items())[:4]])
                print(f"[DEBUG] EventData sample: {sample}")
        except Exception:
            pass
    
    # Delegate to specialized processors
    if event_id == 4688:
        print(f"[DEBUG] Calling _process_4688_event for EventID 4688")
        _process_4688_event(evt)
        return

    # Check for "reg save" detection across ALL event types and channels
    # This catches reg.exe from any parent process (cmd, PowerShell, etc.)
    try:
        eventdata_vals = []
        if isinstance(eventdata, dict):
            for v in eventdata.values():
                if v:
                    eventdata_vals.append(str(v))
        combined = " ".join(eventdata_vals).lower()
        
        # reg save - detect "reg save" or "reg.exe save" anywhere in eventdata
        if re.search(r'\breg(?:\.exe)?\s+save\b', combined, re.IGNORECASE):
            print(f"[DEBUG] Found 'reg save' in non-4688 event - combined text: {combined[:300]}")
            detection_time = system.get("TimeCreated") or _utc_now_iso()
            proc_id = None
            try:
                proc_id = int(system.get("ProcessID")) if system.get("ProcessID") else None
            except Exception:
                proc_id = None
            child_meta = _lookup_process_metadata(proc_id) if proc_id else {}
            
            # Extract output path - look for arguments after "reg save"
            path_match = re.search(r'\breg(?:\.exe)?\s+save\s+([A-Za-z0-9\\:\\_\\.-/]+\s+)?([A-Za-z0-9\\:\\_\\.-/:\\"]+)', combined, re.IGNORECASE)
            out_path = path_match.group(2) if path_match else None
            m = re.search(r'\breg(?:\.exe)?\s+save\b.*', combined, re.IGNORECASE)
            context = {
                "detection_time": detection_time,
                "process_creation_time": child_meta.get("creation_time") or detection_time,
                "user": child_meta.get("owner") or "",
                "domain": "",
                "process_creator": child_meta.get("owner") or "",
                "command_line": m.group(0) if m else combined,
                "output_path": out_path,
                "parent_process": child_meta.get("executable") or "",
                "parent_pid": child_meta.get("parent_pid") or "",
                "reg_pid": system.get("ProcessID"),
                "event_record_id": system.get("EventRecordID"),
                "processes": [
                    {
                        "role": "Child",
                        "path": child_meta.get("executable") or "reg.exe",
                        "pid": system.get("ProcessID"),
                        "parent_pid": child_meta.get("parent_pid") or "",
                        "command": m.group(0) if m else combined,
                        "creation_time": child_meta.get("creation_time") or "",
                    }
                ],
                "raw_event": evt,
            }
            # Execute mitigation actions and build timeline/kill_chain
            actions, timeline = _execute_mitigation_actions(context)
            context["actions"] = actions
            context["timeline"] = timeline
            _report_incident(context, event_type="reg")
            return
    except Exception as e:
        print(f"[ERROR] Exception in reg save detection: {e}")
        pass

    # Some Windows versions/providers emit firewall-related EventIDs in
    # the 2000-range (e.g., 2082) even though the provider channel is the
    # Firewall provider. Previously we only checked for a small set of
    # numeric IDs (4950/4946/4947/4948) which caused those events to be
    # ignored. If the event comes from the Firewall channel, treat it as
    # a firewall event regardless of numeric EventID so it will be parsed
    # and turned into a JSON report.
    try:
        if is_firewall_event:
            _handle_firewall_event(evt)
            return
    except Exception:
        # Fallback: if anything goes wrong with channel checks, avoid
        # crashing the callback
        pass

    # If event comes from the PowerShell operational channel, inspect
    # EventData for explicit script names (e.g. EnableAllTokenPrivs.ps1)
    try:
        if isinstance(channel, str) and "powershell" in channel.lower():
            eventdata_vals = []
            if isinstance(eventdata, dict):
                for v in eventdata.values():
                    if v:
                        eventdata_vals.append(str(v))
            combined = " ".join(eventdata_vals).lower()
            # Match script filename or SetTokenPriv/EnablePrivilege calls mentioning SeBackupPrivilege
            # also support bracket-call syntax like: [SetTokenPriv]::EnablePrivilege("SeBackupPrivilege")
            try:
                bracket_call = bool(re.search(r"\[\s*settokenpriv\s*\]\s*::\s*enableprivilege\s*\(.*sebackupprivilege.*\)", combined, re.IGNORECASE))
            except Exception:
                bracket_call = False

            # Quick detections for common commands in the test script
            # Build a small context helper
            detection_time = system.get("TimeCreated") or _utc_now_iso()
            user_from_event = None
            try:
                m_user = re.search(r"User\s*=\s*([^\n\r]+)", combined, re.IGNORECASE)
                if m_user:
                    user_from_event = m_user.group(1).strip()
            except Exception:
                user_from_event = None

            proc_id = None
            try:
                proc_id = int(system.get("ProcessID")) if system.get("ProcessID") else None
            except Exception:
                proc_id = None

            child_meta = _lookup_process_metadata(proc_id) if proc_id else {}

            # Patterns and handlers (ordered)
            # Set-MpPreference -DisableRealTimeMonitoring
            if re.search(r"\bSet-?MpPreference\b.*-DisableRealTimeMonitoring\b", combined, re.IGNORECASE):
                context = {
                    "detection_time": detection_time,
                    "process_creation_time": child_meta.get("creation_time") or detection_time,
                    "user": (user_from_event or child_meta.get("owner") or ""),
                    "domain": "",
                    "process_creator": child_meta.get("owner") or "",
                    "command_line": "Set-MpPreference -DisableRealTimeMonitoring",
                    "event_record_id": system.get("EventRecordID"),
                    "processes": [
                        {
                            "role": "Child",
                            "path": child_meta.get("executable") or "powershell",
                            "pid": system.get("ProcessID"),
                            "parent_pid": child_meta.get("parent_pid") or "",
                            "command": "Set-MpPreference -DisableRealTimeMonitoring",
                            "creation_time": child_meta.get("creation_time") or "",
                        }
                    ],
                    "raw_event": evt,
                    "summary": "Disable Windows Defender real-time monitoring",
                    "attack_type": "Impair Defenses",
                }
                _report_incident(context, event_type="generic")
                return

            # sc stop <service>
            m_sc = re.search(r"\bsc\s+stop\s+([A-Za-z0-9_-]+)", combined, re.IGNORECASE)
            if m_sc:
                svc = m_sc.group(1)
                context = {
                    "detection_time": detection_time,
                    "process_creation_time": child_meta.get("creation_time") or detection_time,
                    "user": (user_from_event or child_meta.get("owner") or ""),
                    "domain": "",
                    "process_creator": child_meta.get("owner") or "",
                    "command_line": m_sc.group(0),
                    "event_record_id": system.get("EventRecordID"),
                    "processes": [
                        {
                            "role": "Child",
                            "path": child_meta.get("executable") or "powershell",
                            "pid": system.get("ProcessID"),
                            "parent_pid": child_meta.get("parent_pid") or "",
                            "command": m_sc.group(0),
                            "creation_time": child_meta.get("creation_time") or "",
                        }
                    ],
                    "raw_event": evt,
                    "summary": f"Service stop requested: {svc}",
                    "attack_type": "Impair Defenses",
                }
                _report_incident(context, event_type="generic")
                return

            # klist purge / klist get
            if re.search(r"\bklist\s+purge\b", combined, re.IGNORECASE) or re.search(r"\bklist\s+get\b", combined, re.IGNORECASE):
                cmd = "klist purge" if "purge" in combined else "klist get"
                context = {
                    "detection_time": detection_time,
                    "process_creation_time": child_meta.get("creation_time") or detection_time,
                    "user": (user_from_event or child_meta.get("owner") or ""),
                    "domain": "",
                    "process_creator": child_meta.get("owner") or "",
                    "command_line": cmd,
                    "event_record_id": system.get("EventRecordID"),
                    "processes": [
                        {
                            "role": "Child",
                            "path": child_meta.get("executable") or "powershell",
                            "pid": system.get("ProcessID"),
                            "parent_pid": child_meta.get("parent_pid") or "",
                            "command": cmd,
                            "creation_time": child_meta.get("creation_time") or "",
                        }
                    ],
                    "raw_event": evt,
                    "summary": "Kerberos ticket operation detected",
                    "attack_type": "Access Token Manipulation",
                }
                _report_incident(context, event_type="generic")
                return

            # whoami /groups and Get-MpComputerStatus (benign/low priority) — still report for visibility
            if re.search(r"\bwhoami\b.*groups", combined, re.IGNORECASE) or re.search(r"\bGet-MpComputerStatus\b", combined, re.IGNORECASE):
                cmd = "whoami /groups" if "whoami" in combined else "Get-MpComputerStatus"
                context = {
                    "detection_time": detection_time,
                    "process_creation_time": child_meta.get("creation_time") or detection_time,
                    "user": (user_from_event or child_meta.get("owner") or ""),
                    "domain": "",
                    "process_creator": child_meta.get("owner") or "",
                    "command_line": cmd,
                    "event_record_id": system.get("EventRecordID"),
                    "processes": [
                        {
                            "role": "Child",
                            "path": child_meta.get("executable") or "powershell",
                            "pid": system.get("ProcessID"),
                            "parent_pid": child_meta.get("parent_pid") or "",
                            "command": cmd,
                            "creation_time": child_meta.get("creation_time") or "",
                        }
                    ],
                    "raw_event": evt,
                    "summary": "Reconnaissance/administrative command",
                    "attack_type": "Reconnaissance",
                }
                _report_incident(context, event_type="generic")
                return

            # Continue with previous PrivEsc detections
            if (
                "enablealltokenprivs.ps1" in combined
                or "\\enablealltokenprivs.ps1" in combined
                or "./enablealltokenprivs.ps1" in combined
                or ("settokenpriv" in combined and "sebackupprivilege" in combined)
                or ("enableprivilege" in combined and "sebackupprivilege" in combined)
                or bracket_call
                # relaxed fallback: match the privilege or enable call anywhere
                or ("sebackupprivilege" in combined)
                or ("enableprivilege(" in combined)
            ):
                # Build a best-effort context and report as privilege escalation
                raw_command = combined
                # Attempt to look up process metadata if ProcessID is available
                proc_id = None
                try:
                    proc_id = int(system.get("ProcessID")) if system.get("ProcessID") else None
                except Exception:
                    proc_id = None
                child_meta = _lookup_process_metadata(proc_id) if proc_id else {}
                parent_meta = _resolve_parent_meta(eventdata, child_meta)
                context = parse_priv_esc_event(evt, raw_command, child_meta or {}, parent_meta)
                context.setdefault("actions", [])
                context.setdefault("timeline", [{
                    "timestamp": context.get("process_creation_time"),
                    "event": "Privilege escalation script detected (PowerShell channel)",
                    "details": raw_command,
                }])
                _report_incident(context, event_type="priv")
                return
    except Exception:
        pass
    except Exception:
        # Fallback: if anything goes wrong with channel checks, avoid
        # crashing the callback
        pass


def _process_4688_event(evt: dict) -> None:
    """Process Security 4688 events: cache process metadata and handle reg.exe save."""
    data = evt.get("Event", {}).get("EventData", {})
    
    # Debug: Print EVERY 4688 with all fields
    print(f"\n[4688 EVENT] ===== Event Keys/Values =====")
    for key, value in data.items():
        if value:
            print(f"  {key}: {str(value)[:150]}")
    
    # Cache process creation metadata
    try:
        _cache_process_creation(data, evt)
    except Exception:
        pass

    # STEP 1: Extract commandline from EventData (try multiple field names)
    commandline = (
        data.get("CommandLine") 
        or data.get("ProcessCommandLine")
        or ""
    )
    
    print(f"[4688] Extracted CommandLine: '{commandline}'")
    
    # STEP 2: Check if "reg save" or "reg.exe save" appears ANYWHERE in the commandline
    if re.search(r'\breg(?:\.exe)?\s+save\b', commandline, re.IGNORECASE):
        print(f"[4688] *** DETECTED 'reg save' in commandline ***")
        print(f"[4688] Full command: {commandline}")
        
        # Extract metadata for reporting
        reg_pid_hex = data.get("NewProcessId")
        reg_pid_int = _parse_pid(reg_pid_hex)
        child_meta = _lookup_process_metadata(reg_pid_int)
        
        # Resolve parent metadata
        parent_meta = _resolve_parent_meta(data, child_meta)

        # Parse and report
        context = parse_reg_event(evt, commandline, child_meta, parent_meta)
        actions, timeline = _execute_mitigation_actions(context)
        context["actions"] = actions
        context["timeline"] = timeline
        _report_incident(context, event_type="reg")
        return

    # Detect PowerShell-based Privilege Escalation script execution
    # Look for PowerShell executables and the EnableAllTokenPrivs.ps1 script in the command line
    process_path = (data.get("NewProcessName") or "").lower()
    basename = process_path.rsplit("\\", 1)[-1]
    
    if any(p in basename for p in ("powershell.exe", "powershell", "pwsh.exe", "pwsh")):
        # Try to resolve process metadata similar to reg handling
        new_pid_hex = data.get("NewProcessId")
        new_pid_int = _parse_pid(new_pid_hex)
        child_meta = _lookup_process_metadata(new_pid_int)
        raw_command = data.get("CommandLine") or child_meta.get("command_line") or ""

        # Detect execution of the EnableAllTokenPrivs.ps1 script or API call (case-insensitive)
        if raw_command:
            lc_cmd = raw_command.lower()
            try:
                bracket_call_cmd = bool(re.search(r"\[\s*settokenpriv\s*\]\s*::\s*enableprivilege\s*\(.*sebackupprivilege.*\)", raw_command, re.IGNORECASE))
            except Exception:
                bracket_call_cmd = False

            match_priv = (
                ("enablealltokenprivs.ps1" in lc_cmd)
                or ("settokenpriv" in lc_cmd and "sebackupprivilege" in lc_cmd)
                or ("enableprivilege" in lc_cmd and "sebackupprivilege" in lc_cmd)
                or bracket_call_cmd
                # relaxed fallback: match the privilege name or enable call anywhere
                or ("sebackupprivilege" in lc_cmd)
                or ("enableprivilege(" in lc_cmd)
            )
        else:
            match_priv = False

        if match_priv:
            parent_meta = _resolve_parent_meta(data, child_meta)
            context = parse_priv_esc_event(evt, raw_command, child_meta, parent_meta)
            # No automated mitigation implemented for this detection yet; include empty actions/timeline
            context.setdefault("actions", [])
            context.setdefault("timeline", [{
                "timestamp": context.get("process_creation_time"),
                "event": "Privilege escalation script detected",
                "details": raw_command,
            }])
            _report_incident(context, event_type="priv")
            return

    # Not a detection we care about
    return


def start_subscription() -> bool:
    global g_min_record_id, g_subscription_handles, g_processed_event_ids, g_firewall_state
    g_min_record_id = 0
    g_subscription_handles = []
    g_processed_event_ids = set()  # Clear deduplication set on fresh start
    g_firewall_state = {"Domain": None, "Private": None, "Public": None}  # Reset firewall state
    print("[INFO] Subscribing in real-time to Security and Firewall channels...")
    # Subscribe to Security channel (no server-side filter) so we can filter in callback
    try:
        h1 = win32evtlog.EvtSubscribe(
            SECURITY_CHANNEL,
            win32evtlog.EvtSubscribeToFutureEvents,
            None,
            _handle_event_callback,
            None,
        )
        g_subscription_handles.append(h1)
        try:
            print(f"[DEBUG] Security subscription handle: {h1}")
        except Exception:
            pass
    except pywintypes.error as exc:
        print(f"ERROR: Unable to subscribe to Security log: {exc}")

    # Subscribe to Firewall provider channel (no server-side filter)
    try:
        h2 = win32evtlog.EvtSubscribe(
            "Microsoft-Windows-Windows Firewall With Advanced Security/Firewall",
            win32evtlog.EvtSubscribeToFutureEvents,
            None,
            _handle_event_callback,
            None,
        )
        g_subscription_handles.append(h2)
        try:
            print(f"[DEBUG] Firewall subscription handle: {h2}")
        except Exception:
            pass
    except pywintypes.error as exc:
        print(f"WARN: Unable to subscribe to Firewall channel: {exc}")

    # Subscribe to PowerShell operational channel to capture script/execution events
    try:
        h3 = win32evtlog.EvtSubscribe(
            POWER_SHELL_CHANNEL,
            win32evtlog.EvtSubscribeToFutureEvents,
            None,
            _handle_event_callback,
            None,
        )
        g_subscription_handles.append(h3)
        try:
            print(f"[DEBUG] PowerShell subscription handle: {h3}")
        except Exception:
            pass
    except pywintypes.error as exc:
        print(f"WARN: Unable to subscribe to PowerShell channel: {exc}")

    if not g_subscription_handles:
        print("ERROR: No subscriptions could be created")
        return False
    return True


def stop_subscription() -> None:
    global g_subscription_handles
    if g_subscription_handles:
        for h in g_subscription_handles:
            try:
                _safe_evt_close(h)
            except Exception:
                pass
        g_subscription_handles = []


def run_watcher() -> None:
    if not start_subscription():
        return
    try:
        while True:
            time.sleep(WATCH_INTERVAL_SECONDS)
    except KeyboardInterrupt:
        print("Stopping watcher (KeyboardInterrupt)")
    finally:
        stop_subscription()


def main() -> None:
    global DEBUG
    parser = argparse.ArgumentParser(description="Watcher for registry and firewall events")
    parser.add_argument("-d", "--debug", action="store_true", help="Enable debug mode (skip uploads)")
    args = parser.parse_args()
    DEBUG = bool(args.debug)
    if DEBUG:
        print("[INFO] Starting watcher in DEBUG mode: uploads will be skipped")
    run_watcher()


if __name__ == "__main__":
    main()