#include "xrbx/process.h"

HANDLE Process::openProcess() {
    if (pid == 0) {
        return nullptr;
    }

    CLIENT_ID cid{};
    cid.UniqueProcess = (HANDLE)pid;
    cid.UniqueThread = nullptr;

    OBJECT_ATTRIBUTES oa{};
    oa.Length = sizeof oa;

    HANDLE proc;

    if(Nt::resolve<Nt::OpenProcess>("NtOpenProcess")(&proc, PROCESS_ALL_ACCESS, &oa, &cid) < 0) {
        throw std::runtime_error("Failed to open process");
    }

    HANDLE duplicate;
    if(Nt::resolve<Nt::DuplicateObject>("NtDuplicateObject")(GetCurrentProcess(), proc, GetCurrentProcess(), &duplicate, PROCESS_ALL_ACCESS, 0, 0) < 0) {
        Nt::resolve<Nt::Close>("NtClose")(proc);
        throw std::runtime_error("Failed to duplicate process");
    }
    Nt::resolve<Nt::Close>("NtClose")(proc);

    return duplicate;
}

BYTE *Process::getImageBase(void) {
    Nt::PROCESS_BASIC_INFORMATION pbi{};

    if (Nt::resolve<Nt::QueryInformationProcess>("NtQueryInformationProcess")(proc, ProcessBasicInformation, &pbi, sizeof(pbi), nullptr) < 0) {
        throw std::runtime_error("Failed to get process information");
    }

    return (BYTE *)read<Nt::PEB_PARTIAL>((BYTE *)pbi.PebBaseAddress).ImageBaseAddress;
}

DWORD Process::findProcessByName(const std::wstring &name) {
    auto NtQuerySystemInformation = Nt::resolve<Nt::QuerySystemInformation>("NtQuerySystemInformation");

    ULONG size = 0;
    NtQuerySystemInformation(SystemProcessInformation, nullptr, 0, &size);

    std::vector<BYTE> buffer(size);
    NTSTATUS status = NtQuerySystemInformation(SystemProcessInformation, buffer.data(), size, &size);
    if (status < 0) return 0;

    auto entry = reinterpret_cast<PSYSTEM_PROCESS_INFORMATION>(buffer.data());

    while (true) {
        if (entry->ImageName.Buffer) {
            if (_wcsicmp(entry->ImageName.Buffer, name.c_str()) == 0) {
                return DWORD(ULONG_PTR(entry->UniqueProcessId));
            }
        }

        if (entry->NextEntryOffset == 0)
            break;

        entry = reinterpret_cast<PSYSTEM_PROCESS_INFORMATION>(
            reinterpret_cast<BYTE*>(entry)
            + entry->NextEntryOffset
        );
    }

    return 0;
}