/*
Module : PJNMD5.h
Purpose: Defines the interface for some MFC class which encapsulate calculating MD5 hashes and HMACS using the MS CryptoAPI
Created: PJN / 23-04-2005
History: PJN / 18-05-2005 1. Fixed a compiler warning when compiled using Visual Studio .NET 2003. Thanks to Alexey Kuznetsov
                          for reporting this issue.
         PJN / 29-09-2005 1. Format method now allows you to specify an uppercase or lower case string for the hash. This is 
                          necessary since the CRAM-MD5 authentication mechanism requires a lowercase MD5 hash. Thanks to 
                          Jian Peng for reporting this issue.
         PJN / 16-01-2015 1. Updated copright details
         PJN / 11-08-2016 1. Removed defunct comments about VC 6 in CPJNMD5 class.
                          2. Reworked the class to optionally compile without MFC. By default the class now uses STL 
                          classes and idioms but if you define CPJNMD5_MFC_EXTENSIONS the classes will revert back to the 
                          MFC behaviour.
         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 / 22-08-2018 1. Fixed a number of compiler warnings when using VS 2018 15.8.1
         PJN / 07-04-2019 1. Updated copyright details.
                          2. Removed code which supports the now defunct CPJNMD5_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. 

*/


////////////////////////////// Macros / Defines ///////////////////////////////

#pragma once

#ifndef __PJNMD5_H__
#define __PJNMD5_H__


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

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


////////////////////////////// Classes ////////////////////////////////////////

// A simple wrapper class which contains a MD5 hash i.e. a 16 byte blob
class CPJNMD5Hash : public std::vector<BYTE>
{
public:
//typedefs
#ifdef _UNICODE
  using String = std::wstring;
#else
  using String = std::string;
#endif //#ifdef _UNICODE

//Constructors / Destructors
#pragma warning(suppress: 26455)
  CPJNMD5Hash()
  {
    resize(16);
  }
  ~CPJNMD5Hash() = default;
  CPJNMD5Hash(const CPJNMD5Hash& hash) = default;
  CPJNMD5Hash(CPJNMD5Hash&&) = default;

//Methods
  CPJNMD5Hash& operator=(const CPJNMD5Hash&) = default;
  CPJNMD5Hash& operator=(CPJNMD5Hash&&) = default;

  String Format(_In_ bool bUppercase)
  {
    //What will be the return value from this method
    String sRet;
    sRet.reserve(33);
#pragma warning(suppress: 26429 26446)
    LPTSTR pString = &(sRet[0]);

    for (size_t i=0; i<size(); i++)
    {
#pragma warning(suppress: 26482 26446)
      int nChar = (operator[](i) & 0xF0) >> 4;
      if (nChar <= 9)
#pragma warning(suppress: 26472 26481)
        pString[i*2] = static_cast<TCHAR>(nChar + _T('0'));
      else
#pragma warning(suppress: 26472 26481)
        pString[i*2] = static_cast<TCHAR>(nChar - 10 + (bUppercase ? _T('A') : _T('a')));

#pragma warning(suppress: 26482 26446)
      nChar = operator[](i) & 0x0F;
      if (nChar <= 9)
#pragma warning(suppress: 26472 26481)
        pString[i*2 + 1] = static_cast<TCHAR>(nChar + _T('0'));
      else
#pragma warning(suppress: 26472 26481)
        pString[i*2 + 1] = static_cast<TCHAR>(nChar - 10 + (bUppercase ? _T('A') : _T('a')));
    }
#pragma warning(suppress: 26481)
    pString[32] = _T('\0');

    return sRet;
  }
};


// Class which calculates MD5 Hashes and MD5 HMACS
class CPJNMD5
{
public:
//Constructors / Destructors
  CPJNMD5() noexcept
  {
    //Acquire the RSA CSP
    m_hProv = NULL;
    CryptAcquireContext(&m_hProv, nullptr, nullptr, PROV_RSA_FULL, CRYPT_VERIFYCONTEXT | CRYPT_SILENT);
  }

  CPJNMD5(const CPJNMD5&) = delete;
  CPJNMD5(CPJNMD5&&) = delete;

  virtual ~CPJNMD5()
  {
    //Release the CSP
    if (m_hProv)
      CryptReleaseContext(m_hProv, 0);
  }

  CPJNMD5& operator=(const CPJNMD5&) = delete;
  CPJNMD5& operator=(CPJNMD5&&) = delete;

//Methods
  //Simple MD5 hash function which hashes a blob of data and returns
  //the results encapsulated in the "hash" parameter
  bool Hash(_In_reads_bytes_(dwDataSize) const BYTE* pbyData, _In_ DWORD dwDataSize, _Inout_ CPJNMD5Hash& hash)
  {
    //Call the Helper Hash function which returns a HCRYPTHASH handle
    //which contains the hash
    HCRYPTHASH hHash = NULL;
    bool bSuccess = Hash(pbyData, dwDataSize, hHash);
    if (bSuccess)
    {
      //Get the actual hash value
      hash.resize(16);
#pragma warning(suppress: 26472)
      auto dwHashSize = static_cast<DWORD>(hash.size());
      bSuccess = CryptGetHashParam(hHash, HP_HASHVAL, hash.data(), &dwHashSize, 0);
    }
    //Free up the hash handle now that we are finished with it
    if (hHash)
      CryptDestroyHash(hHash);

    return bSuccess;
  }

  //Function which returns a hash encoded as a HCRYPTHASH. If hHash is
  //NULL on entry then a new HCRYPTHASH is created. If hHash is non-null
  //then this existing hash is used and the data to hash is included into
  //the current hash. The client is responsible for freeing the hash handle
  bool Hash(_In_reads_bytes_(dwDataSize) const BYTE* pbyData, DWORD dwDataSize, _Inout_ HCRYPTHASH& hHash) noexcept
  {
    //Create the hash object if required to
    if (hHash == NULL)
      if (!CryptCreateHash(m_hProv, CALG_MD5, 0, 0, &hHash)) 
        return false;

    //Hash in the data
    return CryptHashData(hHash, pbyData, dwDataSize, 0) != 0;
  }

  //Function which calculates a keyed MD5 HMAC according to RFC 2104
  bool HMAC(_In_reads_bytes_(dwDataSize) const BYTE* pbyData, _In_ DWORD dwDataSize, _In_reads_bytes_(dwKeySize) const BYTE* pbyKey, _In_ DWORD dwKeySize, _Inout_ CPJNMD5Hash& hash)
  {
    //if key is longer than 64 bytes then reset it to the MD5 hash of the key
    const BYTE* pbyLocalKey = nullptr;
    DWORD dwLocalKeySize = 0;
    CPJNMD5Hash keyHash;
    if (dwKeySize > 64)
    {
      if (!Hash(pbyKey, dwKeySize, keyHash))
        return false;
#pragma warning(suppress: 26472)
      dwLocalKeySize = static_cast<DWORD>(keyHash.size());
      pbyLocalKey = keyHash.data();
    }
    else
    {
      dwLocalKeySize = dwKeySize;
      pbyLocalKey = pbyKey;
    }

    //start out by storing key in pads
    std::vector<BYTE> k_ipad;
    k_ipad.resize(64);
    std::vector<BYTE> k_opad;
    k_opad.resize(64);
    memcpy_s(k_ipad.data(), k_ipad.size(), pbyLocalKey, dwLocalKeySize);
    memcpy_s(k_opad.data(), k_opad.size(), pbyLocalKey, dwLocalKeySize);

    //XOR key with ipad value
    for (auto& element : k_ipad)
      element ^= 0x36;

    //XOR key with opad
    for (auto& element : k_opad)
      element ^= 0x5c;

    //perform inner MD5
    HCRYPTHASH hInnerHash = NULL;
  #pragma warning(suppress: 6385)
#pragma warning(suppress: 26472)
    if (!Hash(k_ipad.data(), static_cast<DWORD>(k_ipad.size()), hInnerHash))
    {
      //Free up the hash before we return
      if (hInnerHash)
        CryptDestroyHash(hInnerHash);
      return false;
    }
    if (!Hash(pbyData, dwDataSize, hInnerHash))
    {
      //Free up the hash before we return
      CryptDestroyHash(hInnerHash);
      return false;
    }

    //Get the inner hash result
    CPJNMD5Hash InnerHash;
#pragma warning(suppress: 26472)
    auto dwHashSize = static_cast<DWORD>(InnerHash.size());
    if (!CryptGetHashParam(hInnerHash, HP_HASHVAL, InnerHash.data(), &dwHashSize, 0))
    {
      //Free up the hash before we return
      CryptDestroyHash(hInnerHash);
      return false;
    }

    //Free up the inner hash now that we are finished with it
    CryptDestroyHash(hInnerHash);

    //Perform outter MD5
    HCRYPTHASH hOuterHash = NULL;
#pragma warning(suppress: 26472)
    if (!Hash(k_opad.data(), static_cast<DWORD>(k_opad.size()), hOuterHash))
    {
      //Free up the hash before we return
      if (hOuterHash)
        CryptDestroyHash(hOuterHash);
      return false;
    }
#pragma warning(suppress: 26472)
    if (!Hash(InnerHash.data(), static_cast<DWORD>(InnerHash.size()), hOuterHash))
    {
      //Free up the hash before we return
      CryptDestroyHash(hOuterHash);
      return false;
    }

    //Finally get the hash result
#pragma warning(suppress: 26472)
    dwHashSize = static_cast<DWORD>(hash.size());
    const bool bSuccess = CryptGetHashParam(hOuterHash, HP_HASHVAL, hash.data(), &dwHashSize, 0) != 0;

    //Free up the hash before we return
    CryptDestroyHash(hOuterHash);

    return bSuccess;
  }

protected:
//Member variables
  HCRYPTPROV m_hProv;
};

#endif //#ifndef __PJNMD5_H__
