agent.cpp

// agent.cpp
//
// Builds to vgagent32.exe / vgagent64.exe - the process a compliance check
// (e.g. a Citrix Gateway EPA "process running" scan) should look for.
//
// Watches for processes matching a set of target executable names (and,
// optionally, their entire descendant process tree), and injects
// vghook.dll into each match exactly once. vghook.dll then keeps
// re-applying WDA_EXCLUDEFROMCAPTURE to that process's own windows on its
// own. (The hook DLL is named separately from this agent so a process
// listing doesn't show two things both called "vgagent" and so an EPA policy
// checking for the agent process can't be confused with the DLL file.)
//
// Usage:
//   vgagent64.exe [--tree] <dll_path> <exe_name_1> [<exe_name_2> ...]
//   vgagent64.exe --list
//   vgagent64.exe --config <config_path>
//
// Examples:
//   vgagent64.exe vghook64.dll CDViewer.exe wfica32.exe
//   vgagent64.exe --tree vghook64.dll msrdc.exe
//   vgagent64.exe --list        (prints every visible top-level window and
//                               the exe that owns it - use this to find out
//                               which process ACTUALLY owns the window
//                               you're trying to blank, instead of guessing)
//   vgagent64.exe --config C:\path\targets.txt
//                              (reads target names + tree flag from a config
//                               file instead of the command line; resolves
//                               the DLL to vghookNN.dll next to this exe
//                               by convention. This is the mode vgtray.exe
//                               launches at logon - it hides its own console
//                               window since it's meant to run unattended in
//                               the background.)
//
// Config file format (one directive/name per line, '#' comments allowed):
//   TREE=1
//   CDViewer.exe
//   Citrix.DesktopViewer.App.exe
//
// --tree also injects into every descendant process of a matched process,
// even if the descendant's own exe name isn't in the list. Useful when the
// client spawns helper/render subprocesses under generic names.
//
// Notes:
//   - This only works against a target process whose bitness (32/64-bit)
//     matches this agent's own bitness. Build/run the matching variant
//     for the client you're targeting (see build instructions).
//   - No admin rights are required as long as the target process runs as
//     the same user, same session, same integrity level as this agent -
//     which is the normal case for a user's own Citrix/Windows App client.
//   - Expect antivirus/EDR to flag this. CreateRemoteThread + WriteProcessMemory
//     is a well-known pattern, and most AV engines heuristically flag it
//     regardless of intent. That's expected, not a bug.

#include "common.h"

#include <fstream>
#include <iostream>
#include <iterator>
#include <unordered_set>

using namespace vg_common;

namespace {

void EnableDebugPrivilege() {
    // Best-effort only. Not needed for same-user/same-session targets, but
    // helps if this is ever run elevated against a wider target set.
    HANDLE hToken = nullptr;
    if (!OpenProcessToken(GetCurrentProcess(),
                           TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
        return;
    }
    LUID luid;
    if (LookupPrivilegeValueW(nullptr, SE_DEBUG_NAME, &luid)) {
        TOKEN_PRIVILEGES tp;
        tp.PrivilegeCount = 1;
        tp.Privileges[0].Luid = luid;
        tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
        AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(tp), nullptr, nullptr);
    }
    CloseHandle(hToken);
}

enum class InjectResult { Ok, ArchMismatch, OpenFailed, InjectFailed };

bool IsTargetWow64(HANDLE hProcess, bool* isWow64) {
    BOOL result = FALSE;
    if (!IsWow64Process(hProcess, &result)) {
        return false;
    }
    *isWow64 = (result != FALSE);
    return true;
}

InjectResult InjectDll(DWORD pid, const std::wstring& dllPath) {
    HANDLE hProcess = OpenProcess(
        PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION |
            PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ,
        FALSE, pid);
    if (!hProcess) {
        return InjectResult::OpenFailed;
    }

    // Bitness sanity check: a 64-bit injector can only usefully inject a
    // 64-bit DLL into a 64-bit (non-WOW64) process, and vice versa.
#if defined(_WIN64)
    bool selfIsWow64 = false;
#else
    bool selfIsWow64 = true;
#endif
    bool targetIsWow64 = false;
    if (IsTargetWow64(hProcess, &targetIsWow64)) {
        if (targetIsWow64 != selfIsWow64) {
            CloseHandle(hProcess);
            return InjectResult::ArchMismatch;
        }
    }

    size_t bytes = (dllPath.size() + 1) * sizeof(wchar_t);
    LPVOID remoteMem = VirtualAllocEx(hProcess, nullptr, bytes,
                                       MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    if (!remoteMem) {
        CloseHandle(hProcess);
        return InjectResult::InjectFailed;
    }

    if (!WriteProcessMemory(hProcess, remoteMem, dllPath.c_str(), bytes, nullptr)) {
        VirtualFreeEx(hProcess, remoteMem, 0, MEM_RELEASE);
        CloseHandle(hProcess);
        return InjectResult::InjectFailed;
    }

    HMODULE hKernel32 = GetModuleHandleW(L"kernel32.dll");
    auto pLoadLibraryW = reinterpret_cast<LPTHREAD_START_ROUTINE>(
        reinterpret_cast<void*>(GetProcAddress(hKernel32, "LoadLibraryW")));
    if (!pLoadLibraryW) {
        VirtualFreeEx(hProcess, remoteMem, 0, MEM_RELEASE);
        CloseHandle(hProcess);
        return InjectResult::InjectFailed;
    }

    HANDLE hRemoteThread = CreateRemoteThread(
        hProcess, nullptr, 0, pLoadLibraryW, remoteMem, 0, nullptr);
    if (!hRemoteThread) {
        VirtualFreeEx(hProcess, remoteMem, 0, MEM_RELEASE);
        CloseHandle(hProcess);
        return InjectResult::InjectFailed;
    }

    WaitForSingleObject(hRemoteThread, 5000);
    DWORD exitCode = 0;
    GetExitCodeThread(hRemoteThread, &exitCode); // remote HMODULE, 0 = failure

    CloseHandle(hRemoteThread);
    VirtualFreeEx(hProcess, remoteMem, 0, MEM_RELEASE);
    CloseHandle(hProcess);

    return exitCode != 0 ? InjectResult::Ok : InjectResult::InjectFailed;
}

// Diagnostic mode: print every visible top-level window, its class, and the
// exe name that owns it. This is the fast way to find out which process
// ACTUALLY owns the window you're trying to blank, instead of guessing.
void PrintWindowList() {
    std::vector<WindowInfo> windows = ListVisibleWindows();

    std::vector<ProcInfo> procs = SnapshotProcesses();
    std::unordered_map<DWORD, std::wstring> pidToExe;
    for (const auto& p : procs) pidToExe[p.pid] = p.exeName;

    std::wcout << L"Visible top-level windows (title, owning exe, pid, class):\n\n";
    for (const auto& w : windows) {
        auto it = pidToExe.find(w.pid);
        std::wstring exe = (it != pidToExe.end()) ? it->second : L"?";
        std::wcout << L"  \"" << w.title << L"\"\n"
                   << L"      exe: " << exe << L"   pid: " << w.pid
                   << L"   class: " << w.className << L"\n";
    }
    std::wcout << L"\nWhichever exe owns the window you want blanked is the "
                  L"name to pass to vgagent64.exe/vgagent32.exe.\n";
}

struct Config {
    bool treeMode = false;
    std::unordered_set<std::wstring> targetNames;
};

bool LoadConfig(const std::wstring& path, Config* outConfig) {
    std::wifstream f(path.c_str());
    if (!f.is_open()) return false;

    std::wstring line;
    while (std::getline(f, line)) {
        // Trim trailing CR (files saved with CRLF read via wifstream can
        // leave a stray \r), and surrounding whitespace.
        while (!line.empty() && (line.back() == L'\r' || iswspace(line.back()))) {
            line.pop_back();
        }
        size_t start = line.find_first_not_of(L" \t");
        if (start == std::wstring::npos) continue; // blank line
        line = line.substr(start);
        if (line.empty() || line[0] == L'#') continue;

        if (line == L"TREE=1") {
            outConfig->treeMode = true;
        } else if (line == L"TREE=0") {
            outConfig->treeMode = false;
        } else {
            outConfig->targetNames.insert(ToLower(line));
        }
    }
    return true;
}

// The actual watch/inject loop, shared by direct CLI invocation and
// --config mode. Never returns under normal operation.
void RunWatchLoop(const std::unordered_set<std::wstring>& targetNames,
                   bool treeMode, const std::wstring& dllPath) {
    EnableDebugPrivilege();

    std::wcout << L"Watching for: ";
    for (const auto& n : targetNames) std::wcout << n << L" ";
    std::wcout << (treeMode ? L"(and their process trees)" : L"") << L"\n";
    std::wcout << L"DLL: " << dllPath << L"\n";
    std::wcout << L"Press Ctrl+C to stop.\n\n";

    std::unordered_set<DWORD> injectedPids;
    std::unordered_set<DWORD> openFailWarned; // avoid re-printing every 2s

    for (;;) {
        std::vector<ProcInfo> procs = SnapshotProcesses();
        std::unordered_map<DWORD, ProcInfo> byPid;
        for (const auto& p : procs) byPid[p.pid] = p;

        // Pass 1: direct name matches.
        std::unordered_map<DWORD, bool> matchedRoots;
        for (const auto& p : procs) {
            if (targetNames.count(ToLower(p.exeName))) {
                matchedRoots[p.pid] = true;
            }
        }

        // Pass 2: everything, decide final match set.
        for (const auto& p : procs) {
            bool isMatch = matchedRoots.count(p.pid) != 0;
            if (!isMatch && treeMode) {
                isMatch = DescendsFrom(p.pid, byPid, matchedRoots);
            }
            if (!isMatch) continue;
            if (injectedPids.count(p.pid)) continue;

            InjectResult result = InjectDll(p.pid, dllPath);
            switch (result) {
                case InjectResult::Ok:
                    std::wcout << L"[+] Injected into " << p.exeName
                               << L" (pid " << p.pid << L")\n";
                    injectedPids.insert(p.pid);
                    break;
                case InjectResult::ArchMismatch:
                    std::wcout << L"[!] " << p.exeName << L" (pid " << p.pid
                               << L") is a different bitness than this agent - "
                                  L"the other bitness variant should pick it up.\n";
                    injectedPids.insert(p.pid); // don't retry every loop
                    break;
                case InjectResult::OpenFailed:
                    // Print once per pid, not every 2s. A failure here means
                    // we found the process but couldn't get a handle to it -
                    // this is NOT the same as success.
                    if (!openFailWarned.count(p.pid)) {
                        std::wcout << L"[!] Could not open " << p.exeName
                                   << L" (pid " << p.pid << L") - OpenProcess"
                                      L" failed with error " << GetLastError()
                                   << L". This process was NOT injected.\n";
                        openFailWarned.insert(p.pid);
                    }
                    break;
                case InjectResult::InjectFailed:
                    std::wcout << L"[!] Injection failed for " << p.exeName
                               << L" (pid " << p.pid << L")\n";
                    break;
                default:
                    break;
            }
        }

        // Forget PIDs that have exited so a relaunch gets re-injected / re-diagnosed.
        for (auto it = injectedPids.begin(); it != injectedPids.end();) {
            it = (byPid.find(*it) == byPid.end()) ? injectedPids.erase(it) : std::next(it);
        }
        for (auto it = openFailWarned.begin(); it != openFailWarned.end();) {
            it = (byPid.find(*it) == byPid.end()) ? openFailWarned.erase(it) : std::next(it);
        }

        Sleep(2000);
    }
}

void HideOwnConsoleWindow() {
    HWND console = GetConsoleWindow();
    if (console) {
        ShowWindow(console, SW_HIDE);
    }
}

} // namespace

int wmain(int argc, wchar_t* argv[]) {
    if (argc >= 2 && std::wstring(argv[1]) == L"--list") {
        PrintWindowList();
        return 0;
    }

    if (argc >= 3 && std::wstring(argv[1]) == L"--config") {
        // Unattended mode: this is what vgwizard.exe's scheduled tasks run.
        // No console window should be left flashing on screen every logon.
        HideOwnConsoleWindow();

        std::wstring configPath = ResolveFullPath(argv[2]);
        Config cfg;
        if (!LoadConfig(configPath, &cfg) || cfg.targetNames.empty()) {
            // Nothing sensible to do - exit quietly, there's no console to
            // report to anyway in this mode.
            return 1;
        }

        // DLL resolved by convention: vghookNN.dll next to this exe, matching
        // this exe's own bitness.
#if defined(_WIN64)
        std::wstring dllPath = SelfDirectory() + L"vghook64.dll";
#else
        std::wstring dllPath = SelfDirectory() + L"vghook32.dll";
#endif
        if (GetFileAttributesW(dllPath.c_str()) == INVALID_FILE_ATTRIBUTES) {
            return 1;
        }

        RunWatchLoop(cfg.targetNames, cfg.treeMode, dllPath);
        return 0;
    }

    if (argc < 3) {
        std::wcerr << L"Usage: vgagentNN.exe [--tree] <dll_path> <exe_name_1> [<exe_name_2> ...]\n";
        std::wcerr << L"       vgagentNN.exe --list\n";
        std::wcerr << L"       vgagentNN.exe --config <config_path>\n";
        return 1;
    }

    int argi = 1;
    bool treeMode = false;
    if (std::wstring(argv[argi]) == L"--tree") {
        treeMode = true;
        ++argi;
    }
    if (argi + 1 >= argc) {
        std::wcerr << L"Usage: vgagentNN.exe [--tree] <dll_path> <exe_name_1> [<exe_name_2> ...]\n";
        return 1;
    }

    std::wstring dllPath = ResolveFullPath(argv[argi++]);
    std::unordered_set<std::wstring> targetNames;
    for (; argi < argc; ++argi) {
        targetNames.insert(ToLower(argv[argi]));
    }

    if (GetFileAttributesW(dllPath.c_str()) == INVALID_FILE_ATTRIBUTES) {
        std::wcerr << L"DLL not found: " << dllPath << L"\n";
        return 1;
    }

    RunWatchLoop(targetNames, treeMode, dllPath);
    return 0;
}