/*
Module : CryptoWrappersStore.h
Purpose: Defines the interface for a C++ wrapper class for a CryptoAPI HCERTSTORE
History: PJN / 01-08-2014 1. Initial public release
         PJN / 12-05-2016 1. Updated copyright details.
                          2. Updated code to be /analyze clean on VC 2015
         PJN / 21-12-2017 1. Replaced NULL throughout the codebase with nullptr.
         PJN / 12-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 2017 15.8.1

Copyright (c) 2014 - 2018 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 __CRYPTOWRAPPERSSTORE_H__
#define __CRYPTOWRAPPERSSTORE_H__

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


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

#ifndef __WINCRYPT_H__
#pragma message("To avoid this message, please put wincrypt.h in your pre compiled header (normally stdafx.h)")
#include <wincrypt.h>
#endif //#ifndef __WINCRYPT_H__
#include "CryptoWrappersCertificate.h"
#include "CryptoWrappersCRL.h"
#include "CryptoWrappersCTL.h"
#include "CryptoWrappersChain.h"


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

namespace CryptoWrappers
{


//Encapsulates a CryptoAPI HCERTSTORE
class CStore
{
public:
//Constructors / Destructors
  CStore() noexcept : m_hStore(nullptr)
  {
  }

  CStore(_In_ const CStore& store) noexcept : m_hStore(nullptr)
  {
    if (!store.Duplicate(*this))
      m_hStore = nullptr;
  }

  CStore(_In_ CStore&& store) noexcept : m_hStore(nullptr)
  {
    Attach(store.Detach());
  }

  explicit CStore(_In_opt_ HCERTSTORE hStore) noexcept : m_hStore(hStore)
  {
  }

  ~CStore() noexcept
  {
    if (m_hStore != nullptr)
      Close();
  }

//Methods
  CStore& operator=(_In_ const CStore& store) noexcept
  {
    if (this != &store)
    {
      if (m_hStore != nullptr)
        Close();
      if (!store.Duplicate(*this))
        m_hStore = nullptr;
    }
    return *this;
  }

  CStore& operator=(_In_ CStore&& store) noexcept
  {
    if (m_hStore != nullptr)
      Close();
    Attach(store.Detach());
    return *this;
  }

  HCERTSTORE Handle() const noexcept
  {
    return m_hStore;
  }

  void Attach(_In_opt_ HCERTSTORE hStore) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore == nullptr);

    m_hStore = hStore;
  }

  HCERTSTORE Detach() noexcept
  {
    HCERTSTORE hStore = m_hStore;
    m_hStore = nullptr;
    return hStore;
  }

  _Must_inspect_result_
  BOOL Open(_In_ LPCSTR lpszStoreProvider, _In_ DWORD dwEncodingType, _In_ DWORD dwFlags, _In_opt_ const void* pvPara) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore == nullptr);

    m_hStore = CertOpenStore(lpszStoreProvider, dwEncodingType, NULL, dwFlags, pvPara);
    return (m_hStore != nullptr);
  }

  _Must_inspect_result_
  BOOL OpenSystemStore(_In_ LPCTSTR szSubsystemProtocol) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore == nullptr);

    m_hStore = CertOpenSystemStore(NULL, szSubsystemProtocol);
    return (m_hStore != nullptr);
  }

  _Success_(return != FALSE)
  BOOL Close(_In_ DWORD dwFlags = 0) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    const BOOL bResult = CertCloseStore(m_hStore, dwFlags);
    m_hStore = nullptr;
    return bResult;
  }

  _Success_(return != FALSE)
  BOOL Control(_In_ DWORD dwFlags, _In_ DWORD dwCtrlType, _In_opt_ void const* pvCtrlPara) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    return CertControlStore(m_hStore, dwFlags, dwCtrlType, pvCtrlPara);
  }

  _Success_(return != FALSE)
  BOOL Duplicate(_Inout_ CStore& store) const noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);
#pragma warning(suppress: 26477)
    ATLASSERT(store.m_hStore == nullptr);

    store.m_hStore = CertDuplicateStore(m_hStore);
    return (store.m_hStore != nullptr);
  }

  _Success_(return != FALSE)
  BOOL AddStoreToCollection(_In_opt_ HCERTSTORE hSiblingStore, _In_ DWORD dwUpdateFlags, _In_ DWORD dwPriority) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    return CertAddStoreToCollection(m_hStore, hSiblingStore, dwUpdateFlags, dwPriority);
  }

  void RemoveStoreToCollection(_In_ const CStore& siblingStore) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CertRemoveStoreFromCollection(m_hStore, siblingStore.m_hStore);
  }

  _Success_(return != FALSE)
  BOOL Save(_In_ DWORD dwEncodingType, _In_ DWORD dwSaveAs, _In_ DWORD dwSaveTo, _Inout_ void* pvSaveToPara) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    return CertSaveStore(m_hStore, dwEncodingType, dwSaveAs, dwSaveTo, pvSaveToPara, 0);
  }

  _Success_(return != FALSE)
  BOOL SetProperty(_In_ DWORD dwPropId, _In_opt_ const void* pvData) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    return CertSetStoreProperty(m_hStore, dwPropId, 0, pvData);
  }

  _Success_(return != FALSE)
  BOOL GetProperty(_In_ DWORD dwPropId, _Out_writes_bytes_to_opt_(*pcbData, *pcbData) void* pvData, _Inout_ DWORD* pcbData) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    return CertGetStoreProperty(m_hStore, dwPropId, pvData, pcbData);
  }

  _Success_(return != FALSE)
  BOOL AddSerializedElement(_In_reads_bytes_(cbElement) const BYTE *pbElement, _In_ DWORD cbElement, _In_ DWORD dwAddDisposition, _In_ DWORD dwContextTypeFlags, _Out_opt_ DWORD* pdwContextType, _Outptr_opt_ const void** ppvContext) noexcept
  {
    return CertAddSerializedElementToStore(m_hStore, pbElement, cbElement, dwAddDisposition, 0, dwContextTypeFlags, pdwContextType, ppvContext);
  }

  _Success_(return != FALSE)
  BOOL AddCertifcate(_In_ const CCertificate& cert, _In_ DWORD dwAddDisposition, _Inout_opt_ CCertificate* pStoreCertificate) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);
    if (pStoreCertificate != nullptr)
#pragma warning(suppress: 26477)
      ATLASSERT(pStoreCertificate->m_pCert == nullptr);

    return CertAddCertificateContextToStore(m_hStore, cert.m_pCert, dwAddDisposition, pStoreCertificate ? &pStoreCertificate->m_pCert : nullptr);
  }

  _Success_(return != FALSE)
  BOOL AddCRL(_In_ const CCRL& crl, _In_ DWORD dwAddDisposition, _Inout_opt_ CCRL* pStoreCRL) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);
    if (pStoreCRL != nullptr)
#pragma warning(suppress: 26477)
      ATLASSERT(pStoreCRL->m_pCRL == nullptr);

    return CertAddCRLContextToStore(m_hStore, crl.m_pCRL, dwAddDisposition, pStoreCRL ? &pStoreCRL->m_pCRL : nullptr);
  }

  _Success_(return != FALSE)
  BOOL AddCTL(_In_ const CCTL& ctl, _In_ DWORD dwAddDisposition, _Inout_opt_ CCTL* pStoreCTL) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);
    if (pStoreCTL != nullptr)
#pragma warning(suppress: 26477)
      ATLASSERT(pStoreCTL->m_pCTL == nullptr);

    return CertAddCTLContextToStore(m_hStore, ctl.m_pCTL, dwAddDisposition, pStoreCTL ? &pStoreCTL->m_pCTL : nullptr);
  }

  CCertificate SelectCertificate(_In_opt_ HWND hwnd, _In_opt_ LPCWSTR pwszTitle, _In_opt_ LPCWSTR pwszDisplayString, _In_ DWORD dwDontUseColumn) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

  #pragma warning(suppress: 6387) //There are SAL annotation errors on the seventh parameter to CryptUIDlgSelectCertificateFromStore even in the latest Windows 10 SDK
    CCertificate certificate(CryptUIDlgSelectCertificateFromStore(m_hStore, hwnd, pwszTitle, pwszDisplayString, dwDontUseColumn, 0, nullptr));
#pragma warning(suppress: 26486)
    return certificate;
  }

  CCertificate EnumCertificates(_In_ CCertificate& prevCertificate) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CCertificate certificate(CertEnumCertificatesInStore(m_hStore, prevCertificate.m_pCert));
    prevCertificate.m_pCert = nullptr; //Required because CertEnumCertificatesInStore auto frees any non nullptr pPrevCertContext parameter
#pragma warning(suppress: 26486)
    return certificate;
  }

  CCRL EnumCRLs(_In_ CCRL& prevCRL) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CCRL crl(CertEnumCRLsInStore(m_hStore, prevCRL.m_pCRL));
    prevCRL.m_pCRL = nullptr; //Required because CertEnumCRLsInStore auto frees any non nullptr pPrevCrlContext parameter
#pragma warning(suppress: 26486)
    return crl;
  }

  CCTL EnumCTLs(_In_ CCTL& prevCTL) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CCTL crl(CertEnumCTLsInStore(m_hStore, prevCTL.m_pCTL));
    prevCTL.m_pCTL = nullptr; //Required because CertEnumCTLsInStore auto frees any non nullptr pPrevCtlContext parameter
#pragma warning(suppress: 26486)
    return crl;
  }

  CCertificate FindCertificate(_In_ DWORD dwCertEncodingType, _In_ DWORD dwFindFlags, _In_ DWORD dwFindType, _In_opt_ const void* pvFindPara, _In_ CCertificate& prevCertificate) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CCertificate certificate(CertFindCertificateInStore(m_hStore, dwCertEncodingType, dwFindFlags, dwFindType, pvFindPara, prevCertificate.m_pCert));
    prevCertificate.m_pCert = nullptr; //Required because CertFindCertificateInStore auto frees any non nullptr pPrevCertContext parameter
#pragma warning(suppress: 26486)
    return certificate;
  }

  CCRL FindCRL(_In_ DWORD dwCertEncodingType, _In_ DWORD dwFindFlags, _In_ DWORD dwFindType, _In_opt_ const void* pvFindPara, _In_ CCRL& prevCRL) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CCRL crl(CertFindCRLInStore(m_hStore, dwCertEncodingType, dwFindFlags, dwFindType, pvFindPara, prevCRL.m_pCRL));
    prevCRL.m_pCRL = nullptr; //Required because CertFindCRLInStore auto frees any non nullptr pPrevCertContext parameter
#pragma warning(suppress: 26486)
    return crl;
  }

  CCTL FindCTL(_In_ DWORD dwMsgAndCertEncodingType, _In_ DWORD dwFindFlags, _In_ DWORD dwFindType, _In_opt_ const void* pvFindPara, _In_ CCTL& prevCTL) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CCTL ctl(CertFindCTLInStore(m_hStore, dwMsgAndCertEncodingType, dwFindFlags, dwFindType, pvFindPara, prevCTL.m_pCTL));
    prevCTL.m_pCTL = nullptr; //Required because CertFindCRLInStore auto frees any non nullptr pPrevCertContext parameter
#pragma warning(suppress: 26486)
    return ctl;
  }

  CChain FindChain(_In_ DWORD dwCertEncodingType, _In_ DWORD dwFindFlags, _In_ DWORD dwFindType, _In_opt_ const void* pvFindPara, _In_  CChain& prevChain) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CChain chain(CertFindChainInStore(m_hStore, dwCertEncodingType, dwFindFlags, dwFindType, pvFindPara, prevChain.m_pChain));
    prevChain.m_pChain = nullptr; //Required because CertFindChainInStore auto frees any non nullptr pPrevCertContext parameter
#pragma warning(suppress: 26486)
    return chain;
  }

  _Success_(return != FALSE)
  BOOL AddCertificateLink(_In_ const CCertificate& certificate, _In_ DWORD dwAddDisposition, _Inout_opt_ CCertificate* pStoreCertificate) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);
    if (pStoreCertificate != nullptr)
#pragma warning(suppress: 26477)
      ATLASSERT(pStoreCertificate->m_pCert == nullptr);

    return CertAddCertificateLinkToStore(m_hStore, certificate.m_pCert, dwAddDisposition, pStoreCertificate ? &pStoreCertificate->m_pCert : nullptr);
  }

  _Success_(return != FALSE)
  BOOL AddCRLLink(_In_ const CCRL& crl, _In_ DWORD dwAddDisposition, _Inout_opt_ CCRL* pStoreCRL) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);
    if (pStoreCRL != nullptr)
#pragma warning(suppress: 26477)
      ATLASSERT(pStoreCRL->m_pCRL == nullptr);

    return CertAddCRLLinkToStore(m_hStore, crl.m_pCRL, dwAddDisposition, pStoreCRL ? &pStoreCRL->m_pCRL : nullptr);
  }

  _Success_(return != FALSE)
  BOOL AddCTLLink(_In_ const CCTL& ctl, _In_ DWORD dwAddDisposition, _Inout_opt_ CCTL* pStoreCTL) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);
    if (pStoreCTL != nullptr)
#pragma warning(suppress: 26477)
      ATLASSERT(pStoreCTL->m_pCTL == nullptr);

    return CertAddCTLLinkToStore(m_hStore, ctl.m_pCTL, dwAddDisposition, pStoreCTL ? &pStoreCTL->m_pCTL : nullptr);
  }

  _Success_(return != FALSE)
  BOOL AddEncodedCertificate(_In_ DWORD dwCertEncodingType, _In_reads_bytes_(cbCertEncoded) const BYTE* pbCertEncoded, _In_ DWORD cbCertEncoded, _In_ DWORD dwAddDisposition, _Inout_opt_ CCertificate* pCertificate) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);
    if (pCertificate != nullptr)
#pragma warning(suppress: 26477)
      ATLASSERT(pCertificate->m_pCert == nullptr);

    return CertAddEncodedCertificateToStore(m_hStore, dwCertEncodingType, pbCertEncoded, cbCertEncoded, dwAddDisposition, (pCertificate != nullptr) ? &pCertificate->m_pCert : nullptr);
  }

  _Success_(return != FALSE)
  BOOL AddEncodedCRL(_In_ DWORD dwCertEncodingType, _In_reads_bytes_(cbCrlEncoded) const BYTE* pbCrlEncoded, _In_ DWORD cbCrlEncoded, _In_ DWORD dwAddDisposition, _Inout_opt_ CCRL* pCRL) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);
    if (pCRL != nullptr)
#pragma warning(suppress: 26477)
      ATLASSERT(pCRL->m_pCRL == nullptr);

    return CertAddEncodedCRLToStore(m_hStore, dwCertEncodingType, pbCrlEncoded, cbCrlEncoded, dwAddDisposition, (pCRL != nullptr) ? &pCRL->m_pCRL : nullptr);
  }

  _Success_(return != FALSE)
  BOOL AddEncodedCTL(_In_ DWORD dwMsgAndCertEncodingType, _In_reads_bytes_(cbCtlEncoded) const BYTE* pbCtlEncoded, _In_ DWORD cbCtlEncoded, _In_ DWORD dwAddDisposition, _Inout_opt_ CCTL* pCTL) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);
    if (pCTL != nullptr)
#pragma warning(suppress: 26477)
      ATLASSERT(pCTL->m_pCTL == nullptr);

    return CertAddEncodedCTLToStore(m_hStore, dwMsgAndCertEncodingType, pbCtlEncoded, cbCtlEncoded, dwAddDisposition, (pCTL != nullptr) ? &pCTL->m_pCTL : nullptr);
  }

  CCertificate GetIssuerCertificate(_In_ const CCertificate& subjectContext, _In_ CCertificate& prevIssuerContext, _Inout_ DWORD* pdwFlags) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CCertificate certificate(CertGetIssuerCertificateFromStore(m_hStore, subjectContext.Handle(), prevIssuerContext.m_pCert, pdwFlags));
    prevIssuerContext.m_pCert = nullptr; //Required because CertGetIssuerCertificateFromStore auto frees any non nullptr pPrevCertContext parameter
#pragma warning(suppress: 26486)
    return certificate;
  }

  CCertificate GetSubjectCertificate(_In_ DWORD dwCertEncodingType, _In_ PCERT_INFO pCertId) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CCertificate certificate(CertGetSubjectCertificateFromStore(m_hStore, dwCertEncodingType, pCertId));
#pragma warning(suppress: 26486)
    return certificate;
  }

__if_exists(CertSelectCertificateChains)
{
  _Success_(return != FALSE)
  BOOL SelectCertificateChains(_In_opt_ LPCGUID pSelectionContext, _In_ DWORD dwFlags, _In_opt_ PCCERT_SELECT_CHAIN_PARA pChainParameters, _In_ DWORD cCriteria,
                               _In_opt_ PCCERT_SELECT_CRITERIA rgpCriteria, _Out_ PDWORD pcSelection, _Out_ PCCERT_CHAIN_CONTEXT** pprgpSelection) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    return CertSelectCertificateChains(pSelectionContext, dwFlags, pChainParameters, cCriteria, rgpCriteria, m_hStore, pcSelection, pprgpSelection);
  }
}

  CCRL GetCRL(_In_ const CCertificate& issuer, _In_ const CCRL& prevCRL, _Inout_ DWORD* pdwFlags) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

    CCRL crl;
    crl.m_pCRL = CertGetCRLFromStore(m_hStore, issuer.Handle(), prevCRL.Handle(), pdwFlags);
    return crl;
  }

  _Success_(return != FALSE)
  BOOL PFXExport(_Inout_ CRYPT_DATA_BLOB* pPFX, _In_ LPCWSTR szPassword, _In_opt_ void* pvPara, _In_ DWORD dwFlags) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore != nullptr);

  #pragma warning(suppress: 6387)
    return PFXExportCertStoreEx(m_hStore, pPFX, szPassword, pvPara, dwFlags);
  }

  _Success_(return != FALSE)
  BOOL PFXImport(_In_ CRYPT_DATA_BLOB* pPFX, _In_ LPCWSTR szPassword, _In_ DWORD dwFlags) noexcept
  {
    //Validate our parameters
#pragma warning(suppress: 26477)
    ATLASSERT(m_hStore == nullptr);

    m_hStore = PFXImportCertStore(pPFX, szPassword, dwFlags);
    return (m_hStore != nullptr);
  }

//Member variables
  HCERTSTORE m_hStore;
};


}; //namespace CryptoWrappers

#endif //#ifndef __CRYPTOWRAPPERSSTORE_H__
