PE File Buffer and Memory Buffer

// FileMemBufferTest.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include "pch.h"
#include <iostream>
#include <Windows.h>

/*
* 获取文件路径
*/
void GetFile(TCHAR* szPath)
{
#if 1
    if (!GetModuleFileName(NULL, szPath, MAX_PATH))
    {
        printf("Cannot install service (%d)\n", GetLastError());
        return;
    }
    char* pPathEnd = (char*)szPath + strlen(szPath);
    while (pPathEnd--)
    {
        if (*pPathEnd == '\\')
        {
            *pPathEnd = '\0';
            break;
        }
    }
    // user32test.dll ntoskrnltest.exe
    //strcat_s(szPath, "\\ntdlltest.dll");

    // special? ntoskrnltest.exe
    strcat_s(szPath, MAX_PATH, "\\MyPE.exe");
#else
    strcpy_s(szPath, "D:\\test.dll");
#endif
}

/*
* 读取文件到内存
*/
bool ReadFile(const char* name, char** p)
{
    bool bRet = false;

    FILE* fp = NULL;
    errno_t error = fopen_s(&fp, name, "rb");
    if (error == 0)
    {
        fseek(fp, 0, SEEK_END);
        DWORD bytes = ftell(fp);
        // do not forget to free memory later
        *p = (char*)malloc(bytes);

        fseek(fp, 0, SEEK_SET);
        fread_s(*p, bytes, sizeof(char), bytes, fp);

        fclose(fp);
        bRet = true;
    }
    else
    {
        printf("ReadFile %s error %d\n", name, error);
    }

    return bRet;
}

BYTE* FileToMemory(BYTE* fileBuffer)
{
    PIMAGE_DOS_HEADER pDosHeader;
    PIMAGE_NT_HEADERS32 pNTHeaders;
    PIMAGE_SECTION_HEADER pSectionHeader;

    pDosHeader = (PIMAGE_DOS_HEADER)fileBuffer;
    pNTHeaders = (PIMAGE_NT_HEADERS32)(fileBuffer + pDosHeader->e_lfanew);
    pSectionHeader = (PIMAGE_SECTION_HEADER)((char*)pNTHeaders + sizeof(IMAGE_NT_HEADERS32));

    // get image size
    int sizeOfImage = pNTHeaders->OptionalHeader.SizeOfImage;
    // allocate memory
    BYTE* memBuffer = new BYTE[sizeOfImage];
    // zero memory
    memset(memBuffer, 0, sizeOfImage);

    int sizeOfFileHead = 0;
    sizeOfFileHead += pDosHeader->e_lfanew;
    sizeOfFileHead += sizeof(IMAGE_NT_HEADERS32);
    sizeOfFileHead += sizeof(IMAGE_SECTION_HEADER) * pNTHeaders->FileHeader.NumberOfSections;

    // copy header from file buffer
    memcpy_s(memBuffer, sizeOfFileHead, fileBuffer, sizeOfFileHead);

    char* Name[IMAGE_SIZEOF_SHORT_NAME + 1] = { 0 };
    int sectionNum = pNTHeaders->FileHeader.NumberOfSections;
    for (int i = 0; i < sectionNum; i++)
    {
        memcpy_s(Name, IMAGE_SIZEOF_SHORT_NAME, pSectionHeader->Name, IMAGE_SIZEOF_SHORT_NAME);
        printf("%d %s\nVirtualAddress: %08X, VirtualSize: %08X ==> %08X\nPointerToRawData: %08X, SizeOfRawData: %08X ==> %08X\n",
            i, Name,
            pSectionHeader->VirtualAddress,
            pSectionHeader->Misc.VirtualSize,
            pSectionHeader->VirtualAddress + pSectionHeader->Misc.VirtualSize,
            pSectionHeader->PointerToRawData,
            pSectionHeader->SizeOfRawData,
            pSectionHeader->PointerToRawData + pSectionHeader->SizeOfRawData
        );

        if (pSectionHeader->SizeOfRawData == 0)
        {
            // fill with zero
            pSectionHeader->VirtualAddress;
            pSectionHeader->Misc.VirtualSize;
        }
        else
        {
            // memory -> VirtualSize
            int dataLen = pSectionHeader->SizeOfRawData > pSectionHeader->Misc.VirtualSize ? pSectionHeader->Misc.VirtualSize : pSectionHeader->SizeOfRawData;
            char* dest = (char*)(memBuffer + pSectionHeader->VirtualAddress);
            char* src = (char*)(fileBuffer + pSectionHeader->PointerToRawData);
            memcpy_s(dest, pSectionHeader->Misc.VirtualSize, src, dataLen);
        }

        pSectionHeader++;
    }

    // release heap memory later
    return memBuffer;
}

BYTE* MemoryToFile(BYTE* memBuffer, DWORD& sizeOfPERawData)
{
    PIMAGE_DOS_HEADER pDosHeader;
    PIMAGE_NT_HEADERS32 pNTHeaders;
    PIMAGE_SECTION_HEADER pSectionHeader;

    pDosHeader = (PIMAGE_DOS_HEADER)memBuffer;
    pNTHeaders = (PIMAGE_NT_HEADERS32)(memBuffer + pDosHeader->e_lfanew);
    pSectionHeader = (PIMAGE_SECTION_HEADER)((char*)pNTHeaders + sizeof(IMAGE_NT_HEADERS32));

    // get image size
    int sizeOfImage = pNTHeaders->OptionalHeader.SizeOfImage;
    
    int sizeOfFileHead = 0;
    sizeOfFileHead += pDosHeader->e_lfanew;
    sizeOfFileHead += sizeof(IMAGE_NT_HEADERS32);
    sizeOfFileHead += sizeof(IMAGE_SECTION_HEADER) * pNTHeaders->FileHeader.NumberOfSections;

    // create new file buffer
    BYTE* newFileBuffer = new BYTE[sizeOfImage];
    memset(newFileBuffer, 0, sizeOfImage);

    memcpy_s(newFileBuffer, sizeOfFileHead, memBuffer, sizeOfFileHead);

    pSectionHeader = (PIMAGE_SECTION_HEADER)((char*)pNTHeaders + sizeof(IMAGE_NT_HEADERS32));
    char Name[IMAGE_SIZEOF_SHORT_NAME + 1] = { 0 };
    memset(Name, 0, IMAGE_SIZEOF_SHORT_NAME + 1);
    int sectionNum = pNTHeaders->FileHeader.NumberOfSections;
    for (int i = 0; i < sectionNum; i++)
    {
        memcpy_s(Name, IMAGE_SIZEOF_SHORT_NAME, pSectionHeader->Name, IMAGE_SIZEOF_SHORT_NAME);
        printf("%d %s\nVirtualAddress: %08X, VirtualSize: %08X ==> %08X\nPointerToRawData: %08X, SizeOfRawData: %08X ==> %08X\n",
            i, Name,
            pSectionHeader->VirtualAddress,
            pSectionHeader->Misc.VirtualSize,
            pSectionHeader->VirtualAddress + pSectionHeader->Misc.VirtualSize,
            pSectionHeader->PointerToRawData,
            pSectionHeader->SizeOfRawData,
            pSectionHeader->PointerToRawData + pSectionHeader->SizeOfRawData
        );

        if (pSectionHeader->SizeOfRawData == 0)
        {
            // fill with zero
            pSectionHeader->VirtualAddress;
            pSectionHeader->Misc.VirtualSize;
        }
        else
        {
            // File -> SizeOfRawData
            int dataLen = pSectionHeader->SizeOfRawData > pSectionHeader->Misc.VirtualSize ? pSectionHeader->SizeOfRawData : pSectionHeader->Misc.VirtualSize;
            char* dest = (char*)(newFileBuffer + pSectionHeader->PointerToRawData);
            char* src = (char*)(memBuffer + pSectionHeader->VirtualAddress);
            memcpy_s(dest, dataLen, src, pSectionHeader->Misc.VirtualSize);
        }

        pSectionHeader++;
    }

    pSectionHeader--;
    sizeOfPERawData = pSectionHeader->PointerToRawData + pSectionHeader->SizeOfRawData;

    // release heap memory later
    return newFileBuffer;
}

int main(int argc, char** argv)
{
    TCHAR szPath[MAX_PATH];
    if (argc <= 1)
    {
        printf("Please use [ReadFile.exe] [file]\n");

        // use default test
        GetFile(szPath);
    }
    else
    {
        printf("%d %s\n", argc, argv[1]);
        strncpy_s(szPath, argv[1], strlen(argv[1]));
    }

    char* fileBuffer = NULL;
    bool bRet = ReadFile(szPath, &fileBuffer);
    if (bRet)
    {
        BYTE* memBuffer = FileToMemory((BYTE*)fileBuffer);

        DWORD sizeOfPERawData = 0;
        BYTE* newFileBuffer = MemoryToFile((BYTE*)memBuffer, sizeOfPERawData);

        FILE* fpNewFile = nullptr;
        errno_t error = fopen_s(&fpNewFile, "New.exe", "wb");
        if (error == 0)
        {
            fwrite(newFileBuffer, 1, sizeOfPERawData, fpNewFile);
            fclose(fpNewFile);
        }

        // release memory
        if (newFileBuffer)
        {
            delete newFileBuffer;
            newFileBuffer = nullptr;
        }

        // release memory
        if (memBuffer)
        {
            delete memBuffer;
            memBuffer = nullptr;
        }
    }
}