win_iter 0.1.4

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

use crate::{
    modules::Module,
    winapi_bindings::{
        CloseHandle, CreateToolhelp32Snapshot, INVALID_HANDLE_VALUE, IsWow64Process, PROCESS_QUERY_INFORMATION, PROCESSENTRY32,
        Process32First, Process32Next, TH32CS_SNAPPROCESS,
    },
    winapi_wrappers::wrappers::{
        close_handle, create_tool_help32_snapshot, cstr_to_str, is_valid_handle, open_process, process32_first, process32_next,
    },
};

pub struct Process {
    pub name: String,
    pub id: u32,
    pub architecture: Architecture,
}

pub enum Architecture {
    X64 = 64,
    X86 = 86,
    UNKNOWN = 0,
}

impl fmt::Display for Process {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "x{} | \"{}\" [{}]", self.architecture, self.name, self.id)
    }
}

impl fmt::Display for Architecture {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let val = match self {
            Architecture::X64 => 64,
            Architecture::X86 => 86,
            Architecture::UNKNOWN => 0,
        };
        write!(f, "{}", val)
    }
}

pub fn fetch_all_processes() -> Vec<Process> {
    let mut procs = Vec::new();
    let mut pe = PROCESSENTRY32::default();

    pe.dwSize = size_of::<PROCESSENTRY32>() as u32;

    let Some(snap) = create_tool_help32_snapshot(TH32CS_SNAPPROCESS, 0) else {
        return procs;
    };

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

    if !process32_first(snap, &mut pe) {
        close_handle(snap);
        return procs;
    }

    loop {
        let pe_name = cstr_to_str(pe.szExeFile.as_ptr());
        let arch = fetch_process_architecture(pe.th32ProcessID);
        procs.push(Process {
            name: pe_name,
            id: pe.th32ProcessID,
            architecture: arch,
        });

        if !process32_next(snap, &mut pe) {
            break;
        }
    }

    close_handle(snap);

    procs
}

pub fn fetch_process(process_name: &str) -> Option<Process> {
    let mut pe = PROCESSENTRY32::default();

    pe.dwSize = size_of::<PROCESSENTRY32>() as u32;

    let Some(snap) = create_tool_help32_snapshot(TH32CS_SNAPPROCESS, 0) else {
        return None;
    };

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

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

    loop {
        let pe_name = cstr_to_str(pe.szExeFile.as_ptr());
        if pe_name.as_str() == process_name {
            return Some(Process {
                name: pe_name,
                id: pe.th32ProcessID,
                architecture: fetch_process_architecture(pe.th32ProcessID),
            });
        }

        if !process32_next(snap, &mut pe) {
            break;
        }
    }

    close_handle(snap);

    None
}

pub fn fetch_process_architecture(process_id: u32) -> Architecture {
    let handle = match open_process(PROCESS_QUERY_INFORMATION, false, process_id) {
        Some(h) => h,
        None => return Architecture::UNKNOWN,
    };
    let mut is_x64 = 0;

    unsafe { IsWow64Process(handle, &mut is_x64) };

    close_handle(handle);

    if is_x64 == 0 {
        return Architecture::X64;
    } else {
        return Architecture::X86;
    }
}