tray.cpp

// tray.cpp
//
// Builds to vgtray.exe - the piece that actually keeps protection running
// across a login session, and gives it a visible on/off switch. vgwizard.exe
// (after detecting/selecting which apps to protect and writing
// targets.txt) registers this exe in the HKCU Run key and launches it once
// immediately; from then on Windows starts it automatically at every logon.
//
// What it does:
//   - Adds a tray icon. Right-click (or left-click) for a menu: toggle
//     protection on/off, re-detect the session window, open the install
//     folder, or exit.
//   - Launches vgagent64.exe and vgagent32.exe as child processes (--config
//     mode, pointed at targets.txt next to this exe) so the actual
//     detection/injection logic stays in one place (agent.cpp) - this
//     is just the persistent shell around it with a UI.
//   - Owns a named, manual-reset event (see kProtectionEventName in
//     common.h) that every injected vghook.dll instance polls: signaled
//     means "hide," unsignaled means "reveal." Toggling Enable/Disable here
//     flips that event, so it affects windows that are *already* hidden
//     within about one poll interval (750ms) - not just future injections.
//   - Puts its child processes in a Job Object with
//     JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, so if this tray app exits (or
//     crashes) its children don't linger as orphans.
//
// This exists instead of a Scheduled Task (an earlier version of this
// project used one) because schtasks.exe registration reliably hit "ERROR:
// Access is denied" in testing once these binaries had already been
// flagged once by Windows Defender - Task Scheduler persistence is a
// well-known technique behavior-monitoring watches for. An HKCU Run key
// entry is a plain registry write, the same mechanism ordinary consumer
// "start with Windows" checkboxes use, and involves no schtasks.exe/Task
// Scheduler API call at all. See ARCHITECTURE.md for the full writeup.

#include <windows.h>
#include <shellapi.h>

#include "common.h"

#include <string>
#include <unordered_set>
#include <vector>

using namespace vg_common;

namespace {

constexpr UINT WM_TRAYICON = WM_APP + 1;
constexpr UINT ID_TRAY_TOGGLE = 1001;
constexpr UINT ID_TRAY_REDETECT = 1002;
constexpr UINT ID_TRAY_OPENFOLDER = 1003;
constexpr UINT ID_TRAY_EXIT = 1004;

HWND g_hwnd = nullptr;
NOTIFYICONDATAW g_nid = {};
HANDLE g_job = nullptr;
HANDLE g_protectionEvent = nullptr;
bool g_protectionEnabled = true;
UINT g_taskbarCreatedMsg = 0;

std::wstring g_installDir;
std::wstring g_configPath;

PROCESS_INFORMATION g_agent64 = {};
PROCESS_INFORMATION g_agent32 = {};
bool g_agent64Running = false;
bool g_agent32Running = false;

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

void TerminateChildren() {
    if (g_agent64Running) {
        TerminateProcess(g_agent64.hProcess, 0);
        CloseHandle(g_agent64.hThread);
        CloseHandle(g_agent64.hProcess);
        g_agent64Running = false;
    }
    if (g_agent32Running) {
        TerminateProcess(g_agent32.hProcess, 0);
        CloseHandle(g_agent32.hThread);
        CloseHandle(g_agent32.hProcess);
        g_agent32Running = false;
    }
}

// Watcher paths and targets.txt never contain literal quote characters (they
// come from %LOCALAPPDATA%, a fixed install location), so simple wrapping
// is enough here - no need for the full argv-escaping algorithm wizard.cpp
// uses for arbitrary strings.
bool LaunchChild(const std::wstring& exePath, PROCESS_INFORMATION* outPi) {
    std::wstring cmdline = L"\"" + exePath + L"\" --config \"" + g_configPath + L"\"";
    std::vector<wchar_t> mutableCmd(cmdline.begin(), cmdline.end());
    mutableCmd.push_back(L'\0');

    STARTUPINFOW si = {};
    si.cb = sizeof(si);
    BOOL ok = CreateProcessW(exePath.c_str(), mutableCmd.data(), nullptr, nullptr,
                              FALSE, CREATE_NO_WINDOW, nullptr, nullptr, &si, outPi);
    if (ok && g_job) {
        AssignProcessToJobObject(g_job, outPi->hProcess);
    }
    return ok != FALSE;
}

void StartChildren() {
    TerminateChildren(); // safe even if nothing was running
    if (!FileExists(g_configPath)) return; // nothing configured yet
    g_agent64Running = LaunchChild(g_installDir + L"vgagent64.exe", &g_agent64);
    g_agent32Running = LaunchChild(g_installDir + L"vgagent32.exe", &g_agent32);
}

void UpdateTrayVisual() {
    const wchar_t* tip = g_protectionEnabled
                              ? L"ViewGuard - Protection ON"
                              : L"ViewGuard - Protection OFF (paused)";
    lstrcpynW(g_nid.szTip, tip, ARRAYSIZE(g_nid.szTip));
    g_nid.hIcon = LoadIconW(nullptr, g_protectionEnabled ? IDI_APPLICATION : IDI_WARNING);
    Shell_NotifyIconW(NIM_MODIFY, &g_nid);
}

void ToggleProtection() {
    g_protectionEnabled = !g_protectionEnabled;
    if (g_protectionEvent) {
        if (g_protectionEnabled) {
            SetEvent(g_protectionEvent);
        } else {
            ResetEvent(g_protectionEvent);
        }
    }
    UpdateTrayVisual();
}

void RunReDetect() {
    std::vector<AppEntry> detected = DetectKnownApps(g_installDir);
    if (detected.empty()) {
        MessageBoxW(g_hwnd,
                    L"No known Citrix/AVD session window is currently open.\n\n"
                    L"Open your session and try Re-detect again, or add its exe "
                    L"name to known_names.txt next to vgtray.exe if it's a "
                    L"client version this tool doesn't recognize yet.",
                    L"ViewGuard - Re-detect", MB_OK | MB_ICONINFORMATION);
        return;
    }

    std::unordered_set<std::wstring> names;
    std::wstring list;
    for (const auto& app : detected) {
        names.insert(app.exeName);
        list += L"  - " + app.exeName + L"\n";
    }

    if (!WriteTargetsConfig(g_configPath, names, /*treeMode=*/true)) {
        MessageBoxW(g_hwnd, L"Detected apps, but failed to write targets.txt.",
                    L"ViewGuard - Re-detect", MB_OK | MB_ICONERROR);
        return;
    }

    StartChildren(); // restart so the watchers pick up the new list - they don't hot-reload

    std::wstring msg = L"Now protecting:\n\n" + list;
    MessageBoxW(g_hwnd, msg.c_str(), L"ViewGuard - Re-detected", MB_OK | MB_ICONINFORMATION);
}

void ShowContextMenu(HWND hwnd) {
    POINT pt;
    GetCursorPos(&pt);

    HMENU menu = CreatePopupMenu();
    AppendMenuW(menu, MF_STRING | (g_protectionEnabled ? MF_CHECKED : MF_UNCHECKED),
                ID_TRAY_TOGGLE, L"Protection Enabled");
    AppendMenuW(menu, MF_SEPARATOR, 0, nullptr);
    AppendMenuW(menu, MF_STRING, ID_TRAY_REDETECT, L"Re-detect now");
    AppendMenuW(menu, MF_STRING, ID_TRAY_OPENFOLDER, L"Open install folder");
    AppendMenuW(menu, MF_SEPARATOR, 0, nullptr);
    AppendMenuW(menu, MF_STRING, ID_TRAY_EXIT, L"Exit (stop protection)");

    // Required so the popup menu dismisses correctly if the user clicks
    // elsewhere - documented Windows tray-icon quirk, not optional.
    SetForegroundWindow(hwnd);
    TrackPopupMenu(menu, TPM_RIGHTBUTTON, pt.x, pt.y, 0, hwnd, nullptr);
    PostMessageW(hwnd, WM_NULL, 0, 0);
    DestroyMenu(menu);
}

LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam) {
    if (msg == g_taskbarCreatedMsg && g_taskbarCreatedMsg != 0) {
        // Explorer restarted and wiped every tray icon - re-add ours.
        Shell_NotifyIconW(NIM_ADD, &g_nid);
        return 0;
    }

    switch (msg) {
        case WM_TRAYICON:
            if (lparam == WM_RBUTTONUP || lparam == WM_LBUTTONUP || lparam == WM_CONTEXTMENU) {
                ShowContextMenu(hwnd);
            }
            return 0;

        case WM_COMMAND:
            switch (LOWORD(wparam)) {
                case ID_TRAY_TOGGLE:
                    ToggleProtection();
                    break;
                case ID_TRAY_REDETECT:
                    RunReDetect();
                    break;
                case ID_TRAY_OPENFOLDER:
                    ShellExecuteW(nullptr, L"open", g_installDir.c_str(), nullptr, nullptr,
                                  SW_SHOWNORMAL);
                    break;
                case ID_TRAY_EXIT:
                    DestroyWindow(hwnd);
                    break;
                default:
                    break;
            }
            return 0;

        case WM_DESTROY:
            TerminateChildren();
            Shell_NotifyIconW(NIM_DELETE, &g_nid);
            if (g_protectionEvent) CloseHandle(g_protectionEvent);
            if (g_job) CloseHandle(g_job); // JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE as a backstop
            PostQuitMessage(0);
            return 0;

        default:
            break;
    }
    return DefWindowProcW(hwnd, msg, wparam, lparam);
}

} // namespace

int APIENTRY wWinMain(HINSTANCE hInstance, HINSTANCE, PWSTR, int) {
    // Single instance: the Run key launches this at every logon, and
    // vgwizard.exe also launches it directly after setup - without this guard
    // a re-run of vgwizard.exe would leave two tray icons and two sets of
    // watcher children behind.
    HANDLE instanceMutex = CreateMutexW(nullptr, TRUE, L"Local\\ViewGuardTrayMutex");
    if (!instanceMutex || GetLastError() == ERROR_ALREADY_EXISTS) {
        return 0;
    }

    g_installDir = SelfDirectory();
    g_configPath = g_installDir + L"targets.txt";

    // Manual-reset, initially signaled = protection on by default. Every
    // vghook.dll instance opens this by the same well-known name and
    // polls it - see hook.cpp.
    g_protectionEvent = CreateEventW(nullptr, TRUE, TRUE, kProtectionEventName);

    g_job = CreateJobObjectW(nullptr, nullptr);
    if (g_job) {
        JOBOBJECT_EXTENDED_LIMIT_INFORMATION jeli = {};
        jeli.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
        SetInformationJobObject(g_job, JobObjectExtendedLimitInformation, &jeli, sizeof(jeli));
    }

    g_taskbarCreatedMsg = RegisterWindowMessageW(L"TaskbarCreated");

    WNDCLASSW wc = {};
    wc.lpfnWndProc = WndProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = kTrayWindowClassName;
    RegisterClassW(&wc);

    // Never shown - it only exists to own the tray icon and receive its
    // callback/menu messages.
    g_hwnd = CreateWindowW(kTrayWindowClassName, L"ViewGuard", WS_OVERLAPPEDWINDOW, 0, 0, 0, 0,
                            nullptr, nullptr, hInstance, nullptr);

    g_nid.cbSize = sizeof(g_nid);
    g_nid.hWnd = g_hwnd;
    g_nid.uID = 1;
    g_nid.uFlags = NIF_ICON | NIF_MESSAGE | NIF_TIP;
    g_nid.uCallbackMessage = WM_TRAYICON;
    Shell_NotifyIconW(NIM_ADD, &g_nid);
    UpdateTrayVisual();

    StartChildren();

    MSG msg;
    BOOL got;
    while ((got = GetMessageW(&msg, nullptr, 0, 0)) > 0) {
        TranslateMessage(&msg);
        DispatchMessageW(&msg);
    }

    ReleaseMutex(instanceMutex);
    CloseHandle(instanceMutex);
    return 0;
}