refactor: iat related

phases/compiler.py

@@ -9,6 +9,7 @@ from observer import observer
 from model import *
 from phases.masmshc import process_file, Params
 from phases.datareuse import *
+from model.carrier import Carrier
 
 logger = logging.getLogger("Compiler")
 use_templates = True
@@ -166,7 +167,7 @@ def get_function_stubs(asm_in: FilePath) -> List[str]:
     return functions
 
 
-def fixup_iat_reuse(filename: FilePath, exe_host):
+def fixup_iat_reuse(filename: FilePath, carrier: Carrier):
     with open(filename, 'r', encoding='utf-8') as asmfile:
         lines = asmfile.readlines()
 
@@ -180,7 +181,7 @@ def fixup_iat_reuse(filename: FilePath, exe_host):
             randbytes: bytes = os.urandom(6)
             lines[idx] = bytes_to_asm_db(randbytes) + " ; IAT Reuse for {}".format(func_name)
             lines[idx] += "\n"
-            exe_host.add_iat_resolve(func_name, randbytes)
+            carrier.add_iat_request(func_name, randbytes)
 
             logger.info("    > Replace func name: {} with {}".format(
                 func_name, randbytes.hex()))

phases/injector.py

@@ -6,6 +6,7 @@ import time
 import tempfile
 import logging
 
+from model.carrier import Carrier
 from peparser.pehelper import *
 from model.exehost import *
 from observer import observer
@@ -57,26 +58,30 @@ def inject_exe(
         raise Exception("Shellcode injection error")
 
 
-def injected_fix_iat(exe_out: FilePath, exe_host: ExeHost):
+def injected_fix_iat(exe_out: FilePath, carrier: Carrier, exe_host: ExeHost):
     """replace IAT in shellcode in code and re-implant it"""
 
     # get code section of exe_out
     code = extract_code_from_exe(exe_out)
-    for cap in exe_host.get_all_iat_resolvs().values():
-        if not cap.id in code:
-            raise Exception("IatResolve ID {} not found, abort".format(cap.id))
+
+    for iatEntry in carrier.get_all_iat_requests():
+        if not iatEntry.placeholder in code:
+            raise Exception("IatResolve ID {} not found, abort".format(iatEntry.placeholder))
+        addr = exe_host.get_addr_of_iat_function(iatEntry.name)
+        if addr == None:
+            raise Exception("IatResolve: Function {} not found".format(iatEntry.name))
         
-        off = code.index(cap.id)
+        off = code.index(iatEntry.placeholder)
         current_address = off + exe_host.image_base + exe_host.code_virtaddr
         #current_address += 2
-        destination_address = cap.addr
+        destination_address = addr
         logger.info("    Replace at 0x{:x} with call to 0x{:x}".format(
             current_address, destination_address
         ))
         jmp = assemble_and_disassemble_jump(
             current_address, destination_address
         )
-        code = code.replace(cap.id, jmp)
+        code = code.replace(iatEntry.placeholder, jmp)
 
     # write back our patched code into the exe
     write_code_section(exe_file=exe_out, new_data=code)