win_iter 0.1.2

use std::{ffi::CStr, fmt};

use crate::winapi_bindings::{
    CloseHandle, CreateToolhelp32Snapshot, INVALID_HANDLE_VALUE, MODULEENTRY32, Module32First, Module32Next,
    TH32CS_SNAPMODULE,
};

pub struct Module {
    pub name: String,
    pub base_size: u32,
    pub base_address: *mut u8,
}

impl fmt::Display for Module {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "\"{}\" | 0x{:X}", self.name, self.base_address as usize)
    }
}

/// Fetches all currently loaded modules inside of a process
///
/// # Arguments
/// * `process_id` - The process id of which you want to search in.
///
/// # Returns
/// A `Vec` containing all captured modules in the form of a `Module` struct, or an empty `Vec`.
///
/// # Panics
/// If the snapshot is invalid or if module iteration fails.
///
/// # Example
/// ```no_run
/// use win_iter::{fetch_process_id, fetch_all_modules};
///
/// let pid = fetch_process_id("Chrome.exe".to_owned());
///
/// if pid == 0 {
///     panic!("PID was 0");
/// }
///
/// let modules = fetch_all_modules(pid);
/// if modules.is_empty() {
///     panic!("There was no captured modules!");
/// } else {
///     for module in modules.iter() {
///         println!("{}", module);
/// }
/// ```
pub fn fetch_all_modules(process_id: u32) -> Vec<Module> {
    let mut modules = Vec::new();

    let mut pe32 = MODULEENTRY32::default();
    pe32.dwSize = size_of::<MODULEENTRY32>() as u32;

    let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, process_id) };
    if snapshot == INVALID_HANDLE_VALUE {
        panic!("CreateToolhelp32Snaphot() failed!");
    }

    unsafe {
        if Module32First(snapshot, &mut pe32) == 0 {
            CloseHandle(snapshot);
            panic!("Module32First() failed!");
        } else {
            loop {
                let mod_name = CStr::from_ptr(pe32.szModule.as_ptr()).to_string_lossy();
                modules.push(Module {
                    name: mod_name.to_string(),
                    base_size: pe32.modBaseSize,
                    base_address: pe32.modBaseAddr,
                });

                if Module32Next(snapshot, &mut pe32) == 0 {
                    break;
                }
            }
        }
    }

    modules
}

/// Searches through all loaded modules inside of a process to fetch the desired one
///
/// # Arguments
/// * `process_id` - The process id of which you want to search in.
/// * `module_name` - The name of the module you'd like to search for.
///
/// # Returns
/// A `u32` containing the module base address or `0` if it wasn't found.
///
/// # Panics
/// If the snapshot is invalid or if module iteration fails.
///
/// # Example
/// ```no_run
/// use win_iter::{fetch_process_id, fetch_module_address};
///
/// let pid = fetch_process_id("Chrome.exe".to_owned());
///
/// if pid == 0 {
///     panic!("PID was 0");
/// }
///
/// let ntdll = fetch_module_address(pid, "ntdll.dll".to_owned());
/// if ntdll == 0 {
///     panic!("Failed to fetch ntdll!");
/// } else {
///     println!("Address of ntdll: {:x}", ntdll);
/// }
/// ```
pub fn fetch_module_address(process_id: u32, module_name: String) -> usize {
    let mut pe32 = MODULEENTRY32::default();
    pe32.dwSize = size_of::<MODULEENTRY32>() as u32;

    let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, process_id) };
    if snapshot == INVALID_HANDLE_VALUE {
        panic!("CreateToolhelp32Snaphot() failed!");
    }

    unsafe {
        if Module32First(snapshot, &mut pe32) == 0 {
            CloseHandle(snapshot);
            panic!("Module32First() failed!");
        } else {
            loop {
                let mod_name = CStr::from_ptr(pe32.szModule.as_ptr()).to_string_lossy();
                if mod_name == module_name {
                    return pe32.modBaseAddr as usize;
                }

                if Module32Next(snapshot, &mut pe32) == 0 {
                    break;
                }
            }
        }
    }
    0
}