pe_buffer.cpp

Go to the documentation of this file.

#include "pe_buffer.h"
 
#include <iostream>
#include "../scanners/artefact_scanner.h"
#include "../utils/artefacts_util.h"
 
size_t pesieve::PeBuffer::calcRemoteImgSize(ULONGLONG modBaseAddr) const
{
    const size_t hdr_buffer_size = PAGE_SIZE;
    BYTE hdr_buffer[hdr_buffer_size] = { 0 };
    size_t pe_vsize = 0;
 
    PIMAGE_SECTION_HEADER hdr_ptr = NULL;
    if (peconv::read_remote_pe_header(this->processHndl, (BYTE*)modBaseAddr, hdr_buffer, hdr_buffer_size)) {
        hdr_ptr = peconv::get_section_hdr(hdr_buffer, hdr_buffer_size, 0);
    }
    if (!hdr_ptr) {
        pe_vsize = peconv::fetch_region_size(this->processHndl, (PBYTE)modBaseAddr);
        //std::cout << "[!] Image size at: " << std::hex << modBaseAddr << " undetermined, using region size instead: " << pe_vsize << std::endl;
        return pe_vsize;
    }
    pe_vsize = ArtefactScanner::calcImgSize(this->processHndl, (HMODULE)modBaseAddr, hdr_buffer, hdr_buffer_size, hdr_ptr);
    //std::cout << "[!] Image size at: " << std::hex << modBaseAddr << " undetermined, using calculated img size: " << pe_vsize << std::endl;
    return pe_vsize;
}
 
bool pesieve::PeBuffer::readRemote(ULONGLONG module_base, size_t pe_vsize)
{
    if (pe_vsize == 0) {
        // if not size supplied, try with the size fetched from the header
        pe_vsize = peconv::get_remote_image_size(processHndl, (BYTE*)module_base);
    }
    if (_readRemote(module_base, pe_vsize)) {
        return true; //success
    }
    // try with the calculated size
    pe_vsize = calcRemoteImgSize(module_base);
#ifdef _DEBUG
    std::cout << "[!] Image size at: " << std::hex << module_base << " undetermined, using calculated size: " << pe_vsize << std::endl;
#endif
    return _readRemote(module_base, pe_vsize);
}
 
bool pesieve::PeBuffer::fillFromBuffer(ULONGLONG module_base, util::ByteBuffer& data_cache)
{
    size_t cached_size = data_cache.getDataSize();
    if (!cached_size) {
        return false;
    }
    if (!allocBuffer(cached_size)) {
        return false;
    }
    this->moduleBase = module_base;
    this->relocBase = module_base; //by default set the same as module base
 
    ::memcpy(this->vBuf, data_cache.getData(), cached_size);
    return true;
}
 
bool  pesieve::PeBuffer::_readRemote(const ULONGLONG module_base, size_t pe_vsize)
{
    if (pe_vsize == 0) {
        return false;
    }
    if (!allocBuffer(pe_vsize)) {
        return false;
    }
 
    // store the base no matter if reading succeeded or failed
    this->moduleBase = module_base;
    this->relocBase = module_base; //by default set the same as module base
 
    const bool can_force_access = this->isRefl ? true : false;
    size_t read_size = peconv::read_remote_area(processHndl, (BYTE*)this->moduleBase, vBuf, pe_vsize, can_force_access);
    if (read_size != pe_vsize) {
#ifdef _DEBUG
        std::cout << "[!] Failed reading Image at: " << std::hex << this->moduleBase << " img size: " << pe_vsize << std::endl;
#endif
        freeBuffer();
        return false;
    }
    return true;
}
 
bool pesieve::PeBuffer::resizeBuffer(size_t new_size)
{
    if (!vBuf) return false;
 
    BYTE *new_buf = peconv::alloc_aligned(new_size, PAGE_READWRITE);
    if (!new_buf) {
        return false;
    }
 
    size_t smaller_size = (vBufSize < new_size) ? vBufSize : new_size;
    memcpy(new_buf, this->vBuf, smaller_size);
    freeBuffer();
 
    this->vBuf = new_buf;
    this->vBufSize = new_size;
    return true;
}
 
bool  pesieve::PeBuffer::resizeLastSection(size_t new_img_size)
{
    if (!vBuf) return false;
 
    PIMAGE_SECTION_HEADER last_sec = peconv::get_last_section(vBuf, vBufSize, false);
    if (!last_sec) {
        return false;
    }
 
    if (new_img_size < last_sec->VirtualAddress) {
        return false;
    }
 
    const size_t new_sec_vsize = new_img_size - last_sec->VirtualAddress;
    const size_t new_sec_rsize = new_sec_vsize;
 
    if (last_sec->VirtualAddress + new_sec_vsize > this->vBufSize) {
        //buffer too small
        return false;
    }
 
    if (!peconv::update_image_size(vBuf, MASK_TO_DWORD(new_img_size))) {
        return false;
    }
 
    last_sec->Misc.VirtualSize = MASK_TO_DWORD(new_sec_vsize);
    last_sec->SizeOfRawData = MASK_TO_DWORD(new_sec_rsize);
    return true;
}
 
bool pesieve::PeBuffer::dumpPeToFile(
    IN std::string dumpFileName,
    IN OUT peconv::t_pe_dump_mode &dumpMode,
    IN OPTIONAL const peconv::ExportsMapper* exportsMap,
    OUT OPTIONAL peconv::ImpsNotCovered *notCovered
)
{
    if (!vBuf || !isValidPe()) return false;
#ifdef _DEBUG
    std::cout << "Dumping using relocBase: " << std::hex << relocBase << "\n";
#endif
    if (exportsMap != nullptr) {
        const bool fixed = peconv::fix_imports(this->vBuf, this->vBufSize, *exportsMap, notCovered);
#ifdef _DEBUG
        if (!fixed) {
            std::cerr << "[-] Unable to fix imports!" << std::endl;
        }
#endif
    }
    if (dumpMode == peconv::PE_DUMP_AUTO) {
        bool is_raw_alignment_valid = peconv::is_valid_sections_alignment(vBuf, vBufSize, true);
        bool is_virtual_alignment_valid = peconv::is_valid_sections_alignment(vBuf, vBufSize, false);
#ifdef _DEBUG
        std::cout << "Is raw alignment valid: " << is_raw_alignment_valid << std::endl;
        std::cout << "Is virtual alignment valid: " << is_virtual_alignment_valid << std::endl;
#endif
        if (!is_raw_alignment_valid && is_virtual_alignment_valid) {
            //in case if raw alignment is invalid and virtual valid, try to dump using Virtual Alignment first
            dumpMode = peconv::PE_DUMP_REALIGN;
            bool is_dumped = peconv::dump_pe(dumpFileName.c_str(), this->vBuf, this->vBufSize, this->relocBase, dumpMode);
            if (is_dumped) {
                return is_dumped;
            }
            dumpMode = peconv::PE_DUMP_AUTO; //revert and try again
        }
    }
    // dump PE in a given dump mode:
    return peconv::dump_pe(dumpFileName.c_str(), this->vBuf, this->vBufSize, this->relocBase, dumpMode);
}
 
bool pesieve::PeBuffer::dumpToFile(IN std::string dumpFileName)
{
    if (!vBuf) return false;
    return peconv::dump_to_file(dumpFileName.c_str(), vBuf, vBufSize);
}
 
bool pesieve::PeBuffer::isCode()
{
    if (!vBuf) return false;
    return pesieve::util::is_code(vBuf, vBufSize);
}