#include "Injection.h"
#include <iostream>
#include <string>

bool SetDebugPrivilege(bool Enable)
{
    HANDLE hToken = nullptr;
    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken))
        return false;

    TOKEN_PRIVILEGES TokenPrivileges = { 0 };
    TokenPrivileges.PrivilegeCount = 1;
    TokenPrivileges.Privileges[0].Attributes = Enable ? SE_PRIVILEGE_ENABLED : 0;

    if (!LookupPrivilegeValueA(nullptr, "SeDebugPrivilege", &TokenPrivileges.Privileges[0].Luid))
    {
        CloseHandle(hToken);
        return false;
    }

    if (!AdjustTokenPrivileges(hToken, FALSE, &TokenPrivileges, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr))
    {
        CloseHandle(hToken);
        return false;
    }

    CloseHandle(hToken);

    return true;
}

bool Is64BitProcess(HANDLE hProc)
{
    bool Is64BitWin = false;
    BOOL Out = 0;
    IsWow64Process(GetCurrentProcess(), &Out);
    if (Out)
        Is64BitWin = true;

    if (!IsWow64Process(hProc, &Out))
        return false;

    if (Is64BitWin && !Out)
        return true;

    return false;
}

int main(UINT argc, char * argv[])
{
    if (argc < 5)
        return 0;

    bool ThreadHijacking = false;
    bool Unlink = false;

    char * szDll = nullptr;
    DWORD PID = 0;

    DWORD HeaderOption = 0;
    DWORD InjectionMethod = 0;

    for (UINT i = 1; i < argc; ++i)
    {
        if (!lstrcmpA(argv[i], "/p"))
        {
            if (i + 1 < argc)
                PID = strtoul(argv[i + 1], nullptr, 10);
            else
                return 0;

            i++;
        }
        else if (!lstrcmpA(argv[i], "/f"))
        {
            if (i + 1 < argc)
                szDll = argv[i + 1];
            else
                return 0;

            i++;
        }
        else if (!lstrcmpA(argv[i], "/m"))
        {
            if (i + 1 < argc)
            {
                DWORD val = strtoul(argv[i + 1], nullptr, 10);
                if (val <= 2)
                    InjectionMethod = val;
                else
                    InjectionMethod = 0;
            }
            else
                return 0;

            i++;
        }
        else if (!lstrcmpA(argv[i], "/h"))
        {
            if (i + 1 < argc)
            {
                DWORD val = strtoul(argv[i + 1], nullptr, 10);
                if (val <= 2)
                    HeaderOption = val;
                else
                    HeaderOption = 0;
            }
            else
                return 0;

            i++;
        }
        else if (!lstrcmpA(argv[i], "/t"))
        {
            ThreadHijacking = true;
        }
        else if (!lstrcmpA(argv[i], "/u"))
        {
            Unlink = true;
        }
    }

    INJECTION_MODE im = (INJECTION_MODE)InjectionMethod;
    DWORD Flags = 0;
    if (Unlink)
        Flags = INJ_UNLINK_FROM_PEB;
    Flags |= HeaderOption;

    SetDebugPrivilege(true);

    if (PID)
    {
        HANDLE hProc = OpenProcess(PROCESS_ALL_ACCESS, FALSE, PID);
        if (!hProc)
            return 0;

        #ifndef _WIN64

        if(!Is64BitProcess(hProc))
            InjectDLL(szDll, hProc, im, ThreadHijacking, Flags);
        else
        {
            std::string AllArgs("\"GH Injector - x64.exe\" ");
            for (UINT i = 1; i != argc; ++i)
            {
                if (!strcmp(argv[i], "/f"))
                {
                    AllArgs += argv[i];
                    AllArgs += ' ';
                    AllArgs += '"';
                    AllArgs += argv[i + 1];
                    AllArgs += '"';
                    AllArgs += ' ';
                    ++i;
                    continue;
                }
                AllArgs += argv[i];
                AllArgs += ' ';
            }

            PROCESS_INFORMATION pi{ 0 };
            STARTUPINFO si{ 0 };
            si.cb = sizeof(STARTUPINFO);

            if (!CreateProcessA(nullptr, const_cast<char*>(AllArgs.c_str()), nullptr, nullptr, FALSE, 0, nullptr, nullptr, &si, &pi))
                return 0;

            CloseHandle(pi.hProcess);
            CloseHandle(pi.hThread);

        }

        #else
        InjectDLL(szDll, hProc, im, ThreadHijacking, Flags);
        #endif
        CloseHandle(hProc);
    }
    return 0;
}