refactor: cleanup 1

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, 
-            inline_strings=False,  # not for DATA_REUSE
-            remove_crt=True, 
-            append_rsp_stub=True)  # required atm
-        process_file(params)
-    else:
-        run_process_checkret([
-            config.get("path_masmshc"),
-            asm_out,
-            asm_clean_file,
-        ])
+    params = Params(asm_out, asm_clean_file, 
+        inline_strings=False,  # not for DATA_REUSE
+        remove_crt=True, 
+        append_rsp_stub=True)  # required atm
+    process_file(params)
+
     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)