Sec.h

/*
    SPDX-License-Identifier: MIT
    Copyright © 1991-2022 Amebis
    Copyright © 2016 GÉANT
*/
 
#pragma once
 
#include "Common.h"
#include <Security.h>
#include <string>
 
 
#if defined(SECURITY_WIN32) || defined(SECURITY_KERNEL)
 
template<class _Traits, class _Ax>
static BOOLEAN GetUserNameExA(_In_ EXTENDED_NAME_FORMAT NameFormat, _Inout_ std::basic_string<char, _Traits, _Ax> &sName)
{
    assert(0); // TODO: Test this code.
 
    char szStackBuffer[WINSTD_STACK_BUFFER_BYTES/sizeof(char)];
    ULONG ulSize = _countof(szStackBuffer);
 
    // Try with stack buffer first.
    if (::GetUserNameExA(NameFormat, szStackBuffer, &ulSize)) {
        // Copy from stack.
        sName.assign(szStackBuffer, ulSize);
        return TRUE;
    } else {
        if (::GetLastError() == ERROR_MORE_DATA) {
            // Allocate buffer on heap and retry.
            std::unique_ptr<char[]> szBuffer(new char[ulSize]);
            if (::GetUserNameExA(NameFormat, szBuffer.get(), &ulSize)) {
                sName.assign(szBuffer.get(), ulSize);
                return TRUE;
            }
        }
    }
 
    return FALSE;
}
 
template<class _Traits, class _Ax>
static BOOLEAN GetUserNameExW(_In_ EXTENDED_NAME_FORMAT NameFormat, _Inout_ std::basic_string<wchar_t, _Traits, _Ax> &sName)
{
    assert(0); // TODO: Test this code.
 
    wchar_t szStackBuffer[WINSTD_STACK_BUFFER_BYTES/sizeof(wchar_t)];
    ULONG ulSize = _countof(szStackBuffer);
 
    // Try with stack buffer first.
    if (::GetUserNameExW(NameFormat, szStackBuffer, &ulSize)) {
        // Copy from stack.
        sName.assign(szStackBuffer, ulSize);
        return TRUE;
    } else {
        if (::GetLastError() == ERROR_MORE_DATA) {
            // Allocate buffer on heap and retry.
            std::unique_ptr<wchar_t[]> szBuffer(new wchar_t[ulSize]);
            if (::GetUserNameExW(NameFormat, szBuffer.get(), &ulSize)) {
                sName.assign(szBuffer.get(), ulSize);
                return TRUE;
            }
        }
    }
 
    return FALSE;
}
 
#endif
 
 
namespace winstd
{
 
    class sec_credentials : public handle<PCredHandle, NULL>
    {
        WINSTD_NONCOPYABLE(sec_credentials)
 
    public:
        sec_credentials()
        {
            m_expires.QuadPart = -1;
        }
 
        sec_credentials(_In_opt_ handle_type h, _In_ const TimeStamp expires) :
            m_expires(expires),
            handle(h)
        {
        }
 
        sec_credentials(_Inout_ sec_credentials &&h) noexcept :
            m_expires(std::move(h.m_expires)),
            handle<PCredHandle, NULL>(std::move(h))
        {
        }
 
        virtual ~sec_credentials()
        {
            if (m_h != invalid)
                free_internal();
        }
 
        sec_credentials& operator=(_Inout_ sec_credentials &&h) noexcept
        {
            if (this != std::addressof(h)) {
                *(handle<handle_type, NULL>*)this = std::move(h);
                m_expires = std::move(h.m_expires);
            }
            return *this;
        }
 
        SECURITY_STATUS acquire(
            _In_opt_ LPTSTR         pszPrincipal,
            _In_     LPTSTR         pszPackage,
            _In_     unsigned long  fCredentialUse,
            _In_opt_ void           *pvLogonId,
            _In_opt_ void           *pAuthData,
            _In_opt_ SEC_GET_KEY_FN pGetKeyFn = NULL,
            _In_opt_ void           *pvGetKeyArgument = NULL)
        {
            handle_type h = new CredHandle;
            TimeStamp exp;
            SECURITY_STATUS res = AcquireCredentialsHandle(pszPrincipal, pszPackage, fCredentialUse, pvLogonId, pAuthData, pGetKeyFn, pvGetKeyArgument, h, &exp);
            if (SUCCEEDED(res)) {
                attach(h);
                m_expires = exp;
            } else
                delete h;
            return res;
        }
 
    protected:
        void free_internal() noexcept override
        {
            FreeCredentialsHandle(m_h);
            delete m_h;
        }
 
    public:
        TimeStamp m_expires;    
    };
 
    class sec_context : public handle<PCtxtHandle, NULL>
    {
    public:
        sec_context() :
            m_attrib(0),
            handle<PCtxtHandle, NULL>()
        {
            m_expires.QuadPart = -1;
        }
 
        sec_context(_Inout_ sec_context &&h) noexcept :
            m_attrib (std::move(h.m_attrib )),
            m_expires(std::move(h.m_expires)),
            handle<PCtxtHandle, NULL>(std::move(h))
        {
        }
 
        virtual ~sec_context()
        {
            if (m_h != invalid)
                free_internal();
        }
 
        sec_context& operator=(_Inout_ sec_context &&h) noexcept
        {
            if (this != std::addressof(h)) {
                *(handle<handle_type, NULL>*)this = std::move(h);
                m_attrib  = std::move(h.m_attrib);
                m_expires = std::move(h.m_expires);
            }
            return *this;
        }
 
        SECURITY_STATUS initialize(
            _In_opt_    PCredHandle    phCredential,
            _In_opt_z_  LPCTSTR        pszTargetName,
            _In_        ULONG          fContextReq,
            _In_        ULONG          TargetDataRep,
            _In_opt_    PSecBufferDesc pInput,
            _Inout_opt_ PSecBufferDesc pOutput)
        {
            handle_type h = new CtxtHandle;
            h->dwUpper = 0;
            h->dwLower = 0;
            ULONG attr;
            TimeStamp exp;
            SECURITY_STATUS res = InitializeSecurityContext(phCredential, NULL, const_cast<LPTSTR>(pszTargetName), fContextReq, 0, TargetDataRep, pInput, 0, h, pOutput, &attr, &exp);
            if (SUCCEEDED(res)) {
                attach(h);
                m_attrib  = attr;
                m_expires = exp;
            } else
                delete h;
            return res;
        }
 
        SECURITY_STATUS process(
            _In_opt_    PCredHandle    phCredential,
            _In_opt_z_  LPCTSTR        pszTargetName,
            _In_        ULONG          fContextReq,
            _In_        ULONG          TargetDataRep,
            _In_opt_    PSecBufferDesc pInput,
            _Inout_opt_ PSecBufferDesc pOutput)
        {
            return InitializeSecurityContext(phCredential, m_h, const_cast<LPTSTR>(pszTargetName), fContextReq, 0, TargetDataRep, pInput, 0, NULL, pOutput, &m_attrib, &m_expires);
        }
 
    protected:
        void free_internal() noexcept override
        {
            DeleteSecurityContext(m_h);
            delete m_h;
        }
 
    public:
        ULONG     m_attrib;     
        TimeStamp m_expires;    
    };
 
    class sec_buffer_desc : public SecBufferDesc
    {
    public:
        sec_buffer_desc(_Inout_count_(count) PSecBuffer buf, ULONG count, _In_ ULONG version = SECBUFFER_VERSION)
        {
            ulVersion = version;
            cBuffers  = count;
            pBuffers  = buf;
        }
 
        virtual ~sec_buffer_desc()
        {
            for (ULONG i = 0; i < cBuffers; i++) {
                if (pBuffers[i].pvBuffer)
                    FreeContextBuffer(pBuffers[i].pvBuffer);
            }
        }
    };
 
 
 
    class sec_runtime_error : public num_runtime_error<SECURITY_STATUS>
    {
    public:
        sec_runtime_error(_In_ error_type num, _In_ const std::string& msg) : num_runtime_error<SECURITY_STATUS>(num, msg)
        {
        }
 
        sec_runtime_error(_In_ error_type num, _In_opt_z_ const char *msg = nullptr) : num_runtime_error<SECURITY_STATUS>(num, msg)
        {
        }
 
        sec_runtime_error(const sec_runtime_error &other) : num_runtime_error<SECURITY_STATUS>(other)
        {
        }
    };
 
}