diff --git a/model.py b/model.py index ee162f5..232ca7d 100644 --- a/model.py +++ b/model.py @@ -74,10 +74,12 @@ class ExeInfo(): if hasattr(pe, 'DIRECTORY_ENTRY_BASERELOC'): for base_reloc in pe.DIRECTORY_ENTRY_BASERELOC: for entry in base_reloc.entries: - entry_rva = entry.rva + rva = entry.rva + base_rva = entry.base_rva reloc_type = pefile.RELOCATION_TYPE[entry.type][0] self.base_relocs.append({ - 'rva': entry_rva, + 'rva': rva, + 'base_rva': base_rva, 'type': reloc_type, }) diff --git a/phases/datareuse.py b/phases/datareuse.py new file mode 100644 index 0000000..9821485 --- /dev/null +++ b/phases/datareuse.py @@ -0,0 +1,92 @@ +import sys +import pefile +from intervaltree import Interval, IntervalTree +from typing import List + + +class PeSection(): + def __init__(self, pefile_section): + self.name: str = pefile_section.Name.rstrip(b'\x00').decode("utf-8") + self.raw_addr: int = pefile_section.PointerToRawData + self.raw_size: int = pefile_section.SizeOfRawData + self.virt_addr: int = pefile_section.VirtualAddress + self.virt_size: int = pefile_section.Misc_VirtualSize + #self.permissions = pefile_section.Characteristics + + +class PeRelocation(): + def __init__(self, reloc): + self.rva: int = reloc.rva + self.base_rva: int = reloc.base_rva + self.offset: int = reloc.rva - reloc.base_rva + self.type: str = pefile.RELOCATION_TYPE[reloc.type][0] + + +class DataReuser(): + def __init__(self, filepath): + self.pe = pefile.PE(filepath) + self.pe_sections: List[PeSection] = [] + self.base_relocs: List[PeRelocation] = [] + + + def init(self): + # Sections + for section in self.pe.sections: + self.pe_sections.append(PeSection(section)) + + # Base Relocations + if hasattr(self.pe, 'DIRECTORY_ENTRY_BASERELOC'): + for base_reloc in self.pe.DIRECTORY_ENTRY_BASERELOC: + for entry in base_reloc.entries: + self.base_relocs.append(PeRelocation(entry)) + + + def get_section_by_name(self, name: str) -> PeSection: + for section in self.pe_sections: + print("{} {}".format(section.name, name)) + if section.name == name: + return section + return None + + + def get_relocations_all(self) -> List[PeRelocation]: + return self.base_relocs + + + def get_relocations_for_section(self, section_name) -> List[PeRelocation]: + section = self.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] + + + def get_reloc_largest_gap(self, section_name=".rdata"): + tree = IntervalTree() + section = self.get_section_by_name(section_name) + + for reloc in self.base_relocs: + if reloc.base_rva == section.virt_addr: + 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 diff --git a/tests/test_datareuse.py b/tests/test_datareuse.py new file mode 100644 index 0000000..fe332eb --- /dev/null +++ b/tests/test_datareuse.py @@ -0,0 +1,45 @@ +import shutil +from typing import List +import unittest +import logging + +from model import ExeInfo +from defs import * +from helper import hexdump +from observer import observer + +from phases.datareuse import * + + + +class DataReuseTest(unittest.TestCase): + def test_relocation_list(self): + data_reuser = DataReuser("exes/7z.exe") + data_reuser.init() + + relocs = data_reuser.get_relocations_for_section(".rdata") + self.assertEqual(30, len(relocs)) + reloc = relocs[0] + self.assertEqual(393216, reloc.base_rva) + self.assertEqual(394296, reloc.rva) + self.assertEqual(1080, reloc.offset) + self.assertEqual("I", reloc.type) + + + def test_largestgap(self): + data_reuser = DataReuser("exes/7z.exe") + data_reuser.init() + + size, start, stop = data_reuser.get_reloc_largest_gap(".rdata") + self.assertEqual(129395, size) + self.assertEqual(3807, start) + self.assertEqual(133203, stop) + + + def test_rdata_overwrite(self): + pass + + + def test_asm_lea_create(self): + pass +