refactor: datareuse fixups into project.carrier as transport

phases/compiler.py

@@ -19,6 +19,7 @@ def compile(
     c_in: FilePath, 
     asm_out: FilePath,
     payload_len: int,
+    carrier: Carrier,
     short_call_patching: bool = False
 ):
     logger.info("--[ Compile C to ASM: {} -> {} ".format(c_in, asm_out))
@@ -38,12 +39,10 @@ def compile(
     observer.add_text("carrier_asm_orig", file_readall_text(asm_out))
 
     # DataReuse first
-    asmFileParser = AsmFileParser(asm_out)
+    asmFileParser = ReusedataAsmFileParser(asm_out)
     asmFileParser.init()
-    data_fixups = asmFileParser.fixup_data_reuse()
-    data_fixup_entries = asmFileParser.get_data_reuse_entries()
-    config.data_fixups = data_fixups
-    config.data_fixup_entries = data_fixup_entries
+    asmFileParser.process()
+    carrier.set_datareuse_fixups(asmFileParser.get_reusedata_fixups())
     asmFileParser.write_lines_to(asm_out)
 
     # Assembly text fixup (SuperMega)

phases/datareuse.py

@@ -1,9 +1,10 @@
 import sys
 import pefile
 from intervaltree import Interval, IntervalTree
-from typing import List
+from typing import List, Dict
 import os
 
+from model.exehost import DataReuseEntry
 
 class PeSection():
     def __init__(self, pefile_section):
@@ -31,10 +32,15 @@ def bytes_to_asm_db(byte_data: bytes) -> bytes:
     return "\tDB " + formatted_string
 
 
-class AsmFileParser():
+class ReusedataAsmFileParser():
     def __init__(self, filepath):
         self.filepath = filepath
         self.lines = []
+        self.fixups: Dict[str, DataReuseEntry] = {}
+
+
+    def get_reusedata_fixups(self) -> List[DataReuseEntry]:
+        return list(self.fixups.values())
 
 
     def init(self):
@@ -43,7 +49,12 @@ class AsmFileParser():
         self.lines = [line.rstrip() for line in self.lines]
 
 
-    def fixup_data_reuse(self):
+    def process(self):
+        self.fixup_data_reuse_code()
+        self.fixup_data_reuse_data()
+
+
+    def fixup_data_reuse_code(self):
         fixups = []
         # lea	rcx, OFFSET FLAT:$SG72513
         for idx, line in enumerate(self.lines):
@@ -51,18 +62,13 @@ class AsmFileParser():
                 string_ref = line.split("OFFSET FLAT:")[1]
                 register = line.split("lea\t")[1].split(",")[0]
                 randbytes: bytes = os.urandom(7) # lea is 7 bytes
-                fixups.append({
-                    "string_ref": string_ref,
-                    "register": register,
-                    "randbytes": randbytes,
-                })
+                self.fixups[string_ref] = DataReuseEntry(string_ref, register, randbytes)
                 self.lines[idx] = bytes_to_asm_db(randbytes) + " ; .rdata Reuse for {} ({})".format(
                     string_ref, register)
         return fixups
 
 
-    def get_data_reuse_entries(self) -> List[str]:
-        entries = {}
+    def fixup_data_reuse_data(self) -> List[str]:
         current_entry_name = ""
 
         for line in self.lines:
@@ -79,7 +85,11 @@ class AsmFileParser():
                     elif part.endswith('H') or part.endswith('H,'):
                         hex = part.split('H')[0]
                         value += bytes.fromhex(hex)
-                entries[name] = value
+
+                if not name in self.fixups:
+                    raise Exception("DataReuse: Entry {} not found in fixups".format(name))
+                self.fixups[name].data = value
+                
 
             elif line.startswith("\tDB"):
                 if current_entry_name == "":
@@ -97,12 +107,13 @@ class AsmFileParser():
                         #print("---> {}".format(hex))
                         value += bytes.fromhex(hex)
 
-                entries[current_entry_name] += value
+                if not name in self.fixups:
+                    raise Exception("DataReuse: Entry {} not found in fixups".format(name))
+                self.fixups[name].data += value
+
             else:
                 current_entry_name = ""
                 
-        return entries
-
 
     def write_lines_to(self, filename):
         with open(filename, 'w',) as asmfile:

phases/injector.py

@@ -64,14 +64,14 @@ def injected_fix_iat(exe_out: FilePath, carrier: Carrier, exe_host: ExeHost):
     # get code section of exe_out
     code = extract_code_from_exe(exe_out)
 
-    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)
+    for IatRequest in carrier.get_all_iat_requests():
+        if not IatRequest.placeholder in code:
+            raise Exception("IatResolve ID {} not found, abort".format(IatRequest.placeholder))
+        addr = exe_host.get_addr_of_iat_function(IatRequest.name)
         if addr == None:
-            raise Exception("IatResolve: Function {} not found".format(iatEntry.name))
+            raise Exception("IatResolve: Function {} not found".format(IatRequest.name))
         
-        off = code.index(iatEntry.placeholder)
+        off = code.index(IatRequest.placeholder)
         current_address = off + exe_host.image_base + exe_host.code_virtaddr
         #current_address += 2
         destination_address = addr
@@ -81,13 +81,13 @@ def injected_fix_iat(exe_out: FilePath, carrier: Carrier, exe_host: ExeHost):
         jmp = assemble_and_disassemble_jump(
             current_address, destination_address
         )
-        code = code.replace(iatEntry.placeholder, jmp)
+        code = code.replace(IatRequest.placeholder, jmp)
 
     # write back our patched code into the exe
     write_code_section(exe_file=exe_out, new_data=code)
 
 
-def injected_fix_data(exe_path, data_fixups, data_fixup_entries, exe_host):
+def injected_fix_data(exe_path, carrier: Carrier, exe_host: ExeHost):
     data_reuser = DataReuser(exe_path)
     data_reuser.init()
     #ret = data_reuser.get_reloc_largest_gap(".rdata")
@@ -105,9 +105,11 @@ def injected_fix_data(exe_path, data_fixups, data_fixup_entries, exe_host):
     print("Write into .data:".format())
     data_reuser.pe.close()
 
+    reusedata_fixups: List[DataReuseEntry] = carrier.get_all_reusedata_fixups()
+
     with open(exe_path, "r+b") as f:
-        for fixup in data_fixups:
-            var_data = data_fixup_entries[fixup["string_ref"]]
+        for datareuse_fixup in reusedata_fixups:
+            var_data = datareuse_fixup.data
 
             print("    Addr: {} / 0x{:X}  Data: {}".format(
                 addr, addr, len(var_data)))
@@ -118,7 +120,7 @@ def injected_fix_data(exe_path, data_fixups, data_fixup_entries, exe_host):
             f.write(var_data)
             #f.write(b"AAAAAAAAAAAAAAAAAAAAAAAAAAA")
             print("ADD: 0x{:X} 0x{:X} 0x{:X}".format(addr, sect.virt_addr, exe_host.image_base))
-            fixup["addr"] = addr + sect.virt_addr + exe_host.image_base - sect.raw_addr
+            datareuse_fixup.addr = addr + sect.virt_addr + exe_host.image_base - sect.raw_addr
             addr += len(var_data) + 8
     #data_reuser.pe.write(exe_path + ".tmp")
     #data_reuser.pe.close()
@@ -126,20 +128,21 @@ def injected_fix_data(exe_path, data_fixups, data_fixup_entries, exe_host):
 
     # patch code section
     code = extract_code_from_exe(exe_path)
-    for fixup in data_fixups:
-        if not fixup["randbytes"] in code:
-            raise Exception("DataResuse: ID {} not found, abort".format(fixup["randbytes"]))
+    for datareuse_fixup in reusedata_fixups:
+        if not datareuse_fixup.randbytes in code:
+            raise Exception("DataResuse: ID {} not found, abort".format(
+                datareuse_fixup.randbytes))
         
-        off = code.index(fixup["randbytes"])
+        off = code.index(datareuse_fixup.randbytes)
         current_address = off + exe_host.image_base + exe_host.code_virtaddr
-        destination_address = fixup["addr"]
+        destination_address = datareuse_fixup.addr
         logger.info("    Replace at 0x{:x} with call to 0x{:x}".format(
             current_address, destination_address
         ))
         lea = assemble_lea(
-            current_address, destination_address, fixup["register"]
+            current_address, destination_address, datareuse_fixup.register
         )
-        code = code.replace(fixup["randbytes"], lea)
+        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)