#include <windows.h>
#include <TlHelp32.h>
#include <Psapi.h>
#include <iostream>
#include <string>

using namespace std;

bool Inject(DWORD pid, LPTHREAD_START_ROUTINE callRoutine)
{
    PIMAGE_DOS_HEADER pIDH;
    PIMAGE_NT_HEADERS pINH;
    PIMAGE_BASE_RELOCATION pIBR;

    HANDLE hProcess, hThread;
    PUSHORT TypeOffset;

    PVOID ImageBase, Buffer, mem;
    ULONG i, Count, Delta, *p;

    bool result = true;

    printf("\n# Opening target process");
    hProcess = OpenProcess(
        PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION | PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_VM_WRITE,
        FALSE,
        pid);

    if (!hProcess)
    {
        printf("\n# Error: Unable to open target process (%u)", GetLastError());

        result = false;
    }

    ImageBase = GetModuleHandle(NULL);
    printf("\n# Image base in current process: %#x", ImageBase);

    pIDH = (PIMAGE_DOS_HEADER)ImageBase;
    pINH = (PIMAGE_NT_HEADERS)((PUCHAR)ImageBase + pIDH->e_lfanew);

    printf("\n# Allocating memory in target process");
    mem = VirtualAllocEx(hProcess, NULL, pINH->OptionalHeader.SizeOfImage, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);

    if (!mem)
    {
        printf("\n# Error: Unable to allocate memory in target process (%u)", GetLastError());

        CloseHandle(hProcess);

        result = false;
    }

    printf("\n# Memory allocated at %#x", mem);

    Buffer = VirtualAlloc(NULL, pINH->OptionalHeader.SizeOfImage, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    memcpy(Buffer, ImageBase, pINH->OptionalHeader.SizeOfImage);

    printf("\n# Relocating image");

    pIBR = (PIMAGE_BASE_RELOCATION)((PUCHAR)Buffer + pINH->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);
    Delta = (ULONG)mem - (ULONG)ImageBase;

    printf("\n# Delta: %#x", Delta);

    while (pIBR->VirtualAddress)
    {
        if (pIBR->SizeOfBlock >= sizeof(IMAGE_BASE_RELOCATION))
        {
            Count = (pIBR->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(USHORT);
            TypeOffset = (PUSHORT)(pIBR + 1);

            for (i = 0; i<Count; i++)
            {
                if (TypeOffset[i])
                {
                    p = (PULONG)((PUCHAR)Buffer + pIBR->VirtualAddress + (TypeOffset[i] & 0xFFF));
                    *p += Delta;
                }
            }
        }

        pIBR = (PIMAGE_BASE_RELOCATION)((PUCHAR)pIBR + pIBR->SizeOfBlock);
    }

    printf("\n# Writing relocated image into target process");

    if (!WriteProcessMemory(hProcess, mem, Buffer, pINH->OptionalHeader.SizeOfImage, NULL))
    {
        printf("\n# Error: Unable to write process memory (%u)", GetLastError());

        VirtualFreeEx(hProcess, mem, 0, MEM_RELEASE);
        CloseHandle(hProcess);

        result = false;
    }

    VirtualFree(Buffer, 0, MEM_RELEASE);

    printf("\n# Creating thread in target process");
    hThread = CreateRemoteThread(hProcess, NULL, 0, (LPTHREAD_START_ROUTINE)((PUCHAR)callRoutine + Delta), NULL, 0, NULL);

    if (!hThread)
    {
        printf("\n# Error: Unable to create thread in target process (%u)", GetLastError());

        VirtualFreeEx(hProcess, mem, 0, MEM_RELEASE);
        CloseHandle(hProcess);

        result = false;
    }

    printf("\n# Waiting for the thread to terminate");
    WaitForSingleObject(hThread, INFINITE);

    printf("\n# Thread terminated\n# Freeing allocated memory");

    VirtualFreeEx(hProcess, mem, 0, MEM_RELEASE);
    CloseHandle(hProcess);

    return result;
}

DWORD GetProcessesByName(const char *ProcName)
{
    PROCESSENTRY32   pe32;
    HANDLE         hSnapshot = NULL;

    pe32.dwSize = sizeof(PROCESSENTRY32);
    hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if (Process32First(hSnapshot, &pe32))
    {
        do{
            if (strcmp(pe32.szExeFile, ProcName) == 0)
                break;
        } while (Process32Next(hSnapshot, &pe32));
    }

    if (hSnapshot != INVALID_HANDLE_VALUE)
        CloseHandle(hSnapshot);

    return pe32.th32ProcessID;
}

string GetProcessPath(DWORD ProcessId)
{
    string path;
    char buffer[MAX_PATH];

    MessageBox(NULL, "ProcessPath() called", "Error", MB_ICONERROR);
    HANDLE hProcess;

    hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, ProcessId);

    if (GetModuleFileNameEx(hProcess, NULL, buffer, sizeof(buffer) / sizeof(char)))
    {
        path = string(buffer);
    }
    else {
        MessageBox(NULL, "Failed to get process path.", "Error", MB_ICONERROR);
    }

    return path;
}

DWORD WINAPI MainThread(PVOID p)
{
    MessageBox(NULL, "Main thread called.", "Information", MB_ICONINFORMATION);

    string path = GetProcessPath(GetProcessesByName("Spotify.exe"));
    MessageBox(NULL, path.c_str(), "Information", MB_ICONINFORMATION);
    // Calling GetProcessPath() here crashes the slave process cause of using string without runtime library

    bool MAIN_THREAD = true;

    while (MAIN_THREAD)
    {

        if (GetAsyncKeyState(VK_ESCAPE) & 0x8000)
        {
            MessageBox(NULL, "Main thread closed!", "Information", MB_ICONINFORMATION);
            MAIN_THREAD = false;
        }

        Sleep(1);

    }

    return 0;
}

int main()
{
    //string path = GetProcessPath(GetProcessesByName("Spotify.exe"));
    //MessageBox(NULL, path.c_str(), "Information", MB_ICONINFORMATION);
    // Calling GetProcessPath() here works fine, because of runtime library still valid
    Inject(GetProcessesByName("Spotify.exe"), MainThread);

    return EXIT_SUCCESS;
}