diff --git a/model/exehost.py b/model/exehost.py index 911aa9a..6ea183c 100644 --- a/model/exehost.py +++ b/model/exehost.py @@ -7,6 +7,7 @@ from model.defs import * import pe.pehelper as pehelper from pe.superpe import SuperPe, PeSection from model.carrier import Carrier +from model.rangemanager import RangeManager logger = logging.getLogger("ExeHost") @@ -36,6 +37,7 @@ class ExeHost(): self.iat: Dict[str, IatEntry] = {} self.base_relocs: List[PeRelocEntry] = [] + self.base_reloc_ranges: RangeManager = None self.image_base: int = 0 self.dynamic_base: bool = False @@ -127,43 +129,29 @@ class ExeHost(): section: PeSection = self.superpe.get_section_by_name(section_name) if section is None: return [] - return [reloc for reloc in self.base_relocs if reloc.base_rva == section.virt_addr] + ret = [] + + #return [reloc for reloc in self.base_relocs if reloc.base_rva == section.virt_addr] + for reloc in self.base_relocs: + reloc_addr = reloc.rva + if reloc_addr >= section.virt_addr and reloc_addr < section.virt_addr + section.virt_size: + ret.append(reloc) + return ret - def get_reloc_largest_gap(self, section_name: str = ".rdata"): - tree = IntervalTree() - section: PeSection = self.superpe.get_section_by_name(section_name) + def get_rdata_relocmanager(self) -> RangeManager: + section = self.superpe.get_section_by_name(".rdata") + relocs = self.get_relocations_for_section(".rdata") + #print("Relocs for .rdata: {} of {}".format(len(relocs), len(self.base_relocs))) - print("-- Relocations: {}".format(len(self.base_relocs))) - print("-- section: 0x{:x}".format(section.virt_addr)) + rm = RangeManager(section.virt_addr, section.virt_addr + section.virt_size) + for reloc in relocs: + # Reloc destination is probably 8 bytes + # But i add another 8 to skip over small holes (common in .rdata) + rm.add_range(reloc.rva, reloc.rva + 8 + 8) + rm.merge_overlaps() + return rm - for reloc in self.base_relocs: - if reloc.base_rva == section.virt_addr: - print("--- Adding reloc: {} {}".format(reloc.offset, reloc.offset + 8)) - tree.add(Interval(reloc.offset, reloc.offset + 8)) - tree.add(Interval(section.virt_size, section.virt_size + 1)) - - # Initialize variables to track the largest gap and its bounds - max_gap = 0 - gap_start = None - gap_end = None - - # Sort intervals for sequential comparison - sorted_intervals = sorted(tree) - for i in range(len(sorted_intervals) - 1): - current_end = sorted_intervals[i].end - next_start = sorted_intervals[i + 1].begin - gap = next_start - current_end - if gap > max_gap: - max_gap = gap - gap_start = current_end # Adjusted for the actual start of the gap - gap_end = next_start - 1 # Adjusted for the actual end of the gap - - # Adjust for the artificial +1 in interval ends - if gap_start is not None and gap_end is not None: - gap_start -= 1 - - return max_gap - 1, gap_start, gap_end def has_all_carrier_functions(self, carrier: Carrier): is_ok = True diff --git a/model/rangemanager.py b/model/rangemanager.py new file mode 100644 index 0000000..7cfc31d --- /dev/null +++ b/model/rangemanager.py @@ -0,0 +1,43 @@ +from intervaltree import Interval, IntervalTree + +class RangeManager: + def __init__(self, min=0, max=1000): + self.intervals = IntervalTree() + self.min = min + self.max = max + + + def merge_overlaps(self): + self.intervals.merge_overlaps(strict=False) + + + def add_range(self, start, end): + if start < self.min or end > self.max: + raise ValueError("Ranges must be within 0x{:X} and 0x{:X}, not 0x{:X}/0x{:X}".format( + self.min, self.max, start, end + )) + self.intervals.add(Interval(start, end)) + + + def find_hole(self, hole_size): + sorted_intervals = sorted(self.intervals) + last_end = self.min + for interval in sorted_intervals: + start, end = interval.begin, interval.end + if start - last_end >= hole_size: + return (last_end + 1, start - 1) + last_end = max(last_end, end) + + + def find_holes(self, hole_size): + sorted_intervals = sorted(self.intervals) + holes = [] + last_end = self.min + for interval in sorted_intervals: + start, end = interval.begin, interval.end + if start - last_end >= hole_size: + holes.append((last_end + 1, start - 1)) + last_end = max(last_end, end) + if last_end < self.max and self.max - last_end >= hole_size: + holes.append((last_end + 1, self.max)) + return holes \ No newline at end of file diff --git a/phases/injector.py b/phases/injector.py index 16eae60..8bd73b5 100644 --- a/phases/injector.py +++ b/phases/injector.py @@ -110,18 +110,30 @@ def injected_fix_data(superpe: SuperPe, carrier: Carrier, exe_host: ExeHost): 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)) - fixup_offset_rdata = peSection.raw_addr + string_off + + rm = exe_host.get_rdata_relocmanager() + if False: # seems i dont need this, even tho i dont understand why + 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)) + rm.add_range(peSection.virt_addr, peSection.virt_addr + string_off) + # Do all .rdata patches for datareuse_fixup in reusedata_fixups: + # get a hole in the .rdata section to put our data + hole = rm.find_hole(len(datareuse_fixup.data)) + if hole == None: + raise Exception("No hole found in .rdata section, abort") + fixup_offset_rdata = hole[0] # the start address of the hole (from start of .rdata) + rm.add_range(hole[0], hole[1]) # mark it as used var_data = datareuse_fixup.data 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 + logging.info(" Add data to .rdata at 0x{:X} (off: {}): {}".format( + datareuse_fixup.addr, fixup_offset_rdata, var_data.decode('utf-16le'))) fixup_offset_rdata += len(var_data) + 8 # patch code section diff --git a/tests/test_rangemanager.py b/tests/test_rangemanager.py new file mode 100644 index 0000000..689d15a --- /dev/null +++ b/tests/test_rangemanager.py @@ -0,0 +1,33 @@ +from typing import List +import unittest +import logging + +from pe.superpe import SuperPe +from model.exehost import ExeHost +from model.rangemanager import RangeManager + + +class RangeManagerTest(unittest.TestCase): + def test_rangemanager(self): + rm = RangeManager(0, 100) + rm.add_range(0, 10) + rm.add_range(20, 30) + rm.add_range(50, 60) + + hole = rm.find_hole(10) + self.assertEqual((11, 19), hole) + + holes = rm.find_holes(20) + self.assertEqual([(31, 49), (61, 100)], holes) + + + def test_relocmanager(self): + exehost = ExeHost("data/exes/procexp64.exe") + exehost.init() + section = exehost.superpe.get_section_by_name(".rdata") + rm = exehost.get_rdata_relocmanager() + self.assertEqual(69, len(rm.intervals)) + # 0x1ab0 is magic currently (should use find_first_utf16_string_offset() + #rm.add_range(section.virt_addr, section.virt_addr + 0x1AB0) + hole = rm.find_hole(20) + self.assertEqual(hole, (1167361, 1173015))