diff --git a/Cargo.toml b/Cargo.toml index a4c9024..1ec3392 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,7 @@ windows = { version = "0.62.0", features = [ "Win32_NetworkManagement_WindowsFirewall", "Win32_Security", "Win32_System_Com", + "Win32_System_Diagnostics_ToolHelp", "Win32_System_Threading", "Win32_UI_Input_KeyboardAndMouse", "Win32_UI_Shell", diff --git a/src/util.rs b/src/util.rs index 547cb56..8459943 100644 --- a/src/util.rs +++ b/src/util.rs @@ -2,4 +2,5 @@ pub mod consts; pub mod countdown; pub mod log; pub mod persistent_state; +pub mod system_info; pub mod win; diff --git a/src/util/system_info.rs b/src/util/system_info.rs new file mode 100644 index 0000000..29ed1ca --- /dev/null +++ b/src/util/system_info.rs @@ -0,0 +1,149 @@ +use std::{ + os::windows::process::CommandExt, + path::{Path, PathBuf}, + process::Command, +}; +use windows::{ + Win32::System::{ + Diagnostics::ToolHelp::{ + CreateToolhelp32Snapshot, PROCESSENTRY32, Process32First, Process32Next, + TH32CS_SNAPPROCESS, + }, + Threading::{ + CREATE_NO_WINDOW, OpenProcess, PROCESS_NAME_WIN32, PROCESS_QUERY_LIMITED_INFORMATION, + QueryFullProcessImageNameW, + }, + }, + core::PWSTR, +}; + +#[derive(Clone, Debug)] +struct Process { + pid: u32, + name: String, + exe: Option, +} + +impl Process { + fn pid(&self) -> u32 { + self.pid + } + + fn name(&self) -> &str { + &self.name + } + + fn exe(&self) -> Option<&Path> { + self.exe.as_deref() + } + + fn kill(&self) -> bool { + let mut taskkill = Command::new("taskkill.exe"); + taskkill.creation_flags(CREATE_NO_WINDOW.0); + taskkill.arg("/F").arg("/PID").arg(&self.pid.to_string()); + match taskkill.output() { + Ok(output) => output.status.success(), + Err(_) => false, + } + } +} + +#[derive(Debug)] +struct SystemInfo { + processes: Vec, +} + +impl SystemInfo { + fn new() -> Self { + Self { + processes: Vec::new(), + } + } + + // TODO: unfuck this retarded function + fn refresh(&mut self) { + let mut processes = Vec::new(); + let snapshot_handle = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0).unwrap() }; + let mut process_entry = PROCESSENTRY32::default(); + process_entry.dwSize = size_of::() as u32; + unsafe { Process32First(snapshot_handle, &mut process_entry).unwrap() }; + let process_handle_result = unsafe { + OpenProcess( + PROCESS_QUERY_LIMITED_INFORMATION, + false, + process_entry.th32ProcessID, + ) + }; + let mut exename = [0u16; 260]; + let mut dwsize = exename.len() as u32; + let exe_full_path = if let Ok(process_handle) = process_handle_result { + let image_name_result = unsafe { + QueryFullProcessImageNameW( + process_handle, + PROCESS_NAME_WIN32, + PWSTR(exename.as_mut_ptr()), + &mut dwsize, + ) + }; + match image_name_result { + Ok(_) => Some(PathBuf::from( + unsafe { PWSTR(exename.as_mut_ptr()).to_string() }.unwrap(), + )), + Err(_) => None, + } + } else { + None + }; + processes.push(Process { + pid: process_entry.th32ProcessID, + name: c_char_arr_to_string(&process_entry.szExeFile), + exe: exe_full_path, + }); + while unsafe { Process32Next(snapshot_handle, &mut process_entry) }.is_ok() { + let process_handle_result = unsafe { + OpenProcess( + PROCESS_QUERY_LIMITED_INFORMATION, + false, + process_entry.th32ProcessID, + ) + }; + let mut exename = [0u16; 260]; + let mut dwsize = exename.len() as u32; + let exe_full_path = if let Ok(process_handle) = process_handle_result { + let image_name_result = unsafe { + QueryFullProcessImageNameW( + process_handle, + PROCESS_NAME_WIN32, + PWSTR(exename.as_mut_ptr()), + &mut dwsize, + ) + }; + match image_name_result { + Ok(_) => Some(PathBuf::from( + unsafe { PWSTR(exename.as_mut_ptr()).to_string() }.unwrap(), + )), + Err(_) => None, + } + } else { + None + }; + processes.push(Process { + pid: process_entry.th32ProcessID, + name: c_char_arr_to_string(&process_entry.szExeFile), + exe: exe_full_path, + }); + } + self.processes = processes; + } + + fn processes(&self) -> &[Process] { + &self.processes + } +} + +fn c_char_arr_to_string(arr: &[i8]) -> String { + arr.iter() + .take_while(|&&b| b != 0) + .map(|&b| b as u8 as char) + .collect() +}