#include <inject.h>

static inline void write_protected_process_memory(HANDLE process, void *address, const void *buf, size_t size) {
    DWORD oldProtect;
    VirtualProtectEx(process, address, size, PAGE_EXECUTE_READWRITE, &oldProtect);

    size_t bytesWritten;
    WriteProcessMemory(process, address, buf, size, &bytesWritten);

    VirtualProtectEx(process, address, size, oldProtect, &oldProtect);
}

void inject(HANDLE process, const void *payload, size_t payloadSize, const wchar_t *dllPath) {
    size_t _; // Contrary to the docs, {Write,Read}ProcessMemory likes to crash if the last arg is NULL

    // Inject the loader into the module
    size_t dllPathLen = (wcslen(dllPath) + 1) * sizeof(wchar_t);

    char *remoteAlloc = VirtualAllocEx(process, NULL, payloadSize + dllPathLen, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
    WriteProcessMemory(process, remoteAlloc, payload, payloadSize, &_);
    WriteProcessMemory(process, remoteAlloc + payloadSize, dllPath, dllPathLen, &_);

    // Find the EXE header in the process
    char exeHeader[1024];
    IMAGE_DOS_HEADER *dosHeader = NULL;
    IMAGE_NT_HEADERS64 *ntHeaders = NULL;

    MEMORY_BASIC_INFORMATION memoryInfo;
    char *currentAddress = 0x0;
    while (VirtualQueryEx(process, currentAddress, &memoryInfo, sizeof(memoryInfo))) {
        ReadProcessMemory(process, currentAddress, exeHeader, sizeof(exeHeader), &_);

        dosHeader = (IMAGE_DOS_HEADER*)exeHeader;

        // DOS header magic "MZ"
        if (dosHeader->e_magic != 0x5A4D) {
            goto cont;
        } 

        ntHeaders = (IMAGE_NT_HEADERS64*)(exeHeader + dosHeader->e_lfanew);

        // NT header signature "PE"
        if (ntHeaders->Signature != 0x4550) {
            goto cont;
        }

        // Skip DLLs
        if ((ntHeaders->FileHeader.Characteristics | IMAGE_FILE_DLL) == IMAGE_FILE_DLL) {
            goto cont;
        }

        // Skip potential headers without an entry point 
        // I have no idea how and why they exist, but apparently they do
        if (ntHeaders->OptionalHeader.AddressOfEntryPoint == 0) {
            goto cont;
        }

        // Found EXE header
        break;

    cont:
        currentAddress += memoryInfo.RegionSize;
    }

    char *exe = (char*)memoryInfo.BaseAddress;

    // Replace the entry point with a jump to the loader
    char *entryPoint = exe + ntHeaders->OptionalHeader.AddressOfEntryPoint;

    const unsigned char JUMP_INST[] = { 0xFF, 0x25, 0x00, 0x00, 0x00, 0x00 };

    write_protected_process_memory(process, entryPoint, JUMP_INST, sizeof(JUMP_INST));
    write_protected_process_memory(process, entryPoint + sizeof(JUMP_INST), &remoteAlloc, sizeof(remoteAlloc));

    // Break the import table to prevent any dlls from being loaded
    // Step 1: break the first import descriptor
    char *importDescriptors = exe + ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
    IMAGE_IMPORT_DESCRIPTOR firstDescriptor;
    ZeroMemory(&firstDescriptor, sizeof(firstDescriptor));

    write_protected_process_memory(process, importDescriptors, &firstDescriptor, sizeof(firstDescriptor));

    // Step 2: break the image data directory entry
    size_t ddOffset = ((char*)&(ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size)) - exeHeader;
    DWORD newSize = 0;

    write_protected_process_memory(process, exe + ddOffset, &newSize, sizeof(newSize));
}