#include "injector.h"

#include <cassert>
#include <windows.h>

namespace {

DWORD ExecuteFunctionInProcess(HANDLE process, void const* function, void const* argument, size_t argument_size)
{
    void* argument_buffer{nullptr};
    if (argument != nullptr && argument_size > 0) {
        argument_buffer = VirtualAllocEx(process, nullptr, argument_size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
        WriteProcessMemory(process, argument_buffer, argument, argument_size, nullptr);
    }

    auto thread = CreateRemoteThread(
        process,
        nullptr,
        0,
        (LPTHREAD_START_ROUTINE)function,
        argument_buffer,
        0,
        nullptr);

    DWORD result{0};
    WaitForSingleObject(thread, INFINITE);
    GetExitCodeThread(thread, &result);

    if (argument_buffer != nullptr) {
        VirtualFreeEx(process, argument_buffer, 0, MEM_RELEASE);
        argument_buffer = nullptr;
    }

    return result;
}

}

uint32_t LaunchInjected(
    wchar_t const* command_line, 
    wchar_t const* working_directory, 
    wchar_t const* injected_dll, 
    char const* initialize_function)
{
    if (!command_line || !working_directory || !injected_dll || !initialize_function) {
        return 0;
    }

    PROCESS_INFORMATION process_info{};
    STARTUPINFO startup_info{};
    startup_info.cb = sizeof(startup_info);

    // or, you know, i could just const_cast and :pray:...
    auto command_length = wcslen(command_line);
    wchar_t* mutable_command_line = new wchar_t[command_length + 1];
    wcsncpy(mutable_command_line, command_line, command_length);
    mutable_command_line[command_length] = L'\0';

    auto process_launched = CreateProcess(
        nullptr, 
        mutable_command_line, 
        nullptr, 
        nullptr, 
        false, 
        CREATE_SUSPENDED, 
        nullptr, 
        working_directory, 
        &startup_info, 
        &process_info);
    if (!process_launched) {
        return 0;
    }

    delete mutable_command_line;
    mutable_command_line = nullptr;

    // technically wrong, but oh-so-useful. the base address of kernel32.dll is
    // identical for all processes in any given session.
    void* load_library = GetProcAddress(GetModuleHandle(L"kernel32.dll"), "LoadLibraryW");
    assert(load_library);

    void* remote_base_address = (void*)ExecuteFunctionInProcess(
        process_info.hProcess, 
        load_library, 
        injected_dll, 
        sizeof(wchar_t) * wcslen(injected_dll)
    );
    if (!remote_base_address) {
        TerminateProcess(process_info.hProcess, 0);
        return 0;
    }

    void* local_base_address = (void*)LoadLibrary(injected_dll);
    ptrdiff_t offset =
        (uintptr_t)GetProcAddress((HMODULE)local_base_address, initialize_function) - (uintptr_t)local_base_address;
    FreeLibrary((HMODULE)local_base_address);

    void* initialize_address = (void*)((uintptr_t)remote_base_address + offset);
    ExecuteFunctionInProcess(process_info.hProcess, initialize_address, nullptr, 0);

    ResumeThread(process_info.hThread);

    CloseHandle(process_info.hThread);
    CloseHandle(process_info.hProcess);

    return static_cast<uint32_t>(process_info.dwProcessId);
}

int WINAPI DllMain(HMODULE, DWORD, LPVOID)
{
    return true;
}