refactor: cleanup, create unittest, fix get_physical_address

model/defs.py

@@ -51,8 +51,17 @@ class PeRelocEntry():
         self.type: str = type
 
 
+    def __str__(self):
+        return "PeRelocEntry: rva: 0x{:X} base_rva: 0x{:X} offset: 0x{:X} type: {}".format(
+            self.rva, self.base_rva, self.offset, self.type)
+
+
 class IatEntry():
     def __init__(self, dll_name: str, func_name: str, iat_vaddr: int):
         self.dll_name: str = dll_name
         self.func_name: str = func_name
         self.iat_vaddr: int = iat_vaddr
+
+    def __str__(self):
+        return "IatEntry: dll_name: {} func_name: {} iat_vaddr: 0x{:X}".format(
+            self.dll_name, self.func_name, self.iat_vaddr)

model/exehost.py

@@ -29,9 +29,6 @@ class ExeHost():
         self.code_section = None
         self.rwx_section = None
 
-        self.ep = None
-        self.ep_raw = None
-
 
     def init(self):
         logger.info("--[ Analyzing: {}".format(self.filepath))
@@ -40,9 +37,6 @@ class ExeHost():
         if not self.superpe.is_64():
             raise Exception("Binary is not 64bit: {}".format(self.filepath))
 
-        self.ep = self.superpe.get_entrypoint()
-        self.ep_raw = self.superpe.get_physical_address(self.ep)
-
         # image base
         self.image_base = self.superpe.pe.OPTIONAL_HEADER.ImageBase
 

pe/pehelper.py

@@ -20,28 +20,19 @@ def extract_code_from_exe_file_ep(exe_file: FilePath, len: int) -> bytes:
     section = get_code_section(pe)
     data: bytes = section.get_data()
     data = remove_trailing_null_bytes(data)
-
     ep = pe.OPTIONAL_HEADER.AddressOfEntryPoint
-    ep_raw = get_physical_address(pe, ep)
-
+    ep_raw = get_physical_address_tmp(pe, ep)
     data = data[ep_raw:ep_raw+len]
-
     pe.close()
-
     return data
 
 
-def get_physical_address(pe, virtual_address):
-    # Iterate through the section headers to find which section contains the VA
+def get_physical_address_tmp(pe, virtual_address):
     for section in pe.sections:
-        # Check if the VA is within the range of this section
         if section.VirtualAddress <= virtual_address < section.VirtualAddress + section.Misc_VirtualSize:
-            # Calculate the difference between the VA and the section's virtual address
             virtual_offset = virtual_address - section.VirtualAddress
-            # Add the difference to the section's pointer to raw data
-            return virtual_offset
-            #physical_address = section.PointerToRawData + virtual_offset
-            #return physical_address
+            physical_address = section.PointerToRawData + virtual_offset
+            return physical_address
     return None
 
 

pe/superpe.py

@@ -73,7 +73,8 @@ class SuperPe():
 
     ## Section Access
 
-    def get_code_section(self):
+    def get_code_section(self) -> pefile.SectionStructure:
+        """Return the section that contains the entrypoint and is executable"""
         entrypoint = self.pe.OPTIONAL_HEADER.AddressOfEntryPoint
         for sect in self.pe.sections:
             if sect.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_EXECUTE']:
@@ -87,7 +88,7 @@ class SuperPe():
         return bytes(sect.get_data())
     
 
-    def get_rwx_section(self):
+    def get_rwx_section(self) -> pefile.SectionStructure:
         # rwx section
         entrypoint = self.pe.OPTIONAL_HEADER.AddressOfEntryPoint
         for section in self.pe.sections:
@@ -223,7 +224,7 @@ class SuperPe():
             res.append(e.name.decode())
         return res
     
-    
+
     def get_exports_full(self):
         """Return a list of exported functions (names) from the PE file"""
         d = [pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_EXPORT"]]
@@ -262,28 +263,8 @@ class SuperPe():
                 # Calculate the difference between the VA and the section's virtual address
                 virtual_offset = virtual_address - section.VirtualAddress
                 # Add the difference to the section's pointer to raw data
-
-                #print("0x{:X}  0x{:X}  -> 0x{:X}".format(virtual_offset, section.PointerToRawData, virtual_offset + section.PointerToRawData ))
-                return virtual_offset
-                #physical_address = section.PointerToRawData + virtual_offset
-                #return physical_address
-        return None
-    
-
-    def get_physical_address2(self, virtual_address) -> int:
-        """Convert a virtual address to a physical address in the PE file"""
-        # Iterate through the section headers to find which section contains the VA
-        for section in self.pe.sections:
-            # Check if the VA is within the range of this section
-            if section.VirtualAddress <= virtual_address < section.VirtualAddress + section.Misc_VirtualSize:
-                # Calculate the difference between the VA and the section's virtual address
-                virtual_offset = virtual_address - section.VirtualAddress
-                # Add the difference to the section's pointer to raw data
-
-                logger.info("0x{:X}  0x{:X}  -> 0x{:X}".format(virtual_offset, section.PointerToRawData, virtual_offset + section.PointerToRawData ))
-                return virtual_offset + section.PointerToRawData
-                #physical_address = section.PointerToRawData + virtual_offset
-                #return physical_address
+                physical_address = section.PointerToRawData + virtual_offset
+                return physical_address
         return None
     
 

phases/injector.py

@@ -49,7 +49,7 @@ def inject_exe(
         # Special case. put it at the beginning of the exported DLL function
         logger.info("--[ Overwrite DLL function {} with shellcode".format(settings.dllfunc))
         rva = pe_backdoorer.getExportEntryPoint(settings.dllfunc)
-        shellcode_offset = superpe.get_physical_address2(rva)
+        shellcode_offset = superpe.get_physical_address(rva)
         logger.info(f'---[ Using DLL Export "{settings.dllfunc}" at RVA 0x{rva:X} offset 0x{shellcode_offset:X} to overwrite')
         superpe.pe.set_bytes_at_offset(shellcode_offset, main_shc)
 

tester.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, List
 
 from helper import *
 from config import config
@@ -14,7 +14,7 @@ def main():
     logger.info("Super Mega Tester")
     config.load()
 
-    #test_exe()
+    test_exe()
     test_dll()
 
 

tests/test_superpe.py

@@ -0,0 +1,68 @@
+from typing import List, Dict
+import unittest
+import pefile
+
+from pe.superpe import SuperPe, PeSection
+from model.defs import *
+
+
+class SuperPeTest(unittest.TestCase):
+
+    def test_exe(self):
+        pass
+
+
+    def test_dll(self):
+        dll_filepath = PATH_EXES + "libbz2-1.dll"
+        superpe = SuperPe(dll_filepath)
+
+        # Properties
+        self.assertTrue(superpe.is_dll())
+        self.assertTrue(superpe.is_64())
+        self.assertFalse(superpe.is_dotnet())
+        self.assertEqual(superpe.get_entrypoint(), 0x1350)
+        self.assertIsNone(superpe.get_rwx_section())
+
+        # Text Section 1 (pefile SectionStructure)
+        code_sect: pefile.SectionStructure = superpe.get_code_section()
+        self.assertEqual(code_sect.Name.decode(), ".text\x00\x00\x00")
+        self.assertEqual(code_sect.VirtualAddress, 0x1000)
+        self.assertEqual(code_sect.Misc_VirtualSize, 0x12D08)
+
+        # Text Section 2 (PeSection)
+        code_pesect: PeSection = superpe.get_section_by_name(".text")
+        self.assertEqual(code_pesect.name, ".text")
+        self.assertEqual(code_pesect.virt_addr, 0x1000)
+        self.assertEqual(code_pesect.virt_size, 0x12D08)
+
+        # Relocations
+        base_relocs: List[PeRelocEntry] = superpe.get_base_relocs()
+        self.assertEqual(len(base_relocs), 54)
+        base_reloc = base_relocs[0]
+        self.assertEqual(base_reloc.rva, 0x13CE8)
+        self.assertEqual(base_reloc.base_rva, 0x13000)
+        self.assertEqual(base_reloc.offset, 0xCE8)
+
+        # IAT
+        iat_entries: Dict[str, IatEntry] = superpe.get_iat_entries()
+        self.assertEqual(len(iat_entries), 2)
+        self.assertTrue("KERNEL32.dll" in iat_entries)
+        self.assertTrue("msvcrt.dll" in iat_entries)
+        kernel32_entries = iat_entries["KERNEL32.dll"]
+        self.assertEqual(len(kernel32_entries), 12)
+        entry = kernel32_entries[0]
+        self.assertEqual(entry.dll_name, "KERNEL32.dll")
+        self.assertEqual(entry.func_name, "DeleteCriticalSection")
+        self.assertEqual(entry.iat_vaddr, 0x1f13db1c4)
+
+        # Exports
+        exports = superpe.get_exports_full()
+        self.assertEqual(len(exports), 35)
+        export = exports[0]
+        self.assertEqual(export["name"], "BZ2_blockSort")
+        self.assertEqual(export["addr"], 0x2FC0)
+        self.assertEqual(export["size"], 416)
+
+        # VRA/Virt to Phys/Raw
+        raw = superpe.get_physical_address(0xD690)  # BZ2_bzdopen export
+        self.assertEqual(raw, 0xCA90)