refactor: make .rdata offset grabber better

pe/superpe.py

@@ -16,6 +16,7 @@ class PeSection():
         self.raw_size: int = pefile_section.SizeOfRawData
         self.virt_addr: int = pefile_section.VirtualAddress
         self.virt_size: int = pefile_section.Misc_VirtualSize
+        self.pefile_section: pefile.SectionStructure = pefile_section
 
 
 class SuperPe():
@@ -44,13 +45,6 @@ class SuperPe():
         if self.arch == 'x64': 
             self.ptrSize = 8
     
-    ##################
-    def get_section_by_name(self, name: str) -> PeSection:
-        for section in self.pe_sections:
-            if section.name == name:
-                return section
-        return None
-    
 
     def get_physical_address(self, virtual_address):
         # Iterate through the section headers to find which section contains the VA
@@ -90,15 +84,10 @@ class SuperPe():
         return bytes(sect.get_data())
 
 
-    def get_section_data(self, sect_name) -> bytes:
-        sect = self.get_section_by_name_b(sect_name)
-        return bytes(sect.get_data())
-
-
-    def get_section_by_name_b(self, name):
-        for sect in self.pe.sections:
-            if sect.Name.decode().lower().startswith(name.lower()):
-                return sect
+    def get_section_by_name(self, name: str) -> PeSection:
+        for section in self.pe_sections:
+            if section.name == name:
+                return section
         return None
     
 

phases/injector.py

@@ -106,27 +106,23 @@ def injected_fix_data(superpe: SuperPe, carrier: Carrier, exe_host: ExeHost):
         # nothing todo
         return
     
-    # Offset of strings in .rodata
-    sect = exe_host.superpe.get_section_by_name_b(".rdata")
-    sect_data = sect.get_data()
-    string_off = find_first_utf16_string_offset(sect_data)
+    # Put stuff into .rdata section in the PE
+    peSection = exe_host.superpe.get_section_by_name(".rdata")
+    if peSection == None:
+        raise Exception("No .rdata section found, abort")
+    sect_data_copy = peSection.pefile_section.get_data()
+    string_off = find_first_utf16_string_offset(sect_data_copy)
     if string_off == None:
         raise Exception("Strings not found in .rdata section, abort")
     if string_off < 100:
         logging.warn("weird: Strings in .rdata section at offset {} < 100".format(string_off))
-    
-    sect = exe_host.superpe.get_section_by_name(".rdata")
-    addr = sect.raw_addr + string_off
-
+    fixup_offset_rdata = peSection.raw_addr + string_off
+    # Do all .rdata patches
     for datareuse_fixup in reusedata_fixups:
         var_data = datareuse_fixup.data
-        #print("    Addr: {} / 0x{:X}  Data: {}".format(
-        #    addr, addr, len(var_data)))
-        superpe.pe.set_bytes_at_offset(addr, var_data)
-        #f.seek(addr)
-        #f.write(var_data)
-        datareuse_fixup.addr = addr + sect.virt_addr + exe_host.image_base - sect.raw_addr
-        addr += len(var_data) + 8
+        superpe.pe.set_bytes_at_offset(fixup_offset_rdata, var_data)
+        datareuse_fixup.addr = fixup_offset_rdata + peSection.virt_addr + exe_host.image_base - peSection.raw_addr
+        fixup_offset_rdata += len(var_data) + 8
 
     # patch code section
     # replace the placeholder with a LEA instruction to the data we written above