/*
Start by refactoring the DLL. Then the DLL will be used by this loader.

Task 1) Self-injection with loading imports payload-side
We need to refactor this loader in such a way, that it will be loading imports payload-side.
Hints to follow:
1. This time we need to load the DLL manually, and then call its exports (also manually).  Copy the relevant function (get_func_by_name) from the shellcode task.
2. Remove the function load_imports, and instead of this, call the function SelfLoadImports from the DLL (you will fetch it via exports lookup)
3. Then, call fetch and call the function SayHello. If the MessageBox pops up, it means the task is solved.

Task 2) Remote-injection with loading imports payload-side
In case of remote injection of manually loaded PEs, imports have to be loaded payload-side.
Hints to follow:
In addition to the hints from the Task 1, there are some additional steps:
1. Create a new process to which you will be injecting.
2. This time we need to also allocate memory in the remote process, using VirtualAllocEx.
3. First, we load the image locally: yet, relocating it to the remote base. Then, we copy it to the remote memory.
4. We will call the exported functions via remote threads:
fetch the export from a local image, but remeber that in the remote image it is at the different base, so you need to rebase it manually.
*/

#include <windows.h>
#include <iostream>

#define RELOC_32BIT_FIELD 3
#define RELOC_64BIT_FIELD 0xA

#ifdef _WIN64
#define RELOC_FIELD RELOC_64BIT_FIELD
typedef ULONG_PTR FIELD_PTR;
#else
#define RELOC_FIELD RELOC_32BIT_FIELD
typedef  DWORD_PTR FIELD_PTR;
#endif

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

inline void manual_map(BYTE* image, BYTE *rawPE, PIMAGE_NT_HEADERS nt)
{
    memcpy(image, rawPE, nt->OptionalHeader.SizeOfHeaders);

    // map sections
    PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(nt);
    for (WORD i = 0; i < nt->FileHeader.NumberOfSections; i++) {

        memcpy((BYTE*)(image)+section[i].VirtualAddress, (BYTE*)(rawPE)+section[i].PointerToRawData, section[i].SizeOfRawData);
    }
}

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;
}

/*
in the payload:
You need to copy this function, paste it into the payload and refactor in such a way, that the payload will be able to load its own imports.

in the injector:
This function will be replaced, as we no longer want to load the imports of the payload by the injector.
Instead, we want to call the function inside the payload, that will load its own imports manually.
*/
inline bool load_imports(BYTE* image, PIMAGE_NT_HEADERS nt)
{
    IMAGE_DATA_DIRECTORY importsDirectory = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (importsDirectory.VirtualAddress == 0) {
        return false;
    }
    PIMAGE_IMPORT_DESCRIPTOR importDescriptor = (PIMAGE_IMPORT_DESCRIPTOR)(importsDirectory.VirtualAddress + (FIELD_PTR)image);

    while (importDescriptor->Name != NULL)
    {
        LPCSTR libraryName = (LPCSTR)importDescriptor->Name + (FIELD_PTR)image;
        HMODULE library = LoadLibraryA(libraryName);
#ifdef _DEBUG
        std::cout << "Loading: " << libraryName << "\n";
#endif
        if (library)
        {
            PIMAGE_THUNK_DATA thunk = NULL;
            thunk = (PIMAGE_THUNK_DATA)((FIELD_PTR)image + importDescriptor->FirstThunk);

            while (thunk->u1.AddressOfData != NULL)
            {
                FIELD_PTR functionAddress = NULL;
                if (IMAGE_SNAP_BY_ORDINAL(thunk->u1.Ordinal))
                {
                    LPCSTR functionOrdinal = (LPCSTR)IMAGE_ORDINAL(thunk->u1.Ordinal);
                    functionAddress = (FIELD_PTR)GetProcAddress(library, functionOrdinal);
                }
                else
                {
                    PIMAGE_IMPORT_BY_NAME functionName = (PIMAGE_IMPORT_BY_NAME)((FIELD_PTR)image + thunk->u1.AddressOfData);
                    functionAddress = (FIELD_PTR)GetProcAddress(library, functionName->Name);
                }
                thunk->u1.Function = functionAddress;
                ++thunk;
            }
        }

        importDescriptor++;
    }
    return true;
}

PIMAGE_NT_HEADERS get_nt_hdr(BYTE *rawPE)
{
    //get header
    IMAGE_DOS_HEADER* DOSHeader = PIMAGE_DOS_HEADER(rawPE);
    if (DOSHeader->e_magic != IMAGE_DOS_SIGNATURE) {
        return NULL;
    }
    PIMAGE_NT_HEADERS nt = PIMAGE_NT_HEADERS((char*)(rawPE)+DOSHeader->e_lfanew);
    if (nt->Signature != IMAGE_NT_SIGNATURE) {
        return NULL;
    }
    return nt;
}

bool load_and_run(BYTE* rawPE, size_t rawSize)
{
    //get header
    IMAGE_DOS_HEADER* DOSHeader = PIMAGE_DOS_HEADER(rawPE);
    PIMAGE_NT_HEADERS nt = PIMAGE_NT_HEADERS((char*)(rawPE)+DOSHeader->e_lfanew);
    if (!nt) {
        std::cerr << "Not a PE file\n";
        return false;
    }
    //write file to memory
    BYTE* image = (BYTE*)VirtualAlloc(NULL, nt->OptionalHeader.SizeOfImage, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
    if (!image) {
        std::cerr << "Allocating image has failed\n";
        return false;
    }
#ifdef _DEBUG
    std::cout << "Allocated at:" << std::hex << image << "\n";
#endif
    manual_map(image, rawPE, nt);
    if (!relocate(image, nt, (FIELD_PTR)image)) {
        std::cerr << "Relocating image has failed\n";
        return false;
    }
    load_imports(image, nt);

    // call new Entry Point
    ULONG_PTR EntryPoint = nt->OptionalHeader.AddressOfEntryPoint + (ULONG_PTR)image;
#ifdef NEW_THREAD
    HANDLE hThread = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)EntryPoint, NULL, 0, NULL);
    if (hThread) {
        std::cout << "Thread run at: " << std::hex << EntryPoint << "\n";
        WaitForSingleObject(hThread, INFINITE);
    }
#else
    std::cout << "Calling Entry Point at: " << std::hex << EntryPoint << "\n";
    int (*new_main)() = (int(*)())EntryPoint;
    //call the Entry Point of the manually loaded PE:
    new_main();
#endif
    return true;
}

BYTE* read_file(IN const LPSTR sourcePath, OUT DWORD &loaded_size)
{
    HANDLE hFile = CreateFileA(sourcePath, GENERIC_READ, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (!hFile || hFile == INVALID_HANDLE_VALUE) {
        std::cerr << "Opening the file: " << sourcePath << " has failed! Error: " << std::dec << GetLastError() << "\n";
        return NULL;
    }
    DWORD size = GetFileSize(hFile, NULL);
    if (size == INVALID_FILE_SIZE) {
        CloseHandle(hFile);
        return NULL;
    }
    //read the file
    BYTE* rawPE = new BYTE[size];
    if (!ReadFile(hFile, rawPE, size, NULL, NULL)) {
        std::cerr << "[ERROR] Reading the file has failed!\n";
        delete[]rawPE;
        rawPE = NULL;
    }
    CloseHandle(hFile);
    return rawPE;
}

int main(int argc, char* argv[])
{
    //get source process
    if (argc < 2) {
        std::cout << "Args: <exe_path>" << "\n";
        return 1;
    }
    const LPSTR sourcePath = argv[1];
    DWORD size = 0;
    BYTE* rawPE = read_file(sourcePath, size);
    if (!rawPE) {
        return 0;
    }
    load_and_run(rawPE, size);
    return 0;
}

