#pragma once
#include <iostream>
#include "pe32_64.h"

#define RELOC_32BIT_FIELD 3
#define RELOC_64BIT_FIELD 0xA

#ifdef _WIN64
#define RELOC_FIELD RELOC_64BIT_FIELD
#else
#define RELOC_FIELD RELOC_32BIT_FIELD
#endif

typedef struct _BASE_RELOCATION_ENTRY {
    WORD Offset : 12;
    WORD Type : 4;
} BASE_RELOCATION_ENTRY;

inline bool relocate(BYTE* image, PIMAGE_NT_HEADERS nt, FIELD_PTR newImgBase)
{
    IMAGE_DATA_DIRECTORY relocationsDirectory = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
    if (relocationsDirectory.VirtualAddress == 0) {
        return false;
    }
    PIMAGE_BASE_RELOCATION ProcessBReloc = (PIMAGE_BASE_RELOCATION)(relocationsDirectory.VirtualAddress + (FIELD_PTR)image);
    // apply relocations:
    while (ProcessBReloc->VirtualAddress != 0)
    {
        DWORD page = ProcessBReloc->VirtualAddress;
#ifdef _DEBUG
        std::cout << "page: " << std::hex << page << std::endl;
#endif
        if (ProcessBReloc->SizeOfBlock >= sizeof(IMAGE_BASE_RELOCATION))
        {
            size_t count = (ProcessBReloc->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
            BASE_RELOCATION_ENTRY* list = (BASE_RELOCATION_ENTRY*)(LPWORD)(ProcessBReloc + 1);
#ifdef _DEBUG
            std::cout << "Count: " << count << ":\n";
#endif
            for (size_t i = 0; i < count; i++)
            {
                if (list[i].Type & RELOC_FIELD)
                {
                    DWORD rva = list[i].Offset + page;
#ifdef _DEBUG
                    std::cout << "RVA : " << std::hex << rva << "\n";
#endif
                    PULONG_PTR p = (PULONG_PTR)((LPBYTE)image + rva);
                    //relocate the address
                    *p = ((*p) - nt->OptionalHeader.ImageBase) + (FIELD_PTR)newImgBase;
                }
            }
        }
#ifdef _DEBUG
        std::cout << "---\n";
#endif
        ProcessBReloc = (PIMAGE_BASE_RELOCATION)((LPBYTE)ProcessBReloc + ProcessBReloc->SizeOfBlock);
    }
    return true;
}
