leviathan-driver 0.2.0

//! Threat Detection Engine
//!
//! Real-time behavioral analysis and rule-based threat detection
//! for identifying malicious activity patterns.
//!
//! # Detection Methods
//! - **Rule Engine**: Pattern matching on events
//! - **Behavioral Analysis**: Activity correlation
//! - **Anomaly Detection**: Baseline deviation
//! - **Heuristics**: Known attack patterns
//!
//! # Detection Categories
//! - Process injection (CreateRemoteThread, APC, etc.)
//! - Credential theft (LSASS access, SAM dumping)
//! - Persistence (registry, services, scheduled tasks)
//! - Defense evasion (unhooking, AMSI bypass)
//! - Ransomware indicators (mass encryption)

pub mod rules;
pub mod behavior;
pub mod heuristics;

use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use core::sync::atomic::{AtomicU64, Ordering};
use wdk::println;

/// Detection alert severity levels
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(u8)]
pub enum Severity {
    /// Informational - normal activity
    Info = 0,
    /// Low - minor suspicious indicator
    Low = 1,
    /// Medium - moderate threat indicator
    Medium = 2,
    /// High - significant threat
    High = 3,
    /// Critical - active attack/malware
    Critical = 4,
}

/// MITRE ATT&CK tactics
#[derive(Debug, Clone, Copy, PartialEq)]
#[repr(u16)]
pub enum MitreTactic {
    /// Initial Access (TA0001)
    InitialAccess = 1,
    /// Execution (TA0002)
    Execution = 2,
    /// Persistence (TA0003)
    Persistence = 3,
    /// Privilege Escalation (TA0004)
    PrivilegeEscalation = 4,
    /// Defense Evasion (TA0005)
    DefenseEvasion = 5,
    /// Credential Access (TA0006)
    CredentialAccess = 6,
    /// Discovery (TA0007)
    Discovery = 7,
    /// Lateral Movement (TA0008)
    LateralMovement = 8,
    /// Collection (TA0009)
    Collection = 9,
    /// Command and Control (TA0011)
    CommandAndControl = 11,
    /// Exfiltration (TA0010)
    Exfiltration = 10,
    /// Impact (TA0040)
    Impact = 40,
}

/// Detection alert
#[derive(Debug, Clone)]
pub struct Alert {
    /// Unique alert ID
    pub id: u64,
    /// Detection rule ID
    pub rule_id: u32,
    /// Alert severity
    pub severity: Severity,
    /// MITRE ATT&CK tactic
    pub tactic: MitreTactic,
    /// MITRE ATT&CK technique ID (e.g., T1055)
    pub technique_id: [u8; 8],
    /// Alert title
    pub title: [u8; 128],
    /// Description
    pub description: [u8; 512],
    /// Source process ID
    pub source_pid: u32,
    /// Target process ID (if applicable)
    pub target_pid: Option<u32>,
    /// Timestamp
    pub timestamp: u64,
    /// Associated event IDs
    pub event_ids: Vec<u64>,
    /// Recommended action
    pub action: RecommendedAction,
}

/// Recommended response action
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum RecommendedAction {
    /// Log only, no action needed
    Log,
    /// Alert security team
    Alert,
    /// Block the operation
    Block,
    /// Terminate the process
    Terminate,
    /// Quarantine the file
    Quarantine,
    /// Isolate the endpoint
    Isolate,
}

/// Detection context for events
#[derive(Debug, Clone)]
pub struct DetectionContext {
    /// Process ID
    pub pid: u32,
    /// Parent process ID
    pub ppid: u32,
    /// Process name
    pub process_name: [u8; 64],
    /// Image path
    pub image_path: [u16; 260],
    /// Command line
    pub command_line: [u16; 512],
    /// User SID
    pub user_sid: [u8; 68],
    /// Is elevated/admin
    pub is_elevated: bool,
    /// Is system process
    pub is_system: bool,
    /// Session ID
    pub session_id: u32,
}

/// Detection engine
pub struct DetectionEngine {
    /// Loaded detection rules
    rules: Vec<rules::DetectionRule>,
    /// Behavioral analyzers
    analyzers: Vec<Box<dyn behavior::BehaviorAnalyzer>>,
    /// Alert counter
    alert_counter: AtomicU64,
    /// Alert callback
    alert_callback: Option<fn(&Alert)>,
    /// Statistics
    stats: DetectionStats,
    /// Process activity cache
    process_cache: BTreeMap<u32, ProcessActivity>,
}

/// Detection statistics
#[derive(Debug, Default)]
pub struct DetectionStats {
    /// Total events processed
    pub events_processed: AtomicU64,
    /// Alerts generated
    pub alerts_generated: AtomicU64,
    /// Blocked operations
    pub operations_blocked: AtomicU64,
    /// False positives (if tracked)
    pub false_positives: AtomicU64,
}

/// Cached process activity for correlation
#[derive(Debug, Clone)]
pub struct ProcessActivity {
    /// Process ID
    pub pid: u32,
    /// Process creation time
    pub create_time: u64,
    /// Child processes spawned
    pub children: Vec<u32>,
    /// Threads created
    pub threads_created: u32,
    /// Remote threads created
    pub remote_threads: u32,
    /// Images loaded
    pub images_loaded: u32,
    /// Files accessed
    pub files_accessed: u32,
    /// Registry operations
    pub registry_ops: u32,
    /// Network connections
    pub network_conns: u32,
    /// Suspicious indicators count
    pub suspicious_count: u32,
    /// Last activity timestamp
    pub last_activity: u64,
}

impl DetectionEngine {
    /// Create a new detection engine
    pub fn new() -> Self {
        Self {
            rules: Vec::new(),
            analyzers: Vec::new(),
            alert_counter: AtomicU64::new(1),
            alert_callback: None,
            stats: DetectionStats::default(),
            process_cache: BTreeMap::new(),
        }
    }

    /// Load default detection rules
    pub fn load_default_rules(&mut self) {
        // Process injection rules
        self.add_rule(rules::DetectionRule {
            id: 1,
            name: *b"RemoteThreadInjection\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
            severity: Severity::High,
            tactic: MitreTactic::DefenseEvasion,
            technique_id: *b"T1055   ",
            enabled: true,
            rule_type: rules::RuleType::ThreadCreation,
            conditions: Vec::new(),
        });

        // Credential access rules
        self.add_rule(rules::DetectionRule {
            id: 2,
            name: *b"LsassAccess\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
            severity: Severity::Critical,
            tactic: MitreTactic::CredentialAccess,
            technique_id: *b"T1003   ",
            enabled: true,
            rule_type: rules::RuleType::ProcessAccess,
            conditions: Vec::new(),
        });

        // Persistence rules
        self.add_rule(rules::DetectionRule {
            id: 3,
            name: *b"RegistryRunKey\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
            severity: Severity::Medium,
            tactic: MitreTactic::Persistence,
            technique_id: *b"T1547   ",
            enabled: true,
            rule_type: rules::RuleType::RegistryMod,
            conditions: Vec::new(),
        });

        // Defense evasion rules
        self.add_rule(rules::DetectionRule {
            id: 4,
            name: *b"AmsiBypass\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
            severity: Severity::High,
            tactic: MitreTactic::DefenseEvasion,
            technique_id: *b"T1562   ",
            enabled: true,
            rule_type: rules::RuleType::MemoryMod,
            conditions: Vec::new(),
        });

        // Ransomware rules
        self.add_rule(rules::DetectionRule {
            id: 5,
            name: *b"RansomwareIndicator\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
            severity: Severity::Critical,
            tactic: MitreTactic::Impact,
            technique_id: *b"T1486   ",
            enabled: true,
            rule_type: rules::RuleType::FileOp,
            conditions: Vec::new(),
        });

        println!("[Leviathan] Loaded {} detection rules", self.rules.len());
    }

    /// Add a detection rule
    pub fn add_rule(&mut self, rule: rules::DetectionRule) {
        self.rules.push(rule);
    }

    /// Set alert callback
    pub fn set_alert_callback(&mut self, callback: fn(&Alert)) {
        self.alert_callback = Some(callback);
    }

    /// Process an event through detection engine
    pub fn process_event(
        &mut self,
        event_type: EventType,
        context: &DetectionContext,
        event_data: &[u8],
    ) -> Option<Alert> {
        self.stats.events_processed.fetch_add(1, Ordering::Relaxed);

        // Update process activity cache
        self.update_activity(context.pid, event_type);

        // Run through rules
        for rule in &self.rules {
            if !rule.enabled {
                continue;
            }

            if let Some(alert) = self.evaluate_rule(rule, event_type, context, event_data) {
                self.stats.alerts_generated.fetch_add(1, Ordering::Relaxed);

                // Call alert callback if set
                if let Some(callback) = self.alert_callback {
                    callback(&alert);
                }

                return Some(alert);
            }
        }

        // Run behavioral analysis
        for analyzer in &self.analyzers {
            if let Some(alert) = analyzer.analyze(context, &self.process_cache) {
                self.stats.alerts_generated.fetch_add(1, Ordering::Relaxed);
                return Some(alert);
            }
        }

        None
    }

    /// Evaluate a single rule
    fn evaluate_rule(
        &self,
        rule: &rules::DetectionRule,
        event_type: EventType,
        context: &DetectionContext,
        _event_data: &[u8],
    ) -> Option<Alert> {
        // Check if rule applies to this event type
        if !rule.matches_event_type(event_type) {
            return None;
        }

        // Evaluate conditions
        if !rule.evaluate_conditions(context) {
            return None;
        }

        // Generate alert
        let alert_id = self.alert_counter.fetch_add(1, Ordering::SeqCst);

        Some(Alert {
            id: alert_id,
            rule_id: rule.id,
            severity: rule.severity,
            tactic: rule.tactic,
            technique_id: rule.technique_id,
            title: {
                let mut title = [0u8; 128];
                let len = core::cmp::min(rule.name.len(), 128);
                title[..len].copy_from_slice(&rule.name[..len]);
                title
            },
            description: [0u8; 512],
            source_pid: context.pid,
            target_pid: None,
            timestamp: get_timestamp(),
            event_ids: Vec::new(),
            action: severity_to_action(rule.severity),
        })
    }

    /// Update process activity cache
    fn update_activity(&mut self, pid: u32, event_type: EventType) {
        let activity = self.process_cache.entry(pid).or_insert_with(|| ProcessActivity {
            pid,
            create_time: get_timestamp(),
            children: Vec::new(),
            threads_created: 0,
            remote_threads: 0,
            images_loaded: 0,
            files_accessed: 0,
            registry_ops: 0,
            network_conns: 0,
            suspicious_count: 0,
            last_activity: 0,
        });

        activity.last_activity = get_timestamp();

        match event_type {
            EventType::ProcessCreate => activity.children.push(pid),
            EventType::ThreadCreate => activity.threads_created += 1,
            EventType::ImageLoad => activity.images_loaded += 1,
            EventType::FileOp => activity.files_accessed += 1,
            EventType::RegistryOp => activity.registry_ops += 1,
            EventType::NetworkOp => activity.network_conns += 1,
            _ => {}
        }
    }

    /// Clean up old entries from process cache
    pub fn cleanup_cache(&mut self, max_age_ms: u64) {
        let now = get_timestamp();
        let threshold = now.saturating_sub(max_age_ms * 10000); // Convert to 100ns

        self.process_cache.retain(|_, activity| {
            activity.last_activity > threshold
        });
    }

    /// Get detection statistics
    pub fn get_stats(&self) -> &DetectionStats {
        &self.stats
    }
}

/// Event types for detection
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum EventType {
    ProcessCreate,
    ProcessExit,
    ThreadCreate,
    ThreadExit,
    ImageLoad,
    FileOp,
    RegistryOp,
    NetworkOp,
    MemoryOp,
    HandleOp,
}

/// Map severity to recommended action
fn severity_to_action(severity: Severity) -> RecommendedAction {
    match severity {
        Severity::Info => RecommendedAction::Log,
        Severity::Low => RecommendedAction::Log,
        Severity::Medium => RecommendedAction::Alert,
        Severity::High => RecommendedAction::Block,
        Severity::Critical => RecommendedAction::Terminate,
    }
}

/// Get current timestamp
fn get_timestamp() -> u64 {
    // Would use KeQuerySystemTimePrecise
    0
}