pe_reconstructor.cpp

Go to the documentation of this file.

#include "pe_reconstructor.h"
 
#include "../utils/workingset_enum.h"
 
#include <fstream>
 
namespace pesieve {
    inline bool shift_artefacts(PeArtefacts& artefacts, size_t shift_size)
    {
        artefacts.ntFileHdrsOffset += shift_size;
        artefacts.secHdrsOffset += shift_size;
        return true;
    }
}; //namespace pesieve
 
 
//WARNING: this function shifts also offsets saved in the artefacts
size_t pesieve::PeReconstructor::shiftPeHeader()
{
    BYTE *vBuf = this->peBuffer.vBuf;
    const size_t vBufSize = this->peBuffer.vBufSize;
    if (vBuf == nullptr) return 0;
 
    if (!this->artefacts.hasNtHdrs()) return 0;
 
    const size_t dos_pe_size = sizeof(IMAGE_DOS_HEADER) + sizeof(IMAGE_NT_SIGNATURE);
    const size_t nt_offset = this->artefacts.dropPeBase(this->artefacts.ntFileHdrsOffset);
    if (nt_offset == INVALID_OFFSET || nt_offset >= dos_pe_size) {
        return 0;
    }
    //TODO: shift the header
    if (!this->artefacts.hasSectionHdrs()) return 0; //cannot proceed
 
    size_t shift_size = dos_pe_size - nt_offset;
    size_t hdrs_end = this->artefacts.secHdrsOffset + (this->artefacts.secCount + 1)* sizeof(IMAGE_SECTION_HEADER);
    if (!peconv::is_padding(vBuf + hdrs_end, shift_size, 0)) {
        return 0; // no empty space, cannot proceed
    }
    size_t hdrs_size = this->artefacts.dropPeBase(hdrs_end);
    BYTE *new_nt_ptr = vBuf + this->artefacts.peBaseOffset + shift_size;
    if (!peconv::validate_ptr(vBuf, vBufSize, new_nt_ptr, hdrs_size)) {
        return 0;
    }
    if (nt_offset < sizeof(IMAGE_NT_SIGNATURE)) {
        return 0;
    }
    const size_t pe_offset = nt_offset - sizeof(IMAGE_NT_SIGNATURE);
    IMAGE_DOS_HEADER dos_template = { 0 };
    dos_template.e_magic = IMAGE_DOS_SIGNATURE;
    dos_template.e_lfanew = LONG(pe_offset + shift_size);
 
    //check mz signature:
    BYTE *mz_ptr = vBuf + this->artefacts.peBaseOffset;
    if (!peconv::validate_ptr(vBuf, vBufSize, mz_ptr, sizeof(IMAGE_DOS_HEADER))) {
        return 0;
    }
    //check PE signature:
    DWORD* pe_ptr = (DWORD*)(vBuf + this->artefacts.peBaseOffset + dos_template.e_lfanew);
    if (!peconv::validate_ptr(vBuf, vBufSize, pe_ptr, sizeof(DWORD))) {
        return 0;
    }
    //all checks passed, do the actual headers shift:
    memmove(new_nt_ptr, (vBuf + this->artefacts.peBaseOffset), hdrs_size);
 
    //write the DOS header:
    memcpy(mz_ptr, &dos_template, sizeof(IMAGE_DOS_HEADER));
 
    //write the PE signature:
    *pe_ptr = IMAGE_NT_SIGNATURE;
 
    shift_artefacts(this->artefacts, shift_size);
    return shift_size;
}
 
bool pesieve::PeReconstructor::reconstruct()
{
    this->artefacts = origArtefacts;
 
    ULONGLONG moduleBase = artefacts.regionStart + artefacts.peBaseOffset;
    if (!peBuffer.readRemote(moduleBase, artefacts.calculatedImgSize)) {
        return false;
    }
    size_t shift_size = shiftPeHeader();
    if (shift_size) {
        std::cout << "[!] The PE header was shifted by: " << std::hex << shift_size << std::endl;
    }
    bool is_pe_hdr = false;
    if (this->artefacts.hasNtHdrs() && reconstructFileHdr()) {
        is_pe_hdr = reconstructPeHdr();
    }
    if (!is_pe_hdr) {
        return false;
    }
 
    //do not modify section headers if the PE is in raw format, or no unmapping requested
    if (!peconv::is_pe_raw(peBuffer.vBuf, peBuffer.vBufSize)) {
 
        if (!fixSectionsVirtualSize(peBuffer.processHndl) || !fixSectionsCharacteristics(peBuffer.processHndl)) {
            return false;
        }
        const ULONGLONG base_candidate = peconv::find_base_candidate(peBuffer.vBuf, peBuffer.vBufSize);
        if (base_candidate) {
            peconv::update_image_base(peBuffer.vBuf, (ULONGLONG)base_candidate);
        }
    }
    return peBuffer.isValidPe();
}
 
bool pesieve::PeReconstructor::fixSectionsVirtualSize(HANDLE processHandle)
{
    BYTE *vBuf = this->peBuffer.vBuf;
    const size_t vBufSize = this->peBuffer.vBufSize;
    if (!vBuf) return false;
 
    if (!this->artefacts.hasSectionHdrs()) {
        return false;
    }
    ULONGLONG sec_offset = this->artefacts.dropPeBase(this->artefacts.secHdrsOffset);
    BYTE *hdr_ptr = (sec_offset + vBuf);
 
    size_t max_sec_size = 0;
 
    IMAGE_SECTION_HEADER* prev_sec = nullptr;
    IMAGE_SECTION_HEADER* curr_sec = (IMAGE_SECTION_HEADER*)(hdr_ptr);
 
    const ULONGLONG pe_img_base = (ULONGLONG)artefacts.peImageBase();
 
    const size_t hdr_sec_count = peconv::get_sections_count(vBuf, vBufSize);
    const size_t sec_count = artefacts.secCount > hdr_sec_count ? artefacts.secCount : hdr_sec_count;
 
    size_t i = 0;
    for (i = 0; i < sec_count; i++, curr_sec++) {
        if (!peconv::validate_ptr(vBuf, vBufSize, curr_sec, sizeof(IMAGE_SECTION_HEADER))) {
            // buffer finished
            break;
        }
        const DWORD sec_rva = curr_sec->VirtualAddress;
        const DWORD sec_size = curr_sec->Misc.VirtualSize;
 
        const ULONGLONG sec_va = pe_img_base + sec_rva;
        size_t real_sec_size = peconv::fetch_region_size(processHandle, (PBYTE)sec_va);
 
        // if the RVA is out of scope, and the calculated size is 0:
        if (!peconv::validate_ptr(vBuf, vBufSize, vBuf + sec_rva, sizeof(BYTE)) && !real_sec_size) {
#ifdef _DEBUG
            std::cout << i << "# Invalid section found: " << std::hex
                << sec_rva << " of size: " << sec_size << std::endl;
#endif
            break;
        }
 
        if (sec_size > real_sec_size) {
            curr_sec->Misc.VirtualSize = DWORD(real_sec_size);
#ifdef _DEBUG
            std::cout << i << "# Fixed section " << std::hex << sec_rva  << " size: " << std::hex
                << sec_size << " vs real: " << real_sec_size << std::endl;
#endif
        }
 
        max_sec_size = (real_sec_size > max_sec_size) ? real_sec_size : max_sec_size;
 
        if (prev_sec && curr_sec->Misc.VirtualSize > 0) {
            const ULONGLONG prev_sec_end = static_cast<ULONGLONG>(prev_sec->VirtualAddress) + static_cast<ULONGLONG>(prev_sec->Misc.VirtualSize);
            if (prev_sec_end > curr_sec->VirtualAddress) {
                if (curr_sec->VirtualAddress > prev_sec->VirtualAddress) {
                    DWORD diff = curr_sec->VirtualAddress - prev_sec->VirtualAddress;
                    prev_sec->Misc.VirtualSize = diff;
#ifdef _DEBUG
                    std::cout << "Trimmed section: " << std::dec << i << std::endl;
#endif
                }
            }
        }
        if (curr_sec->Misc.VirtualSize > 0) {
            prev_sec = curr_sec;
        }
    }
    
    IMAGE_FILE_HEADER* file_hdr = const_cast<IMAGE_FILE_HEADER*>(peconv::get_file_hdr(vBuf, vBufSize));
    if (file_hdr && (i > 0)) {
        // set the actual number of valid sections:
        file_hdr->NumberOfSections = MASK_TO_WORD(i);
    }
 
    if (max_sec_size == 0) {
        return false;
    }
    return true;
}
 
bool pesieve::PeReconstructor::fixSectionsCharacteristics(HANDLE processHandle)
{
    BYTE *vBuf = this->peBuffer.vBuf;
    const size_t vBufSize = this->peBuffer.vBufSize;
    if (!vBuf) return false;
 
    if (!this->artefacts.hasSectionHdrs()) {
        return false;
    }
 
    ULONGLONG sec_offset = this->artefacts.dropPeBase(this->artefacts.secHdrsOffset);
    const BYTE *hdr_ptr = (sec_offset + vBuf);
    IMAGE_SECTION_HEADER* curr_sec = (IMAGE_SECTION_HEADER*)(hdr_ptr);
 
    const DWORD sec_all_flags = IMAGE_SCN_TYPE_NO_PAD
        | IMAGE_SCN_CNT_CODE | IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_CNT_UNINITIALIZED_DATA
        | IMAGE_SCN_LNK_NRELOC_OVFL | IMAGE_SCN_MEM_DISCARDABLE | IMAGE_SCN_MEM_NOT_CACHED
        | IMAGE_SCN_MEM_NOT_PAGED | IMAGE_SCN_MEM_SHARED | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ
        | IMAGE_SCN_MEM_WRITE
        | IMAGE_SCN_NO_DEFER_SPEC_EXC | IMAGE_SCN_GPREL;
 
    for (size_t i = 0; i < artefacts.secCount; i++, curr_sec++) {
        if (!is_valid_section(vBuf, vBufSize, (BYTE*)curr_sec, 0)) {
            break;
        }
        //leave only the flags that are valid
        const DWORD charact = curr_sec->Characteristics;
        curr_sec->Characteristics = charact & sec_all_flags;
#ifdef DEBUG
        if (charact != curr_sec->Characteristics) {
            std::cout << "Section characteristics overwriten\n";
        }
#endif
    }
    return true;
}
 
bool pesieve::PeReconstructor::reconstructFileHdr()
{
    BYTE *vBuf = this->peBuffer.vBuf;
    const size_t vBufSize = this->peBuffer.vBufSize;
    if (!vBuf) return false;
 
    if (!this->artefacts.hasNtHdrs()) {
        return false;
    }
    size_t nt_offset = this->artefacts.dropPeBase(this->artefacts.ntFileHdrsOffset);
    BYTE* nt_ptr = (BYTE*)((ULONGLONG)vBuf + nt_offset);
    if (is_valid_file_hdr(vBuf, vBufSize, nt_ptr, 0)) {
        return true;
    }
    IMAGE_FILE_HEADER* hdr_candidate = (IMAGE_FILE_HEADER*)nt_ptr;
    if (!peconv::validate_ptr(vBuf, vBufSize, hdr_candidate, sizeof(IMAGE_FILE_HEADER))) {
        // probably buffer finished
        return false;
    }
 
    size_t opt_hdr_size = 0;
    if (artefacts.is64bit) {
        hdr_candidate->Machine = IMAGE_FILE_MACHINE_AMD64;
        opt_hdr_size = sizeof(IMAGE_OPTIONAL_HEADER64);
    }
    else {
        hdr_candidate->Machine = IMAGE_FILE_MACHINE_I386;
        opt_hdr_size = sizeof(IMAGE_OPTIONAL_HEADER32);
    }
    if (this->artefacts.secHdrsOffset) {
        const size_t sec_offset = this->artefacts.dropPeBase(this->artefacts.secHdrsOffset);
        size_t calc_offset = sec_offset - (nt_offset + sizeof(IMAGE_FILE_HEADER));
 
        if (calc_offset != opt_hdr_size) {
            std::cout << "[WARNING] Calculated sections header offset is different than the saved one!\n";
        }
        hdr_candidate->NumberOfSections = WORD(this->artefacts.secCount);
        hdr_candidate->SizeOfOptionalHeader = WORD(calc_offset);
    }
    hdr_candidate->NumberOfSymbols = 0;
    hdr_candidate->PointerToSymbolTable = 0;
    return true;
}
 
bool pesieve::PeReconstructor::reconstructPeHdr()
{
    BYTE *vBuf = this->peBuffer.vBuf;
    const size_t vBufSize = this->peBuffer.vBufSize;
    if (!vBuf) return false;
 
    if (!this->artefacts.hasNtHdrs()) {
        return false;
    }
    ULONGLONG nt_offset = this->artefacts.dropPeBase(this->artefacts.ntFileHdrsOffset);
    BYTE* nt_ptr = (BYTE*)((ULONGLONG)vBuf + nt_offset);
    BYTE *pe_ptr = nt_ptr - sizeof(DWORD);
 
    if (!peconv::validate_ptr(vBuf, vBufSize, pe_ptr, sizeof(DWORD))) {
        return false;
    }
    IMAGE_NT_HEADERS32 *nt32 = (IMAGE_NT_HEADERS32*)pe_ptr;
    //write signature:
    nt32->Signature = IMAGE_NT_SIGNATURE;
    IMAGE_FILE_HEADER *file_hdr = &nt32->FileHeader;
 
    bool is64bit = (file_hdr->Machine == IMAGE_FILE_MACHINE_AMD64) ? true : false;
 
    if (nt32->FileHeader.SizeOfOptionalHeader == 0) {
        nt32->FileHeader.SizeOfOptionalHeader = is64bit ? sizeof(IMAGE_OPTIONAL_HEADER64) : sizeof(IMAGE_OPTIONAL_HEADER32);
    }
    LONG pe_offset = LONG((ULONGLONG)pe_ptr - (ULONGLONG)vBuf);
    IMAGE_DOS_HEADER* dosHdr = (IMAGE_DOS_HEADER*) vBuf;
    dosHdr->e_magic = IMAGE_DOS_SIGNATURE;
    dosHdr->e_lfanew = pe_offset;
 
    bool is_fixed = false;
    if (is64bit) {
        is_fixed = overwrite_opt_hdr<IMAGE_OPTIONAL_HEADER64>(vBuf, vBufSize, (IMAGE_OPTIONAL_HEADER64*)&nt32->OptionalHeader, this->artefacts);
    }
    else {
        is_fixed = overwrite_opt_hdr<IMAGE_OPTIONAL_HEADER32>(vBuf, vBufSize, &nt32->OptionalHeader, this->artefacts);
    }
    if (!is_fixed) {
        return false;
    }
    if (!peconv::get_nt_hdrs(vBuf, vBufSize)) {
        return false;
    }
    return true;
}