#include <windows.h>

#define EPFX "__JADEITE_"

const char ENV_EXE_PATH[] = EPFX"TARGET_EXE_PATH";
const char ENV_DLL_PATH[] = EPFX"INJECT_DLL_PATH";
const char ENV_PROC_CMD[] = EPFX"PROCESS_COMMAND";

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

static inline void inject(HANDLE process, const void *payload, size_t payloadSize, const char *dllPath) {
    size_t _;

    // Inject the loader into the module
    size_t dllPathLen = strlen(dllPath) + 1;

    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
    ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size = 0;

    write_protected_process_memory(process, exe, exeHeader, sizeof(exeHeader));
}