/*
Module : PJNNTLMAuth.cpp
Purpose: Implementation for a simple wrapper class for client NTLM authentication via SSPI
Created: PJN / 06-04-2005
History: PJN / 05-09-2005 1. Function pointer to CompleteAuthToken is now constructed at runtime. This means
                          that NTLM client authentication will work correctly on Win9x as used by the class
                          CPJNSMTPConnection and any other class instead of bringing up a "Failed to load
                          due to missing export..." message from the Windows loader. If you want to remove
                          the NTLM support from CPJNSMTPConnection then you can continue to use the CPJNSMTP_NONTLM
                          preprocessor define. Thanks to "Lior" for reporting this problem.
         PJN / 29-09-2005 1. All SDK function pointers used by the class are now implemented using GetProcAddress.
                          Thanks to Emir Kapic for this update.
         PJN / 29-06-2006 1. Combined the functionality of the _PJNNTLMCLIENTAUTH_DATA class into the main 
                          CNTLMClientAuth class.
         PJN / 14-07-2006 1. Updated CNTLMClientAuth to use newer C++ style casts.
                          2. Made the NTLMAuthenticate function "exception" safe, that is derived classes are free to throw 
                          exceptions in their implementations of NTLMAuthPhase1, NTLMAuthPhase2 or NTLMAuthPhase3 and the
                          CNTLMClientAuth class will properly clean up its resources. This is achieved by making the CredHandle
                          and SecHandle values member variables of the class.
                          3. Replaced all calls to ZeroMemory with memset.
                          4. Now the public NTLMAuthenticate function allows you to pass a user name and password pair to 
                          do NTLM authentication using a specified account instead of the credentials of the current user. If 
                          you leave these values as NULL, then you will get the old behaviour which is to use the current
                          user credentials. Thanks to Wouter Demuynck for this very nice addition.
         PJN / 13-10-2006 1. Code now initializes the Domain name value in NTLMAuthenticate to an empty string instead of a NULL
                          string. This avoids NTLM authentication issues when authenticating as a non-domain user with certain
                          mail servers. Thanks to Wouter Demuynck for reporting this issue.
         PJN / 31-05-2008 1. Code now compiles cleanly using Code Analysis (/analyze)
         PJN / 30-09-2012 1. Updated the code to avoid DLL planting security issues when calling LoadLibrary. Thanks to Mat 
                          Berchtold for reporting this issue.
         PJN / 15-11-2014 1. CNTLMClientAuth now takes a static dependency on Secure32.dll instead of using GetProcAddress.
                          2. The CNTLMClientAuth now calls QuerySecurityPackageInfo to determine the buffer size to use 
                          3. The SEC_SUCCESS macro has now been made a method called "SEC_SUCCESS"
                          4. General cleanup of the CNTLMClientAuth::GenClientContext method
                          5. Made the CNTLMClientAuth destructor virtual
         PJN / 16-01-2015 1. Updated copright details
         PJN / 11-08-2016 1. Minor review and tidy up of the CNTLMClientAuth code
                          2. Reworked the class to optionally compile without MFC. By default the class now uses STL 
                          classes and idioms but if you define CNTLMCLIENTAUTH_MFC_EXTENSIONS the class will revert back to the
                          MFC behaviour.
         PJN / 14-11-2016 1. Updated the non MFC code path in CNTLMClientAuth::NTLMAuthenticate to handle a NULL pszUsername
                          parameter being passed. Thanks to Christopher Craft for reporting this issue.
         PJN / 28-04-2017 1. Updated the code to compile cleanly using /permissive-.
         PJN / 02-05-2017 1. Fixed a bug in the CNTLMClientAuth::NTLMAuthenticate method when CNTLMCLIENTAUTH_MFC_EXTENSIONS
                          is defined. Thanks to Darron Redman for reporting this issue.
         PJN / 24-06-2018 1. Updated copyright details.
                          2. Fixed a number of C++ core guidelines compiler warnings. These changes mean that
                          the code will now only compile on VC 2017 or later.
         PJN / 25-08-2018 1. Fixed a number of compiler warnings when using VS 2017 15.8.1
                          2. Added some additional protective #includes to "PJNNTLMAuth.cpp" module
         PJN / 07-04-2019 1. Updated copyright details.
                          2. Removed code which supports the now defunct CNTLMCLIENTAUTH_MFC_EXTENSIONS codepath
         PJN / 14-04-2019 1. Replaced BOOL with bool throughout the code.

Copyright (c) 2005 - 2020 by PJ Naughter (Web: www.naughter.com, Email: pjna@naughter.com)

All rights reserved.

Copyright / Usage Details:

You are allowed to include the source code in any product (commercial, shareware, freeware or otherwise) 
when your product is released in binary form. You are allowed to modify the source code in any way you want 
except you cannot modify the copyright details at the top of each module. If you want to distribute source 
code with your application, then you are only allowed to distribute versions released by the author. This is 
to maintain a single distribution point for the source code. 

*/


//////////////// Includes /////////////////////////////////////////////////////

#include "stdafx.h"
#include "PJNNTLMAuth.h"
#ifndef __SSPI_H__
#pragma message("To avoid this message, please put Sspi.h in your pre compiled header (usually stdafx.h)")
#include <Sspi.h>
#endif //#ifndef __SSPI_H__

#ifndef _STRING_
#pragma message("To avoid this message, please put string in your pre compiled header (normally stdafx.h)")
#include <string>
#endif //#ifndef _STRING_

#ifndef _VECTOR_
#pragma message("To avoid this message, please put vector in your pre compiled header (normally stdafx.h)")
#include <vector>
#endif //#ifndef _VECTOR_


//////////////// Macros / Locals //////////////////////////////////////////////

#pragma comment(lib, "Secur32.lib") //Automatically link in the Secur32 dll


//////////////// Implementation ///////////////////////////////////////////////

CNTLMClientAuth::CNTLMClientAuth() noexcept : m_hCred{},
                                              m_hContext{},
                                              m_dwBufferSize(0)
{
}

CNTLMClientAuth::~CNTLMClientAuth() noexcept
{
  ReleaseHandles();
}

void CNTLMClientAuth::ReleaseHandles() noexcept
{
  //Free up the security context if valid
  if (m_hContext.dwLower != 0 || m_hContext.dwUpper != 0)
  {
    DeleteSecurityContext(&m_hContext);
    m_hContext = {};
  }

  //Free up the credentials handle if valid
  if (m_hCred.dwLower != 0 || m_hCred.dwUpper != 0) 
  {
    FreeCredentialsHandle(&m_hCred);
    m_hCred = {};
  }
}

SECURITY_STATUS CNTLMClientAuth::NTLMAuthenticate(_In_opt_z_ LPCTSTR pszUsername, _In_opt_z_ LPCTSTR pszPassword)
{
  //Allow "pszUsername" parameter to be of the format "DomainName\UserName"
  LPCTSTR pszDomain = _T("");
#ifdef _UNICODE
#pragma warning(suppress: 26485 26489)
  std::wstring sUsername(pszUsername != nullptr ? pszUsername : L"");
  const std::wstring::size_type nSlashSeparatorOffset = sUsername.find(_T('\\'));
  std::wstring sDomain;
  std::wstring sUsernameWithoutDomain;
  if (nSlashSeparatorOffset != std::wstring::npos)
#else
#pragma warning(suppress: 26489)
  std::string sUsername(pszUsername != nullptr ? pszUsername : "");
  const std::string::size_type nSlashSeparatorOffset = sUsername.find(_T('\\'));
  std::string sDomain;
  std::string sUsernameWithoutDomain;
  if (nSlashSeparatorOffset != std::string::npos)
#endif //#ifdef _UNICODE
  {
    sDomain = sUsername.substr(0, nSlashSeparatorOffset);
    pszDomain = sDomain.c_str();
#pragma warning(suppress: 26489)
    sUsernameWithoutDomain = sUsername.substr(nSlashSeparatorOffset + 1);
    pszUsername = sUsernameWithoutDomain.c_str();
  }

  //Release the handles before we try to authenticate (we do this here to ensure any previous calls
  //to NTLMAuthenticate which throw exceptions are cleaned up prior to any new calls to DoNTLMAuthentication)
  ReleaseHandles();

  //Get the buffer size to use for the NTLM SSPI
  PSecPkgInfo pSecInfo = nullptr;
#ifdef _UNICODE
#pragma warning(suppress: 26465 26486)
  SECURITY_STATUS ss = QuerySecurityPackageInfo(const_cast<LPTSTR>(NTLMSP_NAME), &pSecInfo);
#else
#pragma warning(suppress: 26465 26486)
  SECURITY_STATUS ss = QuerySecurityPackageInfo(const_cast<LPTSTR>(NTLMSP_NAME_A), &pSecInfo);
#endif //#ifdef _UNICODE
  if (ss != SEC_E_OK)
    return ss;
  m_dwBufferSize = pSecInfo->cbMaxToken;

  //Call the helper function which does all the work
  ss = DoNTLMAuthentication(pszUsername, pszPassword, pszDomain);

  //Now free up the handles now that we are finished the authentication (note it is not critical that this code is
  //called since the various NTLMAuthPhase(*) functions may throw exceptions)
  ReleaseHandles();

  return ss;
}

SECURITY_STATUS CNTLMClientAuth::DoNTLMAuthentication(_In_opt_z_ LPCTSTR pszUsername, _In_opt_z_ LPCTSTR pszPassword, _In_opt_z_ LPCTSTR pszDomain)
{
  //Allocate the in and out buffers for SSPI
  std::vector<BYTE> inBuf;
  inBuf.resize(m_dwBufferSize);
  std::vector<BYTE> outBuf;
  outBuf.resize(m_dwBufferSize);

  const DWORD cbMaxMessage = m_dwBufferSize;
  DWORD cbOut = cbMaxMessage;
  bool bDone = false;
  SECURITY_STATUS ss = GenClientContext(nullptr, 0, outBuf.data(), &cbOut, &bDone, pszUsername, pszPassword, pszDomain);
  if (!SEC_SUCCESS(ss))
    return ss;

  ss = NTLMAuthPhase1(outBuf.data(), cbOut);
  if (!SEC_SUCCESS(ss))
    return ss;

  DWORD cbIn = 0;
  while (!bDone)
  {
    ss = NTLMAuthPhase2(inBuf.data(), m_dwBufferSize, &cbIn);
    if (!SEC_SUCCESS(ss))
      return ss;

    cbOut = cbMaxMessage;

    ss = GenClientContext(inBuf.data(), cbIn, outBuf.data(), &cbOut, &bDone, pszUsername, pszPassword, pszDomain);
    if (!SEC_SUCCESS(ss))
      return ss;

    ss = NTLMAuthPhase3(outBuf.data(), cbOut);
    if (!SEC_SUCCESS(ss))
      return ss;
  }

  return ss;
}

SECURITY_STATUS CNTLMClientAuth::GenClientContext(_In_reads_bytes_opt_(cbIn) BYTE* pIn, _In_ DWORD cbIn, _Out_writes_bytes_(*pcbOut) BYTE* pOut, _Inout_ DWORD* pcbOut, _Inout_ bool* pfDone, _In_opt_z_ LPCTSTR pszUserName, _In_opt_z_ LPCTSTR pszPassword, _In_opt_z_ LPCTSTR pszDomain) noexcept
{
  //Validate our parameters
#pragma warning(suppress: 26477)
  ATLASSUME(pcbOut != nullptr);
#pragma warning(suppress: 26477)
  ATLASSUME(pfDone != nullptr);

  SECURITY_STATUS ss = SEC_E_OK;
  if (pIn == nullptr)
  {
#pragma warning(suppress: 26477)
    SEC_WINNT_AUTH_IDENTITY authInfo{};
    void* pvLogonID = nullptr;
    size_t nUserNameLength = 0;
    if (pszUserName != nullptr)
      nUserNameLength = _tcslen(pszUserName);
#pragma warning(suppress: 6387)
    std::unique_ptr<TCHAR, decltype(&std::free)> authInfoUser(_tcsdup(pszUserName), &std::free);
#pragma warning(suppress: 6387)
    std::unique_ptr<TCHAR, decltype(&std::free)> authInfoDomain(_tcsdup(pszDomain), &std::free);
#pragma warning(suppress: 6387)
    std::unique_ptr<TCHAR, decltype(&std::free)> authInfoPassword(_tcsdup(pszPassword), &std::free);
    if (nUserNameLength)
    {
#pragma warning(suppress: 26490)
      authInfo.User = reinterpret_cast<decltype(authInfo.User)>(authInfoUser.get());
#pragma warning(suppress: 26472)
      authInfo.UserLength = static_cast<unsigned long>(nUserNameLength);

#pragma warning(suppress: 26477)
      ATLASSUME(pszDomain != nullptr);
#pragma warning(suppress: 26490)
      authInfo.Domain = reinterpret_cast<decltype(authInfo.Domain)>(authInfoDomain.get());
#pragma warning(suppress: 26472)
      authInfo.DomainLength = static_cast<unsigned long>(_tcslen(pszDomain));

#pragma warning(suppress: 26477)
      ATLASSUME(pszPassword != nullptr);
#pragma warning(suppress: 26490)
      authInfo.Password = reinterpret_cast<decltype(authInfo.Password)>(authInfoPassword.get());
#pragma warning(suppress: 26472)
      authInfo.PasswordLength = static_cast<unsigned long>(_tcslen(pszPassword));
    #ifdef _UNICODE
      authInfo.Flags = SEC_WINNT_AUTH_IDENTITY_UNICODE;
    #else
      authInfo.Flags = SEC_WINNT_AUTH_IDENTITY_ANSI;
    #endif //#ifdef _UNICODE
      pvLogonID = &authInfo;
    }

    TimeStamp Lifetime;
#pragma warning(suppress: 26465)
    ss = AcquireCredentialsHandle(nullptr, const_cast<LPTSTR>(_T("NTLM")), SECPKG_CRED_OUTBOUND, nullptr, pvLogonID, nullptr, nullptr, &m_hCred, &Lifetime);
    if (ss != SEC_E_OK)
      return ss;
  }

  //Prepare the buffers
  SecBufferDesc OutBuffDesc{};
  OutBuffDesc.cBuffers = 1;
  SecBuffer OutSecBuff{};
  OutBuffDesc.pBuffers = &OutSecBuff;
  OutSecBuff.cbBuffer = *pcbOut;
  OutSecBuff.BufferType = SECBUFFER_TOKEN;
  OutSecBuff.pvBuffer = pOut;

  //The input buffer is created only if a message has been received from the server
  SecBufferDesc InBuffDesc{};
  SecBuffer InSecBuff{};
  if (pIn != nullptr)
  {
    InBuffDesc.cBuffers = 1;
    InBuffDesc.pBuffers = &InSecBuff;
    InSecBuff.cbBuffer = cbIn;
    InSecBuff.BufferType = SECBUFFER_TOKEN;
    InSecBuff.pvBuffer = pIn;
    ULONG ContextAttributes = 0;
    ss = InitializeSecurityContext(&m_hCred, &m_hContext, nullptr, 0, 0, SECURITY_NATIVE_DREP, &InBuffDesc, 0, &m_hContext, &OutBuffDesc, &ContextAttributes, nullptr);
  }
  else
  {
    ULONG ContextAttributes = 0;
    ss = InitializeSecurityContext(&m_hCred, nullptr, nullptr, 0, 0, SECURITY_NATIVE_DREP, nullptr, 0, &m_hContext, &OutBuffDesc, &ContextAttributes, nullptr);
  }

  //If necessary, complete the token.
  if ((ss == SEC_I_COMPLETE_NEEDED) || (ss == SEC_I_COMPLETE_AND_CONTINUE))  
  {
    ss = CompleteAuthToken(&m_hContext, &OutBuffDesc);
    if (ss != SEC_E_OK)  
      return ss;
  }
  else if ((ss != SEC_I_CONTINUE_NEEDED) && (ss != SEC_E_OK))
    return ss;

  //Finally update the output parameters
  *pcbOut = OutSecBuff.cbBuffer;
  *pfDone = !((ss == SEC_I_CONTINUE_NEEDED) || (ss == SEC_I_COMPLETE_AND_CONTINUE));

  return ss;
}
