#include <Windows.h>
#include <psapi.h>
#include <string>
#include <TlHelp32.h>
#include <Dbghelp.h>
#include <iostream>
#include <fstream>

#define HEADER_SIZE 1024

HANDLE get_module_of_process(const std::string& module_name, const HANDLE process)
{
    const HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, GetProcessId(process));

    if (snapshot != INVALID_HANDLE_VALUE) {
        MODULEENTRY32 entry = {0};
        entry.dwSize = sizeof(MODULEENTRY32);

        for (BOOL status = Module32First(snapshot, &entry); status; status = Module32Next(snapshot, &entry)) {
            if (!strcmp(entry.szModule, module_name.c_str())) {
                CloseHandle(snapshot);
                return entry.hModule;
            }
        }
    }

    CloseHandle(snapshot);
    throw std::runtime_error("Failed to find Module: " + module_name);
}

HANDLE get_process_handle_by_name(const std::string& name)
{
    const HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

    if(snapshot != INVALID_HANDLE_VALUE) {
        PROCESSENTRY32 entry {0};
        entry.dwSize = sizeof(PROCESSENTRY32);

        for (BOOL status = Process32First(snapshot, &entry); status; status = Process32Next(snapshot, &entry)) {
            if (!strcmp(entry.szExeFile, name.c_str())) {

                CloseHandle(snapshot);

                const HANDLE handle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, entry.th32ProcessID);

                if(handle != INVALID_HANDLE_VALUE)
                    return handle;

                throw std::runtime_error("Failed to open a handle to the Process");
            }
        }
    }

    CloseHandle(snapshot);
    throw std::runtime_error("Failed to find the process");
}

void load_library_remote(const std::string& process_name, const std::string& module_name)
{
    char path[MAX_PATH];

    const unsigned long size = GetFullPathNameA(module_name.c_str(), MAX_PATH, path, nullptr) + 1;

    if(size - 1 <= 0)
        throw std::runtime_error("GetFullPathNameA failed with Error Code " + std::to_string(GetLastError()));

    const HANDLE process = get_process_handle_by_name(process_name);
    const HANDLE process_module = get_module_of_process(process_name, process);

    char header[HEADER_SIZE];
    ReadProcessMemory(process, reinterpret_cast<const void*>(process_module), header, HEADER_SIZE, nullptr);

    IMAGE_DOS_HEADER* dos_header = reinterpret_cast<IMAGE_DOS_HEADER*>(header);
    IMAGE_NT_HEADERS* nt_headers = ImageNtHeader(header);

    if (dos_header->e_magic != IMAGE_DOS_SIGNATURE)
        throw std::runtime_error("Process is not a valid DOS Image");

    std::ifstream s(module_name, std::ios::binary | std::ios::ate);

    char module_header[HEADER_SIZE];

    if (s) {
        const std::streamoff file_size = s.tellg();

        if (file_size > HEADER_SIZE) {
            s.seekg(0, std::ios::beg);
            s.read(module_header, HEADER_SIZE);
            s.close();
        }
        else {
            throw std::runtime_error("Module " + module_name + " has an invalid size.");
        }
    }
    else {
        throw std::runtime_error("Failed to open a file stream to the Module " + module_name + ".");
    }

    IMAGE_DOS_HEADER* module_dos_header = reinterpret_cast<IMAGE_DOS_HEADER*>(header);
    IMAGE_NT_HEADERS* module_nt_header = ImageNtHeader(header);

    if (module_dos_header->e_magic != IMAGE_DOS_SIGNATURE)
        throw std::runtime_error("The module " + module_name + " is no valid DOS Image");

    if (module_nt_header->FileHeader.Machine != nt_headers->FileHeader.Machine)
        throw std::runtime_error("The Architectures of the Module and the Process don't match.");

    FARPROC load_library_ex = GetProcAddress(GetModuleHandleA("kernel32"), "LoadLibraryExA");

    if (load_library_ex) {

        uint8_t byte_code[] = {
            0x6A, 0x00, // push 0
            0x6A, 0x00, // push 0
            0x68, 0xBE, 0xBA, 0xFE, 0xCA,  // push 0xCAFEBABE
            0xB8, 0xBE, 0xBA, 0xFE, 0xCA, // mov eax, 0xCAFEBABE
            0xFF, 0xD0, // call eax
            0xC3 // ret
        };

        // patch the function address
        *reinterpret_cast<unsigned long*>(reinterpret_cast<unsigned long>(byte_code) + 10) =
            reinterpret_cast<unsigned long>(load_library_ex);

        void* remote_memory_path = VirtualAllocEx(process, nullptr, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);

        if (remote_memory_path) {

            // write the path to the process memory and let the stub pass it to LoadLibraryEx as Parameter
            if (WriteProcessMemory(process, remote_memory_path, path, size, nullptr)) {

                // replace the path address in the stub
                *reinterpret_cast<unsigned long*>(reinterpret_cast<unsigned long>(byte_code) + 5) =
                    reinterpret_cast<unsigned long>(remote_memory_path);

                void* remote_memory_byte_code = VirtualAllocEx(process, nullptr, sizeof byte_code, MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);

                if (remote_memory_byte_code) {

                    // Write the Bytecode and call it
                    if (WriteProcessMemory(process, remote_memory_byte_code, byte_code, sizeof byte_code, nullptr)) {
                        const HANDLE h = CreateRemoteThread(process, nullptr, 0, reinterpret_cast<LPTHREAD_START_ROUTINE>(remote_memory_byte_code), nullptr, 0, nullptr);

                        if (h == INVALID_HANDLE_VALUE) {
                            throw std::runtime_error("Failed to create a Thread on the Stub. Error Code: " + std::to_string(GetLastError()));
                        }

                        WaitForSingleObject(h, INFINITE);
                    }

                    // 0 because MEM_RELEASE requires dwSize to be 0
                    if (!VirtualFreeEx(process, remote_memory_byte_code, 0, MEM_RELEASE))
                        throw std::runtime_error("Error while freeing the Memory of the Byte Code. Error Code: " + std::to_string(GetLastError()));
                }
                else {
                    throw std::runtime_error("Failed to allocate Memory for the Byte Code. Error Code: " + std::to_string(GetLastError()));
                }
            }
            else {
                throw std::runtime_error("Failed to write the Path to the Remote Process Memory. Error Code: " + std::to_string(GetLastError()));
            }

            // 0 because MEM_RELEASE requires dwSize to be 0
            if (!VirtualFreeEx(process, remote_memory_path, 0, MEM_RELEASE))
                throw std::runtime_error("Error while freeing the Memory of the Path. Error Code: " + std::to_string(GetLastError()));
        }
    }
}

int main()
{
    try {
        load_library_remote("csgo.exe", "r3csgo.dll");
    }
    catch (std::runtime_error& e) {
        MessageBoxA(nullptr, e.what(), "Error", MB_OK);
        ExitProcess(EXIT_FAILURE);
    }

    return EXIT_SUCCESS;
}