INT RoutineItself(LPVOID)
{
    if (!MessageBoxA(NULL, "Hello, World from Metro game!", "Injected message", MB_OK))
        return ERROR_INVALID_FUNCTION;

    return ERROR_SUCCESS;
}

bool StartRoutine()
{
    if (HANDLE hThread = CreateThread(NULL, 0, LPTHREAD_START_ROUTINE(RoutineItself), 0, 0, NULL))
    {
        return true;
    }

    return false;
}

BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved)
{
    switch (fdwReason) {
    case DLL_PROCESS_DETACH:
        break;
    case DLL_PROCESS_ATTACH:
        if (!StartRoutine())
            return FALSE;

        break;
    case DLL_THREAD_ATTACH:
        break;
    case DLL_THREAD_DETACH:
        break;
    }

    return TRUE;
}

#include <algorithm>
#include <cassert>
#include <cstdio>
#include <iostream>

#include <Windows.h>
#include <psapi.h>
#include <shlwapi.h>
#include <tlhelp32.h>

#pragma comment(lib, "Shlwapi.lib")

enum ErrorCodes
{
    ERROR_LOAD_LIBRARY = -3,
    ERROR_OPEN_PROCESS = -2,
    ERROR_DEBUG_PRIVILEGE = -1,
    ERROR_SUCCEEDED = 0
};

struct HandleHolder
{
    HandleHolder() : handle() {}
    HandleHolder(HANDLE handle) : handle(handle) {}
    HandleHolder(const HandleHolder &) = delete;
    HandleHolder(HandleHolder &&rhs) : handle(NULL) { std::swap(handle, rhs.handle); }
    ~HandleHolder() { CloseHandle(handle); }
    operator HANDLE() { return handle; }
    HandleHolder &operator = (const HandleHolder &) = delete;
    HandleHolder &operator = (HandleHolder &&rhs) { std::swap(handle, rhs.handle); }
    HANDLE handle;
};

struct VirtualAllocHolder
{
    VirtualAllocHolder(const VirtualAllocHolder &) = delete;
    VirtualAllocHolder(HANDLE process, LPVOID address) : process(process), address(address) {}
    ~VirtualAllocHolder() { VirtualFreeEx(process, address, 0, MEM_RELEASE);  }
    VirtualAllocHolder &operator = (const VirtualAllocHolder &) = delete;
    operator LPVOID() { return address; }
    HANDLE process;
    LPVOID address;
};

struct CmdLine
{
    CmdLine() { argv = CommandLineToArgvW(GetCommandLineW(), &argc); }
    INT GetArgumentCount() const { return argc; }
    LPCWSTR GetArgument(INT index) const { assert(index < argc); return argv[index]; }
    INT argc;
    LPWSTR *argv;
};

HandleHolder OpenProcess(LPCWSTR processName)
{
#define OPEN_PROCESS_ERROR() \
    do { \
        std::fwprintf(stderr, \
        L"Failed to open system process '%s' with error code '%d'", \
        processName, \
        GetLastError()); \
    } \
    while (0)

    HandleHolder snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (snapshot == INVALID_HANDLE_VALUE) {
        OPEN_PROCESS_ERROR();
        return NULL;
    }

    PROCESSENTRY32 processEntry;
    if (Process32First(snapshot, &processEntry) == FALSE) {
        if (GetLastError() != ERROR_NO_MORE_FILES)
            OPEN_PROCESS_ERROR();
        return NULL;
    }

    do {
        HandleHolder process = OpenProcess(PROCESS_ALL_ACCESS, FALSE, processEntry.th32ProcessID);
        if (process != NULL) {
            WCHAR imagePath[MAX_PATH];
            if (GetProcessImageFileNameW(process, imagePath, MAX_PATH) != 0) {
                LPCWSTR moduleName = PathFindFileNameW(imagePath);
                if (lstrcmpiW(processName, moduleName) == 0)
                    return process;
            }
        }
    }
    while (Process32Next(snapshot, &processEntry) != FALSE);

    return NULL;

# undef OPEN_PROCESS_ERROR
}

BOOL InjectDLL(HANDLE process, LPCWSTR libraryPath)
{
#define LOAD_LIBRARY_ERROR() \
    do { \
        std::fwprintf(stderr, \
        L"Failed to load library '%s' with error code '%d'", \
        libraryPath, \
        GetLastError()); \
    } \
    while (0)

    FARPROC startRoutine = GetProcAddress(GetModuleHandle("KERNEL32.DLL"), "LoadLibraryW");
    if (startRoutine == NULL) {
        LOAD_LIBRARY_ERROR();
        return FALSE;
    }

    CONST SIZE_T length = (lstrlenW(libraryPath) + 1) * sizeof(WCHAR);
    VirtualAllocHolder argument(process, VirtualAllocEx(process, NULL, length, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE));
    if (argument == NULL) {
        LOAD_LIBRARY_ERROR();
        return FALSE;
    }

    SIZE_T bytesWritten;
    if (WriteProcessMemory(process, argument, libraryPath, length, &bytesWritten) == FALSE) {
        LOAD_LIBRARY_ERROR();
        return FALSE;
    }

    HandleHolder remoteThread = CreateRemoteThread(process, NULL, 0, LPTHREAD_START_ROUTINE(startRoutine), argument, 0, NULL);
    if (remoteThread == NULL) {
        LOAD_LIBRARY_ERROR();
        return FALSE;
    }

    if (WaitForSingleObject(remoteThread, INFINITE) != WAIT_OBJECT_0)
        return FALSE;

    return TRUE;

#undef LOAD_LIBRARY_ERROR
}

BOOL SetDebugPrivilege(BOOL onoff)
{
#define SET_DEBUG_PRIVILEGE_ERROR() \
    do { \
        std::fwprintf(stderr, \
        L"Failed to set debug privilege with error code '%d'", \
        GetLastError()); \
    } \
    while (0)

    HandleHolder token;
    if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &token.handle) == FALSE) {
        SET_DEBUG_PRIVILEGE_ERROR();
        return FALSE;
    }

    LUID luid;
    if (LookupPrivilegeValueA(NULL, "SeDebugPrivilege", &luid) == FALSE) {
        SET_DEBUG_PRIVILEGE_ERROR();
        return FALSE;
    }

    TOKEN_PRIVILEGES privileges;
    privileges.PrivilegeCount = 1;
    privileges.Privileges[0].Luid = luid;
    privileges.Privileges[0].Attributes =
        (onoff) ? SE_PRIVILEGE_ENABLED : SE_PRIVILEGE_REMOVED;

    if (AdjustTokenPrivileges(token, FALSE, &privileges, 0, NULL, NULL) == FALSE) {
        SET_DEBUG_PRIVILEGE_ERROR();
        return FALSE;
    }

    return TRUE;

#undef SET_DEBUG_PRIVILEGE_ERROR
}

int main()
{
    CmdLine cmdLine;
    if (cmdLine.GetArgumentCount() != 3) {
        std::puts("USAGE: dllinject <process name> <DLL module full path>");
        return ERROR_SUCCEEDED;
    }

    if (SetDebugPrivilege(TRUE) == FALSE)
        return ERROR_DEBUG_PRIVILEGE;

    LPCWSTR processName = cmdLine.GetArgument(1);
    HandleHolder process = OpenProcess(processName);
    if (process == NULL)
        return ERROR_LOAD_LIBRARY;

    LPCWSTR dllPath = cmdLine.GetArgument(2);
    if (InjectDLL(process, dllPath) == FALSE)
        return ERROR_OPEN_PROCESS;

    return ERROR_SUCCEEDED;
}

#define _CRT_SECURE_NO_WARNINGS

#include <fcntl.h>
#include <io.h>
#include <sys/stat.h>

#include <cstdlib>

#include <Windows.h>
#include <tlhelp32.h>

#define PROCESS_NAME "MetroExodus.exe"

struct ReaderData
{
    int FileDescriptor;
    long FilePointer;
    long FileLength;
};

BYTE *FindTargetFunctionAddress(BYTE *mem, INT length)
{
    CONST BYTE pattern[] =
    {
        '\x83', '\xEC', '\x20', '\x48', '\x89', '\xD6', '\x41', '\xF7',
        '\xD8', '\x48', '\x89', '\xCB', '\x48', '\x89', '\xF1', '\x19',
        '\xFF', '\x83', '\xE7', '\x20', '\x0F', '\xBA', '\xEF', '\x0F',
        '\x89', '\xFA', '\xFF', '\x15', '\xB4', '\xF9', '\x6A', '\x04',
        '\x89', '\x03', '\x83', '\xF8', '\xFF', '\x75', '\x35', '\xFF',
        '\x15', '\x1F', '\xF9', '\x6A', '\x04', '\x83', '\x38', '\x02',
    };

    INT i, j;
    CONST INT count = sizeof(pattern);
    for (i = 0; i <= length - count; ++i)
    {
        for (j = 0; j < count; ++j)
        {
            if (mem[i + j] != pattern[j])
            {
                break;
            }
        }

        if (j == count)
        {
            return mem + i;
        }
    }

    return NULL;
}

LPVOID GetTargetFunctionAddress()
{
    MODULEENTRY32 entry = MODULEENTRY32();
    entry.dwSize = sizeof(MODULEENTRY32);
    HANDLE hProcess = GetCurrentProcess();
    DWORD processId = GetProcessId(hProcess);
    HANDLE hModuleSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, processId);

    if (Module32First(hModuleSnapshot, &entry) == FALSE)
    {
        return NULL;
    }

    do
    {
        if (strcmp(entry.szModule, PROCESS_NAME) == 0)
        {
            return FindTargetFunctionAddress(entry.modBaseAddr, entry.modBaseSize);
        }
    }
    while (Module32Next(hModuleSnapshot, &entry) == TRUE);

    return NULL;
}

void Metro_OpenFile(ReaderData *reader, const char *path)
{
    int fd = _open(path, _O_BINARY, 0);
    if (fd == -1)
    {
        if (errno == ENOENT)
        {
            return;
        }

        Sleep(10);
        fd = _open(path, _O_BINARY, 0);
        if (fd == -1)
        {
            return;
        }
    }

    long length = _filelength(fd);
    if (length == -1)
    {
        _close(fd);
        return;
    }

    reader->FileDescriptor = fd;
    reader->FileLength = length;
}

BOOL SetHook(LPVOID address, LPVOID detour)
{
    DWORD oldProtection;
    CONST SIZE_T length = 32;
    CONST DWORD64 offset = reinterpret_cast<DWORD64>(detour) - reinterpret_cast<DWORD64>(address) - 9;

    if (VirtualProtect(address, length, PAGE_EXECUTE_READWRITE, &oldProtection) != TRUE)
    {
        return NULL;
    }

    memset(address, 0x90, length);
    CHAR *p = reinterpret_cast<CHAR*>(address);
    *reinterpret_cast<BYTE*>(p++) = 0xEB;
    *reinterpret_cast<DWORD64*>(p) = offset;

    VirtualProtect(address, length, oldProtection, &oldProtection);

    return TRUE;
}

INT RoutineItself(LPVOID)
{
    LPVOID detour = Metro_OpenFile;
    LPVOID target = GetTargetFunctionAddress();

    if (target == NULL)
    {
        std::abort();
    }

    if (SetHook(target, detour) == FALSE)
    {
        return ERROR_INVALID_FUNCTION;
    }

    return ERROR_SUCCESS;
}

bool StartRoutine()
{
    if (HANDLE hThread = CreateThread(NULL, 0, LPTHREAD_START_ROUTINE(RoutineItself), 0, 0, NULL))
    {
        return true;
    }

    return false;
}

BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved)
{
    switch (fdwReason) {
    case DLL_PROCESS_DETACH:
        break;
    case DLL_PROCESS_ATTACH:
        if (!StartRoutine())
        {
            return FALSE;
        }
        break;
    case DLL_THREAD_ATTACH:
        break;
    case DLL_THREAD_DETACH:
        break;
    }

    return TRUE;
}