From c1c6f860f49cd4c4fb6e166b7e61d02a15b182ea 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 1/6] 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 0ba802b95bc999f9307098708e5f3f5e18f0d6fc 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 2/6] 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 6549e83bc0e15e12ec0604f448665fc9d72886b5 Mon Sep 17 00:00:00 2001 From: sud0woodo <40278342+sud0woodo@users.noreply.github.com> Date: Sat, 9 Mar 2024 16:48:59 +0200 Subject: [PATCH 3/6] Add empty __init__.py in tests/ --- tests/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/__init__.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 From e42030efb3f027ce3cf5dd0dd45a99e2ac695690 Mon Sep 17 00:00:00 2001 From: sud0woodo <40278342+sud0woodo@users.noreply.github.com> Date: Sat, 9 Mar 2024 16:55:54 +0200 Subject: [PATCH 4/6] Fix util import in test_dump.py --- tests/test_dump.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_dump.py b/tests/test_dump.py index 8cf0be7..65ef531 100644 --- a/tests/test_dump.py +++ b/tests/test_dump.py @@ -2,7 +2,7 @@ from pathlib import Path import pytest -from util import data_file +from .util import data_file from dissect.executable import ELF From 137d17b1f3cd8cdbefc46b1927699520da3e4109 Mon Sep 17 00:00:00 2001 From: sud0woodo <40278342+sud0woodo@users.noreply.github.com> Date: Sat, 9 Mar 2024 17:08:28 +0200 Subject: [PATCH 5/6] Fix linting --- tests/test_dump.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_dump.py b/tests/test_dump.py index 65ef531..d611596 100644 --- a/tests/test_dump.py +++ b/tests/test_dump.py @@ -2,10 +2,11 @@ from pathlib import Path import pytest -from .util import data_file from dissect.executable import ELF +from .util import data_file + @pytest.mark.parametrize( "file_name", From b48a82c29fa63b705e8e8c1355661367e0e1b8ca 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 6/6] 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 | 41 ++++-- 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_elf.py | 6 +- tests/test_pe.py | 47 +++++-- tests/test_pe_builder.py | 56 +++++--- tests/test_pe_modifications.py | 34 +++-- tests/test_section.py | 28 ++-- tests/test_segment.py | 6 +- tests/test_segment_table.py | 10 +- 19 files changed, 569 insertions(+), 322 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 ff043a1..f79ed58 100644 --- a/dissect/executable/elf/elf.py +++ b/dissect/executable/elf/elf.py @@ -42,7 +42,9 @@ def __init__(self, fh: BinaryIO): self.header = self.c_elf.Ehdr(fh) self.segments = SegmentTable.from_elf(self) self.section_table = SectionTable.from_elf(self) - self.symbol_tables: list[SymbolTable] = self.section_table.by_type([SHT.SYMTAB, SHT.DYNSYM]) + self.symbol_tables: list[SymbolTable] = self.section_table.by_type( + [SHT.SYMTAB, SHT.DYNSYM] + ) def __repr__(self) -> str: return str(self.header) @@ -98,7 +100,9 @@ def find(self, condition: Callable, **kwargs) -> list[T]: class Section: - def __init__(self, fh: BinaryIO, idx: Optional[int] = None, c_elf: cstruct = c_elf_64): + def __init__( + self, fh: BinaryIO, idx: Optional[int] = None, c_elf: cstruct = c_elf_64 + ): self.fh = fh self.idx = idx @@ -224,7 +228,9 @@ def dump_data(self) -> list[tuple[int, bytes]]: class Segment: - def __init__(self, fh: BinaryIO, idx: Optional[int] = None, c_elf: cstruct = c_elf_64): + def __init__( + self, fh: BinaryIO, idx: Optional[int] = None, c_elf: cstruct = c_elf_64 + ): self.fh = fh self.idx = idx self.c_elf = c_elf @@ -246,7 +252,9 @@ def __repr__(self) -> str: return repr(self.header) @classmethod - def from_segment_table(cls, table: SegmentTable, idx: Optional[int] = None) -> Segment: + def from_segment_table( + cls, table: SegmentTable, idx: Optional[int] = None + ) -> Segment: fh = table.fh return cls(fh, idx, table.c_elf) @@ -277,7 +285,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 @@ -297,7 +312,9 @@ def from_elf(cls, elf: ELF) -> SegmentTable: offset = header.e_phoff entries = header.e_phnum size = header.e_phentsize - return cls(fh=elf.fh, offset=offset, entries=entries, size=size, c_elf=elf.c_elf) + return cls( + fh=elf.fh, offset=offset, entries=entries, size=size, c_elf=elf.c_elf + ) def related_segments(self, section: Section) -> list[Segment]: return self.find(lambda x: x.is_related(section)) @@ -318,7 +335,9 @@ def dump_table(self) -> tuple[int, bytearray]: class StringTable(Section): - def __init__(self, fh: BinaryIO, idx: Optional[int] = None, c_elf: cstruct = c_elf_64): + def __init__( + self, fh: BinaryIO, idx: Optional[int] = None, c_elf: cstruct = c_elf_64 + ): super().__init__(fh, idx, c_elf) self._get_string = lru_cache(256)(self._get_string) @@ -333,7 +352,9 @@ def _get_string(self, index: int) -> str: class Symbol: - def __init__(self, fh: BinaryIO, idx: Optional[int] = None, c_elf: cstruct = c_elf_64): + def __init__( + self, fh: BinaryIO, idx: Optional[int] = None, c_elf: cstruct = c_elf_64 + ): self.symbol = c_elf.Sym(fh) self.idx = idx self.c_elf = c_elf @@ -388,7 +409,9 @@ def value_based_on_shndx(self, table: SectionTable) -> int: class SymbolTable(Section, Table[Symbol]): - def __init__(self, fh: BinaryIO, idx: Optional[int] = None, c_elf: cstruct = c_elf_64): + def __init__( + self, fh: BinaryIO, idx: Optional[int] = None, c_elf: cstruct = c_elf_64 + ): # Initializes Section info Section.__init__(self, fh, idx, c_elf) count = self.size // self.entry_size 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_elf.py b/tests/test_elf.py index 67da65e..bd9853d 100644 --- a/tests/test_elf.py +++ b/tests/test_elf.py @@ -6,10 +6,10 @@ from dissect.executable.exception import InvalidSignatureError -def test_elf_invalid_signature(): +def test_elf_invalid_signature() -> None: with pytest.raises(InvalidSignatureError): ELF(BytesIO(b"\x20ELF" + b"\x00" * 0x40)) -def test_elf_valid_signature(): - ELF(BytesIO(b"\x7FELF" + b"\x00" * 0x40)) +def test_elf_valid_signature() -> None: + ELF(BytesIO(b"\x7fELF" + b"\x00" * 0x40)) 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 6a1e5f7..db4ddef 100644 --- a/tests/test_section.py +++ b/tests/test_section.py @@ -14,7 +14,7 @@ @pytest.fixture -def section_table(entries: int): +def section_table(entries: int) -> SectionTable: """Creates a SectionTable without a StringTable attached to it.""" elf = Mock() elf.header.e_shnum = entries @@ -23,7 +23,11 @@ def section_table(entries: int): 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 @@ -33,13 +37,13 @@ def mock_section_table(section_data: bytes) -> Mock: @pytest.mark.parametrize("entries", [0]) -def test_section_unknown_index(section_table: SectionTable): +def test_section_unknown_index(section_table: SectionTable) -> None: with pytest.raises(IndexError): assert section_table[1] @pytest.mark.parametrize("entries", [20]) -def test_section_selector(section_table: SectionTable, entries: int): +def test_section_selector(section_table: SectionTable, entries: int) -> None: with patch.object(SectionTable, "_create_item") as mocked_section: assert section_table.items == [None] * entries assert section_table[0] == mocked_section.return_value @@ -47,7 +51,7 @@ def test_section_selector(section_table: SectionTable, entries: int): assert list(section_table) == [mocked_section.return_value] * entries -def test_string_table(): +def test_string_table() -> None: STRING_TABLE = b"\x00hello\x00world\x00" mocked_table = mock_section_table(STRING_TABLE) @@ -59,14 +63,14 @@ def test_string_table(): assert string_table[7] == "world" -def test_symboltable(): +def test_symboltable() -> None: mocked_table = mock_section_table(b"hello") with patch.object(Symbol, "from_symbol_table") as mocked_symbol: symbol_table = SymbolTable.from_section_table(mocked_table, 0) assert symbol_table[0] == mocked_symbol.return_value -def test_table_symbol_creation(): +def test_table_symbol_creation() -> None: symbol_bytes = c_elf_64.Sym().dumps() mocked_table = mock_section_table(symbol_bytes) @@ -84,7 +88,7 @@ def test_table_symbol_creation(): assert symbol.name == ".hello" -def test_symboltable_filter(): +def test_symboltable_filter() -> None: symbol_bytes = c_elf_64.Sym(st_info=0x16).dumps() mocked_table = mock_section_table(symbol_bytes) @@ -106,8 +110,10 @@ def test_symboltable_filter(): ("COMMON", 100, 100), ], ) -def test_symbol_value(section_index, value, expected_output): - symbol_bytes = c_elf_64.Sym(st_value=value, st_shndx=c_elf_64.SHN[section_index].value).dumps() +def test_symbol_value(section_index, value, expected_output) -> None: + symbol_bytes = c_elf_64.Sym( + st_value=value, st_shndx=c_elf_64.SHN[section_index].value + ).dumps() symbol = Symbol(BytesIO(symbol_bytes), 0, c_elf_64) assert symbol.value == expected_output @@ -121,7 +127,7 @@ def test_symbol_value(section_index, value, expected_output): (20, 100, 100), ], ) -def test_symbol_value_from_shndex(section_index, table_offset, expected_output): +def test_symbol_value_from_shndex(section_index, table_offset, expected_output) -> None: symbol_bytes = c_elf_64.Sym(st_shndx=section_index).dumps() symbol = Symbol(BytesIO(symbol_bytes), 0, c_elf_64) diff --git a/tests/test_segment.py b/tests/test_segment.py index 0118627..606ecfd 100644 --- a/tests/test_segment.py +++ b/tests/test_segment.py @@ -4,12 +4,14 @@ 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) -def test_segment(): +def test_segment() -> None: orig_data = b"hello_world" segment = create_segment(orig_data) assert segment.offset == len(c_elf_64.Phdr) diff --git a/tests/test_segment_table.py b/tests/test_segment_table.py index 92c3e0a..a2de52c 100644 --- a/tests/test_segment_table.py +++ b/tests/test_segment_table.py @@ -23,7 +23,7 @@ def test_segment_table_unknown_index(segment_table: SegmentTable): @pytest.mark.parametrize("entries", [1]) -def test_segment_table_known(segment_table: SegmentTable): +def test_segment_table_known(segment_table: SegmentTable) -> None: with patch("dissect.executable.elf.elf.Segment") as mocked_segment: assert segment_table[0] == mocked_segment.from_segment_table.return_value @@ -32,7 +32,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) @@ -40,7 +42,7 @@ def create_segment_table(amount: int, random_data: bytes) -> SegmentTable: return SegmentTable(segment_data, 0, amount, len(c_elf_64.Phdr)) -def test_dump_data(): +def test_dump_data() -> None: segment_table = create_segment_table(2, b"hello_world") segment_table[0].patch(b"new_data") data = bytes() @@ -49,7 +51,7 @@ def test_dump_data(): assert len(data) == len(b"hello_world" + b"new_data") -def test_dump_table(): +def test_dump_table() -> None: segment_table = create_segment_table(2, b"hello_world") segment_table[0].patch(b"new_data") assert len(segment_table.dump_table()[1]) == 2 * len(c_elf_64.Phdr)