prefilter.py

#!/usr/bin/env python3
"""
Cthaeh Pre-filter - Fast PE import check before Ghidra analysis

Uses pefile to quickly check driver imports. Skips drivers that lack
interesting attack surface (no IOCTL handling, no device creation).
Runs in milliseconds per driver vs minutes for Ghidra.

Requires: pip install pefile

Optional: pip install requests (for LOLDrivers cross-reference)
"""

import os
import sys
import json
import time
import hashlib
from concurrent.futures import ThreadPoolExecutor, as_completed

# Boot-phase awareness (Jiří Vinopal's EDR Phase 0 blind spots research)
# Boot-start drivers (Start=0) load before ANY EDR kernel callbacks.
# System-start drivers (Start=1) load before most EDR user-mode components.
try:
    import winreg
    HAS_WINREG = True
except ImportError:
    HAS_WINREG = False

try:
    import pefile
except ImportError:
    print("ERROR: pefile not installed. Run: pip install pefile")
    sys.exit(1)


# --- WDAC Block Policy ---
POLICIES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "policies")

def load_wdac_block_hashes():
    """Load SHA256 (and SHA1) deny hashes from WDAC block policy JSONs."""
    hashes = set()
    for fname in ("Win10_MicrosoftDriverBlockPolicy.json", "Win11_MicrosoftDriverBlockPolicy.json"):
        fpath = os.path.join(POLICIES_DIR, fname)
        if not os.path.exists(fpath):
            continue
        try:
            with open(fpath, "r") as f:
                data = json.load(f)
            for rule in data.get("file_rules", []):
                if rule.get("action") == "deny" and "hash" in rule:
                    hashes.add(rule["hash"].lower())
        except Exception:
            pass
    return hashes


def load_wdac_filename_rules():
    """Load filename+version deny rules from WDAC block policy JSONs."""
    rules = []  # list of (filename_lower, max_version_str)
    for fname in ("Win10_MicrosoftDriverBlockPolicy.json", "Win11_MicrosoftDriverBlockPolicy.json"):
        fpath = os.path.join(POLICIES_DIR, fname)
        if not os.path.exists(fpath):
            continue
        try:
            with open(fpath, "r") as f:
                data = json.load(f)
            for rule in data.get("file_rules", []):
                if rule.get("action") == "deny" and "file_name" in rule:
                    rules.append((
                        rule["file_name"].lower(),
                        rule.get("maximum_file_version", "65535.65535.65535.65535"),
                    ))
        except Exception:
            pass
    return rules


def _parse_version(v):
    """Parse dotted version string into tuple of ints for comparison."""
    try:
        return tuple(int(x) for x in v.split("."))
    except (ValueError, AttributeError):
        return (0,)


def load_holygrail_loldrivers():
    """Load LOLDrivers SHA256 hashes from HolyGrail's lol_drivers.json."""
    lol_path = os.path.join(POLICIES_DIR, "lol_drivers.json")
    hashes = {}  # sha256_lower -> driver tag/name
    if not os.path.exists(lol_path):
        return hashes
    try:
        with open(lol_path, "r") as f:
            data = json.load(f)
        for entry in data:
            tag = entry.get("Tags", ["unknown"])[0] if entry.get("Tags") else "unknown"
            for sample in entry.get("KnownVulnerableSamples", []):
                sha = sample.get("SHA256", "")
                if sha:
                    hashes[sha.lower()] = tag
    except Exception:
        pass
    return hashes


# --- Boot Phase Awareness ---
# Reference: Jiří Vinopal's research on EDR Phase 0 blind spots.
# During boot Phase 0/1, only the SYSTEM hive is loaded. Drivers accessing
# HKLM\SOFTWARE get NAME NOT FOUND, revealing they operate in an unmonitored window.

BOOT_PHASE_NAMES = {
    0: "BOOT_START",
    1: "SYSTEM_START",
    2: "AUTO_START",
    3: "MANUAL",
    4: "DISABLED",
}

BOOT_PHASE_BONUS = {
    0: 15,  # Phase 0: EDR completely blind — no callbacks, no hooks, no telemetry
    1: 10,  # Phase 1: EDR kernel loading but user-mode agent not ready
}

BOOT_LOG_BLIND_SPOT_BONUS = 5  # Confirmed SOFTWARE hive access during boot


def get_driver_start_type(driver_path):
    """Look up driver start type from Windows registry.

    Uses the driver filename (minus extension) to query:
    HKLM\\SYSTEM\\CurrentControlSet\\Services\\{name}\\Start

    Returns: (start_type: int or None, phase_name: str)
    """
    if not HAS_WINREG:
        return None, "UNKNOWN"

    svc_name = os.path.splitext(os.path.basename(driver_path))[0]

    try:
        key = winreg.OpenKey(
            winreg.HKEY_LOCAL_MACHINE,
            f"SYSTEM\\CurrentControlSet\\Services\\{svc_name}"
        )
        start_val, _ = winreg.QueryValueEx(key, "Start")
        winreg.CloseKey(key)
        return start_val, BOOT_PHASE_NAMES.get(start_val, f"UNKNOWN({start_val})")
    except Exception:
        return None, "UNKNOWN"


def load_boot_log(boot_log_path):
    """Load boot_analyzer.py JSON output for cross-reference.

    Returns: dict of driver_name_lower -> boot log entry
    """
    if not boot_log_path or not os.path.exists(boot_log_path):
        return {}
    try:
        with open(boot_log_path, "r") as f:
            data = json.load(f)
        return {
            entry["driver"].lower(): entry
            for entry in data.get("drivers", [])
        }
    except Exception:
        return {}


# Imports that indicate interesting attack surface
INTERESTING_IMPORTS = {
    # Device creation (required for user-accessible attack surface)
    "IoCreateDevice",
    "IoCreateDeviceSecure",
    "WdfDeviceCreate",
    # IRP handling (IOCTL attack surface)
    "IofCompleteRequest",
    "IoCompleteRequest",
    # WMI (additional attack surface)
    "IoWMIRegistrationControl",
}

# Imports that indicate higher risk
HIGH_RISK_IMPORTS = {
    "MmMapIoSpace",
    "MmMapLockedPagesSpecifyCache",
    "MmMapLockedPagesWithReservedMapping",
    "ZwMapViewOfSection",
    "ExAllocatePool",
    "ExAllocatePoolWithTag",
    "ExAllocatePool2",
    # Physical/MDL (from HolyGrail)
    "MmGetPhysicalAddress",
    "MmCopyMemory",
    "MmCopyVirtualMemory",
    "MmAllocatePagesForMdl",
    "IoAllocateMdl",
    # Section/VM (from HolyGrail)
    "ZwOpenSection",
    "ZwReadVirtualMemory",
    "ZwWriteVirtualMemory",
    # Process (from HolyGrail)
    "KeStackAttachProcess",
}

# User-mode communication bridge primitives
# Drivers WITH comms capability are more interesting (attackable from userspace)
COMMS_IMPORTS = {
    "IoCreateDevice",
    "IoCreateSymbolicLink",
    "FltRegisterFilter",
    "FltCreateCommunicationPort",
    "IofCompleteRequest",
}

# BYOVD process killer pairs - if a driver imports BOTH an opener and terminator,
# it can be weaponized to kill AV/EDR processes
BYOVD_OPENERS = {
    "ZwOpenProcess",
    "NtOpenProcess",
    "ObOpenObjectByPointer",
    "PsLookupProcessByProcessId",
}

BYOVD_TERMINATORS = {
    "ZwTerminateProcess",
    "NtTerminateProcess",
}

# Physical memory R/W pairs - drivers with both mapping + view = potential phys mem access
PHYS_MEM_INDICATORS = {
    "MmMapIoSpace",
    "ZwMapViewOfSection",
    "MmMapLockedPagesSpecifyCache",
    "ZwOpenSection",
    "ZwOpenPhysicalMemory",  # rare but critical
}

# Token stealing / EPROCESS manipulation indicators
TOKEN_STEAL_IMPORTS = {
    "PsLookupProcessByProcessId",
    "PsReferencePrimaryToken",
    "SePrivilegeCheck",
    "ZwOpenProcessTokenEx",
    "NtOpenProcessToken",
}

# Registry manipulation from kernel (persistence vector)
REGISTRY_IMPORTS = {
    "ZwCreateKey",
    "ZwSetValueKey",
    "ZwOpenKey",
    "ZwDeleteKey",
}

# DSE bypass related strings (checked in string scan, not imports)
DSE_STRINGS = {
    "CI.dll",
    "g_CiOptions",
    "CiValidateImageHeader",
    "CiInitialize",
}

# WinIO/WinRing0 codebase indicators (strings)
WINIO_STRINGS = {
    "WinIo",
    "WinRing0",
    "\\Device\\WinIo",
    "\\DosDevices\\WinRing0",
    "\\Device\\WinRing0",
    "WINIO_MAPPHYSTOLIN",
}

# Firmware/SPI flash access indicators
FIRMWARE_IMPORTS = {
    "HalGetBusDataByOffset",
    "HalSetBusDataByOffset",
}

# Disk direct access strings
DISK_ACCESS_STRINGS = {
    "\\Device\\Harddisk",
    "PhysicalDrive",
    "RawDisk",
}

# Good security practice imports (negative scoring - reduce risk_hint)
GOOD_PRACTICE_IMPORTS = {
    "ProbeForRead",           # Proper input validation
    "ProbeForWrite",          # Proper input validation
    "IoCreateDeviceSecure",   # Secure device creation (vs plain IoCreateDevice)
    "SeAccessCheck",          # Authorization enforcement
    "SeSinglePrivilegeCheck", # Authorization enforcement
    "WdfDriverCreate",        # KMDF safety framework
    "WdfDeviceCreate",        # KMDF safety framework
    "ObReferenceObjectByHandleWithTag",  # Proper handle validation
}

# Skip drivers larger than this (huge drivers = slow Ghidra analysis)
MAX_SIZE_BYTES = 5 * 1024 * 1024  # 5MB default


# --- YAML Scoring Config ---
def _load_prefilter_scoring_yaml():
    """Load prefilter scoring from YAML config file.

    Search path: __file__ dir, cwd, CTHAEH_SCORING_PATH env var.
    Returns dict of prefilter scoring values, or empty dict if unavailable.
    """
    try:
        import yaml
    except ImportError:
        return {}

    candidates = []

    # 1. Same directory as this script
    try:
        candidates.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "scoring_rules.yaml"))
    except Exception:
        pass

    # 2. Current working directory
    candidates.append(os.path.join(os.getcwd(), "scoring_rules.yaml"))

    # 3. Environment variable override (highest priority, inserted first)
    env_path = os.environ.get("CTHAEH_SCORING_PATH")
    if env_path:
        candidates.insert(0, env_path)

    for yaml_path in candidates:
        try:
            with open(yaml_path, "r") as f:
                data = yaml.safe_load(f)
            if data and "prefilter" in data:
                return data["prefilter"]
        except Exception:
            continue

    return {}


_PREFILTER_YAML = _load_prefilter_scoring_yaml()


def _pf_score(key, default):
    """Get a prefilter scoring value from YAML config or hardcoded default."""
    return _PREFILTER_YAML.get(key, default)

# LOLDrivers cache file
LOLDRIVERS_CACHE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".loldrivers_cache.json")
LOLDRIVERS_URL = "https://www.loldrivers.io/api/drivers.json"


def load_loldrivers_hashes(force_refresh=False):
    """Load known vulnerable driver hashes from LOLDrivers."""
    cache_valid = False

    if not force_refresh and os.path.exists(LOLDRIVERS_CACHE):
        try:
            with open(LOLDRIVERS_CACHE, "r") as f:
                cache = json.load(f)
            # Cache valid for 7 days
            if time.time() - cache.get("fetched", 0) < 7 * 86400:
                return set(cache.get("hashes", [])), cache.get("names", {})
        except:
            pass

    try:
        import requests
        print("  Fetching LOLDrivers database...", end="", flush=True)
        resp = requests.get(LOLDRIVERS_URL, timeout=15)
        resp.raise_for_status()
        drivers = resp.json()

        hashes = set()
        names = {}  # hash -> driver name

        for driver in drivers:
            driver_name = driver.get("Tags", ["unknown"])[0] if driver.get("Tags") else "unknown"
            for sample in driver.get("KnownVulnerableSamples", []):
                for hash_type in ["SHA256", "SHA1", "MD5"]:
                    h = sample.get(hash_type, "")
                    if h:
                        h_lower = h.lower()
                        hashes.add(h_lower)
                        names[h_lower] = driver_name

        # Cache it
        with open(LOLDRIVERS_CACHE, "w") as f:
            json.dump({"fetched": time.time(), "hashes": list(hashes), "names": names}, f)

        print(f" {len(hashes)} hashes loaded")
        return hashes, names

    except ImportError:
        print("  LOLDrivers check skipped (install requests: pip install requests)")
        return set(), {}
    except Exception as e:
        print(f"  LOLDrivers fetch failed: {e}")
        return set(), {}


def get_file_hashes(filepath):
    """Calculate SHA256, SHA1, MD5 of a file."""
    sha256 = hashlib.sha256()
    sha1 = hashlib.sha1()
    md5 = hashlib.md5()

    with open(filepath, "rb") as f:
        while True:
            chunk = f.read(8192)
            if not chunk:
                break
            sha256.update(chunk)
            sha1.update(chunk)
            md5.update(chunk)

    return {
        "sha256": sha256.hexdigest(),
        "sha1": sha1.hexdigest(),
        "md5": md5.hexdigest(),
    }


def classify_driver_class(imports, import_dlls=None):
    """Classify driver by type based on imports and DLLs.

    Returns dict with 'class' (CRITICAL/HIGH/MEDIUM/LOW/UNKNOWN),
    'category' description, 'exploitability' notes, and 'framework'
    identifying the specific driver technology used.
    """
    if import_dlls is None:
        import_dlls = set()
    import_names_lower = {i.lower() for i in imports}

    has_iocreatedevice = "IoCreateDevice" in imports or "iocreatedevice" in import_names_lower
    has_wdfdrivercreate = "WdfDriverCreate" in imports or "wdfdrivercreate" in import_names_lower
    has_fltregisterfilter = "FltRegisterFilter" in imports or "fltregisterfilter" in import_names_lower

    # --- Framework detection (most specific first) ---

    # minifilter: FltRegisterFilter
    if has_fltregisterfilter:
        return {
            "class": "CRITICAL",
            "category": "File system filter",
            "exploitability": "FS filters intercept all file I/O; bugs = system-wide impact",
            "framework": "minifilter",
        }

    # wfp_callout: FwpsCalloutRegister0/1/2/3
    wfp_imports = {"FwpsCalloutRegister0", "FwpsCalloutRegister1", "FwpsCalloutRegister2", "FwpsCalloutRegister3"}
    if imports & wfp_imports or import_names_lower & {i.lower() for i in wfp_imports}:
        return {
            "class": "HIGH",
            "category": "WFP callout driver",
            "exploitability": "Network packet inspection in kernel; complex parsing surface",
            "framework": "wfp_callout",
        }

    # ndis_miniport: NdisMRegisterMiniportDriver
    if "NdisMRegisterMiniportDriver" in imports or "ndismregisterminiportdriver" in import_names_lower:
        return {
            "class": "HIGH",
            "category": "NDIS miniport driver",
            "exploitability": "Network packet parsing in kernel; remote attack surface",
            "framework": "ndis_miniport",
        }

    # ndis_filter: NdisFRegisterFilterDriver
    if "NdisFRegisterFilterDriver" in imports or "ndisfregisterfilterdriver" in import_names_lower:
        return {
            "class": "HIGH",
            "category": "NDIS filter driver",
            "exploitability": "Network packet interception in kernel; remote attack surface",
            "framework": "ndis_filter",
        }

    # ndis_protocol: NdisRegisterProtocolDriver
    if "NdisRegisterProtocolDriver" in imports or "ndisregisterprotocoldriver" in import_names_lower:
        return {
            "class": "HIGH",
            "category": "NDIS protocol driver",
            "exploitability": "Raw network protocol handling in kernel; remote attack surface",
            "framework": "ndis_protocol",
        }

    bt_dlls = {"bthport.sys", "bthhfp.sys"}
    if import_dlls & bt_dlls:
        return {
            "class": "HIGH",
            "category": "Bluetooth driver",
            "exploitability": "BT stack in kernel; proximity-based attack surface",
            "framework": "bluetooth",
        }

    usb_imports = {"USBD_CreateConfigurationRequestEx", "WdfUsbTargetDeviceSendControlTransferSynchronously"}
    if imports & usb_imports:
        return {
            "class": "HIGH",
            "category": "USB function driver",
            "exploitability": "USB request handling in kernel; physical/logical attack surface",
            "framework": "usb_function",
        }

    # storport: StorPortInitialize
    if "StorPortInitialize" in imports or "storportinitialize" in import_names_lower:
        return {
            "class": "MEDIUM",
            "category": "StorPort miniport driver",
            "exploitability": "Storage stack driver; data corruption risk on bugs",
            "framework": "storport",
        }

    # class_video: VideoPortInitialize
    if "VideoPortInitialize" in imports or "videoportinitialize" in import_names_lower:
        return {
            "class": "MEDIUM",
            "category": "Video miniport driver",
            "exploitability": "Legacy video port driver; limited modern attack surface",
            "framework": "class_video",
        }

    # ks_minidriver: KsCreateFilterFactory or KsInitializeDriver
    if "KsCreateFilterFactory" in imports or "kscreatefilterfactory" in import_names_lower or \
       "KsInitializeDriver" in imports or "ksinitializedriver" in import_names_lower:
        return {
            "class": "MEDIUM",
            "category": "KS minidriver",
            "exploitability": "Kernel streaming driver; media processing surface",
            "framework": "ks_minidriver",
        }

    # kmdf: WdfDriverCreate (check before wdm_raw since it has IoCreateDevice sometimes too)
    if has_wdfdrivercreate:
        return {
            "class": "MEDIUM",
            "category": "WDF/KMDF driver",
            "exploitability": "WDF provides safety rails but bugs still possible",
            "framework": "kmdf",
        }

    if "DxgkInitialize" in imports or "dxgkinitialize" in import_names_lower:
        return {
            "class": "MEDIUM",
            "category": "Display/GPU driver",
            "exploitability": "Complex IOCTL surface but often well-audited",
            "framework": "display_miniport",
        }

    # portclass_audio: PcRegisterSubdevice or PortClsCreate
    if "PortClsCreate" in imports or "portclscreate" in import_names_lower or \
       "PcRegisterSubdevice" in imports or "pcregistersubdevice" in import_names_lower:
        return {
            "class": "LOW",
            "category": "Audio (PortCls) driver",
            "exploitability": "Minimal direct user IOCTL surface",
            "framework": "portclass_audio",
        }

    hid_imports = {"HidRegisterMinidriver", "hidregisterminidriver"}
    if imports & hid_imports or import_names_lower & hid_imports:
        return {
            "class": "LOW",
            "category": "HID minidriver",
            "exploitability": "Limited attack surface through HID stack",
            "framework": "hid_minidriver",
        }

    printer_dlls = {"pjlmon.dll", "tcpmon.dll", "usbmon.dll"}
    if import_dlls & printer_dlls:
        return {
            "class": "LOW",
            "category": "Printer driver",
            "exploitability": "Typically sandboxed print pipeline",
            "framework": "printer",
        }

    # wdm_raw: IoCreateDevice without any framework
    if has_iocreatedevice:
        return {
            "class": "CRITICAL",
            "category": "Raw WDM driver",
            "exploitability": "No WDF safety rails; manual IRP handling prone to bugs",
            "framework": "wdm_raw",
        }

    return {
        "class": "UNKNOWN",
        "category": "Unclassified",
        "exploitability": "Manual review needed",
        "framework": None,
    }


def check_driver(driver_path, max_size=MAX_SIZE_BYTES, lol_hashes=None, lol_names=None,
                  wdac_hashes=None, wdac_filename_rules=None, holygrail_lol=None):
    """
    Quick PE import check on a driver.
    Returns: (should_analyze, reason, risk_hint, flags)
    """
    name = os.path.basename(driver_path)
    size = os.path.getsize(driver_path)
    flags = []

    # Size check
    if max_size and size > max_size:
        return False, f"too large ({size // 1024}KB)", 0, flags, None, None

    # Compute hashes once
    file_hashes = get_file_hashes(driver_path)

    # WDAC block policy check (skip drivers already blocked by Microsoft)
    if wdac_hashes:
        for h in (file_hashes["sha256"], file_hashes["sha1"]):
            if h in wdac_hashes:
                flags.append("WDAC_BLOCKED")
                return False, "blocked by WDAC driver block policy", 0, flags, None, None

    # WDAC filename+version rules
    if wdac_filename_rules:
        name_lower = name.lower()
        for rule_name, rule_max_ver in wdac_filename_rules:
            if name_lower == rule_name:
                flags.append(f"WDAC_FILENAME_BLOCKED:{rule_max_ver}")
                return False, f"blocked by WDAC filename rule ({name} <= {rule_max_ver})", 0, flags, None, None

    # LOLDrivers check (loldrivers.io)
    if lol_hashes:
        for h in file_hashes.values():
            if h in lol_hashes:
                lol_name = lol_names.get(h, "unknown") if lol_names else "unknown"
                flags.append(f"KNOWN_VULN:{lol_name}")
                break

    # HolyGrail LOLDrivers cross-reference by SHA256 (flag but don't skip)
    if holygrail_lol:
        sha256 = file_hashes["sha256"]
        if sha256 in holygrail_lol:
            lol_tag = holygrail_lol[sha256]
            flags.append(f"HOLYGRAIL_LOL:{lol_tag}")

    try:
        pe = pefile.PE(driver_path, fast_load=True)
        pe.parse_data_directories(
            directories=[
                pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IMPORT"],
                pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_RESOURCE"],
            ]
        )
    except Exception as e:
        return False, f"PE parse error: {e}", 0, flags, None, None

    # Extract import names and DLL names
    imports = set()
    import_dlls = set()
    if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            dll_name = entry.dll.decode("utf-8", errors="ignore").lower() if entry.dll else ""
            import_dlls.add(dll_name)
            for imp in entry.imports:
                if imp.name:
                    imports.add(imp.name.decode("utf-8", errors="ignore"))

    # Extract signer / company name from version info
    signer = None
    try:
        if hasattr(pe, "VS_VERSIONINFO"):
            for finfo in pe.FileInfo:
                for entry in finfo:
                    if hasattr(entry, "StringTable"):
                        for st in entry.StringTable:
                            for key, val in st.entries.items():
                                key_str = key.decode("utf-8", errors="ignore")
                                if key_str == "CompanyName" and val:
                                    signer = val.decode("utf-8", errors="ignore").strip()
    except Exception:
        pass

    # Classify driver class based on imports and DLLs
    driver_class = classify_driver_class(imports, import_dlls)

    pe.close()

    # Add signer and driver class to flags for downstream use
    if signer:
        flags.append(f"SIGNER:{signer}")
    if driver_class and driver_class["class"] != "UNKNOWN":
        flags.append(f"CLASS:{driver_class['class']}:{driver_class['category']}")
    if driver_class and driver_class.get("framework"):
        flags.append(f"FRAMEWORK:{driver_class['framework']}")

    # Must have at least one interesting import
    has_interesting = bool(imports & INTERESTING_IMPORTS)
    if not has_interesting:
        return False, "no device/IOCTL imports", 0, flags, signer, driver_class

    # Count high-risk imports as a hint (each worth 1 point from YAML or default)
    high_risk_count = len(imports & HIGH_RISK_IMPORTS) * _pf_score("high_risk_bonus", 1)

    # Communication capability detection (user-mode bridge)
    comms_found = {i for i in imports if i in COMMS_IMPORTS}
    if len(comms_found) >= 2:
        flags.append(f"COMMS:{'+'.join(sorted(comms_found))}")
        high_risk_count += _pf_score("comms_bonus", 2)

    # BYOVD process killer detection
    has_opener = bool(imports & BYOVD_OPENERS)
    has_terminator = bool(imports & BYOVD_TERMINATORS)
    if has_opener and has_terminator:
        flags.append("BYOVD_CANDIDATE")
        high_risk_count += _pf_score("byovd_candidate_bonus", 3)

    # PPL killer potential: ZwTerminateProcess + (ZwOpenProcess | PsLookupProcessByProcessId)
    has_zw_terminate = "ZwTerminateProcess" in imports
    has_zw_open = "ZwOpenProcess" in imports
    has_ps_lookup = "PsLookupProcessByProcessId" in imports
    if has_zw_terminate and (has_zw_open or has_ps_lookup):
        flags.append("PPL_KILLER")
        high_risk_count += _pf_score("ppl_killer_bonus", 4)

    # Physical memory R/W detection
    phys_mem_count = len(imports & PHYS_MEM_INDICATORS)
    if phys_mem_count >= 2:
        flags.append("PHYS_MEM_RW")
        high_risk_count += _pf_score("phys_mem_rw_bonus", 2)

    # MmMapIoSpace alone is notable
    if "MmMapIoSpace" in imports:
        if "PHYS_MEM_RW" not in flags:
            flags.append("MMIO_MAP")

    # Token stealing / EPROCESS manipulation
    token_imports = imports & TOKEN_STEAL_IMPORTS
    if len(token_imports) >= 2:
        flags.append("TOKEN_STEAL")
        high_risk_count += _pf_score("token_steal_bonus", 2)
    elif "PsLookupProcessByProcessId" in imports:
        flags.append("PROCESS_LOOKUP")
        high_risk_count += _pf_score("process_lookup_bonus", 1)

    # Registry manipulation from kernel
    reg_imports = imports & REGISTRY_IMPORTS
    if len(reg_imports) >= 2:
        flags.append("REGISTRY_RW")
        high_risk_count += _pf_score("registry_rw_bonus", 1)

    # Firmware/SPI access
    fw_imports = imports & FIRMWARE_IMPORTS
    if fw_imports:
        flags.append("FIRMWARE_ACCESS")
        high_risk_count += _pf_score("firmware_access_bonus", 2)

    # --- Boot phase scoring ---
    # Drivers that load early in boot operate before EDR can install hooks.
    # Ref: Jiří Vinopal EDR Phase 0 blind spots
    start_type, boot_phase = get_driver_start_type(driver_path)
    if start_type == 0:
        boot_bonus = _pf_score("boot_phase_0_bonus", 15)
    elif start_type == 1:
        boot_bonus = _pf_score("boot_phase_1_bonus", 10)
    else:
        boot_bonus = 0
    if boot_bonus:
        high_risk_count += boot_bonus
        flags.append(f"BOOT_PHASE:{boot_phase}")

    # --- Negative scoring: good security practices reduce risk ---
    if "ProbeForRead" in imports:
        high_risk_count += _pf_score("probe_for_read_bonus", -2)
        flags.append("HAS_PROBES")
    if "ProbeForWrite" in imports:
        high_risk_count += _pf_score("probe_for_write_bonus", -2)
        if "HAS_PROBES" not in flags:
            flags.append("HAS_PROBES")
    if "IoCreateDeviceSecure" in imports:
        high_risk_count += _pf_score("secure_device_creation_bonus", -3)
        flags.append("SECURE_DEVICE")
    if "SeAccessCheck" in imports:
        high_risk_count += _pf_score("se_access_check_bonus", -2)
        flags.append("HAS_ACCESS_CHECK")
    if "SeSinglePrivilegeCheck" in imports:
        high_risk_count += _pf_score("se_privilege_check_bonus", -2)
        if "HAS_ACCESS_CHECK" not in flags:
            flags.append("HAS_ACCESS_CHECK")
    if "WdfDriverCreate" in imports:
        high_risk_count += _pf_score("wdf_driver_create_bonus", -2)
        flags.append("KMDF_FRAMEWORK")
    if "WdfDeviceCreate" in imports:
        high_risk_count += _pf_score("wdf_device_create_bonus", -2)
        if "KMDF_FRAMEWORK" not in flags:
            flags.append("KMDF_FRAMEWORK")
    if "ObReferenceObjectByHandleWithTag" in imports:
        high_risk_count += _pf_score("ob_ref_by_handle_tag_bonus", -1)
        flags.append("VALIDATED_HANDLES")

    # Floor at 0
    if high_risk_count < 0:
        high_risk_count = 0

    return True, "has attack surface", high_risk_count, flags, signer, driver_class


def prefilter_directory(drivers_dir, max_size=MAX_SIZE_BYTES, check_loldrivers=False, byovd_only=False, boot_log_data=None):
    """
    Pre-filter all .sys files in a directory.
    Returns list of drivers worth sending to Ghidra.
    """
    results = {"analyze": [], "skip": [], "known_vuln": [], "byovd_candidates": [], "wdac_blocked": []}

    # Load LOLDrivers if requested
    lol_hashes = set()
    lol_names = {}
    if check_loldrivers:
        lol_hashes, lol_names = load_loldrivers_hashes()

    # Load WDAC block policy hashes and filename rules
    wdac_hashes = load_wdac_block_hashes()
    wdac_filename_rules = load_wdac_filename_rules()
    if wdac_hashes:
        print(f"  WDAC block policy: {len(wdac_hashes)} deny hashes, {len(wdac_filename_rules)} filename rules loaded")

    # Load HolyGrail LOLDrivers for SHA256 cross-reference
    holygrail_lol = load_holygrail_loldrivers()
    if holygrail_lol:
        print(f"  HolyGrail LOLDrivers: {len(holygrail_lol)} SHA256 hashes loaded")

    sys_files = []
    for root, dirs, files in os.walk(drivers_dir):
        for f in files:
            if f.lower().endswith(".sys"):
                sys_files.append(os.path.join(root, f))

    start = time.time()

    def _check_one(path):
        name = os.path.basename(path)
        should_analyze, reason, risk_hint, flags, signer, driver_class = check_driver(
            path, max_size, lol_hashes, lol_names,
            wdac_hashes=wdac_hashes, wdac_filename_rules=wdac_filename_rules,
            holygrail_lol=holygrail_lol,
        )
        # Boot phase info
        start_type, boot_phase = get_driver_start_type(path)

        entry = {
            "name": name,
            "path": path,
            "size": os.path.getsize(path),
            "risk_hint": risk_hint,
            "flags": flags,
            "boot_phase": boot_phase,
            "start_type": start_type,
            "_should_analyze": should_analyze,
            "_reason": reason,
        }
        if signer:
            entry["signer"] = signer
        if driver_class:
            entry["driver_class"] = driver_class
        return entry

    # Parallelize pefile checks with threads (I/O bound + GIL released during file reads)
    worker_count = min(8, max(1, os.cpu_count() or 2))
    with ThreadPoolExecutor(max_workers=worker_count) as pool:
        entries = list(pool.map(_check_one, sys_files))

    for entry in entries:
        should_analyze = entry.pop("_should_analyze")
        reason = entry.pop("_reason")

        # Track special categories
        is_known_vuln = any(f.startswith("KNOWN_VULN") for f in entry["flags"])
        is_wdac_blocked = any(f.startswith("WDAC_") for f in entry["flags"])
        if is_known_vuln:
            results["known_vuln"].append(entry)
        if is_wdac_blocked:
            results["wdac_blocked"].append(entry)
        if "BYOVD_CANDIDATE" in entry["flags"]:
            results["byovd_candidates"].append(entry)

        # Skip WDAC-blocked drivers — useless for research
        if is_wdac_blocked:
            entry["skip_reason"] = reason
            results["skip"].append(entry)
            continue

        # Skip known LOLDrivers — we're hunting 0-days, not rediscovering old bugs
        if is_known_vuln:
            entry["skip_reason"] = "known vulnerable (LOLDrivers) — skipping for novel research"
            results["skip"].append(entry)
            continue

        if should_analyze:
            # In BYOVD-only mode, only keep BYOVD candidates
            if byovd_only and "BYOVD_CANDIDATE" not in entry["flags"]:
                entry["skip_reason"] = "not a BYOVD candidate"
                results["skip"].append(entry)
            else:
                results["analyze"].append(entry)
        else:
            entry["skip_reason"] = reason
            results["skip"].append(entry)

    elapsed = time.time() - start

    # Boot log cross-reference (from boot_analyzer.py output)
    if boot_log_data:
        for entry in results["analyze"]:
            driver_lower = entry["name"].lower()
            if driver_lower in boot_log_data:
                boot_entry = boot_log_data[driver_lower]
                entry["risk_hint"] += _pf_score("boot_log_blind_spot_bonus", BOOT_LOG_BLIND_SPOT_BONUS)
                entry["boot_blind_spot"] = True
                entry["boot_log_hits"] = boot_entry.get("name_not_found_count", 0)
                if not any(f.startswith("BOOT_BLIND_SPOT") for f in entry["flags"]):
                    entry["flags"].append("BOOT_BLIND_SPOT")

    # Sort analyzable drivers by risk hint (highest first)
    results["analyze"].sort(key=lambda x: x["risk_hint"], reverse=True)

    total = len(sys_files)
    kept = len(results["analyze"])
    skipped = len(results["skip"])

    print(f"\n🌳 Cthaeh Pre-filter ({elapsed:.1f}s for {total} drivers)")
    print(f"  ✅ Analyze: {kept} drivers ({kept*100//max(total,1)}%)")
    print(f"  ⏭️  Skip:    {skipped} drivers")

    # Show skip breakdown
    skip_reasons = {}
    for s in results["skip"]:
        r = s.get("skip_reason", "unknown")
        skip_reasons[r] = skip_reasons.get(r, 0) + 1
    for reason, count in sorted(skip_reasons.items(), key=lambda x: -x[1]):
        print(f"      {reason}: {count}")

    # Show special findings
    if results["known_vuln"]:
        print(f"\n  ⚠️  Known vulnerable (LOLDrivers): {len(results['known_vuln'])}")
        for d in results["known_vuln"]:
            vuln_tag = [f for f in d["flags"] if f.startswith("KNOWN_VULN")][0]
            print(f"      {d['name']} ({vuln_tag})")

    if results["wdac_blocked"]:
        print(f"\n  🚫 WDAC blocked (skipped): {len(results['wdac_blocked'])}")

    # HolyGrail LOLDrivers cross-reference
    holygrail_flagged = [d for d in results["analyze"] if any(f.startswith("HOLYGRAIL_LOL") for f in d.get("flags", []))]
    if holygrail_flagged:
        print(f"\n  🔍 HolyGrail LOLDrivers match (kept for variant research): {len(holygrail_flagged)}")
        for d in holygrail_flagged:
            tag = [f for f in d["flags"] if f.startswith("HOLYGRAIL_LOL")][0]
            print(f"      {d['name']} ({tag})")

    if results["byovd_candidates"]:
        print(f"\n  🎯 BYOVD candidates (process killer imports): {len(results['byovd_candidates'])}")
        for d in results["byovd_candidates"]:
            extra = [f for f in d["flags"] if f != "BYOVD_CANDIDATE"]
            extra_str = f" [{', '.join(extra)}]" if extra else ""
            print(f"      {d['name']}{extra_str}")

    # PPL killer candidates
    ppl_killers = [d for d in results["analyze"] if "PPL_KILLER" in d.get("flags", [])]
    if ppl_killers:
        print(f"\n  ☠️  PPL killer potential: {len(ppl_killers)}")
        for d in ppl_killers:
            print(f"      {d['name']}")

    # Communication capability
    comms_drivers = [d for d in results["analyze"] if any(f.startswith("COMMS:") for f in d.get("flags", []))]
    if comms_drivers:
        print(f"\n  📡 User-mode comms bridge: {len(comms_drivers)}")

    phys_mem = [d for d in results["analyze"] if "PHYS_MEM_RW" in d.get("flags", [])]
    if phys_mem:
        print(f"\n  🔓 Physical memory R/W candidates: {len(phys_mem)}")
        for d in phys_mem:
            print(f"      {d['name']}")

    # Boot blind spot drivers
    blind_spot_drivers = [d for d in results["analyze"] if d.get("boot_blind_spot")]
    if blind_spot_drivers:
        print(f"\n  🔇 Boot blind spot drivers: {len(blind_spot_drivers)}")
        for d in blind_spot_drivers:
            print(f"      {d['name']} (boot hits: {d.get('boot_log_hits', 0)}, phase: {d.get('boot_phase', 'UNKNOWN')})")

    # Boot phase breakdown
    boot_phases = {}
    for d in results["analyze"]:
        bp = d.get("boot_phase", "UNKNOWN")
        if bp in ("BOOT_START", "SYSTEM_START"):
            boot_phases[bp] = boot_phases.get(bp, 0) + 1
    if boot_phases:
        print(f"\n  ⏱️  Early-boot drivers (EDR blind window):")
        for phase in ("BOOT_START", "SYSTEM_START"):
            if phase in boot_phases:
                print(f"      {phase}: {boot_phases[phase]}")

    # Show driver class breakdown
    class_counts = {}
    for d in results["analyze"]:
        dc = d.get("driver_class", {})
        cls = dc.get("class", "UNKNOWN") if dc else "UNKNOWN"
        class_counts[cls] = class_counts.get(cls, 0) + 1
    if class_counts:
        print(f"\n  📊 Driver class breakdown:")
        for cls in ["CRITICAL", "HIGH", "MEDIUM", "LOW", "UNKNOWN"]:
            if cls in class_counts:
                print(f"      {cls}: {class_counts[cls]}")

    return results


def main():
    import argparse
    parser = argparse.ArgumentParser(
        description="🌳 Cthaeh Pre-filter - Fast PE import check"
    )
    parser.add_argument("drivers_dir", help="Directory containing .sys files")
    parser.add_argument("--max-size", type=int, default=5,
                        help="Max driver size in MB (default: 5)")
    parser.add_argument("--output", help="Write filtered list to JSON file")
    parser.add_argument("--list", action="store_true",
                        help="Print list of drivers to analyze")
    parser.add_argument("--loldrivers", action="store_true",
                        help="Cross-reference against LOLDrivers database")
    parser.add_argument("--byovd", action="store_true",
                        help="Only show BYOVD process killer candidates")
    parser.add_argument("--refresh-lol", action="store_true",
                        help="Force refresh LOLDrivers cache")
    parser.add_argument("--boot-log",
                        help="Boot analyzer JSON report (from boot_analyzer.py) for cross-reference")

    args = parser.parse_args()

    if args.refresh_lol:
        load_loldrivers_hashes(force_refresh=True)

    # Load boot log for cross-reference if provided
    boot_log_data = {}
    if args.boot_log:
        boot_log_data = load_boot_log(args.boot_log)
        if boot_log_data:
            print(f"  Boot log loaded: {len(boot_log_data)} drivers")

    max_bytes = args.max_size * 1024 * 1024
    results = prefilter_directory(
        args.drivers_dir, max_bytes,
        check_loldrivers=args.loldrivers,
        byovd_only=args.byovd,
        boot_log_data=boot_log_data,
    )

    if args.list:
        print("\nDrivers to analyze:")
        for d in results["analyze"]:
            flags_str = f" [{', '.join(d['flags'])}]" if d["flags"] else ""
            hint = f" (risk:{d['risk_hint']})" if d["risk_hint"] else ""
            print(f"  {d['name']}{hint}{flags_str}")

    if args.output:
        paths = [d["path"] for d in results["analyze"]]
        with open(args.output, "w") as f:
            json.dump(paths, f, indent=2)
        print(f"\nFiltered list written to: {args.output}")


if __name__ == "__main__":
    main()