From 3e5de0e60806d9ee91024a73f856e2e3158fea7e Mon Sep 17 00:00:00 2001 From: Dobin Date: Wed, 6 Mar 2024 20:10:14 +0000 Subject: [PATCH] refactor: move some stuff from exehost to superpe --- model/defs.py | 16 ++++++++++++++ model/exehost.py | 54 ++++++------------------------------------------ pe/superpe.py | 43 +++++++++++++++++++++++++++++++++++++- 3 files changed, 64 insertions(+), 49 deletions(-) diff --git a/model/defs.py b/model/defs.py index fc619fe..77ce312 100644 --- a/model/defs.py +++ b/model/defs.py @@ -37,6 +37,22 @@ class SourceStyle(Enum): iat_reuse = "iat_reuse" + +class PeRelocEntry(): + def __init__(self, rva: int, base_rva: int, type: str): + self.rva: int = rva + self.base_rva: int = base_rva + self.offset: int = rva - base_rva + self.type: str = type + + +class IatEntry(): + def __init__(self, dll_name: str, func_name: str, iat_vaddr: int): + self.dll_name: str = dll_name + self.func_name: str = func_name + self.iat_vaddr: int = iat_vaddr + + # no slash at end build_dir = "working/build" logs_dir = "working/logs" diff --git a/model/exehost.py b/model/exehost.py index 1f40750..cd8eb30 100644 --- a/model/exehost.py +++ b/model/exehost.py @@ -12,21 +12,6 @@ from model.rangemanager import RangeManager logger = logging.getLogger("ExeHost") -class PeRelocEntry(): - def __init__(self, rva: int, base_rva: int, type: str): - self.rva: int = rva - self.base_rva: int = base_rva - self.offset: int = rva - base_rva - self.type: str = type - - -class IatEntry(): - def __init__(self, dll_name: str, func_name: str, iat_vaddr: int): - self.dll_name: str = dll_name - self.func_name: str = func_name - self.iat_vaddr: int = iat_vaddr - - class ExeHost(): def __init__(self, filepath: FilePath): self.filepath: FilePath = filepath @@ -67,48 +52,21 @@ class ExeHost(): else: self.dynamic_base = False - # Info output: .text virtual address + # code section we inject to, usually .text self.code_section = self.superpe.get_code_section() logger.info("---[ Injectable: Chosen code section: {} at 0x{:X} size: {}".format( self.code_section.Name.decode().rstrip('\x00'), self.code_section.VirtualAddress, self.code_section.Misc_VirtualSize)) + + # if there is a rwx section, None otherwise + self.rwx_section = self.superpe.get_rwx_section() # relocs - if hasattr(self.superpe.pe, 'DIRECTORY_ENTRY_BASERELOC'): - for base_reloc in self.superpe.pe.DIRECTORY_ENTRY_BASERELOC: - for entry in base_reloc.entries: - rva = entry.rva - base_rva = entry.base_rva - reloc_type = pefile.RELOCATION_TYPE[entry.type][0] - self.base_relocs.append(PeRelocEntry(rva, base_rva, reloc_type)) - - # rwx section - entrypoint = self.superpe.pe.OPTIONAL_HEADER.AddressOfEntryPoint - for section in self.superpe.pe.sections: - if (section.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_READ'] and - section.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_WRITE'] and - section.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_EXECUTE'] - ): - if entrypoint > section.VirtualAddress and entrypoint < section.VirtualAddress + section.Misc_VirtualSize: - self.rwx_section = section - - # If the PE file was loaded using the fast_load=True argument, we will need to parse the data directories: - #pe.parse_data_directories() + self.base_relocs = self.superpe.get_base_relocs() # IAT - for entry in self.superpe.pe.DIRECTORY_ENTRY_IMPORT: - for imp in entry.imports: - dll_name = entry.dll.decode('utf-8') - if imp.name == None: - continue - imp_name = imp.name.decode('utf-8') - imp_addr = imp.address - - if not dll_name in self.iat: - self.iat[dll_name] = [] - - self.iat[dll_name].append(IatEntry(dll_name, imp_name, imp_addr)) + self.iat = self.superpe.get_iat_entries() def get_vaddr_of_iatentry(self, func_name: str) -> int: diff --git a/pe/superpe.py b/pe/superpe.py index 8989748..905e7a4 100644 --- a/pe/superpe.py +++ b/pe/superpe.py @@ -5,6 +5,7 @@ import logging from typing import List from helper import hexdump +from model.defs import * logger = logging.getLogger("superpe") @@ -77,11 +78,24 @@ class SuperPe(): if entrypoint >= sect.VirtualAddress and entrypoint <= sect.VirtualAddress + sect.Misc_VirtualSize: return sect return None - + def get_code_section_data(self) -> bytes: sect = self.get_code_section() return bytes(sect.get_data()) + + + def get_rwx_section(self): + # rwx section + entrypoint = self.pe.OPTIONAL_HEADER.AddressOfEntryPoint + for section in self.pe.sections: + if (section.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_READ'] and + section.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_WRITE'] and + section.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_EXECUTE'] + ): + if entrypoint > section.VirtualAddress and entrypoint < section.VirtualAddress + section.Misc_VirtualSize: + return section + return None def get_section_by_name(self, name: str) -> PeSection: @@ -91,6 +105,33 @@ class SuperPe(): return None + def get_base_relocs(self): + base_relocs = [] + if hasattr(self.pe, 'DIRECTORY_ENTRY_BASERELOC'): + for base_reloc in self.pe.DIRECTORY_ENTRY_BASERELOC: + for entry in base_reloc.entries: + rva = entry.rva + base_rva = entry.base_rva + reloc_type = pefile.RELOCATION_TYPE[entry.type][0] + base_relocs.append(PeRelocEntry(rva, base_rva, reloc_type)) + return base_relocs + + + def get_iat_entries(self): + iat = {} + for entry in self.pe.DIRECTORY_ENTRY_IMPORT: + for imp in entry.imports: + dll_name = entry.dll.decode('utf-8') + if imp.name == None: + continue + imp_name = imp.name.decode('utf-8') + imp_addr = imp.address + + if not dll_name in iat: + iat[dll_name] = [] + iat[dll_name].append(IatEntry(dll_name, imp_name, imp_addr)) + return iat + def write_code_section_data(self, data: bytes): sect = self.get_code_section() self.pe.set_bytes_at_offset(sect.PointerToRawData, data)