refactor: datareuse fixups into project.carrier as transport

model/carrier.py

@@ -1,12 +1,12 @@
 from typing import Dict, List
 import logging
 
-from model.exehost import ExeHost
+from model.exehost import DataReuseEntry
 
-logger = logging.getLogger("ExeHost")
+logger = logging.getLogger("Carrier")
 
 
-class IatEntry():
+class IatRequest():
     def __init__(self, name: str, placeholder: bytes):
         self.name: str = name           # Function Name, like "VirtualAlloc"
         self.placeholder: bytes = placeholder    # Random bytes as placeholder
@@ -15,7 +15,8 @@ class IatEntry():
 
 class Carrier():
     def __init__(self):
-        self.iat_requests: List[IatEntry] = []
+        self.iat_requests: List[IatRequest] = []
+        self.reusedata_fixups: List[DataReuseEntry] = []
 
 
     def init(self):
@@ -23,8 +24,14 @@ class Carrier():
 
 
     def add_iat_request(self, func_name: str, placeholder: bytes):
-        self.iat_requests.append(IatEntry(func_name, placeholder))
+        self.iat_requests.append(IatRequest(func_name, placeholder))
 
-
+    def get_all_iat_requests(self) -> List[IatRequest]:
-    def get_all_iat_requests(self) -> List[IatEntry]:
         return self.iat_requests
+
+
+    def set_datareuse_fixups(self, fixups: List[DataReuseEntry]):
+        self.reusedata_fixups = fixups
+
+    def get_all_reusedata_fixups(self) -> List[DataReuseEntry]:
+        return self.reusedata_fixups

model/exehost.py

@@ -10,6 +10,28 @@ from peparser.misc import get_physical_address
 logger = logging.getLogger("ExeHost")
 
 
+class RelocEntry():
+    def __init__(self, rva: int, base_rva: int, type: str):
+        self.rva: int = rva
+        self.base_rva: int = base_rva
+        self.type: str = type
+
+
+class IatEntry():
+    def __init__(self, dll_name, func_name, func_addr):
+        self.dll_name = dll_name
+        self.func_name = func_name
+        self.func_addr = func_addr
+
+
+class DataReuseEntry():
+    def __init__(self, string_ref: str, register: str, randbytes: bytes):
+        self.string_ref = string_ref
+        self.register = register
+        self.randbytes = randbytes
+        self.data = b''
+        self.addr = 0
+
 
 class ExeHost():
     def __init__(self, filepath: FilePath):
@@ -21,7 +43,7 @@ class ExeHost():
         self.superpe: SuperPe = None
 
         self.iat = {}  # Dict[str, List[Dict[str, str]]]
-        
+        self.base_relocs = []
 
         self.image_base: int = 0
         self.dynamic_base: bool = False
@@ -30,7 +52,6 @@ class ExeHost():
         self.code_size: int = 0
         self.code_section = None
         
-        self.base_relocs = []
         self.rwx_section = None
 
         self.ep = None

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()
+    asmFileParser.process()
-    data_fixup_entries = asmFileParser.get_data_reuse_entries()
+    carrier.set_datareuse_fixups(asmFileParser.get_reusedata_fixups())
-    config.data_fixups = data_fixups
-    config.data_fixup_entries = data_fixup_entries
     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({
+                self.fixups[string_ref] = DataReuseEntry(string_ref, register, randbytes)
-                    "string_ref": string_ref,
-                    "register": register,
-                    "randbytes": 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]:
+    def fixup_data_reuse_data(self) -> List[str]:
-        entries = {}
         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():
+    for IatRequest in carrier.get_all_iat_requests():
-        if not iatEntry.placeholder in code:
+        if not IatRequest.placeholder in code:
-            raise Exception("IatResolve ID {} not found, abort".format(iatEntry.placeholder))
+            raise Exception("IatResolve ID {} not found, abort".format(IatRequest.placeholder))
-        addr = exe_host.get_addr_of_iat_function(iatEntry.name)
+        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:
+        for datareuse_fixup in reusedata_fixups:
-            var_data = data_fixup_entries[fixup["string_ref"]]
+            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:
+    for datareuse_fixup in reusedata_fixups:
-        if not fixup["randbytes"] in code:
+        if not datareuse_fixup.randbytes in code:
-            raise Exception("DataResuse: ID {} not found, abort".format(fixup["randbytes"]))
+            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)

supermega.py

@@ -167,6 +167,7 @@ def start(settings: Settings):
             c_in = main_c_file, 
             asm_out = main_asm_file, 
             payload_len = project.payload.len,
+            carrier = project.carrier,
             short_call_patching = project.settings.short_call_patching)
 
     # Decide if we can use IAT_REUSE (all function calls available as import)
@@ -262,8 +263,7 @@ def start(settings: Settings):
             # TODO IF?
             phases.injector.injected_fix_data(
                 settings.inject_exe_out, 
-                config.data_fixups,
+                project.carrier,
-                config.data_fixup_entries,
                 project.exe_host)
             
             code = extract_code_from_exe(settings.inject_exe_out)