impl get_syscall_table, and fixed search_proc

crate/pe-tools/src/shared.rs

@@ -180,58 +180,3 @@ pub fn from_wide_ptr(ptr: *const u16) -> String {
         OsString::from_wide(slice).to_string_lossy().into_owned()
     }
 }
-
-pub fn search_syscall<T>(syscall_name: impl Into<String>, is_64bit: bool) -> Result<T> {
-    let syscall_name = syscall_name.into();
-
-    let ntdll = if is_64bit {
-        "c:\\windows\\system32\\ntdll.dll"
-    } else {
-        "c:\\windows\\SysWOW64\\ntdll.dll"
-    };
-
-    let ntdll_buffer = get_binary_from_file(ntdll).with_context(|| format!("could not find ntdll. path: {}", ntdll))?;
-
-    let pe = goblin::pe::PE::parse(&ntdll_buffer).with_context(|| format!("could not parse ntdll. path: {}", ntdll))?;
-
-    for e in pe.exports {
-        match e.name {
-            Some(symbol) => {
-                if symbol == syscall_name {
-                    if pe.is_64 {
-                        if ntdll_buffer[e.offset .. e.offset + 3] == [0x4C, 0x8B, 0xD1] {
-                            let syscall_number = ntdll_buffer[e.offset + 4];
-                            //print!("number: 0x{:x} ", syscall_number);
-                            /*
-                             *    mov    r10, rcx
-                             *    mov    eax, syscall_number
-                             *    syscall
-                             *    ret
-                             */
-                            let b = [
-                                0x4C, 0x8B, 0xD1,
-                                0xB8, syscall_number,  0x0, 0x0, 0x0,
-                                0xF, 0x5,
-                                0xC3
-                            ];
-
-                            return Ok(ptr_to_fn::<T>(&b[0] as *const _ as _))
-                        } else {
-                            bail!("found function, but could not get syscall number. function name: {}", syscall_name);
-                        }
-                    } else {
-                        //println!("fs: 0x{:x}", unsafe { __readfsdword(0xC0) });
-                        //let b = [0];
-                    }
-                    //println!("0x{:x}", e.offset);
-                    //for i in 0..12 {
-                    //    print!("0x{:x} ", ntdll_buffer[e.offset + i]);
-                    //}
-                }
-            },
-            None => { continue }
-        }
-    }
-
-    bail!("not found function. function name: {}", syscall_name);
-}

crate/pe-tools/src/x64.rs

@@ -1,6 +1,7 @@
 #[allow(non_snake_case)]
 use crate::shared::*;
 
+use std::collections::HashMap;
 use std::mem::size_of;
 use std::ptr::null_mut;
 
@@ -24,7 +25,6 @@ use winapi::{
     um::{
         memoryapi::{ReadProcessMemory, WriteProcessMemory},
         winnt::{
-            IMAGE_DIRECTORY_ENTRY_BASERELOC, IMAGE_DIRECTORY_ENTRY_IMPORT,
             IMAGE_ORDINAL64, IMAGE_IMPORT_BY_NAME, IMAGE_IMPORT_DESCRIPTOR,
             IMAGE_DOS_HEADER, IMAGE_NT_HEADERS64,
             IMAGE_SNAP_BY_ORDINAL64, IMAGE_THUNK_DATA64,
@@ -328,7 +328,7 @@ impl PE_Container<'_> {
     }
 }
 
-pub fn search_proc_address(function_name: impl Into<String>) -> anyhow::Result<*mut c_void> {
+pub fn search_proc_from_loaded_module<T>(function_name: impl Into<String>) -> anyhow::Result<T> {
     unsafe {
         let function_name = function_name.into();
         let ppeb = __readgsqword(0x60) as *mut PEB;
@@ -338,16 +338,26 @@ pub fn search_proc_address(function_name: impl Into<String>) -> anyhow::Result<*
 
         while !(*module_list).DllBase.is_null() {
             let dll_base = (*module_list).DllBase;
+            let size = (*module_list).SizeOfImage as usize;
 
             module_list = (*module_list).InLoadOrderLinks.Flink as *mut LDR_DATA_TABLE_ENTRY;
 
-            let dll_container = PE_Container::new(0x0 as _, dll_base)?;
+            let buffer = std::slice::from_raw_parts::<u8>(dll_base as _, size);
 
-            for e in &dll_container.pe.exports {
+            let opts = goblin::pe::options::ParseOptions{ resolve_rva: false };
+            let res = goblin::pe::PE::parse_with_opts(buffer, &opts);
+
+            // temp solution
+            if res.is_err() { continue }
+            let parsed = res?;
+
+            if !parsed.is_lib { continue }
+
+            for e in parsed.exports {
                 match e.name {
                     Some(symbol) => {
                         if symbol == function_name {
-                            return Ok(dll_container.to_va(e.offset as _));
+                            return Ok(ptr_to_fn::<T>((dll_base as usize + e.offset) as _));
                         }
                     },
                     None => {}
@@ -358,3 +368,64 @@ pub fn search_proc_address(function_name: impl Into<String>) -> anyhow::Result<*
         bail!("could not find {}", function_name);
     }
 }
+
+pub fn get_syscall_table() -> anyhow::Result<HashMap<String, usize>> {
+    unsafe {
+        let mut function_table = HashMap::new();
+        let ppeb = __readgsqword(0x60) as *mut PEB;
+
+        let p_peb_ldr_data = (*ppeb).Ldr;
+        let mut module_list = (*p_peb_ldr_data).InLoadOrderModuleList.Flink as *mut LDR_DATA_TABLE_ENTRY;
+
+        while !(*module_list).DllBase.is_null() {
+            let dll_base = (*module_list).DllBase;
+            let size = (*module_list).SizeOfImage as usize;
+
+            module_list = (*module_list).InLoadOrderLinks.Flink as *mut LDR_DATA_TABLE_ENTRY;
+
+            let buffer = std::slice::from_raw_parts::<u8>(dll_base as _, size);
+
+            let opts = goblin::pe::options::ParseOptions{ resolve_rva: false };
+            let res = goblin::pe::PE::parse_with_opts(buffer, &opts);
+
+            // temp solution
+            if res.is_err() { continue }
+            let parsed = res?;
+
+            if !parsed.is_lib && parsed.name.unwrap() != "ntdll.dll" { continue }
+
+            for e in parsed.exports {
+                match e.name {
+                    Some(symbol) => {
+                        if symbol.len() < 2 { continue }
+
+                        let sym_0 = symbol.chars().nth(0).unwrap();
+                        let sym_1 = symbol.chars().nth(1).unwrap();
+
+                        if (sym_0 == 'Z') && sym_1 == 'w' {
+                            let mut function_name = symbol.to_string();
+
+                            if symbol.chars().nth(0).unwrap() != 'N' { function_name.replace_range(0..2, "Nt"); };
+
+                            function_table.insert(function_name, e.offset);
+                        }
+                    },
+                    None => {}
+                }
+            }
+
+            let mut function_table_vec: Vec<(String, usize)> = function_table.into_iter().collect();
+            function_table_vec.sort_by(|x, y| x.1.cmp(&y.1));
+
+            let mut syscall_table = HashMap::new();
+
+            for (i, val) in function_table_vec.into_iter().enumerate() {
+                syscall_table.insert(val.0, i);
+            }
+
+            return Ok(syscall_table);
+        }
+
+        bail!("could not find ntdll");
+    }
+}

crate/pe-tools/src/x86.rs

@@ -1,6 +1,7 @@
 #[allow(non_snake_case)]
 use crate::shared::*;
 
+use std::collections::HashMap;
 use std::mem::size_of;
 use std::ptr::null_mut;
 
@@ -24,7 +25,6 @@ use winapi::{
     um::{
         memoryapi::{ReadProcessMemory, WriteProcessMemory},
         winnt::{
-            IMAGE_DIRECTORY_ENTRY_BASERELOC, IMAGE_DIRECTORY_ENTRY_IMPORT,
             IMAGE_ORDINAL32, IMAGE_IMPORT_BY_NAME, IMAGE_IMPORT_DESCRIPTOR,
             IMAGE_DOS_HEADER, IMAGE_NT_HEADERS32,
             IMAGE_SNAP_BY_ORDINAL32, IMAGE_THUNK_DATA32,
@@ -328,26 +328,36 @@ impl PE_Container<'_> {
     }
 }
 
-pub fn search_proc_address(function_name: impl Into<String>) -> anyhow::Result<*mut c_void> {
+pub fn search_proc_from_loaded_module<T>(function_name: impl Into<String>) -> anyhow::Result<T> {
     unsafe {
         let function_name = function_name.into();
         let ppeb = __readfsdword(0x30) as *mut PEB32;
 
-        let p_peb_ldr_data = (*ppeb).Ldr as *mut PEB_LDR_DATA32;;
+        let p_peb_ldr_data = (*ppeb).Ldr as *mut PEB_LDR_DATA32;
         let mut module_list = (*p_peb_ldr_data).InLoadOrderModuleList.Flink as *mut LDR_DATA_TABLE_ENTRY32;
 
         while !((*module_list).DllBase as PVOID).is_null() {
             let dll_base = (*module_list).DllBase;
+            let size = (*module_list).SizeOfImage as usize;
 
             module_list = (*module_list).InLoadOrderLinks.Flink as *mut LDR_DATA_TABLE_ENTRY32;
 
-            let dll_container = PE_Container::new(0x0 as _, dll_base as _)?;
+            let buffer = std::slice::from_raw_parts::<u8>(dll_base as _, size);
 
-            for e in &dll_container.pe.exports {
+            let opts = goblin::pe::options::ParseOptions{ resolve_rva: false };
+            let res = goblin::pe::PE::parse_with_opts(buffer, &opts);
+
+            // temp solution
+            if res.is_err() { continue }
+            let parsed = res?;
+
+            if !parsed.is_lib { continue }
+
+            for e in parsed.exports {
                 match e.name {
                     Some(symbol) => {
                         if symbol == function_name {
-                            return Ok(dll_container.to_va(e.offset as _));
+                            return Ok(ptr_to_fn::<T>((dll_base as usize + e.offset) as _));
                         }
                     },
                     None => {}
@@ -358,3 +368,64 @@ pub fn search_proc_address(function_name: impl Into<String>) -> anyhow::Result<*
         bail!("could not find {}", function_name);
     }
 }
+
+pub fn get_syscall_table() -> anyhow::Result<HashMap<String, usize>> {
+    unsafe {
+        let mut function_table = HashMap::new();
+        let ppeb = __readfsdword(0x30) as *mut PEB32;
+
+        let p_peb_ldr_data = (*ppeb).Ldr as *mut PEB_LDR_DATA32;
+        let mut module_list = (*p_peb_ldr_data).InLoadOrderModuleList.Flink as *mut LDR_DATA_TABLE_ENTRY32;
+
+        while !((*module_list).DllBase as PVOID).is_null() {
+            let dll_base = (*module_list).DllBase;
+            let size = (*module_list).SizeOfImage as usize;
+
+            module_list = (*module_list).InLoadOrderLinks.Flink as *mut LDR_DATA_TABLE_ENTRY32;
+
+            let buffer = std::slice::from_raw_parts::<u8>(dll_base as _, size);
+
+            let opts = goblin::pe::options::ParseOptions{ resolve_rva: false };
+            let res = goblin::pe::PE::parse_with_opts(buffer, &opts);
+
+            // temp solution
+            if res.is_err() { continue }
+            let parsed = res?;
+
+            if !parsed.is_lib && parsed.name.unwrap() != "ntdll.dll" { continue }
+
+            for e in parsed.exports {
+                match e.name {
+                    Some(symbol) => {
+                        if symbol.len() < 2 { continue }
+
+                        let sym_0 = symbol.chars().nth(0).unwrap();
+                        let sym_1 = symbol.chars().nth(1).unwrap();
+
+                        if (sym_0 == 'Z') && sym_1 == 'w' {
+                            let mut function_name = symbol.to_string();
+
+                            if symbol.chars().nth(0).unwrap() != 'N' { function_name.replace_range(0..2, "Nt"); };
+
+                            function_table.insert(function_name, e.offset);
+                        }
+                    },
+                    None => {}
+                }
+            }
+
+            let mut function_table_vec: Vec<(String, usize)> = function_table.into_iter().collect();
+            function_table_vec.sort_by(|x, y| x.1.cmp(&y.1));
+
+            let mut syscall_table = HashMap::new();
+
+            for (i, val) in function_table_vec.into_iter().enumerate() {
+                syscall_table.insert(val.0, i);
+            }
+
+            return Ok(syscall_table);
+        }
+
+        bail!("could not find ntdll");
+    }
+}