feature: datareuse initial commit

This commit is contained in:
Dobin
2024-02-24 08:51:58 +00:00
parent b032c3f9db
commit 334b17e7a5
3 changed files with 141 additions and 2 deletions
+4 -2
View File
@@ -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,
})
+92
View File
@@ -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
+45
View File
@@ -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