refactor: remove exehost code_virtaddr and code_size with code_section

model/exehost.py

@@ -41,11 +41,7 @@ class ExeHost():
 
         self.image_base: int = 0
         self.dynamic_base: bool = False
-
-        self.code_virtaddr: int = 0
-        self.code_size: int = 0
         self.code_section = None
-        
         self.rwx_section = None
 
         self.ep = None
@@ -71,14 +67,12 @@ class ExeHost():
         else:
             self.dynamic_base = False
 
-        # .text virtual address
-        self.code_section = pehelper.get_code_section(self.superpe.pe)
-        self.code_virtaddr = self.code_section.VirtualAddress
-        self.code_size = self.code_section.Misc_VirtualSize
+        # Info output: .text virtual address
+        self.code_section = self.superpe.get_code_section()
         logger.info("---[ Injectable: Chosen code section: {} at 0x{:X} size: {}".format(
             self.code_section.Name.decode().rstrip('\x00'),
-            self.code_virtaddr,
-            self.code_size))
+            self.code_section.VirtualAddress,
+            self.code_section.Misc_VirtualSize))
 
         # relocs
         if hasattr(self.superpe.pe, 'DIRECTORY_ENTRY_BASERELOC'):

phases/injector.py

@@ -34,9 +34,9 @@ def inject_exe(
     # And check if it fits into the target code section
     main_shc = file_readall_binary(main_shc_file)
     l = len(main_shc)
-    if l + 128 > project.exe_host.code_size:
+    if l + 128 > project.exe_host.code_section.Misc_VirtualSize:
         logger.error("Error: Shellcode {}+128 too small for target code section {}".format(
-            l, project.exe_host.code_size
+            l, project.exe_host.code_section.Misc_VirtualSize
         ))
         return False
 
@@ -85,7 +85,7 @@ def injected_fix_iat(superpe: SuperPe, carrier: Carrier, exe_host: ExeHost):
             raise Exception("IatResolve: Function {} not found".format(iatRequest.name))
         
         offset_from_code = code.index(iatRequest.placeholder)
-        instruction_virtual_address = offset_from_code + exe_host.image_base + exe_host.code_virtaddr
+        instruction_virtual_address = offset_from_code + exe_host.image_base + exe_host.code_section.VirtualAddress
         logger.info("    Replace {} at VA 0x{:X} with call to IAT at VA 0x{:X}".format(
             iatRequest.placeholder.hex(), instruction_virtual_address, destination_virtual_address
         ))
@@ -145,7 +145,7 @@ def injected_fix_data(superpe: SuperPe, carrier: Carrier, exe_host: ExeHost):
                 datareuse_fixup.randbytes))
         
         offset_from_datasection = code.index(datareuse_fixup.randbytes)
-        instruction_virtual_address = offset_from_datasection + exe_host.image_base + exe_host.code_virtaddr
+        instruction_virtual_address = offset_from_datasection + exe_host.image_base + exe_host.code_section.VirtualAddress
         destination_virtual_address = datareuse_fixup.addr
         logger.info("    Replace {} at VA 0x{:X} with .rdata LEA at VA 0x{:X}".format(
             datareuse_fixup.randbytes.hex(), instruction_virtual_address, destination_virtual_address