winapi_main.rs

winapi_main.rs

use winapi::um::winuser::{FindWindowA, GetWindowThreadProcessId};
use std::{ptr};
use std::ffi::{CString, CStr};
use winapi::shared::windef::HWND;
use winapi::shared::minwindef::{DWORD, FALSE, TRUE};
use winapi::um::processthreadsapi::OpenProcess;
use winapi::um::winnt::{PROCESS_ALL_ACCESS, HANDLE, CHAR};
use std::mem::{MaybeUninit, size_of};
use std::io::{Error, ErrorKind};
use winapi::um::tlhelp32::{CreateToolhelp32Snapshot, TH32CS_SNAPMODULE, MODULEENTRY32, Module32First, Module32Next};
use winapi::um::handleapi::{INVALID_HANDLE_VALUE, CloseHandle};
use std::collections::HashMap;

#[derive(Debug)]
struct WinHandle {
    handle: HANDLE,
    process_id: DWORD,
    modules: HashMap<String, WinModule>,
}

#[derive(Debug)]
struct WinModule {
    handle: HANDLE,
    name: String,
    base_address: usize,
}

unsafe fn str_from_null_terminated_utf8(s: &[u8]) -> &str {
    CStr::from_ptr(s.as_ptr() as *const _).to_str().unwrap()
}

fn str_from_null_terminated_utf8_safe(s: &[u8]) -> &str {
    if s.iter().any(|&x| x == 0) {
        unsafe { str_from_null_terminated_utf8(s) }
    } else {
        std::str::from_utf8(s).unwrap()
    }
}

unsafe fn char_array_to_string(char_array: &[CHAR]) -> String {
    let data: Vec<u8> = char_array.iter().map(|&c| c as u8).collect();
    let data = str_from_null_terminated_utf8_safe(&data);
    data.to_ascii_lowercase()
}

impl WinHandle {
    fn new(window_name: &str) -> Result<WinHandle, Error> {
        let cstring_window_name = CString::new(window_name)?;
        let hwnd: HWND = unsafe { FindWindowA(ptr::null_mut(), cstring_window_name.as_ptr()) };
        if hwnd.is_null() {
            return Err(Error::new(ErrorKind::InvalidData, format!("Couldn't find window with name: {}", window_name)));
        }
        let mut process_id: DWORD = 0;
        unsafe { GetWindowThreadProcessId(hwnd, &mut process_id) };
        if process_id == 0 {
            return Err(Error::new(ErrorKind::InvalidData, format!("Couldn't find process_id for window with name: {}", window_name)));
        }
        let handle = unsafe { OpenProcess(PROCESS_ALL_ACCESS, FALSE, process_id) };
        if handle.is_null() {
            return Err(Error::new(ErrorKind::InvalidData, format!("Couldn't open process for window with name: {}", window_name)));
        }
        let modules = unsafe { WinHandle::take_snapshot(handle, process_id) }.ok_or(Error::new(ErrorKind::InvalidData, "Failed to take snapshot of modules"))?;
        Ok(WinHandle { handle, process_id, modules })
    }

    unsafe fn take_snapshot(handle: HANDLE, process_id: DWORD) -> Option<HashMap<String, WinModule>> {
        println!("{:?}", process_id);
        let module_snap = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, process_id);
        println!("{:?}", module_snap);
        let mut module_entry = MaybeUninit::<MODULEENTRY32>::uninit();
        if module_snap == INVALID_HANDLE_VALUE {
            return None;
        }
        module_entry.assume_init().dwSize = size_of::<MODULEENTRY32>() as u32;
        println!("{:?}", module_entry);
        let result = Module32First(module_snap, module_entry.as_mut_ptr() as *mut _);
        if result != TRUE {
            println!("AA");
            CloseHandle(module_snap);
            return None;
        }
        let mut modules: HashMap<String, WinModule> = HashMap::new();
        loop {
            let module = module_entry.assume_init();
            let module_name = char_array_to_string(&module.szModule);
            modules.insert(module_name.clone(), WinModule {
                handle,
                name: module_name.clone(),
                base_address: module.modBaseAddr as usize,
            });
            let next_result = Module32Next(module_snap, module_entry.as_mut_ptr() as *mut _);
            if next_result != TRUE {
                break;
            }
        }
        return Some(modules);
    }
}

fn main() {
    let process = WinHandle::new("Steam").unwrap();
    println!("PROCESS: {:?}", process);
    unsafe {
        for x in &process.modules {
            println!("{:?}", x);
        }
    }
}