// Author:   Brandon LaCombe
// Date:     February 3, 2006
// License:  Public Domain
#include "CFileHeaderRipper.h"
#include "..\..\FileTools.h"
#include "..\..\remem.h"

// global variables
BYTE g_bPEStub[] = {0x0E, 0x1F, 0xBA, 0x0E, 0x00, 0xB4, 0x09, 0xCD, 0x21, 0xB8, 0x01, 0x4C, 0xCD, 0x21, 0x54, 0x68,
                    0x69, 0x73, 0x20, 0x70, 0x72, 0x6F, 0x67, 0x72, 0x61, 0x6D, 0x20, 0x63, 0x61, 0x6E, 0x6E, 0x6F,
                    0x74, 0x20, 0x62, 0x65, 0x20, 0x72, 0x75, 0x6E, 0x20, 0x69, 0x6E, 0x20, 0x44, 0x4F, 0x53, 0x20,
                    0x6D, 0x6F, 0x64, 0x65, 0x2E, 0x0D, 0x0D, 0x0A, 0x24, 0x00, 0x00, 0x00};

// code start

// Class constructor.
CFileHeaderRipper::CFileHeaderRipper()
{
    m_dwHeaderSize = 0;
    m_pbHeader = NULL;
    m_dwFileAlignment = 0x200;
    m_pszSectionName = NULL;
    m_bNewHeader = FALSE;
    m_bCombine = FALSE;
}

// Class destructor.
CFileHeaderRipper::~CFileHeaderRipper()
{
    if(m_pbHeader)
        VirtualFree(m_pbHeader, 0, MEM_RELEASE);
}

// Sets the flag that tells the ripper to combine all existing sections into
// one section. (aka max packing mode)
VOID CFileHeaderRipper::CombineSections(BOOL bCombine)
{
    m_bCombine = bCombine;
}

// Creates a new header with the sections combined (max new header)
VOID CFileHeaderRipper::CreateHeader(PBYTE pbFile)
{
    PBYTE                 pbOutputPtr;
    PIMAGE_NT_HEADERS     pInNt,
                          pOutNt;
    PIMAGE_SECTION_HEADER pOutSec;
    PIMAGE_DOS_HEADER     pOutDos;
    DWORD                 x;

    // allocate header memory
    m_dwHeaderSize = sizeof(IMAGE_DOS_HEADER) +
                     sizeof(g_bPEStub) +
                     sizeof(IMAGE_NT_HEADERS) +
                     sizeof(IMAGE_SECTION_HEADER) * 2;
    m_dwHeaderSize = align(m_dwHeaderSize, m_dwFileAlignment);

    m_pbHeader = (PBYTE)VirtualAlloc(NULL, m_dwHeaderSize, MEM_COMMIT, PAGE_READWRITE);
    ZeroMemory(m_pbHeader, m_dwHeaderSize);
    pbOutputPtr = m_pbHeader;

    // fill out dos header
    pOutDos = PIMAGE_DOS_HEADER(pbOutputPtr);
    pbOutputPtr += sizeof(IMAGE_DOS_HEADER);
    pOutDos->e_magic = IMAGE_DOS_SIGNATURE;
    pOutDos->e_cblp = 0x90;
    pOutDos->e_cp = 0x3;
    pOutDos->e_cparhdr = 0x4;
    pOutDos->e_maxalloc = 0xFFFF;
    pOutDos->e_sp = 0xB8;
    pOutDos->e_lfarlc = 0x40;
    pOutDos->e_lfanew = sizeof(IMAGE_DOS_HEADER) + sizeof(g_bPEStub);

    // copy dos stub
    CopyMemory(pbOutputPtr, g_bPEStub, sizeof(g_bPEStub));
    pbOutputPtr += sizeof(g_bPEStub);

    // fill out nt header
    pInNt = PIMAGE_NT_HEADERS(pbFile + PIMAGE_DOS_HEADER(pbFile)->e_lfanew);
    pOutNt = PIMAGE_NT_HEADERS(pbOutputPtr);
    pbOutputPtr += sizeof(IMAGE_NT_HEADERS);
    *pOutNt = *pInNt;
    pOutNt->FileHeader.NumberOfSections = 2;
    pOutNt->OptionalHeader.FileAlignment = m_dwFileAlignment;
    pOutNt->OptionalHeader.SizeOfHeaders = m_dwHeaderSize;

    // fill out section table
    pOutSec = PIMAGE_SECTION_HEADER(pbOutputPtr);
    pbOutputPtr += (sizeof(IMAGE_SECTION_HEADER) * 2);
    if(m_pszSectionName)
    {
        for(x = 0; m_pszSectionName[x] && x < IMAGE_SIZEOF_SHORT_NAME; x++)
        {
            pOutSec[0].Name[x] = (BYTE)m_pszSectionName[x];
            pOutSec[1].Name[x] = (BYTE)m_pszSectionName[x];
        }
    }
    pOutSec[0].PointerToRawData = m_dwHeaderSize;
    pOutSec[0].Characteristics = IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
    pOutSec[1].Characteristics = IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
    pOutSec[0].VirtualAddress = align(m_dwHeaderSize, pOutNt->OptionalHeader.SectionAlignment);
    pOutSec[1].VirtualAddress = CalculateNewSectionRva(pbFile);
    pOutSec[0].Misc.VirtualSize = pOutSec[1].VirtualAddress - pOutSec[0].VirtualAddress;
}

// Exports file header.
VOID CFileHeaderRipper::Export(PVOID pvOutput)
{
    CopyMemory(pvOutput, m_pbHeader, m_dwHeaderSize);
}

// Returns file header size.
DWORD CFileHeaderRipper::GetSize()
{
    return m_dwHeaderSize;
}

// Tells the class to create a new header.
VOID CFileHeaderRipper::NewHeader(BOOL bNewHeader)
{
    m_bNewHeader = bNewHeader;
}

// Driver for creating the file header.
VOID CFileHeaderRipper::Rip(PVOID pvFile)
{
    // if we are in max new header mode then create a new header
    // otherwise derive the file header from the existing header
    if(m_bCombine && m_bNewHeader)
        CreateHeader((PBYTE)pvFile);
    else
        RipHeader((PBYTE)pvFile);
}

// Rips the original file header and adds a section to it. If we are in combine
// mode it will also compress the existing sections into one section.
VOID CFileHeaderRipper::RipHeader(PBYTE pbFile)
{
    DWORD                 dwActualHeaderSize,
                          dwNtHeaderEnd,
                          dwSectionCount,
                          x;
    PIMAGE_NT_HEADERS     pInNt,
                          pOutNt;
    PIMAGE_SECTION_HEADER pOutSec;

    // allocate memory for the file header
    pInNt = PIMAGE_NT_HEADERS(pbFile + PIMAGE_DOS_HEADER(pbFile)->e_lfanew);
    dwActualHeaderSize = TrimMemorySize(pbFile, pInNt->OptionalHeader.SizeOfHeaders);
    m_dwHeaderSize = align(dwActualHeaderSize, m_dwFileAlignment);

    // the end of the nt header differs depending on if we are combining sections or not
    if(m_bCombine)
        dwNtHeaderEnd = PIMAGE_DOS_HEADER(pbFile)->e_lfanew + sizeof(IMAGE_NT_HEADERS) + sizeof(IMAGE_SECTION_HEADER) * 2;
    else
        dwNtHeaderEnd = PIMAGE_DOS_HEADER(pbFile)->e_lfanew + sizeof(IMAGE_NT_HEADERS) + sizeof(IMAGE_SECTION_HEADER) * (pInNt->FileHeader.NumberOfSections + 1);

    if(dwNtHeaderEnd > m_dwHeaderSize)
        m_dwHeaderSize = align(dwNtHeaderEnd, m_dwFileAlignment);
    m_pbHeader = (PBYTE)VirtualAlloc(NULL, m_dwHeaderSize, MEM_COMMIT, PAGE_READWRITE);
    ZeroMemory(m_pbHeader, m_dwHeaderSize);
    CopyMemory(m_pbHeader, pbFile, dwActualHeaderSize);

    // adjust nt header
    pOutNt = PIMAGE_NT_HEADERS(m_pbHeader + PIMAGE_DOS_HEADER(m_pbHeader)->e_lfanew);
    pOutSec = PIMAGE_SECTION_HEADER(pOutNt + 1);
    pOutNt->OptionalHeader.FileAlignment = m_dwFileAlignment;
    pOutNt->OptionalHeader.SizeOfHeaders = m_dwHeaderSize;

    // finish adjusting nt header and section header depending on which mode we are in
    if(m_bCombine)
    {
        pOutNt->FileHeader.NumberOfSections = 2;
        ZeroMemory(pOutSec, sizeof(IMAGE_SECTION_HEADER) * 2);
        pOutSec[0].PointerToRawData = m_dwHeaderSize;
        pOutSec[0].Characteristics = IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
        pOutSec[1].Characteristics = IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
        pOutSec[0].VirtualAddress = align(m_dwHeaderSize, pOutNt->OptionalHeader.SectionAlignment);
        pOutSec[1].VirtualAddress = CalculateNewSectionRva(pbFile);
        pOutSec[0].Misc.VirtualSize = pOutSec[1].VirtualAddress - pOutSec[0].VirtualAddress;

    }
    else
    {
        dwSectionCount = pOutNt->FileHeader.NumberOfSections++;
        ZeroMemory(&pOutSec[dwSectionCount], sizeof(IMAGE_SECTION_HEADER));
        pOutSec[dwSectionCount].Characteristics = IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
        pOutSec[dwSectionCount].VirtualAddress = CalculateNewSectionRva(pbFile);
    }

    // set section name
    if(m_pszSectionName)
    {
        for(x = 0; m_pszSectionName[x] && x < IMAGE_SIZEOF_SHORT_NAME; x++)
        {
            // which sections get the section name depend on if we are combining
            if(m_bCombine)
            {
                pOutSec[0].Name[x] = (BYTE)m_pszSectionName[x];
                pOutSec[1].Name[x] = (BYTE)m_pszSectionName[x];
            }
            else
                pOutSec[dwSectionCount].Name[x] = (BYTE)m_pszSectionName[x];
        }
    }
}

// Sets desired file alignment.
VOID CFileHeaderRipper::SetFileAlignment(DWORD dwFileAlignment)
{
    m_dwFileAlignment = dwFileAlignment;
}

// Sets desired section name.
VOID CFileHeaderRipper::SetSectionName(PTSTR pszSectionName)
{
    m_pszSectionName = pszSectionName;
}