refactor: ExeCapabilities -> ExeInfo

phases/assembler.py

@@ -4,6 +4,7 @@ from model import *
 from config import config
 from observer import observer
 from pehelper import *
+from helper import *
 
 logger = logging.getLogger("Assembler")
 

phases/compiler.py

@@ -16,7 +16,7 @@ def compile(
     c_in: FilePath, 
     asm_out: FilePath,
     payload_len: int, 
-    exe_capabilities: ExeCapabilities
+    exe_info: ExeInfo
 ):
     logger.info("--[ Compile C to ASM: {} -> {} ".format(c_in, asm_out))
 
@@ -36,7 +36,7 @@ def compile(
 
     # Phase 1.2: Assembly fixup
     logger.info("---[ Fixup  : {} ".format(asm_out))
-    if not fixup_asm_file(asm_out, payload_len, exe_capabilities):
+    if not fixup_asm_file(asm_out, payload_len, exe_info):
         raise Exception("Error: Fixup failed")
     observer.add_text("payload_asm_fixup", file_readall_text(asm_out))
 
@@ -63,7 +63,7 @@ def bytes_to_asm_db(byte_data: bytes) -> bytes:
     return "\tDB " + formatted_string
 
 
-def fixup_asm_file(filename: FilePath, payload_len: int, capabilities: ExeCapabilities):
+def fixup_asm_file(filename: FilePath, payload_len: int, capabilities: ExeInfo):
     with open(filename, 'r', encoding='utf-8') as asmfile:
         lines = asmfile.readlines()
 

phases/injector.py

@@ -14,7 +14,7 @@ def inject_exe(
     shellcode_in: FilePath,
     exe_in: FilePath,
     exe_out: FilePath,
-    exe_capabilities: ExeCapabilities,
+    exe_info: ExeInfo,
 ):
     logger.info("--[ Injecting: {} into: {} -> {} ".format(
         shellcode_in, exe_in, exe_out
@@ -38,12 +38,12 @@ def inject_exe(
     if project.source_style == SourceStyle.iat_reuse:
         # get code section of exe_out
         code = extract_code_from_exe(exe_out)
-        for cap in exe_capabilities.get_all().values():
+        for cap in exe_info.get_all().values():
             if not cap.id in code:
                 raise Exception("Capability ID {} not found, abort".format(cap.id))
             
             off = code.index(cap.id)
-            current_address = off + exe_capabilities.image_base + exe_capabilities.text_virtaddr
+            current_address = off + exe_info.image_base + exe_info.code_virtaddr
             destination_address = cap.addr
             logger.info("    Replace at 0x{:x} with call to 0x{:x}".format(
                 current_address, destination_address