Rename util directory to os

src/os/address.rs

@@ -0,0 +1,137 @@
+use std::fmt::{LowerHex, UpperHex};
+use std::ops::{Add, AddAssign, Sub, SubAssign};
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq, PartialOrd)]
+#[repr(transparent)]
+pub struct Address(pub usize);
+
+impl Address {
+    #[inline]
+    pub fn add(&self, value: usize) -> Self {
+        Self(self.0 + value)
+    }
+
+    #[inline]
+    pub fn is_zero(&self) -> bool {
+        self.0 == 0
+    }
+
+    #[inline]
+    pub fn sub(&self, value: usize) -> Self {
+        Self(self.0 - value)
+    }
+
+    #[inline]
+    pub fn as_ptr<T>(&self) -> *const T {
+        self.0 as *const T
+    }
+
+    #[inline]
+    pub fn as_mut_ptr<T>(&self) -> *mut T {
+        self.0 as *mut T
+    }
+}
+
+impl From<usize> for Address {
+    fn from(value: usize) -> Self {
+        Self(value)
+    }
+}
+
+impl From<*const u8> for Address {
+    fn from(value: *const u8) -> Self {
+        Self(value as usize)
+    }
+}
+
+impl From<*mut u8> for Address {
+    fn from(value: *mut u8) -> Self {
+        Self(value as usize)
+    }
+}
+
+impl From<Address> for usize {
+    fn from(value: Address) -> Self {
+        value.0
+    }
+}
+
+impl From<Address> for *const u8 {
+    fn from(value: Address) -> Self {
+        value.0 as *const u8
+    }
+}
+
+impl From<Address> for *mut u8 {
+    fn from(value: Address) -> Self {
+        value.0 as *mut u8
+    }
+}
+
+impl Add<usize> for Address {
+    type Output = Self;
+
+    fn add(self, rhs: usize) -> Self::Output {
+        Self(self.0 + rhs)
+    }
+}
+
+impl Add<Address> for Address {
+    type Output = Self;
+
+    fn add(self, rhs: Address) -> Self::Output {
+        Self(self.0 + rhs.0)
+    }
+}
+
+impl AddAssign<usize> for Address {
+    fn add_assign(&mut self, rhs: usize) {
+        self.0 += rhs;
+    }
+}
+
+impl AddAssign<Address> for Address {
+    fn add_assign(&mut self, rhs: Address) {
+        self.0 += rhs.0;
+    }
+}
+
+impl Sub<usize> for Address {
+    type Output = Self;
+
+    fn sub(self, rhs: usize) -> Self::Output {
+        Self(self.0 - rhs)
+    }
+}
+
+impl Sub<Address> for Address {
+    type Output = Self;
+
+    fn sub(self, rhs: Address) -> Self::Output {
+        Self(self.0 - rhs.0)
+    }
+}
+
+impl SubAssign<usize> for Address {
+    fn sub_assign(&mut self, rhs: usize) {
+        self.0 -= rhs;
+    }
+}
+
+impl SubAssign<Address> for Address {
+    fn sub_assign(&mut self, rhs: Address) {
+        self.0 -= rhs.0;
+    }
+}
+
+impl UpperHex for Address {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{:#X}", self.0)
+    }
+}
+
+impl LowerHex for Address {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{:#x}", self.0)
+    }
+}

src/os/mod.rs

@@ -0,0 +1,7 @@
+pub use address::Address;
+pub use module::Module;
+pub use process::Process;
+
+pub mod address;
+pub mod module;
+pub mod process;

src/os/module.rs

@@ -0,0 +1,77 @@
+use anyhow::Result;
+
+use goblin::pe::export::Export;
+use goblin::pe::import::Import;
+use goblin::pe::options::ParseOptions;
+use goblin::pe::section_table::SectionTable;
+use goblin::pe::PE;
+
+use super::Address;
+
+pub struct Module<'a> {
+    pub name: &'a str,
+    pub data: &'a [u8],
+    pub pe: PE<'a>,
+}
+
+impl<'a> Module<'a> {
+    pub fn parse(name: &'a str, data: &'a [u8]) -> Result<Self> {
+        let pe = PE::parse_with_opts(
+            data,
+            &ParseOptions {
+                parse_attribute_certificates: false,
+                resolve_rva: false,
+            },
+        )?;
+
+        Ok(Self { name, data, pe })
+    }
+
+    #[inline]
+    pub fn base(&self) -> Address {
+        self.pe.image_base.into()
+    }
+
+    #[inline]
+    pub fn exports(&self) -> &[Export] {
+        &self.pe.exports
+    }
+
+    #[inline]
+    pub fn imports(&self) -> &[Import] {
+        &self.pe.imports
+    }
+
+    #[inline]
+    pub fn export_by_name(&self, name: &str) -> Option<Address> {
+        self.pe
+            .exports
+            .iter()
+            .find(|e| e.name.unwrap() == name)
+            .map(|e| (self.pe.image_base + e.rva).into())
+    }
+
+    #[inline]
+    pub fn import_by_name(&self, name: &str) -> Option<Address> {
+        self.pe
+            .imports
+            .iter()
+            .find(|i| i.name.to_string() == name)
+            .map(|i| (self.pe.image_base + i.rva).into())
+    }
+
+    #[inline]
+    pub fn sections(&self) -> &[SectionTable] {
+        &self.pe.sections
+    }
+
+    #[inline]
+    pub fn size(&self) -> u32 {
+        self.pe
+            .header
+            .optional_header
+            .expect("optional header not found")
+            .windows_fields
+            .size_of_image
+    }
+}

src/os/process.rs

@@ -0,0 +1,228 @@
+use std::collections::HashMap;
+use std::ffi::{c_void, CStr};
+use std::mem;
+use std::ptr;
+
+use anyhow::{bail, Result};
+
+use windows::Win32::Foundation::{CloseHandle, HANDLE};
+use windows::Win32::System::Diagnostics::Debug::ReadProcessMemory;
+use windows::Win32::System::Diagnostics::ToolHelp::*;
+use windows::Win32::System::Threading::{OpenProcess, PROCESS_ALL_ACCESS};
+
+use super::{Address, Module};
+
+#[derive(Debug)]
+pub struct Process {
+    id: u32,
+    handle: HANDLE,
+    modules: HashMap<String, Vec<u8>>,
+}
+
+impl Process {
+    pub fn new(name: &str) -> Result<Self> {
+        let id = Self::get_process_id_by_name(name)?;
+
+        let handle = unsafe { OpenProcess(PROCESS_ALL_ACCESS, false, id) }?;
+
+        let mut process = Self {
+            id,
+            handle,
+            modules: HashMap::new(),
+        };
+
+        process.parse_loaded_modules()?;
+
+        Ok(process)
+    }
+
+    pub fn find_pattern(&self, module_name: &str, pattern: &str) -> Option<Address> {
+        let module = self.get_module_by_name(module_name)?;
+
+        let pattern_bytes = Self::pattern_to_bytes(pattern);
+
+        for (i, window) in module.data.windows(pattern_bytes.len()).enumerate() {
+            if window
+                .iter()
+                .zip(&pattern_bytes)
+                .all(|(&x, &y)| x == y as u8 || y == -1)
+            {
+                return Some(module.base() + i);
+            }
+        }
+
+        None
+    }
+
+    pub fn get_module_by_name<'a>(&'a self, name: &'a str) -> Option<Module<'a>> {
+        self.modules
+            .get(name)
+            .map(|data| Module::parse(name, data).unwrap())
+    }
+
+    pub fn modules(&self) -> Result<Vec<Module>> {
+        let mut modules = Vec::new();
+
+        for (name, data) in &self.modules {
+            modules.push(Module::parse(name, data)?);
+        }
+
+        Ok(modules)
+    }
+
+    pub fn read_memory<T>(&self, address: Address) -> Result<T> {
+        let mut buffer: T = unsafe { mem::zeroed() };
+
+        self.read_memory_raw(
+            address,
+            &mut buffer as *const _ as *mut _,
+            mem::size_of::<T>(),
+        )?;
+
+        Ok(buffer)
+    }
+
+    pub fn read_memory_raw(
+        &self,
+        address: Address,
+        buffer: *mut c_void,
+        size: usize,
+    ) -> Result<()> {
+        unsafe {
+            ReadProcessMemory(
+                self.handle,
+                address.as_ptr(),
+                buffer,
+                size,
+                Some(ptr::null_mut()),
+            )
+        }
+        .map_err(|e| e.into())
+    }
+
+    pub fn read_string(&self, address: Address) -> Result<String> {
+        let mut buffer = Vec::new();
+
+        for i in 0.. {
+            match self.read_memory::<u8>(address + i) {
+                Ok(byte) if byte != 0 => buffer.push(byte),
+                _ => break,
+            }
+        }
+
+        Ok(String::from_utf8(buffer)?)
+    }
+
+    pub fn read_string_length(&self, address: Address, length: usize) -> Result<String> {
+        let mut buffer = vec![0; length];
+
+        self.read_memory_raw(address, buffer.as_mut_ptr() as *mut _, length)?;
+
+        if let Some(end) = buffer.iter().position(|&x| x == 0) {
+            buffer.truncate(end);
+        }
+
+        Ok(String::from_utf8(buffer)?)
+    }
+
+    pub fn resolve_jmp(
+        &self,
+        address: Address,
+        offset: Option<usize>,
+        length: Option<usize>,
+    ) -> Result<Address> {
+        // The displacement value can be negative.
+        let displacement = self.read_memory::<i32>(address.add(offset.unwrap_or(0x1)))?;
+
+        Ok(address
+            .add(length.unwrap_or(0x5))
+            .add(displacement as usize))
+    }
+
+    pub fn resolve_rip(
+        &self,
+        address: Address,
+        offset: Option<usize>,
+        length: Option<usize>,
+    ) -> Result<Address> {
+        // The displacement value can be negative.
+        let displacement = self.read_memory::<i32>(address.add(offset.unwrap_or(0x3)))?;
+
+        Ok(address
+            .add(length.unwrap_or(0x7))
+            .add(displacement as usize))
+    }
+
+    fn get_process_id_by_name(process_name: &str) -> Result<u32> {
+        let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) }?;
+
+        let mut entry = PROCESSENTRY32 {
+            dwSize: mem::size_of::<PROCESSENTRY32>() as u32,
+            ..Default::default()
+        };
+
+        unsafe {
+            Process32First(snapshot, &mut entry)?;
+
+            while Process32Next(snapshot, &mut entry).is_ok() {
+                let name = CStr::from_ptr(&entry.szExeFile as *const _ as *const _).to_str()?;
+
+                if name == process_name {
+                    return Ok(entry.th32ProcessID);
+                }
+            }
+        }
+
+        bail!("Process not found: {}", process_name)
+    }
+
+    fn parse_loaded_modules(&mut self) -> Result<()> {
+        let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, self.id) }?;
+
+        let mut entry = MODULEENTRY32 {
+            dwSize: mem::size_of::<MODULEENTRY32>() as u32,
+            ..Default::default()
+        };
+
+        unsafe {
+            Module32First(snapshot, &mut entry)?;
+
+            while Module32Next(snapshot, &mut entry).is_ok() {
+                let name = CStr::from_ptr(&entry.szModule as *const _ as *const _).to_str()?;
+
+                let mut data = vec![0; entry.modBaseSize as usize];
+
+                if let Ok(_) = self.read_memory_raw(
+                    entry.modBaseAddr.into(),
+                    data.as_mut_ptr() as *mut _,
+                    data.len(),
+                ) {
+                    self.modules.insert(name.to_string(), data);
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    fn pattern_to_bytes(pattern: &str) -> Vec<i32> {
+        pattern
+            .split_whitespace()
+            .map(|s| {
+                if s == "?" {
+                    -1
+                } else {
+                    i32::from_str_radix(s, 16).unwrap_or(0)
+                }
+            })
+            .collect()
+    }
+}
+
+impl Drop for Process {
+    fn drop(&mut self) {
+        if !self.handle.is_invalid() {
+            unsafe { CloseHandle(self.handle).unwrap() }
+        }
+    }
+}