wizard.cpp

// wizard.cpp
//
// Default behavior (no arguments): AUTO-DETECTS Citrix/AVD session windows
// currently open, using the known-process-name list in common.h (extendable
// via known_names.txt - see there), and sets up protection for whatever it
// finds with no prompts. This is meant to be the "one tool, run it, done"
// path - it does not require the user to know or pick a process name.
//
// `--manual` falls back to the original behavior: lists EVERY visible
// window (not just recognized Citrix/AVD ones) and lets you pick, for
// cases where you want to protect something outside the known list, or the
// known list hasn't caught up with a newer client version yet.
//
// Either path writes a config file, then hands off persistence and live
// control to vgtray.exe: a small tray-icon app (see tray.cpp) that
// launches at logon via an HKCU Run key, keeps vgagent32.exe/vgagent64.exe
// running as its children, and lets protection be toggled on/off with a
// right-click - no scheduled tasks involved.
//
// Usage:
//   vgwizard.exe              auto-detect known Citrix/AVD windows, no prompts
//   vgwizard.exe --manual     pick from every visible window instead
//   vgwizard.exe --uninstall  removes the Run key entry and installed files
//
// Deliberately NOT a literal Windows Service (SCM-registered, Session-0
// background service). A true service runs in Session 0, which is isolated
// from your interactive desktop and cannot see or inject into windows in
// your logon session at all (this has been true since Session 0 isolation
// was introduced) - making it useless for this purpose without a much more
// involved cross-session token/impersonation dance. vgtray.exe launching
// via an HKCU Run key runs inside your own interactive session instead,
// which is both simpler and is what actually works here.
//
// Also deliberately NOT a Scheduled Task (an earlier version of this tool
// used one): registering a task via schtasks.exe turned out to reliably hit
// "ERROR: Access is denied" on a real test machine once these binaries had
// already been flagged once by Windows Defender - Task Scheduler
// registration is exactly the kind of persistence mechanism behavior
// monitoring watches for (it's a well-known malware technique, MITRE
// ATT&CK T1053.005), and a folder-level AV exclusion didn't stop it. A Run
// key value is a plain registry write - the same "start with Windows"
// mechanism ordinary consumer software uses - and doesn't go through
// schtasks.exe or the Task Scheduler API at all.

#include "common.h"

#include <iostream>
#include <unordered_set>
#include <vector>

using namespace vg_common;

namespace {

const wchar_t* kRunKeyPath = L"Software\\Microsoft\\Windows\\CurrentVersion\\Run";
const wchar_t* kRunValueName = L"ViewGuard";

std::wstring GetEnvVar(const wchar_t* name) {
    wchar_t buf[MAX_PATH] = {0};
    DWORD len = GetEnvironmentVariableW(name, buf, MAX_PATH);
    if (len == 0 || len >= MAX_PATH) return L"";
    return std::wstring(buf, len);
}

// Standard Windows argv-quoting algorithm: appends `arg` to `out` as a
// single correctly-escaped command-line token, including embedded quotes
// and backslashes. This is the inverse of CommandLineToArgvW's parsing.
void AppendQuotedArg(std::wstring& out, const std::wstring& arg) {
    if (!out.empty()) out += L' ';
    if (!arg.empty() && arg.find_first_of(L" \t\n\v\"") == std::wstring::npos) {
        out += arg;
        return;
    }
    out += L'"';
    for (auto it = arg.begin();; ++it) {
        size_t backslashes = 0;
        while (it != arg.end() && *it == L'\\') {
            ++backslashes;
            ++it;
        }
        if (it == arg.end()) {
            out.append(backslashes * 2, L'\\');
            break;
        } else if (*it == L'"') {
            out.append(backslashes * 2 + 1, L'\\');
            out += L'"';
        } else {
            out.append(backslashes, L'\\');
            out += *it;
        }
    }
    out += L'"';
}

// Writes/overwrites the HKCU Run value that launches vgtray.exe at
// logon. HKCU never requires elevation, unlike HKLM's Run key.
bool SetRunKeyEntry(const std::wstring& trayExePath) {
    HKEY key = nullptr;
    if (RegCreateKeyExW(HKEY_CURRENT_USER, kRunKeyPath, 0, nullptr, 0,
                         KEY_SET_VALUE, nullptr, &key, nullptr) != ERROR_SUCCESS) {
        return false;
    }
    std::wstring quoted;
    AppendQuotedArg(quoted, trayExePath);
    LONG rc = RegSetValueExW(
        key, kRunValueName, 0, REG_SZ,
        reinterpret_cast<const BYTE*>(quoted.c_str()),
        static_cast<DWORD>((quoted.size() + 1) * sizeof(wchar_t)));
    RegCloseKey(key);
    return rc == ERROR_SUCCESS;
}

void RemoveRunKeyEntry() {
    HKEY key = nullptr;
    if (RegOpenKeyExW(HKEY_CURRENT_USER, kRunKeyPath, 0, KEY_SET_VALUE, &key) == ERROR_SUCCESS) {
        RegDeleteValueW(key, kRunValueName);
        RegCloseKey(key);
    }
}

// Fire-and-forget launch: vgtray.exe is a GUI-subsystem app, so this
// doesn't open any visible window besides the tray icon it adds itself.
bool LaunchTray(const std::wstring& trayExePath) {
    std::wstring cmdline;
    AppendQuotedArg(cmdline, trayExePath);
    std::vector<wchar_t> mutableCmd(cmdline.begin(), cmdline.end());
    mutableCmd.push_back(L'\0');

    STARTUPINFOW si = {};
    si.cb = sizeof(si);
    PROCESS_INFORMATION pi = {};
    BOOL ok = CreateProcessW(trayExePath.c_str(), mutableCmd.data(), nullptr, nullptr,
                              FALSE, 0, nullptr, nullptr, &si, &pi);
    if (ok) {
        CloseHandle(pi.hThread);
        CloseHandle(pi.hProcess);
    }
    return ok != FALSE;
}

// Asks a currently-running vgtray.exe to exit cleanly (which also tears
// down its child vgagent32.exe/vgagent64.exe watchers via the job object it
// puts them in - see tray.cpp). Used during --uninstall so nothing
// keeps running out of a folder that's about to be deleted.
void CloseRunningTray() {
    HWND wnd = FindWindowW(kTrayWindowClassName, nullptr);
    if (wnd) {
        PostMessageW(wnd, WM_CLOSE, 0, 0);
        Sleep(500); // give it a moment to tear down its children before we delete files
    }
}

std::vector<int> ParseSelection(const std::wstring& line, int maxIndex) {
    std::vector<int> result;
    std::wstring current;
    auto flush = [&]() {
        if (current.empty()) return;
        try {
            int n = std::stoi(current);
            if (n >= 1 && n <= maxIndex) result.push_back(n);
        } catch (...) {
        }
        current.clear();
    };
    for (wchar_t c : line) {
        if (c == L',' || c == L' ') {
            flush();
        } else {
            current += c;
        }
    }
    flush();
    return result;
}

bool FileExists(const std::wstring& path) {
    return GetFileAttributesW(path.c_str()) != INVALID_FILE_ATTRIBUTES;
}

// Shared by both auto-detect and manual flows: copies the runtime files to
// a stable per-user location, writes the config, registers the Run key, and
// launches vgtray.exe so protection is active immediately. Returns true
// on success.
bool InstallAndActivate(const std::unordered_set<std::wstring>& chosenNames,
                         bool treeMode) {
    std::wstring selfDir = SelfDirectory();
    std::vector<std::wstring> required = {L"vgagent32.exe", L"vgagent64.exe",
                                           L"vghook32.dll", L"vghook64.dll",
                                           L"vgtray.exe"};
    for (const auto& f : required) {
        if (!FileExists(selfDir + f)) {
            std::wcout << L"\n[!] Missing " << f << L" next to vgwizard.exe ("
                       << selfDir << L") - run the wizard from the extracted "
                          L"ViewGuard folder with all files present.\n";
            return false;
        }
    }

    std::wstring localAppData = GetEnvVar(L"LOCALAPPDATA");
    if (localAppData.empty()) {
        std::wcout << L"\n[!] Could not resolve %LOCALAPPDATA% - aborting.\n";
        return false;
    }
    std::wstring installDir = localAppData + L"\\ViewGuard\\";
    CreateDirectoryW(installDir.c_str(), nullptr); // ignores ERROR_ALREADY_EXISTS fine

    for (const auto& f : required) {
        if (!CopyFileW((selfDir + f).c_str(), (installDir + f).c_str(), FALSE)) {
            std::wcout << L"\n[!] Failed to copy " << f << L" to " << installDir
                       << L" (error " << GetLastError() << L")\n";
            return false;
        }
    }

    std::wstring configPath = installDir + L"targets.txt";
    if (!WriteTargetsConfig(configPath, chosenNames, treeMode)) {
        std::wcout << L"\n[!] Failed to write " << configPath << L" - aborting.\n";
        return false;
    }

    std::wcout << L"\nInstalling to " << installDir << L" ...\n";

    // Restart vgtray.exe if one's already running, so it picks up the
    // fresh config instead of continuing to watch the old target list.
    CloseRunningTray();

    std::wstring trayPath = installDir + L"vgtray.exe";
    bool runKeyOk = SetRunKeyEntry(trayPath);
    std::wcout << L"  Run at logon: " << (runKeyOk ? L"registered" : L"FAILED") << L"\n";
    if (!runKeyOk) {
        std::wcout << L"    Could not write the HKCU Run key (error " << GetLastError()
                   << L"). Protection won't start automatically at your next logon - "
                      L"you can still start it manually by running vgtray.exe from "
                   << installDir << L".\n";
    }

    bool launchOk = LaunchTray(trayPath);
    std::wcout << L"  Starting protection now: " << (launchOk ? L"done" : L"FAILED") << L"\n";
    if (!launchOk) {
        std::wcout << L"    Could not launch vgtray.exe (error " << GetLastError()
                   << L"). Try running it directly from " << installDir << L".\n";
    }

    if (!runKeyOk && !launchOk) {
        std::wcout << L"\n[!] Nothing is protected right now. See the errors above.\n";
        return false;
    }

    std::wcout << L"\nDone. Look for the ViewGuard icon in your system tray - "
                  L"right-click it any time to toggle protection on/off, re-detect, "
                  L"or exit. It'll relaunch automatically every time you log on.\n";
    std::wcout << L"Config: " << configPath << L"\n";
    std::wcout << L"Run `vgwizard.exe` again any time to re-detect and update your "
                  L"selection, `--manual` to pick something else, or `--uninstall` "
                  L"to remove everything.\n";
    return true;
}

// Default flow: no prompts. Detect known Citrix/AVD windows and protect them.
void RunAutoSetup() {
    std::wcout << L"=== ViewGuard - automatic setup ===\n\n";

    std::wstring selfDir = SelfDirectory();
    std::vector<AppEntry> detected = DetectKnownApps(selfDir);

    if (detected.empty()) {
        std::wcout
            << L"No known Citrix/AVD session window is currently open, so "
               L"there's nothing to protect yet.\n\n"
               L"Open your Citrix or AVD session and run vgwizard.exe again - "
               L"it'll pick it up automatically. If it's a client version "
               L"this tool doesn't recognize, add its exe name to "
               L"known_names.txt next to vgwizard.exe (one name per line), or "
               L"run `vgwizard.exe --manual` to pick from every open window "
               L"yourself.\n";
        return;
    }

    std::wcout << L"Detected:\n";
    std::unordered_set<std::wstring> chosenNames;
    for (const auto& app : detected) {
        chosenNames.insert(app.exeName);
        std::wcout << L"  - " << app.exeName << L"   \"" << app.sampleTitle << L"\"\n";
    }
    std::wcout << L"\nProtecting these automatically (no confirmation needed - "
                  L"re-run with --manual if you want different apps).\n";

    // Tree mode defaults on here: we don't know in advance whether a given
    // client version renders through a helper/child process, and auto mode
    // is meant to just work without you having to know that.
    InstallAndActivate(chosenNames, /*treeMode=*/true);
}

// --manual: original behavior - list everything, let the user pick.
void RunManualSetup() {
    std::wcout << L"=== ViewGuard - manual setup ===\n\n";

    std::vector<AppEntry> apps = DistinctVisibleApps();
    if (apps.empty()) {
        std::wcout << L"No visible windows found. Is anything open?\n";
        return;
    }

    std::wcout << L"Currently visible apps:\n\n";
    for (size_t i = 0; i < apps.size(); ++i) {
        std::wcout << L"  " << (i + 1) << L") " << apps[i].exeName
                   << L"   \"" << apps[i].sampleTitle << L"\"";
        if (apps[i].windowCount > 1) {
            std::wcout << L"  (+" << (apps[i].windowCount - 1) << L" more window"
                       << (apps[i].windowCount - 1 > 1 ? L"s" : L"") << L")";
        }
        std::wcout << L"\n";
    }

    std::wcout << L"\nEnter the numbers of the apps to protect, comma-separated "
                  L"(e.g. 1,3,5): ";
    std::wstring selectionLine;
    std::getline(std::wcin, selectionLine);
    std::vector<int> selected = ParseSelection(selectionLine, static_cast<int>(apps.size()));

    if (selected.empty()) {
        std::wcout << L"Nothing selected - exiting without changes.\n";
        return;
    }

    std::wcout << L"\nAlso protect child/helper processes these apps spawn? "
                  L"[y/N]: ";
    std::wstring treeAnswer;
    std::getline(std::wcin, treeAnswer);
    bool treeMode = !treeAnswer.empty() && (treeAnswer[0] == L'y' || treeAnswer[0] == L'Y');

    std::wcout << L"\nSelected:\n";
    std::unordered_set<std::wstring> chosenNames;
    for (int idx : selected) {
        const auto& entry = apps[idx - 1];
        chosenNames.insert(entry.exeName);
        std::wcout << L"  - " << entry.exeName << L"\n";
    }

    InstallAndActivate(chosenNames, treeMode);
}

void RunUninstall() {
    std::wcout << L"This will stop protection, remove the ViewGuard tray icon and "
                  L"its logon entry, and delete installed files under "
                  L"%LOCALAPPDATA%\\ViewGuard. Continue? [y/N]: ";
    std::wstring answer;
    std::getline(std::wcin, answer);
    if (answer.empty() || (answer[0] != L'y' && answer[0] != L'Y')) {
        std::wcout << L"Cancelled.\n";
        return;
    }

    CloseRunningTray();
    RemoveRunKeyEntry();

    std::wstring localAppData = GetEnvVar(L"LOCALAPPDATA");
    if (!localAppData.empty()) {
        std::wstring installDir = localAppData + L"\\ViewGuard\\";
        for (const wchar_t* f :
             {L"vgagent32.exe", L"vgagent64.exe", L"vghook32.dll", L"vghook64.dll",
              L"vgtray.exe", L"targets.txt"}) {
            DeleteFileW((installDir + f).c_str());
        }
        RemoveDirectoryW(installDir.c_str());
    }

    std::wcout << L"Uninstalled.\n";
}

// vgwizard.exe is a console-subsystem app. When launched by the installer's
// finish page (or by double-clicking it in Explorer) rather than from an
// already-open terminal, Windows allocates it a brand-new console window -
// and that window is destroyed the instant the process exits. Without this,
// RunAutoSetup()/RunManualSetup() can finish in well under a second, so the
// window opens and closes too fast to ever read, which looks identical to
// "nothing happened." Pausing for a keypress keeps it on screen.
//
// Deliberately NOT applied to --uninstall: that path is invoked
// synchronously by the NSIS uninstaller (ExecWait) as one step in an
// automated flow, and already has its own y/n confirmation prompt: adding a
// second blocking pause after it finishes would leave the uninstaller stuck
// waiting on a keypress nobody's watching for.
void PauseBeforeExit() {
    std::wcout << L"\nPress Enter to close this window...";
    std::wstring dummy;
    std::getline(std::wcin, dummy);
}

} // namespace

int wmain(int argc, wchar_t* argv[]) {
    std::wstring mode = argc >= 2 ? argv[1] : L"";
    if (mode == L"--uninstall") {
        RunUninstall();
    } else if (mode == L"--manual") {
        RunManualSetup();
        PauseBeforeExit();
    } else {
        RunAutoSetup();
        PauseBeforeExit();
    }
    return 0;
}