From 81bcb99129defffeca67724553920ca3124aeeba Mon Sep 17 00:00:00 2001 From: sud0woodo <40278342+sud0woodo@users.noreply.github.com> Date: Sat, 9 Mar 2024 12:24:35 +0100 Subject: [PATCH 01/43] Initial commit for dissect.executable.pe First version for the dissect PE format parser. Currently the parser supports standard PE files, as well as some patching of executables and building an executable from scratch. It is expected there will be some issues that will be found during usage and it will need quite some refactoring. The API as it currently exists might need some revising too :) --- dissect/executable/__init__.py | 2 + dissect/executable/exception.py | 25 + dissect/executable/pe/__init__.py | 25 + dissect/executable/pe/helpers/__init__.py | 0 dissect/executable/pe/helpers/builder.py | 470 +++++++++++++++ dissect/executable/pe/helpers/c_pe.py | 466 ++++++++++++++ dissect/executable/pe/helpers/exports.py | 84 +++ dissect/executable/pe/helpers/imports.py | 326 ++++++++++ dissect/executable/pe/helpers/patcher.py | 344 +++++++++++ dissect/executable/pe/helpers/relocations.py | 53 ++ dissect/executable/pe/helpers/resources.py | 418 +++++++++++++ dissect/executable/pe/helpers/sections.py | 258 ++++++++ dissect/executable/pe/helpers/tls.py | 118 ++++ dissect/executable/pe/helpers/utils.py | 40 ++ dissect/executable/pe/pe.py | 602 +++++++++++++++++++ tests/data/testexe.exe | Bin 0 -> 31336 bytes tests/test_pe.py | 78 +++ tests/test_pe_builder.py | 69 +++ tests/test_pe_modifications.py | 94 +++ 19 files changed, 3472 insertions(+) create mode 100644 dissect/executable/pe/helpers/__init__.py create mode 100644 dissect/executable/pe/helpers/builder.py create mode 100755 dissect/executable/pe/helpers/c_pe.py create mode 100644 dissect/executable/pe/helpers/exports.py create mode 100644 dissect/executable/pe/helpers/imports.py create mode 100644 dissect/executable/pe/helpers/patcher.py create mode 100644 dissect/executable/pe/helpers/relocations.py create mode 100644 dissect/executable/pe/helpers/resources.py create mode 100644 dissect/executable/pe/helpers/sections.py create mode 100644 dissect/executable/pe/helpers/tls.py create mode 100644 dissect/executable/pe/helpers/utils.py create mode 100644 dissect/executable/pe/pe.py create mode 100644 tests/data/testexe.exe create mode 100644 tests/test_pe.py create mode 100644 tests/test_pe_builder.py create mode 100644 tests/test_pe_modifications.py diff --git a/dissect/executable/__init__.py b/dissect/executable/__init__.py index 43d2a88..564a688 100644 --- a/dissect/executable/__init__.py +++ b/dissect/executable/__init__.py @@ -1,5 +1,7 @@ from dissect.executable.elf import ELF +from dissect.executable.pe import PE __all__ = [ "ELF", + "PE", ] diff --git a/dissect/executable/exception.py b/dissect/executable/exception.py index b0fb678..0043c51 100644 --- a/dissect/executable/exception.py +++ b/dissect/executable/exception.py @@ -4,3 +4,28 @@ class Error(Exception): class InvalidSignatureError(Error): """Exception that occurs if the magic in the header does not match.""" + + +class InvalidPE(Error): + """Exception that occurs if the PE signature does not match.""" + + +class InvalidVA(Error): + """Exception that occurs when a virtual address is not found within the PE sections.""" + + +class InvalidAddress(Error): + """Exception that occurs when a raw address is not found within the PE file when translating from a virtual + address.""" + + +class InvalidArchitecture(Error): + """Exception that occurs when an invalid value is encountered for the PE architecture types.""" + + +class BuildSectionException(Error): + """Exception that occurs when the section to be build contains an error.""" + + +class ResourceException(Error): + """Exception that occurs when an error is thrown parsing the resources.""" diff --git a/dissect/executable/pe/__init__.py b/dissect/executable/pe/__init__.py index e69de29..f7e9849 100644 --- a/dissect/executable/pe/__init__.py +++ b/dissect/executable/pe/__init__.py @@ -0,0 +1,25 @@ +from dissect.executable.pe.helpers.builder import Builder +from dissect.executable.pe.helpers.exports import ExportFunction, ExportManager +from dissect.executable.pe.helpers.imports import ( + ImportFunction, + ImportManager, + ImportModule, +) +from dissect.executable.pe.helpers.patcher import Patcher +from dissect.executable.pe.helpers.resources import Resource, ResourceManager +from dissect.executable.pe.helpers.sections import PESection +from dissect.executable.pe.pe import PE + +__all__ = [ + "Builder", + "ExportFunction", + "ExportManager", + "ImportFunction", + "ImportManager", + "ImportModule", + "Patcher", + "PE", + "PESection", + "Resource", + "ResourceManager", +] diff --git a/dissect/executable/pe/helpers/__init__.py b/dissect/executable/pe/helpers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dissect/executable/pe/helpers/builder.py b/dissect/executable/pe/helpers/builder.py new file mode 100644 index 0000000..acc1484 --- /dev/null +++ b/dissect/executable/pe/helpers/builder.py @@ -0,0 +1,470 @@ +from __future__ import annotations + +from datetime import datetime +from io import BytesIO +from typing import TYPE_CHECKING + +from dissect.executable.exception import BuildSectionException +from dissect.executable.pe.helpers import utils +from dissect.executable.pe.helpers.c_pe import pestruct + +# Local imports +from dissect.executable.pe.pe import PE + +if TYPE_CHECKING: + from dissect.cstruct.cstruct import cstruct + + +STUB = b"\x0e\x1f\xba\x0e\x00\xb4\t\xcd!\xb8\x01L\xcd!This program is made with dissect.pe <3 kusjesvanSRT <3.\x0D\x0D\x0A$\x00\x00" # noqa: E501 + + +class Builder: + """Base class for building the PE file with the user applied patches. + + Args: + pe: A `PE` object. + arch: The architecture to use for the new PE. + dll: Whether the new PE should be a DLL or not. + subsystem: The subsystem to use for the new PE default uses IMAGE_SUBSYSTEM_WINDOWS_GUI. + """ + + def __init__( + self, + arch: str = "x64", + dll: bool = False, + subsystem: int = 0x2, + ): + self.arch = ( + pestruct.MachineType.IMAGE_FILE_MACHINE_AMD64 + if arch == "x64" + else pestruct.MachineType.IMAGE_FILE_MACHINE_I386 + ) + self.dll = dll + self.subsystem = subsystem + + self.pe = None + + def new(self): + """Build the PE file from scratch. + + This function will build a new PE that consists of a single dummy section. It will not contain any imports, + exports, code, etc. + """ + + new_pe = BytesIO() + + # Generate the MZ header + self.mz_header = self.gen_mz_header() + + image_characteristics = self.get_characteristics() + # Generate the file header + self.file_header = self.gen_file_header(machine=self.arch, characteristics=image_characteristics) + + # Generate the optional header + self.optional_header = self.gen_optional_header() + + # Add a dummy section header to the new PE, we need at least 1 section to parse the PE + dummy_data = b"<3kusjesvanSRT<3" + dummy_multiplier = 0x400 // len(b"<3kusjesvanSRT<3") + + section_header_offset = self.optional_header.SizeOfHeaders + pointer_to_raw_data = utils.align_int( + integer=section_header_offset + pestruct.IMAGE_SECTION_HEADER.size, blocksize=self.file_alignment + ) + dummy_section = self.section( + pointer_to_raw_data=pointer_to_raw_data, + virtual_address=self.optional_header.BaseOfCode, + virtual_size=dummy_multiplier, + raw_size=dummy_multiplier, + characteristics=pestruct.SectionFlags.IMAGE_SCN_CNT_CODE + | pestruct.SectionFlags.IMAGE_SCN_MEM_EXECUTE + | pestruct.SectionFlags.IMAGE_SCN_MEM_READ + | pestruct.SectionFlags.IMAGE_SCN_MEM_NOT_PAGED, + ) + # Update the number of sections in the file header + self.file_header.NumberOfSections += 1 + + # Write the headers into the new PE + new_pe.write(self.mz_header.dumps()) + new_pe.write(STUB) + new_pe.seek(self.mz_header.e_lfanew) + new_pe.write(b"PE\x00\x00") + new_pe.write(self.file_header.dumps()) + new_pe.write(self.optional_header.dumps()) + + # Write the dummy section header + new_pe.write(dummy_section.dumps()) + + # Write the data of the section + new_pe.seek(dummy_section.PointerToRawData) + new_pe.write(dummy_data * dummy_multiplier) + + self.pe = PE(pe_file=new_pe) + + # Fix our SizeOfImage field in the optional header + self.pe.optional_header.SizeOfImage = self.pe_size + + def gen_mz_header( + self, + e_magic: int = 0x5A4D, + e_cblp: int = 0, + e_cp: int = 1, + e_crlc: int = 0, + e_cparhdr: int = 4, + e_minalloc: int = 0, + e_maxalloc: int = 0, + e_ss: int = 0, + e_sp: int = 0, + e_csum: int = 0, + e_ip: int = 0, + e_cs: int = 0, + e_lfarlc: int = 64, + e_ovno: int = 0, + e_res: list = [0, 0, 0, 0], + e_oemid: int = 0, + e_oeminfo: int = 0, + e_res2: int = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + e_lfanew: int = 0, + ) -> cstruct: + """Generate the MZ header for the new PE file. + + Args: + e_magic: The magic number for the MZ header. + e_cblp: The number of bytes on the last page of the file. + e_cp: The number of pages in the file. + e_crlc: The number of relocations. + e_cparhdr: The number of paragraphs in the header. + e_minalloc: The minimum number of paragraphs in the program. + e_maxalloc: The maximum number of paragraphs in the program. + e_ss: The relative value of the stack segment. + e_sp: The initial value of the stack pointer. + e_csum: The checksum. + e_ip: The initial value of the instruction pointer. + e_cs: The relative value of the code segment. + e_lfarlc: The file address of the relocation table. + e_ovno: The overlay number. + e_res: The reserved words. + e_oemid: The OEM identifier. + e_oeminfo: The OEM information. + e_res2: The reserved words. + e_lfanew: The file address of the new exe header. + + Returns: + The MZ header as a `cstruct` object. + """ + + mz_header = pestruct.IMAGE_DOS_HEADER() + + mz_header.e_magic = e_magic + mz_header.e_cblp = e_cblp + mz_header.e_cp = e_cp + mz_header.e_crlc = e_crlc + mz_header.e_cparhdr = e_cparhdr + mz_header.e_minalloc = e_minalloc + mz_header.e_maxalloc = e_maxalloc + mz_header.e_ss = e_ss + mz_header.e_sp = e_sp + mz_header.e_csum = e_csum + mz_header.e_ip = e_ip + mz_header.e_cs = e_cs + mz_header.e_lfarlc = e_lfarlc + mz_header.e_ovno = e_ovno + mz_header.e_res = e_res + mz_header.e_oemid = e_oemid + mz_header.e_oeminfo = e_oeminfo + mz_header.e_res2 = e_res2 + + # Calculate the start of the NT headers by checking the location and size of the relocation table + # within the MZ header + start_of_nt_header = (mz_header.e_lfarlc + (mz_header.e_crlc * 4)) + len(STUB) + mz_header.e_lfanew = start_of_nt_header if not e_lfanew else e_lfanew + # Align the e_lfanew value + mz_header.e_lfanew = mz_header.e_lfanew + (mz_header.e_lfanew % 2) + + return mz_header + + def get_characteristics(self) -> int: + """Function to retreive the characteristics that are set based on the kind of PE file that needs to be + generated. + + For now it will only contain the main characteristics of a PE file, like if it's an executable image and/or a + DLL. + + Returns: + The characteristics of the PE file. + """ + + if self.arch == pestruct.MachineType.IMAGE_FILE_MACHINE_AMD64: + if self.dll: + return ( + pestruct.ImageCharacteristics.IMAGE_FILE_EXECUTABLE_IMAGE + | pestruct.ImageCharacteristics.IMAGE_FILE_DLL + ) + else: + return pestruct.ImageCharacteristics.IMAGE_FILE_EXECUTABLE_IMAGE + else: + if self.dll: + return ( + pestruct.ImageCharacteristics.IMAGE_FILE_EXECUTABLE_IMAGE + | pestruct.ImageCharacteristics.IMAGE_FILE_DLL + | pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE + ) + else: + return ( + pestruct.ImageCharacteristics.IMAGE_FILE_EXECUTABLE_IMAGE + | pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE + ) + + def gen_file_header( + self, + time_date_stamp: int = 0, + pointer_to_symbol_table: int = 0, + number_of_symbols: int = 0, + size_of_optional_header: int = 0, + characteristics: int = 0, + machine: int = 0x8664, + number_of_sections: int = 0, + ) -> cstruct: + """Generate the file header for the new PE file. + + Args: + machine: The machine type. + number_of_sections: The number of sections. + time_date_stamp: The time and date the file was created. + pointer_to_symbol_table: The file pointer to the COFF symbol table. + number_of_symbols: The number of entries in the symbol table. + size_of_optional_header: The size of the optional header. + characteristics: The characteristics of the file. + + Returns: + The file header as a `cstruct` object. + """ + + # Set the size of the optional header if not given + if not size_of_optional_header: + if machine == 0x8664: + size_of_optional_header = len(pestruct.IMAGE_OPTIONAL_HEADER64) + self.machine = 0x8664 + else: + size_of_optional_header = len(pestruct.IMAGE_OPTIONAL_HEADER) + self.machine = 0x14C + + # Set the timestamp to now if not given + if not time_date_stamp: + time_date_stamp = int(datetime.utcnow().timestamp()) + + file_header = pestruct.IMAGE_FILE_HEADER() + file_header.Machine = machine + file_header.NumberOfSections = number_of_sections + file_header.TimeDateStamp = time_date_stamp + file_header.PointerToSymbolTable = pointer_to_symbol_table + file_header.NumberOfSymbols = number_of_symbols + file_header.SizeOfOptionalHeader = size_of_optional_header + file_header.Characteristics = characteristics + + return file_header + + def gen_optional_header( + self, + magic: int = 0, + major_linker_version: int = 0xE, + minor_linker_version: int = 0, + size_of_code: int = 0, + size_of_initialized_data: int = 0, + size_of_uninitialized_data: int = 0, + address_of_entrypoint: int = 0, + base_of_code: int = 0x1000, + imagebase: int = 0x69000, + section_alignment: int = 0x1000, + file_alignment: int = 0x200, + major_os_version: int = 0x5, + minor_os_version: int = 0x2, + major_image_version: int = 0, + minor_image_version: int = 0, + major_subsystem_version: int = 0x5, + minor_subsystem_version: int = 0x2, + win32_version_value: int = 0, + size_of_image: int = 0, + size_of_headers: int = 0x400, + checksum: int = 0, + subsystem: int = 0x2, + dll_characteristics: int = 0, + size_of_stack_reserve: int = 0x1000, + size_of_stack_commit: int = 0x1000, + size_of_heap_reserve: int = 0x1000, + size_of_heap_commit: int = 0x1000, + loaderflags: int = 0, + number_of_rva_and_sizes: int = pestruct.IMAGE_NUMBEROF_DIRECTORY_ENTRIES, + datadirectory: list = [ + pestruct.IMAGE_DATA_DIRECTORY(BytesIO(b"\x00" * len(pestruct.IMAGE_DATA_DIRECTORY))) + for _ in range(pestruct.IMAGE_NUMBEROF_DIRECTORY_ENTRIES) + ], + ) -> cstruct: + """Generate the optional header for the new PE file. + + Args: + magic: The magic number for the optional header, this indicates the architecture for the PE. + major_linker_version: The major version of the linker. + minor_linker_version: The minor version of the linker. + size_of_code: The size of the code section. + size_of_initialized_data: The size of the initialized data section. + size_of_uninitialized_data: The size of the uninitialized data section. + address_of_entrypoint: The address of the entry point. + base_of_code: The base of the code section. + imagebase: The base address of the image. + section_alignment: The alignment of sections in memory. + file_alignment: The alignment of sections in the file. + major_os_version: The major version of the operating system. + minor_os_version: The minor version of the operating system. + major_image_version: The major version of the image. + minor_image_version: The minor version of the image. + major_subsystem_version: The major version of the subsystem. + minor_subsystem_version: The minor version of the subsystem. + win32_version_value: The Win32 version value. + size_of_image: The size of the image. + size_of_headers: The size of the headers. + checksum: The checksum of the image. + subsystem: The subsystem of the image. + dll_characteristics: The DLL characteristics of the image. + size_of_stack_reserve: The size of the stack to reserve. + size_of_stack_commit: The size of the stack to commit. + size_of_heap_reserve: The size of the heap to reserve. + size_of_heap_commit: The size of the heap to commit. + loaderflags: The loader flags. + number_of_rva_and_sizes: The number of RVA and sizes. + datadirectory: The data directory entries, initialized as nullbyte directories. + + Returns: + The optional header as a `cstruct` object. + """ + + if self.machine == 0x8664: + optional_header = pestruct.IMAGE_OPTIONAL_HEADER64() + optional_header.Magic = 0x20B if not magic else magic + else: + optional_header = pestruct.IMAGE_OPTIONAL_HEADER() + optional_header.Magic = 0x10B if not magic else magic + + self.file_alignment = file_alignment + self.section_alignment = section_alignment + + # Calculate the SizeOfHeaders field, we add the length of a section header because we know there's going to be + # at least 1 section header + size_of_headers = utils.align_int( + integer=len(self.mz_header) + + len(STUB) + + len(b"PE\x00\x00") + + len(self.file_header) + + len(optional_header) + + len(pestruct.IMAGE_SECTION_HEADER), + blocksize=file_alignment, + ) + + optional_header.MajorLinkerVersion = major_linker_version + optional_header.MinorLinkerVersion = minor_linker_version + optional_header.SizeOfCode = size_of_code + optional_header.SizeOfInitializedData = size_of_initialized_data + optional_header.SizeOfUninitializedData = size_of_uninitialized_data + optional_header.AddressOfEntryPoint = address_of_entrypoint + optional_header.BaseOfCode = base_of_code + optional_header.ImageBase = imagebase + optional_header.SectionAlignment = section_alignment + optional_header.FileAlignment = file_alignment + optional_header.MajorOperatingSystemVersion = major_os_version + optional_header.MinorOperatingSystemVersion = minor_os_version + optional_header.MajorImageVersion = major_image_version + optional_header.MinorImageVersion = minor_image_version + optional_header.MajorSubsystemVersion = major_subsystem_version + optional_header.MinorSubsystemVersion = minor_subsystem_version + optional_header.Win32VersionValue = win32_version_value + optional_header.SizeOfImage = size_of_image + optional_header.SizeOfHeaders = size_of_headers + optional_header.CheckSum = checksum + optional_header.Subsystem = subsystem + optional_header.DllCharacteristics = dll_characteristics + optional_header.SizeOfStackReserve = size_of_stack_reserve + optional_header.SizeOfStackCommit = size_of_stack_commit + optional_header.SizeOfHeapReserve = size_of_heap_reserve + optional_header.SizeOfHeapCommit = size_of_heap_commit + optional_header.LoaderFlags = loaderflags + optional_header.NumberOfRvaAndSizes = number_of_rva_and_sizes + optional_header.DataDirectory = datadirectory + + return optional_header + + def section( + self, + pointer_to_raw_data: int, + name: str | bytes = b".dissect", + virtual_size: int = 0x1000, + virtual_address: int = 0x1000, + raw_size: int = 0x200, + pointer_to_relocations: int = 0, + pointer_to_linenumbers: int = 0, + number_of_relocations: int = 0, + number_of_linenumbers: int = 0, + characteristics: int = 0x68000020, + ) -> cstruct: + """Build a new section for the PE. + + The default characteristics of the new section will be: + - IMAGE_SCN_CNT_CODE + - IMAGE_SCN_MEM_EXECUTE + - IMAGE_SCN_MEM_READ + - IMAGE_SCN_MEM_NOT_PAGED + + Args: + pointer_to_raw_data: The file pointer to the raw data of the new section. + name: The new section name, default: .dissect + virtual_size: The virtual size of the new section data. + virtual_address: The virtual address where the new section is located. + raw_size: The size of the section data. + pointer_to_relocations: The file pointer to the relocation table. + pointer_to_linenumbers: The file pointer to the line number table. + number_of_relocations: The number of relocations. + number_of_linenumbers: The number of line numbers. + characteristics: The characteristics of the new section. + + Returns: + The new section header as a `cstruct` object. + """ + + if len(name) > 8: + raise BuildSectionException("section names can't be longer than 8 characters") + + if isinstance(name, str): + name = name.encode() + + section_header = pestruct.IMAGE_SECTION_HEADER() + + pointer_to_raw_data = utils.align_int(integer=pointer_to_raw_data, blocksize=self.file_alignment) + + section_header.Name = name + utils.pad(size=8 - len(name)) + section_header.VirtualSize = virtual_size + section_header.VirtualAddress = virtual_address + section_header.SizeOfRawData = raw_size + section_header.PointerToRawData = pointer_to_raw_data + section_header.PointerToRelocations = pointer_to_relocations + section_header.PointerToLinenumbers = pointer_to_linenumbers + section_header.NumberOfRelocations = number_of_relocations + section_header.NumberOfLinenumbers = number_of_linenumbers + section_header.Characteristics = characteristics + + return section_header + + @property + def pe_size(self) -> int: + """Calculate the new PE size. + + We can calculate the new size of the PE by adding the virtual address and virtual size of the last section + together. + + Returns: + The size of the PE. + """ + + last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] + va = last_section.virtual_address + size = last_section.virtual_size + + return utils.align_int(integer=(va + size), blocksize=self.section_alignment) diff --git a/dissect/executable/pe/helpers/c_pe.py b/dissect/executable/pe/helpers/c_pe.py new file mode 100755 index 0000000..606fd6c --- /dev/null +++ b/dissect/executable/pe/helpers/c_pe.py @@ -0,0 +1,466 @@ +from dissect.cstruct import cstruct + +pe_def = """ +#define IMAGE_NUMBEROF_DIRECTORY_ENTRIES 16 +#define IMAGE_SIZEOF_SHORT_NAME 8 + +#define IMAGE_DIRECTORY_ENTRY_EXPORT 0 // Export Directory +#define IMAGE_DIRECTORY_ENTRY_IMPORT 1 // Import Directory +#define IMAGE_DIRECTORY_ENTRY_RESOURCE 2 // Resource Directory +#define IMAGE_DIRECTORY_ENTRY_EXCEPTION 3 // Exception Directory +#define IMAGE_DIRECTORY_ENTRY_SECURITY 4 // Security Directory +#define IMAGE_DIRECTORY_ENTRY_BASERELOC 5 // Base Relocation Table +#define IMAGE_DIRECTORY_ENTRY_DEBUG 6 // Debug Directory +// IMAGE_DIRECTORY_ENTRY_COPYRIGHT 7 // (X86 usage) +#define IMAGE_DIRECTORY_ENTRY_ARCHITECTURE 7 // Architecture Specific Data +#define IMAGE_DIRECTORY_ENTRY_GLOBALPTR 8 // RVA of GP +#define IMAGE_DIRECTORY_ENTRY_TLS 9 // TLS Directory +#define IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG 10 // Load Configuration Directory +#define IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT 11 // Bound Import Directory in headers +#define IMAGE_DIRECTORY_ENTRY_IAT 12 // Import Address Table +#define IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT 13 // Delay Load Import Descriptors +#define IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR 14 // COM Runtime descriptor + +// --- PE HEADERS --- + +typedef struct IMAGE_DOS_HEADER { + WORD e_magic; + WORD e_cblp; + WORD e_cp; + WORD e_crlc; + WORD e_cparhdr; + WORD e_minalloc; + WORD e_maxalloc; + WORD e_ss; + WORD e_sp; + WORD e_csum; + WORD e_ip; + WORD e_cs; + WORD e_lfarlc; + WORD e_ovno; + WORD e_res[4]; + WORD e_oemid; + WORD e_oeminfo; + WORD e_res2[10]; + LONG e_lfanew; +}; + +typedef enum MachineType : WORD { + IMAGE_FILE_MACHINE_UNKNOWN = 0x0, + IMAGE_FILE_MACHINE_AM33 = 0x1d3, + IMAGE_FILE_MACHINE_AMD64 = 0x8664, + IMAGE_FILE_MACHINE_ARM = 0x1c0, + IMAGE_FILE_MACHINE_ARM64 = 0xaa64, + IMAGE_FILE_MACHINE_ARMNT = 0x1c4, + IMAGE_FILE_MACHINE_EBC = 0xebc, + IMAGE_FILE_MACHINE_I386 = 0x14c, + IMAGE_FILE_MACHINE_IA64 = 0x200, + IMAGE_FILE_MACHINE_M32R = 0x9041, + IMAGE_FILE_MACHINE_MIPS16 = 0x266, + IMAGE_FILE_MACHINE_MIPSFPU = 0x366, + IMAGE_FILE_MACHINE_MIPSFPU16 = 0x466, + IMAGE_FILE_MACHINE_POWERPC = 0x1f0, + IMAGE_FILE_MACHINE_POWERPCFP = 0x1f1, + IMAGE_FILE_MACHINE_R4000 = 0x166, + IMAGE_FILE_MACHINE_RISCV32 = 0x5032, + IMAGE_FILE_MACHINE_RISCV64 = 0x5064, + IMAGE_FILE_MACHINE_RISCV128 = 0x5128, + IMAGE_FILE_MACHINE_SH3 = 0x1a2, + IMAGE_FILE_MACHINE_SH3DSP = 0x1a3, + IMAGE_FILE_MACHINE_SH4 = 0x1a6, + IMAGE_FILE_MACHINE_SH5 = 0x1a8, + IMAGE_FILE_MACHINE_THUMB = 0x1c2, + IMAGE_FILE_MACHINE_WCEMIPSV2 = 0x169, +}; + +flag ImageCharacteristics : WORD { + IMAGE_FILE_RELOCS_STRIPPED = 0x0001, + IMAGE_FILE_EXECUTABLE_IMAGE = 0x0002, + IMAGE_FILE_LINE_NUMS_STRIPPED = 0x0004, + IMAGE_FILE_LOCAL_SYMS_STRIPPED = 0x0008, + IMAGE_FILE_AGGRESSIVE_WS_TRIM = 0x0010, + IMAGE_FILE_LARGE_ADDRESS_AWARE = 0x0020, + Reserved = 0x0040, + IMAGE_FILE_BYTES_REVERSED_LO = 0x0080, + IMAGE_FILE_32BIT_MACHINE = 0x0100, + IMAGE_FILE_DEBUG_STRIPPED = 0x0200, + IMAGE_FILE_REMOVABLE_RUN_FROM_SWAP = 0x0400, + IMAGE_FILE_NET_RUN_FROM_SWAP = 0x0800, + IMAGE_FILE_SYSTEM = 0x1000, + IMAGE_FILE_DLL = 0x2000, + IMAGE_FILE_UP_SYSTEM_ONLY = 0x4000, + IMAGE_FILE_BYTES_REVERSED_HI = 0x8000, +}; + +typedef struct IMAGE_FILE_HEADER { + MachineType Machine; + WORD NumberOfSections; + DWORD TimeDateStamp; + DWORD PointerToSymbolTable; + DWORD NumberOfSymbols; + WORD SizeOfOptionalHeader; + ImageCharacteristics Characteristics; +}; + +typedef struct IMAGE_DATA_DIRECTORY { + ULONG VirtualAddress; + ULONG Size; +}; + +typedef enum WindowsSubsystem : WORD { + IMAGE_SUBSYSTEM_UNKNOWN = 0, + IMAGE_SUBSYSTEM_NATIVE = 1, + IMAGE_SUBSYSTEM_WINDOWS_GUI = 2, + IMAGE_SUBSYSTEM_WINDOWS_CUI = 3, + IMAGE_SUBSYSTEM_OS2_CUI = 5, + IMAGE_SUBSYSTEM_POSIX_CUI = 7, + IMAGE_SUBSYSTEM_NATIVE_WINDOWS = 8, + IMAGE_SUBSYSTEM_WINDOWS_CE_GUI = 9, + IMAGE_SUBSYSTEM_EFI_APPLICATION = 10, + IMAGE_SUBSYSTEM_EFI_BOOT_SERVICE_DRIVER = 11, + IMAGE_SUBSYSTEM_EFI_RUNTIME_DRIVER = 12, + IMAGE_SUBSYSTEM_EFI_ROM = 13, + IMAGE_SUBSYSTEM_XBOX = 14, + IMAGE_SUBSYSTEM_WINDOWS_BOOT_APPLICATION = 16, +}; + +typedef enum DLLCharacteristics : WORD { + IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA = 0x0020, + IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE = 0x0040, + IMAGE_DLLCHARACTERISTICS_FORCE_INTEGRITY = 0x0080, + IMAGE_DLLCHARACTERISTICS_NX_COMPAT = 0x0100, + IMAGE_DLLCHARACTERISTICS_NO_ISOLATION = 0x0200, + IMAGE_DLLCHARACTERISTICS_NO_SEH = 0x0400, + IMAGE_DLLCHARACTERISTICS_NO_BIND = 0x0800, + IMAGE_DLLCHARACTERISTICS_APPCONTAINER = 0x1000, + IMAGE_DLLCHARACTERISTICS_WDM_DRIVER = 0x2000, + IMAGE_DLLCHARACTERISTICS_GUARD_CF = 0x4000, + IMAGE_DLLCHARACTERISTICS_TERMINAL_SERVER_AWARE = 0x8000, +}; + +typedef struct IMAGE_OPTIONAL_HEADER { + WORD Magic; + BYTE MajorLinkerVersion; + BYTE MinorLinkerVersion; + DWORD SizeOfCode; + DWORD SizeOfInitializedData; + DWORD SizeOfUninitializedData; + DWORD AddressOfEntryPoint; + DWORD BaseOfCode; + DWORD BaseOfData; + DWORD ImageBase; + DWORD SectionAlignment; + DWORD FileAlignment; + WORD MajorOperatingSystemVersion; + WORD MinorOperatingSystemVersion; + WORD MajorImageVersion; + WORD MinorImageVersion; + WORD MajorSubsystemVersion; + WORD MinorSubsystemVersion; + DWORD Win32VersionValue; + DWORD SizeOfImage; + DWORD SizeOfHeaders; + DWORD CheckSum; + WindowsSubsystem Subsystem; + DLLCharacteristics DllCharacteristics; + DWORD SizeOfStackReserve; + DWORD SizeOfStackCommit; + DWORD SizeOfHeapReserve; + DWORD SizeOfHeapCommit; + DWORD LoaderFlags; + DWORD NumberOfRvaAndSizes; + IMAGE_DATA_DIRECTORY DataDirectory[IMAGE_NUMBEROF_DIRECTORY_ENTRIES]; +}; + +typedef struct IMAGE_OPTIONAL_HEADER64 { + WORD Magic; + BYTE MajorLinkerVersion; + BYTE MinorLinkerVersion; + DWORD SizeOfCode; + DWORD SizeOfInitializedData; + DWORD SizeOfUninitializedData; + DWORD AddressOfEntryPoint; + DWORD BaseOfCode; + ULONGLONG ImageBase; + DWORD SectionAlignment; + DWORD FileAlignment; + WORD MajorOperatingSystemVersion; + WORD MinorOperatingSystemVersion; + WORD MajorImageVersion; + WORD MinorImageVersion; + WORD MajorSubsystemVersion; + WORD MinorSubsystemVersion; + DWORD Win32VersionValue; + DWORD SizeOfImage; + DWORD SizeOfHeaders; + DWORD CheckSum; + WORD Subsystem; + WORD DllCharacteristics; + ULONGLONG SizeOfStackReserve; + ULONGLONG SizeOfStackCommit; + ULONGLONG SizeOfHeapReserve; + ULONGLONG SizeOfHeapCommit; + DWORD LoaderFlags; + DWORD NumberOfRvaAndSizes; + IMAGE_DATA_DIRECTORY DataDirectory[IMAGE_NUMBEROF_DIRECTORY_ENTRIES]; +}; + +typedef struct IMAGE_NT_HEADERS { + DWORD Signature; + IMAGE_FILE_HEADER FileHeader; + IMAGE_OPTIONAL_HEADER OptionalHeader; +}; + +typedef struct IMAGE_NT_HEADERS64 { + DWORD Signature; + IMAGE_FILE_HEADER FileHeader; + IMAGE_OPTIONAL_HEADER64 OptionalHeader; +}; + +flag SectionFlags : DWORD { + IMAGE_SCN_TYPE_NO_PAD = 0x00000008, + IMAGE_SCN_CNT_CODE = 0x00000020, + IMAGE_SCN_CNT_INITIALIZED_DATA = 0x00000040, + IMAGE_SCN_CNT_UNINITIALIZED_DATA = 0x00000080, + IMAGE_SCN_LNK_OTHER = 0x00000100, + IMAGE_SCN_LNK_INFO = 0x00000200, + IMAGE_SCN_LNK_REMOVE = 0x00000800, + IMAGE_SCN_LNK_COMDAT = 0x00001000, + IMAGE_SCN_NO_DEFER_SPEC_EXC = 0x00004000, + IMAGE_SCN_GPREL = 0x00008000, + IMAGE_SCN_MEM_FARDATA = 0x00008000, + IMAGE_SCN_MEM_PURGEABLE = 0x00020000, + IMAGE_SCN_MEM_16BIT = 0x00020000, + IMAGE_SCN_MEM_LOCKED = 0x00040000, + IMAGE_SCN_MEM_PRELOAD = 0x00080000, + IMAGE_SCN_ALIGN_1BYTES = 0x00100000, + IMAGE_SCN_ALIGN_2BYTES = 0x00200000, + IMAGE_SCN_ALIGN_4BYTES = 0x00300000, + IMAGE_SCN_ALIGN_8BYTES = 0x00400000, + IMAGE_SCN_ALIGN_16BYTES = 0x00500000, + IMAGE_SCN_ALIGN_32BYTES = 0x00600000, + IMAGE_SCN_ALIGN_64BYTES = 0x00700000, + IMAGE_SCN_ALIGN_128BYTES = 0x00800000, + IMAGE_SCN_ALIGN_256BYTES = 0x00900000, + IMAGE_SCN_ALIGN_512BYTES = 0x00A00000, + IMAGE_SCN_ALIGN_1024BYTES = 0x00B00000, + IMAGE_SCN_ALIGN_2048BYTES = 0x00C00000, + IMAGE_SCN_ALIGN_4096BYTES = 0x00D00000, + IMAGE_SCN_ALIGN_8192BYTES = 0x00E00000, + IMAGE_SCN_LNK_NRELOC_OVFL = 0x01000000, + IMAGE_SCN_MEM_DISCARDABLE = 0x02000000, + IMAGE_SCN_MEM_NOT_CACHED = 0x04000000, + IMAGE_SCN_MEM_NOT_PAGED = 0x08000000, + IMAGE_SCN_MEM_SHARED = 0x10000000, + IMAGE_SCN_MEM_EXECUTE = 0x20000000, + IMAGE_SCN_MEM_READ = 0x40000000, + IMAGE_SCN_MEM_WRITE = 0x80000000 +}; + +typedef struct IMAGE_SECTION_HEADER { + char Name[IMAGE_SIZEOF_SHORT_NAME]; + ULONG VirtualSize; + ULONG VirtualAddress; + ULONG SizeOfRawData; + ULONG PointerToRawData; + ULONG PointerToRelocations; + ULONG PointerToLinenumbers; + USHORT NumberOfRelocations; + USHORT NumberOfLinenumbers; + SectionFlags Characteristics; +}; + +// --- END OF PE HEADERS + +// --- EXPORTS + +typedef struct IMAGE_EXPORT_DIRECTORY { + ULONG Characteristics; + ULONG TimeDateStamp; + USHORT MajorVersion; + USHORT MinorVersion; + ULONG Name; + ULONG Base; + ULONG NumberOfFunctions; + ULONG NumberOfNames; + ULONG AddressOfFunctions; // RVA from base of image + ULONG AddressOfNames; // RVA from base of image + ULONG AddressOfNameOrdinals; // RVA from base of image +}; + +// --- END OF EXPORTS + +// --- IMPORTS + +typedef struct IMAGE_IMPORT_DESCRIPTOR { + DWORD OriginalFirstThunk; + DWORD TimeDateStamp; + DWORD ForwarderChain; + DWORD Name; + DWORD FirstThunk; +}; + +typedef struct IMAGE_IMPORT_BY_NAME { + uint16 Hint; + // char Name; +}; + +typedef struct IMAGE_THUNK_DATA32 { + union { + DWORD ForwarderString; + DWORD Function; + DWORD Ordinal; + DWORD AddressOfData; + } u1; +}; + +typedef struct IMAGE_THUNK_DATA64 { + union { + ULONGLONG ForwarderString; + ULONGLONG Function; + ULONGLONG Ordinal; + ULONGLONG AddressOfData; + } u1; +} + +// --- END OF IMPORTS + +// --- RESOURCE DIRECTORY + +enum ResourceID : WORD { + Cursor = 0x1, + Bitmap = 0x2, + Icon = 0x3, + Menu = 0x4, + Dialog = 0x5, + String = 0x6, + FontDirectory = 0x7, + Font = 0x8, + Accelerator = 0x9, + RcData = 0xA, + MessageTable = 0xB, + Version = 0x10, + DlgInclude = 0x11, + PlugAndPlay = 0x13, + VXD = 0x14, + AnimatedCursor = 0x15, + AnimatedIcon = 0x16, + HTML = 0x17, + Manifest = 0x18, +}; + +struct IMAGE_RESOURCE_DIRECTORY_ENTRY { + union { + struct { + DWORD NameOffset:31; + DWORD NameIsString:1; + }; + DWORD Name; + WORD Id; + }; + union { + DWORD OffsetToData; + struct { + DWORD OffsetToDirectory:31; + DWORD DataIsDirectory:1; + }; + }; +} + +struct IMAGE_RESOURCE_DIRECTORY { + uint32 Characteristics; + uint32 TimeDateStamp; + ushort MajorVersion; + ushort MinorVersion; + ushort NumberOfNamedEntries; + ushort NumberOfIdEntries; +}; + +/* +struct IMAGE_RESOURCE_DIRECTORY_ENTRY { + uint32 Name; + uint32 OffsetToDirectory:31; + uint32 DataIsDirectory:1; +}; +*/ + +typedef struct IMAGE_RESOURCE_DATA_ENTRY { + uint32 OffsetToData; + uint32 Size; + uint32 CodePage; + uint32 Reserved; +}; + +// --- END OF RESOURCE DIRECTORY + +// --- DEBUG DIRECTORY + +typedef struct IMAGE_DEBUG_DIRECTORY { + DWORD Characteristics; + DWORD TimeDateStamp; + WORD MajorVersion; + WORD MinorVersion; + DWORD Type; + DWORD SizeOfData; + DWORD AddressOfRawData; + DWORD PointerToRawData; +}; + +// --- END OF DEBUG DIRECTORY + +// --- RELOCATION DIRECTORY + +typedef struct _IMAGE_BASE_RELOCATION { + DWORD VirtualAddress; + DWORD SizeOfBlock; +// WORD TypeOffset[1]; +} IMAGE_BASE_RELOCATION; + +// --- END OF RELOCATION DIRECTORY + +// --- TLS DIRECTORY + +typedef struct _IMAGE_TLS_DIRECTORY32 { + DWORD StartAddressOfRawData; + DWORD EndAddressOfRawData; + DWORD AddressOfIndex; // PDWORD + DWORD AddressOfCallBacks; // PIMAGE_TLS_CALLBACK * + DWORD SizeOfZeroFill; + DWORD Characteristics; +} IMAGE_TLS_DIRECTORY32; + +typedef struct _IMAGE_TLS_DIRECTORY64 { + ULONGLONG StartAddressOfRawData; + ULONGLONG EndAddressOfRawData; + ULONGLONG AddressOfIndex; + ULONGLONG AddressOfCallBacks; + DWORD SizeOfZeroFill; + DWORD Characteristics; +} IMAGE_TLS_DIRECTORY64; + +// --- END OF TLS DIRECTORY +""" + + +pestruct = cstruct() +pestruct.load(pe_def) + + +cv_info_def = """ +struct GUID { + DWORD Data1; + WORD Data2; + WORD Data3; + char Data4[8]; +}; + +struct CV_INFO_PDB70 { + DWORD CvSignature; + GUID Signature; // unique identifier + DWORD Age; // an always-incrementing value + char PdbFileName[]; // zero terminated string with the name of the PDB file +}; +""" + +cv_info_struct = cstruct() +cv_info_struct.load(cv_info_def) diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py new file mode 100644 index 0000000..f7e4715 --- /dev/null +++ b/dissect/executable/pe/helpers/exports.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from collections import OrderedDict +from io import BytesIO +from typing import TYPE_CHECKING + +# Local imports +from dissect.executable.pe.helpers.c_pe import pestruct + +if TYPE_CHECKING: + from dissect.executable.pe.helpers.sections import PESection + from dissect.executable.pe.pe import PE + + +class ExportFunction: + """Object to store the information belonging to export functions. + + Args: + ordinal: The ordinal of the export function. + address: The export function address. + name: The name of the function, if available. + """ + + def __init__(self, ordinal: int, address: int, name: bytes = b""): + self.ordinal = ordinal + self.address = address + self.name = name + + def __str__(self) -> str: + return self.name.decode() if self.name else self.ordinal + + def __repr__(self) -> str: + return f"" if self.name else f"" + + +class ExportManager: + def __init__(self, pe: PE, section: PESection): + self.pe = pe + self.section = section + self.exports = OrderedDict() + + self.parse_exports() + + def parse_exports(self): + """Parse the export directory of the PE file. + + This function will store every export function within the PE file as an `ExportFunction` object containing the + name (if available), the call ordinal, and the function address. + """ + + export_entry_va = self.pe.directory_va(pestruct.IMAGE_DIRECTORY_ENTRY_EXPORT) + export_entry = BytesIO(self.pe.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_EXPORT)) + export_directory = pestruct.IMAGE_EXPORT_DIRECTORY(export_entry) + + # Seek to the offset of the export name + export_entry.seek(export_directory.Name - export_entry_va) + self.export_name = pestruct.char[None](export_entry) + + # Create a list of adresses for the exported functions + export_entry.seek(export_directory.AddressOfFunctions - export_entry_va) + export_addresses = pestruct.uint32[export_directory.NumberOfFunctions].read(export_entry) + # Create a list of addresses for the exported functions that have associated names + export_entry.seek(export_directory.AddressOfNames - export_entry_va) + export_names = pestruct.uint32[export_directory.NumberOfNames].read(export_entry) + # Create a list of addresses for the ordinals associated with the functions + export_entry.seek(export_directory.AddressOfNameOrdinals - export_entry_va) + export_ordinals = pestruct.uint16[export_directory.NumberOfNames].read(export_entry) + + # Iterate over the export functions and store the information + export_entry.seek(export_directory.AddressOfFunctions - export_entry_va) + for idx, address in enumerate(export_addresses): + if idx in export_ordinals: + export_entry.seek(export_names[export_ordinals.index(idx)] - export_entry_va) + export_name = pestruct.char[None](export_entry) + self.exports[export_name.decode()] = ExportFunction(ordinal=idx + 1, address=address, name=export_name) + else: + export_name = None + self.exports[str(idx + 1)] = ExportFunction(ordinal=idx + 1, address=address, name=export_name) + + def add(self): + raise NotImplementedError + + def delete(self): + raise NotImplementedError diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py new file mode 100644 index 0000000..1880644 --- /dev/null +++ b/dissect/executable/pe/helpers/imports.py @@ -0,0 +1,326 @@ +from __future__ import annotations + +import struct +from collections import OrderedDict +from io import BytesIO +from typing import TYPE_CHECKING, BinaryIO, Generator + +from dissect.executable.pe.helpers import utils + +# Local imports +from dissect.executable.pe.helpers.c_pe import pestruct + +if TYPE_CHECKING: + from dissect.cstruct.cstruct import cstruct + + from dissect.executable.pe.helpers.sections import PESection + from dissect.executable.pe.pe import PE + + +class ImportModule: + """Base class for the import modules, these hold their respective functions. + + Args: + name: The name of the module. + import_descriptor: The import descriptor of the module as a cstruct object. + module_va: The virtual address of the module. + name_va: The virtual address of the name of the module. + first_thunk: The virtual address of the first thunk. + """ + + def __init__(self, name: bytes, import_descriptor: cstruct, module_va: int, name_va: int, first_thunk: int): + self.name = name + self.import_descriptor = import_descriptor + self.module_va = module_va + self.name_va = name_va + self.first_thunk = first_thunk + self.functions = [] + + def __str__(self) -> str: + return self.name.decode() + + def __repr__(self) -> str: + return f"" # noqa: E501 + + +class ImportFunction: + """Base class for the import functions. + + Args: + pe: A `PE` object. + thunkdata: The thunkdata of the import function as a cstruct object. + """ + + def __init__(self, pe: PE, thunkdata: cstruct, name: str = ""): + self.pe = pe + self.thunkdata = thunkdata + self._name = name + + @property + def name(self) -> str: + """Return the name of the import function if available, otherwise return the ordinal of the function. + + Returns: + The name or ordinal of the import function. + """ + + if self._name: + return self._name + + ordinal = self.thunkdata.u1.AddressOfData & self.pe._high_bit + + if not ordinal: + self.pe.seek(self.thunkdata.u1.AddressOfData + 2) + entry = pestruct.char[None](self.pe).decode() + else: + entry = ordinal + + if isinstance(entry, int): + return str(entry) + + return entry + + def __str__(self) -> str: + return self.name + + def __repr__(self) -> str: + return f"" + + +class ImportManager: + """The base class for dealing with the imports that are present within the PE file. + + Args: + pe: A `PE` object. + section: The associated `PESection` object. + """ + + def __init__(self, pe: PE, section: PESection): + self.pe = pe + self.section = section + self.import_directory_rva = 0 + self.import_data = bytearray() + self.new_size_of_image = 0 + self.section_data = bytearray() + self.imports = OrderedDict() + self.thunks = [] + + self.parse_imports() + + def parse_imports(self): + """Parse the imports of the PE file. + + The imports are in turn added to the `imports` attribute so they can be accessed by the user. + """ + + import_data = BytesIO(self.pe.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_IMPORT)) + import_data.seek(0) + + # Loop over the entries + for descriptor_va, import_descriptor in self.import_descriptors(import_data=import_data): + if import_descriptor.Name != 0xFFFFF800 and import_descriptor.Name != 0x0: + self.pe.seek(import_descriptor.Name) + modulename = pestruct.char[None](self.pe) + + # Use the OriginalFirstThunk if available, FirstThunk otherwise + first_thunk = ( + import_descriptor.FirstThunk + if not import_descriptor.OriginalFirstThunk + else import_descriptor.OriginalFirstThunk + ) + module = ImportModule( + name=modulename, + import_descriptor=import_descriptor, + module_va=descriptor_va, + name_va=import_descriptor.Name, + first_thunk=first_thunk, + ) + + for thunkdata in self.parse_thunks(offset=first_thunk): + module.functions.append(ImportFunction(pe=self.pe, thunkdata=thunkdata)) + + self.imports[modulename.decode()] = module + + def import_descriptors(self, import_data: BinaryIO) -> Generator[tuple[int, cstruct], None, None]: + """Parse the import descriptors of the PE file. + + Args: + import_data: The data within the import directory. + + Yields: + The import descriptor as a `cstruct` object. + """ + + while True: + try: + import_descriptor = pestruct.IMAGE_IMPORT_DESCRIPTOR(import_data) + except EOFError: + break + + yield import_data.tell(), import_descriptor + + def parse_thunks(self, offset: int) -> Generator[cstruct, None, None]: + """Parse the import thunks for every module. + + Args: + offset: The offset to the first thunk + + Yields: + The function name or ordinal + """ + + self.pe.seek(offset) + + while True: + thunkdata = self.pe.image_thunk_data(self.pe) + if not thunkdata.u1.Function: + break + + yield thunkdata + + def add(self, dllname: str, functions: list): + """Add the given module and its functions to the PE. + + Args: + dllname: The name of the module to add. + functions: A `list` of function names belonging to the module. + """ + + self.last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] + + # Build a dummy import module + self.imports[dllname] = ImportModule( + name=dllname.encode(), import_descriptor=None, module_va=0, name_va=0, first_thunk=0 + ) + # Build the dummy module functions + for function in functions: + self.pe.imports[dllname].functions.append(ImportFunction(pe=self.pe, thunkdata=None, name=function)) + + # Rebuild the import table with the new import module and functions + self.build_import_table() + + def delete(self, dllname: str, functions: list): + raise NotImplementedError + + def build_import_table(self): + """Function to rebuild the import table after a change has been made to the PE imports. + + Currently we're using the .idata section to store the imports, there might be a better way to do this but for + now this will do. + """ + + # Reset the known thunkdata + self.thunks = [] + + import_descriptors = [] + self.import_data = bytearray() + + for name, module in self.imports.items(): + # Take note of the current offset to store the modulename + name_offset = len(self.import_data) + self.import_data += name.encode() + b"\x00" + + # Build the module imports and get the RVA of the first thunk to generate an import descriptor + first_thunk_rva = self._build_module_imports(functions=module.functions) + import_descriptor = self._build_import_descriptor( + first_thunk_rva=first_thunk_rva, name_rva=self.pe.optional_header.SizeOfImage + name_offset + ) + import_descriptors.append(import_descriptor) + + datadirectory_size = 0 + for idx, descriptor in enumerate(import_descriptors): + if idx == 0: + # Take note of the RVA of the first import descriptor + import_rva = self.pe.optional_header.SizeOfImage + len(self.import_data) + self.import_data += descriptor.dumps() + datadirectory_size += len(descriptor) + + # Create a new section + section_data = utils.align_data(data=self.import_data, blocksize=self.pe.file_alignment) + size = len(self.import_data) + pestruct.IMAGE_SECTION_HEADER.size + self.pe.add_section( + name=".idata", + data=section_data, + datadirectory=pestruct.IMAGE_DIRECTORY_ENTRY_IMPORT, + datadirectory_rva=import_rva, + datadirectory_size=datadirectory_size, + size=size, + ) + + def _build_module_imports(self, functions: list[ImportFunction]) -> int: + """Function to build the imports for a module. + + This function is responsible for building the functions by name, as well as the associated thunkdata that is + used to parse the imports at a later stage. + + Args: + functions: A `list` of `ImportFunction` objects. + + Returns: + The relative virtual address of the first thunk. + """ + + function_offsets = [] + + for idx, function in enumerate(functions): + function_offsets.append(len(self.import_data)) + self.import_data += struct.pack(" bytes: + """Function to build the thunkdata for the new import table. + + Args: + import_rvas: A `list` of relative virtual addresses. + + Returns: + The thunkdata as a `bytes` object. + """ + + thunkdata = bytearray() + for rva in import_rvas: + rva += self.pe.optional_header.SizeOfImage + thunkdata += ( + struct.pack(" cstruct: + """Function to build the import descriptor for the new import table. + + Args: + first_thunk_rva: The relative address of the first piece of thunkdata. + + Returns: + The image import descriptor as a `cstruct` object. + """ + + new_import_descriptor = pestruct.IMAGE_IMPORT_DESCRIPTOR() + + new_import_descriptor.OriginalFirstThunk = first_thunk_rva + new_import_descriptor.TimeDateStamp = 0 + new_import_descriptor.ForwarderChain = 0 + new_import_descriptor.Name = name_rva + new_import_descriptor.FirstThunk = first_thunk_rva + + return new_import_descriptor diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py new file mode 100644 index 0000000..5bdc883 --- /dev/null +++ b/dissect/executable/pe/helpers/patcher.py @@ -0,0 +1,344 @@ +from __future__ import annotations + +import struct +from io import BytesIO +from typing import TYPE_CHECKING + +# Local imports +from dissect.executable.pe.helpers import utils +from dissect.executable.pe.helpers.c_pe import pestruct +from dissect.executable.pe.helpers.sections import PESection + +if TYPE_CHECKING: + from dissect.executable import PE + + +class Patcher: + """Class that is used to patch existing PE files with the changes made by the user. + + Args: + pe: A `PE` object that contains the original PE file. + """ + + def __init__(self, pe: PE): + self.pe = pe + self.patched_pe = BytesIO() + self.functions = [] + + @property + def build(self) -> BytesIO: + """Build the patched PE file. + + This function will return a new PE file as a `BytesIO` object that contains the new PE file. + + Returns: + The patched PE file as a `BytesIO` object. + """ + + # Update the SizeOfImage + self.pe.optional_header.SizeOfImage = self.pe_size + + self.patched_pe.seek(0) + + # Build the section table and add the sections + self._build_section_table() + + # Apply the patches + self._patch_rvas() + + # Add the MZ, File and NT headers + self.patched_pe.seek(0) + self._build_dos_header() + + # Reset the file pointer + self.patched_pe.seek(0) + return self.patched_pe + + @property + def pe_size(self) -> int: + """Calculate the new PE size. + + We can calculate the new size of the PE by looking at the ending of the last section. + + Returns: + The new size of the PE as an `int`. + """ + + last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] + va = last_section.virtual_address + size = last_section.virtual_size + + return utils.align_int(integer=va + size, blocksize=self.pe.optional_header.SectionAlignment) + + def seek(self, address: int): + """Seek that is used to seek to a virtual address in the patched PE file. + + Args: + address: The virtual address to seek to. + """ + + raw_address = self.pe.virtual_address(address=address) + self.patched_pe.seek(raw_address) + + def _build_section_table(self): + """Function to build the section table and add the sections with their data.""" + + if self.patched_pe.tell() < self.pe.section_header_offset: + # Pad the patched file with null bytes until we reach the section header offset + self.patched_pe.write(utils.pad(size=self.pe.section_header_offset - self.patched_pe.tell())) + + # Write the section headers + for section in self.pe.patched_sections.values(): + self.patched_pe.write(section.dump()) + + # Add the data for each section + for section in self.pe.patched_sections.values(): + self.patched_pe.seek(section.pointer_to_raw_data) + self.patched_pe.write(section.data) + + def _build_dos_header(self): + """Function to build the DOS header, NT headers and the DOS stub.""" + + # Add the MZ + self.patched_pe.write(self.pe.mz_header.dumps()) + + # Add the DOS stub + stub_size = self.pe.mz_header.e_lfanew - self.patched_pe.tell() + dos_stub = self.pe.raw_read(offset=self.patched_pe.tell(), size=stub_size) + self.patched_pe.write(dos_stub) + + # Add the NT headers + self.patched_pe.seek(self.pe.mz_header.e_lfanew) + self.patched_pe.write(b"PE\x00\x00") + self.patched_pe.write(self.pe.file_header.dumps()) + self.patched_pe.write(self.pe.optional_header.dumps()) + + def _patch_rvas(self): + """Function to call the different patch functions responsible for patching any kind of relative addressing.""" + + self._patch_import_rvas() + self._patch_export_rvas() + self._patch_rsrc_rvas() + self._patch_tls_rvas() + + def _patch_import_rvas(self): + """Function to patch the RVAs of the import directory and the thunkdata entries.""" + + patched_import_data = bytearray() + + # Get the directory entry virtual adddress, this is the updated address if it has been patched. + directory_va = self.pe.directory_va(pestruct.IMAGE_DIRECTORY_ENTRY_IMPORT) + if not directory_va: + return + + # Get the original VA of the section the import directory is residing in, this value is used to calculate the + # new RVA's + section = self.pe.patched_section(va=directory_va) + directory_offset = directory_va - section.virtual_address + original_directory_va = self.pe.sections[section.name].virtual_address + directory_offset + + # Loop over the imports of the PE to patch the RVA's of the import descriptors and the associated thunkdata + # entries + for name, module in self.pe.imports.items(): + import_descriptor = module.import_descriptor + patched_thunkdata = bytearray() + + if import_descriptor.Name != 0xFFFFF800 and import_descriptor.Name != 0x0: + old_first_thunk = import_descriptor.FirstThunk + + first_thunk_offset = old_first_thunk - original_directory_va + import_descriptor.FirstThunk = abs(directory_va + first_thunk_offset) + + import_descriptor.OriginalFirstThunk = import_descriptor.FirstThunk + + name_offset = import_descriptor.Name - original_directory_va + import_descriptor.Name = abs(directory_va + name_offset) + + for function in module.functions: + thunkdata = function.thunkdata + # Check if we're dealing with an ordinal entry, if it's an ordinal entry we don't need + # to patch since it's not an RVA + if thunkdata.u1.AddressOfData & self.pe._high_bit: + patched_thunkdata += thunkdata.dumps() + continue + + # Check the original RVA associated with the AddressOfData field in the thunkdata, retrieve the + # original VA + # and use it to also select the patched virtual address of this section that the RVA is located in + for name, section in self.pe.sections.items(): + if thunkdata.u1.AddressOfData in range( + section.virtual_address, section.virtual_address + section.virtual_size + ): + virtual_address = section.virtual_address + new_virtual_address = self.pe.patched_sections[name].virtual_address + break + + # Calculate the offset using the VA of the section and update the thunkdata + va_offset = thunkdata.u1.AddressOfData - virtual_address + new_thunkdata = new_virtual_address + va_offset + thunkdata.u1.AddressOfData = new_thunkdata + thunkdata.u1.ForwarderString = new_thunkdata + thunkdata.u1.Function = new_thunkdata + thunkdata.u1.Ordinal = new_thunkdata + + patched_thunkdata += thunkdata.dumps() + + # Write the thunk data into the patched PE + self.seek(import_descriptor.FirstThunk) + self.patched_pe.write(patched_thunkdata) + + patched_import_data += import_descriptor.dumps() + + self.seek(directory_va) + self.patched_pe.write(patched_import_data) + + def _patch_export_rvas(self): + """Function to patch the RVAs of the export directory and the associated function and name RVA's.""" + + directory_va = self.pe.directory_va(pestruct.IMAGE_DIRECTORY_ENTRY_EXPORT) + if not directory_va: + return + + self.seek(directory_va) + export_directory = pestruct.IMAGE_EXPORT_DIRECTORY(self.patched_pe) + + # Get the original VA of the section the import directory is residing in, this value is used to calculate the + # new RVA's + section = self.pe.patched_section(va=directory_va) + directory_offset = directory_va - section.virtual_address + original_directory_va = self.pe.sections[section.name].virtual_address + directory_offset + + name_offset = export_directory.Name - original_directory_va + address_of_functions_offset = export_directory.AddressOfFunctions - original_directory_va + address_of_names_offset = export_directory.AddressOfNames - original_directory_va + address_of_name_ordinals = export_directory.AddressOfNameOrdinals - original_directory_va + + export_directory.Name = directory_va + name_offset + export_directory.AddressOfFunctions = directory_va + address_of_functions_offset + export_directory.AddressOfNames = directory_va + address_of_names_offset + export_directory.AddressOfNameOrdinals = directory_va + address_of_name_ordinals + + # Write the new export directory + self.seek(directory_va) + self.patched_pe.write(export_directory.dumps()) + + # Patch the addresses of the functions + new_function_rvas = [] + function_rvas = bytearray() + self.seek(export_directory.AddressOfFunctions) + export_addresses = pestruct.uint32[export_directory.NumberOfFunctions].read(self.patched_pe) + for address in export_addresses: + section = self.pe.section(va=address) + if not section: + continue + address_offset = address - section.virtual_address + new_address = self.pe.patched_sections[section.name].virtual_address + address_offset + new_function_rvas.append(new_address) + + for rva in new_function_rvas: + function_rvas += struct.pack(" PESection: + """Function to get the section that contains the TLS attribute. + + Args: + va: The virtual address of the TLS attribute. + + Returns: + The section that contains the TLS attribute as a `PESection` object. + """ + + for name, section in self.pe.sections.items(): + if va in range(section.virtual_address, section.virtual_address + section.virtual_size): + return section diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/helpers/relocations.py new file mode 100644 index 0000000..465a66b --- /dev/null +++ b/dissect/executable/pe/helpers/relocations.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from io import BytesIO +from typing import TYPE_CHECKING + +# Local imports +from dissect.executable.pe.helpers.c_pe import pestruct + +if TYPE_CHECKING: + from dissect.executable.pe.helpers.sections import PESection + from dissect.executable.pe.pe import PE + + +class RelocationManager: + """Base class for dealing with the relocations within the PE file. + + Args: + pe: The PE file object. + section: The section object that contains the relocation table. + """ + + def __init__(self, pe: PE, section: PESection): + self.pe = pe + self.section = section + self.relocations = [] + + self.parse_relocations() + + def parse_relocations(self): + """Parse the relocation table of the PE file.""" + + reloc_data = BytesIO(self.pe.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_BASERELOC)) + reloc_data_size = reloc_data.getbuffer().nbytes + while reloc_data.tell() < reloc_data_size: + reloc_directory = pestruct.IMAGE_BASE_RELOCATION(reloc_data) + if not reloc_directory.VirtualAddress: + # End of relocation entries + break + + # Each entry consists of 2 bytes + number_of_entries = (reloc_directory.SizeOfBlock - len(reloc_directory.dumps())) // 2 + entries = [] + for _ in range(0, number_of_entries): + entry = pestruct.uint16(reloc_data) + if entry: + entries.append(entry) + + self.relocations.append( + {"rva:": reloc_directory.VirtualAddress, "number_of_entries": number_of_entries, "entries": entries} + ) + + def add(self): + raise NotImplementedError diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py new file mode 100644 index 0000000..ee011fb --- /dev/null +++ b/dissect/executable/pe/helpers/resources.py @@ -0,0 +1,418 @@ +from __future__ import annotations + +from collections import OrderedDict +from io import BytesIO +from typing import TYPE_CHECKING, Iterator + +from dissect.executable.exception import ResourceException + +# Local imports +from dissect.executable.pe.helpers.c_pe import pestruct + +if TYPE_CHECKING: + from dissect.cstruct.cstruct import cstruct + from dissect.cstruct.types.enum import EnumInstance + + from dissect.executable.pe.helpers.sections import PESection + from dissect.executable.pe.pe import PE + + +class ResourceManager: + """Base class to perform actions regarding the resources within the PE file. + + Args: + pe: A `PE` object. + section: The section object that contains the resource table. + """ + + def __init__(self, pe: PE, section: PESection): + self.pe = pe + self.section = section + self.resources = OrderedDict() + self.raw_resources = [] + + self.parse_rsrc() + + def parse_rsrc(self): + """Parse the resource directory entry of the PE file.""" + + rsrc_data = BytesIO(self.pe.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_RESOURCE)) + self.resources = self._read_resource(rc_type="_root", data=rsrc_data, offset=0, level=1) + + def _read_entries(self, data: bytes, directory: cstruct) -> list: + """Read the entries within the resource directory. + + Args: + data: The data of the resource directory. + directory: The resource directory entry. + + Returns: + A list containing the entries of the resource directory. + """ + + entries = [] + for _ in range(0, directory.NumberOfNamedEntries + directory.NumberOfIdEntries): + entry_offset = data.tell() + entry = pestruct.IMAGE_RESOURCE_DIRECTORY_ENTRY(data) + self.raw_resources.append({"offset": entry_offset, "entry": entry, "data_offset": entry_offset}) + entries.append(entry) + return entries + + def _handle_data_entry(self, data: bytes, entry: cstruct, rc_type: str) -> Resource: + """Handle the data entry of a resource. This is the actual data associated with the directory entry. + + Args: + data: The data of the resource. + entry: The resource directory entry. + + Returns: + The resource that was given by name as a `Resource` object. + """ + + data.seek(entry.OffsetToDirectory) + data_entry = pestruct.IMAGE_RESOURCE_DATA_ENTRY(data) + self.pe.seek(data_entry.OffsetToData) + data = self.pe.read(data_entry.Size) + raw_offset = data_entry.OffsetToData - self.section.virtual_address + rsrc = Resource( + pe=self.pe, + section=self.section, + name=entry.Name, + entry_offset=entry.OffsetToData, + data_entry=data_entry, + rc_type=rc_type, + ) + self.raw_resources.append( + { + "offset": entry.OffsetToDirectory, + "entry": data_entry, + "data": data, + "data_offset": raw_offset, + "resource": rsrc, + } + ) + return rsrc + + def _read_resource(self, data: bytes, offset: int, rc_type: str, level: int = 1) -> dict: + """Recursively read the resources within the PE file. + + Each resource is added to the dictionary that is available to the user, as well as a list of + raw resources that are used to update the section data and size when a resource has been modified. + + Args: + data: The data of the resource. + offset: The offset of the resource. + rc_type: The type of the resource. + level: The depth level of the resource, this dictates the resource type. + + Returns: + A dictionary containing the resources that were found. + """ + + resource = OrderedDict() + + data.seek(offset) + directory = pestruct.IMAGE_RESOURCE_DIRECTORY(data) + self.raw_resources.append({"offset": offset, "entry": directory, "data_offset": offset}) + + entries = self._read_entries(data, directory) + + for entry in entries: + if level == 1: + rc_type = pestruct.ResourceID(entry.Id).name + else: + if entry.NameIsString: + data.seek(entry.NameOffset) + name_len = pestruct.uint16(data) + rc_type = pestruct.wchar[name_len](data) + else: + rc_type = str(entry.Id) + + if entry.DataIsDirectory: + resource[rc_type] = self._read_resource( + data=data, offset=entry.OffsetToDirectory, rc_type=rc_type, level=level + 1 + ) + else: + resource[rc_type] = self._handle_data_entry(data=data, entry=entry, rc_type=rc_type) + + return resource + + def get_resource(self, name: str) -> Resource: + """Retrieve the resource by name. + + Args: + name: The name of the resource to retrieve. + + Returns: + The resource that was given by name as a `Resource` object. + """ + + try: + return self.resources[name] + except KeyError: + raise ResourceException(f"Resource {name} not found!") + + def get_resource_type(self, rsrc_id: str | EnumInstance): + """Yields a generator containing all of the nodes within the resources that contain the requested ID. + + The ID can be either given by name or its value. + + Args: + rsrc_id: The resource ID to find, this can be a cstruct `EnumInstance` or `str`. + + Yields: + All of the nodes that contain the requested type. + """ + + if rsrc_id not in self.resources: + raise ResourceException(f"Resource with ID {rsrc_id} not found in PE!") + + for resource in self.parse_resources(resources=self.resources[rsrc_id]): + yield resource + + def parse_resources(self, resources: dict) -> Iterator[Resource]: + """Parse the resources within the PE file. + + Args: + resources: A `dict` containing the different resources that were found. + + Yields: + All of the resources within the PE file. + """ + + for _, resource in resources.items(): + if type(resource) is not OrderedDict: + yield resource + else: + yield from self.parse_resources(resources=resource) + + def show_resource_tree(self, resources: dict, indent: int = 0): + """Print the resources within the PE as a tree. + + Args: + resources: A `dict` containing the different resources that were found. + indent: The amount of indentation for each child resource. + """ + + for name, resource in resources.items(): + if type(resource) is not OrderedDict: + print(f"{' ' * indent} - name: {name} ID: {resource.rsrc_id}") + else: + print(f"{' ' * indent} + name: {name}") + self.show_resource_tree(resources=resource, indent=indent + 1) + + def show_resource_info(self, resources: dict): + """Print basic information about the resource as well as the header. + + Args: + resources: A `dict` containing the different resources that were found. + """ + + for name, resource in resources.items(): + if type(resource) is not OrderedDict: + print( + f"* resource: {name} offset=0x{resource.offset:02x} size=0x{resource.size:02x} header: {resource.data[:64]}" # noqa: E501 + ) + else: + self.show_resource_info(resources=resource) + + def add_resource(self, name: str, data: bytes): + # TODO + raise NotImplementedError + + def delete_resource(self, name: str): + # TODO + raise NotImplementedError + + def update_section(self, update_offset: int): + """Function to dynamically update the section data and size when a resource has been modified. + + Args: + update_offset: The offset of the resource that was modified. + """ + + new_size = 0 + section_data = b"" + + for idx, resource in enumerate(self.parse_resources(resources=self.pe.resources)): + if idx == 0: + # Use the offset of the first resource to account for the size of the directory header + header_size = resource.offset - self.section.virtual_address + section_data = self.section.data[:header_size] + + # Take note of the previous offset and size so we can update any of these values after changing the data + # within the resource + prev_offset = resource.offset + prev_size = resource.size + + # Update the resource data + section_data += resource.data + + new_size += resource.size + 1 # Account for the id field + + # Skip the resources that are below our own offset + if update_offset >= resource.offset: + continue + + offset = prev_offset + prev_size + 2 + resource.offset = offset + + # Add the header to the total size so we can check if we need to update the section size + new_size += header_size + + # Update the section + self.section.data = section_data + + +class Resource(ResourceManager): + """Base class representing a resource entry in the PE file. + + Args: + pe: A `PE` object. + section: The section object that contains the resource table. + name: The name of the resource. + entry_offset: The offset of the resource entry. + data_entry: The data entry of the resource. + rc_type: The type of the resource. + data: The data of the resource if there was data provided by the user. + """ + + def __init__( + self, + pe: PE, + section: PESection, + name: str | int, + entry_offset: int, + data_entry: cstruct, + rc_type: str, + data: bytes = b"", + ): + self.pe = pe + self.section = section + self.name = name + self.entry_offset = entry_offset + self.entry = data_entry + self.rc_type = rc_type + self.offset = data_entry.OffsetToData + self._size = data_entry.Size + self.codepage = data_entry.CodePage + self._data = self.read_data() if not data else data + + def read_data(self) -> bytes: + """Read the data within the resource. + + Returns: + The resource data. + """ + + return self.pe.virtual_read(address=self.offset, size=self._size) + + @property + def size(self) -> int: + """Function to return the size of the resource. + This needs to be done dynamically in the case that the data is patched by the user. + + Returns: + The size of the data within the resource. + """ + + return len(self.data) + + @size.setter + def size(self, value: int) -> int: + """Setter to set the size of the resource to the specified value. + + Args: + value: The size of the resource. + """ + + self._size = value + self.entry.Size = value + + @property + def offset(self) -> int: + """Return the offset of the resource.""" + return self.entry.OffsetToData + + @offset.setter + def offset(self, value: int): + """Setter to set the offset of the resource to the specified value. + + Args: + value: The offset of the resource. + """ + + self.entry.OffsetToData = value + + @property + def data(self) -> bytes: + """Return the data within the resource.""" + return self._data + + @data.setter + def data(self, value: bytes): + """Setter to set the new data of the resource, but also dynamically update the offset of the resources within + the same directory. + + This function currently also updates the section sizes and alignment. Ideally this would be moved to a more + abstract function that + can handle tasks like these in a more transparant manner. + + Args: + value: The new data of the resource. + """ + + # Set the new data + self._data = value + + if len(value) != self.entry.Size: + self.size = len(value) + + section_data = BytesIO() + + prev_offset = 0 + prev_size = 0 + + for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"]): + entry_offset = rsrc_entry["offset"] + entry = rsrc_entry["entry"] + + if entry._type.name == "IMAGE_RESOURCE_DATA_ENTRY": + rsrc_obj = rsrc_entry["resource"] + data_offset = rsrc_entry["data_offset"] + + # Normally the data is separated by a null byte, increment the new offset by 1 + new_data_offset = prev_offset + prev_size + # if new_data_offset and (new_data_offset > data_offset or new_data_offset < data_offset): + if new_data_offset and new_data_offset != data_offset: + data_offset = new_data_offset + rsrc_entry["data_offset"] = data_offset + rsrc_obj.offset = self.section.virtual_address + data_offset + + data = rsrc_obj.data + + # Write the resource entry data into the section + section_data.seek(data_offset) + section_data.write(data) + + # Take note of the offset and size so we can update any of these values after changing the data within + # the resource + prev_offset = data_offset + prev_size = rsrc_obj.size + + # Write the resource entry into the section + section_data.seek(entry_offset) + section_data.write(entry.dumps()) + + section_data.seek(0) + data = section_data.read() + + # Update the section data and size + self.section.data = data + self.pe.optional_header.DataDirectory[pestruct.IMAGE_DIRECTORY_ENTRY_RESOURCE].Size = len(data) + + def __str__(self) -> str: + return str(self.name) + + def __repr__(self) -> str: + return f"" # noqa: E501 diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py new file mode 100644 index 0000000..792ad82 --- /dev/null +++ b/dissect/executable/pe/helpers/sections.py @@ -0,0 +1,258 @@ +from __future__ import annotations + +from collections import OrderedDict +from typing import TYPE_CHECKING + +from dissect.executable.exception import BuildSectionException + +# Local imports +from dissect.executable.pe.helpers import utils +from dissect.executable.pe.helpers.c_pe import pestruct + +if TYPE_CHECKING: + from dissect.cstruct.cstruct import cstruct + + from dissect.executable.pe.pe import PE + + +class PESection: + """Base class for the PE sections that are present. + + Args: + pe: A `PE` object. + section: A `cstruct` definition holding the information about the section. + offset: The offset of the section within the PE file. + data: The data that should be part of the section, this can be used to add new sections. + """ + + def __init__(self, pe: PE, section: cstruct, offset: int, data: bytes = b""): + self.pe = pe + self.section = section + self.offset = offset + self.name = section.Name.decode().rstrip("\x00") + self._virtual_address = section.VirtualAddress + self._virtual_size = section.VirtualSize + self._pointer_to_raw_data = section.PointerToRawData + self._size_of_raw_data = section.SizeOfRawData + + # Keep track of the directories that are within this section + self.directories = OrderedDict() + + self._data = self.read_data() if not data else data + + def read_data(self) -> bytes: + """Return the data within the section. + + Returns: + The `bytes` contained within the section. + """ + + if self.pe.virtual: + return self.pe.virtual_read(self.virtual_address, self.virtual_size) + + return self.pe.raw_read(self.pointer_to_raw_data, self.size_of_raw_data) + + @property + def size(self) -> int: + """Return the size of the data within the section.""" + return self.virtual_size + + @size.setter + def size(self, value: int): + """Setter to set the size of the data to the specified value. + + This function can be used to update the size of the data, but also dynamically update the offset of the data + within the same directory. + + Args: + value: The size of the data. + """ + + self.virtual_size = value + self.size_of_raw_data = utils.align_int(integer=value, blocksize=self.pe.file_alignment) + + @property + def virtual_address(self) -> int: + """Return the virtual address of the section.""" + return self._virtual_address + + @virtual_address.setter + def virtual_address(self, value: int): + """Setter to set the virtual address of the section to the specified value. + + This function also updates any of the virtual addresses of the directories that are residing within the section + itself. + + Args: + value: The virtual address of the section. + """ + + self._virtual_address = value + self.section.VirtualAddress = value + + # Update the VA of the directory residing within this section + for idx, offset in self.directories.items(): + directory_va = value + offset + self.pe.optional_header.DataDirectory[idx].VirtualAddress = directory_va + + @property + def virtual_size(self) -> int: + """Return the virtual size of the section.""" + return self._virtual_size + + @virtual_size.setter + def virtual_size(self, value: int): + """Setter to set the virtual size of the section to the specified value. + + Args: + value: The virtual size of the section. + """ + + self._virtual_size = value + self.section.VirtualSize = value + + @property + def pointer_to_raw_data(self) -> int: + """Return the pointer to the raw data within the section.""" + return self._pointer_to_raw_data + + @pointer_to_raw_data.setter + def pointer_to_raw_data(self, value: int): + """Setter to set the pointer to the raw data of the section to the specified value. + + Args: + value: The pointer to the raw data of the section. + """ + + self._pointer_to_raw_data = value + self.section.PointerToRawData = value + + @property + def size_of_raw_data(self) -> int: + """Return the size of the raw data within the section. This acounts for section alignment within the PE.""" + return self._size_of_raw_data + + @size_of_raw_data.setter + def size_of_raw_data(self, value: int): + """Setter to set the size of the raw data to the specified value. + + The SizeOfRawData field uses the section alignment to make sure the data within this section is aligned to the + section alignment. + + Args: + value: The size of the data. + """ + + self._size_of_raw_data = utils.align_int(integer=value, blocksize=self.pe.file_alignment) + self.section.SizeOfRawData = utils.align_int(integer=value, blocksize=self.pe.file_alignment) + + @property + def data(self) -> bytes: + """Return the data within the section.""" + return self._data[: self.virtual_size] + + @data.setter + def data(self, value: bytes): + """Setter to set the new data of the resource, but also dynamically update the offset of the resources within + the same directory. + + This function currently also updates the section sizes and alignment. Ideally this would be moved to a more + abstract function that can handle tasks like these in a more transparant manner. + + Args: + value: The new data of the resource. + """ + + # Keep track of the section changes using the patched_sections dictionary + self.pe.patched_sections[self.name]._data = value + self.pe.patched_sections[self.name].size = len(value) + + # Set the new data and size + self._data = value + self.size = len(value) + + # Pad the remainder of the section if the SizeOfRawData is smaller than the VirtualSize + if self.size_of_raw_data < self.virtual_size: + self._data += utils.pad(size=self.virtual_size - self.size_of_raw_data) + + # Take note of the first section as our starting point + first_section = next(iter(self.pe.patched_sections.values())) + + prev_ptr = first_section.pointer_to_raw_data + prev_size = first_section.size_of_raw_data + prev_va = first_section.virtual_address + prev_vsize = first_section.virtual_size + + for name, section in self.pe.patched_sections.items(): + if section.virtual_address == prev_va: + continue + + pointer_to_raw_data = utils.align_int(integer=prev_ptr + prev_size, blocksize=self.pe.file_alignment) + virtual_address = utils.align_int(integer=prev_va + prev_vsize, blocksize=self.pe.section_alignment) + + if section.virtual_address < virtual_address: + """Set the virtual address and raw pointer of the section to the new values, but only do so if the + section virtual address is lower than the previous section. We want to prevent messing up RVA's as + much as possible, this could lead to binaries that are a bit larger than they need to be but that + doesn't really matter.""" + self.pe.patched_sections[name].virtual_address = virtual_address + self.pe.patched_sections[name].pointer_to_raw_data = pointer_to_raw_data + + prev_ptr = pointer_to_raw_data + prev_size = section.size_of_raw_data + prev_va = virtual_address + prev_vsize = section.virtual_size + + def dump(self) -> bytes: + """Return the section header as a `bytes` object.""" + return self.section.dumps() + + def __str__(self) -> str: + return self.name + + def __repr__(self) -> str: + return f"" # noqa: E501 + + +def build_section( + virtual_size: int, + virtual_address: int, + raw_size: int, + pointer_to_raw_data: int, + name: str | bytes = b".dissect", + characteristics: int = 0xC0000040, +) -> cstruct: + """Build a new section for the PE. + + Args: + virtual_size: The virtual size of the new section data. + virtual_address: The virtual address where the new section is located. + raw_size: The size of the section data. + pointer_to_raw_data: The pointer to the raw data of the new section. + characteristics: The characteristics of the new section, default: 0xC0000040 + name: The new section name, default: .dissect + + Returns: + The new section header as a `cstruct` object. + """ + + if len(name) > 8: + raise BuildSectionException("section names can't be longer than 8 characters") + + if isinstance(name, str): + name = name.encode() + + section_header = pestruct.IMAGE_SECTION_HEADER() + + section_header.Name = name + utils.pad(size=8 - len(name)) + section_header.VirtualSize = virtual_size + section_header.VirtualAddress = virtual_address + section_header.SizeOfRawData = raw_size + section_header.PointerToRawData = pointer_to_raw_data + section_header.PointerToRelocations = 0 + section_header.PointerToLinenumbers = 0 + section_header.NumberOfRelocations = 0 + section_header.NumberOfLinenumbers = 0 + section_header.Characteristics = characteristics + + return section_header diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/helpers/tls.py new file mode 100644 index 0000000..3e3929d --- /dev/null +++ b/dissect/executable/pe/helpers/tls.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from io import BytesIO +from typing import TYPE_CHECKING + +# Local imports +from dissect.executable.pe.helpers.c_pe import pestruct + +if TYPE_CHECKING: + from dissect.executable.pe.helpers.sections import PESection + from dissect.executable.pe.pe import PE + + +class TLSManager: + """Base class to manage the TLS entries of a PE file. + + Args: + pe: The PE object to manage the TLS entries for. + """ + + def __init__(self, pe: PE, section: PESection): + self.pe = pe + self.section = section + self.callbacks = [] + self.tls = None + self._data = b"" + + self.parse_tls() + + def parse_tls(self): + """Parse the TLS directory entry of the PE file when present.""" + + tls_data = BytesIO(self.pe.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_TLS)) + self.tls = self.pe.image_tls_directory(tls_data) + + self.pe.seek(self.tls.AddressOfCallBacks - self.pe.optional_header.ImageBase) + + # Parse the TLS callback addresses if present + while True: + callback_address = self.pe.read_address(self.pe) + if not callback_address: + break + self.callbacks.append(callback_address) + + # Read the TLS data + self._data = self.read_data() + + @property + def size(self) -> int: + """Return the size of the TLS data. + + Returns: + The size of the TLS data in bytes. + """ + + return self.tls.EndAddressOfRawData - self.tls.StartAddressOfRawData + + @size.setter + def size(self, value): + """Setter to set the size of the TLS data to the specified value. + + Args: + value: The new size of the TLS data in bytes. + """ + + self.tls.EndAddressOfRawData = self.tls.StartAddressOfRawData + value + + def read_data(self) -> bytes: + """Read the TLS data from the PE file. + + Returns: + The TLS data in bytes. + """ + + return self.pe.virtual_read( + address=self.tls.StartAddressOfRawData - self.pe.optional_header.ImageBase, size=self.size + ) + + @property + def data(self) -> bytes: + """Return the TLS data. + + Returns: + The TLS data in bytes. + """ + + return self._data + + @data.setter + def data(self, value): + """Dynamically update the TLS directory data if the user changes the data. + + Args: + value: The new TLS data to write to the PE file. + """ + + self._data = value + section_data = BytesIO(self.section.data) + + if len(self._data) != self.size: + # Update the size of the TLS data + self.size = len(self._data) + + # Write the new TLS values to the section + section_data.write(self.tls.dumps()) + + # Write the new TLS data to the section + start_address_rva = self.tls.StartAddressOfRawData - self.pe.optional_header.ImageBase + start_address_section_offset = start_address_rva - self.section.virtual_address + section_data.seek(start_address_section_offset) + section_data.write(self._data) + + # Update the section itself + section_data.seek(0) + self.section.data = section_data.read() + + def add(self): + raise NotImplementedError diff --git a/dissect/executable/pe/helpers/utils.py b/dissect/executable/pe/helpers/utils.py new file mode 100644 index 0000000..93ce896 --- /dev/null +++ b/dissect/executable/pe/helpers/utils.py @@ -0,0 +1,40 @@ +def align_data(data: bytes, blocksize: int) -> bytes: + """Align the new data according to the file alignment as specified in the PE header. + + Args: + data: The raw data that needs to be aligned. + blocksize: The alignment to adhere to. + + Returns: + Padded data if the data was not aligned to the blocksize. + """ + + needs_alignment = len(data) % blocksize + return data if not needs_alignment else data + ((blocksize - needs_alignment) * b"\x00") + + +def align_int(integer: int, blocksize: int) -> int: + """Align integer values to the specified section alignment described in the PE header. + + Args: + integer: The address or value that needs to have an aligned value. + blocksize: The alignment to adhere to. + + Returns: + An aligned integer if the integer itself was not aligned yet. + """ + + needs_alignment = integer % blocksize + return integer if not needs_alignment else integer + (blocksize - needs_alignment) + + +def pad(size: int) -> bytes: + """Pad the data with null bytes. + + Args: + size: The amount of null bytes to return. + + Returns: + The null bytes as `bytes`. + """ + return size * b"\x00" diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py new file mode 100644 index 0000000..cbb19fb --- /dev/null +++ b/dissect/executable/pe/pe.py @@ -0,0 +1,602 @@ +from __future__ import annotations + +from collections import OrderedDict +from datetime import datetime +from io import BytesIO +from typing import TYPE_CHECKING, BinaryIO, Tuple + +from dissect.executable.exception import ( + InvalidAddress, + InvalidArchitecture, + InvalidPE, + InvalidVA, + ResourceException, +) +from dissect.executable.pe.helpers import ( + exports, + imports, + patcher, + relocations, + resources, + sections, + tls, + utils, +) + +# Local imports +from dissect.executable.pe.helpers.c_pe import cv_info_struct, pestruct + +if TYPE_CHECKING: + from dissect.cstruct.cstruct import cstruct + from dissect.cstruct.types.enum import EnumInstance + + +class PE: + """Base class for parsing PE files. + + Args: + pe_file: A file-like object of an executable. + virtual: Indicate whether the PE file exists within a memory image. + parse: Indicate if the different sections should be parsed automatically. + """ + + def __init__(self, pe_file: BinaryIO, virtual: bool = False): + pe_file.seek(0) + self.pe_file = BytesIO(pe_file.read()) + self.virtual = virtual + + # Make sure we reset any kind of pointers within the PE file before continueing + self.pe_file.seek(0) + + self.mz_header = None + self.file_header = None + self.nt_headers = None + self.optional_header = None + + self.section_header_offset = 0 + self.last_section_offset = 0 + + self.sections = OrderedDict() + self.patched_sections = OrderedDict() + + self.imports = None + self.exports = None + self.resources = None + self.raw_resources = None + self.relocations = None + self.tls_callbacks = None + + self.directories = OrderedDict() + + # We always want to parse the DOS header and NT headers + self.parse_headers() + + # The offset of the section header is always at the end of the NT headers + self.section_header_offset = self.pe_file.tell() + + self.imagebase = self.optional_header.ImageBase + self.file_alignment = self.optional_header.FileAlignment + self.section_alignment = self.optional_header.SectionAlignment + + self.base_address = self.optional_header.ImageBase + + self.timestamp = datetime.utcfromtimestamp(self.file_header.TimeDateStamp) + + # Parse the section header + self.parse_section_header() + + # Parsing the directories present in the PE + self.parse_directories() + + def _valid(self) -> bool: + """Check if the PE file is a valid PE file. By looking for the "PE" signature at the offset of e_lfanew. + + Returns: + `True` if the file is a valid PE file, `False` otherwise. + """ + + self.pe_file.seek(self.mz_header.e_lfanew) + return True if pestruct.uint32(self.pe_file) == 0x4550 else False + + def parse_headers(self): + """Function to parse the basic PE headers: + - DOS header + - File header (part of NT header) + - Optional header (part of NT header) + + Function also sets some architecture dependent variables. + + Raises: + InvalidPE if the PE file is not a valid PE file. + InvalidArchitecture if the architecture is not supported or unknown. + """ + + self.mz_header = pestruct.IMAGE_DOS_HEADER(self.pe_file) + + if not self._valid(): + raise InvalidPE("file is not a valid PE file") + + self.file_header = pestruct.IMAGE_FILE_HEADER(self.pe_file) + + image_nt_headers_offset = self.mz_header.e_lfanew + self.pe_file.seek(image_nt_headers_offset) + + # Set the architecture specific settings + self._set_pe_architecture() + if self.file_header.Machine == pestruct.MachineType.IMAGE_FILE_MACHINE_AMD64: + self.nt_headers = pestruct.IMAGE_NT_HEADERS64(self.pe_file) + else: + self.nt_headers = pestruct.IMAGE_NT_HEADERS(self.pe_file) + + self.optional_header = self.nt_headers.OptionalHeader + + def _set_pe_architecture(self): + """Set the architecture specific settings. Some of the structs are architecture specific. + + Raises: + InvalidArchitecture if the architecture is not supported or unknown. + """ + + if self.file_header.Machine == pestruct.MachineType.IMAGE_FILE_MACHINE_AMD64: + self.image_thunk_data = pestruct.IMAGE_THUNK_DATA64 + self.image_tls_directory = pestruct.IMAGE_TLS_DIRECTORY64 + self._high_bit = 1 << 63 + self.read_address = pestruct.uint64 + elif self.file_header.Machine == pestruct.MachineType.IMAGE_FILE_MACHINE_I386: + self.image_thunk_data = pestruct.IMAGE_THUNK_DATA32 + self.image_tls_directory = pestruct.IMAGE_TLS_DIRECTORY32 + self._high_bit = 1 << 31 + self.read_address = pestruct.uint32 + else: + raise InvalidArchitecture(f"Invalid architecture found: {self.file_header.Machine:02x}") + + def parse_section_header(self): + """Parse the sections within the PE file.""" + + self.pe_file.seek(self.section_header_offset) + + for _ in range(self.file_header.NumberOfSections): + # Keep track of the last section offset + offset = self.pe_file.tell() + section_header = pestruct.IMAGE_SECTION_HEADER(self) + section_name = section_header.Name.decode().strip("\x00") + # Take note of the sections, keep track of any patches seperately + self.sections[section_name] = sections.PESection(pe=self, section=section_header, offset=offset) + self.patched_sections[section_name] = sections.PESection(pe=self, section=section_header, offset=offset) + + self.last_section_offset = self.sections[next(reversed(self.sections))].offset + + def section(self, va: int = 0, name: str = "") -> sections.PESection: + """Function to retrieve a section based on the given virtual address or name. + + Args: + va: The virtual address to look for within the sections. + name: The name of the section. + + Returns: + A `PESection` object. + """ + + if not name: + for section in self.sections.values(): + if va in range(section.virtual_address, section.virtual_address + section.virtual_size): + return section + else: + return self.sections[name] + + def patched_section(self, va: int = 0, name: str = "") -> sections.PESection: + """Function to retrieve a patched section based on the given virtual address or name. + + Args: + va: The virtual address to look for within the patched sections. + name: The name of the patched section. + + Returns: + A `PESection` object. + """ + + if not name: + for section in self.patched_sections.values(): + if va in range(section.virtual_address, section.virtual_address + section.virtual_size): + return section + else: + return self.patched_sections[name] + + def datadirectory_section(self, index: int) -> sections.PESection: + """Return the section that contains the given virtual address. + + Args: + index: The index of the data directory to find the associated section for. + + Returns: + The section that contains the given virtual address. + """ + + va = self.directory_va(index=index) + for _, section in self.patched_sections.items(): + if va >= section.virtual_address and va < section.virtual_address + section.virtual_size: + return section + + raise InvalidVA(f"VA not found in sections: {va:#04x}") + + def parse_directories(self): + """Parse the different data directories in the PE file and initialize their associated managers. + + For now the following data directories are implemented: + - Import Address Table (IAT) + - Export Directory + - Resources + - Base Relocations + - Thread Local Storage (TLS) Callbacks + """ + + for idx in range(pestruct.IMAGE_NUMBEROF_DIRECTORY_ENTRIES): + if not self.has_directory(index=idx): + continue + + # Take note of the current directory VA so we can dynamically update it when resizing sections + section = self.datadirectory_section(index=idx) + directory_va_offset = self.optional_header.DataDirectory[idx].VirtualAddress - section.virtual_address + section.directories[idx] = directory_va_offset + + # Parse the Import Address Table (IAT) + if idx == pestruct.IMAGE_DIRECTORY_ENTRY_IMPORT: + self.import_mgr = imports.ImportManager(pe=self, section=section) + self.imports = self.import_mgr.imports + + if idx == pestruct.IMAGE_DIRECTORY_ENTRY_EXPORT: + self.export_mgr = exports.ExportManager(pe=self, section=section) + self.exports = self.export_mgr.exports + + # Parse the resources directory entry of the PE file + if idx == pestruct.IMAGE_DIRECTORY_ENTRY_RESOURCE: + self.rsrc_mgr = resources.ResourceManager(pe=self, section=section) + self.resources = self.rsrc_mgr.resources + self.raw_resources = self.rsrc_mgr.raw_resources + + # Parse the relocation directory entry of the PE file + if idx == pestruct.IMAGE_DIRECTORY_ENTRY_BASERELOC: + self.reloc_mgr = relocations.RelocationManager(pe=self, section=section) + self.relocations = self.reloc_mgr.relocations + + # Parse the TLS directory entry of the PE file + if idx == pestruct.IMAGE_DIRECTORY_ENTRY_TLS: + self.tls_mgr = tls.TLSManager(pe=self, section=section) + self.tls_callbacks = self.tls_mgr.callbacks + + def get_resource_type(self, rsrc_id: str | EnumInstance): + """Yields a generator containing all of the nodes within the resources that contain the requested ID. + + The ID can be either given by name or its value. + + Args: + rsrc_id: The resource ID to find, this can be a cstruct `EnumInstance` or `str`. + + Yields: + All of the nodes that contain the requested type. + """ + + if rsrc_id not in self.resources: + raise ResourceException(f"Resource with ID {rsrc_id} not found in PE!") + + for resource in self.rsrc_mgr.parse_resources(resources=self.resources[rsrc_id]): + yield resource + + def virtual_address(self, address: int) -> int: + """Return the virtual address given a (possible) physical address. + + Args: + address: The address to translate. + + Returns: + The virtual address as an` int` + """ + + if self.virtual: + return address + + for _, section in self.patched_sections.items(): + max_address = section.virtual_address + section.virtual_size + if address >= section.virtual_address and address < max_address: + return section.pointer_to_raw_data + (address - section.virtual_address) + + raise InvalidVA(f"VA not found in sections: {address:#04x}") + + def raw_address(self, offset) -> int: + """Return the physical address given a virtual address. + + Args: + offset: The offset to translate into a virtual address. + + Returns: + The physical address as an `int`. + """ + + for _, section in self.patched_sections.items(): + max_address = section.pointer_to_raw_data + section.size_of_raw_data + if offset >= section.pointer_to_raw_data and offset < max_address: + return section.virtual_address + (offset - section.pointer_to_raw_data) + + raise InvalidAddress(f"Raw address not found in sections: {offset:#04x}") + + def virtual_read(self, address: int, size: int) -> bytes: + """Wrapper for reading virtual address offsets within a PE file. + + Args: + address: The virtual address to read from. + size: The amount of bytes to read from the given virtual address. + + Returns: + The bytes that were read. + """ + + physical_address = self.virtual_address(address=address) + if self.virtual: + return self.pe_file.readoffset(offset=physical_address, size=size) + + self.pe_file.seek(physical_address) + return self.pe_file.read(size) + + def raw_read(self, offset: int, size: int) -> bytes: + """Read the amount of bytes denoted by the size argument within the PE file at the given offset. + + Args: + offset: The offset within the file to start reading. + size: The amount of bytes to read within the PE file. + + Returns: + The bytes that were read from the given offset. + """ + + old_offset = self.pe_file.tell() + self.pe_file.seek(offset) + + data = self.pe_file.read(size) + self.pe_file.seek(old_offset) + return data + + def seek(self, address: int): + """Seek to the given virtual address within a PE file. + + Args: + address: The virtual address to seek to. + """ + + raw_address = self.virtual_address(address=address) + self.pe_file.seek(raw_address) + + def tell(self) -> int: + """Returns the current offset within the PE file. + + Returns: + The current offset within the PE file. + """ + + offset = self.pe_file.tell() + return self.raw_address(offset=offset) + + def read(self, size: int) -> bytes: + """Read x amount of bytes of the PE file. + + Args: + size: The amount of bytes to read. + + Returns: + The bytes that were read. + """ + + return self.pe_file.read(size) + + def write(self, data: bytes): + """Write the data to the PE file. + + This write function will also make sure to update the section data. + + Args: + data: The data to write to the PE file. + """ + + offset = self.tell() + + # Write the data to the PE file so we can do a raw_read on this data in the section + self.pe_file.write(data) + print(self.patched_sections) + + # Update the section data + for section in self.patched_sections.values(): + if section.virtual_address <= offset and section.virtual_address + section.virtual_size >= offset: + self.seek(address=section.virtual_address) + section.data = self.read(size=section.virtual_size) + + def read_image_directory(self, index: int) -> bytes: + """Read the PE file image directory entry of a given index. + + Args: + index: The index of the data directory to read. + + Returns: + The bytes of the directory that was read. + """ + + directory_entry = self.optional_header.DataDirectory[index] + return self.virtual_read(address=directory_entry.VirtualAddress, size=directory_entry.Size) + + def directory_va(self, index: int) -> int: + """Returns the virtual address of a directory given its index. + + Args: + index: The index of the data directory to read. + + Returns: + The virtual address of the data directory at the given index. + """ + + return self.optional_header.DataDirectory[index].VirtualAddress + + def has_directory(self, index: int) -> bool: + """Check if a certain data directory exists within the PE file given its index. + + Args: + index: The index of the data directory to check. + + Returns: + `True` if the data directory has a size associated with it, indicating it exists, `False` otherwise. + """ + + return self.optional_header.DataDirectory[index].Size != 0 + + def debug(self) -> cstruct: + """Return the debug directory of the given PE file. + + Returns: + A `cstruct` object of the debug entry within the PE file. + """ + + debug_directory_entry = self.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_DEBUG) + image_directory_size = len(pestruct.IMAGE_DEBUG_DIRECTORY) + + for _ in range(0, len(debug_directory_entry) // image_directory_size): + entry = pestruct.IMAGE_DEBUG_DIRECTORY(debug_directory_entry) + dbg_entry = self.virtual_read(address=entry.AddressOfRawData, size=entry.SizeOfData) + + if entry.Type == 0x2: + return cv_info_struct.CV_INFO_PDB70(dbg_entry) + + def get_section(self, segment_index: int) -> Tuple[str, sections.PESection]: + """Retrieve the section of the PE by index. + + Args: + segment_index: The segment to retrieve based on the order within the PE. + + Returns: + A `Tuple` contianing the section name and attributes as `PESection`. + """ + + sections = list(self.sections.items()) + + idx = 0 if segment_index - 1 == -1 else segment_index + section_name = sections[idx - 1][0] + + return self.sections[section_name] + + def symbol_data(self, symbol: cstruct, size: int) -> bytes: + """Retrieve data from the PE using a PDB symbol. + + Args: + symbol: A `cstruct` object of a symbol. + size: The size to read from the symbol offset. + + Returns: + The bytes that were read from the offset within the PE. + """ + + _section = self.get_section(segment_index=symbol.seg) + address = self.imagebase + _section.virtual_address + symbol.off + + self.pe_file.seek(address) + return self.pe_file.read(size) + + def add_section( + self, + name: str, + data: bytes, + va: int = None, + datadirectory: int = None, + datadirectory_rva: int = None, + datadirectory_size: int = None, + size: int = None, + ): + """Add a new section to the PE file. + + Args: + name: The name of the new section. + data: The data to add to the new section. + datadirectory: Whether this section should be a specific data directory entry. + rva: The RVA of the directory entry if this is different than the virtual address of the section. + size: The size of the entry. + """ + + # Take note of the last section + last_section = self.patched_sections[next(reversed(self.sections))] + + # Calculate the new section size + raw_size = utils.align_int(integer=len(data), blocksize=self.file_alignment) + virtual_size = len(data) if not size else size + + # Use the provided RVA or calculate the new section virtual address + virtual_address = ( + utils.align_int( + integer=last_section.virtual_address + last_section.virtual_size, blocksize=self.section_alignment + ) + if not va + else va + ) + + # Calculate the new section raw address + pointer_to_raw_data = last_section.pointer_to_raw_data + last_section.size_of_raw_data + + # Build the new section + new_section = sections.build_section( + virtual_size=virtual_size, + virtual_address=virtual_address, + raw_size=raw_size, + pointer_to_raw_data=pointer_to_raw_data, + name=name.encode(), + ) + + # Update the last section offset + offset = last_section.offset + pestruct.IMAGE_SECTION_HEADER.size + self.last_section_offset = offset + + # Increment the NumberOfSections field + self.file_header.NumberOfSections += 1 + + # Set the VA and size of the datadirectory entry if this was marked as being such + if datadirectory is not None: + self.optional_header.DataDirectory[datadirectory].VirtualAddress = ( + virtual_address if not datadirectory_rva else datadirectory_rva + ) + self.optional_header.DataDirectory[datadirectory].Size = ( + len(data) if not datadirectory_size else datadirectory_size + ) + + # Add the new section to the PE + self.sections[name] = sections.PESection(pe=self, section=new_section, offset=offset, data=data) + self.patched_sections[name] = sections.PESection(pe=self, section=new_section, offset=offset, data=data) + + # Update the SizeOfImage field + last_section = self.patched_sections[next(reversed(self.patched_sections))] + last_va = last_section.virtual_address + last_size = last_section.virtual_size + + pe_size = utils.align_int(integer=(last_va + last_size), blocksize=self.section_alignment) + self.optional_header.SizeOfImage = pe_size + + # Write the data to the PE file + self.pe_file.seek(pointer_to_raw_data) + if virtual_size > raw_size: + data += utils.pad(virtual_size - raw_size) + + # Pad the data to align the section + padsize = utils.align_int(integer=len(data), blocksize=self.section_alignment) + data += utils.pad(size=padsize) + self.pe_file.write(data) + + # Reparse the directories + self.parse_directories() + + def write_pe(self, filename: str = "out.exe"): + """Write the contents of the PE to a new file. + + This will use the patcher that is part of the project to make sure any kind of relative addressing is also + corrected for the supported data directories. + + Args: + filename: The filename to write the PE to, default out.exe. + """ + + pepatcher = patcher.Patcher(pe=self) + new_pe = pepatcher.build + + with open(filename, "wb") as fhout: + fhout.write(new_pe.read()) diff --git a/tests/data/testexe.exe b/tests/data/testexe.exe new file mode 100644 index 0000000000000000000000000000000000000000..04a7d7259a51db87e133a048233618437781e9d3 GIT binary patch literal 31336 zcmeHwe{fsbb>5XIIq+s}W7eJw#&%Me<0go0lAuUwArEDPmMB9uwSqyJc~rRb=t zn5QPp&#y+FQzAD#!MW@6Xpg?*9gloa2l$TZ+37~iRpP!4G7I4wjViGyZgJ7#&_SDM-*=x~h{|Z#>xp`B zoz`7>pFXAm-}^pxy3eoQ7yo%&DE+h$AHsO1zMJq>-WQI9m#ACPM5m4Y6V%cB*yVeb zKnOQV8i7{g44!xN?xAn;zGxsq-zBX;GlBKkKTAEmj~(OnTF2iW6a}Ux`ekuy;$GYR zPal8RHQF&9?w<}1iocJ?Ii|od1&%3jOo5|Q;D`9o-9Pin`@}ndr|~m{pEmsBN6}}( z&+E?%QH9U`l@AN?fh$7%uk%8D1pSxsvxJ|Mc;b`@-i)Xo)qies)Efwhi($nRRj%Ds z!U50iY1Mxv7*)c{o+X8^%=tXt&@Jcwi@ok^!MWR!sIoE_^@PdsJ^mb1;Fto(6gZ~9 zF$In(a7=+?3LI15m;ygi3OrUtoEV?a|ACHS#%CC38K)S3mhmyhH~IPRGrrFF3Zv^A z`nb!CHyMAG@k@;Vk+IJBSBwT8XN~dwj2~fiFh0u|W=t@?#P}m6>VTKM=XIS9Kg)QI zah7qB@w1G-!uUnT3}czG%2;|$ANLfu|L1?!+YQFo7++%iO~yFmGUGhsCm4O-)W`We zjuH zev0vD7|$?XW_+5_#~5S$0^^q%Ut;_^V~z1I7@LegV0`3T`ujh~XlHza@ma%EL0 zVf+;1B4dp4MaB%{%Z%S-{4Qgg@v)jd{wSlJ@j9cMF~;~JW0vt{#@89&U^MK zMrWtZeEuOu6QhB#{cZjE?=jXGUt#-_Vu-q`Heea46Np38rtuj8YP?`M3J@dqr&f64eKj8VRxV@z@X zuXe68{wKcvdB$@rH+jDPUT#0j*DKurMQ%UJ*MGq6?=!x}_!o>1i#|sk_fu?--^)13 zc+}_b^}amw#Wx=^e^{B{?{$7hy}$Z9JvsUy-bL;5%e_83`lxN}yYExzlkHOPzAqCU z(02Ac`yKRA+gk43_fG5!wHsgU^*w_=YWsex*S8Elq;{(BStoe0zumvaWBvXO+%{OG;!Z>{p|qun_;P|B2p(OP-~GcB_t4bCE1sx76cn$~3}=GTuwRK>2~IBsu+Geh*-+^DRdsr4DYP1l zKH*tW(EC}uKZwerDBd(5n!^*p8~BPV&qzwj(j%sxhk7n40VPVrofeCG+C=fCOqKEw0M`mjXj>mZBNZypSyDH2|1c`sK;;@!Qr590IVdr z6sbo;rzD&w3)!nQ)k7Y~$2*CDI+|tDWcOz@D9_gpr=P zxEc=QT`X}|yyCy}JLi4iT}aP+`tH*EaFh5O>g_f&Qub#3i(wCTFhkd9z-g~Hj2#!^ z&rdJ-y~;(OCp;gLQ~$*6s3N{i_vw{&m;7O@1KYnKOsBE_zLiW3ITVj7m6q7*jEqsK z7usit^hJk0zi@gkpeU;NJEGyGK%hG;X!jgFdo|=;4KStASKzhw-}f^25erR=i}VzI zB^H86y*YGkHL9*gFDZ+wH*|5wI16;wEL2or-oFA_@%P@X08W!Pii+Gh;(M>tuPPYjxA{4_X0)xRB=b5^(n$I%xTIVp{iKA& z_<<{+;xDc$!KK^cZzbJ$*F|nzGPx+Nw z*OqAkqA8!JuU9>)3a-LaHn~$$SBuwpmT~}Gdlyk;FRzK;9m2DHWBwLq0TSv1i2tCs z5Fr9eJ7w=`AV4aKwgi3u=O90z_s@8SOkd37&-eD=o|WEwb;(EH zcK<%i#L#1WTObk*FRiE;BS&-fFW>G{Y`Eix;*zoy@+!`dU+((FQ`bHI2v}@MQKgy^ z*ALfW;OG{~3hzw696WxI2AD(e^t=*Y@dsgLI?N;9eTu(cArr_d9<=>a8eyM&Pm6a9 zp_i9Ha_}zM4h2?Mls$Ig+e2MVF(2xv50M=D^u!qqz;D!ThV(jGOyT)ZDBx29>cpA3 zphu1PLeQhe&I~8cT=s_}(FO5XujdK;HXG<@VC^ioUA#O0pEcv8O^>+6N?<9d0PSSmc zURH;$M0nky@MY9AV-8LzSTYsQ;hNM*FwLF5-;?yGn&nlMNVFjSOK(R=QB6l^u~KEt z_wmTwEUYTjt5GXP}zp3&*>km^` z0b@Z(LbvrU{Q4j)P4I@IH~F19iLry98@d%y;ig5G7sMN4Cb+snvqSMWNFRkFqxb%8 zfBz+g+#qy+4c=7H+gn3K(xca29b)TzB?t>~MS4v$L>(AxG00qgRHJ-4_i z-gbgK9KAIy-rm<*0AKd5kxJO45$C|+fzH!El(W@a5%FhqU2o|2xhtGP{9XULI_s5G zeR6K*ddFXTnYdPGSJ3OXspVXn4-u~rSAbDJ)5k6QycxLYbFc@}28sCT`9OsA`M;*` zn40ug#pkGn8R)HlMJ=i6pT@{w0J!_l!{V(*d{UAmM+#T}waWoq)cgM_{IGrvxNrh~ zVMjxV0P%dk-Isdx-PFrkujA=#6Xm8aI)(U#E;(Ly$;1q`rL5dm74fQmm0bOKw4pmv zC**YKj!&#EFT?Q^;=42&8yp(*jlMRf11}*ccUJm-A|k&_%J&Eq`!<#-Q`$%Qc zZPVhH_p}$pD)azo>i1qGlK=@7zecmYHglEc@Krenr_`6P>GN5E%K*zC#8`i{uV+Dw z(z7$coBnVpxB@k!J3M3y*QLid=+gftd(wB|DNoq%S){qXdB~Lo@m_k*CFl%D3XWhH z{~xG_G*h?dH*nfmQ9WTrp9JptcSEx{ji}xB)AWpVA*3b29k)8v(0j#cvNd$_f%5Ud zg`T~5<&qGy;GfkM-B9T_(hC1r+C|-AxEKuwbY2LHY3b#VJ1E4*<@f7~&f=q_z2~A` zuK7u7(>AebhA@fGNI#F+0&n?i-Ph}Ov~?PJ+07al%Q)^lrob@;`YCWyP@Rhz^)2*g z18hJ*-{*wCiE& z+*AVQ0R4Mm*n=3_O@#8o!{${#!e?+DFAOhx0ug0+tTW2g*rBhO8tb3V)L3^_^z0NU z0oT&M=no+Lc=vVgxt++TPhbj%SQ$IR+`ChXtg31#j749YGhgvu82)F<8SD7?SyyR3{@tv0WsC==d^W$QohThoC7CeDse*cQ)wM$UTAV}RvF?}$w~ z;~nvwnVhtqIdgXT{DkL@15A3(ET6F{6C-CQ7SE4NPEL-GEPBo_jd;(U_fC2yJ?E4O z+Z_iOUq1KvxuwU?jo7Tp_{ij$iSr}p&n_;G*e1Nz<#S%E*Yo(D2e2vD<;OkaXGRv6 z7bma;#WsQkogJC9Eql*eJ>zFR=X*<-8rw(Lsj)p9Aa}6umb>qtWcn_j-jW9h@*}0^ z-rYHf)D1eMPv3&S&>7ShjjA6V8)IP_?F!Q{Sq;@>0-+AM12Pr&-jy)9%Rr72F;rc8}u2N7|&FnqP6 z14enKqZdOfE1}>BokqiF|;dA*UreQeGJT z5U#8)LLLyNmeJ!2!{^R|P0me@Pfji^EsuMbhs^;Tu~y+=%G+6vr^W^@uGdXO0ExN_ z?!tO|dUd_FUSDslH`iP1?R7)Om@#F{8B4~Rv1RNTSH_o#W#XAcCYecP@O*Lv zXIUX?+x6|nc5}P6-QE^Eh8^RMX~(={*|F}}cI-Qj9p{c~$Gzj*QFmfH@twp@awoN; z?c{b!JLR3qPIafYQ{QRqG)dtks=M*s z1Z^O{9%!bK07= zr=4kcT205($+VU(r7P)Lx{+?B1$NJj-LtPb*W7FBT6`_JrmdB*1nWOI!*3lXf@e}Yf)-8C84dW;9 zh7{hB12>fq@RP{w=O_o>?83W!czchtv`j8j%9JyeOf^%>G&0RhE7Q)14a0_U!?a=E zuxwa2Y#a6s$A)vmwc+0IZKxZujrc}lBe{{<&^B@#rH%4NWuv-L+o*3eHkuo)jrNAf z8nVW$DQnJJvev9EE4B^W#%CK});F7*c-THF8h}c!nycj!TdA$wR(Y$sRo`lEwYLnQn*}s; zfL=b(DzTl~&TW^stIz~ZXaNH>fCbX;faLoi^$AFP4$@wQq&Feu21vLC((Qm``ykZ` zNOTU;T!kbz^X<~2&8Nk*gqG5B(4bYVt~Ir`W=NaTmb5MHNW0R$bS#}nr_#A}IbBWH)6H}{ zZCEowgWA>{YpyllT5K(`mRifLm7#O%Yt6Oxnql1p>9DOk)?Mp9=-&k7BL~^2LeDm( zjx|BQ+AIW1SpRls?TTq`HGj9ca{>y~}Xx#iwcx8hsLEp4l`RoSX-HMUw?0&;AI4BH{U zZpdzYJGrfGm$ob0wSNRs+`79(P)Eh1*o(b#tiC0Cpj$)IJ#A|qKPF$bGihE=S>F4kcg zo3M?L%=oByg?>2$evA>aX@+!KVIdvTM*5_cOh`MKgQYCPR#v68Y{Fi)rNuN!n`wj9 zbii)9U^#uToiSL?1ng%D7BmMNT80&^!j9HqNt>{xZCFzS?5PPB)gnD2J7n4ksdhuI zRY*43A+k1PSxV5Pm8=!|(+&;lgbsB>i>lD0acI&cbg2eyT7o{UK%>^6Qyb8#t<5%g z-2iSkf!{6QcpG@$0j_s}?|tC>82CO1&M$-aN$=Ld|4m3h8!}*k1lXhuxF7{S$UzK} zkbo?tq%@Qv4^>D+9Wv2`RJ0)%hC5ogIP`WBnp=bJE*wPh7x#EraxfgmQIV%V z)FXLV=U^N2unV3Zad8aNO?KCYfNBEU-&=lH#`_!(_iwca&tZXBKJE z9MY!km*yO-+hOu-KEz+urTM|ja_g?AvYsrL`c7m=m*iyE3RR>HYV=tkv$Q{UczVt~ z)<=WCSL(Apdo9j>Sxp>jbA&c3o`I|y`mD`hS=Hbdmbw*e(GvsV(%o2he@IC$M&SJGExpLI7gXpZpFBJ5mAOHl=U&OPBVKaQ`z*cXm}03LO(d_gc@Oc#(h(wwgC*~+`<-vi z=S;x*l@2*M9wu|@o-<@_fX`_u5a%<3&sze0~-TQfP<``HXwd;Fm~hSs7=q_RSv{$^pjq{$JvY8bgW(D zK*YwrFB=gb5Ty|Za&7JdS+>f7{Mx~IO#DDB#Gn72m@&5(BF)4)?fvkq7mzxu%EVbL~GF_1iGox&Ew^C|SiXVU72Hnm$TIS&wX0 zWL^j_tysmI0P`VPr+$xYaLm4EclDEvK~_mvWrc)GR!78TWkgO^Mbu!jSfS)Na656Jf=xuI+colc4JBsDTq zQiCpp&r7jav&_Qlu^}I4`N-d;h-;NId~MF}(W#VtT{@AfGAEyuSv`sb)ns12C9`=Z zSsg>Ud8f?c#bkvH<>yN>_d$^#@?D661G#wGRL0i<`t_l&)$?7MT*8)WUk#IYh_ef{XzNm zoXjBBWCch|#y(B5cE%=CL(dA5xLA%_irsO|njlszmhM0Yy}jvQk9P9Mq-F zrfLzR%vRetSM8HkTU0Zm$#@Fsibkhyq@l|z0y^qdsUMaCM{HX8DtfoRaTieWwljI)_PH;Nl8{&QBTi>^ z9#nCn%DSo);{K!{$HhFrx?SBdqu#0UAgbZa<6;tH-l6v4K#hs8>(5xF#P-gjL;?Lg zs@JmZ@oAjnB1UtUT<$>ZUqz1GFrYrfBze)kukJ8G-hr%+AZ^$mA+UEdQ;B^UW%3Dy zdz|X;g&UdI3bc(GGpw&0uopg5`cy&l80vXQW9l(IH+YXYPpIP}h39IhL()(BCd9$* z;GbSxjPz4Z?zFyBq+b%UHm59UMpZYYD{Qh#!G-#S6i&6e+&8(Zz=%8~Rn<_rom#xcAnQIjFsA36~ZzjZ@UC?wnu52|x zx7p!4HS$(e+|&wYSk6k(2u}6mC~t->s6hLx(D7#MYz16yMqN}A)g4yUFxGYiw5tYx z)Q&134Kp&rrpGZSfjPyY?;5C%n-EQ$QH;pds5~13Voc5-liIhcw2|#i11zLP+D8}S zI5Ff8Q>eZvgFou1x*@x0f^^ux7cSJ>#87LK`thoQG){;TSkLf}vRN32@khRTaREL>R-|2=G&cR~C*Jb-7-_=p9Go~5kRaN|71g8ZV literal 0 HcmV?d00001 diff --git a/tests/test_pe.py b/tests/test_pe.py new file mode 100644 index 0000000..095bec4 --- /dev/null +++ b/tests/test_pe.py @@ -0,0 +1,78 @@ +from io import BytesIO + +import pytest + +from dissect.executable.exception import InvalidPE +from dissect.executable.pe.pe import PE + + +def test_pe_valid_signature(): + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + assert pe._valid() is True + + +def test_pe_invalid_signature(): + with pytest.raises(InvalidPE): + PE(BytesIO(b"MZ" + b"\x00" * 400)) + + +def test_pe_sections(): + known_sections = [".dissect", ".text", ".rdata", ".idata", ".rsrc", ".reloc", ".tls"] + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + assert known_sections == [section for section in pe.sections.keys()] + + +def test_pe_imports(): + known_imports = [ + "SHELL32.dll", + "ole32.dll", + "OLEAUT32.dll", + "ADVAPI32.dll", + "WTSAPI32.dll", + "SHLWAPI.dll", + "VERSION.dll", + "KERNEL32.dll", + "USER32.dll", + ] + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + assert known_imports == [import_ for import_ in pe.imports.keys()] + + +def test_pe_exports(): + # Too much export functions to put in a list + known_exports = ["1", "2", "CreateOverlayApiInterface", "CreateShadowPlayApiInterface", "ShadowPlayOnSystemStart"] + + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + assert known_exports == [export_ for export_ in pe.exports.keys()] + + +def test_pe_resources(): + known_resource_types = ["RcData", "Manifest"] + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + assert known_resource_types == [resource for resource in pe.resources.keys()] + + +def test_pe_relocations(): + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + assert len(pe.relocations) == 9 + + +def test_pe_tls_callbacks(): + known_callbacks = [430080, 434176, 438272, 442368, 446464, 450560, 454656, 458752, 462848, 466944] + + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + assert pe.tls_callbacks == known_callbacks diff --git a/tests/test_pe_builder.py b/tests/test_pe_builder.py new file mode 100644 index 0000000..5d8dd6d --- /dev/null +++ b/tests/test_pe_builder.py @@ -0,0 +1,69 @@ +from dissect.executable import PE +from dissect.executable.pe import Builder, Patcher +from dissect.executable.pe.helpers.c_pe import pestruct + + +def test_build_new_pe_lfanew(): + builder = Builder() + builder.new() + pe = builder.pe + + assert pe.mz_header.e_lfanew == 0x8C + + +def test_build_new_x86_pe_exe(): + builder = Builder(arch="x86") + builder.new() + pe = builder.pe + + pe.pe_file.seek(len(pe.mz_header)) + stub = pe.pe_file.read(pe.mz_header.e_lfanew - len(pe.mz_header)) + assert stub[14 : 73 - 4] == b"This program is made with dissect.pe <3 kusjesvanSRT <3" + + assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE == 0x0100 + + +def test_build_new_x64_pe_exe(): + builder = Builder(arch="x64") + builder.new() + pe = builder.pe + + pe.pe_file.seek(len(pe.mz_header)) + stub = pe.pe_file.read(pe.mz_header.e_lfanew - len(pe.mz_header)) + assert stub[14 : 73 - 4] == b"This program is made with dissect.pe <3 kusjesvanSRT <3" + + assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE != 0x0100 + + +def test_build_new_x86_pe_dll(): + builder = Builder(arch="x86", dll=True) + builder.new() + pe = builder.pe + + assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE == 0x0100 + assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_DLL == 0x2000 + + +def test_build_new_x64_pe_dll(): + builder = Builder(arch="x64", dll=True) + builder.new() + pe = builder.pe + + assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE != 0x0100 + assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_DLL == 0x2000 + + +def test_build_new_pe_with_custom_section(): + builder = Builder() + builder.new() + pe = builder.pe + + pe.add_section(name=".SRT", data=b"kusjesvanSRT") + + patcher = Patcher(pe=pe) + + new_pe = PE(pe_file=patcher.build) + + assert new_pe.sections[".SRT"].name == ".SRT" + assert new_pe.sections[".SRT"].size == 12 + assert new_pe.sections[".SRT"].data == b"kusjesvanSRT" diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py new file mode 100644 index 0000000..e1fcc06 --- /dev/null +++ b/tests/test_pe_modifications.py @@ -0,0 +1,94 @@ +# Local imports +from dissect.executable import PE +from dissect.executable.pe import Patcher + + +def test_add_imports(): + dllname = "kusjesvanSRT.dll" + functions = ["PressButtons", "LooseLips"] + + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + pe.import_mgr.add(dllname=dllname, functions=functions) + + patcher = Patcher(pe=pe) + new_pe = PE(pe_file=patcher.build) + + assert "kusjesvanSRT.dll" in new_pe.imports + + custom_dll_imports = [i.name for i in new_pe.imports["kusjesvanSRT.dll"].functions] + assert "PressButtons" in custom_dll_imports + assert "LooseLips" in custom_dll_imports + + +def test_resize_section_smaller(): + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + pe.sections[".text"].data = b"kusjesvanSRT, patched with dissect" + + patcher = Patcher(pe=pe) + new_pe = PE(pe_file=patcher.build) + + assert new_pe.sections[".text"].size == len(b"kusjesvanSRT, patched with dissect") + assert ( + new_pe.sections[".text"].data[: len(b"kusjesvanSRT, patched with dissect")] + == b"kusjesvanSRT, patched with dissect" + ) + + +def test_resize_section_bigger(): + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + original_size = pe.sections[".rdata"].size + + pe.patched_sections[".rdata"].data += b"kusjesvanSRT, patched with dissect" * 100 + + patcher = Patcher(pe=pe) + new_pe = PE(pe_file=patcher.build) + + assert new_pe.sections[".rdata"].size == original_size + len(b"kusjesvanSRT, patched with dissect" * 100) + + +def test_resize_resource_smaller(): + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + for e in pe.get_resource_type(rsrc_id="Manifest"): + e.data = b"kusjesvanSRT, patched with dissect" + + patcher = Patcher(pe=pe) + new_pe = PE(pe_file=patcher.build) + + assert [patched.data for patched in new_pe.get_resource_type(rsrc_id="Manifest")] == [ + b"kusjesvanSRT, patched with dissect" + ] + + +def test_resize_resource_bigger(): + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + + for e in pe.get_resource_type(rsrc_id="Manifest"): + e.data = b"kusjesvanSRT, patched with dissect" + e.data + + patcher = Patcher(pe=pe) + new_pe = PE(pe_file=patcher.build) + + assert [ + patched.data[: len(b"kusjesvanSRT, patched with dissect")] + for patched in new_pe.get_resource_type(rsrc_id="Manifest") + ] == [b"kusjesvanSRT, patched with dissect"] + + +def test_add_section(): + with open("tests/data/testexe.exe", "rb") as pe_fh: + pe = PE(pe_file=pe_fh) + pe.add_section(name=".SRT", data=b"kusjesvanSRT") + + patcher = Patcher(pe=pe) + new_pe = PE(pe_file=patcher.build) + + assert ".SRT" in new_pe.sections + assert new_pe.sections[".SRT"].data == b"kusjesvanSRT" From d23b36d250c0590205f3aea7982e9789f6559d56 Mon Sep 17 00:00:00 2001 From: sud0woodo <40278342+sud0woodo@users.noreply.github.com> Date: Sat, 9 Mar 2024 12:35:51 +0100 Subject: [PATCH 02/43] Fix tabs to spaces --- dissect/executable/pe/helpers/c_pe.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/dissect/executable/pe/helpers/c_pe.py b/dissect/executable/pe/helpers/c_pe.py index 606fd6c..8e631f4 100755 --- a/dissect/executable/pe/helpers/c_pe.py +++ b/dissect/executable/pe/helpers/c_pe.py @@ -306,21 +306,21 @@ }; typedef struct IMAGE_THUNK_DATA32 { - union { - DWORD ForwarderString; - DWORD Function; - DWORD Ordinal; - DWORD AddressOfData; - } u1; + union { + DWORD ForwarderString; + DWORD Function; + DWORD Ordinal; + DWORD AddressOfData; + } u1; }; typedef struct IMAGE_THUNK_DATA64 { - union { - ULONGLONG ForwarderString; - ULONGLONG Function; - ULONGLONG Ordinal; - ULONGLONG AddressOfData; - } u1; + union { + ULONGLONG ForwarderString; + ULONGLONG Function; + ULONGLONG Ordinal; + ULONGLONG AddressOfData; + } u1; } // --- END OF IMPORTS From c307872f0d86efdd9d3007c8febc374e67d73883 Mon Sep 17 00:00:00 2001 From: sud0woodo <40278342+sud0woodo@users.noreply.github.com> Date: Sun, 27 Oct 2024 12:15:38 +0200 Subject: [PATCH 03/43] Style refactor * Refactored based on comments given * Ran `ruff` for formatting * Applied style and formatting on `dissect.executable.elf` as well (hope you don't mind) --- dissect/executable/elf/elf.py | 9 +- dissect/executable/pe/{helpers => }/c_pe.py | 54 ++++---- dissect/executable/pe/helpers/builder.py | 105 +++++++-------- dissect/executable/pe/helpers/exports.py | 43 +++--- dissect/executable/pe/helpers/imports.py | 96 +++++++++----- dissect/executable/pe/helpers/patcher.py | 107 ++++++++++----- dissect/executable/pe/helpers/relocations.py | 21 ++- dissect/executable/pe/helpers/resources.py | 59 ++++++--- dissect/executable/pe/helpers/sections.py | 26 ++-- dissect/executable/pe/helpers/tls.py | 14 +- dissect/executable/pe/helpers/utils.py | 6 +- dissect/executable/pe/pe.py | 132 ++++++++++++------- tests/test_pe.py | 47 +++++-- tests/test_pe_builder.py | 56 +++++--- tests/test_pe_modifications.py | 34 +++-- tests/test_section.py | 6 +- tests/test_segment.py | 4 +- tests/test_segment_table.py | 4 +- 18 files changed, 526 insertions(+), 297 deletions(-) rename dissect/executable/pe/{helpers => }/c_pe.py (90%) diff --git a/dissect/executable/elf/elf.py b/dissect/executable/elf/elf.py index f2b827f..8a8963c 100644 --- a/dissect/executable/elf/elf.py +++ b/dissect/executable/elf/elf.py @@ -287,7 +287,14 @@ def patch(self, new_data: bytes) -> None: class SegmentTable(Table[Segment]): - def __init__(self, fh: BinaryIO, offset: int, entries: int, size: int, c_elf: cstruct = c_elf_64): + def __init__( + self, + fh: BinaryIO, + offset: int, + entries: int, + size: int, + c_elf: cstruct = c_elf_64, + ): super().__init__(entries) self.fh = fh self.offset = offset diff --git a/dissect/executable/pe/helpers/c_pe.py b/dissect/executable/pe/c_pe.py similarity index 90% rename from dissect/executable/pe/helpers/c_pe.py rename to dissect/executable/pe/c_pe.py index 8e631f4..ff5e449 100755 --- a/dissect/executable/pe/helpers/c_pe.py +++ b/dissect/executable/pe/c_pe.py @@ -1,6 +1,6 @@ from dissect.cstruct import cstruct -pe_def = """ +c_pe_def = """ #define IMAGE_NUMBEROF_DIRECTORY_ENTRIES 16 #define IMAGE_SIZEOF_SHORT_NAME 8 @@ -93,12 +93,12 @@ }; typedef struct IMAGE_FILE_HEADER { - MachineType Machine; - WORD NumberOfSections; - DWORD TimeDateStamp; - DWORD PointerToSymbolTable; - DWORD NumberOfSymbols; - WORD SizeOfOptionalHeader; + MachineType Machine; + WORD NumberOfSections; + DWORD TimeDateStamp; + DWORD PointerToSymbolTable; + DWORD NumberOfSymbols; + WORD SizeOfOptionalHeader; ImageCharacteristics Characteristics; }; @@ -212,9 +212,9 @@ }; typedef struct IMAGE_NT_HEADERS64 { - DWORD Signature; - IMAGE_FILE_HEADER FileHeader; - IMAGE_OPTIONAL_HEADER64 OptionalHeader; + DWORD Signature; + IMAGE_FILE_HEADER FileHeader; + IMAGE_OPTIONAL_HEADER64 OptionalHeader; }; flag SectionFlags : DWORD { @@ -258,16 +258,16 @@ }; typedef struct IMAGE_SECTION_HEADER { - char Name[IMAGE_SIZEOF_SHORT_NAME]; - ULONG VirtualSize; - ULONG VirtualAddress; - ULONG SizeOfRawData; - ULONG PointerToRawData; - ULONG PointerToRelocations; - ULONG PointerToLinenumbers; - USHORT NumberOfRelocations; - USHORT NumberOfLinenumbers; - SectionFlags Characteristics; + char Name[IMAGE_SIZEOF_SHORT_NAME]; + ULONG VirtualSize; + ULONG VirtualAddress; + ULONG SizeOfRawData; + ULONG PointerToRawData; + ULONG PointerToRelocations; + ULONG PointerToLinenumbers; + USHORT NumberOfRelocations; + USHORT NumberOfLinenumbers; + SectionFlags Characteristics; }; // --- END OF PE HEADERS @@ -277,8 +277,8 @@ typedef struct IMAGE_EXPORT_DIRECTORY { ULONG Characteristics; ULONG TimeDateStamp; - USHORT MajorVersion; - USHORT MinorVersion; + USHORT MajorVersion; + USHORT MinorVersion; ULONG Name; ULONG Base; ULONG NumberOfFunctions; @@ -441,12 +441,9 @@ // --- END OF TLS DIRECTORY """ +c_pe = cstruct().load(c_pe_def) -pestruct = cstruct() -pestruct.load(pe_def) - - -cv_info_def = """ +c_cv_info_def = """ struct GUID { DWORD Data1; WORD Data2; @@ -462,5 +459,4 @@ }; """ -cv_info_struct = cstruct() -cv_info_struct.load(cv_info_def) +c_cv_info = cstruct().load(c_cv_info_def) diff --git a/dissect/executable/pe/helpers/builder.py b/dissect/executable/pe/helpers/builder.py index acc1484..261e831 100644 --- a/dissect/executable/pe/helpers/builder.py +++ b/dissect/executable/pe/helpers/builder.py @@ -5,17 +5,11 @@ from typing import TYPE_CHECKING from dissect.executable.exception import BuildSectionException +from dissect.executable.pe.c_pe import c_pe from dissect.executable.pe.helpers import utils -from dissect.executable.pe.helpers.c_pe import pestruct - -# Local imports from dissect.executable.pe.pe import PE -if TYPE_CHECKING: - from dissect.cstruct.cstruct import cstruct - - -STUB = b"\x0e\x1f\xba\x0e\x00\xb4\t\xcd!\xb8\x01L\xcd!This program is made with dissect.pe <3 kusjesvanSRT <3.\x0D\x0D\x0A$\x00\x00" # noqa: E501 +STUB = b"\x0e\x1f\xba\x0e\x00\xb4\t\xcd!\xb8\x01L\xcd!This program is made with dissect.pe <3 kusjesvanSRT <3.\x0d\x0d\x0a$\x00\x00" # noqa: E501 class Builder: @@ -35,16 +29,16 @@ def __init__( subsystem: int = 0x2, ): self.arch = ( - pestruct.MachineType.IMAGE_FILE_MACHINE_AMD64 + c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64 if arch == "x64" - else pestruct.MachineType.IMAGE_FILE_MACHINE_I386 + else c_pe.MachineType.IMAGE_FILE_MACHINE_I386 ) self.dll = dll self.subsystem = subsystem self.pe = None - def new(self): + def new(self) -> None: """Build the PE file from scratch. This function will build a new PE that consists of a single dummy section. It will not contain any imports, @@ -58,7 +52,9 @@ def new(self): image_characteristics = self.get_characteristics() # Generate the file header - self.file_header = self.gen_file_header(machine=self.arch, characteristics=image_characteristics) + self.file_header = self.gen_file_header( + machine=self.arch, characteristics=image_characteristics + ) # Generate the optional header self.optional_header = self.gen_optional_header() @@ -69,17 +65,18 @@ def new(self): section_header_offset = self.optional_header.SizeOfHeaders pointer_to_raw_data = utils.align_int( - integer=section_header_offset + pestruct.IMAGE_SECTION_HEADER.size, blocksize=self.file_alignment + integer=section_header_offset + c_pe.IMAGE_SECTION_HEADER.size, + blocksize=self.file_alignment, ) dummy_section = self.section( pointer_to_raw_data=pointer_to_raw_data, virtual_address=self.optional_header.BaseOfCode, virtual_size=dummy_multiplier, raw_size=dummy_multiplier, - characteristics=pestruct.SectionFlags.IMAGE_SCN_CNT_CODE - | pestruct.SectionFlags.IMAGE_SCN_MEM_EXECUTE - | pestruct.SectionFlags.IMAGE_SCN_MEM_READ - | pestruct.SectionFlags.IMAGE_SCN_MEM_NOT_PAGED, + characteristics=c_pe.SectionFlags.IMAGE_SCN_CNT_CODE + | c_pe.SectionFlags.IMAGE_SCN_MEM_EXECUTE + | c_pe.SectionFlags.IMAGE_SCN_MEM_READ + | c_pe.SectionFlags.IMAGE_SCN_MEM_NOT_PAGED, ) # Update the number of sections in the file header self.file_header.NumberOfSections += 1 @@ -125,7 +122,7 @@ def gen_mz_header( e_oeminfo: int = 0, e_res2: int = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], e_lfanew: int = 0, - ) -> cstruct: + ) -> c_pe.IMAGE_DOS_HEADER: """Generate the MZ header for the new PE file. Args: @@ -153,7 +150,7 @@ def gen_mz_header( The MZ header as a `cstruct` object. """ - mz_header = pestruct.IMAGE_DOS_HEADER() + mz_header = c_pe.IMAGE_DOS_HEADER() mz_header.e_magic = e_magic mz_header.e_cblp = e_cblp @@ -194,26 +191,14 @@ def get_characteristics(self) -> int: The characteristics of the PE file. """ - if self.arch == pestruct.MachineType.IMAGE_FILE_MACHINE_AMD64: - if self.dll: - return ( - pestruct.ImageCharacteristics.IMAGE_FILE_EXECUTABLE_IMAGE - | pestruct.ImageCharacteristics.IMAGE_FILE_DLL - ) - else: - return pestruct.ImageCharacteristics.IMAGE_FILE_EXECUTABLE_IMAGE - else: - if self.dll: - return ( - pestruct.ImageCharacteristics.IMAGE_FILE_EXECUTABLE_IMAGE - | pestruct.ImageCharacteristics.IMAGE_FILE_DLL - | pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE - ) - else: - return ( - pestruct.ImageCharacteristics.IMAGE_FILE_EXECUTABLE_IMAGE - | pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE - ) + output = c_pe.ImageCharacteristics.IMAGE_FILE_EXECUTABLE_IMAGE + if self.arch != c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64: + output |= c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE + + if self.dll: + output |= c_pe.ImageCharacteristics.IMAGE_FILE_DLL + + return output def gen_file_header( self, @@ -224,17 +209,17 @@ def gen_file_header( characteristics: int = 0, machine: int = 0x8664, number_of_sections: int = 0, - ) -> cstruct: + ) -> c_pe.IMAGE_FILE_HEADER: """Generate the file header for the new PE file. Args: - machine: The machine type. - number_of_sections: The number of sections. time_date_stamp: The time and date the file was created. pointer_to_symbol_table: The file pointer to the COFF symbol table. number_of_symbols: The number of entries in the symbol table. size_of_optional_header: The size of the optional header. characteristics: The characteristics of the file. + machine: The machine type. + number_of_sections: The number of sections. Returns: The file header as a `cstruct` object. @@ -243,17 +228,17 @@ def gen_file_header( # Set the size of the optional header if not given if not size_of_optional_header: if machine == 0x8664: - size_of_optional_header = len(pestruct.IMAGE_OPTIONAL_HEADER64) + size_of_optional_header = len(c_pe.IMAGE_OPTIONAL_HEADER64) self.machine = 0x8664 else: - size_of_optional_header = len(pestruct.IMAGE_OPTIONAL_HEADER) + size_of_optional_header = len(c_pe.IMAGE_OPTIONAL_HEADER) self.machine = 0x14C # Set the timestamp to now if not given if not time_date_stamp: time_date_stamp = int(datetime.utcnow().timestamp()) - file_header = pestruct.IMAGE_FILE_HEADER() + file_header = c_pe.IMAGE_FILE_HEADER() file_header.Machine = machine file_header.NumberOfSections = number_of_sections file_header.TimeDateStamp = time_date_stamp @@ -294,12 +279,12 @@ def gen_optional_header( size_of_heap_reserve: int = 0x1000, size_of_heap_commit: int = 0x1000, loaderflags: int = 0, - number_of_rva_and_sizes: int = pestruct.IMAGE_NUMBEROF_DIRECTORY_ENTRIES, + number_of_rva_and_sizes: int = c_pe.IMAGE_NUMBEROF_DIRECTORY_ENTRIES, datadirectory: list = [ - pestruct.IMAGE_DATA_DIRECTORY(BytesIO(b"\x00" * len(pestruct.IMAGE_DATA_DIRECTORY))) - for _ in range(pestruct.IMAGE_NUMBEROF_DIRECTORY_ENTRIES) + c_pe.IMAGE_DATA_DIRECTORY(BytesIO(b"\x00" * len(c_pe.IMAGE_DATA_DIRECTORY))) + for _ in range(c_pe.IMAGE_NUMBEROF_DIRECTORY_ENTRIES) ], - ) -> cstruct: + ) -> c_pe.IMAGE_OPTIONAL_HEADER | c_pe.IMAGE_OPTIONAL_HEADER64: """Generate the optional header for the new PE file. Args: @@ -339,10 +324,10 @@ def gen_optional_header( """ if self.machine == 0x8664: - optional_header = pestruct.IMAGE_OPTIONAL_HEADER64() + optional_header = c_pe.IMAGE_OPTIONAL_HEADER64() optional_header.Magic = 0x20B if not magic else magic else: - optional_header = pestruct.IMAGE_OPTIONAL_HEADER() + optional_header = c_pe.IMAGE_OPTIONAL_HEADER() optional_header.Magic = 0x10B if not magic else magic self.file_alignment = file_alignment @@ -356,7 +341,7 @@ def gen_optional_header( + len(b"PE\x00\x00") + len(self.file_header) + len(optional_header) - + len(pestruct.IMAGE_SECTION_HEADER), + + len(c_pe.IMAGE_SECTION_HEADER), blocksize=file_alignment, ) @@ -404,7 +389,7 @@ def section( number_of_relocations: int = 0, number_of_linenumbers: int = 0, characteristics: int = 0x68000020, - ) -> cstruct: + ) -> c_pe.IMAGE_SECTION_HEADER: """Build a new section for the PE. The default characteristics of the new section will be: @@ -430,14 +415,18 @@ def section( """ if len(name) > 8: - raise BuildSectionException("section names can't be longer than 8 characters") + raise BuildSectionException( + "section names can't be longer than 8 characters" + ) if isinstance(name, str): name = name.encode() - section_header = pestruct.IMAGE_SECTION_HEADER() + section_header = c_pe.IMAGE_SECTION_HEADER() - pointer_to_raw_data = utils.align_int(integer=pointer_to_raw_data, blocksize=self.file_alignment) + pointer_to_raw_data = utils.align_int( + integer=pointer_to_raw_data, blocksize=self.file_alignment + ) section_header.Name = name + utils.pad(size=8 - len(name)) section_header.VirtualSize = virtual_size @@ -463,7 +452,9 @@ def pe_size(self) -> int: The size of the PE. """ - last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] + last_section = self.pe.patched_sections[ + next(reversed(self.pe.patched_sections)) + ] va = last_section.virtual_address size = last_section.virtual_size diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py index f7e4715..7017c9a 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/helpers/exports.py @@ -4,8 +4,7 @@ from io import BytesIO from typing import TYPE_CHECKING -# Local imports -from dissect.executable.pe.helpers.c_pe import pestruct +from dissect.executable.pe.c_pe import c_pe if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection @@ -30,7 +29,11 @@ def __str__(self) -> str: return self.name.decode() if self.name else self.ordinal def __repr__(self) -> str: - return f"" if self.name else f"" + return ( + f"" + if self.name + else f"" + ) class ExportManager: @@ -41,41 +44,51 @@ def __init__(self, pe: PE, section: PESection): self.parse_exports() - def parse_exports(self): + def parse_exports(self) -> None: """Parse the export directory of the PE file. This function will store every export function within the PE file as an `ExportFunction` object containing the name (if available), the call ordinal, and the function address. """ - export_entry_va = self.pe.directory_va(pestruct.IMAGE_DIRECTORY_ENTRY_EXPORT) - export_entry = BytesIO(self.pe.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_EXPORT)) - export_directory = pestruct.IMAGE_EXPORT_DIRECTORY(export_entry) + export_entry_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) + export_entry = BytesIO( + self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) + ) + export_directory = c_pe.IMAGE_EXPORT_DIRECTORY(export_entry) # Seek to the offset of the export name export_entry.seek(export_directory.Name - export_entry_va) - self.export_name = pestruct.char[None](export_entry) + self.export_name = c_pe.char[None](export_entry) # Create a list of adresses for the exported functions export_entry.seek(export_directory.AddressOfFunctions - export_entry_va) - export_addresses = pestruct.uint32[export_directory.NumberOfFunctions].read(export_entry) + export_addresses = c_pe.uint32[export_directory.NumberOfFunctions].read( + export_entry + ) # Create a list of addresses for the exported functions that have associated names export_entry.seek(export_directory.AddressOfNames - export_entry_va) - export_names = pestruct.uint32[export_directory.NumberOfNames].read(export_entry) + export_names = c_pe.uint32[export_directory.NumberOfNames].read(export_entry) # Create a list of addresses for the ordinals associated with the functions export_entry.seek(export_directory.AddressOfNameOrdinals - export_entry_va) - export_ordinals = pestruct.uint16[export_directory.NumberOfNames].read(export_entry) + export_ordinals = c_pe.uint16[export_directory.NumberOfNames].read(export_entry) # Iterate over the export functions and store the information export_entry.seek(export_directory.AddressOfFunctions - export_entry_va) for idx, address in enumerate(export_addresses): if idx in export_ordinals: - export_entry.seek(export_names[export_ordinals.index(idx)] - export_entry_va) - export_name = pestruct.char[None](export_entry) - self.exports[export_name.decode()] = ExportFunction(ordinal=idx + 1, address=address, name=export_name) + export_entry.seek( + export_names[export_ordinals.index(idx)] - export_entry_va + ) + export_name = c_pe.char[None](export_entry) + self.exports[export_name.decode()] = ExportFunction( + ordinal=idx + 1, address=address, name=export_name + ) else: export_name = None - self.exports[str(idx + 1)] = ExportFunction(ordinal=idx + 1, address=address, name=export_name) + self.exports[str(idx + 1)] = ExportFunction( + ordinal=idx + 1, address=address, name=export_name + ) def add(self): raise NotImplementedError diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index 1880644..84e3213 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -3,13 +3,11 @@ import struct from collections import OrderedDict from io import BytesIO -from typing import TYPE_CHECKING, BinaryIO, Generator +from typing import TYPE_CHECKING, BinaryIO, Iterator +from dissect.executable.pe.c_pe import c_pe from dissect.executable.pe.helpers import utils -# Local imports -from dissect.executable.pe.helpers.c_pe import pestruct - if TYPE_CHECKING: from dissect.cstruct.cstruct import cstruct @@ -28,7 +26,14 @@ class ImportModule: first_thunk: The virtual address of the first thunk. """ - def __init__(self, name: bytes, import_descriptor: cstruct, module_va: int, name_va: int, first_thunk: int): + def __init__( + self, + name: bytes, + import_descriptor: c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT, + module_va: int, + name_va: int, + first_thunk: int, + ): self.name = name self.import_descriptor = import_descriptor self.module_va = module_va @@ -51,7 +56,12 @@ class ImportFunction: thunkdata: The thunkdata of the import function as a cstruct object. """ - def __init__(self, pe: PE, thunkdata: cstruct, name: str = ""): + def __init__( + self, + pe: PE, + thunkdata: c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64, + name: str = "", + ): self.pe = pe self.thunkdata = thunkdata self._name = name @@ -71,7 +81,7 @@ def name(self) -> str: if not ordinal: self.pe.seek(self.thunkdata.u1.AddressOfData + 2) - entry = pestruct.char[None](self.pe).decode() + entry = c_pe.char[None](self.pe).decode() else: entry = ordinal @@ -107,20 +117,24 @@ def __init__(self, pe: PE, section: PESection): self.parse_imports() - def parse_imports(self): + def parse_imports(self) -> None: """Parse the imports of the PE file. The imports are in turn added to the `imports` attribute so they can be accessed by the user. """ - import_data = BytesIO(self.pe.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_IMPORT)) + import_data = BytesIO( + self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT) + ) import_data.seek(0) # Loop over the entries - for descriptor_va, import_descriptor in self.import_descriptors(import_data=import_data): - if import_descriptor.Name != 0xFFFFF800 and import_descriptor.Name != 0x0: + for descriptor_va, import_descriptor in self.import_descriptors( + import_data=import_data + ): + if import_descriptor.Name not in [0xFFFFF800, 0x0]: self.pe.seek(import_descriptor.Name) - modulename = pestruct.char[None](self.pe) + modulename = c_pe.char[None](self.pe) # Use the OriginalFirstThunk if available, FirstThunk otherwise first_thunk = ( @@ -137,11 +151,15 @@ def parse_imports(self): ) for thunkdata in self.parse_thunks(offset=first_thunk): - module.functions.append(ImportFunction(pe=self.pe, thunkdata=thunkdata)) + module.functions.append( + ImportFunction(pe=self.pe, thunkdata=thunkdata) + ) self.imports[modulename.decode()] = module - def import_descriptors(self, import_data: BinaryIO) -> Generator[tuple[int, cstruct], None, None]: + def import_descriptors( + self, import_data: BinaryIO + ) -> Iterator[tuple[int, c_pe.IMAGE_IMPORT_DESCRIPTOR], None, None]: """Parse the import descriptors of the PE file. Args: @@ -153,13 +171,15 @@ def import_descriptors(self, import_data: BinaryIO) -> Generator[tuple[int, cstr while True: try: - import_descriptor = pestruct.IMAGE_IMPORT_DESCRIPTOR(import_data) + import_descriptor = c_pe.IMAGE_IMPORT_DESCRIPTOR(import_data) except EOFError: break yield import_data.tell(), import_descriptor - def parse_thunks(self, offset: int) -> Generator[cstruct, None, None]: + def parse_thunks( + self, offset: int + ) -> Iterator[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64, None, None]: """Parse the import thunks for every module. Args: @@ -178,7 +198,7 @@ def parse_thunks(self, offset: int) -> Generator[cstruct, None, None]: yield thunkdata - def add(self, dllname: str, functions: list): + def add(self, dllname: str, functions: list[str]) -> None: """Add the given module and its functions to the PE. Args: @@ -186,15 +206,23 @@ def add(self, dllname: str, functions: list): functions: A `list` of function names belonging to the module. """ - self.last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] + self.last_section = self.pe.patched_sections[ + next(reversed(self.pe.patched_sections)) + ] # Build a dummy import module self.imports[dllname] = ImportModule( - name=dllname.encode(), import_descriptor=None, module_va=0, name_va=0, first_thunk=0 + name=dllname.encode(), + import_descriptor=None, + module_va=0, + name_va=0, + first_thunk=0, ) # Build the dummy module functions for function in functions: - self.pe.imports[dllname].functions.append(ImportFunction(pe=self.pe, thunkdata=None, name=function)) + self.pe.imports[dllname].functions.append( + ImportFunction(pe=self.pe, thunkdata=None, name=function) + ) # Rebuild the import table with the new import module and functions self.build_import_table() @@ -202,7 +230,7 @@ def add(self, dllname: str, functions: list): def delete(self, dllname: str, functions: list): raise NotImplementedError - def build_import_table(self): + def build_import_table(self) -> None: """Function to rebuild the import table after a change has been made to the PE imports. Currently we're using the .idata section to store the imports, there might be a better way to do this but for @@ -223,25 +251,28 @@ def build_import_table(self): # Build the module imports and get the RVA of the first thunk to generate an import descriptor first_thunk_rva = self._build_module_imports(functions=module.functions) import_descriptor = self._build_import_descriptor( - first_thunk_rva=first_thunk_rva, name_rva=self.pe.optional_header.SizeOfImage + name_offset + first_thunk_rva=first_thunk_rva, + name_rva=self.pe.optional_header.SizeOfImage + name_offset, ) import_descriptors.append(import_descriptor) datadirectory_size = 0 - for idx, descriptor in enumerate(import_descriptors): - if idx == 0: - # Take note of the RVA of the first import descriptor - import_rva = self.pe.optional_header.SizeOfImage + len(self.import_data) + + # Take note of the RVA of the first import descriptor + import_rva = self.pe.optional_header.SizeOfImage + len(self.import_data) + for descriptor in import_descriptors: self.import_data += descriptor.dumps() datadirectory_size += len(descriptor) # Create a new section - section_data = utils.align_data(data=self.import_data, blocksize=self.pe.file_alignment) - size = len(self.import_data) + pestruct.IMAGE_SECTION_HEADER.size + section_data = utils.align_data( + data=self.import_data, blocksize=self.pe.file_alignment + ) + size = len(self.import_data) + c_pe.IMAGE_SECTION_HEADER.size self.pe.add_section( name=".idata", data=section_data, - datadirectory=pestruct.IMAGE_DIRECTORY_ENTRY_IMPORT, + datadirectory=c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT, datadirectory_rva=import_rva, datadirectory_size=datadirectory_size, size=size, @@ -290,14 +321,15 @@ def _build_thunkdata(self, import_rvas: list[int]) -> bytes: rva += self.pe.optional_header.SizeOfImage thunkdata += ( struct.pack(" cstru The image import descriptor as a `cstruct` object. """ - new_import_descriptor = pestruct.IMAGE_IMPORT_DESCRIPTOR() + new_import_descriptor = c_pe.IMAGE_IMPORT_DESCRIPTOR() new_import_descriptor.OriginalFirstThunk = first_thunk_rva new_import_descriptor.TimeDateStamp = 0 diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index 5bdc883..c3fc4ee 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -4,9 +4,8 @@ from io import BytesIO from typing import TYPE_CHECKING -# Local imports +from dissect.executable.pe.c_pe import c_pe from dissect.executable.pe.helpers import utils -from dissect.executable.pe.helpers.c_pe import pestruct from dissect.executable.pe.helpers.sections import PESection if TYPE_CHECKING: @@ -64,11 +63,15 @@ def pe_size(self) -> int: The new size of the PE as an `int`. """ - last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] + last_section = self.pe.patched_sections[ + next(reversed(self.pe.patched_sections)) + ] va = last_section.virtual_address size = last_section.virtual_size - return utils.align_int(integer=va + size, blocksize=self.pe.optional_header.SectionAlignment) + return utils.align_int( + integer=va + size, blocksize=self.pe.optional_header.SectionAlignment + ) def seek(self, address: int): """Seek that is used to seek to a virtual address in the patched PE file. @@ -85,7 +88,9 @@ def _build_section_table(self): if self.patched_pe.tell() < self.pe.section_header_offset: # Pad the patched file with null bytes until we reach the section header offset - self.patched_pe.write(utils.pad(size=self.pe.section_header_offset - self.patched_pe.tell())) + self.patched_pe.write( + utils.pad(size=self.pe.section_header_offset - self.patched_pe.tell()) + ) # Write the section headers for section in self.pe.patched_sections.values(): @@ -127,7 +132,7 @@ def _patch_import_rvas(self): patched_import_data = bytearray() # Get the directory entry virtual adddress, this is the updated address if it has been patched. - directory_va = self.pe.directory_va(pestruct.IMAGE_DIRECTORY_ENTRY_IMPORT) + directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT) if not directory_va: return @@ -135,7 +140,9 @@ def _patch_import_rvas(self): # new RVA's section = self.pe.patched_section(va=directory_va) directory_offset = directory_va - section.virtual_address - original_directory_va = self.pe.sections[section.name].virtual_address + directory_offset + original_directory_va = ( + self.pe.sections[section.name].virtual_address + directory_offset + ) # Loop over the imports of the PE to patch the RVA's of the import descriptors and the associated thunkdata # entries @@ -167,10 +174,13 @@ def _patch_import_rvas(self): # and use it to also select the patched virtual address of this section that the RVA is located in for name, section in self.pe.sections.items(): if thunkdata.u1.AddressOfData in range( - section.virtual_address, section.virtual_address + section.virtual_size + section.virtual_address, + section.virtual_address + section.virtual_size, ): virtual_address = section.virtual_address - new_virtual_address = self.pe.patched_sections[name].virtual_address + new_virtual_address = self.pe.patched_sections[ + name + ].virtual_address break # Calculate the offset using the VA of the section and update the thunkdata @@ -195,23 +205,31 @@ def _patch_import_rvas(self): def _patch_export_rvas(self): """Function to patch the RVAs of the export directory and the associated function and name RVA's.""" - directory_va = self.pe.directory_va(pestruct.IMAGE_DIRECTORY_ENTRY_EXPORT) + directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) if not directory_va: return self.seek(directory_va) - export_directory = pestruct.IMAGE_EXPORT_DIRECTORY(self.patched_pe) + export_directory = c_pe.IMAGE_EXPORT_DIRECTORY(self.patched_pe) # Get the original VA of the section the import directory is residing in, this value is used to calculate the # new RVA's section = self.pe.patched_section(va=directory_va) directory_offset = directory_va - section.virtual_address - original_directory_va = self.pe.sections[section.name].virtual_address + directory_offset + original_directory_va = ( + self.pe.sections[section.name].virtual_address + directory_offset + ) name_offset = export_directory.Name - original_directory_va - address_of_functions_offset = export_directory.AddressOfFunctions - original_directory_va - address_of_names_offset = export_directory.AddressOfNames - original_directory_va - address_of_name_ordinals = export_directory.AddressOfNameOrdinals - original_directory_va + address_of_functions_offset = ( + export_directory.AddressOfFunctions - original_directory_va + ) + address_of_names_offset = ( + export_directory.AddressOfNames - original_directory_va + ) + address_of_name_ordinals = ( + export_directory.AddressOfNameOrdinals - original_directory_va + ) export_directory.Name = directory_va + name_offset export_directory.AddressOfFunctions = directory_va + address_of_functions_offset @@ -226,13 +244,17 @@ def _patch_export_rvas(self): new_function_rvas = [] function_rvas = bytearray() self.seek(export_directory.AddressOfFunctions) - export_addresses = pestruct.uint32[export_directory.NumberOfFunctions].read(self.patched_pe) + export_addresses = c_pe.uint32[export_directory.NumberOfFunctions].read( + self.patched_pe + ) for address in export_addresses: section = self.pe.section(va=address) if not section: continue address_offset = address - section.virtual_address - new_address = self.pe.patched_sections[section.name].virtual_address + address_offset + new_address = ( + self.pe.patched_sections[section.name].virtual_address + address_offset + ) new_function_rvas.append(new_address) for rva in new_function_rvas: @@ -245,11 +267,13 @@ def _patch_export_rvas(self): new_name_rvas = [] name_rvas = bytearray() self.seek(export_directory.AddressOfNames) - export_names = pestruct.uint32[export_directory.NumberOfNames].read(self.patched_pe) + export_names = c_pe.uint32[export_directory.NumberOfNames].read(self.patched_pe) for name_address in export_names: section = self.pe.section(va=name_address) address_offset = name_address - section.virtual_address - new_address = self.pe.patched_sections[section.name].virtual_address + address_offset + new_address = ( + self.pe.patched_sections[section.name].virtual_address + address_offset + ) new_name_rvas.append(new_address) for name_rva in new_name_rvas: @@ -257,19 +281,21 @@ def _patch_export_rvas(self): self.seek(export_directory.AddressOfNames) self.patched_pe.write(name_rvas) - # self.pe.optional_header.DataDirectory[pestruct.IMAGE_DIRECTORY_ENTRY_EXPORT].Size = len(name_rvas) + # self.pe.optional_header.DataDirectory[c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT].Size = len(name_rvas) def _patch_rsrc_rvas(self): """Function to patch the RVAs of the resource directory and the associated resource data RVA's.""" - directory_va = self.pe.directory_va(pestruct.IMAGE_DIRECTORY_ENTRY_RESOURCE) + directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE) if not directory_va: return section_data = BytesIO() self.seek(directory_va) - for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"]): + for rsrc_entry in sorted( + self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"] + ): entry_offset = rsrc_entry["offset"] entry = rsrc_entry["entry"] @@ -295,7 +321,7 @@ def _patch_rsrc_rvas(self): def _patch_tls_rvas(self): """Function to patch the RVAs of the TLS directory and the associated TLS callbacks.""" - directory_va = self.pe.directory_va(pestruct.IMAGE_DIRECTORY_ENTRY_TLS) + directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_TLS) if not directory_va: return @@ -306,24 +332,39 @@ def _patch_tls_rvas(self): # Patch the TLS StartAddressOfRawData and EndAddressOfRawData section = self.pe.section(va=tls_directory.StartAddressOfRawData - image_base) - start_address_offset = tls_directory.StartAddressOfRawData - section.virtual_address + start_address_offset = ( + tls_directory.StartAddressOfRawData - section.virtual_address + ) tls_directory.StartAddressOfRawData = ( - self.pe.patched_sections[section.name].virtual_address + start_address_offset + self.pe.patched_sections[section.name].virtual_address + + start_address_offset + ) + end_address_offset = ( + tls_directory.EndAddressOfRawData - tls_directory.StartAddressOfRawData + ) + tls_directory.EndAddressOfRawData = ( + tls_directory.StartAddressOfRawData + end_address_offset ) - end_address_offset = tls_directory.EndAddressOfRawData - tls_directory.StartAddressOfRawData - tls_directory.EndAddressOfRawData = tls_directory.StartAddressOfRawData + end_address_offset # Patch the TLS callbacks address section = self.pe.section(va=tls_directory.AddressOfCallBacks - image_base) - address_of_callbacks_offset = tls_directory.AddressOfCallBacks - section.virtual_address + address_of_callbacks_offset = ( + tls_directory.AddressOfCallBacks - section.virtual_address + ) tls_directory.AddressOfCallBacks = ( - self.pe.patched_sections[section.name].virtual_address + address_of_callbacks_offset + self.pe.patched_sections[section.name].virtual_address + + address_of_callbacks_offset ) # Patch the TLS AddressOfIndex section = self.pe.section(va=tls_directory.AddressOfIndex - image_base) - address_of_index_offset = tls_directory.AddressOfIndex - self.pe.sections[section.name].virtual_address - tls_directory.AddressOfIndex = self.pe.sections[section.name].virtual_address + address_of_index_offset + address_of_index_offset = ( + tls_directory.AddressOfIndex + - self.pe.sections[section.name].virtual_address + ) + tls_directory.AddressOfIndex = ( + self.pe.sections[section.name].virtual_address + address_of_index_offset + ) # Write the patched TLS directory to the new PE self.seek(directory_va) @@ -340,5 +381,7 @@ def _get_tls_attribute_section(self, va: int) -> PESection: """ for name, section in self.pe.sections.items(): - if va in range(section.virtual_address, section.virtual_address + section.virtual_size): + if va in range( + section.virtual_address, section.virtual_address + section.virtual_size + ): return section diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/helpers/relocations.py index 465a66b..ef4d16c 100644 --- a/dissect/executable/pe/helpers/relocations.py +++ b/dissect/executable/pe/helpers/relocations.py @@ -3,8 +3,7 @@ from io import BytesIO from typing import TYPE_CHECKING -# Local imports -from dissect.executable.pe.helpers.c_pe import pestruct +from dissect.executable.pe.c_pe import c_pe if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection @@ -29,24 +28,32 @@ def __init__(self, pe: PE, section: PESection): def parse_relocations(self): """Parse the relocation table of the PE file.""" - reloc_data = BytesIO(self.pe.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_BASERELOC)) + reloc_data = BytesIO( + self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC) + ) reloc_data_size = reloc_data.getbuffer().nbytes while reloc_data.tell() < reloc_data_size: - reloc_directory = pestruct.IMAGE_BASE_RELOCATION(reloc_data) + reloc_directory = c_pe.IMAGE_BASE_RELOCATION(reloc_data) if not reloc_directory.VirtualAddress: # End of relocation entries break # Each entry consists of 2 bytes - number_of_entries = (reloc_directory.SizeOfBlock - len(reloc_directory.dumps())) // 2 + number_of_entries = ( + reloc_directory.SizeOfBlock - len(reloc_directory.dumps()) + ) // 2 entries = [] for _ in range(0, number_of_entries): - entry = pestruct.uint16(reloc_data) + entry = c_pe.uint16(reloc_data) if entry: entries.append(entry) self.relocations.append( - {"rva:": reloc_directory.VirtualAddress, "number_of_entries": number_of_entries, "entries": entries} + { + "rva:": reloc_directory.VirtualAddress, + "number_of_entries": number_of_entries, + "entries": entries, + } ) def add(self): diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index ee011fb..8b23afb 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -5,9 +5,7 @@ from typing import TYPE_CHECKING, Iterator from dissect.executable.exception import ResourceException - -# Local imports -from dissect.executable.pe.helpers.c_pe import pestruct +from dissect.executable.pe.c_pe import c_pe if TYPE_CHECKING: from dissect.cstruct.cstruct import cstruct @@ -36,8 +34,12 @@ def __init__(self, pe: PE, section: PESection): def parse_rsrc(self): """Parse the resource directory entry of the PE file.""" - rsrc_data = BytesIO(self.pe.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_RESOURCE)) - self.resources = self._read_resource(rc_type="_root", data=rsrc_data, offset=0, level=1) + rsrc_data = BytesIO( + self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE) + ) + self.resources = self._read_resource( + rc_type="_root", data=rsrc_data, offset=0, level=1 + ) def _read_entries(self, data: bytes, directory: cstruct) -> list: """Read the entries within the resource directory. @@ -53,8 +55,10 @@ def _read_entries(self, data: bytes, directory: cstruct) -> list: entries = [] for _ in range(0, directory.NumberOfNamedEntries + directory.NumberOfIdEntries): entry_offset = data.tell() - entry = pestruct.IMAGE_RESOURCE_DIRECTORY_ENTRY(data) - self.raw_resources.append({"offset": entry_offset, "entry": entry, "data_offset": entry_offset}) + entry = c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY(data) + self.raw_resources.append( + {"offset": entry_offset, "entry": entry, "data_offset": entry_offset} + ) entries.append(entry) return entries @@ -70,7 +74,7 @@ def _handle_data_entry(self, data: bytes, entry: cstruct, rc_type: str) -> Resou """ data.seek(entry.OffsetToDirectory) - data_entry = pestruct.IMAGE_RESOURCE_DATA_ENTRY(data) + data_entry = c_pe.IMAGE_RESOURCE_DATA_ENTRY(data) self.pe.seek(data_entry.OffsetToData) data = self.pe.read(data_entry.Size) raw_offset = data_entry.OffsetToData - self.section.virtual_address @@ -93,7 +97,9 @@ def _handle_data_entry(self, data: bytes, entry: cstruct, rc_type: str) -> Resou ) return rsrc - def _read_resource(self, data: bytes, offset: int, rc_type: str, level: int = 1) -> dict: + def _read_resource( + self, data: bytes, offset: int, rc_type: str, level: int = 1 + ) -> dict: """Recursively read the resources within the PE file. Each resource is added to the dictionary that is available to the user, as well as a list of @@ -112,28 +118,35 @@ def _read_resource(self, data: bytes, offset: int, rc_type: str, level: int = 1) resource = OrderedDict() data.seek(offset) - directory = pestruct.IMAGE_RESOURCE_DIRECTORY(data) - self.raw_resources.append({"offset": offset, "entry": directory, "data_offset": offset}) + directory = c_pe.IMAGE_RESOURCE_DIRECTORY(data) + self.raw_resources.append( + {"offset": offset, "entry": directory, "data_offset": offset} + ) entries = self._read_entries(data, directory) for entry in entries: if level == 1: - rc_type = pestruct.ResourceID(entry.Id).name + rc_type = c_pe.ResourceID(entry.Id).name else: if entry.NameIsString: data.seek(entry.NameOffset) - name_len = pestruct.uint16(data) - rc_type = pestruct.wchar[name_len](data) + name_len = c_pe.uint16(data) + rc_type = c_pe.wchar[name_len](data) else: rc_type = str(entry.Id) if entry.DataIsDirectory: resource[rc_type] = self._read_resource( - data=data, offset=entry.OffsetToDirectory, rc_type=rc_type, level=level + 1 + data=data, + offset=entry.OffsetToDirectory, + rc_type=rc_type, + level=level + 1, ) else: - resource[rc_type] = self._handle_data_entry(data=data, entry=entry, rc_type=rc_type) + resource[rc_type] = self._handle_data_entry( + data=data, entry=entry, rc_type=rc_type + ) return resource @@ -234,7 +247,9 @@ def update_section(self, update_offset: int): new_size = 0 section_data = b"" - for idx, resource in enumerate(self.parse_resources(resources=self.pe.resources)): + for idx, resource in enumerate( + self.parse_resources(resources=self.pe.resources) + ): if idx == 0: # Use the offset of the first resource to account for the size of the directory header header_size = resource.offset - self.section.virtual_address @@ -264,7 +279,7 @@ def update_section(self, update_offset: int): self.section.data = section_data -class Resource(ResourceManager): +class Resource: """Base class representing a resource entry in the PE file. Args: @@ -373,7 +388,9 @@ def data(self, value: bytes): prev_offset = 0 prev_size = 0 - for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"]): + for rsrc_entry in sorted( + self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"] + ): entry_offset = rsrc_entry["offset"] entry = rsrc_entry["entry"] @@ -409,7 +426,9 @@ def data(self, value: bytes): # Update the section data and size self.section.data = data - self.pe.optional_header.DataDirectory[pestruct.IMAGE_DIRECTORY_ENTRY_RESOURCE].Size = len(data) + self.pe.optional_header.DataDirectory[ + c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE + ].Size = len(data) def __str__(self) -> str: return str(self.name) diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 792ad82..6385246 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -4,10 +4,8 @@ from typing import TYPE_CHECKING from dissect.executable.exception import BuildSectionException - -# Local imports +from dissect.executable.pe.c_pe import c_pe from dissect.executable.pe.helpers import utils -from dissect.executable.pe.helpers.c_pe import pestruct if TYPE_CHECKING: from dissect.cstruct.cstruct import cstruct @@ -69,7 +67,9 @@ def size(self, value: int): """ self.virtual_size = value - self.size_of_raw_data = utils.align_int(integer=value, blocksize=self.pe.file_alignment) + self.size_of_raw_data = utils.align_int( + integer=value, blocksize=self.pe.file_alignment + ) @property def virtual_address(self) -> int: @@ -143,8 +143,12 @@ def size_of_raw_data(self, value: int): value: The size of the data. """ - self._size_of_raw_data = utils.align_int(integer=value, blocksize=self.pe.file_alignment) - self.section.SizeOfRawData = utils.align_int(integer=value, blocksize=self.pe.file_alignment) + self._size_of_raw_data = utils.align_int( + integer=value, blocksize=self.pe.file_alignment + ) + self.section.SizeOfRawData = utils.align_int( + integer=value, blocksize=self.pe.file_alignment + ) @property def data(self) -> bytes: @@ -187,8 +191,12 @@ def data(self, value: bytes): if section.virtual_address == prev_va: continue - pointer_to_raw_data = utils.align_int(integer=prev_ptr + prev_size, blocksize=self.pe.file_alignment) - virtual_address = utils.align_int(integer=prev_va + prev_vsize, blocksize=self.pe.section_alignment) + pointer_to_raw_data = utils.align_int( + integer=prev_ptr + prev_size, blocksize=self.pe.file_alignment + ) + virtual_address = utils.align_int( + integer=prev_va + prev_vsize, blocksize=self.pe.section_alignment + ) if section.virtual_address < virtual_address: """Set the virtual address and raw pointer of the section to the new values, but only do so if the @@ -242,7 +250,7 @@ def build_section( if isinstance(name, str): name = name.encode() - section_header = pestruct.IMAGE_SECTION_HEADER() + section_header = c_pe.IMAGE_SECTION_HEADER() section_header.Name = name + utils.pad(size=8 - len(name)) section_header.VirtualSize = virtual_size diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/helpers/tls.py index 3e3929d..a8fa402 100644 --- a/dissect/executable/pe/helpers/tls.py +++ b/dissect/executable/pe/helpers/tls.py @@ -3,8 +3,7 @@ from io import BytesIO from typing import TYPE_CHECKING -# Local imports -from dissect.executable.pe.helpers.c_pe import pestruct +from dissect.executable.pe.c_pe import c_pe if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection @@ -30,7 +29,9 @@ def __init__(self, pe: PE, section: PESection): def parse_tls(self): """Parse the TLS directory entry of the PE file when present.""" - tls_data = BytesIO(self.pe.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_TLS)) + tls_data = BytesIO( + self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_TLS) + ) self.tls = self.pe.image_tls_directory(tls_data) self.pe.seek(self.tls.AddressOfCallBacks - self.pe.optional_header.ImageBase) @@ -73,7 +74,8 @@ def read_data(self) -> bytes: """ return self.pe.virtual_read( - address=self.tls.StartAddressOfRawData - self.pe.optional_header.ImageBase, size=self.size + address=self.tls.StartAddressOfRawData - self.pe.optional_header.ImageBase, + size=self.size, ) @property @@ -105,7 +107,9 @@ def data(self, value): section_data.write(self.tls.dumps()) # Write the new TLS data to the section - start_address_rva = self.tls.StartAddressOfRawData - self.pe.optional_header.ImageBase + start_address_rva = ( + self.tls.StartAddressOfRawData - self.pe.optional_header.ImageBase + ) start_address_section_offset = start_address_rva - self.section.virtual_address section_data.seek(start_address_section_offset) section_data.write(self._data) diff --git a/dissect/executable/pe/helpers/utils.py b/dissect/executable/pe/helpers/utils.py index 93ce896..b73ee29 100644 --- a/dissect/executable/pe/helpers/utils.py +++ b/dissect/executable/pe/helpers/utils.py @@ -10,7 +10,11 @@ def align_data(data: bytes, blocksize: int) -> bytes: """ needs_alignment = len(data) % blocksize - return data if not needs_alignment else data + ((blocksize - needs_alignment) * b"\x00") + return ( + data + if not needs_alignment + else data + ((blocksize - needs_alignment) * b"\x00") + ) def align_int(integer: int, blocksize: int) -> int: diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index cbb19fb..02da94b 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections import OrderedDict -from datetime import datetime +from datetime import datetime, timezone from io import BytesIO from typing import TYPE_CHECKING, BinaryIO, Tuple @@ -12,6 +12,7 @@ InvalidVA, ResourceException, ) +from dissect.executable.pe.c_pe import c_cv_info, c_pe from dissect.executable.pe.helpers import ( exports, imports, @@ -23,9 +24,6 @@ utils, ) -# Local imports -from dissect.executable.pe.helpers.c_pe import cv_info_struct, pestruct - if TYPE_CHECKING: from dissect.cstruct.cstruct import cstruct from dissect.cstruct.types.enum import EnumInstance @@ -80,7 +78,9 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.base_address = self.optional_header.ImageBase - self.timestamp = datetime.utcfromtimestamp(self.file_header.TimeDateStamp) + self.timestamp = datetime.fromtimestamp( + self.file_header.TimeDateStamp, tz=timezone.utc + ) # Parse the section header self.parse_section_header() @@ -96,7 +96,7 @@ def _valid(self) -> bool: """ self.pe_file.seek(self.mz_header.e_lfanew) - return True if pestruct.uint32(self.pe_file) == 0x4550 else False + return True if c_pe.uint32(self.pe_file) == 0x4550 else False def parse_headers(self): """Function to parse the basic PE headers: @@ -111,22 +111,22 @@ def parse_headers(self): InvalidArchitecture if the architecture is not supported or unknown. """ - self.mz_header = pestruct.IMAGE_DOS_HEADER(self.pe_file) + self.mz_header = c_pe.IMAGE_DOS_HEADER(self.pe_file) if not self._valid(): raise InvalidPE("file is not a valid PE file") - self.file_header = pestruct.IMAGE_FILE_HEADER(self.pe_file) + self.file_header = c_pe.IMAGE_FILE_HEADER(self.pe_file) image_nt_headers_offset = self.mz_header.e_lfanew self.pe_file.seek(image_nt_headers_offset) # Set the architecture specific settings self._set_pe_architecture() - if self.file_header.Machine == pestruct.MachineType.IMAGE_FILE_MACHINE_AMD64: - self.nt_headers = pestruct.IMAGE_NT_HEADERS64(self.pe_file) + if self.file_header.Machine == c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64: + self.nt_headers = c_pe.IMAGE_NT_HEADERS64(self.pe_file) else: - self.nt_headers = pestruct.IMAGE_NT_HEADERS(self.pe_file) + self.nt_headers = c_pe.IMAGE_NT_HEADERS(self.pe_file) self.optional_header = self.nt_headers.OptionalHeader @@ -137,18 +137,20 @@ def _set_pe_architecture(self): InvalidArchitecture if the architecture is not supported or unknown. """ - if self.file_header.Machine == pestruct.MachineType.IMAGE_FILE_MACHINE_AMD64: - self.image_thunk_data = pestruct.IMAGE_THUNK_DATA64 - self.image_tls_directory = pestruct.IMAGE_TLS_DIRECTORY64 + if self.file_header.Machine == c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64: + self.image_thunk_data = c_pe.IMAGE_THUNK_DATA64 + self.image_tls_directory = c_pe.IMAGE_TLS_DIRECTORY64 self._high_bit = 1 << 63 - self.read_address = pestruct.uint64 - elif self.file_header.Machine == pestruct.MachineType.IMAGE_FILE_MACHINE_I386: - self.image_thunk_data = pestruct.IMAGE_THUNK_DATA32 - self.image_tls_directory = pestruct.IMAGE_TLS_DIRECTORY32 + self.read_address = c_pe.uint64 + elif self.file_header.Machine == c_pe.MachineType.IMAGE_FILE_MACHINE_I386: + self.image_thunk_data = c_pe.IMAGE_THUNK_DATA32 + self.image_tls_directory = c_pe.IMAGE_TLS_DIRECTORY32 self._high_bit = 1 << 31 - self.read_address = pestruct.uint32 + self.read_address = c_pe.uint32 else: - raise InvalidArchitecture(f"Invalid architecture found: {self.file_header.Machine:02x}") + raise InvalidArchitecture( + f"Invalid architecture found: {self.file_header.Machine:02x}" + ) def parse_section_header(self): """Parse the sections within the PE file.""" @@ -158,11 +160,15 @@ def parse_section_header(self): for _ in range(self.file_header.NumberOfSections): # Keep track of the last section offset offset = self.pe_file.tell() - section_header = pestruct.IMAGE_SECTION_HEADER(self) + section_header = c_pe.IMAGE_SECTION_HEADER(self.pe_file) section_name = section_header.Name.decode().strip("\x00") # Take note of the sections, keep track of any patches seperately - self.sections[section_name] = sections.PESection(pe=self, section=section_header, offset=offset) - self.patched_sections[section_name] = sections.PESection(pe=self, section=section_header, offset=offset) + self.sections[section_name] = sections.PESection( + pe=self, section=section_header, offset=offset + ) + self.patched_sections[section_name] = sections.PESection( + pe=self, section=section_header, offset=offset + ) self.last_section_offset = self.sections[next(reversed(self.sections))].offset @@ -179,7 +185,10 @@ def section(self, va: int = 0, name: str = "") -> sections.PESection: if not name: for section in self.sections.values(): - if va in range(section.virtual_address, section.virtual_address + section.virtual_size): + if va in range( + section.virtual_address, + section.virtual_address + section.virtual_size, + ): return section else: return self.sections[name] @@ -197,7 +206,10 @@ def patched_section(self, va: int = 0, name: str = "") -> sections.PESection: if not name: for section in self.patched_sections.values(): - if va in range(section.virtual_address, section.virtual_address + section.virtual_size): + if va in range( + section.virtual_address, + section.virtual_address + section.virtual_size, + ): return section else: return self.patched_sections[name] @@ -214,7 +226,10 @@ def datadirectory_section(self, index: int) -> sections.PESection: va = self.directory_va(index=index) for _, section in self.patched_sections.items(): - if va >= section.virtual_address and va < section.virtual_address + section.virtual_size: + if ( + va >= section.virtual_address + and va < section.virtual_address + section.virtual_size + ): return section raise InvalidVA(f"VA not found in sections: {va:#04x}") @@ -230,37 +245,40 @@ def parse_directories(self): - Thread Local Storage (TLS) Callbacks """ - for idx in range(pestruct.IMAGE_NUMBEROF_DIRECTORY_ENTRIES): + for idx in range(c_pe.IMAGE_NUMBEROF_DIRECTORY_ENTRIES): if not self.has_directory(index=idx): continue # Take note of the current directory VA so we can dynamically update it when resizing sections section = self.datadirectory_section(index=idx) - directory_va_offset = self.optional_header.DataDirectory[idx].VirtualAddress - section.virtual_address + directory_va_offset = ( + self.optional_header.DataDirectory[idx].VirtualAddress + - section.virtual_address + ) section.directories[idx] = directory_va_offset # Parse the Import Address Table (IAT) - if idx == pestruct.IMAGE_DIRECTORY_ENTRY_IMPORT: + if idx == c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT: self.import_mgr = imports.ImportManager(pe=self, section=section) self.imports = self.import_mgr.imports - if idx == pestruct.IMAGE_DIRECTORY_ENTRY_EXPORT: + if idx == c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT: self.export_mgr = exports.ExportManager(pe=self, section=section) self.exports = self.export_mgr.exports # Parse the resources directory entry of the PE file - if idx == pestruct.IMAGE_DIRECTORY_ENTRY_RESOURCE: + if idx == c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE: self.rsrc_mgr = resources.ResourceManager(pe=self, section=section) self.resources = self.rsrc_mgr.resources self.raw_resources = self.rsrc_mgr.raw_resources # Parse the relocation directory entry of the PE file - if idx == pestruct.IMAGE_DIRECTORY_ENTRY_BASERELOC: + if idx == c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC: self.reloc_mgr = relocations.RelocationManager(pe=self, section=section) self.relocations = self.reloc_mgr.relocations # Parse the TLS directory entry of the PE file - if idx == pestruct.IMAGE_DIRECTORY_ENTRY_TLS: + if idx == c_pe.IMAGE_DIRECTORY_ENTRY_TLS: self.tls_mgr = tls.TLSManager(pe=self, section=section) self.tls_callbacks = self.tls_mgr.callbacks @@ -279,7 +297,9 @@ def get_resource_type(self, rsrc_id: str | EnumInstance): if rsrc_id not in self.resources: raise ResourceException(f"Resource with ID {rsrc_id} not found in PE!") - for resource in self.rsrc_mgr.parse_resources(resources=self.resources[rsrc_id]): + for resource in self.rsrc_mgr.parse_resources( + resources=self.resources[rsrc_id] + ): yield resource def virtual_address(self, address: int) -> int: @@ -404,7 +424,10 @@ def write(self, data: bytes): # Update the section data for section in self.patched_sections.values(): - if section.virtual_address <= offset and section.virtual_address + section.virtual_size >= offset: + if ( + section.virtual_address <= offset + and section.virtual_address + section.virtual_size >= offset + ): self.seek(address=section.virtual_address) section.data = self.read(size=section.virtual_size) @@ -419,7 +442,9 @@ def read_image_directory(self, index: int) -> bytes: """ directory_entry = self.optional_header.DataDirectory[index] - return self.virtual_read(address=directory_entry.VirtualAddress, size=directory_entry.Size) + return self.virtual_read( + address=directory_entry.VirtualAddress, size=directory_entry.Size + ) def directory_va(self, index: int) -> int: """Returns the virtual address of a directory given its index. @@ -452,15 +477,19 @@ def debug(self) -> cstruct: A `cstruct` object of the debug entry within the PE file. """ - debug_directory_entry = self.read_image_directory(index=pestruct.IMAGE_DIRECTORY_ENTRY_DEBUG) - image_directory_size = len(pestruct.IMAGE_DEBUG_DIRECTORY) + debug_directory_entry = self.read_image_directory( + index=c_pe.IMAGE_DIRECTORY_ENTRY_DEBUG + ) + image_directory_size = len(c_pe.IMAGE_DEBUG_DIRECTORY) for _ in range(0, len(debug_directory_entry) // image_directory_size): - entry = pestruct.IMAGE_DEBUG_DIRECTORY(debug_directory_entry) - dbg_entry = self.virtual_read(address=entry.AddressOfRawData, size=entry.SizeOfData) + entry = c_pe.IMAGE_DEBUG_DIRECTORY(debug_directory_entry) + dbg_entry = self.virtual_read( + address=entry.AddressOfRawData, size=entry.SizeOfData + ) if entry.Type == 0x2: - return cv_info_struct.CV_INFO_PDB70(dbg_entry) + return c_cv_info.CV_INFO_PDB70(dbg_entry) def get_section(self, segment_index: int) -> Tuple[str, sections.PESection]: """Retrieve the section of the PE by index. @@ -526,14 +555,17 @@ def add_section( # Use the provided RVA or calculate the new section virtual address virtual_address = ( utils.align_int( - integer=last_section.virtual_address + last_section.virtual_size, blocksize=self.section_alignment + integer=last_section.virtual_address + last_section.virtual_size, + blocksize=self.section_alignment, ) if not va else va ) # Calculate the new section raw address - pointer_to_raw_data = last_section.pointer_to_raw_data + last_section.size_of_raw_data + pointer_to_raw_data = ( + last_section.pointer_to_raw_data + last_section.size_of_raw_data + ) # Build the new section new_section = sections.build_section( @@ -545,7 +577,7 @@ def add_section( ) # Update the last section offset - offset = last_section.offset + pestruct.IMAGE_SECTION_HEADER.size + offset = last_section.offset + c_pe.IMAGE_SECTION_HEADER.size self.last_section_offset = offset # Increment the NumberOfSections field @@ -561,15 +593,21 @@ def add_section( ) # Add the new section to the PE - self.sections[name] = sections.PESection(pe=self, section=new_section, offset=offset, data=data) - self.patched_sections[name] = sections.PESection(pe=self, section=new_section, offset=offset, data=data) + self.sections[name] = sections.PESection( + pe=self, section=new_section, offset=offset, data=data + ) + self.patched_sections[name] = sections.PESection( + pe=self, section=new_section, offset=offset, data=data + ) # Update the SizeOfImage field last_section = self.patched_sections[next(reversed(self.patched_sections))] last_va = last_section.virtual_address last_size = last_section.virtual_size - pe_size = utils.align_int(integer=(last_va + last_size), blocksize=self.section_alignment) + pe_size = utils.align_int( + integer=(last_va + last_size), blocksize=self.section_alignment + ) self.optional_header.SizeOfImage = pe_size # Write the data to the PE file diff --git a/tests/test_pe.py b/tests/test_pe.py index 095bec4..141fff2 100644 --- a/tests/test_pe.py +++ b/tests/test_pe.py @@ -6,27 +6,35 @@ from dissect.executable.pe.pe import PE -def test_pe_valid_signature(): +def test_pe_valid_signature() -> None: with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) assert pe._valid() is True -def test_pe_invalid_signature(): +def test_pe_invalid_signature() -> None: with pytest.raises(InvalidPE): PE(BytesIO(b"MZ" + b"\x00" * 400)) -def test_pe_sections(): - known_sections = [".dissect", ".text", ".rdata", ".idata", ".rsrc", ".reloc", ".tls"] +def test_pe_sections() -> None: + known_sections = [ + ".dissect", + ".text", + ".rdata", + ".idata", + ".rsrc", + ".reloc", + ".tls", + ] with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) assert known_sections == [section for section in pe.sections.keys()] -def test_pe_imports(): +def test_pe_imports() -> None: known_imports = [ "SHELL32.dll", "ole32.dll", @@ -44,9 +52,15 @@ def test_pe_imports(): assert known_imports == [import_ for import_ in pe.imports.keys()] -def test_pe_exports(): +def test_pe_exports() -> None: # Too much export functions to put in a list - known_exports = ["1", "2", "CreateOverlayApiInterface", "CreateShadowPlayApiInterface", "ShadowPlayOnSystemStart"] + known_exports = [ + "1", + "2", + "CreateOverlayApiInterface", + "CreateShadowPlayApiInterface", + "ShadowPlayOnSystemStart", + ] with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) @@ -54,7 +68,7 @@ def test_pe_exports(): assert known_exports == [export_ for export_ in pe.exports.keys()] -def test_pe_resources(): +def test_pe_resources() -> None: known_resource_types = ["RcData", "Manifest"] with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) @@ -62,15 +76,26 @@ def test_pe_resources(): assert known_resource_types == [resource for resource in pe.resources.keys()] -def test_pe_relocations(): +def test_pe_relocations() -> None: with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) assert len(pe.relocations) == 9 -def test_pe_tls_callbacks(): - known_callbacks = [430080, 434176, 438272, 442368, 446464, 450560, 454656, 458752, 462848, 466944] +def test_pe_tls_callbacks() -> None: + known_callbacks = [ + 430080, + 434176, + 438272, + 442368, + 446464, + 450560, + 454656, + 458752, + 462848, + 466944, + ] with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) diff --git a/tests/test_pe_builder.py b/tests/test_pe_builder.py index 5d8dd6d..9c71ac6 100644 --- a/tests/test_pe_builder.py +++ b/tests/test_pe_builder.py @@ -1,9 +1,9 @@ from dissect.executable import PE from dissect.executable.pe import Builder, Patcher -from dissect.executable.pe.helpers.c_pe import pestruct +from dissect.executable.pe.c_pe import c_pe -def test_build_new_pe_lfanew(): +def test_build_new_pe_lfanew() -> None: builder = Builder() builder.new() pe = builder.pe @@ -11,49 +11,75 @@ def test_build_new_pe_lfanew(): assert pe.mz_header.e_lfanew == 0x8C -def test_build_new_x86_pe_exe(): +def test_build_new_x86_pe_exe() -> None: builder = Builder(arch="x86") builder.new() pe = builder.pe pe.pe_file.seek(len(pe.mz_header)) stub = pe.pe_file.read(pe.mz_header.e_lfanew - len(pe.mz_header)) - assert stub[14 : 73 - 4] == b"This program is made with dissect.pe <3 kusjesvanSRT <3" + assert ( + stub[14 : 73 - 4] == b"This program is made with dissect.pe <3 kusjesvanSRT <3" + ) - assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE == 0x0100 + assert ( + pe.file_header.Characteristics + & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE + == 0x0100 + ) -def test_build_new_x64_pe_exe(): +def test_build_new_x64_pe_exe() -> None: builder = Builder(arch="x64") builder.new() pe = builder.pe pe.pe_file.seek(len(pe.mz_header)) stub = pe.pe_file.read(pe.mz_header.e_lfanew - len(pe.mz_header)) - assert stub[14 : 73 - 4] == b"This program is made with dissect.pe <3 kusjesvanSRT <3" + assert ( + stub[14 : 73 - 4] == b"This program is made with dissect.pe <3 kusjesvanSRT <3" + ) - assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE != 0x0100 + assert ( + pe.file_header.Characteristics + & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE + != 0x0100 + ) -def test_build_new_x86_pe_dll(): +def test_build_new_x86_pe_dll() -> None: builder = Builder(arch="x86", dll=True) builder.new() pe = builder.pe - assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE == 0x0100 - assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_DLL == 0x2000 + assert ( + pe.file_header.Characteristics + & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE + == 0x0100 + ) + assert ( + pe.file_header.Characteristics & c_pe.ImageCharacteristics.IMAGE_FILE_DLL + == 0x2000 + ) -def test_build_new_x64_pe_dll(): +def test_build_new_x64_pe_dll() -> None: builder = Builder(arch="x64", dll=True) builder.new() pe = builder.pe - assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE != 0x0100 - assert pe.file_header.Characteristics & pestruct.ImageCharacteristics.IMAGE_FILE_DLL == 0x2000 + assert ( + pe.file_header.Characteristics + & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE + != 0x0100 + ) + assert ( + pe.file_header.Characteristics & c_pe.ImageCharacteristics.IMAGE_FILE_DLL + == 0x2000 + ) -def test_build_new_pe_with_custom_section(): +def test_build_new_pe_with_custom_section() -> None: builder = Builder() builder.new() pe = builder.pe diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py index e1fcc06..8cc7fc1 100644 --- a/tests/test_pe_modifications.py +++ b/tests/test_pe_modifications.py @@ -3,7 +3,7 @@ from dissect.executable.pe import Patcher -def test_add_imports(): +def test_add_imports() -> None: dllname = "kusjesvanSRT.dll" functions = ["PressButtons", "LooseLips"] @@ -16,12 +16,14 @@ def test_add_imports(): assert "kusjesvanSRT.dll" in new_pe.imports - custom_dll_imports = [i.name for i in new_pe.imports["kusjesvanSRT.dll"].functions] + custom_dll_imports = [ + i.name for i in new_pe.imports["kusjesvanSRT.dll"].functions + ] assert "PressButtons" in custom_dll_imports assert "LooseLips" in custom_dll_imports -def test_resize_section_smaller(): +def test_resize_section_smaller() -> None: with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) @@ -30,28 +32,34 @@ def test_resize_section_smaller(): patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build) - assert new_pe.sections[".text"].size == len(b"kusjesvanSRT, patched with dissect") + assert new_pe.sections[".text"].size == len( + b"kusjesvanSRT, patched with dissect" + ) assert ( new_pe.sections[".text"].data[: len(b"kusjesvanSRT, patched with dissect")] == b"kusjesvanSRT, patched with dissect" ) -def test_resize_section_bigger(): +def test_resize_section_bigger() -> None: with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) original_size = pe.sections[".rdata"].size - pe.patched_sections[".rdata"].data += b"kusjesvanSRT, patched with dissect" * 100 + pe.patched_sections[".rdata"].data += ( + b"kusjesvanSRT, patched with dissect" * 100 + ) patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build) - assert new_pe.sections[".rdata"].size == original_size + len(b"kusjesvanSRT, patched with dissect" * 100) + assert new_pe.sections[".rdata"].size == original_size + len( + b"kusjesvanSRT, patched with dissect" * 100 + ) -def test_resize_resource_smaller(): +def test_resize_resource_smaller() -> None: with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) @@ -61,12 +69,12 @@ def test_resize_resource_smaller(): patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build) - assert [patched.data for patched in new_pe.get_resource_type(rsrc_id="Manifest")] == [ - b"kusjesvanSRT, patched with dissect" - ] + assert [ + patched.data for patched in new_pe.get_resource_type(rsrc_id="Manifest") + ] == [b"kusjesvanSRT, patched with dissect"] -def test_resize_resource_bigger(): +def test_resize_resource_bigger() -> None: with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) @@ -82,7 +90,7 @@ def test_resize_resource_bigger(): ] == [b"kusjesvanSRT, patched with dissect"] -def test_add_section(): +def test_add_section() -> None: with open("tests/data/testexe.exe", "rb") as pe_fh: pe = PE(pe_file=pe_fh) pe.add_section(name=".SRT", data=b"kusjesvanSRT") diff --git a/tests/test_section.py b/tests/test_section.py index cd0becf..1282267 100644 --- a/tests/test_section.py +++ b/tests/test_section.py @@ -25,7 +25,11 @@ def section_table(entries: int) -> SectionTable: def mock_section_table(section_data: bytes) -> Mock: - shdr = c_elf_64.Shdr(sh_offset=len(c_elf_64.Shdr), sh_size=len(section_data), sh_entsize=len(section_data)) + shdr = c_elf_64.Shdr( + sh_offset=len(c_elf_64.Shdr), + sh_size=len(section_data), + sh_entsize=len(section_data), + ) mocked_table = Mock() mocked_table.fh = BytesIO(shdr.dumps() + section_data) mocked_table.offset = 0 diff --git a/tests/test_segment.py b/tests/test_segment.py index 1d2d945..3dda28a 100644 --- a/tests/test_segment.py +++ b/tests/test_segment.py @@ -6,7 +6,9 @@ def create_segment(segment_data: bytes) -> Segment: - c_segment = c_elf_64.Phdr(p_offset=len(c_elf_64.Phdr), p_filesz=len(segment_data)).dumps() + c_segment = c_elf_64.Phdr( + p_offset=len(c_elf_64.Phdr), p_filesz=len(segment_data) + ).dumps() fh = BytesIO(c_segment + segment_data) return Segment(fh, 0) diff --git a/tests/test_segment_table.py b/tests/test_segment_table.py index e666c80..8e8ac0b 100644 --- a/tests/test_segment_table.py +++ b/tests/test_segment_table.py @@ -34,7 +34,9 @@ def create_segment_table(amount: int, random_data: bytes) -> SegmentTable: data_size = len(random_data) segments_data = [] for idx in range(amount): - data = c_elf_64.Phdr(p_offset=len(c_elf_64.Phdr) * amount + idx * data_size, p_filesz=data_size).dumps() + data = c_elf_64.Phdr( + p_offset=len(c_elf_64.Phdr) * amount + idx * data_size, p_filesz=data_size + ).dumps() segments_data.append(data) segments_data.append(random_data * amount) From a58b9be99d087dae7538da0d22523d7cf6f35365 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Mon, 17 Feb 2025 10:16:00 +0000 Subject: [PATCH 04/43] Fix ruff formatting from main --- dissect/executable/pe/__init__.py | 4 +- dissect/executable/pe/helpers/builder.py | 50 +++--- dissect/executable/pe/helpers/exports.py | 4 +- dissect/executable/pe/helpers/imports.py | 37 ++-- dissect/executable/pe/helpers/patcher.py | 110 ++++-------- dissect/executable/pe/helpers/relocations.py | 14 +- dissect/executable/pe/helpers/resources.py | 82 ++++----- dissect/executable/pe/helpers/sections.py | 34 ++-- dissect/executable/pe/helpers/tls.py | 16 +- dissect/executable/pe/pe.py | 167 ++++++++----------- tests/test_pe.py | 24 +-- tests/test_pe_modifications.py | 36 ++-- 12 files changed, 220 insertions(+), 358 deletions(-) diff --git a/dissect/executable/pe/__init__.py b/dissect/executable/pe/__init__.py index f7e9849..7667b2f 100644 --- a/dissect/executable/pe/__init__.py +++ b/dissect/executable/pe/__init__.py @@ -11,15 +11,15 @@ from dissect.executable.pe.pe import PE __all__ = [ + "PE", "Builder", "ExportFunction", "ExportManager", "ImportFunction", "ImportManager", "ImportModule", - "Patcher", - "PE", "PESection", + "Patcher", "Resource", "ResourceManager", ] diff --git a/dissect/executable/pe/helpers/builder.py b/dissect/executable/pe/helpers/builder.py index 261e831..d5a9a79 100644 --- a/dissect/executable/pe/helpers/builder.py +++ b/dissect/executable/pe/helpers/builder.py @@ -1,8 +1,7 @@ from __future__ import annotations -from datetime import datetime +from datetime import datetime, timezone from io import BytesIO -from typing import TYPE_CHECKING from dissect.executable.exception import BuildSectionException from dissect.executable.pe.c_pe import c_pe @@ -29,9 +28,7 @@ def __init__( subsystem: int = 0x2, ): self.arch = ( - c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64 - if arch == "x64" - else c_pe.MachineType.IMAGE_FILE_MACHINE_I386 + c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64 if arch == "x64" else c_pe.MachineType.IMAGE_FILE_MACHINE_I386 ) self.dll = dll self.subsystem = subsystem @@ -52,9 +49,7 @@ def new(self) -> None: image_characteristics = self.get_characteristics() # Generate the file header - self.file_header = self.gen_file_header( - machine=self.arch, characteristics=image_characteristics - ) + self.file_header = self.gen_file_header(machine=self.arch, characteristics=image_characteristics) # Generate the optional header self.optional_header = self.gen_optional_header() @@ -117,10 +112,10 @@ def gen_mz_header( e_cs: int = 0, e_lfarlc: int = 64, e_ovno: int = 0, - e_res: list = [0, 0, 0, 0], + e_res: list[int] | None = None, e_oemid: int = 0, e_oeminfo: int = 0, - e_res2: int = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + e_res2: list[int] | None = None, e_lfanew: int = 0, ) -> c_pe.IMAGE_DOS_HEADER: """Generate the MZ header for the new PE file. @@ -166,15 +161,15 @@ def gen_mz_header( mz_header.e_cs = e_cs mz_header.e_lfarlc = e_lfarlc mz_header.e_ovno = e_ovno - mz_header.e_res = e_res + mz_header.e_res = e_res or [0, 0, 0, 0] mz_header.e_oemid = e_oemid mz_header.e_oeminfo = e_oeminfo - mz_header.e_res2 = e_res2 + mz_header.e_res2 = e_res2 or [0, 0, 0, 0, 0, 0, 0, 0, 0, 0] # Calculate the start of the NT headers by checking the location and size of the relocation table # within the MZ header start_of_nt_header = (mz_header.e_lfarlc + (mz_header.e_crlc * 4)) + len(STUB) - mz_header.e_lfanew = start_of_nt_header if not e_lfanew else e_lfanew + mz_header.e_lfanew = e_lfanew if e_lfanew else start_of_nt_header # Align the e_lfanew value mz_header.e_lfanew = mz_header.e_lfanew + (mz_header.e_lfanew % 2) @@ -236,7 +231,7 @@ def gen_file_header( # Set the timestamp to now if not given if not time_date_stamp: - time_date_stamp = int(datetime.utcnow().timestamp()) + time_date_stamp = int(datetime.now(tz=timezone.utc).timestamp()) file_header = c_pe.IMAGE_FILE_HEADER() file_header.Machine = machine @@ -280,10 +275,7 @@ def gen_optional_header( size_of_heap_commit: int = 0x1000, loaderflags: int = 0, number_of_rva_and_sizes: int = c_pe.IMAGE_NUMBEROF_DIRECTORY_ENTRIES, - datadirectory: list = [ - c_pe.IMAGE_DATA_DIRECTORY(BytesIO(b"\x00" * len(c_pe.IMAGE_DATA_DIRECTORY))) - for _ in range(c_pe.IMAGE_NUMBEROF_DIRECTORY_ENTRIES) - ], + datadirectory: list[c_pe.IMAGE_DATA_DIRECTORY] | None = None, ) -> c_pe.IMAGE_OPTIONAL_HEADER | c_pe.IMAGE_OPTIONAL_HEADER64: """Generate the optional header for the new PE file. @@ -325,11 +317,12 @@ def gen_optional_header( if self.machine == 0x8664: optional_header = c_pe.IMAGE_OPTIONAL_HEADER64() - optional_header.Magic = 0x20B if not magic else magic + _magic = 0x20B else: optional_header = c_pe.IMAGE_OPTIONAL_HEADER() - optional_header.Magic = 0x10B if not magic else magic + _magic = 0x10B + optional_header.Magic = magic or _magic self.file_alignment = file_alignment self.section_alignment = section_alignment @@ -373,7 +366,10 @@ def gen_optional_header( optional_header.SizeOfHeapCommit = size_of_heap_commit optional_header.LoaderFlags = loaderflags optional_header.NumberOfRvaAndSizes = number_of_rva_and_sizes - optional_header.DataDirectory = datadirectory + optional_header.DataDirectory = datadirectory or [ + c_pe.IMAGE_DATA_DIRECTORY(BytesIO(b"\x00" * len(c_pe.IMAGE_DATA_DIRECTORY))) + for _ in range(c_pe.IMAGE_NUMBEROF_DIRECTORY_ENTRIES) + ] return optional_header @@ -415,18 +411,14 @@ def section( """ if len(name) > 8: - raise BuildSectionException( - "section names can't be longer than 8 characters" - ) + raise BuildSectionException("section names can't be longer than 8 characters") if isinstance(name, str): name = name.encode() section_header = c_pe.IMAGE_SECTION_HEADER() - pointer_to_raw_data = utils.align_int( - integer=pointer_to_raw_data, blocksize=self.file_alignment - ) + pointer_to_raw_data = utils.align_int(integer=pointer_to_raw_data, blocksize=self.file_alignment) section_header.Name = name + utils.pad(size=8 - len(name)) section_header.VirtualSize = virtual_size @@ -452,9 +444,7 @@ def pe_size(self) -> int: The size of the PE. """ - last_section = self.pe.patched_sections[ - next(reversed(self.pe.patched_sections)) - ] + last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] va = last_section.virtual_address size = last_section.virtual_size diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py index 7017c9a..bad29e9 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/helpers/exports.py @@ -90,8 +90,8 @@ def parse_exports(self) -> None: ordinal=idx + 1, address=address, name=export_name ) - def add(self): + def add(self) -> None: raise NotImplementedError - def delete(self): + def delete(self) -> None: raise NotImplementedError diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index 84e3213..1127fe7 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -3,12 +3,14 @@ import struct from collections import OrderedDict from io import BytesIO -from typing import TYPE_CHECKING, BinaryIO, Iterator +from typing import TYPE_CHECKING, BinaryIO from dissect.executable.pe.c_pe import c_pe from dissect.executable.pe.helpers import utils if TYPE_CHECKING: + from collections.abc import Iterator + from dissect.cstruct.cstruct import cstruct from dissect.executable.pe.helpers.sections import PESection @@ -123,15 +125,11 @@ def parse_imports(self) -> None: The imports are in turn added to the `imports` attribute so they can be accessed by the user. """ - import_data = BytesIO( - self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT) - ) + import_data = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT)) import_data.seek(0) # Loop over the entries - for descriptor_va, import_descriptor in self.import_descriptors( - import_data=import_data - ): + for descriptor_va, import_descriptor in self.import_descriptors(import_data=import_data): if import_descriptor.Name not in [0xFFFFF800, 0x0]: self.pe.seek(import_descriptor.Name) modulename = c_pe.char[None](self.pe) @@ -151,9 +149,7 @@ def parse_imports(self) -> None: ) for thunkdata in self.parse_thunks(offset=first_thunk): - module.functions.append( - ImportFunction(pe=self.pe, thunkdata=thunkdata) - ) + module.functions.append(ImportFunction(pe=self.pe, thunkdata=thunkdata)) self.imports[modulename.decode()] = module @@ -177,9 +173,7 @@ def import_descriptors( yield import_data.tell(), import_descriptor - def parse_thunks( - self, offset: int - ) -> Iterator[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64, None, None]: + def parse_thunks(self, offset: int) -> Iterator[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64, None, None]: """Parse the import thunks for every module. Args: @@ -206,9 +200,7 @@ def add(self, dllname: str, functions: list[str]) -> None: functions: A `list` of function names belonging to the module. """ - self.last_section = self.pe.patched_sections[ - next(reversed(self.pe.patched_sections)) - ] + self.last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] # Build a dummy import module self.imports[dllname] = ImportModule( @@ -220,14 +212,12 @@ def add(self, dllname: str, functions: list[str]) -> None: ) # Build the dummy module functions for function in functions: - self.pe.imports[dllname].functions.append( - ImportFunction(pe=self.pe, thunkdata=None, name=function) - ) + self.pe.imports[dllname].functions.append(ImportFunction(pe=self.pe, thunkdata=None, name=function)) # Rebuild the import table with the new import module and functions self.build_import_table() - def delete(self, dllname: str, functions: list): + def delete(self, dllname: str, functions: list) -> None: raise NotImplementedError def build_import_table(self) -> None: @@ -265,9 +255,7 @@ def build_import_table(self) -> None: datadirectory_size += len(descriptor) # Create a new section - section_data = utils.align_data( - data=self.import_data, blocksize=self.pe.file_alignment - ) + section_data = utils.align_data(data=self.import_data, blocksize=self.pe.file_alignment) size = len(self.import_data) + c_pe.IMAGE_SECTION_HEADER.size self.pe.add_section( name=".idata", @@ -321,8 +309,7 @@ def _build_thunkdata(self, import_rvas: list[int]) -> bytes: rva += self.pe.optional_header.SizeOfImage thunkdata += ( struct.pack(" int: The new size of the PE as an `int`. """ - last_section = self.pe.patched_sections[ - next(reversed(self.pe.patched_sections)) - ] + last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] va = last_section.virtual_address size = last_section.virtual_size - return utils.align_int( - integer=va + size, blocksize=self.pe.optional_header.SectionAlignment - ) + return utils.align_int(integer=va + size, blocksize=self.pe.optional_header.SectionAlignment) - def seek(self, address: int): + def seek(self, address: int) -> None: """Seek that is used to seek to a virtual address in the patched PE file. Args: @@ -83,14 +79,12 @@ def seek(self, address: int): raw_address = self.pe.virtual_address(address=address) self.patched_pe.seek(raw_address) - def _build_section_table(self): + def _build_section_table(self) -> None: """Function to build the section table and add the sections with their data.""" if self.patched_pe.tell() < self.pe.section_header_offset: # Pad the patched file with null bytes until we reach the section header offset - self.patched_pe.write( - utils.pad(size=self.pe.section_header_offset - self.patched_pe.tell()) - ) + self.patched_pe.write(utils.pad(size=self.pe.section_header_offset - self.patched_pe.tell())) # Write the section headers for section in self.pe.patched_sections.values(): @@ -101,7 +95,7 @@ def _build_section_table(self): self.patched_pe.seek(section.pointer_to_raw_data) self.patched_pe.write(section.data) - def _build_dos_header(self): + def _build_dos_header(self) -> None: """Function to build the DOS header, NT headers and the DOS stub.""" # Add the MZ @@ -118,7 +112,7 @@ def _build_dos_header(self): self.patched_pe.write(self.pe.file_header.dumps()) self.patched_pe.write(self.pe.optional_header.dumps()) - def _patch_rvas(self): + def _patch_rvas(self) -> None: """Function to call the different patch functions responsible for patching any kind of relative addressing.""" self._patch_import_rvas() @@ -126,7 +120,7 @@ def _patch_rvas(self): self._patch_rsrc_rvas() self._patch_tls_rvas() - def _patch_import_rvas(self): + def _patch_import_rvas(self) -> None: """Function to patch the RVAs of the import directory and the thunkdata entries.""" patched_import_data = bytearray() @@ -140,9 +134,7 @@ def _patch_import_rvas(self): # new RVA's section = self.pe.patched_section(va=directory_va) directory_offset = directory_va - section.virtual_address - original_directory_va = ( - self.pe.sections[section.name].virtual_address + directory_offset - ) + original_directory_va = self.pe.sections[section.name].virtual_address + directory_offset # Loop over the imports of the PE to patch the RVA's of the import descriptors and the associated thunkdata # entries @@ -178,9 +170,7 @@ def _patch_import_rvas(self): section.virtual_address + section.virtual_size, ): virtual_address = section.virtual_address - new_virtual_address = self.pe.patched_sections[ - name - ].virtual_address + new_virtual_address = self.pe.patched_sections[name].virtual_address break # Calculate the offset using the VA of the section and update the thunkdata @@ -202,7 +192,7 @@ def _patch_import_rvas(self): self.seek(directory_va) self.patched_pe.write(patched_import_data) - def _patch_export_rvas(self): + def _patch_export_rvas(self) -> None: """Function to patch the RVAs of the export directory and the associated function and name RVA's.""" directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) @@ -216,20 +206,12 @@ def _patch_export_rvas(self): # new RVA's section = self.pe.patched_section(va=directory_va) directory_offset = directory_va - section.virtual_address - original_directory_va = ( - self.pe.sections[section.name].virtual_address + directory_offset - ) + original_directory_va = self.pe.sections[section.name].virtual_address + directory_offset name_offset = export_directory.Name - original_directory_va - address_of_functions_offset = ( - export_directory.AddressOfFunctions - original_directory_va - ) - address_of_names_offset = ( - export_directory.AddressOfNames - original_directory_va - ) - address_of_name_ordinals = ( - export_directory.AddressOfNameOrdinals - original_directory_va - ) + address_of_functions_offset = export_directory.AddressOfFunctions - original_directory_va + address_of_names_offset = export_directory.AddressOfNames - original_directory_va + address_of_name_ordinals = export_directory.AddressOfNameOrdinals - original_directory_va export_directory.Name = directory_va + name_offset export_directory.AddressOfFunctions = directory_va + address_of_functions_offset @@ -244,17 +226,13 @@ def _patch_export_rvas(self): new_function_rvas = [] function_rvas = bytearray() self.seek(export_directory.AddressOfFunctions) - export_addresses = c_pe.uint32[export_directory.NumberOfFunctions].read( - self.patched_pe - ) + export_addresses = c_pe.uint32[export_directory.NumberOfFunctions].read(self.patched_pe) for address in export_addresses: section = self.pe.section(va=address) if not section: continue address_offset = address - section.virtual_address - new_address = ( - self.pe.patched_sections[section.name].virtual_address + address_offset - ) + new_address = self.pe.patched_sections[section.name].virtual_address + address_offset new_function_rvas.append(new_address) for rva in new_function_rvas: @@ -271,9 +249,7 @@ def _patch_export_rvas(self): for name_address in export_names: section = self.pe.section(va=name_address) address_offset = name_address - section.virtual_address - new_address = ( - self.pe.patched_sections[section.name].virtual_address + address_offset - ) + new_address = self.pe.patched_sections[section.name].virtual_address + address_offset new_name_rvas.append(new_address) for name_rva in new_name_rvas: @@ -283,7 +259,7 @@ def _patch_export_rvas(self): self.patched_pe.write(name_rvas) # self.pe.optional_header.DataDirectory[c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT].Size = len(name_rvas) - def _patch_rsrc_rvas(self): + def _patch_rsrc_rvas(self) -> None: """Function to patch the RVAs of the resource directory and the associated resource data RVA's.""" directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE) @@ -293,9 +269,7 @@ def _patch_rsrc_rvas(self): section_data = BytesIO() self.seek(directory_va) - for rsrc_entry in sorted( - self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"] - ): + for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"]): entry_offset = rsrc_entry["offset"] entry = rsrc_entry["entry"] @@ -318,7 +292,7 @@ def _patch_rsrc_rvas(self): self.seek(directory_va) self.patched_pe.write(section_data.read()) - def _patch_tls_rvas(self): + def _patch_tls_rvas(self) -> None: """Function to patch the RVAs of the TLS directory and the associated TLS callbacks.""" directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_TLS) @@ -332,45 +306,30 @@ def _patch_tls_rvas(self): # Patch the TLS StartAddressOfRawData and EndAddressOfRawData section = self.pe.section(va=tls_directory.StartAddressOfRawData - image_base) - start_address_offset = ( - tls_directory.StartAddressOfRawData - section.virtual_address - ) + start_address_offset = tls_directory.StartAddressOfRawData - section.virtual_address tls_directory.StartAddressOfRawData = ( - self.pe.patched_sections[section.name].virtual_address - + start_address_offset - ) - end_address_offset = ( - tls_directory.EndAddressOfRawData - tls_directory.StartAddressOfRawData - ) - tls_directory.EndAddressOfRawData = ( - tls_directory.StartAddressOfRawData + end_address_offset + self.pe.patched_sections[section.name].virtual_address + start_address_offset ) + end_address_offset = tls_directory.EndAddressOfRawData - tls_directory.StartAddressOfRawData + tls_directory.EndAddressOfRawData = tls_directory.StartAddressOfRawData + end_address_offset # Patch the TLS callbacks address section = self.pe.section(va=tls_directory.AddressOfCallBacks - image_base) - address_of_callbacks_offset = ( - tls_directory.AddressOfCallBacks - section.virtual_address - ) + address_of_callbacks_offset = tls_directory.AddressOfCallBacks - section.virtual_address tls_directory.AddressOfCallBacks = ( - self.pe.patched_sections[section.name].virtual_address - + address_of_callbacks_offset + self.pe.patched_sections[section.name].virtual_address + address_of_callbacks_offset ) # Patch the TLS AddressOfIndex section = self.pe.section(va=tls_directory.AddressOfIndex - image_base) - address_of_index_offset = ( - tls_directory.AddressOfIndex - - self.pe.sections[section.name].virtual_address - ) - tls_directory.AddressOfIndex = ( - self.pe.sections[section.name].virtual_address + address_of_index_offset - ) + address_of_index_offset = tls_directory.AddressOfIndex - self.pe.sections[section.name].virtual_address + tls_directory.AddressOfIndex = self.pe.sections[section.name].virtual_address + address_of_index_offset # Write the patched TLS directory to the new PE self.seek(directory_va) self.patched_pe.write(tls_directory.dumps()) - def _get_tls_attribute_section(self, va: int) -> PESection: + def _get_tls_attribute_section(self, va: int) -> PESection | None: """Function to get the section that contains the TLS attribute. Args: @@ -380,8 +339,7 @@ def _get_tls_attribute_section(self, va: int) -> PESection: The section that contains the TLS attribute as a `PESection` object. """ - for name, section in self.pe.sections.items(): - if va in range( - section.virtual_address, section.virtual_address + section.virtual_size - ): + for section in self.pe.sections.values(): + if va in range(section.virtual_address, section.virtual_address + section.virtual_size): return section + return None diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/helpers/relocations.py index ef4d16c..3aeb911 100644 --- a/dissect/executable/pe/helpers/relocations.py +++ b/dissect/executable/pe/helpers/relocations.py @@ -25,12 +25,10 @@ def __init__(self, pe: PE, section: PESection): self.parse_relocations() - def parse_relocations(self): + def parse_relocations(self) -> None: """Parse the relocation table of the PE file.""" - reloc_data = BytesIO( - self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC) - ) + reloc_data = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC)) reloc_data_size = reloc_data.getbuffer().nbytes while reloc_data.tell() < reloc_data_size: reloc_directory = c_pe.IMAGE_BASE_RELOCATION(reloc_data) @@ -39,11 +37,9 @@ def parse_relocations(self): break # Each entry consists of 2 bytes - number_of_entries = ( - reloc_directory.SizeOfBlock - len(reloc_directory.dumps()) - ) // 2 + number_of_entries = (reloc_directory.SizeOfBlock - len(reloc_directory.dumps())) // 2 entries = [] - for _ in range(0, number_of_entries): + for _ in range(number_of_entries): entry = c_pe.uint16(reloc_data) if entry: entries.append(entry) @@ -56,5 +52,5 @@ def parse_relocations(self): } ) - def add(self): + def add(self) -> None: raise NotImplementedError diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index 8b23afb..e59399b 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -2,12 +2,15 @@ from collections import OrderedDict from io import BytesIO -from typing import TYPE_CHECKING, Iterator +from typing import TYPE_CHECKING from dissect.executable.exception import ResourceException from dissect.executable.pe.c_pe import c_pe if TYPE_CHECKING: + from collections.abc import Iterator + from typing import BinaryIO + from dissect.cstruct.cstruct import cstruct from dissect.cstruct.types.enum import EnumInstance @@ -26,22 +29,18 @@ class ResourceManager: def __init__(self, pe: PE, section: PESection): self.pe = pe self.section = section - self.resources = OrderedDict() + self.resources: OrderedDict[str, Resource] = OrderedDict() self.raw_resources = [] self.parse_rsrc() - def parse_rsrc(self): + def parse_rsrc(self) -> None: """Parse the resource directory entry of the PE file.""" - rsrc_data = BytesIO( - self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE) - ) - self.resources = self._read_resource( - rc_type="_root", data=rsrc_data, offset=0, level=1 - ) + rsrc_data = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE)) + self.resources = self._read_resource(rc_type="_root", data=rsrc_data, offset=0, level=1) - def _read_entries(self, data: bytes, directory: cstruct) -> list: + def _read_entries(self, data: BinaryIO, directory: cstruct) -> list[c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY]: """Read the entries within the resource directory. Args: @@ -53,16 +52,14 @@ def _read_entries(self, data: bytes, directory: cstruct) -> list: """ entries = [] - for _ in range(0, directory.NumberOfNamedEntries + directory.NumberOfIdEntries): + for _ in range(directory.NumberOfNamedEntries + directory.NumberOfIdEntries): entry_offset = data.tell() entry = c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY(data) - self.raw_resources.append( - {"offset": entry_offset, "entry": entry, "data_offset": entry_offset} - ) + self.raw_resources.append({"offset": entry_offset, "entry": entry, "data_offset": entry_offset}) entries.append(entry) return entries - def _handle_data_entry(self, data: bytes, entry: cstruct, rc_type: str) -> Resource: + def _handle_data_entry(self, data: BinaryIO, entry: cstruct, rc_type: str) -> Resource: """Handle the data entry of a resource. This is the actual data associated with the directory entry. Args: @@ -76,7 +73,7 @@ def _handle_data_entry(self, data: bytes, entry: cstruct, rc_type: str) -> Resou data.seek(entry.OffsetToDirectory) data_entry = c_pe.IMAGE_RESOURCE_DATA_ENTRY(data) self.pe.seek(data_entry.OffsetToData) - data = self.pe.read(data_entry.Size) + _data = self.pe.read(data_entry.Size) raw_offset = data_entry.OffsetToData - self.section.virtual_address rsrc = Resource( pe=self.pe, @@ -90,7 +87,7 @@ def _handle_data_entry(self, data: bytes, entry: cstruct, rc_type: str) -> Resou { "offset": entry.OffsetToDirectory, "entry": data_entry, - "data": data, + "data": _data, "data_offset": raw_offset, "resource": rsrc, } @@ -98,8 +95,8 @@ def _handle_data_entry(self, data: bytes, entry: cstruct, rc_type: str) -> Resou return rsrc def _read_resource( - self, data: bytes, offset: int, rc_type: str, level: int = 1 - ) -> dict: + self, data: BinaryIO, offset: int, rc_type: str, level: int = 1 + ) -> OrderedDict[str, Resource | dict[str, Resource]]: """Recursively read the resources within the PE file. Each resource is added to the dictionary that is available to the user, as well as a list of @@ -119,9 +116,7 @@ def _read_resource( data.seek(offset) directory = c_pe.IMAGE_RESOURCE_DIRECTORY(data) - self.raw_resources.append( - {"offset": offset, "entry": directory, "data_offset": offset} - ) + self.raw_resources.append({"offset": offset, "entry": directory, "data_offset": offset}) entries = self._read_entries(data, directory) @@ -144,9 +139,7 @@ def _read_resource( level=level + 1, ) else: - resource[rc_type] = self._handle_data_entry( - data=data, entry=entry, rc_type=rc_type - ) + resource[rc_type] = self._handle_data_entry(data=data, entry=entry, rc_type=rc_type) return resource @@ -165,7 +158,7 @@ def get_resource(self, name: str) -> Resource: except KeyError: raise ResourceException(f"Resource {name} not found!") - def get_resource_type(self, rsrc_id: str | EnumInstance): + def get_resource_type(self, rsrc_id: str | EnumInstance) -> Iterator[Resource]: """Yields a generator containing all of the nodes within the resources that contain the requested ID. The ID can be either given by name or its value. @@ -180,10 +173,9 @@ def get_resource_type(self, rsrc_id: str | EnumInstance): if rsrc_id not in self.resources: raise ResourceException(f"Resource with ID {rsrc_id} not found in PE!") - for resource in self.parse_resources(resources=self.resources[rsrc_id]): - yield resource + yield from self.parse_resources(resources=self.resources[rsrc_id]) - def parse_resources(self, resources: dict) -> Iterator[Resource]: + def parse_resources(self, resources: OrderedDict[str, Resource]) -> Iterator[Resource]: """Parse the resources within the PE file. Args: @@ -193,13 +185,13 @@ def parse_resources(self, resources: dict) -> Iterator[Resource]: All of the resources within the PE file. """ - for _, resource in resources.items(): + for resource in resources.values(): if type(resource) is not OrderedDict: yield resource else: yield from self.parse_resources(resources=resource) - def show_resource_tree(self, resources: dict, indent: int = 0): + def show_resource_tree(self, resources: dict, indent: int = 0) -> None: """Print the resources within the PE as a tree. Args: @@ -214,7 +206,7 @@ def show_resource_tree(self, resources: dict, indent: int = 0): print(f"{' ' * indent} + name: {name}") self.show_resource_tree(resources=resource, indent=indent + 1) - def show_resource_info(self, resources: dict): + def show_resource_info(self, resources: dict) -> None: """Print basic information about the resource as well as the header. Args: @@ -229,15 +221,15 @@ def show_resource_info(self, resources: dict): else: self.show_resource_info(resources=resource) - def add_resource(self, name: str, data: bytes): + def add_resource(self, name: str, data: bytes) -> None: # TODO raise NotImplementedError - def delete_resource(self, name: str): + def delete_resource(self, name: str) -> None: # TODO raise NotImplementedError - def update_section(self, update_offset: int): + def update_section(self, update_offset: int) -> None: """Function to dynamically update the section data and size when a resource has been modified. Args: @@ -247,9 +239,7 @@ def update_section(self, update_offset: int): new_size = 0 section_data = b"" - for idx, resource in enumerate( - self.parse_resources(resources=self.pe.resources) - ): + for idx, resource in enumerate(self.parse_resources(resources=self.pe.resources)): if idx == 0: # Use the offset of the first resource to account for the size of the directory header header_size = resource.offset - self.section.virtual_address @@ -311,7 +301,7 @@ def __init__( self.offset = data_entry.OffsetToData self._size = data_entry.Size self.codepage = data_entry.CodePage - self._data = self.read_data() if not data else data + self._data = data or self.read_data() def read_data(self) -> bytes: """Read the data within the resource. @@ -334,7 +324,7 @@ def size(self) -> int: return len(self.data) @size.setter - def size(self, value: int) -> int: + def size(self, value: int) -> None: """Setter to set the size of the resource to the specified value. Args: @@ -350,7 +340,7 @@ def offset(self) -> int: return self.entry.OffsetToData @offset.setter - def offset(self, value: int): + def offset(self, value: int) -> None: """Setter to set the offset of the resource to the specified value. Args: @@ -365,7 +355,7 @@ def data(self) -> bytes: return self._data @data.setter - def data(self, value: bytes): + def data(self, value: bytes) -> None: """Setter to set the new data of the resource, but also dynamically update the offset of the resources within the same directory. @@ -388,9 +378,7 @@ def data(self, value: bytes): prev_offset = 0 prev_size = 0 - for rsrc_entry in sorted( - self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"] - ): + for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"]): entry_offset = rsrc_entry["offset"] entry = rsrc_entry["entry"] @@ -426,9 +414,7 @@ def data(self, value: bytes): # Update the section data and size self.section.data = data - self.pe.optional_header.DataDirectory[ - c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE - ].Size = len(data) + self.pe.optional_header.DataDirectory[c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE].Size = len(data) def __str__(self) -> str: return str(self.name) diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 6385246..591c583 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -36,7 +36,7 @@ def __init__(self, pe: PE, section: cstruct, offset: int, data: bytes = b""): # Keep track of the directories that are within this section self.directories = OrderedDict() - self._data = self.read_data() if not data else data + self._data = data or self.read_data() def read_data(self) -> bytes: """Return the data within the section. @@ -56,7 +56,7 @@ def size(self) -> int: return self.virtual_size @size.setter - def size(self, value: int): + def size(self, value: int) -> None: """Setter to set the size of the data to the specified value. This function can be used to update the size of the data, but also dynamically update the offset of the data @@ -67,9 +67,7 @@ def size(self, value: int): """ self.virtual_size = value - self.size_of_raw_data = utils.align_int( - integer=value, blocksize=self.pe.file_alignment - ) + self.size_of_raw_data = utils.align_int(integer=value, blocksize=self.pe.file_alignment) @property def virtual_address(self) -> int: @@ -77,7 +75,7 @@ def virtual_address(self) -> int: return self._virtual_address @virtual_address.setter - def virtual_address(self, value: int): + def virtual_address(self, value: int) -> None: """Setter to set the virtual address of the section to the specified value. This function also updates any of the virtual addresses of the directories that are residing within the section @@ -101,7 +99,7 @@ def virtual_size(self) -> int: return self._virtual_size @virtual_size.setter - def virtual_size(self, value: int): + def virtual_size(self, value: int) -> None: """Setter to set the virtual size of the section to the specified value. Args: @@ -117,7 +115,7 @@ def pointer_to_raw_data(self) -> int: return self._pointer_to_raw_data @pointer_to_raw_data.setter - def pointer_to_raw_data(self, value: int): + def pointer_to_raw_data(self, value: int) -> None: """Setter to set the pointer to the raw data of the section to the specified value. Args: @@ -133,7 +131,7 @@ def size_of_raw_data(self) -> int: return self._size_of_raw_data @size_of_raw_data.setter - def size_of_raw_data(self, value: int): + def size_of_raw_data(self, value: int) -> None: """Setter to set the size of the raw data to the specified value. The SizeOfRawData field uses the section alignment to make sure the data within this section is aligned to the @@ -143,12 +141,8 @@ def size_of_raw_data(self, value: int): value: The size of the data. """ - self._size_of_raw_data = utils.align_int( - integer=value, blocksize=self.pe.file_alignment - ) - self.section.SizeOfRawData = utils.align_int( - integer=value, blocksize=self.pe.file_alignment - ) + self._size_of_raw_data = utils.align_int(integer=value, blocksize=self.pe.file_alignment) + self.section.SizeOfRawData = utils.align_int(integer=value, blocksize=self.pe.file_alignment) @property def data(self) -> bytes: @@ -156,7 +150,7 @@ def data(self) -> bytes: return self._data[: self.virtual_size] @data.setter - def data(self, value: bytes): + def data(self, value: bytes) -> None: """Setter to set the new data of the resource, but also dynamically update the offset of the resources within the same directory. @@ -191,12 +185,8 @@ def data(self, value: bytes): if section.virtual_address == prev_va: continue - pointer_to_raw_data = utils.align_int( - integer=prev_ptr + prev_size, blocksize=self.pe.file_alignment - ) - virtual_address = utils.align_int( - integer=prev_va + prev_vsize, blocksize=self.pe.section_alignment - ) + pointer_to_raw_data = utils.align_int(integer=prev_ptr + prev_size, blocksize=self.pe.file_alignment) + virtual_address = utils.align_int(integer=prev_va + prev_vsize, blocksize=self.pe.section_alignment) if section.virtual_address < virtual_address: """Set the virtual address and raw pointer of the section to the new values, but only do so if the diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/helpers/tls.py index a8fa402..d757e72 100644 --- a/dissect/executable/pe/helpers/tls.py +++ b/dissect/executable/pe/helpers/tls.py @@ -26,12 +26,10 @@ def __init__(self, pe: PE, section: PESection): self.parse_tls() - def parse_tls(self): + def parse_tls(self) -> None: """Parse the TLS directory entry of the PE file when present.""" - tls_data = BytesIO( - self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_TLS) - ) + tls_data = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_TLS)) self.tls = self.pe.image_tls_directory(tls_data) self.pe.seek(self.tls.AddressOfCallBacks - self.pe.optional_header.ImageBase) @@ -57,7 +55,7 @@ def size(self) -> int: return self.tls.EndAddressOfRawData - self.tls.StartAddressOfRawData @size.setter - def size(self, value): + def size(self, value: int) -> None: """Setter to set the size of the TLS data to the specified value. Args: @@ -89,7 +87,7 @@ def data(self) -> bytes: return self._data @data.setter - def data(self, value): + def data(self, value: bytes) -> None: """Dynamically update the TLS directory data if the user changes the data. Args: @@ -107,9 +105,7 @@ def data(self, value): section_data.write(self.tls.dumps()) # Write the new TLS data to the section - start_address_rva = ( - self.tls.StartAddressOfRawData - self.pe.optional_header.ImageBase - ) + start_address_rva = self.tls.StartAddressOfRawData - self.pe.optional_header.ImageBase start_address_section_offset = start_address_rva - self.section.virtual_address section_data.seek(start_address_section_offset) section_data.write(self._data) @@ -118,5 +114,5 @@ def data(self, value): section_data.seek(0) self.section.data = section_data.read() - def add(self): + def add(self) -> None: raise NotImplementedError diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index 02da94b..26b7afa 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -3,7 +3,8 @@ from collections import OrderedDict from datetime import datetime, timezone from io import BytesIO -from typing import TYPE_CHECKING, BinaryIO, Tuple +from pathlib import Path +from typing import TYPE_CHECKING, BinaryIO from dissect.executable.exception import ( InvalidAddress, @@ -25,9 +26,13 @@ ) if TYPE_CHECKING: + from collections.abc import Iterator + from dissect.cstruct.cstruct import cstruct from dissect.cstruct.types.enum import EnumInstance + from dissect.executable.pe.helpers.resources import Resource + class PE: """Base class for parsing PE files. @@ -54,14 +59,14 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.section_header_offset = 0 self.last_section_offset = 0 - self.sections = OrderedDict() - self.patched_sections = OrderedDict() + self.sections: OrderedDict[str, sections.PESection] = OrderedDict() + self.patched_sections: OrderedDict[str, sections.PESection] = OrderedDict() - self.imports = None - self.exports = None - self.resources = None + self.imports: OrderedDict[str, imports.ImportModule] = None + self.exports: OrderedDict[str, exports.ExportFunction] = None + self.resources: OrderedDict[str, resources.Resource] = None self.raw_resources = None - self.relocations = None + self.relocations: list[dict] = None self.tls_callbacks = None self.directories = OrderedDict() @@ -78,9 +83,7 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.base_address = self.optional_header.ImageBase - self.timestamp = datetime.fromtimestamp( - self.file_header.TimeDateStamp, tz=timezone.utc - ) + self.timestamp = datetime.fromtimestamp(self.file_header.TimeDateStamp, tz=timezone.utc) # Parse the section header self.parse_section_header() @@ -96,9 +99,9 @@ def _valid(self) -> bool: """ self.pe_file.seek(self.mz_header.e_lfanew) - return True if c_pe.uint32(self.pe_file) == 0x4550 else False + return c_pe.uint32(self.pe_file) == 0x4550 - def parse_headers(self): + def parse_headers(self) -> None: """Function to parse the basic PE headers: - DOS header - File header (part of NT header) @@ -130,7 +133,7 @@ def parse_headers(self): self.optional_header = self.nt_headers.OptionalHeader - def _set_pe_architecture(self): + def _set_pe_architecture(self) -> None: """Set the architecture specific settings. Some of the structs are architecture specific. Raises: @@ -148,11 +151,9 @@ def _set_pe_architecture(self): self._high_bit = 1 << 31 self.read_address = c_pe.uint32 else: - raise InvalidArchitecture( - f"Invalid architecture found: {self.file_header.Machine:02x}" - ) + raise InvalidArchitecture(f"Invalid architecture found: {self.file_header.Machine:02x}") - def parse_section_header(self): + def parse_section_header(self) -> None: """Parse the sections within the PE file.""" self.pe_file.seek(self.section_header_offset) @@ -163,16 +164,12 @@ def parse_section_header(self): section_header = c_pe.IMAGE_SECTION_HEADER(self.pe_file) section_name = section_header.Name.decode().strip("\x00") # Take note of the sections, keep track of any patches seperately - self.sections[section_name] = sections.PESection( - pe=self, section=section_header, offset=offset - ) - self.patched_sections[section_name] = sections.PESection( - pe=self, section=section_header, offset=offset - ) + self.sections[section_name] = sections.PESection(pe=self, section=section_header, offset=offset) + self.patched_sections[section_name] = sections.PESection(pe=self, section=section_header, offset=offset) self.last_section_offset = self.sections[next(reversed(self.sections))].offset - def section(self, va: int = 0, name: str = "") -> sections.PESection: + def section(self, va: int = 0, name: str = "") -> sections.PESection | None: """Function to retrieve a section based on the given virtual address or name. Args: @@ -190,10 +187,10 @@ def section(self, va: int = 0, name: str = "") -> sections.PESection: section.virtual_address + section.virtual_size, ): return section - else: - return self.sections[name] + return None + return self.sections[name] - def patched_section(self, va: int = 0, name: str = "") -> sections.PESection: + def patched_section(self, va: int = 0, name: str = "") -> sections.PESection | None: """Function to retrieve a patched section based on the given virtual address or name. Args: @@ -211,8 +208,8 @@ def patched_section(self, va: int = 0, name: str = "") -> sections.PESection: section.virtual_address + section.virtual_size, ): return section - else: - return self.patched_sections[name] + return None + return self.patched_sections[name] def datadirectory_section(self, index: int) -> sections.PESection: """Return the section that contains the given virtual address. @@ -225,16 +222,13 @@ def datadirectory_section(self, index: int) -> sections.PESection: """ va = self.directory_va(index=index) - for _, section in self.patched_sections.items(): - if ( - va >= section.virtual_address - and va < section.virtual_address + section.virtual_size - ): + for section in self.patched_sections.values(): + if va >= section.virtual_address and va < section.virtual_address + section.virtual_size: return section raise InvalidVA(f"VA not found in sections: {va:#04x}") - def parse_directories(self): + def parse_directories(self) -> None: """Parse the different data directories in the PE file and initialize their associated managers. For now the following data directories are implemented: @@ -251,10 +245,7 @@ def parse_directories(self): # Take note of the current directory VA so we can dynamically update it when resizing sections section = self.datadirectory_section(index=idx) - directory_va_offset = ( - self.optional_header.DataDirectory[idx].VirtualAddress - - section.virtual_address - ) + directory_va_offset = self.optional_header.DataDirectory[idx].VirtualAddress - section.virtual_address section.directories[idx] = directory_va_offset # Parse the Import Address Table (IAT) @@ -282,7 +273,7 @@ def parse_directories(self): self.tls_mgr = tls.TLSManager(pe=self, section=section) self.tls_callbacks = self.tls_mgr.callbacks - def get_resource_type(self, rsrc_id: str | EnumInstance): + def get_resource_type(self, rsrc_id: str | EnumInstance) -> Iterator[Resource]: """Yields a generator containing all of the nodes within the resources that contain the requested ID. The ID can be either given by name or its value. @@ -297,10 +288,7 @@ def get_resource_type(self, rsrc_id: str | EnumInstance): if rsrc_id not in self.resources: raise ResourceException(f"Resource with ID {rsrc_id} not found in PE!") - for resource in self.rsrc_mgr.parse_resources( - resources=self.resources[rsrc_id] - ): - yield resource + yield from self.rsrc_mgr.parse_resources(resources=self.resources[rsrc_id]) def virtual_address(self, address: int) -> int: """Return the virtual address given a (possible) physical address. @@ -315,14 +303,14 @@ def virtual_address(self, address: int) -> int: if self.virtual: return address - for _, section in self.patched_sections.items(): + for section in self.patched_sections.values(): max_address = section.virtual_address + section.virtual_size if address >= section.virtual_address and address < max_address: return section.pointer_to_raw_data + (address - section.virtual_address) raise InvalidVA(f"VA not found in sections: {address:#04x}") - def raw_address(self, offset) -> int: + def raw_address(self, offset: int) -> int: """Return the physical address given a virtual address. Args: @@ -332,7 +320,7 @@ def raw_address(self, offset) -> int: The physical address as an `int`. """ - for _, section in self.patched_sections.items(): + for section in self.patched_sections.values(): max_address = section.pointer_to_raw_data + section.size_of_raw_data if offset >= section.pointer_to_raw_data and offset < max_address: return section.virtual_address + (offset - section.pointer_to_raw_data) @@ -375,7 +363,7 @@ def raw_read(self, offset: int, size: int) -> bytes: self.pe_file.seek(old_offset) return data - def seek(self, address: int): + def seek(self, address: int) -> None: """Seek to the given virtual address within a PE file. Args: @@ -407,7 +395,7 @@ def read(self, size: int) -> bytes: return self.pe_file.read(size) - def write(self, data: bytes): + def write(self, data: bytes) -> None: """Write the data to the PE file. This write function will also make sure to update the section data. @@ -424,10 +412,7 @@ def write(self, data: bytes): # Update the section data for section in self.patched_sections.values(): - if ( - section.virtual_address <= offset - and section.virtual_address + section.virtual_size >= offset - ): + if section.virtual_address <= offset and section.virtual_address + section.virtual_size >= offset: self.seek(address=section.virtual_address) section.data = self.read(size=section.virtual_size) @@ -442,9 +427,7 @@ def read_image_directory(self, index: int) -> bytes: """ directory_entry = self.optional_header.DataDirectory[index] - return self.virtual_read( - address=directory_entry.VirtualAddress, size=directory_entry.Size - ) + return self.virtual_read(address=directory_entry.VirtualAddress, size=directory_entry.Size) def directory_va(self, index: int) -> int: """Returns the virtual address of a directory given its index. @@ -470,35 +453,32 @@ def has_directory(self, index: int) -> bool: return self.optional_header.DataDirectory[index].Size != 0 - def debug(self) -> cstruct: + def debug(self) -> cstruct | None: """Return the debug directory of the given PE file. Returns: A `cstruct` object of the debug entry within the PE file. """ - debug_directory_entry = self.read_image_directory( - index=c_pe.IMAGE_DIRECTORY_ENTRY_DEBUG - ) + debug_directory_entry = self.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_DEBUG) image_directory_size = len(c_pe.IMAGE_DEBUG_DIRECTORY) - for _ in range(0, len(debug_directory_entry) // image_directory_size): + for _ in range(len(debug_directory_entry) // image_directory_size): entry = c_pe.IMAGE_DEBUG_DIRECTORY(debug_directory_entry) - dbg_entry = self.virtual_read( - address=entry.AddressOfRawData, size=entry.SizeOfData - ) + dbg_entry = self.virtual_read(address=entry.AddressOfRawData, size=entry.SizeOfData) if entry.Type == 0x2: return c_cv_info.CV_INFO_PDB70(dbg_entry) + return None - def get_section(self, segment_index: int) -> Tuple[str, sections.PESection]: + def get_section(self, segment_index: int) -> tuple[str, sections.PESection]: """Retrieve the section of the PE by index. Args: segment_index: The segment to retrieve based on the order within the PE. Returns: - A `Tuple` contianing the section name and attributes as `PESection`. + A `tuple` contianing the section name and attributes as `PESection`. """ sections = list(self.sections.items()) @@ -529,12 +509,12 @@ def add_section( self, name: str, data: bytes, - va: int = None, - datadirectory: int = None, - datadirectory_rva: int = None, - datadirectory_size: int = None, - size: int = None, - ): + va: int | None = None, + datadirectory: int | None = None, + datadirectory_rva: int | None = None, + datadirectory_size: int | None = None, + size: int | None = None, + ) -> None: """Add a new section to the PE file. Args: @@ -550,22 +530,17 @@ def add_section( # Calculate the new section size raw_size = utils.align_int(integer=len(data), blocksize=self.file_alignment) - virtual_size = len(data) if not size else size + virtual_size = size or len(data) # Use the provided RVA or calculate the new section virtual address - virtual_address = ( - utils.align_int( - integer=last_section.virtual_address + last_section.virtual_size, - blocksize=self.section_alignment, - ) - if not va - else va + + virtual_address = va or utils.align_int( + integer=last_section.virtual_address + last_section.virtual_size, + blocksize=self.section_alignment, ) # Calculate the new section raw address - pointer_to_raw_data = ( - last_section.pointer_to_raw_data + last_section.size_of_raw_data - ) + pointer_to_raw_data = last_section.pointer_to_raw_data + last_section.size_of_raw_data # Build the new section new_section = sections.build_section( @@ -585,29 +560,19 @@ def add_section( # Set the VA and size of the datadirectory entry if this was marked as being such if datadirectory is not None: - self.optional_header.DataDirectory[datadirectory].VirtualAddress = ( - virtual_address if not datadirectory_rva else datadirectory_rva - ) - self.optional_header.DataDirectory[datadirectory].Size = ( - len(data) if not datadirectory_size else datadirectory_size - ) + self.optional_header.DataDirectory[datadirectory].VirtualAddress = datadirectory_rva or virtual_address + self.optional_header.DataDirectory[datadirectory].Size = datadirectory_size or len(data) # Add the new section to the PE - self.sections[name] = sections.PESection( - pe=self, section=new_section, offset=offset, data=data - ) - self.patched_sections[name] = sections.PESection( - pe=self, section=new_section, offset=offset, data=data - ) + self.sections[name] = sections.PESection(pe=self, section=new_section, offset=offset, data=data) + self.patched_sections[name] = sections.PESection(pe=self, section=new_section, offset=offset, data=data) # Update the SizeOfImage field last_section = self.patched_sections[next(reversed(self.patched_sections))] last_va = last_section.virtual_address last_size = last_section.virtual_size - pe_size = utils.align_int( - integer=(last_va + last_size), blocksize=self.section_alignment - ) + pe_size = utils.align_int(integer=(last_va + last_size), blocksize=self.section_alignment) self.optional_header.SizeOfImage = pe_size # Write the data to the PE file @@ -623,7 +588,7 @@ def add_section( # Reparse the directories self.parse_directories() - def write_pe(self, filename: str = "out.exe"): + def write_pe(self, filename: str = "out.exe") -> None: """Write the contents of the PE to a new file. This will use the patcher that is part of the project to make sure any kind of relative addressing is also @@ -635,6 +600,4 @@ def write_pe(self, filename: str = "out.exe"): pepatcher = patcher.Patcher(pe=self) new_pe = pepatcher.build - - with open(filename, "wb") as fhout: - fhout.write(new_pe.read()) + Path(filename).write_bytes(new_pe.read()) diff --git a/tests/test_pe.py b/tests/test_pe.py index 141fff2..1bac7bd 100644 --- a/tests/test_pe.py +++ b/tests/test_pe.py @@ -5,9 +5,11 @@ from dissect.executable.exception import InvalidPE from dissect.executable.pe.pe import PE +from .util import data_file + def test_pe_valid_signature() -> None: - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) assert pe._valid() is True @@ -28,10 +30,10 @@ def test_pe_sections() -> None: ".reloc", ".tls", ] - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - assert known_sections == [section for section in pe.sections.keys()] + assert known_sections == list(pe.sections) def test_pe_imports() -> None: @@ -46,10 +48,10 @@ def test_pe_imports() -> None: "KERNEL32.dll", "USER32.dll", ] - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - assert known_imports == [import_ for import_ in pe.imports.keys()] + assert known_imports == list(pe.imports) def test_pe_exports() -> None: @@ -62,22 +64,22 @@ def test_pe_exports() -> None: "ShadowPlayOnSystemStart", ] - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - assert known_exports == [export_ for export_ in pe.exports.keys()] + assert known_exports == list(pe.exports) def test_pe_resources() -> None: known_resource_types = ["RcData", "Manifest"] - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - assert known_resource_types == [resource for resource in pe.resources.keys()] + assert known_resource_types == list(pe.resources) def test_pe_relocations() -> None: - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) assert len(pe.relocations) == 9 @@ -97,7 +99,7 @@ def test_pe_tls_callbacks() -> None: 466944, ] - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) assert pe.tls_callbacks == known_callbacks diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py index 8cc7fc1..d1ae2a1 100644 --- a/tests/test_pe_modifications.py +++ b/tests/test_pe_modifications.py @@ -2,12 +2,14 @@ from dissect.executable import PE from dissect.executable.pe import Patcher +from .util import data_file + def test_add_imports() -> None: dllname = "kusjesvanSRT.dll" functions = ["PressButtons", "LooseLips"] - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) pe.import_mgr.add(dllname=dllname, functions=functions) @@ -16,15 +18,13 @@ def test_add_imports() -> None: assert "kusjesvanSRT.dll" in new_pe.imports - custom_dll_imports = [ - i.name for i in new_pe.imports["kusjesvanSRT.dll"].functions - ] + custom_dll_imports = [i.name for i in new_pe.imports["kusjesvanSRT.dll"].functions] assert "PressButtons" in custom_dll_imports assert "LooseLips" in custom_dll_imports def test_resize_section_smaller() -> None: - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) pe.sections[".text"].data = b"kusjesvanSRT, patched with dissect" @@ -32,9 +32,7 @@ def test_resize_section_smaller() -> None: patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build) - assert new_pe.sections[".text"].size == len( - b"kusjesvanSRT, patched with dissect" - ) + assert new_pe.sections[".text"].size == len(b"kusjesvanSRT, patched with dissect") assert ( new_pe.sections[".text"].data[: len(b"kusjesvanSRT, patched with dissect")] == b"kusjesvanSRT, patched with dissect" @@ -42,25 +40,21 @@ def test_resize_section_smaller() -> None: def test_resize_section_bigger() -> None: - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) original_size = pe.sections[".rdata"].size - pe.patched_sections[".rdata"].data += ( - b"kusjesvanSRT, patched with dissect" * 100 - ) + pe.patched_sections[".rdata"].data += b"kusjesvanSRT, patched with dissect" * 100 patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build) - assert new_pe.sections[".rdata"].size == original_size + len( - b"kusjesvanSRT, patched with dissect" * 100 - ) + assert new_pe.sections[".rdata"].size == original_size + len(b"kusjesvanSRT, patched with dissect" * 100) def test_resize_resource_smaller() -> None: - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) for e in pe.get_resource_type(rsrc_id="Manifest"): @@ -69,13 +63,13 @@ def test_resize_resource_smaller() -> None: patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build) - assert [ - patched.data for patched in new_pe.get_resource_type(rsrc_id="Manifest") - ] == [b"kusjesvanSRT, patched with dissect"] + assert [patched.data for patched in new_pe.get_resource_type(rsrc_id="Manifest")] == [ + b"kusjesvanSRT, patched with dissect" + ] def test_resize_resource_bigger() -> None: - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) for e in pe.get_resource_type(rsrc_id="Manifest"): @@ -91,7 +85,7 @@ def test_resize_resource_bigger() -> None: def test_add_section() -> None: - with open("tests/data/testexe.exe", "rb") as pe_fh: + with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) pe.add_section(name=".SRT", data=b"kusjesvanSRT") From 325b8ddc6dd2696e478f95badfb584d90a407cf0 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Mon, 17 Feb 2025 11:50:21 +0000 Subject: [PATCH 05/43] Fix tests for cstruct v4 compatibility --- dissect/executable/pe/helpers/builder.py | 4 ++-- dissect/executable/pe/helpers/patcher.py | 2 +- dissect/executable/pe/helpers/resources.py | 3 +-- dissect/executable/pe/helpers/sections.py | 2 +- 4 files changed, 5 insertions(+), 6 deletions(-) diff --git a/dissect/executable/pe/helpers/builder.py b/dissect/executable/pe/helpers/builder.py index d5a9a79..b665498 100644 --- a/dissect/executable/pe/helpers/builder.py +++ b/dissect/executable/pe/helpers/builder.py @@ -358,8 +358,8 @@ def gen_optional_header( optional_header.SizeOfImage = size_of_image optional_header.SizeOfHeaders = size_of_headers optional_header.CheckSum = checksum - optional_header.Subsystem = subsystem - optional_header.DllCharacteristics = dll_characteristics + optional_header.Subsystem = c_pe.WindowsSubsystem(subsystem) + optional_header.DllCharacteristics = c_pe.DLLCharacteristics(dll_characteristics) optional_header.SizeOfStackReserve = size_of_stack_reserve optional_header.SizeOfStackCommit = size_of_stack_commit optional_header.SizeOfHeapReserve = size_of_heap_reserve diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index 75b5687..30282d7 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -273,7 +273,7 @@ def _patch_rsrc_rvas(self) -> None: entry_offset = rsrc_entry["offset"] entry = rsrc_entry["entry"] - if entry._type.name == "IMAGE_RESOURCE_DATA_ENTRY": + if isinstance(entry, c_pe.IMAGE_RESOURCE_DATA_ENTRY): rsrc_obj = rsrc_entry["resource"] data_offset = rsrc_entry["data_offset"] diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index e59399b..4f3298e 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -381,8 +381,7 @@ def data(self, value: bytes) -> None: for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"]): entry_offset = rsrc_entry["offset"] entry = rsrc_entry["entry"] - - if entry._type.name == "IMAGE_RESOURCE_DATA_ENTRY": + if isinstance(entry, c_pe.IMAGE_RESOURCE_DATA_ENTRY): rsrc_obj = rsrc_entry["resource"] data_offset = rsrc_entry["data_offset"] diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 591c583..69bfa84 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -251,6 +251,6 @@ def build_section( section_header.PointerToLinenumbers = 0 section_header.NumberOfRelocations = 0 section_header.NumberOfLinenumbers = 0 - section_header.Characteristics = characteristics + section_header.Characteristics = c_pe.SectionFlags(characteristics) return section_header From 88d852f42e5934050e328e1ea924bfd8bef7e6dc Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 18 Feb 2025 11:36:02 +0000 Subject: [PATCH 06/43] Small cleanup of imports.py --- dissect/executable/pe/helpers/imports.py | 68 +++++++++++------------- 1 file changed, 31 insertions(+), 37 deletions(-) diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index 1127fe7..e52ffd3 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -31,13 +31,13 @@ class ImportModule: def __init__( self, name: bytes, - import_descriptor: c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT, + import_descriptor: int | None, module_va: int, name_va: int, first_thunk: int, ): self.name = name - self.import_descriptor = import_descriptor + self.import_descriptor = import_descriptor or c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT self.module_va = module_va self.name_va = name_va self.first_thunk = first_thunk @@ -61,7 +61,7 @@ class ImportFunction: def __init__( self, pe: PE, - thunkdata: c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64, + thunkdata: c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64 | None, name: str = "", ): self.pe = pe @@ -79,6 +79,10 @@ def name(self) -> str: if self._name: return self._name + if self.thunkdata is None: + # For the case thunkdata is not defined, such as during the `add` + return "" + ordinal = self.thunkdata.u1.AddressOfData & self.pe._high_bit if not ordinal: @@ -109,12 +113,13 @@ class ImportManager: def __init__(self, pe: PE, section: PESection): self.pe = pe + self.image_size: int = self.pe.optional_header.SizeOfImage self.section = section self.import_directory_rva = 0 self.import_data = bytearray() self.new_size_of_image = 0 self.section_data = bytearray() - self.imports = OrderedDict() + self.imports: OrderedDict[str, ImportModule] = OrderedDict() self.thunks = [] self.parse_imports() @@ -135,11 +140,8 @@ def parse_imports(self) -> None: modulename = c_pe.char[None](self.pe) # Use the OriginalFirstThunk if available, FirstThunk otherwise - first_thunk = ( - import_descriptor.FirstThunk - if not import_descriptor.OriginalFirstThunk - else import_descriptor.OriginalFirstThunk - ) + first_thunk = import_descriptor.OriginalFirstThunk or import_descriptor.FirstThunk + module = ImportModule( name=modulename, import_descriptor=import_descriptor, @@ -153,9 +155,7 @@ def parse_imports(self) -> None: self.imports[modulename.decode()] = module - def import_descriptors( - self, import_data: BinaryIO - ) -> Iterator[tuple[int, c_pe.IMAGE_IMPORT_DESCRIPTOR], None, None]: + def import_descriptors(self, import_data: BinaryIO) -> Iterator[tuple[int, c_pe.IMAGE_IMPORT_DESCRIPTOR]]: """Parse the import descriptors of the PE file. Args: @@ -173,7 +173,7 @@ def import_descriptors( yield import_data.tell(), import_descriptor - def parse_thunks(self, offset: int) -> Iterator[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64, None, None]: + def parse_thunks(self, offset: int) -> Iterator[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64]: """Parse the import thunks for every module. Args: @@ -200,7 +200,7 @@ def add(self, dllname: str, functions: list[str]) -> None: functions: A `list` of function names belonging to the module. """ - self.last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] + self.last_section = next(reversed(self.pe.patched_sections.values())) # Build a dummy import module self.imports[dllname] = ImportModule( @@ -211,8 +211,9 @@ def add(self, dllname: str, functions: list[str]) -> None: first_thunk=0, ) # Build the dummy module functions - for function in functions: - self.pe.imports[dllname].functions.append(ImportFunction(pe=self.pe, thunkdata=None, name=function)) + self.imports[dllname].functions.extend( + ImportFunction(pe=self.pe, thunkdata=None, name=function) for function in functions + ) # Rebuild the import table with the new import module and functions self.build_import_table() @@ -230,7 +231,7 @@ def build_import_table(self) -> None: # Reset the known thunkdata self.thunks = [] - import_descriptors = [] + import_descriptors: list[c_pe.IMAGE_IMPORT_DESCRIPTOR] = [] self.import_data = bytearray() for name, module in self.imports.items(): @@ -242,14 +243,14 @@ def build_import_table(self) -> None: first_thunk_rva = self._build_module_imports(functions=module.functions) import_descriptor = self._build_import_descriptor( first_thunk_rva=first_thunk_rva, - name_rva=self.pe.optional_header.SizeOfImage + name_offset, + name_rva=self.image_size + name_offset, ) import_descriptors.append(import_descriptor) datadirectory_size = 0 # Take note of the RVA of the first import descriptor - import_rva = self.pe.optional_header.SizeOfImage + len(self.import_data) + import_rva = self.image_size + len(self.import_data) for descriptor in import_descriptors: self.import_data += descriptor.dumps() datadirectory_size += len(descriptor) @@ -286,7 +287,7 @@ def _build_module_imports(self, functions: list[ImportFunction]) -> int: self.import_data += struct.pack(" bytes: The thunkdata as a `bytes` object. """ - thunkdata = bytearray() - for rva in import_rvas: - rva += self.pe.optional_header.SizeOfImage - thunkdata += ( - struct.pack(" cstruct: + def _build_import_descriptor(self, first_thunk_rva: int, name_rva: int) -> c_pe.IMAGE_IMPORT_DESCRIPTOR: """Function to build the import descriptor for the new import table. Args: From 7e299224035f84792086862441f00d78eff7e29a Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Mon, 24 Feb 2025 15:22:15 +0000 Subject: [PATCH 07/43] Small changes to the ResourceManager class --- dissect/executable/pe/helpers/resources.py | 90 ++++++++++++---------- 1 file changed, 49 insertions(+), 41 deletions(-) diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index 4f3298e..3255a14 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -2,6 +2,8 @@ from collections import OrderedDict from io import BytesIO +from itertools import chain +from textwrap import indent from typing import TYPE_CHECKING from dissect.executable.exception import ResourceException @@ -30,17 +32,19 @@ def __init__(self, pe: PE, section: PESection): self.pe = pe self.section = section self.resources: OrderedDict[str, Resource] = OrderedDict() - self.raw_resources = [] + self.raw_resources: list[dict] = [] - self.parse_rsrc() + self.parse() - def parse_rsrc(self) -> None: + def parse(self) -> None: """Parse the resource directory entry of the PE file.""" - rsrc_data = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE)) - self.resources = self._read_resource(rc_type="_root", data=rsrc_data, offset=0, level=1) + rsrc_data = BytesIO(self.section.data) + self.resources = self._read_resource(data=rsrc_data, offset=0, level=1) - def _read_entries(self, data: BinaryIO, directory: cstruct) -> list[c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY]: + def _read_entries( + self, data: BinaryIO, directory: c_pe.IMAGE_RESOURCE_DIRECTORY + ) -> list[c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY]: """Read the entries within the resource directory. Args: @@ -59,7 +63,7 @@ def _read_entries(self, data: BinaryIO, directory: cstruct) -> list[c_pe.IMAGE_R entries.append(entry) return entries - def _handle_data_entry(self, data: BinaryIO, entry: cstruct, rc_type: str) -> Resource: + def _handle_data_entry(self, data: BinaryIO, entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY, rc_type: str) -> Resource: """Handle the data entry of a resource. This is the actual data associated with the directory entry. Args: @@ -94,9 +98,7 @@ def _handle_data_entry(self, data: BinaryIO, entry: cstruct, rc_type: str) -> Re ) return rsrc - def _read_resource( - self, data: BinaryIO, offset: int, rc_type: str, level: int = 1 - ) -> OrderedDict[str, Resource | dict[str, Resource]]: + def _read_resource(self, data: BinaryIO, offset: int, level: int = 1) -> OrderedDict[str, Resource]: """Recursively read the resources within the PE file. Each resource is added to the dictionary that is available to the user, as well as a list of @@ -120,29 +122,33 @@ def _read_resource( entries = self._read_entries(data, directory) + _rc_type: str = "" + for entry in entries: - if level == 1: - rc_type = c_pe.ResourceID(entry.Id).name - else: - if entry.NameIsString: - data.seek(entry.NameOffset) - name_len = c_pe.uint16(data) - rc_type = c_pe.wchar[name_len](data) - else: - rc_type = str(entry.Id) + _rc_type = self._rc_type(entry, data, level) if entry.DataIsDirectory: - resource[rc_type] = self._read_resource( + resource[_rc_type] = self._read_resource( data=data, offset=entry.OffsetToDirectory, - rc_type=rc_type, level=level + 1, ) else: - resource[rc_type] = self._handle_data_entry(data=data, entry=entry, rc_type=rc_type) + resource[_rc_type] = self._handle_data_entry(data=data, entry=entry, rc_type=_rc_type) return resource + def _rc_type(self, entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY, data: BinaryIO, level: int = 1) -> str: + if level == 1: + return c_pe.ResourceID(entry.Id).name + + if entry.NameIsString: + data.seek(entry.NameOffset) + name_len = c_pe.uint16(data) + return c_pe.wchar[name_len](data) + + return str(entry.Id) + def get_resource(self, name: str) -> Resource: """Retrieve the resource by name. @@ -186,12 +192,12 @@ def parse_resources(self, resources: OrderedDict[str, Resource]) -> Iterator[Res """ for resource in resources.values(): - if type(resource) is not OrderedDict: - yield resource - else: + if isinstance(resource, OrderedDict): yield from self.parse_resources(resources=resource) + else: + yield resource - def show_resource_tree(self, resources: dict, indent: int = 0) -> None: + def show_resource_tree(self, resources: dict, indentation: int = 0) -> None: """Print the resources within the PE as a tree. Args: @@ -200,11 +206,13 @@ def show_resource_tree(self, resources: dict, indent: int = 0) -> None: """ for name, resource in resources.items(): - if type(resource) is not OrderedDict: - print(f"{' ' * indent} - name: {name} ID: {resource.rsrc_id}") + prefix = " " * indentation + + if isinstance(resource, OrderedDict): + print(indent(f"+ name: {name}", prefix=prefix)) + self.show_resource_tree(resources=resource, indentation=indentation + 1) else: - print(f"{' ' * indent} + name: {name}") - self.show_resource_tree(resources=resource, indent=indent + 1) + print(indent(f"- name: {name} ID: {resource.rsrc_id}", prefix=prefix)) def show_resource_info(self, resources: dict) -> None: """Print basic information about the resource as well as the header. @@ -214,18 +222,18 @@ def show_resource_info(self, resources: dict) -> None: """ for name, resource in resources.items(): - if type(resource) is not OrderedDict: + if isinstance(resource, OrderedDict): + self.show_resource_info(resources=resource) + else: print( f"* resource: {name} offset=0x{resource.offset:02x} size=0x{resource.size:02x} header: {resource.data[:64]}" # noqa: E501 ) - else: - self.show_resource_info(resources=resource) - def add_resource(self, name: str, data: bytes) -> None: + def add(self, name: str, data: bytes) -> None: # TODO raise NotImplementedError - def delete_resource(self, name: str) -> None: + def delete(self, name: str) -> None: # TODO raise NotImplementedError @@ -237,14 +245,14 @@ def update_section(self, update_offset: int) -> None: """ new_size = 0 - section_data = b"" - for idx, resource in enumerate(self.parse_resources(resources=self.pe.resources)): - if idx == 0: - # Use the offset of the first resource to account for the size of the directory header - header_size = resource.offset - self.section.virtual_address - section_data = self.section.data[:header_size] + resource_iter = iter(self.parse_resources(resources=self.resources)) + first_resource = next(resource_iter) + + header_size = first_resource.offset - self.section.virtual_address + section_data = self.section.data[:header_size] + for resource in chain([first_resource], resource_iter): # Take note of the previous offset and size so we can update any of these values after changing the data # within the resource prev_offset = resource.offset From 7be6602e1f12b07c2668a365104cb45f371f9dc4 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Mon, 24 Feb 2025 15:30:07 +0000 Subject: [PATCH 08/43] Clean up Export and export function --- dissect/executable/pe/helpers/exports.py | 47 ++++++++++-------------- 1 file changed, 19 insertions(+), 28 deletions(-) diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py index bad29e9..c6f4867 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/helpers/exports.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections import OrderedDict +from dataclasses import dataclass from io import BytesIO from typing import TYPE_CHECKING @@ -11,6 +12,7 @@ from dissect.executable.pe.pe import PE +@dataclass class ExportFunction: """Object to store the information belonging to export functions. @@ -20,27 +22,22 @@ class ExportFunction: name: The name of the function, if available. """ - def __init__(self, ordinal: int, address: int, name: bytes = b""): - self.ordinal = ordinal - self.address = address - self.name = name + ordinal: int + address: int + name: bytes | None = b"" def __str__(self) -> str: - return self.name.decode() if self.name else self.ordinal + return self.name.decode() if self.name else f"#{self.ordinal}" def __repr__(self) -> str: - return ( - f"" - if self.name - else f"" - ) + return f"" class ExportManager: def __init__(self, pe: PE, section: PESection): self.pe = pe self.section = section - self.exports = OrderedDict() + self.exports: OrderedDict[str, ExportFunction] = OrderedDict() self.parse_exports() @@ -52,9 +49,7 @@ def parse_exports(self) -> None: """ export_entry_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) - export_entry = BytesIO( - self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) - ) + export_entry = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT)) export_directory = c_pe.IMAGE_EXPORT_DIRECTORY(export_entry) # Seek to the offset of the export name @@ -63,9 +58,7 @@ def parse_exports(self) -> None: # Create a list of adresses for the exported functions export_entry.seek(export_directory.AddressOfFunctions - export_entry_va) - export_addresses = c_pe.uint32[export_directory.NumberOfFunctions].read( - export_entry - ) + export_addresses = c_pe.uint32[export_directory.NumberOfFunctions].read(export_entry) # Create a list of addresses for the exported functions that have associated names export_entry.seek(export_directory.AddressOfNames - export_entry_va) export_names = c_pe.uint32[export_directory.NumberOfNames].read(export_entry) @@ -76,19 +69,17 @@ def parse_exports(self) -> None: # Iterate over the export functions and store the information export_entry.seek(export_directory.AddressOfFunctions - export_entry_va) for idx, address in enumerate(export_addresses): + _idx = idx + 1 + key = str(_idx) + export_name: str | None = None + if idx in export_ordinals: - export_entry.seek( - export_names[export_ordinals.index(idx)] - export_entry_va - ) + entry_offset = export_names[export_ordinals.index(idx)] - export_entry_va + export_entry.seek(entry_offset) export_name = c_pe.char[None](export_entry) - self.exports[export_name.decode()] = ExportFunction( - ordinal=idx + 1, address=address, name=export_name - ) - else: - export_name = None - self.exports[str(idx + 1)] = ExportFunction( - ordinal=idx + 1, address=address, name=export_name - ) + key = export_name.decode() + + self.exports[key] = ExportFunction(ordinal=export_directory.Base + _idx, address=address, name=export_name) def add(self) -> None: raise NotImplementedError From 21c8fa6bdf6b3ed4212273a58a50327b0ffa44a6 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Mon, 24 Feb 2025 15:33:05 +0000 Subject: [PATCH 09/43] Move entries to use list comprehension --- dissect/executable/pe/helpers/relocations.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/helpers/relocations.py index 3aeb911..ab719d1 100644 --- a/dissect/executable/pe/helpers/relocations.py +++ b/dissect/executable/pe/helpers/relocations.py @@ -21,7 +21,7 @@ class RelocationManager: def __init__(self, pe: PE, section: PESection): self.pe = pe self.section = section - self.relocations = [] + self.relocations: list[dict] = [] self.parse_relocations() @@ -38,11 +38,7 @@ def parse_relocations(self) -> None: # Each entry consists of 2 bytes number_of_entries = (reloc_directory.SizeOfBlock - len(reloc_directory.dumps())) // 2 - entries = [] - for _ in range(number_of_entries): - entry = c_pe.uint16(reloc_data) - if entry: - entries.append(entry) + entries = [entry for _ in range(number_of_entries) if (entry := c_pe.uint16(reloc_data))] self.relocations.append( { From 72b4a5f5935cafa931e0da8e53d7931a4cc1e04c Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 25 Feb 2025 12:27:14 +0000 Subject: [PATCH 10/43] Compile and cache structs --- dissect/executable/pe/helpers/imports.py | 10 +++++----- dissect/executable/pe/helpers/utils.py | 21 ++++++++++++++++----- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index e52ffd3..e393da9 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -1,12 +1,12 @@ from __future__ import annotations -import struct from collections import OrderedDict from io import BytesIO from typing import TYPE_CHECKING, BinaryIO from dissect.executable.pe.c_pe import c_pe from dissect.executable.pe.helpers import utils +from dissect.executable.pe.helpers.utils import create_struct if TYPE_CHECKING: from collections.abc import Iterator @@ -281,10 +281,10 @@ def _build_module_imports(self, functions: list[ImportFunction]) -> int: """ function_offsets = [] - + _hint_struct = create_struct(" bytes: """ packing = " bytes: output = b"".join(thunkdata) - self.thunks.append(self.pe.image_thunk_data(output)) + self.thunks.append(self._thunk_data(output)) return output diff --git a/dissect/executable/pe/helpers/utils.py b/dissect/executable/pe/helpers/utils.py index b73ee29..c557802 100644 --- a/dissect/executable/pe/helpers/utils.py +++ b/dissect/executable/pe/helpers/utils.py @@ -1,3 +1,18 @@ +from __future__ import annotations + +import struct +from functools import lru_cache +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from dissect.executable.pe import PE, PESection + + +@lru_cache +def create_struct(packing: str) -> struct.Struct: + return struct.Struct(packing) + + def align_data(data: bytes, blocksize: int) -> bytes: """Align the new data according to the file alignment as specified in the PE header. @@ -10,11 +25,7 @@ def align_data(data: bytes, blocksize: int) -> bytes: """ needs_alignment = len(data) % blocksize - return ( - data - if not needs_alignment - else data + ((blocksize - needs_alignment) * b"\x00") - ) + return data if not needs_alignment else data + ((blocksize - needs_alignment) * b"\x00") def align_int(integer: int, blocksize: int) -> int: From e70d0067fa3adcab9b99b98ee4e6b3daeaeda9aa Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 25 Feb 2025 12:30:30 +0000 Subject: [PATCH 11/43] Move architecture specific code to the managers - Moves pe high bit to import functions --- dissect/executable/pe/helpers/imports.py | 77 +++++++++++++++--------- dissect/executable/pe/helpers/patcher.py | 5 +- dissect/executable/pe/helpers/tls.py | 27 ++++++--- dissect/executable/pe/pe.py | 25 +++----- 4 files changed, 82 insertions(+), 52 deletions(-) diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index e393da9..31d641e 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -31,7 +31,7 @@ class ImportModule: def __init__( self, name: bytes, - import_descriptor: int | None, + import_descriptor: int | None | c_pe.IMAGE_IMPORT_DESCRIPTOR, module_va: int, name_va: int, first_thunk: int, @@ -41,7 +41,7 @@ def __init__( self.module_va = module_va self.name_va = name_va self.first_thunk = first_thunk - self.functions = [] + self.functions: list[ImportFunction] = [] def __str__(self) -> str: return self.name.decode() @@ -61,13 +61,19 @@ class ImportFunction: def __init__( self, pe: PE, - thunkdata: c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64 | None, + thunkdata: c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64, + high_bit: int, name: str = "", ): self.pe = pe self.thunkdata = thunkdata + self.high_bit = high_bit self._name = name + @property + def ordinal(self) -> int: + return self.thunkdata.u1.AddressOfData & self.high_bit + @property def name(self) -> str: """Return the name of the import function if available, otherwise return the ordinal of the function. @@ -83,13 +89,11 @@ def name(self) -> str: # For the case thunkdata is not defined, such as during the `add` return "" - ordinal = self.thunkdata.u1.AddressOfData & self.pe._high_bit + data_address = self.thunkdata.u1.AddressOfData - if not ordinal: - self.pe.seek(self.thunkdata.u1.AddressOfData + 2) + if not (entry := self.ordinal): + self.pe.seek(data_address + 2) entry = c_pe.char[None](self.pe).decode() - else: - entry = ordinal if isinstance(entry, int): return str(entry) @@ -120,10 +124,22 @@ def __init__(self, pe: PE, section: PESection): self.new_size_of_image = 0 self.section_data = bytearray() self.imports: OrderedDict[str, ImportModule] = OrderedDict() - self.thunks = [] + self.thunks: list[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = [] + + self._thunk_data: type[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = None + self._high_bit: int = None + self.set_architecture(pe) self.parse_imports() + def set_architecture(self, pe: PE) -> None: + if pe.is64bit(): + self._thunk_data = c_pe.IMAGE_THUNK_DATA64 + self._high_bit = 1 << 63 + else: + self._thunk_data = c_pe.IMAGE_THUNK_DATA32 + self._high_bit = 1 << 31 + def parse_imports(self) -> None: """Parse the imports of the PE file. @@ -135,25 +151,28 @@ def parse_imports(self) -> None: # Loop over the entries for descriptor_va, import_descriptor in self.import_descriptors(import_data=import_data): - if import_descriptor.Name not in [0xFFFFF800, 0x0]: - self.pe.seek(import_descriptor.Name) - modulename = c_pe.char[None](self.pe) + if import_descriptor.Name in [0xFFFFF800, 0x0]: + continue - # Use the OriginalFirstThunk if available, FirstThunk otherwise - first_thunk = import_descriptor.OriginalFirstThunk or import_descriptor.FirstThunk + self.pe.seek(import_descriptor.Name) + modulename: bytes = c_pe.char[None](self.pe) - module = ImportModule( - name=modulename, - import_descriptor=import_descriptor, - module_va=descriptor_va, - name_va=import_descriptor.Name, - first_thunk=first_thunk, - ) + # Use the OriginalFirstThunk if available, FirstThunk otherwise + first_thunk = import_descriptor.OriginalFirstThunk or import_descriptor.FirstThunk - for thunkdata in self.parse_thunks(offset=first_thunk): - module.functions.append(ImportFunction(pe=self.pe, thunkdata=thunkdata)) + module = ImportModule( + name=modulename, + import_descriptor=import_descriptor, + module_va=descriptor_va, + name_va=import_descriptor.Name, + first_thunk=first_thunk, + ) - self.imports[modulename.decode()] = module + module.functions.extend( + ImportFunction(pe=self.pe, thunkdata=thunkdata, high_bit=self._high_bit) + for thunkdata in self.parse_thunks(offset=first_thunk) + ) + self.imports[modulename.decode()] = module def import_descriptors(self, import_data: BinaryIO) -> Iterator[tuple[int, c_pe.IMAGE_IMPORT_DESCRIPTOR]]: """Parse the import descriptors of the PE file. @@ -186,7 +205,7 @@ def parse_thunks(self, offset: int) -> Iterator[c_pe.IMAGE_THUNK_DATA32 | c_pe.I self.pe.seek(offset) while True: - thunkdata = self.pe.image_thunk_data(self.pe) + thunkdata = self._thunk_data(self.pe) if not thunkdata.u1.Function: break @@ -203,7 +222,7 @@ def add(self, dllname: str, functions: list[str]) -> None: self.last_section = next(reversed(self.pe.patched_sections.values())) # Build a dummy import module - self.imports[dllname] = ImportModule( + _module = ImportModule( name=dllname.encode(), import_descriptor=None, module_va=0, @@ -211,10 +230,12 @@ def add(self, dllname: str, functions: list[str]) -> None: first_thunk=0, ) # Build the dummy module functions - self.imports[dllname].functions.extend( - ImportFunction(pe=self.pe, thunkdata=None, name=function) for function in functions + _module.functions.extend( + ImportFunction(pe=self.pe, thunkdata=None, high_bit=self._high_bit, name=function) for function in functions ) + self.imports[dllname] = _module + # Rebuild the import table with the new import module and functions self.build_import_table() diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index 30282d7..aa8703f 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -1,5 +1,6 @@ from __future__ import annotations +import copy import struct from io import BytesIO from typing import TYPE_CHECKING @@ -157,7 +158,7 @@ def _patch_import_rvas(self) -> None: thunkdata = function.thunkdata # Check if we're dealing with an ordinal entry, if it's an ordinal entry we don't need # to patch since it's not an RVA - if thunkdata.u1.AddressOfData & self.pe._high_bit: + if function.ordinal: patched_thunkdata += thunkdata.dumps() continue @@ -300,7 +301,7 @@ def _patch_tls_rvas(self) -> None: return self.seek(directory_va) - tls_directory = self.pe.image_tls_directory(self.patched_pe) + tls_directory = self.pe.tls_mgr._tls_directory(self.patched_pe) image_base = self.pe.optional_header.ImageBase diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/helpers/tls.py index d757e72..50f68a1 100644 --- a/dissect/executable/pe/helpers/tls.py +++ b/dissect/executable/pe/helpers/tls.py @@ -21,22 +21,35 @@ def __init__(self, pe: PE, section: PESection): self.pe = pe self.section = section self.callbacks = [] - self.tls = None + self.tls: c_pe._IMAGE_TLS_DIRECTORY32 | c_pe._IMAGE_TLS_DIRECTORY64 = None + + self._read_address: type[c_pe.uint64 | c_pe.uint32] = None + self._tls_directory: type[c_pe._IMAGE_TLS_DIRECTORY32 | c_pe._IMAGE_TLS_DIRECTORY64] = None self._data = b"" + self._image_base = pe.optional_header.ImageBase + self.set_architecture(pe) self.parse_tls() + def set_architecture(self, pe: PE) -> None: + if pe.is64bit(): + self._read_address = c_pe.uint64 + self._tls_directory = c_pe._IMAGE_TLS_DIRECTORY64 + else: + self._read_address = c_pe.uint32 + self._tls_directory = c_pe._IMAGE_TLS_DIRECTORY32 + def parse_tls(self) -> None: """Parse the TLS directory entry of the PE file when present.""" - tls_data = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_TLS)) - self.tls = self.pe.image_tls_directory(tls_data) + tls_data = BytesIO(self.section.data) + self.tls = self._tls_directory(tls_data) - self.pe.seek(self.tls.AddressOfCallBacks - self.pe.optional_header.ImageBase) + self.pe.seek(self.tls.AddressOfCallBacks - self._image_base) # Parse the TLS callback addresses if present while True: - callback_address = self.pe.read_address(self.pe) + callback_address = self._read_address(self.pe) if not callback_address: break self.callbacks.append(callback_address) @@ -72,7 +85,7 @@ def read_data(self) -> bytes: """ return self.pe.virtual_read( - address=self.tls.StartAddressOfRawData - self.pe.optional_header.ImageBase, + address=self.tls.StartAddressOfRawData - self._image_base, size=self.size, ) @@ -105,7 +118,7 @@ def data(self, value: bytes) -> None: section_data.write(self.tls.dumps()) # Write the new TLS data to the section - start_address_rva = self.tls.StartAddressOfRawData - self.pe.optional_header.ImageBase + start_address_rva = self.tls.StartAddressOfRawData - self._image_base start_address_section_offset = start_address_rva - self.section.virtual_address section_data.seek(start_address_section_offset) section_data.write(self._data) diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index 26b7afa..e8e2edd 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -101,6 +101,9 @@ def _valid(self) -> bool: self.pe_file.seek(self.mz_header.e_lfanew) return c_pe.uint32(self.pe_file) == 0x4550 + def is64bit(self) -> bool: + return self.file_header.Machine == c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64 + def parse_headers(self) -> None: """Function to parse the basic PE headers: - DOS header @@ -125,7 +128,7 @@ def parse_headers(self) -> None: self.pe_file.seek(image_nt_headers_offset) # Set the architecture specific settings - self._set_pe_architecture() + self._check_architecture() if self.file_header.Machine == c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64: self.nt_headers = c_pe.IMAGE_NT_HEADERS64(self.pe_file) else: @@ -133,24 +136,16 @@ def parse_headers(self) -> None: self.optional_header = self.nt_headers.OptionalHeader - def _set_pe_architecture(self) -> None: - """Set the architecture specific settings. Some of the structs are architecture specific. + def _check_architecture(self) -> None: + """Check whether the architecture belonging to the binary is one we support. Raises: InvalidArchitecture if the architecture is not supported or unknown. """ - - if self.file_header.Machine == c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64: - self.image_thunk_data = c_pe.IMAGE_THUNK_DATA64 - self.image_tls_directory = c_pe.IMAGE_TLS_DIRECTORY64 - self._high_bit = 1 << 63 - self.read_address = c_pe.uint64 - elif self.file_header.Machine == c_pe.MachineType.IMAGE_FILE_MACHINE_I386: - self.image_thunk_data = c_pe.IMAGE_THUNK_DATA32 - self.image_tls_directory = c_pe.IMAGE_TLS_DIRECTORY32 - self._high_bit = 1 << 31 - self.read_address = c_pe.uint32 - else: + if self.file_header.Machine not in [ + c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64, + c_pe.MachineType.IMAGE_FILE_MACHINE_I386, + ]: raise InvalidArchitecture(f"Invalid architecture found: {self.file_header.Machine:02x}") def parse_section_header(self) -> None: From 9a64a0af320383dfdfc75c815ad42bf88b562449 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 25 Feb 2025 14:17:28 +0000 Subject: [PATCH 12/43] Move _valid check to `parse_headers` --- dissect/executable/pe/pe.py | 19 +++++++------------ tests/test_pe.py | 7 ------- 2 files changed, 7 insertions(+), 19 deletions(-) diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index e8e2edd..772925b 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -91,16 +91,6 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): # Parsing the directories present in the PE self.parse_directories() - def _valid(self) -> bool: - """Check if the PE file is a valid PE file. By looking for the "PE" signature at the offset of e_lfanew. - - Returns: - `True` if the file is a valid PE file, `False` otherwise. - """ - - self.pe_file.seek(self.mz_header.e_lfanew) - return c_pe.uint32(self.pe_file) == 0x4550 - def is64bit(self) -> bool: return self.file_header.Machine == c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64 @@ -118,8 +108,13 @@ def parse_headers(self) -> None: """ self.mz_header = c_pe.IMAGE_DOS_HEADER(self.pe_file) + if self.mz_header.e_magic != 0x5A4D: + raise InvalidPE("File is not a valid PE file, MZ header has wrong signature.") + + self.pe_file.seek(self.mz_header.e_lfanew) + nt_signature = c_pe.uint32(self.pe_file) - if not self._valid(): + if nt_signature != 0x4550: raise InvalidPE("file is not a valid PE file") self.file_header = c_pe.IMAGE_FILE_HEADER(self.pe_file) @@ -129,7 +124,7 @@ def parse_headers(self) -> None: # Set the architecture specific settings self._check_architecture() - if self.file_header.Machine == c_pe.MachineType.IMAGE_FILE_MACHINE_AMD64: + if self.is64bit(): self.nt_headers = c_pe.IMAGE_NT_HEADERS64(self.pe_file) else: self.nt_headers = c_pe.IMAGE_NT_HEADERS(self.pe_file) diff --git a/tests/test_pe.py b/tests/test_pe.py index 1bac7bd..3f83053 100644 --- a/tests/test_pe.py +++ b/tests/test_pe.py @@ -8,13 +8,6 @@ from .util import data_file -def test_pe_valid_signature() -> None: - with data_file("testexe.exe").open("rb") as pe_fh: - pe = PE(pe_file=pe_fh) - - assert pe._valid() is True - - def test_pe_invalid_signature() -> None: with pytest.raises(InvalidPE): PE(BytesIO(b"MZ" + b"\x00" * 400)) From 4ecc96c1cc2e099b25d92fd86b59dbe0416cd726 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 4 Mar 2025 13:07:24 +0000 Subject: [PATCH 13/43] Reduce indentation _patch_import_rvas --- dissect/executable/pe/helpers/patcher.py | 92 +++++++++++++----------- 1 file changed, 49 insertions(+), 43 deletions(-) diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index aa8703f..b26eea0 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -143,52 +143,58 @@ def _patch_import_rvas(self) -> None: import_descriptor = module.import_descriptor patched_thunkdata = bytearray() - if import_descriptor.Name != 0xFFFFF800 and import_descriptor.Name != 0x0: - old_first_thunk = import_descriptor.FirstThunk - - first_thunk_offset = old_first_thunk - original_directory_va - import_descriptor.FirstThunk = abs(directory_va + first_thunk_offset) - - import_descriptor.OriginalFirstThunk = import_descriptor.FirstThunk - - name_offset = import_descriptor.Name - original_directory_va - import_descriptor.Name = abs(directory_va + name_offset) - - for function in module.functions: - thunkdata = function.thunkdata - # Check if we're dealing with an ordinal entry, if it's an ordinal entry we don't need - # to patch since it's not an RVA - if function.ordinal: - patched_thunkdata += thunkdata.dumps() - continue - - # Check the original RVA associated with the AddressOfData field in the thunkdata, retrieve the - # original VA - # and use it to also select the patched virtual address of this section that the RVA is located in - for name, section in self.pe.sections.items(): - if thunkdata.u1.AddressOfData in range( - section.virtual_address, - section.virtual_address + section.virtual_size, - ): - virtual_address = section.virtual_address - new_virtual_address = self.pe.patched_sections[name].virtual_address - break - - # Calculate the offset using the VA of the section and update the thunkdata - va_offset = thunkdata.u1.AddressOfData - virtual_address - new_thunkdata = new_virtual_address + va_offset - thunkdata.u1.AddressOfData = new_thunkdata - thunkdata.u1.ForwarderString = new_thunkdata - thunkdata.u1.Function = new_thunkdata - thunkdata.u1.Ordinal = new_thunkdata + if import_descriptor.Name in [0xFFFFF800, 0x0]: + continue - patched_thunkdata += thunkdata.dumps() + old_first_thunk = import_descriptor.FirstThunk + + first_thunk_offset = old_first_thunk - original_directory_va + import_descriptor.FirstThunk = abs(directory_va + first_thunk_offset) - # Write the thunk data into the patched PE - self.seek(import_descriptor.FirstThunk) - self.patched_pe.write(patched_thunkdata) + import_descriptor.OriginalFirstThunk = import_descriptor.FirstThunk - patched_import_data += import_descriptor.dumps() + name_offset = import_descriptor.Name - original_directory_va + import_descriptor.Name = abs(directory_va + name_offset) + + for function in module.functions: + thunkdata = function.thunkdata + # Check if we're dealing with an ordinal entry, if it's an ordinal entry we don't need + # to patch since it's not an RVA + if function.ordinal: + patched_thunkdata += thunkdata.dumps() + continue + + # Check the original RVA associated with the AddressOfData field in the thunkdata, retrieve the + # original VA + # and use it to also select the patched virtual address of this section that the RVA is located in + for name, section in self.pe.sections.items(): + if ( + section.virtual_address + <= thunkdata.u1.AddressOfData + < section.virtual_address + section.virtual_size + ): + virtual_address = section.virtual_address + new_virtual_address = self.pe.patched_sections[name].virtual_address + break + + # Calculate the offset using the VA of the section and update the thunkdata + va_offset = thunkdata.u1.AddressOfData - virtual_address + new_address = new_virtual_address + va_offset + + tmp_thunkdata = copy.deepcopy(thunkdata) + + tmp_thunkdata.u1.AddressOfData = new_address + tmp_thunkdata.u1.ForwarderString = new_address + tmp_thunkdata.u1.Function = new_address + tmp_thunkdata.u1.Ordinal = new_address + + patched_thunkdata += tmp_thunkdata.dumps() + + # Write the thunk data into the patched PE + self.seek(import_descriptor.FirstThunk) + self.patched_pe.write(patched_thunkdata) + + patched_import_data += import_descriptor.dumps() self.seek(directory_va) self.patched_pe.write(patched_import_data) From 6cb4a742f926691ae1bb2adcd04e1f0d86641c8d Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 4 Mar 2025 13:17:22 +0000 Subject: [PATCH 14/43] Use a dataclass instead of a dictionary for raw_resources --- dissect/executable/pe/helpers/patcher.py | 10 ++-- dissect/executable/pe/helpers/resources.py | 54 +++++++++++++++------- 2 files changed, 42 insertions(+), 22 deletions(-) diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index b26eea0..e8d17b5 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -276,13 +276,13 @@ def _patch_rsrc_rvas(self) -> None: section_data = BytesIO() self.seek(directory_va) - for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"]): - entry_offset = rsrc_entry["offset"] - entry = rsrc_entry["entry"] + for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc.data_offset): + entry_offset = rsrc_entry.offset + entry = rsrc_entry.entry if isinstance(entry, c_pe.IMAGE_RESOURCE_DATA_ENTRY): - rsrc_obj = rsrc_entry["resource"] - data_offset = rsrc_entry["data_offset"] + rsrc_obj = rsrc_entry.resource + data_offset = rsrc_entry.data_offset # Update the offset of the entry to match with the new directory VA rsrc_obj.offset = directory_va + data_offset diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index 3255a14..9d50ff0 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections import OrderedDict +from dataclasses import dataclass from io import BytesIO from itertools import chain from textwrap import indent @@ -20,7 +21,14 @@ from dissect.executable.pe.pe import PE -class ResourceManager: +@dataclass +class RawResource: + offset: int + entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY | c_pe.IMAGE_RESOURCE_DIRECTORY + data_offset: int + data: bytes | None = None + resource: Resource | None = None + """Base class to perform actions regarding the resources within the PE file. Args: @@ -32,7 +40,7 @@ def __init__(self, pe: PE, section: PESection): self.pe = pe self.section = section self.resources: OrderedDict[str, Resource] = OrderedDict() - self.raw_resources: list[dict] = [] + self.raw_resources: list[RawResource] = [] self.parse() @@ -59,7 +67,13 @@ def _read_entries( for _ in range(directory.NumberOfNamedEntries + directory.NumberOfIdEntries): entry_offset = data.tell() entry = c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY(data) - self.raw_resources.append({"offset": entry_offset, "entry": entry, "data_offset": entry_offset}) + self.raw_resources.append( + RawResource( + offset=entry_offset, + entry=entry, + data_offset=entry_offset, + ) + ) entries.append(entry) return entries @@ -88,13 +102,13 @@ def _handle_data_entry(self, data: BinaryIO, entry: c_pe.IMAGE_RESOURCE_DIRECTOR rc_type=rc_type, ) self.raw_resources.append( - { - "offset": entry.OffsetToDirectory, - "entry": data_entry, - "data": _data, - "data_offset": raw_offset, - "resource": rsrc, - } + RawResource( + offset=entry.OffsetToDirectory, + entry=data_entry, + data=_data, + data_offset=raw_offset, + resource=rsrc, + ) ) return rsrc @@ -118,7 +132,13 @@ def _read_resource(self, data: BinaryIO, offset: int, level: int = 1) -> Ordered data.seek(offset) directory = c_pe.IMAGE_RESOURCE_DIRECTORY(data) - self.raw_resources.append({"offset": offset, "entry": directory, "data_offset": offset}) + self.raw_resources.append( + RawResource( + offset=offset, + entry=directory, + data_offset=offset, + ) + ) entries = self._read_entries(data, directory) @@ -386,19 +406,19 @@ def data(self, value: bytes) -> None: prev_offset = 0 prev_size = 0 - for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc["data_offset"]): - entry_offset = rsrc_entry["offset"] - entry = rsrc_entry["entry"] + for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc.data_offset): + entry_offset = rsrc_entry.offset + entry = rsrc_entry.entry if isinstance(entry, c_pe.IMAGE_RESOURCE_DATA_ENTRY): - rsrc_obj = rsrc_entry["resource"] - data_offset = rsrc_entry["data_offset"] + rsrc_obj = rsrc_entry.resource + data_offset = rsrc_entry.data_offset # Normally the data is separated by a null byte, increment the new offset by 1 new_data_offset = prev_offset + prev_size # if new_data_offset and (new_data_offset > data_offset or new_data_offset < data_offset): if new_data_offset and new_data_offset != data_offset: data_offset = new_data_offset - rsrc_entry["data_offset"] = data_offset + rsrc_entry.data_offset = data_offset rsrc_obj.offset = self.section.virtual_address + data_offset data = rsrc_obj.data From 7affc195ea6065e08796a4be179dc0c57dac3735 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 4 Mar 2025 13:23:54 +0000 Subject: [PATCH 15/43] replace struct.pack with create_struct().pack --- dissect/executable/pe/helpers/patcher.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index e8d17b5..950beca 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -242,8 +242,10 @@ def _patch_export_rvas(self) -> None: new_address = self.pe.patched_sections[section.name].virtual_address + address_offset new_function_rvas.append(new_address) + rva_struct = utils.create_struct(" None: new_name_rvas.append(new_address) for name_rva in new_name_rvas: - name_rvas += struct.pack(" Date: Tue, 4 Mar 2025 13:40:07 +0000 Subject: [PATCH 16/43] Make a function to search through sections --- dissect/executable/pe/helpers/exports.py | 2 +- dissect/executable/pe/pe.py | 92 +++++++++++------------- 2 files changed, 42 insertions(+), 52 deletions(-) diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py index c6f4867..a6d124b 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/helpers/exports.py @@ -48,7 +48,7 @@ def parse_exports(self) -> None: name (if available), the call ordinal, and the function address. """ - export_entry_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) + export_entry_va = self.pe.directory_entry_rva(c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) export_entry = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT)) export_directory = c_pe.IMAGE_EXPORT_DIRECTORY(export_entry) diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index 772925b..2329af6 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -26,10 +26,10 @@ ) if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterable, Iterator from dissect.cstruct.cstruct import cstruct - from dissect.cstruct.types.enum import EnumInstance + from dissect.cstruct.types.enum import Enum from dissect.executable.pe.helpers.resources import Resource @@ -51,10 +51,10 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): # Make sure we reset any kind of pointers within the PE file before continueing self.pe_file.seek(0) - self.mz_header = None - self.file_header = None - self.nt_headers = None - self.optional_header = None + self.mz_header: c_pe.IMAGE_DOS_HEADER = None + self.nt_headers: c_pe.IMAGE_NT_HEADERS | c_pe.IMAGE_NT_HEADERS64 = None + self.file_header: c_pe.IMAGE_FILE_HEADER = None + self.optional_header: c_pe.IMAGE_OPTIONAL_HEADER | c_pe.IMAGE_OPTIONAL_HEADER64 = None self.section_header_offset = 0 self.last_section_offset = 0 @@ -65,8 +65,8 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.imports: OrderedDict[str, imports.ImportModule] = None self.exports: OrderedDict[str, exports.ExportFunction] = None self.resources: OrderedDict[str, resources.Resource] = None - self.raw_resources = None - self.relocations: list[dict] = None + self.raw_resources: list[resources.RawResource] = [] + self.relocations: list[relocations.Relocation] = [] self.tls_callbacks = None self.directories = OrderedDict() @@ -77,12 +77,10 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): # The offset of the section header is always at the end of the NT headers self.section_header_offset = self.pe_file.tell() - self.imagebase = self.optional_header.ImageBase - self.file_alignment = self.optional_header.FileAlignment - self.section_alignment = self.optional_header.SectionAlignment - - self.base_address = self.optional_header.ImageBase - + self.imagebase = self.nt_headers.OptionalHeader.ImageBase + self.file_alignment = self.nt_headers.OptionalHeader.FileAlignment + self.section_alignment = self.nt_headers.OptionalHeader.SectionAlignment + self.base_address = self.nt_headers.OptionalHeader.ImageBase self.timestamp = datetime.fromtimestamp(self.file_header.TimeDateStamp, tz=timezone.utc) # Parse the section header @@ -159,6 +157,13 @@ def parse_section_header(self) -> None: self.last_section_offset = self.sections[next(reversed(self.sections))].offset + def _section_in_range(self, address: int, values: Iterable[sections.PESection]) -> sections.PESection | None: + for section in values: + if section.virtual_address <= address < section.virtual_address + section.virtual_size: + return section + + return None + def section(self, va: int = 0, name: str = "") -> sections.PESection | None: """Function to retrieve a section based on the given virtual address or name. @@ -170,15 +175,10 @@ def section(self, va: int = 0, name: str = "") -> sections.PESection | None: A `PESection` object. """ - if not name: - for section in self.sections.values(): - if va in range( - section.virtual_address, - section.virtual_address + section.virtual_size, - ): - return section - return None - return self.sections[name] + if name: + return self.sections[name] + + return self._section_in_range(va, self.sections.values()) def patched_section(self, va: int = 0, name: str = "") -> sections.PESection | None: """Function to retrieve a patched section based on the given virtual address or name. @@ -191,15 +191,10 @@ def patched_section(self, va: int = 0, name: str = "") -> sections.PESection | N A `PESection` object. """ - if not name: - for section in self.patched_sections.values(): - if va in range( - section.virtual_address, - section.virtual_address + section.virtual_size, - ): - return section - return None - return self.patched_sections[name] + if name: + return self.patched_sections[name] + + return self._section_in_range(va, self.patched_sections.values()) def datadirectory_section(self, index: int) -> sections.PESection: """Return the section that contains the given virtual address. @@ -211,10 +206,9 @@ def datadirectory_section(self, index: int) -> sections.PESection: The section that contains the given virtual address. """ - va = self.directory_va(index=index) - for section in self.patched_sections.values(): - if va >= section.virtual_address and va < section.virtual_address + section.virtual_size: - return section + va = self.directory_entry_rva(index=index) + if section := self._section_in_range(va, self.patched_sections.values()): + return section raise InvalidVA(f"VA not found in sections: {va:#04x}") @@ -263,7 +257,7 @@ def parse_directories(self) -> None: self.tls_mgr = tls.TLSManager(pe=self, section=section) self.tls_callbacks = self.tls_mgr.callbacks - def get_resource_type(self, rsrc_id: str | EnumInstance) -> Iterator[Resource]: + def get_resource_type(self, rsrc_id: str | Enum) -> Iterator[Resource]: """Yields a generator containing all of the nodes within the resources that contain the requested ID. The ID can be either given by name or its value. @@ -293,10 +287,8 @@ def virtual_address(self, address: int) -> int: if self.virtual: return address - for section in self.patched_sections.values(): - max_address = section.virtual_address + section.virtual_size - if address >= section.virtual_address and address < max_address: - return section.pointer_to_raw_data + (address - section.virtual_address) + if section := self._section_in_range(address, self.patched_sections.values()): + return section.pointer_to_raw_data + (address - section.virtual_address) raise InvalidVA(f"VA not found in sections: {address:#04x}") @@ -311,8 +303,7 @@ def raw_address(self, offset: int) -> int: """ for section in self.patched_sections.values(): - max_address = section.pointer_to_raw_data + section.size_of_raw_data - if offset >= section.pointer_to_raw_data and offset < max_address: + if section.pointer_to_raw_data <= offset < section.pointer_to_raw_data + section.size_of_raw_data: return section.virtual_address + (offset - section.pointer_to_raw_data) raise InvalidAddress(f"Raw address not found in sections: {offset:#04x}") @@ -401,10 +392,9 @@ def write(self, data: bytes) -> None: print(self.patched_sections) # Update the section data - for section in self.patched_sections.values(): - if section.virtual_address <= offset and section.virtual_address + section.virtual_size >= offset: - self.seek(address=section.virtual_address) - section.data = self.read(size=section.virtual_size) + if section := self._section_in_range(offset, self.patched_sections.values()): + self.seek(address=section.virtual_address) + section.data = self.read(size=section.virtual_size) def read_image_directory(self, index: int) -> bytes: """Read the PE file image directory entry of a given index. @@ -419,7 +409,7 @@ def read_image_directory(self, index: int) -> bytes: directory_entry = self.optional_header.DataDirectory[index] return self.virtual_read(address=directory_entry.VirtualAddress, size=directory_entry.Size) - def directory_va(self, index: int) -> int: + def directory_entry_rva(self, index: int) -> int: """Returns the virtual address of a directory given its index. Args: @@ -451,7 +441,7 @@ def debug(self) -> cstruct | None: """ debug_directory_entry = self.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_DEBUG) - image_directory_size = len(c_pe.IMAGE_DEBUG_DIRECTORY) + image_directory_size = c_pe.IMAGE_DEBUG_DIRECTORY.size for _ in range(len(debug_directory_entry) // image_directory_size): entry = c_pe.IMAGE_DEBUG_DIRECTORY(debug_directory_entry) @@ -461,14 +451,14 @@ def debug(self) -> cstruct | None: return c_cv_info.CV_INFO_PDB70(dbg_entry) return None - def get_section(self, segment_index: int) -> tuple[str, sections.PESection]: + def get_section(self, segment_index: int) -> sections.PESection: """Retrieve the section of the PE by index. Args: segment_index: The segment to retrieve based on the order within the PE. Returns: - A `tuple` contianing the section name and attributes as `PESection`. + A `PESection` corresponding to the segment_index. """ sections = list(self.sections.items()) From 63548256c694711dce29e37f11ae1d21eaf62caa Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Wed, 26 Mar 2025 13:26:20 +0000 Subject: [PATCH 17/43] add some checks during the patching if certain sections do not exist --- dissect/executable/pe/helpers/imports.py | 10 ++- dissect/executable/pe/helpers/patcher.py | 83 +++++++++--------------- 2 files changed, 38 insertions(+), 55 deletions(-) diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index 31d641e..ecdde24 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -70,9 +70,14 @@ def __init__( self.high_bit = high_bit self._name = name + @property + def data_address(self) -> int: + """Shows the AddressOfData of the thunk data.""" + return self.thunkdata.u1.AddressOfData + @property def ordinal(self) -> int: - return self.thunkdata.u1.AddressOfData & self.high_bit + return self.data_address & self.high_bit @property def name(self) -> str: @@ -89,10 +94,9 @@ def name(self) -> str: # For the case thunkdata is not defined, such as during the `add` return "" - data_address = self.thunkdata.u1.AddressOfData if not (entry := self.ordinal): - self.pe.seek(data_address + 2) + self.pe.seek(self.data_address + 2) entry = c_pe.char[None](self.pe).decode() if isinstance(entry, int): diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index 950beca..47a68ac 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -1,7 +1,6 @@ from __future__ import annotations import copy -import struct from io import BytesIO from typing import TYPE_CHECKING @@ -10,7 +9,6 @@ if TYPE_CHECKING: from dissect.executable import PE - from dissect.executable.pe.helpers.sections import PESection class Patcher: @@ -127,7 +125,7 @@ def _patch_import_rvas(self) -> None: patched_import_data = bytearray() # Get the directory entry virtual adddress, this is the updated address if it has been patched. - directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT) + directory_va = self.pe.directory_entry_rva(c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT) if not directory_va: return @@ -139,8 +137,8 @@ def _patch_import_rvas(self) -> None: # Loop over the imports of the PE to patch the RVA's of the import descriptors and the associated thunkdata # entries - for name, module in self.pe.imports.items(): - import_descriptor = module.import_descriptor + for module in self.pe.imports.values(): + import_descriptor: c_pe.IMAGE_IMPORT_DESCRIPTOR = module.import_descriptor patched_thunkdata = bytearray() if import_descriptor.Name in [0xFFFFF800, 0x0]: @@ -157,32 +155,25 @@ def _patch_import_rvas(self) -> None: import_descriptor.Name = abs(directory_va + name_offset) for function in module.functions: - thunkdata = function.thunkdata # Check if we're dealing with an ordinal entry, if it's an ordinal entry we don't need # to patch since it's not an RVA if function.ordinal: - patched_thunkdata += thunkdata.dumps() + patched_thunkdata += function.thunkdata.dumps() continue # Check the original RVA associated with the AddressOfData field in the thunkdata, retrieve the # original VA # and use it to also select the patched virtual address of this section that the RVA is located in - for name, section in self.pe.sections.items(): - if ( - section.virtual_address - <= thunkdata.u1.AddressOfData - < section.virtual_address + section.virtual_size - ): - virtual_address = section.virtual_address - new_virtual_address = self.pe.patched_sections[name].virtual_address - break + if section := self.pe.section(function.data_address): + virtual_address = section.virtual_address + new_virtual_address = self.pe.patched_section(name=section.name).virtual_address # Calculate the offset using the VA of the section and update the thunkdata - va_offset = thunkdata.u1.AddressOfData - virtual_address + va_offset = function.data_address - virtual_address new_address = new_virtual_address + va_offset - tmp_thunkdata = copy.deepcopy(thunkdata) - + # Avoid overwriting the original data + tmp_thunkdata = copy.deepcopy(function.thunkdata) tmp_thunkdata.u1.AddressOfData = new_address tmp_thunkdata.u1.ForwarderString = new_address tmp_thunkdata.u1.Function = new_address @@ -202,7 +193,7 @@ def _patch_import_rvas(self) -> None: def _patch_export_rvas(self) -> None: """Function to patch the RVAs of the export directory and the associated function and name RVA's.""" - directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) + directory_va = self.pe.directory_entry_rva(c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) if not directory_va: return @@ -257,6 +248,9 @@ def _patch_export_rvas(self) -> None: export_names = c_pe.uint32[export_directory.NumberOfNames].read(self.patched_pe) for name_address in export_names: section = self.pe.section(va=name_address) + if section is None: + continue + address_offset = name_address - section.virtual_address new_address = self.pe.patched_sections[section.name].virtual_address + address_offset new_name_rvas.append(new_address) @@ -271,7 +265,7 @@ def _patch_export_rvas(self) -> None: def _patch_rsrc_rvas(self) -> None: """Function to patch the RVAs of the resource directory and the associated resource data RVA's.""" - directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE) + directory_va = self.pe.directory_entry_rva(c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE) if not directory_va: return @@ -304,7 +298,7 @@ def _patch_rsrc_rvas(self) -> None: def _patch_tls_rvas(self) -> None: """Function to patch the RVAs of the TLS directory and the associated TLS callbacks.""" - directory_va = self.pe.directory_va(c_pe.IMAGE_DIRECTORY_ENTRY_TLS) + directory_va = self.pe.directory_entry_rva(c_pe.IMAGE_DIRECTORY_ENTRY_TLS) if not directory_va: return @@ -314,41 +308,26 @@ def _patch_tls_rvas(self) -> None: image_base = self.pe.optional_header.ImageBase # Patch the TLS StartAddressOfRawData and EndAddressOfRawData - section = self.pe.section(va=tls_directory.StartAddressOfRawData - image_base) - start_address_offset = tls_directory.StartAddressOfRawData - section.virtual_address - tls_directory.StartAddressOfRawData = ( - self.pe.patched_sections[section.name].virtual_address + start_address_offset - ) - end_address_offset = tls_directory.EndAddressOfRawData - tls_directory.StartAddressOfRawData - tls_directory.EndAddressOfRawData = tls_directory.StartAddressOfRawData + end_address_offset + if section := self.pe.section(va=tls_directory.StartAddressOfRawData - image_base): + start_address_offset = tls_directory.StartAddressOfRawData - section.virtual_address + tls_directory.StartAddressOfRawData = ( + self.pe.patched_sections[section.name].virtual_address + start_address_offset + ) + end_address_offset = tls_directory.EndAddressOfRawData - tls_directory.StartAddressOfRawData + tls_directory.EndAddressOfRawData = tls_directory.StartAddressOfRawData + end_address_offset # Patch the TLS callbacks address - section = self.pe.section(va=tls_directory.AddressOfCallBacks - image_base) - address_of_callbacks_offset = tls_directory.AddressOfCallBacks - section.virtual_address - tls_directory.AddressOfCallBacks = ( - self.pe.patched_sections[section.name].virtual_address + address_of_callbacks_offset - ) + if section := self.pe.section(va=tls_directory.AddressOfCallBacks - image_base): + address_of_callbacks_offset = tls_directory.AddressOfCallBacks - section.virtual_address + tls_directory.AddressOfCallBacks = ( + self.pe.patched_sections[section.name].virtual_address + address_of_callbacks_offset + ) # Patch the TLS AddressOfIndex - section = self.pe.section(va=tls_directory.AddressOfIndex - image_base) - address_of_index_offset = tls_directory.AddressOfIndex - self.pe.sections[section.name].virtual_address - tls_directory.AddressOfIndex = self.pe.sections[section.name].virtual_address + address_of_index_offset + if section := self.pe.section(va=tls_directory.AddressOfIndex - image_base): + address_of_index_offset = tls_directory.AddressOfIndex - self.pe.sections[section.name].virtual_address + tls_directory.AddressOfIndex = self.pe.sections[section.name].virtual_address + address_of_index_offset # Write the patched TLS directory to the new PE self.seek(directory_va) self.patched_pe.write(tls_directory.dumps()) - - def _get_tls_attribute_section(self, va: int) -> PESection | None: - """Function to get the section that contains the TLS attribute. - - Args: - va: The virtual address of the TLS attribute. - - Returns: - The section that contains the TLS attribute as a `PESection` object. - """ - - for section in self.pe.sections.values(): - if va in range(section.virtual_address, section.virtual_address + section.virtual_size): - return section - return None From 8ed8611901ab355e7635631a0df456a438bbdac7 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Thu, 13 Mar 2025 10:05:38 +0000 Subject: [PATCH 18/43] Change build from a property to a function in Patcher --- dissect/executable/pe/helpers/patcher.py | 6 +++- tests/test_pe_builder.py | 44 +++++------------------- tests/test_pe_modifications.py | 12 +++---- 3 files changed, 20 insertions(+), 42 deletions(-) diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index 47a68ac..df0c69a 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -23,7 +23,6 @@ def __init__(self, pe: PE): self.patched_pe = BytesIO() self.functions = [] - @property def build(self) -> BytesIO: """Build the patched PE file. @@ -132,6 +131,8 @@ def _patch_import_rvas(self) -> None: # Get the original VA of the section the import directory is residing in, this value is used to calculate the # new RVA's section = self.pe.patched_section(va=directory_va) + if section is None: + return directory_offset = directory_va - section.virtual_address original_directory_va = self.pe.sections[section.name].virtual_address + directory_offset @@ -203,6 +204,9 @@ def _patch_export_rvas(self) -> None: # Get the original VA of the section the import directory is residing in, this value is used to calculate the # new RVA's section = self.pe.patched_section(va=directory_va) + if section is None: + return + directory_offset = directory_va - section.virtual_address original_directory_va = self.pe.sections[section.name].virtual_address + directory_offset diff --git a/tests/test_pe_builder.py b/tests/test_pe_builder.py index 9c71ac6..45cb44f 100644 --- a/tests/test_pe_builder.py +++ b/tests/test_pe_builder.py @@ -18,15 +18,9 @@ def test_build_new_x86_pe_exe() -> None: pe.pe_file.seek(len(pe.mz_header)) stub = pe.pe_file.read(pe.mz_header.e_lfanew - len(pe.mz_header)) - assert ( - stub[14 : 73 - 4] == b"This program is made with dissect.pe <3 kusjesvanSRT <3" - ) + assert stub[14 : 73 - 4] == b"This program is made with dissect.pe <3 kusjesvanSRT <3" - assert ( - pe.file_header.Characteristics - & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE - == 0x0100 - ) + assert pe.file_header.Characteristics & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE def test_build_new_x64_pe_exe() -> None: @@ -36,15 +30,9 @@ def test_build_new_x64_pe_exe() -> None: pe.pe_file.seek(len(pe.mz_header)) stub = pe.pe_file.read(pe.mz_header.e_lfanew - len(pe.mz_header)) - assert ( - stub[14 : 73 - 4] == b"This program is made with dissect.pe <3 kusjesvanSRT <3" - ) + assert stub[14 : 73 - 4] == b"This program is made with dissect.pe <3 kusjesvanSRT <3" - assert ( - pe.file_header.Characteristics - & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE - != 0x0100 - ) + assert not (pe.file_header.Characteristics & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE) def test_build_new_x86_pe_dll() -> None: @@ -52,15 +40,8 @@ def test_build_new_x86_pe_dll() -> None: builder.new() pe = builder.pe - assert ( - pe.file_header.Characteristics - & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE - == 0x0100 - ) - assert ( - pe.file_header.Characteristics & c_pe.ImageCharacteristics.IMAGE_FILE_DLL - == 0x2000 - ) + assert pe.file_header.Characteristics & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE + assert pe.file_header.Characteristics & c_pe.ImageCharacteristics.IMAGE_FILE_DLL def test_build_new_x64_pe_dll() -> None: @@ -68,15 +49,8 @@ def test_build_new_x64_pe_dll() -> None: builder.new() pe = builder.pe - assert ( - pe.file_header.Characteristics - & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE - != 0x0100 - ) - assert ( - pe.file_header.Characteristics & c_pe.ImageCharacteristics.IMAGE_FILE_DLL - == 0x2000 - ) + assert not (pe.file_header.Characteristics & c_pe.ImageCharacteristics.IMAGE_FILE_32BIT_MACHINE) + assert pe.file_header.Characteristics & c_pe.ImageCharacteristics.IMAGE_FILE_DLL def test_build_new_pe_with_custom_section() -> None: @@ -88,7 +62,7 @@ def test_build_new_pe_with_custom_section() -> None: patcher = Patcher(pe=pe) - new_pe = PE(pe_file=patcher.build) + new_pe = PE(pe_file=patcher.build()) assert new_pe.sections[".SRT"].name == ".SRT" assert new_pe.sections[".SRT"].size == 12 diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py index d1ae2a1..1a0ecaa 100644 --- a/tests/test_pe_modifications.py +++ b/tests/test_pe_modifications.py @@ -14,7 +14,7 @@ def test_add_imports() -> None: pe.import_mgr.add(dllname=dllname, functions=functions) patcher = Patcher(pe=pe) - new_pe = PE(pe_file=patcher.build) + new_pe = PE(pe_file=patcher.build()) assert "kusjesvanSRT.dll" in new_pe.imports @@ -30,7 +30,7 @@ def test_resize_section_smaller() -> None: pe.sections[".text"].data = b"kusjesvanSRT, patched with dissect" patcher = Patcher(pe=pe) - new_pe = PE(pe_file=patcher.build) + new_pe = PE(pe_file=patcher.build()) assert new_pe.sections[".text"].size == len(b"kusjesvanSRT, patched with dissect") assert ( @@ -48,7 +48,7 @@ def test_resize_section_bigger() -> None: pe.patched_sections[".rdata"].data += b"kusjesvanSRT, patched with dissect" * 100 patcher = Patcher(pe=pe) - new_pe = PE(pe_file=patcher.build) + new_pe = PE(pe_file=patcher.build()) assert new_pe.sections[".rdata"].size == original_size + len(b"kusjesvanSRT, patched with dissect" * 100) @@ -61,7 +61,7 @@ def test_resize_resource_smaller() -> None: e.data = b"kusjesvanSRT, patched with dissect" patcher = Patcher(pe=pe) - new_pe = PE(pe_file=patcher.build) + new_pe = PE(pe_file=patcher.build()) assert [patched.data for patched in new_pe.get_resource_type(rsrc_id="Manifest")] == [ b"kusjesvanSRT, patched with dissect" @@ -76,7 +76,7 @@ def test_resize_resource_bigger() -> None: e.data = b"kusjesvanSRT, patched with dissect" + e.data patcher = Patcher(pe=pe) - new_pe = PE(pe_file=patcher.build) + new_pe = PE(pe_file=patcher.build()) assert [ patched.data[: len(b"kusjesvanSRT, patched with dissect")] @@ -90,7 +90,7 @@ def test_add_section() -> None: pe.add_section(name=".SRT", data=b"kusjesvanSRT") patcher = Patcher(pe=pe) - new_pe = PE(pe_file=patcher.build) + new_pe = PE(pe_file=patcher.build()) assert ".SRT" in new_pe.sections assert new_pe.sections[".SRT"].data == b"kusjesvanSRT" From 6051423237b2a3f596bd5f7ca1c3937b3e018e16 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Thu, 13 Mar 2025 10:08:06 +0000 Subject: [PATCH 19/43] Use specific c_pe types for certain options Additinally remove a redundant definition --- dissect/executable/pe/helpers/builder.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dissect/executable/pe/helpers/builder.py b/dissect/executable/pe/helpers/builder.py index b665498..d436e65 100644 --- a/dissect/executable/pe/helpers/builder.py +++ b/dissect/executable/pe/helpers/builder.py @@ -56,7 +56,7 @@ def new(self) -> None: # Add a dummy section header to the new PE, we need at least 1 section to parse the PE dummy_data = b"<3kusjesvanSRT<3" - dummy_multiplier = 0x400 // len(b"<3kusjesvanSRT<3") + dummy_multiplier = 0x400 // len(dummy_data) section_header_offset = self.optional_header.SizeOfHeaders pointer_to_raw_data = utils.align_int( @@ -234,13 +234,13 @@ def gen_file_header( time_date_stamp = int(datetime.now(tz=timezone.utc).timestamp()) file_header = c_pe.IMAGE_FILE_HEADER() - file_header.Machine = machine + file_header.Machine = c_pe.MachineType(machine) file_header.NumberOfSections = number_of_sections file_header.TimeDateStamp = time_date_stamp file_header.PointerToSymbolTable = pointer_to_symbol_table file_header.NumberOfSymbols = number_of_symbols file_header.SizeOfOptionalHeader = size_of_optional_header - file_header.Characteristics = characteristics + file_header.Characteristics = c_pe.ImageCharacteristics(characteristics) return file_header @@ -429,7 +429,7 @@ def section( section_header.PointerToLinenumbers = pointer_to_linenumbers section_header.NumberOfRelocations = number_of_relocations section_header.NumberOfLinenumbers = number_of_linenumbers - section_header.Characteristics = characteristics + section_header.Characteristics = c_pe.SectionFlags(characteristics) return section_header From 546476d31c22b75d7bd7fec57ff4336dcfbda0b8 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Thu, 13 Mar 2025 10:15:25 +0000 Subject: [PATCH 20/43] Retrieve resources using the resource manager --- dissect/executable/pe/helpers/resources.py | 13 ++++--------- dissect/executable/pe/pe.py | 17 ----------------- tests/test_pe_modifications.py | 8 ++++---- 3 files changed, 8 insertions(+), 30 deletions(-) diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index 9d50ff0..f663c33 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -14,9 +14,6 @@ from collections.abc import Iterator from typing import BinaryIO - from dissect.cstruct.cstruct import cstruct - from dissect.cstruct.types.enum import EnumInstance - from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE @@ -24,7 +21,7 @@ @dataclass class RawResource: offset: int - entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY | c_pe.IMAGE_RESOURCE_DIRECTORY + entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY | c_pe.IMAGE_RESOURCE_DIRECTORY | c_pe.IMAGE_RESOURCE_DATA_ENTRY data_offset: int data: bytes | None = None resource: Resource | None = None @@ -142,8 +139,6 @@ def _read_resource(self, data: BinaryIO, offset: int, level: int = 1) -> Ordered entries = self._read_entries(data, directory) - _rc_type: str = "" - for entry in entries: _rc_type = self._rc_type(entry, data, level) @@ -169,7 +164,7 @@ def _rc_type(self, entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY, data: BinaryIO, l return str(entry.Id) - def get_resource(self, name: str) -> Resource: + def by_name(self, name: str) -> Resource: """Retrieve the resource by name. Args: @@ -184,7 +179,7 @@ def get_resource(self, name: str) -> Resource: except KeyError: raise ResourceException(f"Resource {name} not found!") - def get_resource_type(self, rsrc_id: str | EnumInstance) -> Iterator[Resource]: + def by_type(self, rsrc_id: str | c_pe.ResourceID) -> Iterator[Resource]: """Yields a generator containing all of the nodes within the resources that contain the requested ID. The ID can be either given by name or its value. @@ -316,7 +311,7 @@ def __init__( section: PESection, name: str | int, entry_offset: int, - data_entry: cstruct, + data_entry: c_pe.IMAGE_RESOURCE_DATA_ENTRY, rc_type: str, data: bytes = b"", ): diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index 2329af6..f6d7e01 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -257,23 +257,6 @@ def parse_directories(self) -> None: self.tls_mgr = tls.TLSManager(pe=self, section=section) self.tls_callbacks = self.tls_mgr.callbacks - def get_resource_type(self, rsrc_id: str | Enum) -> Iterator[Resource]: - """Yields a generator containing all of the nodes within the resources that contain the requested ID. - - The ID can be either given by name or its value. - - Args: - rsrc_id: The resource ID to find, this can be a cstruct `EnumInstance` or `str`. - - Yields: - All of the nodes that contain the requested type. - """ - - if rsrc_id not in self.resources: - raise ResourceException(f"Resource with ID {rsrc_id} not found in PE!") - - yield from self.rsrc_mgr.parse_resources(resources=self.resources[rsrc_id]) - def virtual_address(self, address: int) -> int: """Return the virtual address given a (possible) physical address. diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py index 1a0ecaa..9ada1a6 100644 --- a/tests/test_pe_modifications.py +++ b/tests/test_pe_modifications.py @@ -57,13 +57,13 @@ def test_resize_resource_smaller() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - for e in pe.get_resource_type(rsrc_id="Manifest"): + for e in pe.rsrc_mgr.by_type(rsrc_id="Manifest"): e.data = b"kusjesvanSRT, patched with dissect" patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) - assert [patched.data for patched in new_pe.get_resource_type(rsrc_id="Manifest")] == [ + assert [patched.data for patched in new_pe.rsrc_mgr.by_type(rsrc_id="Manifest")] == [ b"kusjesvanSRT, patched with dissect" ] @@ -72,7 +72,7 @@ def test_resize_resource_bigger() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - for e in pe.get_resource_type(rsrc_id="Manifest"): + for e in pe.rsrc_mgr.by_type(rsrc_id="Manifest"): e.data = b"kusjesvanSRT, patched with dissect" + e.data patcher = Patcher(pe=pe) @@ -80,7 +80,7 @@ def test_resize_resource_bigger() -> None: assert [ patched.data[: len(b"kusjesvanSRT, patched with dissect")] - for patched in new_pe.get_resource_type(rsrc_id="Manifest") + for patched in new_pe.rsrc_mgr.by_type(rsrc_id="Manifest") ] == [b"kusjesvanSRT, patched with dissect"] From f97ad4a7f1dae3aeee1316a50f70888c37dfbebe Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Wed, 26 Mar 2025 10:33:30 +0000 Subject: [PATCH 21/43] Read section from pe file I had the wrong assumption that the directory data is the same as the data inside the section. This has been reverted for now --- .gitignore | 2 ++ dissect/executable/pe/helpers/relocations.py | 2 +- dissect/executable/pe/helpers/resources.py | 2 +- dissect/executable/pe/helpers/tls.py | 2 +- dissect/executable/pe/helpers/utils.py | 20 +++++++++++++++++++- 5 files changed, 24 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index 74cecaf..5525039 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,5 @@ __pycache__/ tests/docs/api tests/docs/build .tox/ + +*.pyi \ No newline at end of file diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/helpers/relocations.py index ab719d1..eeb5cc8 100644 --- a/dissect/executable/pe/helpers/relocations.py +++ b/dissect/executable/pe/helpers/relocations.py @@ -28,7 +28,7 @@ def __init__(self, pe: PE, section: PESection): def parse_relocations(self) -> None: """Parse the relocation table of the PE file.""" - reloc_data = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC)) + reloc_data = BytesIO(self.pe.read_image_directory(c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC)) reloc_data_size = reloc_data.getbuffer().nbytes while reloc_data.tell() < reloc_data_size: reloc_directory = c_pe.IMAGE_BASE_RELOCATION(reloc_data) diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index f663c33..3eb7754 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -44,7 +44,7 @@ def __init__(self, pe: PE, section: PESection): def parse(self) -> None: """Parse the resource directory entry of the PE file.""" - rsrc_data = BytesIO(self.section.data) + rsrc_data = BytesIO(self.pe.read_image_directory(c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE)) self.resources = self._read_resource(data=rsrc_data, offset=0, level=1) def _read_entries( diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/helpers/tls.py index 50f68a1..3a90f70 100644 --- a/dissect/executable/pe/helpers/tls.py +++ b/dissect/executable/pe/helpers/tls.py @@ -42,7 +42,7 @@ def set_architecture(self, pe: PE) -> None: def parse_tls(self) -> None: """Parse the TLS directory entry of the PE file when present.""" - tls_data = BytesIO(self.section.data) + tls_data = BytesIO(self.pe.read_image_directory(c_pe.IMAGE_DIRECTORY_ENTRY_TLS)) self.tls = self._tls_directory(tls_data) self.pe.seek(self.tls.AddressOfCallBacks - self._image_base) diff --git a/dissect/executable/pe/helpers/utils.py b/dissect/executable/pe/helpers/utils.py index c557802..6277e7f 100644 --- a/dissect/executable/pe/helpers/utils.py +++ b/dissect/executable/pe/helpers/utils.py @@ -2,7 +2,7 @@ import struct from functools import lru_cache -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Generic, TypeVar if TYPE_CHECKING: from dissect.executable.pe import PE, PESection @@ -53,3 +53,21 @@ def pad(size: int) -> bytes: The null bytes as `bytes`. """ return size * b"\x00" + + +class Manager: + def __init__(self, pe: PE, section: PESection) -> None: + self.pe = pe + self.section = section + + def parse(self) -> None: + raise NotImplementedError + + def add(self, *args, **kwargs) -> None: + raise NotImplementedError + + def delete(self, *args, **kwargs) -> None: + raise NotImplementedError + + def patch(self, *args, **kwargs) -> None: + raise NotImplementedError From d0ef4d96c02d1dfcc731e889e7948926a8f941f9 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Wed, 26 Mar 2025 12:59:22 +0000 Subject: [PATCH 22/43] Grab data from the section, instead of reading it from the PE file The data inside the section is already available, so it made sense to make it possible to get the information from there --- dissect/executable/pe/helpers/exports.py | 2 +- dissect/executable/pe/helpers/imports.py | 3 +-- dissect/executable/pe/helpers/relocations.py | 2 +- dissect/executable/pe/helpers/resources.py | 2 +- dissect/executable/pe/helpers/sections.py | 20 +++++++++++++------- dissect/executable/pe/helpers/tls.py | 2 +- dissect/executable/pe/pe.py | 6 ++++-- 7 files changed, 22 insertions(+), 15 deletions(-) diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py index a6d124b..9df479b 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/helpers/exports.py @@ -49,7 +49,7 @@ def parse_exports(self) -> None: """ export_entry_va = self.pe.directory_entry_rva(c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT) - export_entry = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT)) + export_entry = BytesIO(self.section.directory_data(index=c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT)) export_directory = c_pe.IMAGE_EXPORT_DIRECTORY(export_entry) # Seek to the offset of the export name diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index ecdde24..a5782a3 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -150,8 +150,7 @@ def parse_imports(self) -> None: The imports are in turn added to the `imports` attribute so they can be accessed by the user. """ - import_data = BytesIO(self.pe.read_image_directory(index=c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT)) - import_data.seek(0) + import_data = BytesIO(self.section.directory_data(index=c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT)) # Loop over the entries for descriptor_va, import_descriptor in self.import_descriptors(import_data=import_data): diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/helpers/relocations.py index eeb5cc8..f9104d2 100644 --- a/dissect/executable/pe/helpers/relocations.py +++ b/dissect/executable/pe/helpers/relocations.py @@ -28,7 +28,7 @@ def __init__(self, pe: PE, section: PESection): def parse_relocations(self) -> None: """Parse the relocation table of the PE file.""" - reloc_data = BytesIO(self.pe.read_image_directory(c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC)) + reloc_data = BytesIO(self.section.directory_data(c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC)) reloc_data_size = reloc_data.getbuffer().nbytes while reloc_data.tell() < reloc_data_size: reloc_directory = c_pe.IMAGE_BASE_RELOCATION(reloc_data) diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index 3eb7754..b73c5a5 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -44,7 +44,7 @@ def __init__(self, pe: PE, section: PESection): def parse(self) -> None: """Parse the resource directory entry of the PE file.""" - rsrc_data = BytesIO(self.pe.read_image_directory(c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE)) + rsrc_data = BytesIO(self.section.directory_data(c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE)) self.resources = self._read_resource(data=rsrc_data, offset=0, level=1) def _read_entries( diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 69bfa84..7bc273e 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -34,10 +34,17 @@ def __init__(self, pe: PE, section: cstruct, offset: int, data: bytes = b""): self._size_of_raw_data = section.SizeOfRawData # Keep track of the directories that are within this section - self.directories = OrderedDict() + self.directories: OrderedDict[int, tuple[int, int]] = OrderedDict() self._data = data or self.read_data() + def directory_data(self, index: int) -> bytes: + if (dir_information := self.directories.get(index)) is None: + raise ValueError("Directory not found in PE Section") + + offset, size = dir_information + return self.data[offset : offset + size] + def read_data(self) -> bytes: """Return the data within the section. @@ -89,9 +96,8 @@ def virtual_address(self, value: int) -> None: self.section.VirtualAddress = value # Update the VA of the directory residing within this section - for idx, offset in self.directories.items(): - directory_va = value + offset - self.pe.optional_header.DataDirectory[idx].VirtualAddress = directory_va + for idx, (offset, _) in self.directories.items(): + self.pe.optional_header.DataDirectory[idx].VirtualAddress = value + offset @property def virtual_size(self) -> int: @@ -181,7 +187,7 @@ def data(self, value: bytes) -> None: prev_va = first_section.virtual_address prev_vsize = first_section.virtual_size - for name, section in self.pe.patched_sections.items(): + for section in self.pe.patched_sections.values(): if section.virtual_address == prev_va: continue @@ -193,8 +199,8 @@ def data(self, value: bytes) -> None: section virtual address is lower than the previous section. We want to prevent messing up RVA's as much as possible, this could lead to binaries that are a bit larger than they need to be but that doesn't really matter.""" - self.pe.patched_sections[name].virtual_address = virtual_address - self.pe.patched_sections[name].pointer_to_raw_data = pointer_to_raw_data + section.virtual_address = virtual_address + section.pointer_to_raw_data = pointer_to_raw_data prev_ptr = pointer_to_raw_data prev_size = section.size_of_raw_data diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/helpers/tls.py index 3a90f70..b9f4f49 100644 --- a/dissect/executable/pe/helpers/tls.py +++ b/dissect/executable/pe/helpers/tls.py @@ -42,7 +42,7 @@ def set_architecture(self, pe: PE) -> None: def parse_tls(self) -> None: """Parse the TLS directory entry of the PE file when present.""" - tls_data = BytesIO(self.pe.read_image_directory(c_pe.IMAGE_DIRECTORY_ENTRY_TLS)) + tls_data = BytesIO(self.section.directory_data(c_pe.IMAGE_DIRECTORY_ENTRY_TLS)) self.tls = self._tls_directory(tls_data) self.pe.seek(self.tls.AddressOfCallBacks - self._image_base) diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index f6d7e01..13d0190 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -229,9 +229,11 @@ def parse_directories(self) -> None: # Take note of the current directory VA so we can dynamically update it when resizing sections section = self.datadirectory_section(index=idx) - directory_va_offset = self.optional_header.DataDirectory[idx].VirtualAddress - section.virtual_address - section.directories[idx] = directory_va_offset + section_dir = self.optional_header.DataDirectory[idx] + + directory_va_offset = section_dir.VirtualAddress - section.virtual_address + section.directories[idx] = (directory_va_offset, section_dir.Size) # Parse the Import Address Table (IAT) if idx == c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT: self.import_mgr = imports.ImportManager(pe=self, section=section) From 8197638a1a2da71998091fbc01797b67cf1082b5 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Wed, 26 Mar 2025 14:18:50 +0000 Subject: [PATCH 23/43] Move relocations to a dataclass --- dissect/executable/pe/helpers/relocations.py | 22 +++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/helpers/relocations.py index f9104d2..9964556 100644 --- a/dissect/executable/pe/helpers/relocations.py +++ b/dissect/executable/pe/helpers/relocations.py @@ -1,5 +1,6 @@ from __future__ import annotations +from dataclasses import dataclass from io import BytesIO from typing import TYPE_CHECKING @@ -10,6 +11,13 @@ from dissect.executable.pe.pe import PE +@dataclass +class Relocation: + rva: int + number_of_entries: int + entries: list[int] + + class RelocationManager: """Base class for dealing with the relocations within the PE file. @@ -21,7 +29,7 @@ class RelocationManager: def __init__(self, pe: PE, section: PESection): self.pe = pe self.section = section - self.relocations: list[dict] = [] + self.relocations: list[Relocation] = [] self.parse_relocations() @@ -31,7 +39,7 @@ def parse_relocations(self) -> None: reloc_data = BytesIO(self.section.directory_data(c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC)) reloc_data_size = reloc_data.getbuffer().nbytes while reloc_data.tell() < reloc_data_size: - reloc_directory = c_pe.IMAGE_BASE_RELOCATION(reloc_data) + reloc_directory = c_pe._IMAGE_BASE_RELOCATION(reloc_data) if not reloc_directory.VirtualAddress: # End of relocation entries break @@ -41,11 +49,11 @@ def parse_relocations(self) -> None: entries = [entry for _ in range(number_of_entries) if (entry := c_pe.uint16(reloc_data))] self.relocations.append( - { - "rva:": reloc_directory.VirtualAddress, - "number_of_entries": number_of_entries, - "entries": entries, - } + Relocation( + rva=reloc_directory.VirtualAddress, + number_of_entries=number_of_entries, + entries=entries, + ) ) def add(self) -> None: From 3498bde05049cbb1bd46252113a80745638f480a Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Wed, 26 Mar 2025 14:39:56 +0000 Subject: [PATCH 24/43] Improve typing --- dissect/executable/pe/helpers/exports.py | 10 +++++----- dissect/executable/pe/helpers/imports.py | 10 +++------- dissect/executable/pe/helpers/sections.py | 6 ++---- dissect/executable/pe/pe.py | 8 +++----- 4 files changed, 13 insertions(+), 21 deletions(-) diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py index 9df479b..800d605 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/helpers/exports.py @@ -24,7 +24,7 @@ class ExportFunction: ordinal: int address: int - name: bytes | None = b"" + name: bytes = b"" def __str__(self) -> str: return self.name.decode() if self.name else f"#{self.ordinal}" @@ -58,20 +58,20 @@ def parse_exports(self) -> None: # Create a list of adresses for the exported functions export_entry.seek(export_directory.AddressOfFunctions - export_entry_va) - export_addresses = c_pe.uint32[export_directory.NumberOfFunctions].read(export_entry) + export_addresses: list[int] = c_pe.uint32[export_directory.NumberOfFunctions].read(export_entry) # Create a list of addresses for the exported functions that have associated names export_entry.seek(export_directory.AddressOfNames - export_entry_va) - export_names = c_pe.uint32[export_directory.NumberOfNames].read(export_entry) + export_names: list[int] = c_pe.uint32[export_directory.NumberOfNames].read(export_entry) # Create a list of addresses for the ordinals associated with the functions export_entry.seek(export_directory.AddressOfNameOrdinals - export_entry_va) - export_ordinals = c_pe.uint16[export_directory.NumberOfNames].read(export_entry) + export_ordinals: list[int] = c_pe.uint16[export_directory.NumberOfNames].read(export_entry) # Iterate over the export functions and store the information export_entry.seek(export_directory.AddressOfFunctions - export_entry_va) for idx, address in enumerate(export_addresses): _idx = idx + 1 key = str(_idx) - export_name: str | None = None + export_name: bytes | None = None if idx in export_ordinals: entry_offset = export_names[export_ordinals.index(idx)] - export_entry_va diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index a5782a3..79033ba 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -11,8 +11,6 @@ if TYPE_CHECKING: from collections.abc import Iterator - from dissect.cstruct.cstruct import cstruct - from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE @@ -31,13 +29,13 @@ class ImportModule: def __init__( self, name: bytes, - import_descriptor: int | None | c_pe.IMAGE_IMPORT_DESCRIPTOR, + import_descriptor: c_pe.IMAGE_IMPORT_DESCRIPTOR, module_va: int, name_va: int, first_thunk: int, ): self.name = name - self.import_descriptor = import_descriptor or c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT + self.import_descriptor = import_descriptor self.module_va = module_va self.name_va = name_va self.first_thunk = first_thunk @@ -94,7 +92,6 @@ def name(self) -> str: # For the case thunkdata is not defined, such as during the `add` return "" - if not (entry := self.ordinal): self.pe.seek(self.data_address + 2) entry = c_pe.char[None](self.pe).decode() @@ -131,7 +128,7 @@ def __init__(self, pe: PE, section: PESection): self.thunks: list[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = [] self._thunk_data: type[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = None - self._high_bit: int = None + self._high_bit: int = 0 self.set_architecture(pe) self.parse_imports() @@ -149,7 +146,6 @@ def parse_imports(self) -> None: The imports are in turn added to the `imports` attribute so they can be accessed by the user. """ - import_data = BytesIO(self.section.directory_data(index=c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT)) # Loop over the entries diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 7bc273e..9ff431e 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -8,8 +8,6 @@ from dissect.executable.pe.helpers import utils if TYPE_CHECKING: - from dissect.cstruct.cstruct import cstruct - from dissect.executable.pe.pe import PE @@ -23,7 +21,7 @@ class PESection: data: The data that should be part of the section, this can be used to add new sections. """ - def __init__(self, pe: PE, section: cstruct, offset: int, data: bytes = b""): + def __init__(self, pe: PE, section: c_pe.IMAGE_SECTION_HEADER, offset: int, data: bytes = b""): self.pe = pe self.section = section self.offset = offset @@ -225,7 +223,7 @@ def build_section( pointer_to_raw_data: int, name: str | bytes = b".dissect", characteristics: int = 0xC0000040, -) -> cstruct: +) -> c_pe.IMAGE_SECTION_HEADER: """Build a new section for the PE. Args: diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index 13d0190..afacab5 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -11,7 +11,6 @@ InvalidArchitecture, InvalidPE, InvalidVA, - ResourceException, ) from dissect.executable.pe.c_pe import c_cv_info, c_pe from dissect.executable.pe.helpers import ( @@ -26,12 +25,10 @@ ) if TYPE_CHECKING: - from collections.abc import Iterable, Iterator + from collections.abc import Iterable from dissect.cstruct.cstruct import cstruct - from dissect.cstruct.types.enum import Enum - from dissect.executable.pe.helpers.resources import Resource class PE: @@ -234,6 +231,7 @@ def parse_directories(self) -> None: directory_va_offset = section_dir.VirtualAddress - section.virtual_address section.directories[idx] = (directory_va_offset, section_dir.Size) + # Parse the Import Address Table (IAT) if idx == c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT: self.import_mgr = imports.ImportManager(pe=self, section=section) @@ -418,7 +416,7 @@ def has_directory(self, index: int) -> bool: return self.optional_header.DataDirectory[index].Size != 0 - def debug(self) -> cstruct | None: + def debug(self) -> c_cv_info.CV_INFO_PDB70 | None: """Return the debug directory of the given PE file. Returns: From 72658ad394aff4238076abff661acf44c8ab8cd5 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 1 Apr 2025 09:55:42 +0000 Subject: [PATCH 25/43] Initial PESectionManager --- dissect/executable/pe/helpers/sections.py | 77 +++++++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 9ff431e..1e60c00 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections import OrderedDict +from copy import copy from typing import TYPE_CHECKING from dissect.executable.exception import BuildSectionException @@ -8,9 +9,85 @@ from dissect.executable.pe.helpers import utils if TYPE_CHECKING: + from collections.abc import Iterable + from dissect.executable.pe.pe import PE +class PESectionManager: + def __init__(self) -> None: + self._sections: OrderedDict[str, PESection] = OrderedDict() + self._patched_sections: OrderedDict[str, PESection] = OrderedDict() + + def add(self, name: str, section: PESection) -> None: + self._sections[name] = section + self._patched_sections[name] = PESection(section.pe, section.section, section.offset, copy(section.data)) + + @property + def last_section(self) -> PESection: + return self._sections[next(reversed(self._sections))] + + def get(self, va: int = 0, name: str = "", patch: bool = False) -> PESection | None: + sections = self.sections(patch) + + if name: + return sections.get(name) + + return self._in_virtual_range(va, sections.values()) + + def sections(self, patch: bool = False) -> OrderedDict[str, PESection]: + return self._patched_sections if patch else self._sections + + def in_range(self, va: int, patch: bool = False) -> PESection | None: + """Retrieve a section pof the PE file by virtual address. + + Args: + va: The virtual address to look for + patch: Whether it should look through the patched sections. + + Returns: + a `.PESection` corresponding to the virtual address + + """ + return self.get(va=va, patch=patch) + + def in_raw_range(self, offset: int, patch: bool = False) -> PESection | None: + sections = self.sections(patch) + + for section in sections.values(): + if section.pointer_to_raw_data <= offset < section.pointer_to_raw_data + section.size_of_raw_data: + return section + + return None + + def from_index(self, segment_index: int, patch: bool = False) -> PESection: + """Retrieve the section of the PE by index. + + Args: + segment_index: The segment to retrieve based on the order within the PE. + + TODO: Need to check whether this works for pdb stuff + + Returns: + A `PESection` corresponding to the segment_index. + """ + sections = self.sections(patch) + + sections_items = list(sections.items()) + + idx = 0 if segment_index - 1 == -1 else segment_index + section_name = sections_items[idx - 1][0] + + return sections[section_name] + + def _in_virtual_range(self, va: int, sections: Iterable[PESection]) -> PESection | None: + for section in sections: + if section.virtual_address <= va < section.virtual_address + section.virtual_size: + return section + + return None + + class PESection: """Base class for the PE sections that are present. From d8acc68acabc79a1a70cae6ca1f2571e2758396e Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Thu, 3 Apr 2025 11:18:07 +0000 Subject: [PATCH 26/43] Refer to all sections using the section manager --- dissect/executable/pe/helpers/builder.py | 4 +-- dissect/executable/pe/helpers/imports.py | 2 +- dissect/executable/pe/helpers/patcher.py | 18 ++++++------- dissect/executable/pe/helpers/sections.py | 16 +++++++----- dissect/executable/pe/pe.py | 32 ++++++++++------------- tests/test_pe_modifications.py | 6 ++--- 6 files changed, 38 insertions(+), 40 deletions(-) diff --git a/dissect/executable/pe/helpers/builder.py b/dissect/executable/pe/helpers/builder.py index d436e65..d189ffc 100644 --- a/dissect/executable/pe/helpers/builder.py +++ b/dissect/executable/pe/helpers/builder.py @@ -33,7 +33,7 @@ def __init__( self.dll = dll self.subsystem = subsystem - self.pe = None + self.pe: PE = None def new(self) -> None: """Build the PE file from scratch. @@ -444,7 +444,7 @@ def pe_size(self) -> int: The size of the PE. """ - last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] + last_section = self.pe.section_manager.last_section(patch=True) va = last_section.virtual_address size = last_section.virtual_size diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index 79033ba..dfcc589 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -218,7 +218,7 @@ def add(self, dllname: str, functions: list[str]) -> None: functions: A `list` of function names belonging to the module. """ - self.last_section = next(reversed(self.pe.patched_sections.values())) + self.last_section = self.pe.section_manager.last_section(patch=True) # Build a dummy import module _module = ImportModule( diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index df0c69a..65509a1 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -61,7 +61,7 @@ def pe_size(self) -> int: The new size of the PE as an `int`. """ - last_section = self.pe.patched_sections[next(reversed(self.pe.patched_sections))] + last_section = self.pe.section_manager.last_section() va = last_section.virtual_address size = last_section.virtual_size @@ -85,11 +85,11 @@ def _build_section_table(self) -> None: self.patched_pe.write(utils.pad(size=self.pe.section_header_offset - self.patched_pe.tell())) # Write the section headers - for section in self.pe.patched_sections.values(): + for section in self.pe.section_manager.sections(True).values(): self.patched_pe.write(section.dump()) # Add the data for each section - for section in self.pe.patched_sections.values(): + for section in self.pe.section_manager.sections(True).values(): self.patched_pe.seek(section.pointer_to_raw_data) self.patched_pe.write(section.data) @@ -234,7 +234,7 @@ def _patch_export_rvas(self) -> None: if not section: continue address_offset = address - section.virtual_address - new_address = self.pe.patched_sections[section.name].virtual_address + address_offset + new_address = self.pe.section_manager.get(name=section.name).virtual_address + address_offset new_function_rvas.append(new_address) rva_struct = utils.create_struct(" None: continue address_offset = name_address - section.virtual_address - new_address = self.pe.patched_sections[section.name].virtual_address + address_offset + new_address = self.pe.patched_section(name=section.name).virtual_address + address_offset new_name_rvas.append(new_address) for name_rva in new_name_rvas: @@ -315,7 +315,7 @@ def _patch_tls_rvas(self) -> None: if section := self.pe.section(va=tls_directory.StartAddressOfRawData - image_base): start_address_offset = tls_directory.StartAddressOfRawData - section.virtual_address tls_directory.StartAddressOfRawData = ( - self.pe.patched_sections[section.name].virtual_address + start_address_offset + self.pe.patched_section(name=section.name).virtual_address + start_address_offset ) end_address_offset = tls_directory.EndAddressOfRawData - tls_directory.StartAddressOfRawData tls_directory.EndAddressOfRawData = tls_directory.StartAddressOfRawData + end_address_offset @@ -324,13 +324,13 @@ def _patch_tls_rvas(self) -> None: if section := self.pe.section(va=tls_directory.AddressOfCallBacks - image_base): address_of_callbacks_offset = tls_directory.AddressOfCallBacks - section.virtual_address tls_directory.AddressOfCallBacks = ( - self.pe.patched_sections[section.name].virtual_address + address_of_callbacks_offset + self.pe.patched_section(name=section.name).virtual_address + address_of_callbacks_offset ) # Patch the TLS AddressOfIndex if section := self.pe.section(va=tls_directory.AddressOfIndex - image_base): - address_of_index_offset = tls_directory.AddressOfIndex - self.pe.sections[section.name].virtual_address - tls_directory.AddressOfIndex = self.pe.sections[section.name].virtual_address + address_of_index_offset + address_of_index_offset = tls_directory.AddressOfIndex - section.virtual_address + tls_directory.AddressOfIndex = section.virtual_address + address_of_index_offset # Write the patched TLS directory to the new PE self.seek(directory_va) diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 1e60c00..760a7ae 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -23,9 +23,9 @@ def add(self, name: str, section: PESection) -> None: self._sections[name] = section self._patched_sections[name] = PESection(section.pe, section.section, section.offset, copy(section.data)) - @property - def last_section(self) -> PESection: - return self._sections[next(reversed(self._sections))] + def last_section(self, *, patch: bool = False) -> PESection: + sections = self.sections(patch) + return sections[next(reversed(sections))] def get(self, va: int = 0, name: str = "", patch: bool = False) -> PESection | None: sections = self.sections(patch) @@ -243,8 +243,10 @@ def data(self, value: bytes) -> None: """ # Keep track of the section changes using the patched_sections dictionary - self.pe.patched_sections[self.name]._data = value - self.pe.patched_sections[self.name].size = len(value) + section_manager = self.pe.section_manager + patched_section: PESection = section_manager.get(name=self.name, patch=True) + patched_section._data = value + patched_section.size = len(value) # Set the new data and size self._data = value @@ -255,14 +257,14 @@ def data(self, value: bytes) -> None: self._data += utils.pad(size=self.virtual_size - self.size_of_raw_data) # Take note of the first section as our starting point - first_section = next(iter(self.pe.patched_sections.values())) + first_section = next(iter(section_manager.sections(patch=True).values())) prev_ptr = first_section.pointer_to_raw_data prev_size = first_section.size_of_raw_data prev_va = first_section.virtual_address prev_vsize = first_section.virtual_size - for section in self.pe.patched_sections.values(): + for section in section_manager.sections(patch=True).values(): if section.virtual_address == prev_va: continue diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index afacab5..bdebd38 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -56,6 +56,8 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.section_header_offset = 0 self.last_section_offset = 0 + self.section_manager = sections.PESectionManager() + self.sections: OrderedDict[str, sections.PESection] = OrderedDict() self.patched_sections: OrderedDict[str, sections.PESection] = OrderedDict() @@ -151,6 +153,7 @@ def parse_section_header(self) -> None: # Take note of the sections, keep track of any patches seperately self.sections[section_name] = sections.PESection(pe=self, section=section_header, offset=offset) self.patched_sections[section_name] = sections.PESection(pe=self, section=section_header, offset=offset) + self.section_manager.add(section_name, sections.PESection(pe=self, section=section_header, offset=offset)) self.last_section_offset = self.sections[next(reversed(self.sections))].offset @@ -172,10 +175,7 @@ def section(self, va: int = 0, name: str = "") -> sections.PESection | None: A `PESection` object. """ - if name: - return self.sections[name] - - return self._section_in_range(va, self.sections.values()) + return self.section_manager.get(va, name) def patched_section(self, va: int = 0, name: str = "") -> sections.PESection | None: """Function to retrieve a patched section based on the given virtual address or name. @@ -188,10 +188,7 @@ def patched_section(self, va: int = 0, name: str = "") -> sections.PESection | N A `PESection` object. """ - if name: - return self.patched_sections[name] - - return self._section_in_range(va, self.patched_sections.values()) + return self.section_manager.get(va=va, name=name, patch=True) def datadirectory_section(self, index: int) -> sections.PESection: """Return the section that contains the given virtual address. @@ -204,7 +201,7 @@ def datadirectory_section(self, index: int) -> sections.PESection: """ va = self.directory_entry_rva(index=index) - if section := self._section_in_range(va, self.patched_sections.values()): + if section := self.section_manager.in_range(va, patch=True): return section raise InvalidVA(f"VA not found in sections: {va:#04x}") @@ -270,7 +267,7 @@ def virtual_address(self, address: int) -> int: if self.virtual: return address - if section := self._section_in_range(address, self.patched_sections.values()): + if section := self.section_manager.in_range(address, patch=True): return section.pointer_to_raw_data + (address - section.virtual_address) raise InvalidVA(f"VA not found in sections: {address:#04x}") @@ -284,10 +281,8 @@ def raw_address(self, offset: int) -> int: Returns: The physical address as an `int`. """ - - for section in self.patched_sections.values(): - if section.pointer_to_raw_data <= offset < section.pointer_to_raw_data + section.size_of_raw_data: - return section.virtual_address + (offset - section.pointer_to_raw_data) + if section := self.section_manager.in_raw_range(offset, patch=True): + return section.virtual_address + (offset - section.pointer_to_raw_data) raise InvalidAddress(f"Raw address not found in sections: {offset:#04x}") @@ -375,7 +370,7 @@ def write(self, data: bytes) -> None: print(self.patched_sections) # Update the section data - if section := self._section_in_range(offset, self.patched_sections.values()): + if section := self.section_manager.in_range(offset, patch=True): self.seek(address=section.virtual_address) section.data = self.read(size=section.virtual_size) @@ -489,7 +484,7 @@ def add_section( """ # Take note of the last section - last_section = self.patched_sections[next(reversed(self.sections))] + last_section = self.section_manager.last_section(patch=True) # Calculate the new section size raw_size = utils.align_int(integer=len(data), blocksize=self.file_alignment) @@ -529,9 +524,10 @@ def add_section( # Add the new section to the PE self.sections[name] = sections.PESection(pe=self, section=new_section, offset=offset, data=data) self.patched_sections[name] = sections.PESection(pe=self, section=new_section, offset=offset, data=data) + self.section_manager.add(name, sections.PESection(pe=self, section=new_section, offset=offset, data=data)) # Update the SizeOfImage field - last_section = self.patched_sections[next(reversed(self.patched_sections))] + last_section = self.section_manager.last_section(patch=True) last_va = last_section.virtual_address last_size = last_section.virtual_size @@ -562,5 +558,5 @@ def write_pe(self, filename: str = "out.exe") -> None: """ pepatcher = patcher.Patcher(pe=self) - new_pe = pepatcher.build + new_pe = pepatcher.build() Path(filename).write_bytes(new_pe.read()) diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py index 9ada1a6..a2f7b87 100644 --- a/tests/test_pe_modifications.py +++ b/tests/test_pe_modifications.py @@ -27,14 +27,14 @@ def test_resize_section_smaller() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - pe.sections[".text"].data = b"kusjesvanSRT, patched with dissect" + pe.section_manager.get(name=".text").data = b"kusjesvanSRT, patched with dissect" patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) - assert new_pe.sections[".text"].size == len(b"kusjesvanSRT, patched with dissect") + assert new_pe.section_manager.get(name=".text").size == len(b"kusjesvanSRT, patched with dissect") assert ( - new_pe.sections[".text"].data[: len(b"kusjesvanSRT, patched with dissect")] + new_pe.section_manager.get(name=".text").data[: len(b"kusjesvanSRT, patched with dissect")] == b"kusjesvanSRT, patched with dissect" ) From 6a2677e4aa2798df800140624927e851c5b23a50 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Thu, 3 Apr 2025 12:02:57 +0000 Subject: [PATCH 27/43] Remove references to internal sections inside the PE file --- dissect/executable/pe/helpers/patcher.py | 38 +++++++-------- dissect/executable/pe/helpers/sections.py | 18 +++---- dissect/executable/pe/pe.py | 57 +---------------------- tests/test_pe.py | 2 +- tests/test_pe_builder.py | 9 ++-- tests/test_pe_modifications.py | 12 +++-- 6 files changed, 43 insertions(+), 93 deletions(-) diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index 65509a1..3a8719f 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -21,6 +21,7 @@ class Patcher: def __init__(self, pe: PE): self.pe = pe self.patched_pe = BytesIO() + self._section_manager = pe.section_manager self.functions = [] def build(self) -> BytesIO: @@ -85,11 +86,11 @@ def _build_section_table(self) -> None: self.patched_pe.write(utils.pad(size=self.pe.section_header_offset - self.patched_pe.tell())) # Write the section headers - for section in self.pe.section_manager.sections(True).values(): + for section in self._section_manager.sections(patch=True).values(): self.patched_pe.write(section.dump()) # Add the data for each section - for section in self.pe.section_manager.sections(True).values(): + for section in self._section_manager.sections(patch=True).values(): self.patched_pe.seek(section.pointer_to_raw_data) self.patched_pe.write(section.data) @@ -130,11 +131,11 @@ def _patch_import_rvas(self) -> None: # Get the original VA of the section the import directory is residing in, this value is used to calculate the # new RVA's - section = self.pe.patched_section(va=directory_va) + section = self._section_manager.get(va=directory_va, patch=True) if section is None: return directory_offset = directory_va - section.virtual_address - original_directory_va = self.pe.sections[section.name].virtual_address + directory_offset + original_directory_va = self._section_manager.get(name=section.name).virtual_address + directory_offset # Loop over the imports of the PE to patch the RVA's of the import descriptors and the associated thunkdata # entries @@ -165,9 +166,9 @@ def _patch_import_rvas(self) -> None: # Check the original RVA associated with the AddressOfData field in the thunkdata, retrieve the # original VA # and use it to also select the patched virtual address of this section that the RVA is located in - if section := self.pe.section(function.data_address): + if section := self._section_manager.get(va=function.data_address): virtual_address = section.virtual_address - new_virtual_address = self.pe.patched_section(name=section.name).virtual_address + new_virtual_address = self._section_manager.get(name=section.name, patch=True).virtual_address # Calculate the offset using the VA of the section and update the thunkdata va_offset = function.data_address - virtual_address @@ -203,12 +204,11 @@ def _patch_export_rvas(self) -> None: # Get the original VA of the section the import directory is residing in, this value is used to calculate the # new RVA's - section = self.pe.patched_section(va=directory_va) - if section is None: + if not (section := self._section_manager.get(va=directory_va, patch=True)): return directory_offset = directory_va - section.virtual_address - original_directory_va = self.pe.sections[section.name].virtual_address + directory_offset + original_directory_va = self._section_manager.get(name=section.name).virtual_address + directory_offset name_offset = export_directory.Name - original_directory_va address_of_functions_offset = export_directory.AddressOfFunctions - original_directory_va @@ -230,11 +230,10 @@ def _patch_export_rvas(self) -> None: self.seek(export_directory.AddressOfFunctions) export_addresses = c_pe.uint32[export_directory.NumberOfFunctions].read(self.patched_pe) for address in export_addresses: - section = self.pe.section(va=address) - if not section: + if not (section := self._section_manager.get(va=address)): continue address_offset = address - section.virtual_address - new_address = self.pe.section_manager.get(name=section.name).virtual_address + address_offset + new_address = self._section_manager.get(name=section.name, patch=True).virtual_address + address_offset new_function_rvas.append(new_address) rva_struct = utils.create_struct(" None: self.seek(export_directory.AddressOfNames) export_names = c_pe.uint32[export_directory.NumberOfNames].read(self.patched_pe) for name_address in export_names: - section = self.pe.section(va=name_address) + section = self._section_manager.get(va=name_address) if section is None: continue address_offset = name_address - section.virtual_address - new_address = self.pe.patched_section(name=section.name).virtual_address + address_offset + new_address = self._section_manager.get(name=section.name, patch=True).virtual_address + address_offset new_name_rvas.append(new_address) for name_rva in new_name_rvas: @@ -264,7 +263,6 @@ def _patch_export_rvas(self) -> None: self.seek(export_directory.AddressOfNames) self.patched_pe.write(name_rvas) - # self.pe.optional_header.DataDirectory[c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT].Size = len(name_rvas) def _patch_rsrc_rvas(self) -> None: """Function to patch the RVAs of the resource directory and the associated resource data RVA's.""" @@ -312,23 +310,23 @@ def _patch_tls_rvas(self) -> None: image_base = self.pe.optional_header.ImageBase # Patch the TLS StartAddressOfRawData and EndAddressOfRawData - if section := self.pe.section(va=tls_directory.StartAddressOfRawData - image_base): + if section := self._section_manager.get(va=tls_directory.StartAddressOfRawData - image_base): start_address_offset = tls_directory.StartAddressOfRawData - section.virtual_address tls_directory.StartAddressOfRawData = ( - self.pe.patched_section(name=section.name).virtual_address + start_address_offset + self._section_manager.get(name=section.name, patch=True).virtual_address + start_address_offset ) end_address_offset = tls_directory.EndAddressOfRawData - tls_directory.StartAddressOfRawData tls_directory.EndAddressOfRawData = tls_directory.StartAddressOfRawData + end_address_offset # Patch the TLS callbacks address - if section := self.pe.section(va=tls_directory.AddressOfCallBacks - image_base): + if section := self._section_manager.get(va=tls_directory.AddressOfCallBacks - image_base): address_of_callbacks_offset = tls_directory.AddressOfCallBacks - section.virtual_address tls_directory.AddressOfCallBacks = ( - self.pe.patched_section(name=section.name).virtual_address + address_of_callbacks_offset + self._section_manager.get(name=section.name, patch=True).virtual_address + address_of_callbacks_offset ) # Patch the TLS AddressOfIndex - if section := self.pe.section(va=tls_directory.AddressOfIndex - image_base): + if section := self._section_manager.get(va=tls_directory.AddressOfIndex - image_base, patch=True): address_of_index_offset = tls_directory.AddressOfIndex - section.virtual_address tls_directory.AddressOfIndex = section.virtual_address + address_of_index_offset diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 760a7ae..16a38b4 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -24,21 +24,21 @@ def add(self, name: str, section: PESection) -> None: self._patched_sections[name] = PESection(section.pe, section.section, section.offset, copy(section.data)) def last_section(self, *, patch: bool = False) -> PESection: - sections = self.sections(patch) + sections = self.sections(patch=patch) return sections[next(reversed(sections))] - def get(self, va: int = 0, name: str = "", patch: bool = False) -> PESection | None: - sections = self.sections(patch) + def get(self, va: int = 0, name: str = "", *, patch: bool = False) -> PESection | None: + sections = self.sections(patch=patch) if name: return sections.get(name) return self._in_virtual_range(va, sections.values()) - def sections(self, patch: bool = False) -> OrderedDict[str, PESection]: + def sections(self, *, patch: bool = False) -> OrderedDict[str, PESection]: return self._patched_sections if patch else self._sections - def in_range(self, va: int, patch: bool = False) -> PESection | None: + def in_range(self, va: int, *, patch: bool = False) -> PESection | None: """Retrieve a section pof the PE file by virtual address. Args: @@ -51,8 +51,8 @@ def in_range(self, va: int, patch: bool = False) -> PESection | None: """ return self.get(va=va, patch=patch) - def in_raw_range(self, offset: int, patch: bool = False) -> PESection | None: - sections = self.sections(patch) + def in_raw_range(self, offset: int, *, patch: bool = False) -> PESection | None: + sections = self.sections(patch=patch) for section in sections.values(): if section.pointer_to_raw_data <= offset < section.pointer_to_raw_data + section.size_of_raw_data: @@ -60,7 +60,7 @@ def in_raw_range(self, offset: int, patch: bool = False) -> PESection | None: return None - def from_index(self, segment_index: int, patch: bool = False) -> PESection: + def from_index(self, segment_index: int, *, patch: bool = False) -> PESection: """Retrieve the section of the PE by index. Args: @@ -71,7 +71,7 @@ def from_index(self, segment_index: int, patch: bool = False) -> PESection: Returns: A `PESection` corresponding to the segment_index. """ - sections = self.sections(patch) + sections = self.sections(patch=patch) sections_items = list(sections.items()) diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index bdebd38..55185de 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -30,7 +30,6 @@ from dissect.cstruct.cstruct import cstruct - class PE: """Base class for parsing PE files. @@ -58,9 +57,6 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.section_manager = sections.PESectionManager() - self.sections: OrderedDict[str, sections.PESection] = OrderedDict() - self.patched_sections: OrderedDict[str, sections.PESection] = OrderedDict() - self.imports: OrderedDict[str, imports.ImportModule] = None self.exports: OrderedDict[str, exports.ExportFunction] = None self.resources: OrderedDict[str, resources.Resource] = None @@ -79,7 +75,6 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.imagebase = self.nt_headers.OptionalHeader.ImageBase self.file_alignment = self.nt_headers.OptionalHeader.FileAlignment self.section_alignment = self.nt_headers.OptionalHeader.SectionAlignment - self.base_address = self.nt_headers.OptionalHeader.ImageBase self.timestamp = datetime.fromtimestamp(self.file_header.TimeDateStamp, tz=timezone.utc) # Parse the section header @@ -151,11 +146,9 @@ def parse_section_header(self) -> None: section_header = c_pe.IMAGE_SECTION_HEADER(self.pe_file) section_name = section_header.Name.decode().strip("\x00") # Take note of the sections, keep track of any patches seperately - self.sections[section_name] = sections.PESection(pe=self, section=section_header, offset=offset) - self.patched_sections[section_name] = sections.PESection(pe=self, section=section_header, offset=offset) self.section_manager.add(section_name, sections.PESection(pe=self, section=section_header, offset=offset)) - self.last_section_offset = self.sections[next(reversed(self.sections))].offset + self.last_section_offset = self.section_manager.last_section().offset def _section_in_range(self, address: int, values: Iterable[sections.PESection]) -> sections.PESection | None: for section in values: @@ -164,32 +157,6 @@ def _section_in_range(self, address: int, values: Iterable[sections.PESection]) return None - def section(self, va: int = 0, name: str = "") -> sections.PESection | None: - """Function to retrieve a section based on the given virtual address or name. - - Args: - va: The virtual address to look for within the sections. - name: The name of the section. - - Returns: - A `PESection` object. - """ - - return self.section_manager.get(va, name) - - def patched_section(self, va: int = 0, name: str = "") -> sections.PESection | None: - """Function to retrieve a patched section based on the given virtual address or name. - - Args: - va: The virtual address to look for within the patched sections. - name: The name of the patched section. - - Returns: - A `PESection` object. - """ - - return self.section_manager.get(va=va, name=name, patch=True) - def datadirectory_section(self, index: int) -> sections.PESection: """Return the section that contains the given virtual address. @@ -367,7 +334,6 @@ def write(self, data: bytes) -> None: # Write the data to the PE file so we can do a raw_read on this data in the section self.pe_file.write(data) - print(self.patched_sections) # Update the section data if section := self.section_manager.in_range(offset, patch=True): @@ -429,23 +395,6 @@ def debug(self) -> c_cv_info.CV_INFO_PDB70 | None: return c_cv_info.CV_INFO_PDB70(dbg_entry) return None - def get_section(self, segment_index: int) -> sections.PESection: - """Retrieve the section of the PE by index. - - Args: - segment_index: The segment to retrieve based on the order within the PE. - - Returns: - A `PESection` corresponding to the segment_index. - """ - - sections = list(self.sections.items()) - - idx = 0 if segment_index - 1 == -1 else segment_index - section_name = sections[idx - 1][0] - - return self.sections[section_name] - def symbol_data(self, symbol: cstruct, size: int) -> bytes: """Retrieve data from the PE using a PDB symbol. @@ -457,7 +406,7 @@ def symbol_data(self, symbol: cstruct, size: int) -> bytes: The bytes that were read from the offset within the PE. """ - _section = self.get_section(segment_index=symbol.seg) + _section = self.section_manager.from_index(segment_index=symbol.seg) address = self.imagebase + _section.virtual_address + symbol.off self.pe_file.seek(address) @@ -522,8 +471,6 @@ def add_section( self.optional_header.DataDirectory[datadirectory].Size = datadirectory_size or len(data) # Add the new section to the PE - self.sections[name] = sections.PESection(pe=self, section=new_section, offset=offset, data=data) - self.patched_sections[name] = sections.PESection(pe=self, section=new_section, offset=offset, data=data) self.section_manager.add(name, sections.PESection(pe=self, section=new_section, offset=offset, data=data)) # Update the SizeOfImage field diff --git a/tests/test_pe.py b/tests/test_pe.py index 3f83053..b683134 100644 --- a/tests/test_pe.py +++ b/tests/test_pe.py @@ -26,7 +26,7 @@ def test_pe_sections() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - assert known_sections == list(pe.sections) + assert known_sections == list(pe.section_manager.sections()) def test_pe_imports() -> None: diff --git a/tests/test_pe_builder.py b/tests/test_pe_builder.py index 45cb44f..932a057 100644 --- a/tests/test_pe_builder.py +++ b/tests/test_pe_builder.py @@ -64,6 +64,9 @@ def test_build_new_pe_with_custom_section() -> None: new_pe = PE(pe_file=patcher.build()) - assert new_pe.sections[".SRT"].name == ".SRT" - assert new_pe.sections[".SRT"].size == 12 - assert new_pe.sections[".SRT"].data == b"kusjesvanSRT" + section_manager = new_pe.section_manager + + section = section_manager.get(name=".SRT") + assert section.name == ".SRT" + assert section.size == 12 + assert section.data == b"kusjesvanSRT" diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py index a2f7b87..e044a83 100644 --- a/tests/test_pe_modifications.py +++ b/tests/test_pe_modifications.py @@ -43,14 +43,16 @@ def test_resize_section_bigger() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - original_size = pe.sections[".rdata"].size + original_size = pe.section_manager.get(name=".rdata").size - pe.patched_sections[".rdata"].data += b"kusjesvanSRT, patched with dissect" * 100 + pe.section_manager.get(name=".rdata", patch=True).data += b"kusjesvanSRT, patched with dissect" * 100 patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) - assert new_pe.sections[".rdata"].size == original_size + len(b"kusjesvanSRT, patched with dissect" * 100) + assert new_pe.section_manager.get(name=".rdata").size == original_size + len( + b"kusjesvanSRT, patched with dissect" * 100 + ) def test_resize_resource_smaller() -> None: @@ -92,5 +94,5 @@ def test_add_section() -> None: patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) - assert ".SRT" in new_pe.sections - assert new_pe.sections[".SRT"].data == b"kusjesvanSRT" + assert ".SRT" in new_pe.section_manager.sections() + assert new_pe.section_manager.get(name=".SRT").data == b"kusjesvanSRT" From 36a705428829d41ac54e64c34ca5cf56e6b4b747 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Thu, 3 Apr 2025 13:28:31 +0000 Subject: [PATCH 28/43] Rename pe.section_manager to pe.sections --- dissect/executable/pe/helpers/builder.py | 2 +- dissect/executable/pe/helpers/imports.py | 2 +- dissect/executable/pe/helpers/patcher.py | 4 ++-- dissect/executable/pe/helpers/sections.py | 2 +- dissect/executable/pe/pe.py | 22 +++++++++++----------- tests/test_pe.py | 2 +- tests/test_pe_builder.py | 2 +- tests/test_pe_modifications.py | 16 ++++++++-------- 8 files changed, 26 insertions(+), 26 deletions(-) diff --git a/dissect/executable/pe/helpers/builder.py b/dissect/executable/pe/helpers/builder.py index d189ffc..ace8105 100644 --- a/dissect/executable/pe/helpers/builder.py +++ b/dissect/executable/pe/helpers/builder.py @@ -444,7 +444,7 @@ def pe_size(self) -> int: The size of the PE. """ - last_section = self.pe.section_manager.last_section(patch=True) + last_section = self.pe.sections.last_section(patch=True) va = last_section.virtual_address size = last_section.virtual_size diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index dfcc589..ace18dc 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -218,7 +218,7 @@ def add(self, dllname: str, functions: list[str]) -> None: functions: A `list` of function names belonging to the module. """ - self.last_section = self.pe.section_manager.last_section(patch=True) + self.last_section = self.pe.sections.last_section(patch=True) # Build a dummy import module _module = ImportModule( diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index 3a8719f..ecad7b4 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -21,7 +21,7 @@ class Patcher: def __init__(self, pe: PE): self.pe = pe self.patched_pe = BytesIO() - self._section_manager = pe.section_manager + self._section_manager = pe.sections self.functions = [] def build(self) -> BytesIO: @@ -62,7 +62,7 @@ def pe_size(self) -> int: The new size of the PE as an `int`. """ - last_section = self.pe.section_manager.last_section() + last_section = self.pe.sections.last_section() va = last_section.virtual_address size = last_section.virtual_size diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 16a38b4..d1c3146 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -243,7 +243,7 @@ def data(self, value: bytes) -> None: """ # Keep track of the section changes using the patched_sections dictionary - section_manager = self.pe.section_manager + section_manager = self.pe.sections patched_section: PESection = section_manager.get(name=self.name, patch=True) patched_section._data = value patched_section.size = len(value) diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index 55185de..5d1f0eb 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -55,7 +55,7 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.section_header_offset = 0 self.last_section_offset = 0 - self.section_manager = sections.PESectionManager() + self.sections = sections.PESectionManager() self.imports: OrderedDict[str, imports.ImportModule] = None self.exports: OrderedDict[str, exports.ExportFunction] = None @@ -146,9 +146,9 @@ def parse_section_header(self) -> None: section_header = c_pe.IMAGE_SECTION_HEADER(self.pe_file) section_name = section_header.Name.decode().strip("\x00") # Take note of the sections, keep track of any patches seperately - self.section_manager.add(section_name, sections.PESection(pe=self, section=section_header, offset=offset)) + self.sections.add(section_name, sections.PESection(pe=self, section=section_header, offset=offset)) - self.last_section_offset = self.section_manager.last_section().offset + self.last_section_offset = self.sections.last_section().offset def _section_in_range(self, address: int, values: Iterable[sections.PESection]) -> sections.PESection | None: for section in values: @@ -168,7 +168,7 @@ def datadirectory_section(self, index: int) -> sections.PESection: """ va = self.directory_entry_rva(index=index) - if section := self.section_manager.in_range(va, patch=True): + if section := self.sections.in_range(va, patch=True): return section raise InvalidVA(f"VA not found in sections: {va:#04x}") @@ -234,7 +234,7 @@ def virtual_address(self, address: int) -> int: if self.virtual: return address - if section := self.section_manager.in_range(address, patch=True): + if section := self.sections.in_range(address, patch=True): return section.pointer_to_raw_data + (address - section.virtual_address) raise InvalidVA(f"VA not found in sections: {address:#04x}") @@ -248,7 +248,7 @@ def raw_address(self, offset: int) -> int: Returns: The physical address as an `int`. """ - if section := self.section_manager.in_raw_range(offset, patch=True): + if section := self.sections.in_raw_range(offset, patch=True): return section.virtual_address + (offset - section.pointer_to_raw_data) raise InvalidAddress(f"Raw address not found in sections: {offset:#04x}") @@ -336,7 +336,7 @@ def write(self, data: bytes) -> None: self.pe_file.write(data) # Update the section data - if section := self.section_manager.in_range(offset, patch=True): + if section := self.sections.in_range(offset, patch=True): self.seek(address=section.virtual_address) section.data = self.read(size=section.virtual_size) @@ -406,7 +406,7 @@ def symbol_data(self, symbol: cstruct, size: int) -> bytes: The bytes that were read from the offset within the PE. """ - _section = self.section_manager.from_index(segment_index=symbol.seg) + _section = self.sections.from_index(segment_index=symbol.seg) address = self.imagebase + _section.virtual_address + symbol.off self.pe_file.seek(address) @@ -433,7 +433,7 @@ def add_section( """ # Take note of the last section - last_section = self.section_manager.last_section(patch=True) + last_section = self.sections.last_section(patch=True) # Calculate the new section size raw_size = utils.align_int(integer=len(data), blocksize=self.file_alignment) @@ -471,10 +471,10 @@ def add_section( self.optional_header.DataDirectory[datadirectory].Size = datadirectory_size or len(data) # Add the new section to the PE - self.section_manager.add(name, sections.PESection(pe=self, section=new_section, offset=offset, data=data)) + self.sections.add(name, sections.PESection(pe=self, section=new_section, offset=offset, data=data)) # Update the SizeOfImage field - last_section = self.section_manager.last_section(patch=True) + last_section = self.sections.last_section(patch=True) last_va = last_section.virtual_address last_size = last_section.virtual_size diff --git a/tests/test_pe.py b/tests/test_pe.py index b683134..c524874 100644 --- a/tests/test_pe.py +++ b/tests/test_pe.py @@ -26,7 +26,7 @@ def test_pe_sections() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - assert known_sections == list(pe.section_manager.sections()) + assert known_sections == list(pe.sections.sections()) def test_pe_imports() -> None: diff --git a/tests/test_pe_builder.py b/tests/test_pe_builder.py index 932a057..6370024 100644 --- a/tests/test_pe_builder.py +++ b/tests/test_pe_builder.py @@ -64,7 +64,7 @@ def test_build_new_pe_with_custom_section() -> None: new_pe = PE(pe_file=patcher.build()) - section_manager = new_pe.section_manager + section_manager = new_pe.sections section = section_manager.get(name=".SRT") assert section.name == ".SRT" diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py index e044a83..a343b50 100644 --- a/tests/test_pe_modifications.py +++ b/tests/test_pe_modifications.py @@ -27,14 +27,14 @@ def test_resize_section_smaller() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - pe.section_manager.get(name=".text").data = b"kusjesvanSRT, patched with dissect" + pe.sections.get(name=".text").data = b"kusjesvanSRT, patched with dissect" patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) - assert new_pe.section_manager.get(name=".text").size == len(b"kusjesvanSRT, patched with dissect") + assert new_pe.sections.get(name=".text").size == len(b"kusjesvanSRT, patched with dissect") assert ( - new_pe.section_manager.get(name=".text").data[: len(b"kusjesvanSRT, patched with dissect")] + new_pe.sections.get(name=".text").data[: len(b"kusjesvanSRT, patched with dissect")] == b"kusjesvanSRT, patched with dissect" ) @@ -43,14 +43,14 @@ def test_resize_section_bigger() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - original_size = pe.section_manager.get(name=".rdata").size + original_size = pe.sections.get(name=".rdata").size - pe.section_manager.get(name=".rdata", patch=True).data += b"kusjesvanSRT, patched with dissect" * 100 + pe.sections.get(name=".rdata", patch=True).data += b"kusjesvanSRT, patched with dissect" * 100 patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) - assert new_pe.section_manager.get(name=".rdata").size == original_size + len( + assert new_pe.sections.get(name=".rdata").size == original_size + len( b"kusjesvanSRT, patched with dissect" * 100 ) @@ -94,5 +94,5 @@ def test_add_section() -> None: patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) - assert ".SRT" in new_pe.section_manager.sections() - assert new_pe.section_manager.get(name=".SRT").data == b"kusjesvanSRT" + assert ".SRT" in new_pe.sections.sections() + assert new_pe.sections.get(name=".SRT").data == b"kusjesvanSRT" From 719d47c6bce4cca47076471725155f5ad669cf86 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Mon, 7 Apr 2025 13:11:15 +0000 Subject: [PATCH 29/43] Use utils.Manager for most base classes --- dissect/executable/pe/helpers/exports.py | 7 +-- dissect/executable/pe/helpers/imports.py | 13 +++--- dissect/executable/pe/helpers/relocations.py | 13 +++--- dissect/executable/pe/helpers/resources.py | 45 ++++++++------------ dissect/executable/pe/helpers/tls.py | 23 +++++----- dissect/executable/pe/helpers/utils.py | 2 +- 6 files changed, 44 insertions(+), 59 deletions(-) diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py index 800d605..7799be6 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/helpers/exports.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING from dissect.executable.pe.c_pe import c_pe +from dissect.executable.pe.helpers.utils import Manager if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection @@ -33,15 +34,15 @@ def __repr__(self) -> str: return f"" -class ExportManager: +class ExportManager(Manager): def __init__(self, pe: PE, section: PESection): self.pe = pe self.section = section self.exports: OrderedDict[str, ExportFunction] = OrderedDict() - self.parse_exports() + self.parse() - def parse_exports(self) -> None: + def parse(self) -> None: """Parse the export directory of the PE file. This function will store every export function within the PE file as an `ExportFunction` object containing the diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index ace18dc..58f40e5 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -6,7 +6,7 @@ from dissect.executable.pe.c_pe import c_pe from dissect.executable.pe.helpers import utils -from dissect.executable.pe.helpers.utils import create_struct +from dissect.executable.pe.helpers.utils import Manager, create_struct if TYPE_CHECKING: from collections.abc import Iterator @@ -108,7 +108,7 @@ def __repr__(self) -> str: return f"" -class ImportManager: +class ImportManager(Manager): """The base class for dealing with the imports that are present within the PE file. Args: @@ -117,9 +117,8 @@ class ImportManager: """ def __init__(self, pe: PE, section: PESection): - self.pe = pe - self.image_size: int = self.pe.optional_header.SizeOfImage - self.section = section + super().__init__(pe, section) + self.image_size: int = pe.optional_header.SizeOfImage self.import_directory_rva = 0 self.import_data = bytearray() self.new_size_of_image = 0 @@ -131,7 +130,7 @@ def __init__(self, pe: PE, section: PESection): self._high_bit: int = 0 self.set_architecture(pe) - self.parse_imports() + self.parse() def set_architecture(self, pe: PE) -> None: if pe.is64bit(): @@ -141,7 +140,7 @@ def set_architecture(self, pe: PE) -> None: self._thunk_data = c_pe.IMAGE_THUNK_DATA32 self._high_bit = 1 << 31 - def parse_imports(self) -> None: + def parse(self) -> None: """Parse the imports of the PE file. The imports are in turn added to the `imports` attribute so they can be accessed by the user. diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/helpers/relocations.py index 9964556..d0fb7c2 100644 --- a/dissect/executable/pe/helpers/relocations.py +++ b/dissect/executable/pe/helpers/relocations.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING from dissect.executable.pe.c_pe import c_pe +from dissect.executable.pe.helpers.utils import Manager if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection @@ -18,7 +19,7 @@ class Relocation: entries: list[int] -class RelocationManager: +class RelocationManager(Manager): """Base class for dealing with the relocations within the PE file. Args: @@ -27,13 +28,12 @@ class RelocationManager: """ def __init__(self, pe: PE, section: PESection): - self.pe = pe - self.section = section + super().__init__(pe, section) self.relocations: list[Relocation] = [] - self.parse_relocations() + self.parse() - def parse_relocations(self) -> None: + def parse(self) -> None: """Parse the relocation table of the PE file.""" reloc_data = BytesIO(self.section.directory_data(c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC)) @@ -55,6 +55,3 @@ def parse_relocations(self) -> None: entries=entries, ) ) - - def add(self) -> None: - raise NotImplementedError diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index b73c5a5..f0ed589 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -9,6 +9,7 @@ from dissect.executable.exception import ResourceException from dissect.executable.pe.c_pe import c_pe +from dissect.executable.pe.helpers.utils import Manager if TYPE_CHECKING: from collections.abc import Iterator @@ -26,6 +27,21 @@ class RawResource: data: bytes | None = None resource: Resource | None = None + +def rc_type(entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY, data: BinaryIO, depth: int = 1) -> str: + """Returns the name of the rc type depending on the data and the depth level of the resource""" + if depth == 1: + return c_pe.ResourceID(entry.Id).name + + if entry.NameIsString: + data.seek(entry.NameOffset) + name_len = c_pe.uint16(data) + return c_pe.wchar[name_len](data) + + return str(entry.Id) + + +class ResourceManager(Manager): """Base class to perform actions regarding the resources within the PE file. Args: @@ -140,7 +156,7 @@ def _read_resource(self, data: BinaryIO, offset: int, level: int = 1) -> Ordered entries = self._read_entries(data, directory) for entry in entries: - _rc_type = self._rc_type(entry, data, level) + _rc_type = rc_type(entry, data, level) if entry.DataIsDirectory: resource[_rc_type] = self._read_resource( @@ -153,17 +169,6 @@ def _read_resource(self, data: BinaryIO, offset: int, level: int = 1) -> Ordered return resource - def _rc_type(self, entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY, data: BinaryIO, level: int = 1) -> str: - if level == 1: - return c_pe.ResourceID(entry.Id).name - - if entry.NameIsString: - data.seek(entry.NameOffset) - name_len = c_pe.uint16(data) - return c_pe.wchar[name_len](data) - - return str(entry.Id) - def by_name(self, name: str) -> Resource: """Retrieve the resource by name. @@ -244,14 +249,6 @@ def show_resource_info(self, resources: dict) -> None: f"* resource: {name} offset=0x{resource.offset:02x} size=0x{resource.size:02x} header: {resource.data[:64]}" # noqa: E501 ) - def add(self, name: str, data: bytes) -> None: - # TODO - raise NotImplementedError - - def delete(self, name: str) -> None: - # TODO - raise NotImplementedError - def update_section(self, update_offset: int) -> None: """Function to dynamically update the section data and size when a resource has been modified. @@ -268,11 +265,6 @@ def update_section(self, update_offset: int) -> None: section_data = self.section.data[:header_size] for resource in chain([first_resource], resource_iter): - # Take note of the previous offset and size so we can update any of these values after changing the data - # within the resource - prev_offset = resource.offset - prev_size = resource.size - # Update the resource data section_data += resource.data @@ -282,8 +274,7 @@ def update_section(self, update_offset: int) -> None: if update_offset >= resource.offset: continue - offset = prev_offset + prev_size + 2 - resource.offset = offset + resource.offset = resource.offset + resource.size + 2 # Add the header to the total size so we can check if we need to update the section size new_size += header_size diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/helpers/tls.py index b9f4f49..f47771c 100644 --- a/dissect/executable/pe/helpers/tls.py +++ b/dissect/executable/pe/helpers/tls.py @@ -4,13 +4,14 @@ from typing import TYPE_CHECKING from dissect.executable.pe.c_pe import c_pe +from dissect.executable.pe.helpers.utils import Manager if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE -class TLSManager: +class TLSManager(Manager): """Base class to manage the TLS entries of a PE file. Args: @@ -18,18 +19,17 @@ class TLSManager: """ def __init__(self, pe: PE, section: PESection): - self.pe = pe - self.section = section + super().__init__(pe, section) self.callbacks = [] self.tls: c_pe._IMAGE_TLS_DIRECTORY32 | c_pe._IMAGE_TLS_DIRECTORY64 = None self._read_address: type[c_pe.uint64 | c_pe.uint32] = None self._tls_directory: type[c_pe._IMAGE_TLS_DIRECTORY32 | c_pe._IMAGE_TLS_DIRECTORY64] = None - self._data = b"" - self._image_base = pe.optional_header.ImageBase + self._data: bytes = b"" + self._base_address = pe.optional_header.ImageBase self.set_architecture(pe) - self.parse_tls() + self.parse() def set_architecture(self, pe: PE) -> None: if pe.is64bit(): @@ -39,13 +39,13 @@ def set_architecture(self, pe: PE) -> None: self._read_address = c_pe.uint32 self._tls_directory = c_pe._IMAGE_TLS_DIRECTORY32 - def parse_tls(self) -> None: + def parse(self) -> None: """Parse the TLS directory entry of the PE file when present.""" tls_data = BytesIO(self.section.directory_data(c_pe.IMAGE_DIRECTORY_ENTRY_TLS)) self.tls = self._tls_directory(tls_data) - self.pe.seek(self.tls.AddressOfCallBacks - self._image_base) + self.pe.seek(self.tls.AddressOfCallBacks - self._base_address) # Parse the TLS callback addresses if present while True: @@ -85,7 +85,7 @@ def read_data(self) -> bytes: """ return self.pe.virtual_read( - address=self.tls.StartAddressOfRawData - self._image_base, + address=self.tls.StartAddressOfRawData - self._base_address, size=self.size, ) @@ -118,7 +118,7 @@ def data(self, value: bytes) -> None: section_data.write(self.tls.dumps()) # Write the new TLS data to the section - start_address_rva = self.tls.StartAddressOfRawData - self._image_base + start_address_rva = self.tls.StartAddressOfRawData - self._base_address start_address_section_offset = start_address_rva - self.section.virtual_address section_data.seek(start_address_section_offset) section_data.write(self._data) @@ -126,6 +126,3 @@ def data(self, value: bytes) -> None: # Update the section itself section_data.seek(0) self.section.data = section_data.read() - - def add(self) -> None: - raise NotImplementedError diff --git a/dissect/executable/pe/helpers/utils.py b/dissect/executable/pe/helpers/utils.py index 6277e7f..e22f602 100644 --- a/dissect/executable/pe/helpers/utils.py +++ b/dissect/executable/pe/helpers/utils.py @@ -2,7 +2,7 @@ import struct from functools import lru_cache -from typing import TYPE_CHECKING, Generic, TypeVar +from typing import TYPE_CHECKING if TYPE_CHECKING: from dissect.executable.pe import PE, PESection From 9ad4b5d1b89fee6928234db921637c81d8571860 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Mon, 7 Apr 2025 14:05:27 +0000 Subject: [PATCH 30/43] Create a specific dict manager for the ordered dict datastructure --- dissect/executable/pe/helpers/exports.py | 17 +++---------- dissect/executable/pe/helpers/utils.py | 32 +++++++++++++++++++++++- 2 files changed, 35 insertions(+), 14 deletions(-) diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py index 7799be6..9257e74 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/helpers/exports.py @@ -1,12 +1,11 @@ from __future__ import annotations -from collections import OrderedDict from dataclasses import dataclass from io import BytesIO from typing import TYPE_CHECKING from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers.utils import Manager +from dissect.executable.pe.helpers.utils import DictManager if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection @@ -34,11 +33,9 @@ def __repr__(self) -> str: return f"" -class ExportManager(Manager): +class ExportManager(DictManager[ExportFunction]): def __init__(self, pe: PE, section: PESection): - self.pe = pe - self.section = section - self.exports: OrderedDict[str, ExportFunction] = OrderedDict() + super().__init__(pe, section) self.parse() @@ -80,10 +77,4 @@ def parse(self) -> None: export_name = c_pe.char[None](export_entry) key = export_name.decode() - self.exports[key] = ExportFunction(ordinal=export_directory.Base + _idx, address=address, name=export_name) - - def add(self) -> None: - raise NotImplementedError - - def delete(self) -> None: - raise NotImplementedError + self.elements[key] = ExportFunction(ordinal=export_directory.Base + _idx, address=address, name=export_name) diff --git a/dissect/executable/pe/helpers/utils.py b/dissect/executable/pe/helpers/utils.py index e22f602..55df13f 100644 --- a/dissect/executable/pe/helpers/utils.py +++ b/dissect/executable/pe/helpers/utils.py @@ -1,8 +1,9 @@ from __future__ import annotations import struct +from collections import OrderedDict from functools import lru_cache -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Generic, TypeVar if TYPE_CHECKING: from dissect.executable.pe import PE, PESection @@ -55,6 +56,9 @@ def pad(size: int) -> bytes: return size * b"\x00" +T = TypeVar("T") + + class Manager: def __init__(self, pe: PE, section: PESection) -> None: self.pe = pe @@ -71,3 +75,29 @@ def delete(self, *args, **kwargs) -> None: def patch(self, *args, **kwargs) -> None: raise NotImplementedError + + +class DictManager(Manager, Generic[T]): + elements: OrderedDict[str, T] + + def __init__(self, pe: PE, section: PESection) -> None: + super().__init__(pe, section) + self.elements = OrderedDict() + + def __getitem__(self, key: str) -> T: + return self.elements[key] + + def add(self, name: str, elem: T) -> None: + self._add(name, elem) + self.elements.update({name: elem}) + + def delete(self, name: str) -> None: + if name in self.elements: + self._delete(name) + raise KeyError("Name not inside internal structure.") + + def _add(self, name: str, elem: T) -> None: + raise NotImplementedError + + def _delete(self, name: str) -> None: + raise NotImplementedError From 463fd56ae1344aa62693e47631419fa6590e01b4 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 8 Apr 2025 08:51:17 +0000 Subject: [PATCH 31/43] Move to use different managers --- dissect/executable/pe/helpers/imports.py | 14 ++--- dissect/executable/pe/helpers/patcher.py | 4 +- dissect/executable/pe/helpers/relocations.py | 7 +-- dissect/executable/pe/helpers/resources.py | 61 ++++++++++---------- dissect/executable/pe/helpers/tls.py | 8 +-- dissect/executable/pe/helpers/utils.py | 54 ++++++++++------- dissect/executable/pe/pe.py | 27 ++++----- tests/test_pe.py | 6 +- tests/test_pe_modifications.py | 10 ++-- 9 files changed, 96 insertions(+), 95 deletions(-) diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index 58f40e5..fe4c737 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -6,7 +6,7 @@ from dissect.executable.pe.c_pe import c_pe from dissect.executable.pe.helpers import utils -from dissect.executable.pe.helpers.utils import Manager, create_struct +from dissect.executable.pe.helpers.utils import DictManager, create_struct if TYPE_CHECKING: from collections.abc import Iterator @@ -108,7 +108,7 @@ def __repr__(self) -> str: return f"" -class ImportManager(Manager): +class ImportManager(DictManager[ImportModule]): """The base class for dealing with the imports that are present within the PE file. Args: @@ -123,7 +123,6 @@ def __init__(self, pe: PE, section: PESection): self.import_data = bytearray() self.new_size_of_image = 0 self.section_data = bytearray() - self.imports: OrderedDict[str, ImportModule] = OrderedDict() self.thunks: list[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = [] self._thunk_data: type[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = None @@ -170,7 +169,7 @@ def parse(self) -> None: ImportFunction(pe=self.pe, thunkdata=thunkdata, high_bit=self._high_bit) for thunkdata in self.parse_thunks(offset=first_thunk) ) - self.imports[modulename.decode()] = module + self.elements[modulename.decode()] = module def import_descriptors(self, import_data: BinaryIO) -> Iterator[tuple[int, c_pe.IMAGE_IMPORT_DESCRIPTOR]]: """Parse the import descriptors of the PE file. @@ -232,14 +231,11 @@ def add(self, dllname: str, functions: list[str]) -> None: ImportFunction(pe=self.pe, thunkdata=None, high_bit=self._high_bit, name=function) for function in functions ) - self.imports[dllname] = _module + self.elements[dllname] = _module # Rebuild the import table with the new import module and functions self.build_import_table() - def delete(self, dllname: str, functions: list) -> None: - raise NotImplementedError - def build_import_table(self) -> None: """Function to rebuild the import table after a change has been made to the PE imports. @@ -253,7 +249,7 @@ def build_import_table(self) -> None: import_descriptors: list[c_pe.IMAGE_IMPORT_DESCRIPTOR] = [] self.import_data = bytearray() - for name, module in self.imports.items(): + for name, module in self.elements.items(): # Take note of the current offset to store the modulename name_offset = len(self.import_data) self.import_data += name.encode() + b"\x00" diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index ecad7b4..c8e1c38 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -274,7 +274,7 @@ def _patch_rsrc_rvas(self) -> None: section_data = BytesIO() self.seek(directory_va) - for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc.data_offset): + for rsrc_entry in self.pe.resources.raw(lambda rsrc: rsrc.data_offset): entry_offset = rsrc_entry.offset entry = rsrc_entry.entry @@ -305,7 +305,7 @@ def _patch_tls_rvas(self) -> None: return self.seek(directory_va) - tls_directory = self.pe.tls_mgr._tls_directory(self.patched_pe) + tls_directory = self.pe.tls._tls_directory(self.patched_pe) image_base = self.pe.optional_header.ImageBase diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/helpers/relocations.py index d0fb7c2..dea2f4a 100644 --- a/dissect/executable/pe/helpers/relocations.py +++ b/dissect/executable/pe/helpers/relocations.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers.utils import Manager +from dissect.executable.pe.helpers.utils import ListManager if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection @@ -19,7 +19,7 @@ class Relocation: entries: list[int] -class RelocationManager(Manager): +class RelocationManager(ListManager[Relocation]): """Base class for dealing with the relocations within the PE file. Args: @@ -29,7 +29,6 @@ class RelocationManager(Manager): def __init__(self, pe: PE, section: PESection): super().__init__(pe, section) - self.relocations: list[Relocation] = [] self.parse() @@ -48,7 +47,7 @@ def parse(self) -> None: number_of_entries = (reloc_directory.SizeOfBlock - len(reloc_directory.dumps())) // 2 entries = [entry for _ in range(number_of_entries) if (entry := c_pe.uint16(reloc_data))] - self.relocations.append( + self.elements.append( Relocation( rva=reloc_directory.VirtualAddress, number_of_entries=number_of_entries, diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index f0ed589..f29debd 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -2,6 +2,7 @@ from collections import OrderedDict from dataclasses import dataclass +from functools import partial from io import BytesIO from itertools import chain from textwrap import indent @@ -9,11 +10,11 @@ from dissect.executable.exception import ResourceException from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers.utils import Manager +from dissect.executable.pe.helpers.utils import DictManager if TYPE_CHECKING: from collections.abc import Iterator - from typing import BinaryIO + from typing import BinaryIO, Callable from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE @@ -28,7 +29,7 @@ class RawResource: resource: Resource | None = None -def rc_type(entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY, data: BinaryIO, depth: int = 1) -> str: +def rc_type_name(entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY, data: BinaryIO, depth: int = 1) -> str: """Returns the name of the rc type depending on the data and the depth level of the resource""" if depth == 1: return c_pe.ResourceID(entry.Id).name @@ -41,7 +42,7 @@ def rc_type(entry: c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY, data: BinaryIO, depth: i return str(entry.Id) -class ResourceManager(Manager): +class ResourceManager(DictManager["Resource"]): """Base class to perform actions regarding the resources within the PE file. Args: @@ -50,18 +51,17 @@ class ResourceManager(Manager): """ def __init__(self, pe: PE, section: PESection): - self.pe = pe - self.section = section - self.resources: OrderedDict[str, Resource] = OrderedDict() + super().__init__(pe, section) + self.elements: OrderedDict[str, Resource] = OrderedDict() self.raw_resources: list[RawResource] = [] - + self.values = partial(self._resources, self.elements) self.parse() def parse(self) -> None: """Parse the resource directory entry of the PE file.""" rsrc_data = BytesIO(self.section.directory_data(c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE)) - self.resources = self._read_resource(data=rsrc_data, offset=0, level=1) + self.elements = self._read_resource(data=rsrc_data, offset=0) def _read_entries( self, data: BinaryIO, directory: c_pe.IMAGE_RESOURCE_DIRECTORY @@ -125,7 +125,7 @@ def _handle_data_entry(self, data: BinaryIO, entry: c_pe.IMAGE_RESOURCE_DIRECTOR ) return rsrc - def _read_resource(self, data: BinaryIO, offset: int, level: int = 1) -> OrderedDict[str, Resource]: + def _read_resource(self, data: BinaryIO, offset: int, depth: int = 1) -> OrderedDict[str, Resource]: """Recursively read the resources within the PE file. Each resource is added to the dictionary that is available to the user, as well as a list of @@ -134,8 +134,7 @@ def _read_resource(self, data: BinaryIO, offset: int, level: int = 1) -> Ordered Args: data: The data of the resource. offset: The offset of the resource. - rc_type: The type of the resource. - level: The depth level of the resource, this dictates the resource type. + depth: The depth level of the resource, this dictates the resource type. Returns: A dictionary containing the resources that were found. @@ -153,23 +152,21 @@ def _read_resource(self, data: BinaryIO, offset: int, level: int = 1) -> Ordered ) ) - entries = self._read_entries(data, directory) - - for entry in entries: - _rc_type = rc_type(entry, data, level) + for entry in self._read_entries(data, directory): + rc_name = rc_type_name(entry, data, depth) if entry.DataIsDirectory: - resource[_rc_type] = self._read_resource( + resource[rc_name] = self._read_resource( data=data, offset=entry.OffsetToDirectory, - level=level + 1, + depth=depth + 1, ) else: - resource[_rc_type] = self._handle_data_entry(data=data, entry=entry, rc_type=_rc_type) + resource[rc_name] = self._handle_data_entry(data=data, entry=entry, rc_type=rc_name) return resource - def by_name(self, name: str) -> Resource: + def by_name(self, name: str) -> Resource | OrderedDict: """Retrieve the resource by name. Args: @@ -180,7 +177,7 @@ def by_name(self, name: str) -> Resource: """ try: - return self.resources[name] + return self.elements[name] except KeyError: raise ResourceException(f"Resource {name} not found!") @@ -196,13 +193,13 @@ def by_type(self, rsrc_id: str | c_pe.ResourceID) -> Iterator[Resource]: All of the nodes that contain the requested type. """ - if rsrc_id not in self.resources: + if rsrc_id not in self.elements: raise ResourceException(f"Resource with ID {rsrc_id} not found in PE!") - yield from self.parse_resources(resources=self.resources[rsrc_id]) + yield from self._resources(resources=self.elements[rsrc_id]) - def parse_resources(self, resources: OrderedDict[str, Resource]) -> Iterator[Resource]: - """Parse the resources within the PE file. + def _resources(self, resources: OrderedDict[str, Resource]) -> Iterator[Resource]: + """Iterates throught the resources inside the PE file. Args: resources: A `dict` containing the different resources that were found. @@ -213,11 +210,11 @@ def parse_resources(self, resources: OrderedDict[str, Resource]) -> Iterator[Res for resource in resources.values(): if isinstance(resource, OrderedDict): - yield from self.parse_resources(resources=resource) + yield from self._resources(resources=resource) else: yield resource - def show_resource_tree(self, resources: dict, indentation: int = 0) -> None: + def show_resource_tree(self, resources: OrderedDict[str, OrderedDict | Resource], indentation: int = 0) -> None: """Print the resources within the PE as a tree. Args: @@ -249,6 +246,12 @@ def show_resource_info(self, resources: dict) -> None: f"* resource: {name} offset=0x{resource.offset:02x} size=0x{resource.size:02x} header: {resource.data[:64]}" # noqa: E501 ) + def raw(self, sort_key: Callable | None = None) -> Iterator[RawResource]: + if sort_key: + yield from sorted(self.raw_resources, key=sort_key) + else: + yield from self.raw_resources + def update_section(self, update_offset: int) -> None: """Function to dynamically update the section data and size when a resource has been modified. @@ -258,7 +261,7 @@ def update_section(self, update_offset: int) -> None: new_size = 0 - resource_iter = iter(self.parse_resources(resources=self.resources)) + resource_iter = iter(self._resources(resources=self.elements)) first_resource = next(resource_iter) header_size = first_resource.offset - self.section.virtual_address @@ -392,7 +395,7 @@ def data(self, value: bytes) -> None: prev_offset = 0 prev_size = 0 - for rsrc_entry in sorted(self.pe.raw_resources, key=lambda rsrc: rsrc.data_offset): + for rsrc_entry in self.pe.resources.raw(lambda rsrc: rsrc.data_offset): entry_offset = rsrc_entry.offset entry = rsrc_entry.entry if isinstance(entry, c_pe.IMAGE_RESOURCE_DATA_ENTRY): diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/helpers/tls.py index f47771c..825aa78 100644 --- a/dissect/executable/pe/helpers/tls.py +++ b/dissect/executable/pe/helpers/tls.py @@ -4,14 +4,14 @@ from typing import TYPE_CHECKING from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers.utils import Manager +from dissect.executable.pe.helpers.utils import ListManager if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE -class TLSManager(Manager): +class TLSManager(ListManager[int]): """Base class to manage the TLS entries of a PE file. Args: @@ -20,7 +20,7 @@ class TLSManager(Manager): def __init__(self, pe: PE, section: PESection): super().__init__(pe, section) - self.callbacks = [] + self.tls: c_pe._IMAGE_TLS_DIRECTORY32 | c_pe._IMAGE_TLS_DIRECTORY64 = None self._read_address: type[c_pe.uint64 | c_pe.uint32] = None @@ -52,7 +52,7 @@ def parse(self) -> None: callback_address = self._read_address(self.pe) if not callback_address: break - self.callbacks.append(callback_address) + self.elements.append(callback_address) # Read the TLS data self._data = self.read_data() diff --git a/dissect/executable/pe/helpers/utils.py b/dissect/executable/pe/helpers/utils.py index 55df13f..4ea8c39 100644 --- a/dissect/executable/pe/helpers/utils.py +++ b/dissect/executable/pe/helpers/utils.py @@ -6,6 +6,8 @@ from typing import TYPE_CHECKING, Generic, TypeVar if TYPE_CHECKING: + from collections.abc import Iterable + from dissect.executable.pe import PE, PESection @@ -59,45 +61,53 @@ def pad(size: int) -> bytes: T = TypeVar("T") -class Manager: +class Manager(Generic[T]): + elements: list[T] | OrderedDict[str, T] + def __init__(self, pe: PE, section: PESection) -> None: self.pe = pe self.section = section + def __contains__(self, item: str | T) -> bool: + return item in self.elements + + def __getitem__(self, item: str | int) -> T: + return self.elements[item] + + def __iter__(self): + yield from self.elements + + def __len__(self) -> int: + return len(self.elements) + + def __repr__(self) -> str: + return f"<{self.__class__.__name__} {self.elements}>" + + def values(self) -> Iterable[T]: + raise NotImplementedError + def parse(self) -> None: raise NotImplementedError def add(self, *args, **kwargs) -> None: raise NotImplementedError - def delete(self, *args, **kwargs) -> None: + def delete(self, elem: int | str) -> None: raise NotImplementedError - def patch(self, *args, **kwargs) -> None: + def patch(self, elem: int | str, data: bytes) -> None: raise NotImplementedError class DictManager(Manager, Generic[T]): - elements: OrderedDict[str, T] - def __init__(self, pe: PE, section: PESection) -> None: super().__init__(pe, section) - self.elements = OrderedDict() - - def __getitem__(self, key: str) -> T: - return self.elements[key] - - def add(self, name: str, elem: T) -> None: - self._add(name, elem) - self.elements.update({name: elem}) - - def delete(self, name: str) -> None: - if name in self.elements: - self._delete(name) - raise KeyError("Name not inside internal structure.") + self.elements: OrderedDict[str, T] = OrderedDict() + self.values = self.elements.values - def _add(self, name: str, elem: T) -> None: - raise NotImplementedError - def _delete(self, name: str) -> None: - raise NotImplementedError +class ListManager(Manager, Generic[T]): + def __init__(self, pe: PE, section: PESection) -> None: + super().__init__(pe, section) + self.elements: list[T] = [] + self.values = lambda: self.elements diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index 5d1f0eb..2e5028c 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -57,12 +57,11 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.sections = sections.PESectionManager() - self.imports: OrderedDict[str, imports.ImportModule] = None - self.exports: OrderedDict[str, exports.ExportFunction] = None - self.resources: OrderedDict[str, resources.Resource] = None - self.raw_resources: list[resources.RawResource] = [] - self.relocations: list[relocations.Relocation] = [] - self.tls_callbacks = None + self.imports: imports.ImportManager = None + self.exports: exports.ExportManager = None + self.resources: resources.ResourceManager = None + self.relocations: relocations.RelocationManager = None + self.tls: tls.TLSManager = None self.directories = OrderedDict() @@ -198,28 +197,22 @@ def parse_directories(self) -> None: # Parse the Import Address Table (IAT) if idx == c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT: - self.import_mgr = imports.ImportManager(pe=self, section=section) - self.imports = self.import_mgr.imports + self.imports = imports.ImportManager(pe=self, section=section) if idx == c_pe.IMAGE_DIRECTORY_ENTRY_EXPORT: - self.export_mgr = exports.ExportManager(pe=self, section=section) - self.exports = self.export_mgr.exports + self.exports = exports.ExportManager(pe=self, section=section) # Parse the resources directory entry of the PE file if idx == c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE: - self.rsrc_mgr = resources.ResourceManager(pe=self, section=section) - self.resources = self.rsrc_mgr.resources - self.raw_resources = self.rsrc_mgr.raw_resources + self.resources = resources.ResourceManager(pe=self, section=section) # Parse the relocation directory entry of the PE file if idx == c_pe.IMAGE_DIRECTORY_ENTRY_BASERELOC: - self.reloc_mgr = relocations.RelocationManager(pe=self, section=section) - self.relocations = self.reloc_mgr.relocations + self.relocations = relocations.RelocationManager(pe=self, section=section) # Parse the TLS directory entry of the PE file if idx == c_pe.IMAGE_DIRECTORY_ENTRY_TLS: - self.tls_mgr = tls.TLSManager(pe=self, section=section) - self.tls_callbacks = self.tls_mgr.callbacks + self.tls = tls.TLSManager(pe=self, section=section) def virtual_address(self, address: int) -> int: """Return the virtual address given a (possible) physical address. diff --git a/tests/test_pe.py b/tests/test_pe.py index c524874..7a7110b 100644 --- a/tests/test_pe.py +++ b/tests/test_pe.py @@ -44,7 +44,7 @@ def test_pe_imports() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - assert known_imports == list(pe.imports) + assert known_imports == list(pe.imports.elements) def test_pe_exports() -> None: @@ -60,7 +60,7 @@ def test_pe_exports() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - assert known_exports == list(pe.exports) + assert known_exports == list(pe.exports.elements) def test_pe_resources() -> None: @@ -95,4 +95,4 @@ def test_pe_tls_callbacks() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - assert pe.tls_callbacks == known_callbacks + assert pe.tls.elements == known_callbacks diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py index a343b50..974f704 100644 --- a/tests/test_pe_modifications.py +++ b/tests/test_pe_modifications.py @@ -11,7 +11,7 @@ def test_add_imports() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - pe.import_mgr.add(dllname=dllname, functions=functions) + pe.imports.add(dllname=dllname, functions=functions) patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) @@ -59,13 +59,13 @@ def test_resize_resource_smaller() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - for e in pe.rsrc_mgr.by_type(rsrc_id="Manifest"): + for e in pe.resources.by_type(rsrc_id="Manifest"): e.data = b"kusjesvanSRT, patched with dissect" patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) - assert [patched.data for patched in new_pe.rsrc_mgr.by_type(rsrc_id="Manifest")] == [ + assert [patched.data for patched in new_pe.resources.by_type(rsrc_id="Manifest")] == [ b"kusjesvanSRT, patched with dissect" ] @@ -74,7 +74,7 @@ def test_resize_resource_bigger() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - for e in pe.rsrc_mgr.by_type(rsrc_id="Manifest"): + for e in pe.resources.by_type(rsrc_id="Manifest"): e.data = b"kusjesvanSRT, patched with dissect" + e.data patcher = Patcher(pe=pe) @@ -82,7 +82,7 @@ def test_resize_resource_bigger() -> None: assert [ patched.data[: len(b"kusjesvanSRT, patched with dissect")] - for patched in new_pe.rsrc_mgr.by_type(rsrc_id="Manifest") + for patched in new_pe.resources.by_type(rsrc_id="Manifest") ] == [b"kusjesvanSRT, patched with dissect"] From 15d869012e7dc6d07972ed4f6c42d7425b543825 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 8 Apr 2025 10:46:16 +0000 Subject: [PATCH 32/43] Move patch to PESectionManager --- dissect/executable/pe/helpers/resources.py | 5 +- dissect/executable/pe/helpers/sections.py | 103 ++++++++++----------- dissect/executable/pe/helpers/tls.py | 2 +- dissect/executable/pe/pe.py | 4 +- tests/test_pe_modifications.py | 9 +- 5 files changed, 58 insertions(+), 65 deletions(-) diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index f29debd..2d1710c 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -282,8 +282,7 @@ def update_section(self, update_offset: int) -> None: # Add the header to the total size so we can check if we need to update the section size new_size += header_size - # Update the section - self.section.data = section_data + self.pe.sections.patch(self.section.name, section_data) class Resource: @@ -429,7 +428,7 @@ def data(self, value: bytes) -> None: data = section_data.read() # Update the section data and size - self.section.data = data + self.pe.sections.patch(self.section.name, data) self.pe.optional_header.DataDirectory[c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE].Size = len(data) def __str__(self) -> str: diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index d1c3146..0744fe3 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -2,6 +2,7 @@ from collections import OrderedDict from copy import copy +from itertools import chain from typing import TYPE_CHECKING from dissect.executable.exception import BuildSectionException @@ -15,9 +16,11 @@ class PESectionManager: - def __init__(self) -> None: + def __init__(self, file_alignment: int, section_alignment: int) -> None: self._sections: OrderedDict[str, PESection] = OrderedDict() self._patched_sections: OrderedDict[str, PESection] = OrderedDict() + self._file_alignment = file_alignment + self._section_alignment = section_alignment def add(self, name: str, section: PESection) -> None: self._sections[name] = section @@ -87,6 +90,50 @@ def _in_virtual_range(self, va: int, sections: Iterable[PESection]) -> PESection return None + def patch(self, name: str, data: bytes) -> None: + """Sets the new data of the resource and dynamically updates the other patched sections. + + Args: + name: The section to patch + data: The data to patch it with + """ + patched_section: PESection = self._patched_sections[name] + + # Update the patched section data and size + patched_section._data = data + patched_section.size = len(data) + + if patched_section.size_of_raw_data < patched_section.virtual_size: + patched_section._data += utils.pad(size=patched_section.virtual_size - patched_section.size_of_raw_data) + + iterator = iter(self.sections(patch=True).values()) + first_section = next(iterator) + + prev_ptr = first_section.pointer_to_raw_data + prev_size = first_section.size_of_raw_data + prev_va = first_section.virtual_address + prev_vsize = first_section.virtual_size + + for section in chain([first_section], iterator): + if section.virtual_address == prev_va: + continue + + pointer_to_raw_data = utils.align_int(integer=prev_ptr + prev_size, blocksize=self._file_alignment) + virtual_address = utils.align_int(integer=prev_va + prev_vsize, blocksize=self._section_alignment) + + if section.virtual_address < virtual_address: + """Set the virtual address and raw pointer of the section to the new values, but only do so if the + section virtual address is lower than the previous section. We want to prevent messing up RVA's as + much as possible, this could lead to binaries that are a bit larger than they need to be but that + doesn't really matter.""" + section.virtual_address = virtual_address + section.pointer_to_raw_data = pointer_to_raw_data + + prev_ptr = pointer_to_raw_data + prev_size = section.size_of_raw_data + prev_va = virtual_address + prev_vsize = section.virtual_size + class PESection: """Base class for the PE sections that are present. @@ -230,60 +277,6 @@ def data(self) -> bytes: """Return the data within the section.""" return self._data[: self.virtual_size] - @data.setter - def data(self, value: bytes) -> None: - """Setter to set the new data of the resource, but also dynamically update the offset of the resources within - the same directory. - - This function currently also updates the section sizes and alignment. Ideally this would be moved to a more - abstract function that can handle tasks like these in a more transparant manner. - - Args: - value: The new data of the resource. - """ - - # Keep track of the section changes using the patched_sections dictionary - section_manager = self.pe.sections - patched_section: PESection = section_manager.get(name=self.name, patch=True) - patched_section._data = value - patched_section.size = len(value) - - # Set the new data and size - self._data = value - self.size = len(value) - - # Pad the remainder of the section if the SizeOfRawData is smaller than the VirtualSize - if self.size_of_raw_data < self.virtual_size: - self._data += utils.pad(size=self.virtual_size - self.size_of_raw_data) - - # Take note of the first section as our starting point - first_section = next(iter(section_manager.sections(patch=True).values())) - - prev_ptr = first_section.pointer_to_raw_data - prev_size = first_section.size_of_raw_data - prev_va = first_section.virtual_address - prev_vsize = first_section.virtual_size - - for section in section_manager.sections(patch=True).values(): - if section.virtual_address == prev_va: - continue - - pointer_to_raw_data = utils.align_int(integer=prev_ptr + prev_size, blocksize=self.pe.file_alignment) - virtual_address = utils.align_int(integer=prev_va + prev_vsize, blocksize=self.pe.section_alignment) - - if section.virtual_address < virtual_address: - """Set the virtual address and raw pointer of the section to the new values, but only do so if the - section virtual address is lower than the previous section. We want to prevent messing up RVA's as - much as possible, this could lead to binaries that are a bit larger than they need to be but that - doesn't really matter.""" - section.virtual_address = virtual_address - section.pointer_to_raw_data = pointer_to_raw_data - - prev_ptr = pointer_to_raw_data - prev_size = section.size_of_raw_data - prev_va = virtual_address - prev_vsize = section.virtual_size - def dump(self) -> bytes: """Return the section header as a `bytes` object.""" return self.section.dumps() diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/helpers/tls.py index 825aa78..08be933 100644 --- a/dissect/executable/pe/helpers/tls.py +++ b/dissect/executable/pe/helpers/tls.py @@ -125,4 +125,4 @@ def data(self, value: bytes) -> None: # Update the section itself section_data.seek(0) - self.section.data = section_data.read() + self.pe.sections.patch(self.section.name, section_data.read()) diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index 2e5028c..f1121e4 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -55,8 +55,6 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.section_header_offset = 0 self.last_section_offset = 0 - self.sections = sections.PESectionManager() - self.imports: imports.ImportManager = None self.exports: exports.ExportManager = None self.resources: resources.ResourceManager = None @@ -76,6 +74,7 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.section_alignment = self.nt_headers.OptionalHeader.SectionAlignment self.timestamp = datetime.fromtimestamp(self.file_header.TimeDateStamp, tz=timezone.utc) + self.sections = sections.PESectionManager(self.file_alignment, self.section_alignment) # Parse the section header self.parse_section_header() @@ -189,7 +188,6 @@ def parse_directories(self) -> None: # Take note of the current directory VA so we can dynamically update it when resizing sections section = self.datadirectory_section(index=idx) - section_dir = self.optional_header.DataDirectory[idx] directory_va_offset = section_dir.VirtualAddress - section.virtual_address diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py index 974f704..c73c8f8 100644 --- a/tests/test_pe_modifications.py +++ b/tests/test_pe_modifications.py @@ -27,7 +27,7 @@ def test_resize_section_smaller() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - pe.sections.get(name=".text").data = b"kusjesvanSRT, patched with dissect" + pe.sections.patch(name=".text", data=b"kusjesvanSRT, patched with dissect") patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) @@ -43,9 +43,12 @@ def test_resize_section_bigger() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - original_size = pe.sections.get(name=".rdata").size + section = pe.sections.get(name=".rdata") - pe.sections.get(name=".rdata", patch=True).data += b"kusjesvanSRT, patched with dissect" * 100 + original_size = section.size + + patch_data = section.data + b"kusjesvanSRT, patched with dissect" * 100 + pe.sections.patch(name=".rdata", data=patch_data) patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) From 8572b4db46bd6539014903230055b9f7c1b5f6bc Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 8 Apr 2025 10:47:39 +0000 Subject: [PATCH 33/43] Move assigning section directories to a specific method --- dissect/executable/pe/helpers/sections.py | 3 +++ dissect/executable/pe/pe.py | 5 +---- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 0744fe3..996e7c6 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -167,6 +167,9 @@ def directory_data(self, index: int) -> bytes: offset, size = dir_information return self.data[offset : offset + size] + def add_directory(self, index: int, section_dir: c_pe.IMAGE_DATA_DIRECTORY) -> None: + self.directories[index] = (section_dir.VirtualAddress - self.virtual_address, section_dir.Size) + def read_data(self) -> bytes: """Return the data within the section. diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index f1121e4..b809289 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -61,8 +61,6 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.relocations: relocations.RelocationManager = None self.tls: tls.TLSManager = None - self.directories = OrderedDict() - # We always want to parse the DOS header and NT headers self.parse_headers() @@ -190,8 +188,7 @@ def parse_directories(self) -> None: section = self.datadirectory_section(index=idx) section_dir = self.optional_header.DataDirectory[idx] - directory_va_offset = section_dir.VirtualAddress - section.virtual_address - section.directories[idx] = (directory_va_offset, section_dir.Size) + section.add_directory(idx, section_dir) # Parse the Import Address Table (IAT) if idx == c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT: From 5de7a413c2d7338014f5b7439d7467cf9c3b321f Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 8 Apr 2025 10:48:13 +0000 Subject: [PATCH 34/43] Readability changes --- dissect/executable/pe/helpers/patcher.py | 2 +- dissect/executable/pe/pe.py | 33 +++++++++--------------- 2 files changed, 13 insertions(+), 22 deletions(-) diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index c8e1c38..570e98f 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -75,7 +75,7 @@ def seek(self, address: int) -> None: address: The virtual address to seek to. """ - raw_address = self.pe.virtual_address(address=address) + raw_address = self.pe.virtual_address(physical_address=address) self.patched_pe.seek(raw_address) def _build_section_table(self) -> None: diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index b809289..b8cb3bd 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -25,8 +25,6 @@ ) if TYPE_CHECKING: - from collections.abc import Iterable - from dissect.cstruct.cstruct import cstruct @@ -146,13 +144,6 @@ def parse_section_header(self) -> None: self.last_section_offset = self.sections.last_section().offset - def _section_in_range(self, address: int, values: Iterable[sections.PESection]) -> sections.PESection | None: - for section in values: - if section.virtual_address <= address < section.virtual_address + section.virtual_size: - return section - - return None - def datadirectory_section(self, index: int) -> sections.PESection: """Return the section that contains the given virtual address. @@ -209,7 +200,7 @@ def parse_directories(self) -> None: if idx == c_pe.IMAGE_DIRECTORY_ENTRY_TLS: self.tls = tls.TLSManager(pe=self, section=section) - def virtual_address(self, address: int) -> int: + def virtual_address(self, physical_address: int) -> int: """Return the virtual address given a (possible) physical address. Args: @@ -220,14 +211,14 @@ def virtual_address(self, address: int) -> int: """ if self.virtual: - return address + return physical_address - if section := self.sections.in_range(address, patch=True): - return section.pointer_to_raw_data + (address - section.virtual_address) + if section := self.sections.in_range(physical_address, patch=True): + return section.pointer_to_raw_data + (physical_address - section.virtual_address) - raise InvalidVA(f"VA not found in sections: {address:#04x}") + raise InvalidVA(f"VA not found in sections: {physical_address:#04x}") - def raw_address(self, offset: int) -> int: + def raw_address(self, virtual_address: int) -> int: """Return the physical address given a virtual address. Args: @@ -236,10 +227,10 @@ def raw_address(self, offset: int) -> int: Returns: The physical address as an `int`. """ - if section := self.sections.in_raw_range(offset, patch=True): - return section.virtual_address + (offset - section.pointer_to_raw_data) + if section := self.sections.in_raw_range(virtual_address, patch=True): + return section.virtual_address + (virtual_address - section.pointer_to_raw_data) - raise InvalidAddress(f"Raw address not found in sections: {offset:#04x}") + raise InvalidAddress(f"Raw address not found in sections: {virtual_address:#04x}") def virtual_read(self, address: int, size: int) -> bytes: """Wrapper for reading virtual address offsets within a PE file. @@ -252,7 +243,7 @@ def virtual_read(self, address: int, size: int) -> bytes: The bytes that were read. """ - physical_address = self.virtual_address(address=address) + physical_address = self.virtual_address(physical_address=address) if self.virtual: return self.pe_file.readoffset(offset=physical_address, size=size) @@ -284,7 +275,7 @@ def seek(self, address: int) -> None: address: The virtual address to seek to. """ - raw_address = self.virtual_address(address=address) + raw_address = self.virtual_address(physical_address=address) self.pe_file.seek(raw_address) def tell(self) -> int: @@ -295,7 +286,7 @@ def tell(self) -> int: """ offset = self.pe_file.tell() - return self.raw_address(offset=offset) + return self.raw_address(virtual_address=offset) def read(self, size: int) -> bytes: """Read x amount of bytes of the PE file. From f8f03628b7c001ff609480f81a7d23d161861216 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 8 Apr 2025 12:56:57 +0000 Subject: [PATCH 35/43] Reorder operations for building This is mostly cause I did not expect it updating pe section when it was creating an import table --- dissect/executable/pe/helpers/imports.py | 99 +++++++++++------------- 1 file changed, 47 insertions(+), 52 deletions(-) diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index fe4c737..32a7175 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -1,6 +1,5 @@ from __future__ import annotations -from collections import OrderedDict from io import BytesIO from typing import TYPE_CHECKING, BinaryIO @@ -10,6 +9,7 @@ if TYPE_CHECKING: from collections.abc import Iterator + from struct import Struct from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE @@ -118,6 +118,7 @@ class ImportManager(DictManager[ImportModule]): def __init__(self, pe: PE, section: PESection): super().__init__(pe, section) + self._section_manager = pe.sections self.image_size: int = pe.optional_header.SizeOfImage self.import_directory_rva = 0 self.import_data = bytearray() @@ -126,6 +127,7 @@ def __init__(self, pe: PE, section: PESection): self.thunks: list[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = [] self._thunk_data: type[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = None + self._thunkdata_packing: Struct = None self._high_bit: int = 0 self.set_architecture(pe) @@ -135,9 +137,11 @@ def set_architecture(self, pe: PE) -> None: if pe.is64bit(): self._thunk_data = c_pe.IMAGE_THUNK_DATA64 self._high_bit = 1 << 63 + self._thunkdata_packing = create_struct(" None: """Parse the imports of the PE file. @@ -216,7 +220,7 @@ def add(self, dllname: str, functions: list[str]) -> None: functions: A `list` of function names belonging to the module. """ - self.last_section = self.pe.sections.last_section(patch=True) + self.last_section = self._section_manager.last_section(patch=True) # Build a dummy import module _module = ImportModule( @@ -234,55 +238,56 @@ def add(self, dllname: str, functions: list[str]) -> None: self.elements[dllname] = _module # Rebuild the import table with the new import module and functions - self.build_import_table() + imports, import_rva, directory_size = self.build_import_table() - def build_import_table(self) -> None: + # Create a new section + section_data = utils.align_data(data=imports, blocksize=self.pe.file_alignment) + size = len(imports) + c_pe.IMAGE_SECTION_HEADER.size + self.pe.add_section( + name=".idata", + data=section_data, + datadirectory=c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT, + datadirectory_rva=import_rva, + datadirectory_size=directory_size, + size=size, + ) + + def build_import_table(self) -> tuple[bytearray, int, int]: """Function to rebuild the import table after a change has been made to the PE imports. Currently we're using the .idata section to store the imports, there might be a better way to do this but for now this will do. """ - # Reset the known thunkdata - self.thunks = [] - import_descriptors: list[c_pe.IMAGE_IMPORT_DESCRIPTOR] = [] - self.import_data = bytearray() + import_data = bytearray() for name, module in self.elements.items(): # Take note of the current offset to store the modulename - name_offset = len(self.import_data) - self.import_data += name.encode() + b"\x00" + name_offset = len(import_data) + import_data += name.encode() + b"\x00" # Build the module imports and get the RVA of the first thunk to generate an import descriptor - first_thunk_rva = self._build_module_imports(functions=module.functions) + module_offset = len(import_data) + module_bytes, offsets = self._build_module_imports(module_offset, module.functions) + thunkdata = self._build_thunkdata(offsets) + import_data += module_bytes + thunkdata + import_descriptor = self._build_import_descriptor( - first_thunk_rva=first_thunk_rva, + first_thunk_rva=self.image_size + module_offset + len(module_bytes), name_rva=self.image_size + name_offset, ) import_descriptors.append(import_descriptor) - datadirectory_size = 0 - # Take note of the RVA of the first import descriptor - import_rva = self.image_size + len(self.import_data) - for descriptor in import_descriptors: - self.import_data += descriptor.dumps() - datadirectory_size += len(descriptor) + import_rva = self.image_size + len(import_data) + descriptor_data = b"".join(descriptor.dumps() for descriptor in import_descriptors) + directory_size = len(descriptor_data) + import_data += descriptor_data - # Create a new section - section_data = utils.align_data(data=self.import_data, blocksize=self.pe.file_alignment) - size = len(self.import_data) + c_pe.IMAGE_SECTION_HEADER.size - self.pe.add_section( - name=".idata", - data=section_data, - datadirectory=c_pe.IMAGE_DIRECTORY_ENTRY_IMPORT, - datadirectory_rva=import_rva, - datadirectory_size=datadirectory_size, - size=size, - ) + return import_data, import_rva, directory_size - def _build_module_imports(self, functions: list[ImportFunction]) -> int: + def _build_module_imports(self, function_offset: int, functions: list[ImportFunction]) -> tuple[bytes, list]: """Function to build the imports for a module. This function is responsible for building the functions by name, as well as the associated thunkdata that is @@ -296,19 +301,16 @@ def _build_module_imports(self, functions: list[ImportFunction]) -> int: """ function_offsets = [] - _hint_struct = create_struct(" bytes: """Function to build the thunkdata for the new import table. @@ -320,18 +322,11 @@ def _build_thunkdata(self, import_rvas: list[int]) -> bytes: The thunkdata as a `bytes` object. """ - packing = " c_pe.IMAGE_IMPORT_DESCRIPTOR: """Function to build the import descriptor for the new import table. From dc91018a4f1e647343374ff4ec46efe4addbf031 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 8 Apr 2025 13:01:57 +0000 Subject: [PATCH 36/43] make sure virtual_address gets defined --- dissect/executable/pe/helpers/patcher.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index 570e98f..44afbe3 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -166,9 +166,11 @@ def _patch_import_rvas(self) -> None: # Check the original RVA associated with the AddressOfData field in the thunkdata, retrieve the # original VA # and use it to also select the patched virtual address of this section that the RVA is located in - if section := self._section_manager.get(va=function.data_address): - virtual_address = section.virtual_address - new_virtual_address = self._section_manager.get(name=section.name, patch=True).virtual_address + if (section := self._section_manager.get(va=function.data_address)) is None: + continue + + virtual_address = section.virtual_address + new_virtual_address = self._section_manager.get(name=section.name, patch=True).virtual_address # Calculate the offset using the VA of the section and update the thunkdata va_offset = function.data_address - virtual_address From ce9f60cd4ac4b4904926cddccfcdc1a5e3c0a63b Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 8 Apr 2025 13:13:17 +0000 Subject: [PATCH 37/43] Expose inner import manager variables --- dissect/executable/pe/helpers/imports.py | 19 +++++++++---------- dissect/executable/pe/helpers/patcher.py | 3 +-- 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index 32a7175..535c703 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -124,10 +124,9 @@ def __init__(self, pe: PE, section: PESection): self.import_data = bytearray() self.new_size_of_image = 0 self.section_data = bytearray() - self.thunks: list[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = [] - self._thunk_data: type[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = None - self._thunkdata_packing: Struct = None + self.thunk_struct: type[c_pe.IMAGE_THUNK_DATA32 | c_pe.IMAGE_THUNK_DATA64] = None + self.thunk_pack_struct: Struct = None self._high_bit: int = 0 self.set_architecture(pe) @@ -135,13 +134,13 @@ def __init__(self, pe: PE, section: PESection): def set_architecture(self, pe: PE) -> None: if pe.is64bit(): - self._thunk_data = c_pe.IMAGE_THUNK_DATA64 + self.thunk_struct = c_pe.IMAGE_THUNK_DATA64 self._high_bit = 1 << 63 - self._thunkdata_packing = create_struct(" None: """Parse the imports of the PE file. @@ -206,7 +205,7 @@ def parse_thunks(self, offset: int) -> Iterator[c_pe.IMAGE_THUNK_DATA32 | c_pe.I self.pe.seek(offset) while True: - thunkdata = self._thunk_data(self.pe) + thunkdata = self.thunk_struct(self.pe) if not thunkdata.u1.Function: break @@ -323,8 +322,8 @@ def _build_thunkdata(self, import_rvas: list[int]) -> bytes: """ thunkdata: list[bytes] = [] - thunkdata.extend(self._thunkdata_packing.pack(rva + self.image_size) for rva in import_rvas) - thunkdata.append(self._thunkdata_packing.pack(0)) + thunkdata.extend(self.thunk_pack_struct.pack(rva + self.image_size) for rva in import_rvas) + thunkdata.append(self.thunk_pack_struct.pack(0)) return b"".join(thunkdata) diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index 44afbe3..125e316 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -1,6 +1,5 @@ from __future__ import annotations -import copy from io import BytesIO from typing import TYPE_CHECKING @@ -177,7 +176,7 @@ def _patch_import_rvas(self) -> None: new_address = new_virtual_address + va_offset # Avoid overwriting the original data - tmp_thunkdata = copy.deepcopy(function.thunkdata) + tmp_thunkdata = self.pe.imports.thunk_struct() tmp_thunkdata.u1.AddressOfData = new_address tmp_thunkdata.u1.ForwarderString = new_address tmp_thunkdata.u1.Function = new_address From 98098cccf6387dce0ef4947a78d11baf45d25b19 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Tue, 8 Apr 2025 14:43:00 +0000 Subject: [PATCH 38/43] Move updating resource to ResourceManager.patch --- dissect/executable/pe/helpers/resources.py | 121 ++++++++++----------- tests/test_pe_modifications.py | 7 +- 2 files changed, 63 insertions(+), 65 deletions(-) diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index 2d1710c..0917f42 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -63,6 +63,66 @@ def parse(self) -> None: rsrc_data = BytesIO(self.section.directory_data(c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE)) self.elements = self._read_resource(data=rsrc_data, offset=0) + def patch(self, name: str, data: bytes) -> None: + """Sets the new data of the resource and updates the offsets with the resources within the same directory. + + Resource looks like this: + + | Resource headers (1*...) | + | ------------------------ | + | Resource data (1*...) | + + So it is not important in what order the metadata of the entry gets written. + """ + try: + resource = next(self.by_type(name)) + except StopIteration: + raise ValueError(f"Could not find a resource by type for {name}") + + # TODO: Still rewrites the data to the original instance. Maybe we should change that. + resource._data = data + resource.size = len(data) + + output = BytesIO() + prev_offset = prev_size = 0 + + for rsrc_entry in self.raw(lambda rsrc: rsrc.data_offset): + entry_offset = rsrc_entry.offset + entry = rsrc_entry.entry + + # Write the resource entry into the section + output.seek(entry_offset) + output.write(entry.dumps()) + + if not isinstance(entry, c_pe.IMAGE_RESOURCE_DATA_ENTRY): + continue + + rsrc_obj = rsrc_entry.resource + data_offset = rsrc_entry.data_offset + + # Normally the data is separated by a null byte, increment the new offset by 1 + new_data_offset = prev_offset + prev_size + # if new_data_offset and (new_data_offset > data_offset or new_data_offset < data_offset): + if new_data_offset and new_data_offset != data_offset: + data_offset = new_data_offset + rsrc_entry.data_offset = data_offset + rsrc_obj.offset = self.section.virtual_address + data_offset + + # Write the resource entry data into the section + output.seek(data_offset) + output.write(rsrc_obj.data) + + # Take note of the offset and size so we can update any of these values after changing the data within + # the resource + prev_offset = data_offset + prev_size = rsrc_obj.size + + output.seek(0) + _data = output.read() + + self.pe.sections.patch(self.section.name, _data) + self.pe.optional_header.DataDirectory[c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE].size = len(_data) + def _read_entries( self, data: BinaryIO, directory: c_pe.IMAGE_RESOURCE_DIRECTORY ) -> list[c_pe.IMAGE_RESOURCE_DIRECTORY_ENTRY]: @@ -370,67 +430,6 @@ def data(self) -> bytes: """Return the data within the resource.""" return self._data - @data.setter - def data(self, value: bytes) -> None: - """Setter to set the new data of the resource, but also dynamically update the offset of the resources within - the same directory. - - This function currently also updates the section sizes and alignment. Ideally this would be moved to a more - abstract function that - can handle tasks like these in a more transparant manner. - - Args: - value: The new data of the resource. - """ - - # Set the new data - self._data = value - - if len(value) != self.entry.Size: - self.size = len(value) - - section_data = BytesIO() - - prev_offset = 0 - prev_size = 0 - - for rsrc_entry in self.pe.resources.raw(lambda rsrc: rsrc.data_offset): - entry_offset = rsrc_entry.offset - entry = rsrc_entry.entry - if isinstance(entry, c_pe.IMAGE_RESOURCE_DATA_ENTRY): - rsrc_obj = rsrc_entry.resource - data_offset = rsrc_entry.data_offset - - # Normally the data is separated by a null byte, increment the new offset by 1 - new_data_offset = prev_offset + prev_size - # if new_data_offset and (new_data_offset > data_offset or new_data_offset < data_offset): - if new_data_offset and new_data_offset != data_offset: - data_offset = new_data_offset - rsrc_entry.data_offset = data_offset - rsrc_obj.offset = self.section.virtual_address + data_offset - - data = rsrc_obj.data - - # Write the resource entry data into the section - section_data.seek(data_offset) - section_data.write(data) - - # Take note of the offset and size so we can update any of these values after changing the data within - # the resource - prev_offset = data_offset - prev_size = rsrc_obj.size - - # Write the resource entry into the section - section_data.seek(entry_offset) - section_data.write(entry.dumps()) - - section_data.seek(0) - data = section_data.read() - - # Update the section data and size - self.pe.sections.patch(self.section.name, data) - self.pe.optional_header.DataDirectory[c_pe.IMAGE_DIRECTORY_ENTRY_RESOURCE].Size = len(data) - def __str__(self) -> str: return str(self.name) diff --git a/tests/test_pe_modifications.py b/tests/test_pe_modifications.py index c73c8f8..9e2dbdb 100644 --- a/tests/test_pe_modifications.py +++ b/tests/test_pe_modifications.py @@ -62,8 +62,7 @@ def test_resize_resource_smaller() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - for e in pe.resources.by_type(rsrc_id="Manifest"): - e.data = b"kusjesvanSRT, patched with dissect" + pe.resources.patch("Manifest", b"kusjesvanSRT, patched with dissect") patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) @@ -77,8 +76,8 @@ def test_resize_resource_bigger() -> None: with data_file("testexe.exe").open("rb") as pe_fh: pe = PE(pe_file=pe_fh) - for e in pe.resources.by_type(rsrc_id="Manifest"): - e.data = b"kusjesvanSRT, patched with dissect" + e.data + resource = next(pe.resources.by_type(rsrc_id="Manifest")) + pe.resources.patch("Manifest", b"kusjesvanSRT, patched with dissect" + resource.data) patcher = Patcher(pe=pe) new_pe = PE(pe_file=patcher.build()) From 0c324783431476e6e50b59a20fe92dbacb345226 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Wed, 9 Apr 2025 09:05:24 +0000 Subject: [PATCH 39/43] Small cleanup --- dissect/executable/pe/helpers/exports.py | 2 +- dissect/executable/pe/helpers/imports.py | 1 + dissect/executable/pe/pe.py | 3 --- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py index 9257e74..3c27460 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/helpers/exports.py @@ -24,7 +24,7 @@ class ExportFunction: ordinal: int address: int - name: bytes = b"" + name: bytes | None = b"" def __str__(self) -> str: return self.name.decode() if self.name else f"#{self.ordinal}" diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index 535c703..7a12540 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -293,6 +293,7 @@ def _build_module_imports(self, function_offset: int, functions: list[ImportFunc used to parse the imports at a later stage. Args: + function_offset: The start offset of the functions for this module functions: A `list` of `ImportFunction` objects. Returns: diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index b8cb3bd..749d2dc 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -1,6 +1,5 @@ from __future__ import annotations -from collections import OrderedDict from datetime import datetime, timezone from io import BytesIO from pathlib import Path @@ -51,7 +50,6 @@ def __init__(self, pe_file: BinaryIO, virtual: bool = False): self.optional_header: c_pe.IMAGE_OPTIONAL_HEADER | c_pe.IMAGE_OPTIONAL_HEADER64 = None self.section_header_offset = 0 - self.last_section_offset = 0 self.imports: imports.ImportManager = None self.exports: exports.ExportManager = None @@ -439,7 +437,6 @@ def add_section( # Update the last section offset offset = last_section.offset + c_pe.IMAGE_SECTION_HEADER.size - self.last_section_offset = offset # Increment the NumberOfSections field self.file_header.NumberOfSections += 1 From 0290ec6a523fc645eda21036bbe8dce0f1f9cca3 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Wed, 9 Apr 2025 09:07:59 +0000 Subject: [PATCH 40/43] Move helpers/utils.py to dissect/executable/utils.py --- dissect/executable/pe/helpers/builder.py | 2 +- dissect/executable/pe/helpers/exports.py | 2 +- dissect/executable/pe/helpers/imports.py | 4 ++-- dissect/executable/pe/helpers/patcher.py | 2 +- dissect/executable/pe/helpers/relocations.py | 2 +- dissect/executable/pe/helpers/resources.py | 2 +- dissect/executable/pe/helpers/sections.py | 2 +- dissect/executable/pe/helpers/tls.py | 2 +- dissect/executable/pe/pe.py | 2 +- dissect/executable/{pe/helpers => }/utils.py | 0 10 files changed, 10 insertions(+), 10 deletions(-) rename dissect/executable/{pe/helpers => }/utils.py (100%) diff --git a/dissect/executable/pe/helpers/builder.py b/dissect/executable/pe/helpers/builder.py index ace8105..d1015c0 100644 --- a/dissect/executable/pe/helpers/builder.py +++ b/dissect/executable/pe/helpers/builder.py @@ -3,9 +3,9 @@ from datetime import datetime, timezone from io import BytesIO +from dissect.executable import utils from dissect.executable.exception import BuildSectionException from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers import utils from dissect.executable.pe.pe import PE STUB = b"\x0e\x1f\xba\x0e\x00\xb4\t\xcd!\xb8\x01L\xcd!This program is made with dissect.pe <3 kusjesvanSRT <3.\x0d\x0d\x0a$\x00\x00" # noqa: E501 diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/helpers/exports.py index 3c27460..c08aa90 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/helpers/exports.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers.utils import DictManager +from dissect.executable.utils import DictManager if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/helpers/imports.py index 7a12540..e229c6f 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/helpers/imports.py @@ -3,9 +3,9 @@ from io import BytesIO from typing import TYPE_CHECKING, BinaryIO +from dissect.executable import utils from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers import utils -from dissect.executable.pe.helpers.utils import DictManager, create_struct +from dissect.executable.utils import DictManager, create_struct if TYPE_CHECKING: from collections.abc import Iterator diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/helpers/patcher.py index 125e316..fa9599c 100644 --- a/dissect/executable/pe/helpers/patcher.py +++ b/dissect/executable/pe/helpers/patcher.py @@ -3,8 +3,8 @@ from io import BytesIO from typing import TYPE_CHECKING +from dissect.executable import utils from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers import utils if TYPE_CHECKING: from dissect.executable import PE diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/helpers/relocations.py index dea2f4a..a7dcf83 100644 --- a/dissect/executable/pe/helpers/relocations.py +++ b/dissect/executable/pe/helpers/relocations.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers.utils import ListManager +from dissect.executable.utils import ListManager if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/helpers/resources.py index 0917f42..1ae6427 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/helpers/resources.py @@ -10,7 +10,7 @@ from dissect.executable.exception import ResourceException from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers.utils import DictManager +from dissect.executable.utils import DictManager if TYPE_CHECKING: from collections.abc import Iterator diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/helpers/sections.py index 996e7c6..d587f12 100644 --- a/dissect/executable/pe/helpers/sections.py +++ b/dissect/executable/pe/helpers/sections.py @@ -5,9 +5,9 @@ from itertools import chain from typing import TYPE_CHECKING +from dissect.executable import utils from dissect.executable.exception import BuildSectionException from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers import utils if TYPE_CHECKING: from collections.abc import Iterable diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/helpers/tls.py index 08be933..8562c18 100644 --- a/dissect/executable/pe/helpers/tls.py +++ b/dissect/executable/pe/helpers/tls.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING from dissect.executable.pe.c_pe import c_pe -from dissect.executable.pe.helpers.utils import ListManager +from dissect.executable.utils import ListManager if TYPE_CHECKING: from dissect.executable.pe.helpers.sections import PESection diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index 749d2dc..b72c9b5 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -5,6 +5,7 @@ from pathlib import Path from typing import TYPE_CHECKING, BinaryIO +from dissect.executable import utils from dissect.executable.exception import ( InvalidAddress, InvalidArchitecture, @@ -20,7 +21,6 @@ resources, sections, tls, - utils, ) if TYPE_CHECKING: diff --git a/dissect/executable/pe/helpers/utils.py b/dissect/executable/utils.py similarity index 100% rename from dissect/executable/pe/helpers/utils.py rename to dissect/executable/utils.py From d617e049d0eafaa789c680bcbd86d427574ebea2 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Wed, 9 Apr 2025 09:17:25 +0000 Subject: [PATCH 41/43] Move builder.py out of helpers --- dissect/executable/pe/{helpers => }/builder.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename dissect/executable/pe/{helpers => }/builder.py (100%) diff --git a/dissect/executable/pe/helpers/builder.py b/dissect/executable/pe/builder.py similarity index 100% rename from dissect/executable/pe/helpers/builder.py rename to dissect/executable/pe/builder.py From c72bbbca72816aca21423aadb5c7c3102b55cfd8 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Wed, 9 Apr 2025 09:17:49 +0000 Subject: [PATCH 42/43] Move patcher.py out of helpers --- dissect/executable/pe/{helpers => }/patcher.py | 0 dissect/executable/pe/pe.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename dissect/executable/pe/{helpers => }/patcher.py (100%) diff --git a/dissect/executable/pe/helpers/patcher.py b/dissect/executable/pe/patcher.py similarity index 100% rename from dissect/executable/pe/helpers/patcher.py rename to dissect/executable/pe/patcher.py diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index b72c9b5..4c4122a 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -12,11 +12,11 @@ InvalidPE, InvalidVA, ) +from dissect.executable.pe import patcher from dissect.executable.pe.c_pe import c_cv_info, c_pe from dissect.executable.pe.helpers import ( exports, imports, - patcher, relocations, resources, sections, From aff11da6865176ede99bcb49936eab33b63d7ea5 Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Wed, 9 Apr 2025 09:20:20 +0000 Subject: [PATCH 43/43] Rename helpers to sections --- dissect/executable/pe/__init__.py | 14 +++++++------- dissect/executable/pe/pe.py | 2 +- .../pe/{helpers => sections}/__init__.py | 0 .../executable/pe/{helpers => sections}/exports.py | 2 +- .../executable/pe/{helpers => sections}/imports.py | 2 +- .../pe/{helpers => sections}/relocations.py | 2 +- .../pe/{helpers => sections}/resources.py | 2 +- .../pe/{helpers => sections}/sections.py | 0 dissect/executable/pe/{helpers => sections}/tls.py | 2 +- 9 files changed, 13 insertions(+), 13 deletions(-) rename dissect/executable/pe/{helpers => sections}/__init__.py (100%) rename dissect/executable/pe/{helpers => sections}/exports.py (97%) rename dissect/executable/pe/{helpers => sections}/imports.py (99%) rename dissect/executable/pe/{helpers => sections}/relocations.py (96%) rename dissect/executable/pe/{helpers => sections}/resources.py (99%) rename dissect/executable/pe/{helpers => sections}/sections.py (100%) rename dissect/executable/pe/{helpers => sections}/tls.py (98%) diff --git a/dissect/executable/pe/__init__.py b/dissect/executable/pe/__init__.py index 7667b2f..66eb098 100644 --- a/dissect/executable/pe/__init__.py +++ b/dissect/executable/pe/__init__.py @@ -1,14 +1,14 @@ -from dissect.executable.pe.helpers.builder import Builder -from dissect.executable.pe.helpers.exports import ExportFunction, ExportManager -from dissect.executable.pe.helpers.imports import ( +from dissect.executable.pe.builder import Builder +from dissect.executable.pe.patcher import Patcher +from dissect.executable.pe.pe import PE +from dissect.executable.pe.sections.exports import ExportFunction, ExportManager +from dissect.executable.pe.sections.imports import ( ImportFunction, ImportManager, ImportModule, ) -from dissect.executable.pe.helpers.patcher import Patcher -from dissect.executable.pe.helpers.resources import Resource, ResourceManager -from dissect.executable.pe.helpers.sections import PESection -from dissect.executable.pe.pe import PE +from dissect.executable.pe.sections.resources import Resource, ResourceManager +from dissect.executable.pe.sections.sections import PESection __all__ = [ "PE", diff --git a/dissect/executable/pe/pe.py b/dissect/executable/pe/pe.py index 4c4122a..be70c4d 100644 --- a/dissect/executable/pe/pe.py +++ b/dissect/executable/pe/pe.py @@ -14,7 +14,7 @@ ) from dissect.executable.pe import patcher from dissect.executable.pe.c_pe import c_cv_info, c_pe -from dissect.executable.pe.helpers import ( +from dissect.executable.pe.sections import ( exports, imports, relocations, diff --git a/dissect/executable/pe/helpers/__init__.py b/dissect/executable/pe/sections/__init__.py similarity index 100% rename from dissect/executable/pe/helpers/__init__.py rename to dissect/executable/pe/sections/__init__.py diff --git a/dissect/executable/pe/helpers/exports.py b/dissect/executable/pe/sections/exports.py similarity index 97% rename from dissect/executable/pe/helpers/exports.py rename to dissect/executable/pe/sections/exports.py index c08aa90..68e436e 100644 --- a/dissect/executable/pe/helpers/exports.py +++ b/dissect/executable/pe/sections/exports.py @@ -8,8 +8,8 @@ from dissect.executable.utils import DictManager if TYPE_CHECKING: - from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE + from dissect.executable.pe.sections.sections import PESection @dataclass diff --git a/dissect/executable/pe/helpers/imports.py b/dissect/executable/pe/sections/imports.py similarity index 99% rename from dissect/executable/pe/helpers/imports.py rename to dissect/executable/pe/sections/imports.py index e229c6f..ba95440 100644 --- a/dissect/executable/pe/helpers/imports.py +++ b/dissect/executable/pe/sections/imports.py @@ -11,8 +11,8 @@ from collections.abc import Iterator from struct import Struct - from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE + from dissect.executable.pe.sections.sections import PESection class ImportModule: diff --git a/dissect/executable/pe/helpers/relocations.py b/dissect/executable/pe/sections/relocations.py similarity index 96% rename from dissect/executable/pe/helpers/relocations.py rename to dissect/executable/pe/sections/relocations.py index a7dcf83..03489d6 100644 --- a/dissect/executable/pe/helpers/relocations.py +++ b/dissect/executable/pe/sections/relocations.py @@ -8,8 +8,8 @@ from dissect.executable.utils import ListManager if TYPE_CHECKING: - from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE + from dissect.executable.pe.sections.sections import PESection @dataclass diff --git a/dissect/executable/pe/helpers/resources.py b/dissect/executable/pe/sections/resources.py similarity index 99% rename from dissect/executable/pe/helpers/resources.py rename to dissect/executable/pe/sections/resources.py index 1ae6427..fa0e1dc 100644 --- a/dissect/executable/pe/helpers/resources.py +++ b/dissect/executable/pe/sections/resources.py @@ -16,8 +16,8 @@ from collections.abc import Iterator from typing import BinaryIO, Callable - from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE + from dissect.executable.pe.sections.sections import PESection @dataclass diff --git a/dissect/executable/pe/helpers/sections.py b/dissect/executable/pe/sections/sections.py similarity index 100% rename from dissect/executable/pe/helpers/sections.py rename to dissect/executable/pe/sections/sections.py diff --git a/dissect/executable/pe/helpers/tls.py b/dissect/executable/pe/sections/tls.py similarity index 98% rename from dissect/executable/pe/helpers/tls.py rename to dissect/executable/pe/sections/tls.py index 8562c18..84b37ad 100644 --- a/dissect/executable/pe/helpers/tls.py +++ b/dissect/executable/pe/sections/tls.py @@ -7,8 +7,8 @@ from dissect.executable.utils import ListManager if TYPE_CHECKING: - from dissect.executable.pe.helpers.sections import PESection from dissect.executable.pe.pe import PE + from dissect.executable.pe.sections.sections import PESection class TLSManager(ListManager[int]):