mirror of
https://github.com/dobin/SuperMega
synced 2026-06-02 17:27:10 +00:00
feature: datareuse initial commit
This commit is contained in:
@@ -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,
|
||||
})
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user