setup.cpp

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <shellapi.h>
#include <shlobj.h>
#include <shobjidl.h>
#include <tlhelp32.h>
#include <winreg.h>

#include <cstddef>
#include <filesystem>
#include <fstream>
#include <string>
#include <system_error>
#include <vector>

constexpr int IDR_APP_PAYLOAD = 101;

constexpr wchar_t APP_NAME[] = L"ClickGuard";
constexpr wchar_t APP_EXE_NAME[] = L"ClickGuard.exe";
constexpr wchar_t UNINSTALL_EXE_NAME[] = L"ClickGuard Uninstall.exe";
constexpr wchar_t UNINSTALL_LINK_NAME[] = L"Uninstall ClickGuard.lnk";
constexpr wchar_t UNINSTALL_REG_PATH[] =
    L"Software\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\ClickGuard";

std::wstring GetKnownFolder(REFKNOWNFOLDERID folderId) {
    PWSTR path = nullptr;
    const HRESULT hr = SHGetKnownFolderPath(
        folderId,
        KF_FLAG_DEFAULT,
        nullptr,
        &path
    );

    if (FAILED(hr) || !path) {
        return {};
    }

    std::wstring result(path);
    CoTaskMemFree(path);
    return result;
}

std::wstring GetModulePath() {
    std::wstring path(MAX_PATH, L'\0');
    const DWORD length = GetModuleFileNameW(
        nullptr,
        path.data(),
        static_cast<DWORD>(path.size())
    );

    path.resize(length);
    return path;
}

bool CreateShortcut(
    const std::wstring& shortcutPath,
    const std::wstring& targetPath,
    const std::wstring& arguments,
    const std::wstring& workingDirectory,
    const std::wstring& description
) {
    IShellLinkW* shellLink = nullptr;

    const HRESULT createHr = CoCreateInstance(
        CLSID_ShellLink,
        nullptr,
        CLSCTX_INPROC_SERVER,
        IID_IShellLinkW,
        reinterpret_cast<void**>(&shellLink)
    );

    if (FAILED(createHr) || !shellLink) {
        return false;
    }

    shellLink->SetPath(targetPath.c_str());
    shellLink->SetArguments(arguments.c_str());
    shellLink->SetWorkingDirectory(workingDirectory.c_str());
    shellLink->SetDescription(description.c_str());
    shellLink->SetIconLocation(targetPath.c_str(), 0);
    shellLink->SetShowCmd(SW_SHOWMINNOACTIVE);

    IPersistFile* persistFile = nullptr;
    const HRESULT queryHr = shellLink->QueryInterface(
        IID_IPersistFile,
        reinterpret_cast<void**>(&persistFile)
    );

    bool success = false;

    if (SUCCEEDED(queryHr) && persistFile) {
        success = SUCCEEDED(
            persistFile->Save(shortcutPath.c_str(), TRUE)
        );

        persistFile->Release();
    }

    shellLink->Release();
    return success;
}

void RemoveIfExists(const std::wstring& path) {
    std::error_code error;
    std::filesystem::remove(path, error);
}

bool CopyFileReplace(
    const std::wstring& sourcePath,
    const std::wstring& destinationPath
) {
    std::error_code error;
    std::filesystem::copy_file(
        sourcePath,
        destinationPath,
        std::filesystem::copy_options::overwrite_existing,
        error
    );

    return !error;
}

std::vector<std::byte> LoadPayloadResource() {
    const HRSRC resource = FindResourceW(
        nullptr,
        MAKEINTRESOURCEW(IDR_APP_PAYLOAD),
        RT_RCDATA
    );

    if (!resource) {
        return {};
    }

    const HGLOBAL loaded = LoadResource(nullptr, resource);
    if (!loaded) {
        return {};
    }

    const DWORD size = SizeofResource(nullptr, resource);
    const void* data = LockResource(loaded);

    if (!data || size == 0) {
        return {};
    }

    const auto* bytes = static_cast<const std::byte*>(data);
    return std::vector<std::byte>(bytes, bytes + size);
}

void StopRunningClickGuard() {
    const HANDLE snapshot = CreateToolhelp32Snapshot(
        TH32CS_SNAPPROCESS,
        0
    );

    if (snapshot == INVALID_HANDLE_VALUE) {
        return;
    }

    PROCESSENTRY32W entry{};
    entry.dwSize = sizeof(entry);

    if (!Process32FirstW(snapshot, &entry)) {
        CloseHandle(snapshot);
        return;
    }

    do {
        if (wcscmp(entry.szExeFile, APP_EXE_NAME) != 0) {
            continue;
        }

        HANDLE process = OpenProcess(
            PROCESS_TERMINATE,
            FALSE,
            entry.th32ProcessID
        );

        if (!process) {
            continue;
        }

        TerminateProcess(process, 0);
        CloseHandle(process);
    } while (Process32NextW(snapshot, &entry));

    CloseHandle(snapshot);
    Sleep(500);
}

bool WriteInstalledExe(const std::wstring& installDir) {
    std::error_code fsError;
    std::filesystem::create_directories(installDir, fsError);

    if (fsError) {
        return false;
    }

    const auto payload = LoadPayloadResource();
    if (payload.empty()) {
        return false;
    }

    const std::filesystem::path targetPath =
        std::filesystem::path(installDir) / APP_EXE_NAME;

    std::ofstream output(
        targetPath,
        std::ios::binary | std::ios::trunc
    );

    if (!output) {
        return false;
    }

    output.write(
        reinterpret_cast<const char*>(payload.data()),
        static_cast<std::streamsize>(payload.size())
    );

    return output.good();
}

bool WriteUninstallRegistry(
    const std::wstring& uninstallExePath,
    const std::wstring& installDir
) {
    HKEY uninstallKey = nullptr;
    const LONG createStatus = RegCreateKeyExW(
        HKEY_CURRENT_USER,
        UNINSTALL_REG_PATH,
        0,
        nullptr,
        0,
        KEY_WRITE,
        nullptr,
        &uninstallKey,
        nullptr
    );

    if (createStatus != ERROR_SUCCESS || !uninstallKey) {
        return false;
    }

    const std::wstring displayName = APP_NAME;
    const std::wstring uninstallString =
        L"\"" + uninstallExePath + L"\" /uninstall";
    const std::wstring publisher = L"drawliin";
    const DWORD noModify = 1;
    const DWORD noRepair = 1;

    auto setString = [&](const wchar_t* name, const std::wstring& value) {
        return RegSetValueExW(
            uninstallKey,
            name,
            0,
            REG_SZ,
            reinterpret_cast<const BYTE*>(value.c_str()),
            static_cast<DWORD>((value.size() + 1) * sizeof(wchar_t))
        ) == ERROR_SUCCESS;
    };

    const bool ok =
        setString(L"DisplayName", displayName) &&
        setString(L"DisplayIcon", uninstallExePath) &&
        setString(L"InstallLocation", installDir) &&
        setString(L"Publisher", publisher) &&
        setString(L"UninstallString", uninstallString) &&
        RegSetValueExW(
            uninstallKey,
            L"NoModify",
            0,
            REG_DWORD,
            reinterpret_cast<const BYTE*>(&noModify),
            sizeof(noModify)
        ) == ERROR_SUCCESS &&
        RegSetValueExW(
            uninstallKey,
            L"NoRepair",
            0,
            REG_DWORD,
            reinterpret_cast<const BYTE*>(&noRepair),
            sizeof(noRepair)
        ) == ERROR_SUCCESS;

    RegCloseKey(uninstallKey);
    return ok;
}

void RemoveUninstallRegistry() {
    RegDeleteTreeW(HKEY_CURRENT_USER, UNINSTALL_REG_PATH);
}

bool RunHiddenCommand(const std::wstring& commandLine) {
    STARTUPINFOW startupInfo{};
    startupInfo.cb = sizeof(startupInfo);
    startupInfo.dwFlags = STARTF_USESHOWWINDOW;
    startupInfo.wShowWindow = SW_HIDE;

    PROCESS_INFORMATION processInfo{};
    std::wstring mutableCommandLine = commandLine;

    const BOOL ok = CreateProcessW(
        nullptr,
        mutableCommandLine.data(),
        nullptr,
        nullptr,
        FALSE,
        CREATE_NO_WINDOW,
        nullptr,
        nullptr,
        &startupInfo,
        &processInfo
    );

    if (!ok) {
        return false;
    }

    CloseHandle(processInfo.hThread);
    CloseHandle(processInfo.hProcess);
    return true;
}

bool UninstallClickGuard(
    const std::wstring& installDir,
    const std::wstring& desktop,
    const std::wstring& programs,
    const std::wstring& startup,
    const std::wstring& currentModulePath
) {
    StopRunningClickGuard();

    RemoveIfExists(desktop + L"\\" + APP_NAME + L".lnk");
    RemoveIfExists(programs + L"\\" + APP_NAME + L".lnk");
    RemoveIfExists(programs + L"\\" + UNINSTALL_LINK_NAME);
    RemoveIfExists(startup + L"\\" + APP_NAME + L".lnk");
    RemoveUninstallRegistry();
    RemoveIfExists(installDir + L"\\" + APP_EXE_NAME);

    const std::wstring cleanupCommand =
        L"cmd.exe /c ping 127.0.0.1 -n 2 >nul && del /f /q \"" +
        currentModulePath +
        L"\" && rmdir /s /q \"" +
        installDir +
        L"\"";

    return RunHiddenCommand(cleanupCommand);
}

int WINAPI wWinMain(
    HINSTANCE,
    HINSTANCE,
    PWSTR cmdLine,
    int
) {
    const HRESULT initHr = CoInitializeEx(
        nullptr,
        COINIT_APARTMENTTHREADED
    );

    if (FAILED(initHr)) {
        MessageBoxW(
            nullptr,
            L"Failed to initialize setup.",
            L"ClickGuard Setup",
            MB_OK | MB_ICONERROR
        );
        return 1;
    }

    const std::wstring localAppData =
        GetKnownFolder(FOLDERID_LocalAppData);

    const std::wstring desktop =
        GetKnownFolder(FOLDERID_Desktop);

    const std::wstring programs =
        GetKnownFolder(FOLDERID_Programs);

    const std::wstring startup =
        GetKnownFolder(FOLDERID_Startup);

    if (
        localAppData.empty() ||
        desktop.empty() ||
        programs.empty() ||
        startup.empty()
    ) {
        MessageBoxW(
            nullptr,
            L"Failed to resolve Windows folders for installation.",
            L"ClickGuard Setup",
            MB_OK | MB_ICONERROR
        );
        CoUninitialize();
        return 1;
    }

    const std::wstring installDir =
        localAppData + L"\\" + APP_NAME;

    const std::wstring exePath =
        installDir + L"\\" + APP_EXE_NAME;

    const std::wstring uninstallExePath =
        installDir + L"\\" + UNINSTALL_EXE_NAME;

    const std::wstring selfPath = GetModulePath();
    const std::wstring commandLine =
        cmdLine ? std::wstring(cmdLine) : std::wstring();

    if (
        commandLine.find(L"/uninstall") != std::wstring::npos ||
        commandLine.find(L"-uninstall") != std::wstring::npos
    ) {
        const bool removed = UninstallClickGuard(
            installDir,
            desktop,
            programs,
            startup,
            selfPath
        );

        MessageBoxW(
            nullptr,
            removed
                ? L"ClickGuard uninstall was started.\n\nAny remaining cleanup will finish after this window closes."
                : L"ClickGuard could not be fully removed.",
            L"ClickGuard Uninstall",
            MB_OK | (removed ? MB_ICONINFORMATION : MB_ICONWARNING)
        );

        CoUninitialize();
        return removed ? 0 : 1;
    }

    StopRunningClickGuard();

    if (!WriteInstalledExe(installDir)) {
        MessageBoxW(
            nullptr,
            L"Failed to install ClickGuard.exe.",
            L"ClickGuard Setup",
            MB_OK | MB_ICONERROR
        );
        CoUninitialize();
        return 1;
    }

    if (!CopyFileReplace(selfPath, uninstallExePath)) {
        MessageBoxW(
            nullptr,
            L"Failed to create the ClickGuard uninstaller.",
            L"ClickGuard Setup",
            MB_OK | MB_ICONERROR
        );
        CoUninitialize();
        return 1;
    }

    const bool desktopOk = CreateShortcut(
        desktop + L"\\" + APP_NAME + L".lnk",
        exePath,
        L"",
        installDir,
        L"Launch ClickGuard"
    );

    const bool programsOk = CreateShortcut(
        programs + L"\\" + APP_NAME + L".lnk",
        exePath,
        L"",
        installDir,
        L"Launch ClickGuard"
    );

    const bool startupOk = CreateShortcut(
        startup + L"\\" + APP_NAME + L".lnk",
        exePath,
        L"",
        installDir,
        L"Launch ClickGuard at sign-in"
    );

    const bool uninstallShortcutOk = CreateShortcut(
        programs + L"\\" + UNINSTALL_LINK_NAME,
        uninstallExePath,
        L"/uninstall",
        installDir,
        L"Remove ClickGuard"
    );

    const bool uninstallRegistryOk = WriteUninstallRegistry(
        uninstallExePath,
        installDir
    );

    if (!desktopOk || !programsOk || !startupOk) {
        MessageBoxW(
            nullptr,
            L"ClickGuard was installed, but one or more app shortcuts could not be created.",
            L"ClickGuard Setup",
            MB_OK | MB_ICONWARNING
        );
    }

    if (!uninstallShortcutOk || !uninstallRegistryOk) {
        MessageBoxW(
            nullptr,
            L"ClickGuard was installed, but the uninstall entry could not be fully created.",
            L"ClickGuard Setup",
            MB_OK | MB_ICONWARNING
        );
    }

    ShellExecuteW(
        nullptr,
        L"open",
        exePath.c_str(),
        nullptr,
        installDir.c_str(),
        SW_SHOWMINNOACTIVE
    );

    MessageBoxW(
        nullptr,
        L"ClickGuard was installed.\n\nDesktop, Start Menu, Startup, and uninstall entries were created.",
        L"ClickGuard Setup",
        MB_OK | MB_ICONINFORMATION
    );

    CoUninitialize();
    return 0;
}