pe_hdrs_helper.h

Go to the documentation of this file.

 
#pragma once
 
#include <windows.h>
#include "buffer_util.h"
 
#ifndef PAGE_SIZE
#define PAGE_SIZE 0x1000
#endif
 
namespace peconv {
    const ULONGLONG MAX_HEADER_SIZE = PAGE_SIZE;
 
    template <typename INT_TYPE>
    INT_TYPE round_up_to_unit(const INT_TYPE size, const INT_TYPE unit)
    {
        if (unit == 0) return size;
 
        INT_TYPE rem = size % unit;
        if (rem == 0) return size;
 
        return size + (unit - rem);
    }
 
    DWORD get_image_size(IN const BYTE *payload);
 
    bool update_image_size(IN OUT BYTE* payload, IN DWORD new_img_size);
 
    WORD get_nt_hdr_architecture(IN const BYTE *pe_buffer);
 
    bool is64bit(IN const BYTE *pe_buffer);
 
    BYTE* get_nt_hdrs(
        IN const BYTE *pe_buffer, 
        IN OPTIONAL size_t buffer_size=0 //if buffer_size=0 means size unknown
    );
 
    IMAGE_NT_HEADERS32* get_nt_hdrs32(IN const BYTE *pe_buffer);
 
    IMAGE_NT_HEADERS64* get_nt_hdrs64(IN const BYTE *pe_buffer);
 
    LPVOID get_optional_hdr(IN const BYTE* payload, IN const size_t buffer_size);
 
    const IMAGE_FILE_HEADER* get_file_hdr(
        IN const BYTE* payload,
        IN const size_t buffer_size
    );
 
    DWORD get_hdrs_size(IN const BYTE *pe_buffer);
 
    IMAGE_DATA_DIRECTORY* get_directory_entry(IN const BYTE* pe_buffer, IN DWORD dir_id, IN bool allow_empty = false);
 
    template <typename IMAGE_TYPE_DIRECTORY>
    IMAGE_TYPE_DIRECTORY* get_type_directory(IN HMODULE modulePtr, IN DWORD dir_id)
    {
        IMAGE_DATA_DIRECTORY *my_dir = peconv::get_directory_entry((const BYTE*)modulePtr, dir_id);
        if (!my_dir) return nullptr;
 
        DWORD dir_addr = my_dir->VirtualAddress;
        if (dir_addr == 0) return nullptr;
 
        return (IMAGE_TYPE_DIRECTORY*)(dir_addr + (ULONG_PTR)modulePtr);
    }
 
    IMAGE_EXPORT_DIRECTORY* get_export_directory(IN HMODULE modulePtr);
 
    // Fetch Image Base from Optional Header.
    ULONGLONG get_image_base(IN const BYTE *pe_buffer);
 
    bool update_image_base(IN OUT BYTE* payload, IN ULONGLONG destImageBase);
 
    DWORD get_entry_point_rva(IN const BYTE *pe_buffer);
 
    bool update_entry_point_rva(IN OUT BYTE *pe_buffer, IN DWORD ep);
 
    size_t get_sections_count(
        IN const BYTE* buffer,
        IN const size_t buffer_size
    );
 
    bool is_valid_sections_hdr_offset(IN const BYTE* buffer, IN const size_t buffer_size);
 
    PIMAGE_SECTION_HEADER get_section_hdr(
        IN const BYTE* pe_buffer,
        IN const size_t buffer_size,
        IN size_t section_num
    );
 
    WORD get_file_characteristics(IN const BYTE* payload);
 
    bool is_module_dll(IN const BYTE* payload);
 
    bool is_dot_net(BYTE *pe_buffer, size_t pe_buffer_size);
 
    WORD get_dll_characteristics(IN const BYTE* payload);
 
    bool set_subsystem(IN OUT BYTE* payload, IN WORD subsystem);
 
    WORD get_subsystem(IN const BYTE* payload);
 
    bool has_relocations(IN const BYTE *pe_buffer);
 
    IMAGE_COR20_HEADER* get_dotnet_hdr(
        IN const BYTE* pe_buffer,
        IN size_t const buffer_size,
        IN const IMAGE_DATA_DIRECTORY* dotNetDir
    );
 
    DWORD get_sec_alignment(IN const BYTE* modulePtr, IN bool is_raw);
 
    bool set_sec_alignment(IN OUT BYTE* pe_buffer, IN bool is_raw, IN DWORD new_alignment);
 
    DWORD get_virtual_sec_size(
        IN const BYTE* pe_hdr,
        IN const PIMAGE_SECTION_HEADER sec_hdr,
        IN bool rounded //if set, it rounds it up to the Virtual Alignment
    );
 
    PIMAGE_SECTION_HEADER get_last_section(IN const PBYTE pe_buffer, IN size_t pe_size, IN bool is_raw);
 
    DWORD calc_pe_size(
        IN const PBYTE pe_buffer,
        IN size_t pe_size,
        IN bool is_raw
    );
 
    bool is_valid_sections_alignment(IN const BYTE* buffer, IN const SIZE_T buffer_size, IN bool is_raw);
 
}; // namespace peconv