leviathan-driver 0.1.0

//! Pool Tag Scanning for Kernel Forensics
//!
//! Scans kernel memory pools to find objects by their pool tags.
//! This technique can find hidden/unlinked kernel objects.
//!
//! # How It Works
//! Most kernel allocations start with a POOL_HEADER structure that
//! contains a 4-byte pool tag. By scanning memory for known tags,
//! we can find objects that have been unlinked from their lists.
//!
//! # Common Pool Tags
//! - `Proc` - EPROCESS structures
//! - `Thre` - ETHREAD structures
//! - `MmLd` - Loaded kernel modules
//! - `File` - FILE_OBJECT structures
//! - `Driv` - DRIVER_OBJECT structures
//! - `ObHd` - Object handles
//!
//! # Use Cases
//! - Detect DKOM attacks (hidden processes)
//! - Find unlinked kernel objects
//! - Memory forensics analysis
//! - Malware detection

use alloc::vec::Vec;
use core::ptr;
use wdk::println;
use wdk_sys::{
    ntddk::MmIsAddressValid,
    NTSTATUS, PVOID, STATUS_SUCCESS,
};

/// Pool header structure (simplified)
#[repr(C)]
#[derive(Clone, Copy)]
pub struct PoolHeader {
    /// Previous size in pool blocks
    pub previous_size: u8,
    /// Pool index
    pub pool_index: u8,
    /// Block size in pool blocks
    pub block_size: u8,
    /// Pool type
    pub pool_type: u8,
    /// Pool tag (4 bytes)
    pub pool_tag: [u8; 4],
}

/// Size of pool header
pub const POOL_HEADER_SIZE: usize = 16; // On x64

/// Known pool tags for important structures
pub mod pool_tags {
    /// EPROCESS structure
    pub const PROCESS: [u8; 4] = *b"Proc";
    /// ETHREAD structure
    pub const THREAD: [u8; 4] = *b"Thre";
    /// Loaded module (LDR_DATA_TABLE_ENTRY)
    pub const MODULE: [u8; 4] = *b"MmLd";
    /// FILE_OBJECT
    pub const FILE: [u8; 4] = *b"File";
    /// DRIVER_OBJECT
    pub const DRIVER: [u8; 4] = *b"Driv";
    /// Object handle entry
    pub const HANDLE: [u8; 4] = *b"ObHd";
    /// Registry key
    pub const REGKEY: [u8; 4] = *b"CM10";
    /// Network socket
    pub const SOCKET: [u8; 4] = *b"TcpE";
    /// Raw socket
    pub const RAW_SOCKET: [u8; 4] = *b"RawE";
    /// Token
    pub const TOKEN: [u8; 4] = *b"Toke";
}

/// Result of a pool scan
#[derive(Debug)]
pub struct PoolScanResult {
    /// Address where object was found
    pub address: usize,
    /// Pool tag of the object
    pub pool_tag: [u8; 4],
    /// Size of the pool block
    pub block_size: usize,
    /// Whether this object appears in linked lists
    pub is_linked: bool,
}

/// Pool scanner configuration
pub struct PoolScanner {
    /// Start address for scanning
    pub start_address: usize,
    /// End address for scanning
    pub end_address: usize,
    /// Target pool tag to find
    pub target_tag: [u8; 4],
    /// Maximum results to return
    pub max_results: usize,
}

impl PoolScanner {
    /// Create a new pool scanner
    pub fn new(target_tag: [u8; 4]) -> Self {
        Self {
            // Default to scanning typical kernel pool regions
            start_address: 0xFFFF_8000_0000_0000,
            end_address: 0xFFFF_FFFF_FFFF_FFFF,
            target_tag,
            max_results: 1000,
        }
    }

    /// Set custom address range
    pub fn with_range(mut self, start: usize, end: usize) -> Self {
        self.start_address = start;
        self.end_address = end;
        self
    }

    /// Set maximum results
    pub fn with_max_results(mut self, max: usize) -> Self {
        self.max_results = max;
        self
    }

    /// Perform the pool scan
    ///
    /// # Safety
    /// - Must be at IRQL <= DISPATCH_LEVEL
    /// - Address range must be kernel addresses
    pub unsafe fn scan(&self) -> Vec<PoolScanResult> {
        let mut results = Vec::new();

        println!(
            "[Leviathan] Pool scan starting for tag '{}'",
            core::str::from_utf8(&self.target_tag).unwrap_or("????")
        );

        // In production, this would:
        // 1. Enumerate pool pages using MmGetSystemRoutineAddress
        // 2. Scan each page for pool headers
        // 3. Validate pool header consistency
        // 4. Extract object pointers

        // Simplified scanning approach (demonstration)
        // Real implementation needs proper pool page enumeration

        println!(
            "[Leviathan] Pool scan complete: {} results",
            results.len()
        );

        results
    }
}

/// Scan for EPROCESS structures
///
/// Finds all EPROCESS structures including those unlinked from
/// PsActiveProcessHead (hidden processes).
pub fn scan_for_processes() -> Vec<usize> {
    let scanner = PoolScanner::new(pool_tags::PROCESS);

    let results = unsafe { scanner.scan() };

    results.iter().map(|r| r.address).collect()
}

/// Scan for ETHREAD structures
///
/// Finds all kernel thread structures.
pub fn scan_for_threads() -> Vec<usize> {
    let scanner = PoolScanner::new(pool_tags::THREAD);

    let results = unsafe { scanner.scan() };

    results.iter().map(|r| r.address).collect()
}

/// Scan for loaded modules
///
/// Finds all LDR_DATA_TABLE_ENTRY structures including
/// unlinked/hidden kernel modules.
pub fn scan_for_modules() -> Vec<usize> {
    let scanner = PoolScanner::new(pool_tags::MODULE);

    let results = unsafe { scanner.scan() };

    results.iter().map(|r| r.address).collect()
}

/// Scan for DRIVER_OBJECT structures
pub fn scan_for_drivers() -> Vec<usize> {
    let scanner = PoolScanner::new(pool_tags::DRIVER);

    let results = unsafe { scanner.scan() };

    results.iter().map(|r| r.address).collect()
}

/// Scan for network connections (TCP endpoints)
pub fn scan_for_network_connections() -> Vec<usize> {
    let scanner = PoolScanner::new(pool_tags::SOCKET);

    let results = unsafe { scanner.scan() };

    results.iter().map(|r| r.address).collect()
}

/// Pool tag database for identification
pub struct PoolTagDatabase {
    /// Known pool tags and their descriptions
    entries: Vec<(u32, &'static str, &'static str)>,
}

impl PoolTagDatabase {
    /// Create a new pool tag database with known tags
    pub fn new() -> Self {
        let mut entries = Vec::new();

        // Add known pool tags
        // Format: (tag as u32, owner, description)
        entries.push((u32::from_le_bytes(*b"Proc"), "nt", "EPROCESS structure"));
        entries.push((u32::from_le_bytes(*b"Thre"), "nt", "ETHREAD structure"));
        entries.push((u32::from_le_bytes(*b"MmLd"), "nt", "Loaded module entry"));
        entries.push((u32::from_le_bytes(*b"File"), "nt", "FILE_OBJECT"));
        entries.push((u32::from_le_bytes(*b"Driv"), "nt", "DRIVER_OBJECT"));
        entries.push((u32::from_le_bytes(*b"ObHd"), "nt", "Object handle table"));
        entries.push((u32::from_le_bytes(*b"CM10"), "nt", "Registry key object"));
        entries.push((u32::from_le_bytes(*b"TcpE"), "tcpip", "TCP endpoint"));
        entries.push((u32::from_le_bytes(*b"UdpA"), "tcpip", "UDP endpoint"));
        entries.push((u32::from_le_bytes(*b"RawE"), "tcpip", "Raw socket"));
        entries.push((u32::from_le_bytes(*b"Toke"), "nt", "Token object"));
        entries.push((u32::from_le_bytes(*b"Sect"), "nt", "Section object"));
        entries.push((u32::from_le_bytes(*b"Sema"), "nt", "Semaphore object"));
        entries.push((u32::from_le_bytes(*b"Even"), "nt", "Event object"));
        entries.push((u32::from_le_bytes(*b"Muta"), "nt", "Mutant object"));

        Self { entries }
    }

    /// Look up a pool tag
    pub fn lookup(&self, tag: [u8; 4]) -> Option<(&str, &str)> {
        let tag_u32 = u32::from_le_bytes(tag);

        for (entry_tag, owner, desc) in &self.entries {
            if *entry_tag == tag_u32 {
                return Some((owner, desc));
            }
        }

        None
    }

    /// Get all known tags
    pub fn all_tags(&self) -> &[(u32, &'static str, &'static str)] {
        &self.entries
    }
}

/// Quick scan optimization using pool tracker tables
///
/// The kernel maintains _POOL_TRACKER_BIG_PAGES for large allocations.
/// Scanning this table is faster than full memory scanning.
pub mod quick_scan {
    use super::*;

    /// Pool tracker entry for big allocations
    #[repr(C)]
    pub struct PoolTrackerBigPages {
        /// Virtual address of allocation
        pub va: usize,
        /// Pool tag
        pub key: u32,
        /// Pattern (should be 0x1 for in-use)
        pub pattern: u32,
        /// Number of bytes
        pub number_of_bytes: usize,
    }

    /// Scan big pool allocations table
    ///
    /// This is much faster than full memory scanning.
    #[allow(dead_code)]
    pub fn scan_big_pool_table(_target_tag: [u8; 4]) -> Vec<usize> {
        let results = Vec::new();

        // In production:
        // 1. Find PoolBigPageTable and PoolBigPageTableSize
        // 2. Iterate through entries
        // 3. Filter by target tag

        results
    }
}

/// Validation helpers
pub mod validation {
    use super::*;

    /// Validate a pool header
    pub fn is_valid_pool_header(header: &PoolHeader) -> bool {
        // Basic sanity checks
        if header.block_size == 0 {
            return false;
        }

        // Pool type should have valid flags
        if header.pool_type > 0x7F {
            return false;
        }

        // Tag should be printable ASCII or null
        for byte in &header.pool_tag {
            if *byte != 0 && (*byte < 0x20 || *byte > 0x7E) {
                return false;
            }
        }

        true
    }

    /// Validate EPROCESS structure
    pub fn is_valid_eprocess(_address: usize) -> bool {
        // Check EPROCESS signature fields:
        // - Pcb.Header.Type should be ProcessObject (3)
        // - UniqueProcessId should be reasonable
        // - ActiveProcessLinks should be valid pointers

        true // Placeholder
    }

    /// Validate ETHREAD structure
    pub fn is_valid_ethread(_address: usize) -> bool {
        // Check ETHREAD signature fields:
        // - Tcb.Header.Type should be ThreadObject (6)
        // - Process pointer should be valid

        true // Placeholder
    }
}