win_iter 0.1.2

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

use crate::winapi_bindings::{
    CloseHandle, CreateToolhelp32Snapshot, INVALID_HANDLE_VALUE, PROCESSENTRY32, Process32First, Process32Next,
    TH32CS_SNAPPROCESS,
};

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

/// Fetches all currently running processes
///
/// # Returns
/// A `Vec` containing all captured processes in the form of a `Process` struct, or an empty `Vec`.
///
/// # Example
/// ```no_run
/// use win_iter::fetch_all_processes;
///
/// let processes = fetch_all_processes();
/// if processes.is_empty() {
///     eprintln!("There was no captured processes!");
/// } else {
///     println!("Successfully fetched {} proceseses!", processes.len())
/// }
/// ```
pub fn fetch_all_processes() -> Vec<Process> {
    let mut procs = Vec::new();
    let mut pe32 = PROCESSENTRY32::default();

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

    let snap = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
    if snap == INVALID_HANDLE_VALUE {
        panic!("CreateToolhelp32Snapshot failed!");
    }

    unsafe {
        if Process32First(snap, &mut pe32) == 0 {
            CloseHandle(snap);
            panic!("Process32First failed!");
        }

        loop {
            let pe_name = CStr::from_ptr(pe32.szExeFile.as_ptr()).to_string_lossy();
            procs.push(Process {
                name: pe_name.to_string(),
                id: pe32.th32ProcessID,
                architecture: Architecture::UNKNOWN,
            });

            if Process32Next(snap, &mut pe32) == 0 {
                break;
            }
        }

        CloseHandle(snap);
    }
    procs
}

/// Fetches the process id of the given process name if it exists.
///
/// # Arguments
/// * `process_name` - The process name to look for.
///
/// # Returns
/// A `u32` containing the process id or `0` if it wasn't found.
///
/// # Panics
/// If the snapshot is invalid or if process iteration fails.
///
/// # Example
/// ```no_run
/// use win_iter::fetch_process_id;
///
/// let pid = fetch_process_id("Chrome.exe".to_owned());
///
/// if pid == 0 {
///     eprintln!("Failed to fetch the process id");
/// } else {
///     println!("Successfully found process id: {}", pid);
/// }
/// ```
pub fn fetch_process_id(process_name: String) -> u32 {
    let mut pe = PROCESSENTRY32::default();
    pe.dwSize = size_of::<PROCESSENTRY32>() as u32;
    let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };

    if snapshot == INVALID_HANDLE_VALUE {
        unsafe {
            CloseHandle(snapshot);
        }
        panic!("CreateToolhelp32Snapshot() failed!");
    }

    unsafe {
        if Process32First(snapshot, &mut pe) == 0 {
            unsafe {
                CloseHandle(snapshot);
            }
            panic!("Process32First() failed!");
        } else {
            let pe_name = CStr::from_ptr(pe.szExeFile.as_ptr()).to_string_lossy();
            if pe_name == process_name {
                CloseHandle(snapshot);
                return pe.th32ProcessID;
            }
        }

        loop {
            let pe_name = CStr::from_ptr(pe.szExeFile.as_ptr()).to_string_lossy();
            if pe_name == process_name {
                CloseHandle(snapshot);
                return pe.th32ProcessID;
            }

            if Process32Next(snapshot, &mut pe) == 0 {
                break;
            }
        }
        0
    }
}