PortableExecutable.cpp

#include "PortableExecutable.h"

#include "Log.h"

#include <algorithm>
#include <cassert>

#include <share.h>

// address without the image base!
DWORD PortableExecutable::VirtualToRaw(DWORD const dwAddress) const noexcept
{
    for (auto const& uSection : vecPESections)
    {
        auto const dwDifference = dwAddress - uSection.VirtualAddress;

        if (dwDifference < uSection.SizeOfRawData)
        {
            return uSection.PointerToRawData + dwDifference;
        }
    }

    return 0;
}

bool PortableExecutable::ReadFile()
{
    auto const pFile = this->Handle.get();

    // dos header
    fseek(pFile, 0, SEEK_SET);
    fread(&uDOSHeader, sizeof(IMAGE_DOS_HEADER), 1, pFile);
    if (uDOSHeader.e_magic != IMAGE_DOS_SIGNATURE)
    {
        return false;
    }

    // pe header
    fseek(pFile, uDOSHeader.e_lfanew, SEEK_SET);
    fread(&uPEHeader, sizeof(IMAGE_NT_HEADERS), 1, pFile);
    if (uPEHeader.Signature != IMAGE_NT_SIGNATURE)
    {
        return false;
    }

    // sections
    vecPESections.resize(uPEHeader.FileHeader.NumberOfSections);
    fread(&vecPESections[0], sizeof(IMAGE_SECTION_HEADER), vecPESections.size(), pFile);

    // imports
    auto const& Imports = uPEHeader.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (Imports.Size)
    {
        fseek(pFile, static_cast<long>(VirtualToRaw(Imports.VirtualAddress)), SEEK_SET);

        for (;;)
        {
            IMAGE_IMPORT_DESCRIPTOR import_desc;
            fread(&import_desc, sizeof(IMAGE_IMPORT_DESCRIPTOR), 1, pFile);

            if (!import_desc.Characteristics)
            {
                break;
            }

            vecImports.emplace_back().uDesc = import_desc;
        }
    }

    for (auto& current_import : vecImports)
    {
        constexpr auto Size = 0x100;
        char name_buf[Size];
        fseek(pFile, static_cast<long>(VirtualToRaw(current_import.uDesc.Name)), SEEK_SET);
        fgets(name_buf, Size, pFile);

        current_import.Name = name_buf;

        // thunks
        fseek(pFile, static_cast<long>(VirtualToRaw(current_import.uDesc.FirstThunk)), SEEK_SET);

        for (;;)
        {
            IMAGE_THUNK_DATA thunk_data;
            fread(&thunk_data, sizeof(IMAGE_THUNK_DATA), 1, pFile);

            if (!thunk_data.u1.AddressOfData)
            {
                break;
            }

            current_import.vecThunkData.emplace_back().uThunkData = thunk_data;
        }

        auto thunk_addr = reinterpret_cast<IMAGE_THUNK_DATA*>(current_import.uDesc.FirstThunk);
        for (auto& thunk : current_import.vecThunkData)
        {
            thunk.Address = reinterpret_cast<DWORD>(thunk_addr++);
            thunk.bIsOrdinal = IMAGE_SNAP_BY_ORDINAL(thunk.uThunkData.u1.Ordinal);

            if (thunk.bIsOrdinal)
            {
                thunk.Ordinal = IMAGE_ORDINAL(thunk.uThunkData.u1.Ordinal);
            }
            else
            {
                fseek(pFile, static_cast<long>(VirtualToRaw(thunk.uThunkData.u1.AddressOfData)), SEEK_SET);
                fread(&thunk.wWord, 2, 1, pFile);
                fgets(name_buf, Size, pFile);
                thunk.Name = name_buf;
            }
        }
    }

    return true;
}

bool PortableExecutable::ReadBytes(
    DWORD const dwRawAddress, size_t const Size, void* const Dest) const noexcept
{
    auto const pFile = this->Handle.get();

    if (!fseek(pFile, static_cast<long>(dwRawAddress), SEEK_SET))
    {
        return (fread(Dest, Size, 1, pFile) == 1);
    }

    return false;
}

bool PortableExecutable::ReadCString(
    DWORD const dwRawAddress, std::string& Result) const
{
    auto const pFile = this->Handle.get();

    if (!fseek(pFile, static_cast<long>(dwRawAddress), SEEK_SET))
    {
        constexpr size_t sz = 0x100;
        char tmpBuf[sz];

        tmpBuf[0] = 0;
        if (fread(tmpBuf, 1, sz, pFile) == sz)
        {
            tmpBuf[sz - 1] = 0;
            Result.assign(tmpBuf);
            return true;
        }
    }

    return false;
}

IMAGE_SECTION_HEADER const* PortableExecutable::FindSection(
    std::string_view const name) const noexcept
{
    assert(name.size() <= IMAGE_SIZEOF_SHORT_NAME);
    char buffer[IMAGE_SIZEOF_SHORT_NAME] = {};
    std::memcpy(buffer, name.data(), name.size());

    auto const found = std::find_if(
        vecPESections.cbegin(), vecPESections.cend(),
        [&buffer](auto const& section)
        {
            return !memcmp(section.Name, buffer, std::size(buffer));
        });

    if (found == vecPESections.cend())
    {
        return nullptr;
    }
    else
    {
        return &(*found);
    }
}