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

// user defined types

// Defines the nodes that make up the rva linked list.
struct _RVA_ITEM
{
    PRVA_ITEM pLink;
    DWORD dwRva;
};

// code start

// Class constructor.
CRelocationRipper::CRelocationRipper()
{
    m_hHeap = GetProcessHeap();
    m_pRvaList = NULL;
    m_bOutput = FALSE;
    m_bStrip = FALSE;
}

// Class destructor.
CRelocationRipper::~CRelocationRipper()
{
    DeleteRvaList();
}

// Adds a rva to the rva list.
VOID CRelocationRipper::AddRelocation(DWORD dwRva)
{
    PRVA_ITEM pNewRva;

    pNewRva = CreateRva(dwRva);
    if(pNewRva)
        pNewRva->dwRva = dwRva;
}

// Creates a new rva item and inserts it in the ordered rva list.
PRVA_ITEM CRelocationRipper::CreateRva(DWORD dwRva)
{
    PRVA_ITEM   pNewRva,
              * pPrevLink,
                pRva;

    pRva = m_pRvaList;
    pPrevLink = &m_pRvaList;
    while(pRva)
    {
        if(pRva->dwRva == dwRva)
            return NULL;
        else if(pRva->dwRva > dwRva)
            break;

        pPrevLink = &pRva->pLink;
        pRva = pRva->pLink;
    }
    pNewRva = (PRVA_ITEM)HeapAlloc(m_hHeap, 0, sizeof(RVA_ITEM));
    pNewRva->pLink = pRva;
    *pPrevLink = pNewRva;
    return pNewRva;
}

// Deletes all rva items in the rva list.
VOID CRelocationRipper::DeleteRvaList()
{
    PRVA_ITEM pCurrentRva,
              pNextRva;

    pCurrentRva = m_pRvaList;
    while(pCurrentRva)
    {
        pNextRva = pCurrentRva->pLink;
        HeapFree(m_hHeap, 0, pCurrentRva);
        pCurrentRva = pNextRva;
    }
    m_pRvaList = NULL;
}

// Exports relocation section.
VOID CRelocationRipper::Export(PVOID pvOutput)
{
    DWORD                  dwCurrentBase,
                           dwRvaBase;
    PRVA_ITEM              pCurrentRva;
    PIMAGE_BASE_RELOCATION pRelocHeader;
    PWORD                  pRelocs;

    if(pvOutput && m_bOutput)
    {
        pRelocHeader = (PIMAGE_BASE_RELOCATION)pvOutput;
        pRelocs = (PWORD)pvOutput;
        dwCurrentBase = -1;
        pCurrentRva = m_pRvaList;
        while(pCurrentRva)
        {
            // create a new table if needed
            dwRvaBase = pCurrentRva->dwRva & 0xFFFFF000;
            if(dwRvaBase != dwCurrentBase)
            {
                dwCurrentBase = dwRvaBase;
                pRelocHeader->SizeOfBlock = DWORD((PBYTE)pRelocs - (PBYTE)pRelocHeader);
                pRelocHeader = (PIMAGE_BASE_RELOCATION)pRelocs;
                pRelocHeader->VirtualAddress = dwCurrentBase;
                pRelocs = PWORD((PBYTE)pRelocHeader + IMAGE_SIZEOF_BASE_RELOCATION);
            }
            *pRelocs = WORD((IMAGE_REL_BASED_HIGHLOW << 12) | (pCurrentRva->dwRva & 0xFFF));
            pRelocs++;
            pCurrentRva = pCurrentRva->pLink;
        }
        pRelocHeader->SizeOfBlock = DWORD((PBYTE)pRelocs - (PBYTE)pRelocHeader);
        pRelocHeader = (PIMAGE_BASE_RELOCATION)pRelocs;

        // output terminating header
        pRelocHeader->VirtualAddress = 0;
        pRelocHeader->SizeOfBlock = sizeof(IMAGE_BASE_RELOCATION);
    }
}

// Returns size of exported relocation section.
DWORD CRelocationRipper::GetSize()
{
    DWORD     dwCurrentBase,
              dwRvaBase,
              dwSize;
    PRVA_ITEM pCurrentRva;

    dwSize = 0;
    if(m_bOutput)
    {
        // this basically mirrors the logic of the export function
        dwCurrentBase = -1;
        pCurrentRva = m_pRvaList;
        while(pCurrentRva)
        {
            dwRvaBase = pCurrentRva->dwRva & 0xFFFFF000;
            if(dwRvaBase != dwCurrentBase)
            {
                dwCurrentBase = dwRvaBase;
                dwSize += IMAGE_SIZEOF_BASE_RELOCATION;
            }
            dwSize += sizeof(WORD);
            pCurrentRva = pCurrentRva->pLink;
        }
        dwSize += IMAGE_SIZEOF_BASE_RELOCATION;
    }
    return dwSize;
}

// Loads input file. Returns FALSE if no relocation section exists.
BOOL CRelocationRipper::LoadFile(PBYTE pbFile)
{
    BOOL              bRet;
    DWORD             dwRelocRva;
    PIMAGE_NT_HEADERS pNt;

    bRet = FALSE;
    m_pbFile = pbFile;
    pNt = PIMAGE_NT_HEADERS(pbFile + PIMAGE_DOS_HEADER(pbFile)->e_lfanew);
    dwRelocRva = pNt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress;
    m_pReloc = (PIMAGE_BASE_RELOCATION)RvaToPointer(pbFile, dwRelocRva);
    if(m_pReloc) bRet = TRUE;
    return bRet;
}

// Just sets flags and strips an existing relocation section if necessary.
// No relocations are actually ripped.
VOID CRelocationRipper::Rip(PVOID pvFile)
{
    m_bOutput = FALSE;
    if(LoadFile((PBYTE)pvFile))
    {
        if(m_bStrip)
            StripRelocationsFromFile();
        else
            m_bOutput = TRUE;
    }
}

// Sets strip flag
VOID CRelocationRipper::Strip(BOOL bStrip)
{
    m_bStrip = bStrip;
}

// Strips relocations from the input file.
VOID CRelocationRipper::StripRelocationsFromFile()
{
    DWORD                  dwSize;
    PIMAGE_BASE_RELOCATION pReloc;

    pReloc = m_pReloc;
    while(pReloc->VirtualAddress)
    {
        dwSize = pReloc->SizeOfBlock;
        ZeroMemory(pReloc, dwSize);
        pReloc = PIMAGE_BASE_RELOCATION((PBYTE)pReloc + dwSize);
    }
}