win_iter 0.1.4

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

use crate::{
    Process,
    winapi_bindings::{
        CloseHandle, CreateToolhelp32Snapshot, GetProcAddress, INVALID_HANDLE_VALUE, MODULEENTRY32, Module32First, Module32Next,
        TH32CS_SNAPMODULE,
    },
    winapi_wrappers::wrappers::{
        close_handle, create_tool_help32_snapshot, cstr_to_str, is_valid_handle, module32_first, module32_next,
    },
};

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)
    }
}

impl Process {
    /// Iterates over each loaded module inside of the process
    ///
    /// # Returns
    ///
    /// A `Vec<Module>` containing a `Module` for each loaded module,
    /// or an empty `Vec` if enumeration fails.
    ///
    /// # Example
    /// ```no_run
    /// use win_iter::{fetch_loaded_modules, fetch_process};
    ///
    /// fn main() {
    ///     let proc = fetch_process("brave.exe");
    ///     for loaded_mod in proc.fetch_loaded_modules().iter() {
    ///         println!("Found module: {}, base address: {:#X}",
    ///             loaded_mod.name, loaded_mod.base_address as usize
    ///         );
    ///     }
    /// }
    /// ```
    pub fn fetch_loaded_modules(&self) -> Vec<Module> {
        let mut modules = Vec::new();

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

        let Some(snap) = create_tool_help32_snapshot(TH32CS_SNAPMODULE, self.id) else {
            return modules;
        };

        if !is_valid_handle(snap) {
            close_handle(snap);
            return modules;
        }

        if !module32_first(snap, &mut pe) {
            close_handle(snap);
            return modules;
        }

        loop {
            let module_name = cstr_to_str(pe.szModule.as_ptr());
            modules.push(Module {
                name: module_name,
                base_size: pe.modBaseSize,
                base_address: pe.modBaseAddr,
            });

            if !module32_next(snap, &mut pe) {
                close_handle(snap);
                break;
            }
        }
        modules
    }

    /// Iterates over each loaded module inside of the process to find the target.
    ///
    /// # Arguments
    /// * `module_name` - The module name to search for.
    ///
    /// # Returns
    /// An `Option<Module>` if the module was found or `None`.
    ///
    /// # Example
    /// ```no_run
    /// use win_iter::{fetch_loaded_module, fetch_process};
    ///
    /// fn main() {
    ///     let proc = fetch_process("brave.exe");
    ///     if let Some(ntdll) = proc.fetch_loaded_module("ntdll.dll") {
    ///         println!("ntdll: {:#X}", ntdll.base_address as usize);
    ///     } else {
    ///         println!("ntdll.dll wasn't found!");
    ///     }
    /// }
    /// ```
    pub fn fetch_loaded_module(&self, module_name: &str) -> Option<Module> {
        let mut pe = MODULEENTRY32::default();
        pe.dwSize = size_of::<MODULEENTRY32>() as u32;

        let Some(snap) = create_tool_help32_snapshot(TH32CS_SNAPMODULE, self.id) else {
            return None;
        };

        if !is_valid_handle(snap) {
            close_handle(snap);
            return None;
        }

        if !module32_first(snap, &mut pe) {
            close_handle(snap);
            return None;
        }

        loop {
            let name = cstr_to_str(pe.szModule.as_ptr());
            if name.as_str() == module_name {
                close_handle(snap);
                return Some(Module {
                    name: name,
                    base_size: pe.modBaseSize,
                    base_address: pe.modBaseAddr,
                });
            }

            if !module32_next(snap, &mut pe) {
                close_handle(snap);
                break;
            }
        }
        None
    }
}