use super::encryption::*;
 use super::helpers::*;
 
 pub struct DynamicResolverBuilder {
     dll_name: String,
     function_name: String,
     var_name: String,
     const_name: String,
     dll_array: String,
     fn_array: String,
     decrypt_fn: String,
     expect_clause: String,
     key_const_name: String,
 }
 
 impl DynamicResolverBuilder {
     pub fn new(
         dll_name: &str,
         function_name: &str,
         encryption_type: &str,
         var_name: &str,
         key_const_name: &str,
         encrypted_dll: &[u8],
         encrypted_fn: &[u8],
     ) -> Result<Self, String> {
         Ok(Self {
             dll_name: dll_name.to_string(),
             function_name: function_name.to_string(),
             var_name: var_name.to_string(),
             const_name: var_name.to_uppercase(),
             dll_array: format_bytes_as_hex(encrypted_dll),
             fn_array: format_bytes_as_hex(encrypted_fn),
             decrypt_fn: get_decrypt_function(encryption_type)?.to_string(),
             expect_clause: get_expect_clause(encryption_type).to_string(),
             key_const_name: key_const_name.to_string(),
         })
     }
 
     pub fn build(self) -> String {
         format!(
             r#"// Dynamic resolution of {function_name} from {dll_name}
         const ENCRYPTED_DLL_{const_name}: &[u8] = &{dll_array};
         const ENCRYPTED_FN_{const_name}: &[u8] = &{fn_array};
         
         let decrypted_dll = {decrypt_fn}(ENCRYPTED_DLL_{const_name}, &{key_const_name}){expect_clause};
         let decrypted_fn = {decrypt_fn}(ENCRYPTED_FN_{const_name}, &{key_const_name}){expect_clause};
         
         let dll_name = std::str::from_utf8(&decrypted_dll).expect("Invalid UTF-8 in DLL name");
         let fn_name = std::str::from_utf8(&decrypted_fn).expect("Invalid UTF-8 in function name");
         
         let module_base = internals::get_module_base(dll_name);
         if module_base.is_null() {{
             panic!("Failed to resolve module: {{}}", dll_name);
         }}
         
         let {var_name}_ptr = internals::get_function_address(module_base, fn_name);
         if {var_name}_ptr.is_null() {{
             panic!("Failed to resolve function: {{}}", fn_name);
         }}"#,
             function_name = self.function_name,
             dll_name = self.dll_name,
             const_name = self.const_name,
             dll_array = self.dll_array,
             fn_array = self.fn_array,
             decrypt_fn = self.decrypt_fn,
             expect_clause = self.expect_clause,
             key_const_name = self.key_const_name,
             var_name = self.var_name
         )
     }
 }
 
 pub fn generate_dynamic_function_resolver(
     dll_name: &str,
     function_name: &str,
     encryption_type: &str,
     var_name: &str,
     key_const_name: &str,
     encrypted_dll: &[u8],
     encrypted_fn: &[u8],
 ) -> Result<String, String> {
     let builder = DynamicResolverBuilder::new(
         dll_name,
         function_name,
         encryption_type,
         var_name,
         key_const_name,
         encrypted_dll,
         encrypted_fn,
     )?;
     Ok(builder.build())
 }
 
 /// Generate a type declaration and cast for a dynamically resolved function
 pub fn generate_function_type_and_cast(var_name: &str, function_signature: &str) -> String {
     let type_name = to_upper_camel_case(var_name);
 
     format!(
         r#"type {type_name} = {signature};
         let {var_name}: {type_name} = std::mem::transmute({var_name}_ptr);"#,
         type_name = type_name,
         signature = function_signature,
         var_name = var_name
     )
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     // Helper to create mock encrypted data for testing
     fn mock_encrypt_xor(data: &[u8]) -> Vec<u8> {
         // Simple XOR with a fixed key for testing
         let key = b"test_key";
         data.iter()
             .enumerate()
             .map(|(i, &b)| b ^ key[i % key.len()])
             .collect()
     }
 
     #[test]
     fn test_generate_dynamic_function_resolver_xor() {
         let dll_encrypted = mock_encrypt_xor(b"kernel32.dll");
         let fn_encrypted = mock_encrypt_xor(b"VirtualAlloc");
 
         let result = generate_dynamic_function_resolver(
             "kernel32.dll",
             "VirtualAlloc",
             "xor",
             "virtual_alloc",
             "DECRYPTION_KEY",
             &dll_encrypted,
             &fn_encrypted,
         );
 
         assert!(result.is_ok());
         let code = result.unwrap();
         // Encrypted mode - should have ENCRYPTED prefix and decryption
         assert!(code.contains("ENCRYPTED_DLL_VIRTUAL_ALLOC"));
         assert!(code.contains("ENCRYPTED_FN_VIRTUAL_ALLOC"));
         assert!(code.contains("decrypt_xor"));
         assert!(code.contains("internals::get_module_base"));
         assert!(code.contains("internals::get_function_address"));
         assert!(code.contains("virtual_alloc_ptr"));
     }
 
     #[test]
     fn test_generate_dynamic_function_resolver_aes() {
         let dll_encrypted = mock_encrypt_xor(b"ntdll.dll");
         let fn_encrypted = mock_encrypt_xor(b"RtlCopyMemory");
 
         let result = generate_dynamic_function_resolver(
             "ntdll.dll",
             "RtlCopyMemory",
             "aes",
             "rtl_copy_memory",
             "DECRYPTION_KEY",
             &dll_encrypted,
             &fn_encrypted,
         );
 
         assert!(result.is_ok());
         let code = result.unwrap();
         // Encrypted mode - should have ENCRYPTED prefix and decryption
         assert!(code.contains("ENCRYPTED_DLL_RTL_COPY_MEMORY"));
         assert!(code.contains("ENCRYPTED_FN_RTL_COPY_MEMORY"));
         assert!(code.contains("decrypt_aes"));
     }
 
     #[test]
     fn test_generate_function_type_and_cast() {
         let code = generate_function_type_and_cast(
             "virtual_alloc",
             "unsafe extern \"system\" fn(*const std::ffi::c_void, usize, u32, u32) -> *mut std::ffi::c_void",
         );
 
         assert!(code.contains("type VirtualAllocFn ="));
         assert!(code.contains("let virtual_alloc: VirtualAllocFn"));
         assert!(code.contains("std::mem::transmute(virtual_alloc_ptr)"));
     }
 
     #[test]
     fn test_generate_dynamic_function_resolver_invalid_encryption() {
         let dll_encrypted = mock_encrypt_xor(b"kernel32.dll");
         let fn_encrypted = mock_encrypt_xor(b"VirtualAlloc");
 
         let result = generate_dynamic_function_resolver(
             "kernel32.dll",
             "VirtualAlloc",
             "invalid_type",
             "virtual_alloc",
             "DECRYPTION_KEY",
             &dll_encrypted,
             &fn_encrypted,
         );
 
         assert!(result.is_err());
         assert!(result.unwrap_err().contains("Unknown encryption type"));
     }
 }

© 2024 - 2026 djnn, inc • generated with stagit