From c14fd3dd12e25cae7db3c43959a5f9eabb21ee16 Mon Sep 17 00:00:00 2001 From: Kamela Date: Fri, 8 Nov 2024 23:04:48 +0200 Subject: [PATCH] Switch over to `clap` for command line parsing --- Cargo.lock | 129 ++++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 3 +- src/main.rs | 53 +++++++++++++-------- 3 files changed, 166 insertions(+), 19 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 27855be..aa2a5b9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,21 +2,129 @@ # It is not intended for manual editing. version = 3 +[[package]] +name = "anstream" +version = "0.6.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8acc5369981196006228e28809f761875c0327210a891e941f4c683b3a99529b" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55cc3b69f167a1ef2e161439aa98aed94e6028e5f9a59be9a6ffb47aef1651f9" + +[[package]] +name = "anstyle-parse" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b2d16507662817a6a20a9ea92df6652ee4f94f914589377d69f3b21bc5798a9" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79947af37f4177cfead1110013d678905c37501914fba0efea834c3fe9a8d60c" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2109dbce0e72be3ec00bed26e6a7479ca384ad226efdd66db8fa2e3a38c83125" +dependencies = [ + "anstyle", + "windows-sys", +] + [[package]] name = "cfg-if" version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "clap" +version = "4.5.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97f376d85a664d5837dbae44bf546e6477a679ff6610010f17276f686d867e8" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.5.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19bc80abd44e4bed93ca373a0704ccbd1b710dc5749406201bb018272808dc54" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ac6a0c7b1a9e9a5186361f67dfa1b88213572f427fb9ab038efb2bd8c582dab" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1462739cb27611015575c0c11df5df7601141071f07518d56fcc1be504cbec97" + +[[package]] +name = "colorchoice" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b63caa9aa9397e2d9480a9b13673856c78d8ac123288526c37d7839f2a86990" + [[package]] name = "dll_injector" version = "0.1.0" dependencies = [ + "clap", "tracing", "tracing-subscriber", "windows", ] +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" + [[package]] name = "lazy_static" version = "1.5.0" @@ -90,6 +198,12 @@ version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "syn" version = "2.0.87" @@ -174,6 +288,12 @@ version = "1.0.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "valuable" version = "0.1.0" @@ -266,6 +386,15 @@ dependencies = [ "windows-targets", ] +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-targets" version = "0.52.6" diff --git a/Cargo.toml b/Cargo.toml index c909526..a3e725b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,4 +19,5 @@ windows = { version = "0.58.0", features = [ "Win32_Security" ] } tracing = { version = "0.1.40", features = ["max_level_debug"] } -tracing-subscriber = { version = "0.3.18", features = ["ansi"] } \ No newline at end of file +tracing-subscriber = { version = "0.3.18", features = ["ansi"] } +clap = { version = "4.5.20", features = ["derive"] } \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 5756785..68b6496 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,8 +8,23 @@ use windows::Win32::System::LibraryLoader::{GetModuleHandleA, GetProcAddress}; use windows::Win32::Foundation::{GetLastError, FARPROC, HANDLE, HMODULE}; use windows::Win32::System::Diagnostics::Debug::WriteProcessMemory; use tracing_subscriber::util::SubscriberInitExt; -use windows::core::{s, HRESULT, PCWSTR, Error}; +use windows::core::{s, PCWSTR, Error}; use tracing::{error, info, warn}; +use clap::Parser; + +#[derive(Parser, Debug)] +#[command(version, about, long_about = None)] +struct Args { + /// The name of the target process + #[arg(required = true)] + target: OsString, + /// The path to the DLL file that will be injected + #[arg(required = true)] + dll: OsString, + /// Enables verbose logging for events + #[arg(short, long)] + verbose: bool, +} fn main() -> Result<(), Error> { tracing_subscriber::fmt() @@ -18,27 +33,27 @@ fn main() -> Result<(), Error> { .finish() .init(); - let mut args = std::env::args_os().collect::>(); + let args = Args::parse(); - if args.len() < 3 || args.len() > 3 { - println!("Usage: dll_injector [TARGET_NAME] [PATH_TO_DLL]"); - return Err(Error::new(HRESULT::default(), "Provide one item per field")) - } + let mut critical_args = [args.target, args.dll]; - for arg in &mut args { + for arg in &mut critical_args { arg.push("\0") } - let dll_name = args[2].clone(); + let dll_name = critical_args[1].clone(); - let target_handle = process_enumerate_and_search(PCWSTR::from_raw( - args[1].clone() - .to_string_lossy() - .replace('\u{FFFD}', "") - .encode_utf16() - .collect::>() - .as_mut_ptr() - ))?; + let target_handle = process_enumerate_and_search( + PCWSTR::from_raw( + critical_args[0].clone() + .to_string_lossy() + .replace('\u{FFFD}', "") + .encode_utf16() + .collect::>() + .as_mut_ptr() + ), + args.verbose + )?; inject_dll( PCWSTR::from_raw( @@ -55,7 +70,7 @@ fn main() -> Result<(), Error> { Ok(()) } -fn process_enumerate_and_search(process_name: PCWSTR) -> Result { +fn process_enumerate_and_search(process_name: PCWSTR, verbose: bool) -> Result { let mut process_handle: Option = None; let snapshot_handle: HANDLE; let mut process_entry: PROCESSENTRY32 = unsafe { std::mem::zeroed() }; @@ -83,7 +98,9 @@ fn process_enumerate_and_search(process_name: PCWSTR) -> Result { CStr::from_ptr(process_entry.szExeFile.clone().as_ptr() as *mut c_char) }; - info!("Checking: {sz_exe_file:?}"); + if verbose { + info!("Checking: {sz_exe_file:?}"); + } if unsafe { PCWSTR::from_raw({ let mut str = sz_exe_file