// BSSO.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include <iostream>
#include <wrl.h>
#include <windows.foundation.h>
#include <windows.security.authentication.web.core.h>
#include <lm.h>
#include <proofofpossessioncookieinfo.h>

#pragma comment(lib, "runtimeobject.lib")
#pragma comment(lib, "netapi32.lib")

using namespace Microsoft::WRL;
using namespace Microsoft::WRL::Wrappers;
using namespace ABI::Windows::Security::Authentication::Web::Core;
using namespace ABI::Windows::Security::Credentials;

bool IsWindowsSSOCapable();
bool IsHostNameSSOCapable(LPCWSTR hostWithScheme);
void DumpSSOArtifacts(LPCWSTR fullUrl);

int main()
{
    std::cout << "IsWindowsSSOCapable: " << std::boolalpha << IsWindowsSSOCapable() << std::endl;
    std::cout << "IsHostName SSO Capable (https://www.office.com): " << IsHostNameSSOCapable(L"https://www.office.com") << std::endl;
    std::cout << "IsHostName SSO Capable (https://login.live.com): " << IsHostNameSSOCapable(L"https://login.live.com") << std::endl;
    std::cout << "IsHostName SSO Capable (https://login.microsoftonline.com): " << IsHostNameSSOCapable(L"https://login.microsoftonline.com") << std::endl;
    DumpSSOArtifacts(L"https://login.microsoftonline.com/common/oauth2/authorize?client_id=4345a7b9-9a63-4910-a426-35363201d503&redirect_uri=https%3A%2F%2Fwww.office.com%2Flanding&response_type=code%20id_token&scope=openid%20profile&response_mode=form_post&nonce=637606044728492219.NzE4ZTM3MzMtNzFkZi00ZjMyLThmYzYtMGUxMTdhNWNkMmUzYWE4OTM0N2YtZjJkMS00ZmJjLWI0YzEtNGJhNTMyNjlhMTg1&ui_locales=en-US&mkt=en-US&state=W0fX569Ah_dknKwvc2R-tgcZBfwBqHpVBSW5HRHXSg6-Nt_7fFL5AQaqjQAMUJYRdgPFb2H-xFToII8YmJONvgBNWx8sheflTJGpPR1A3hYKd_gmJEk2Td_kDcquzsbS_g5jHQ_D01AeEWFYgzyRUEPLfhsdB7-JqgNdIZ99NKGodX5Kf8-hMTz7I6CP34xnZ1QYnIvP7dhrq-3rqthF04J3i-n7-rKNGJcLIXHkdFyG3rnAGt1bepwVauuC6wEt3iBZ5zj9Xt5m_ndIoy96MQ&x-client-SKU=ID_NETSTANDARD2_0&x-client-ver=6.8.0.0#");
    DumpSSOArtifacts(L"https://login.live.com/login.srf?wa=wsignin1.0&rpsnv=13&ct=1625007720&rver=7.0.6737.0&wp=MBI_SSL&wreply=https%3a%2f%2foutlook.live.com%2fowa%2f%3fnlp%3d1%26RpsCsrfState%3dae3545d5-d1b4-517a-a807-9a503f122ce9&id=292841&aadredir=1&CBCXT=out&lw=1&fl=dob%2cflname%2cwld&cobrandid=90015");
}

bool IsMSASSOCapable();
bool IsAADSSOCabable();

bool IsWindowsSSOCapable()
{
    return IsMSASSOCapable() || IsAADSSOCabable();
}

class CoInitializeWrapper
{
    HRESULT _hr;
public:
    CoInitializeWrapper(DWORD flags)
    {

        _hr = ::CoInitializeEx(nullptr, flags);
    }

    ~CoInitializeWrapper()
    {
        if (SUCCEEDED(_hr))
        {
            ::CoUninitialize();
        }
    }

    operator HRESULT()
    {
        return _hr;
    }
};

bool IsHostNameSSOCapable(LPCWSTR hostWithScheme)
{
    CoInitializeWrapper coInit(COINIT_MULTITHREADED);

    DWORD foundTokenCount = 0;
    ProofOfPossessionCookieInfo* tokens = nullptr;
    ComPtr<IProofOfPossessionCookieInfoManager> pCookieInfoManager;
    HRESULT hr = CoCreateInstance(__uuidof(ProofOfPossessionCookieInfoManager), nullptr, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&pCookieInfoManager));
    if (SUCCEEDED(hr))
    {
        hr = pCookieInfoManager->GetCookieInfoForUri(hostWithScheme, &foundTokenCount, &tokens);
        if (SUCCEEDED(hr))
        {
            FreeProofOfPossessionCookieInfoArray(tokens, foundTokenCount);
            return foundTokenCount != 0;
        }
    }

    return false;
}


void DumpSSOArtifacts(LPCWSTR fullUrl)
{
    CoInitializeWrapper coInit(COINIT_MULTITHREADED);

    DWORD foundTokenCount = 0;
    ProofOfPossessionCookieInfo* tokens = nullptr;
    ComPtr<IProofOfPossessionCookieInfoManager> pCookieInfoManager;
    HRESULT hr = CoCreateInstance(__uuidof(ProofOfPossessionCookieInfoManager), nullptr, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&pCookieInfoManager));
    if (SUCCEEDED(hr))
    {
        hr = pCookieInfoManager->GetCookieInfoForUri(fullUrl, &foundTokenCount, &tokens);
        if (SUCCEEDED(hr))
        {
            for (DWORD i = 0; i < foundTokenCount; i++)
            {
                if (_wcsnicmp(tokens[i].name, L"x-ms-", _countof(L"x-ms-")) == 0)
                {
                    std::wcout << "Add as a header: \n\t" << tokens[i].name << ": " << tokens[i].data << std::endl;
                }
                else
                {
                    std::wcout << "Add as a cookie: \n\t" << tokens[i].name << "==" << tokens[i].data << tokens[i].p3pHeader << std::endl;
                    std::wcout << "\tP3P: " << tokens[i].p3pHeader << std::endl;
                    std::wcout << "\tdwFlags for InternetSetCookieEx: "<< std::hex << tokens[i].flags<< std::endl;
                    // https://docs.microsoft.com/en-us/windows/win32/api/wininet/nf-wininet-internetsetcookieexa
                }
            }

            FreeProofOfPossessionCookieInfoArray(tokens, foundTokenCount);
        }
    }
}

template<
    typename InnerType,
    typename OutThing>
    HRESULT AwaitOperationResult(
        ABI::Windows::Foundation::IAsyncOperation<InnerType> *pOperation,
        OutThing &&out);

bool IsMSASSOCapable()
{
    Microsoft::WRL::Wrappers::RoInitializeWrapper initialize(RO_INIT_MULTITHREADED);

    HRESULT hr;
    ComPtr<IWebAuthenticationCoreManagerStatics> authenticationCoreManager;
    ComPtr<ABI::Windows::Foundation::IAsyncOperation<WebAccountProvider*>> findAccountProviderOperation;
    ComPtr<IWebAccountProvider> webProvider;
    ComPtr<IWebAccountProvider2> webProvider2;

    hr = ABI::Windows::Foundation::GetActivationFactory(
        HStringReference(RuntimeClass_Windows_Security_Authentication_Web_Core_WebAuthenticationCoreManager).Get(),
        &authenticationCoreManager);
    if (FAILED(hr))
    {
        return false;
    }

    hr = authenticationCoreManager->FindAccountProviderAsync(
        HStringReference(L"https://login.windows.local").Get(),
        &findAccountProviderOperation);
    
    if (FAILED(hr))
    {
        return false;
    }

    hr = AwaitOperationResult(findAccountProviderOperation.Get(), &webProvider);
    if (FAILED(hr))
    {
        return false;
    }

    if (webProvider == nullptr)
    {
        return false;
    }

    webProvider.As(&webProvider2);

    if (webProvider2 == nullptr)
    {
        return false;
    }

    HString value;
    hr = webProvider2->get_Authority(value.GetAddressOf());
    
    if (FAILED(hr))
    {
        return false;
    }

    LPCWSTR authority = value.GetRawBuffer(nullptr);

    return authority != nullptr && _wcsicmp(L"consumers", authority) == 0;
}

bool IsAADSSOCabable()
{
    PDSREG_JOIN_INFO pJoinInfo =  nullptr;
    HRESULT hr;
    hr = NetGetAadJoinInformation(nullptr, &pJoinInfo);

    if (hr != S_OK)
    {
        return false;
    }

    bool isAADJoined = pJoinInfo->joinType != DSREG_UNKNOWN_JOIN;

    if (pJoinInfo)
    {
        NetFreeAadJoinInformation(pJoinInfo);
    }

    return isAADJoined;
}

template<typename T> struct OperationCallback : public Microsoft::WRL::RuntimeClass<
    Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::ClassicCom>,
    ABI::Windows::Foundation::IAsyncOperationCompletedHandler<T>>
{
public:

    HRESULT RuntimeClassInitialize()
    {
        _completedEvent = CreateEventEx(nullptr, nullptr, 0, EVENT_ALL_ACCESS);

        return S_OK;
    }

    ~OperationCallback()
    {
        if (_completedEvent != NULL)
        {
            CloseHandle(_completedEvent);
            _completedEvent = NULL;
        }
    }

    IFACEMETHOD(Invoke)(ABI::Windows::Foundation::IAsyncOperation<T> *pOp, ABI::Windows::Foundation::AsyncStatus status)
    {
        UNREFERENCED_PARAMETER(pOp);
        UNREFERENCED_PARAMETER(status);
        SetEvent(_completedEvent);
        return S_OK;
    }

    HRESULT Wait()
    {
        DWORD index = 0;
        // Wait until Get Token UI flow is completed.
        HRESULT hr = CoWaitForMultipleObjects(CWMO_DISPATCH_CALLS, INFINITE, 1, &_completedEvent, &index);
        return hr;
    }

    HANDLE _completedEvent;
};

template<typename InnerType, typename OutThing>
HRESULT AwaitOperationResult(
    ABI::Windows::Foundation::IAsyncOperation<InnerType> *pOperation,
    OutThing &&out)
{
    HRESULT hr = S_OK;
    Microsoft::WRL::ComPtr<OperationCallback<InnerType>> completer;
    Microsoft::WRL::MakeAndInitialize<OperationCallback<InnerType>>(&completer);

    pOperation->put_Completed(completer.Get());
    hr = completer->Wait();
    if (FAILED(hr))
    {
        return hr;
    }

    hr = pOperation->GetResults(out);
    return hr;
}