/////////////////////////////////////////////////////////////////////////////// 
// Copyright (C) 2002-2025, Open Design Alliance (the "Alliance"). 
// All rights reserved. 
// 
// This software and its documentation and related materials are owned by 
// the Alliance. The software may only be incorporated into application 
// programs owned by members of the Alliance, subject to a signed 
// Membership Agreement and Supplemental Software License Agreement with the
// Alliance. The structure and organization of this software are the valuable  
// trade secrets of the Alliance and its suppliers. The software is also 
// protected by copyright law and international treaty provisions. Application  
// programs incorporating this software must include the following statement 
// with their copyright notices:
//   
//   This application incorporates Open Design Alliance software pursuant to a license 
//   agreement with Open Design Alliance.
//   Open Design Alliance Copyright (C) 2002-2025 by Open Design Alliance. 
//   All rights reserved.
//
// By use of this software, its documentation or related materials, you 
// acknowledge and accept the above terms.
///////////////////////////////////////////////////////////////////////////////

#ifndef _WIN32_WINNT 
#define _WIN32_WINNT 0x0601
#endif

#include "OdAlternativeCertificateStoreWinImpl.h"
#include "OdCryptoServices/OdCertificateObjectWinImpl.h"

#include <wincrypt.h>

#define STL_USING_ALL
#include "OdaSTL.h"

#include <openssl/bio.h>
#include <openssl/pem.h>
#include <openssl/pkcs12.h>

#include "RxObjectImpl.h"

ODRX_CONS_DEFINE_MEMBERS(OdAlternativeCertificateStoreWinImpl, OdAlternativeCertificateStore, RXIMPL_CONSTR);

OdAlternativeCertificateStoreWinImpl::OdAlternativeCertificateStoreWinImpl()
  : m_hMemoryStore(nullptr), m_loaded(false), m_allowSelfSignedCerts(false)
{
}

OdAlternativeCertificateStoreWinImpl::~OdAlternativeCertificateStoreWinImpl()
{
  if (m_hMemoryStore)
    CertCloseStore(m_hMemoryStore, 0);
}

OdUInt32 OdAlternativeCertificateStoreWinImpl::getPersonalCertsWithTrustedStatus(OdArray<OdCertificateDescription>& certificates) const
{
  ensureLoaded();

  OdUInt32 result = 0;
  for (OdUInt32 i = 0; i < m_certObjects.size(); ++i)
  {
    OdCertificateObjectPtr certObj = m_certObjects[i];
    if (certObj.isNull())
      continue;

    if (m_allowSelfSignedCerts || isTrusted(certObj))
    {
      OdCertificateDescription desc = certObj->getCertDescription();
      certificates.push_back(desc);
      result++;
    }
  }
  return result;
}

OdCertificateObjectPtr OdAlternativeCertificateStoreWinImpl::getCertObjByShortDesc(const OdCertificateShortDesc& certShortDesc) const
{
  ensureLoaded();

  for (OdUInt32 i = 0; i < m_certObjects.size(); ++i)
  {
    OdCertificateObjectPtr certObj = m_certObjects[i];
    if (certObj.isNull())
      continue;

    OdCertificateDescription desc = certObj->getCertDescription();

    if ((desc.m_CertSubject == certShortDesc.m_CertSubject) &&
      (desc.m_CertIssuer == certShortDesc.m_CertIssuer) &&
      (desc.m_CertSerialNum == certShortDesc.m_CertSerialNum))
    {
      if (m_allowSelfSignedCerts || isTrusted(certObj))
        return certObj;
    }
  }
  return NULL;
}

bool OdAlternativeCertificateStoreWinImpl::isTrusted(OdCertificateObjectPtr cert) const
{
  ensureLoaded();

  if (!cert.get())
    return false;

  PCCERT_CONTEXT pCertContext = ((OdCertificateObjectWinImpl*)cert.get())->getCertContext();
  if (!pCertContext)
    return false;

  bool res = false;
  HCERTCHAINENGINE hChainEngine = NULL;
  CERT_CHAIN_ENGINE_CONFIG ChainConfig = {};
  ChainConfig.cbSize = sizeof(CERT_CHAIN_ENGINE_CONFIG);
  ChainConfig.hRestrictedRoot = NULL;
  ChainConfig.cAdditionalStore = 1;
  HCERTSTORE additionalStores[1] = { m_hMemoryStore };
  ChainConfig.rghAdditionalStore = additionalStores;
  ChainConfig.dwFlags = CERT_CHAIN_CACHE_END_CERT;

  PCCERT_CHAIN_CONTEXT pChainContext = NULL;

  if (CertCreateCertificateChainEngine(&ChainConfig, &hChainEngine))
  {
    CERT_ENHKEY_USAGE EnhkeyUsage = {};
    CERT_USAGE_MATCH  CertUsage = {};
    CERT_CHAIN_PARA   ChainPara = {};
    EnhkeyUsage.cUsageIdentifier = 0;
    EnhkeyUsage.rgpszUsageIdentifier = NULL;
    CertUsage.dwType = USAGE_MATCH_TYPE_AND;
    CertUsage.Usage = EnhkeyUsage;
    ChainPara.cbSize = sizeof(CERT_CHAIN_PARA);
    ChainPara.RequestedUsage = CertUsage;
    DWORD dwFlags = 0;

    if (CertGetCertificateChain(hChainEngine, pCertContext, NULL,
      m_hMemoryStore, &ChainPara, dwFlags, NULL, &pChainContext))
    {
      CERT_TRUST_STATUS certTrustStatus = pChainContext->TrustStatus;
      if (certTrustStatus.dwErrorStatus == CERT_TRUST_NO_ERROR)
        res = true;
    }
  }

  if (hChainEngine)
    CertFreeCertificateChainEngine(hChainEngine);
  if (pChainContext)
    CertFreeCertificateChain(pChainContext);

  return res;
}

void OdAlternativeCertificateStoreWinImpl::setCertsDirectory(const OdString& dir)
{
  m_certsDir = dir;
}

OdString OdAlternativeCertificateStoreWinImpl::certsDirectory() const
{
  return m_certsDir;
}

void OdAlternativeCertificateStoreWinImpl::setPrivateKeysDirectory(const OdString& dir)
{
  m_privateKeysDir = dir;
}

OdString OdAlternativeCertificateStoreWinImpl::privateKeysDirectory() const
{
  return m_privateKeysDir;
}

void OdAlternativeCertificateStoreWinImpl::setCaDirectory(const OdString& dir)
{
  m_caDirectory = dir;
}

OdString OdAlternativeCertificateStoreWinImpl::caDirectory() const
{
  return m_caDirectory;
}

void OdAlternativeCertificateStoreWinImpl::setCaBundleFile(const OdString& fullPathToCaBundle)
{
  m_caBundleFile = fullPathToCaBundle;
}

OdString OdAlternativeCertificateStoreWinImpl::caBundleFile() const
{
  return m_caBundleFile;
}

void OdAlternativeCertificateStoreWinImpl::ensureLoaded() const
{
  if (!m_loaded)
  {
    if (!m_hMemoryStore)
      m_hMemoryStore = CertOpenStore(CERT_STORE_PROV_MEMORY, 0, 0, 0, 0);
    
    loadCertificates();
    loadCACertificates();
    loadCABundle();

    m_loaded = true;
  }
}

static void findFiles(const OdString& dirPath, OdArray<OdString>& fileList, const OdString& ext)
{
  fileList.clear();
  OdString searchPath;
  searchPath.format(OD_T("%ls\\*.%ls"), dirPath.c_str(), ext.c_str());

  WIN32_FIND_DATA findFileData;
  HANDLE hFind = FindFirstFile(searchPath.c_str(), &findFileData);

  if (hFind != INVALID_HANDLE_VALUE)
  {
    do
    {
      if (!(findFileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY))
      {
        OdString filePath;
        filePath.format(OD_T("%ls\\%ls"), dirPath.c_str(), findFileData.cFileName);
        fileList.push_back(filePath);
      }
    } while (FindNextFile(hFind, &findFileData));
    FindClose(hFind);
  }
}

bool certificateExistsInStore(HCERTSTORE hStore, PCCERT_CONTEXT pCertContext)
{
  BYTE hash[20];
  DWORD hashLen = sizeof(hash);
  if (!CertGetCertificateContextProperty(pCertContext, CERT_HASH_PROP_ID, hash, &hashLen))
    return false;

  PCCERT_CONTEXT pEnum = NULL;
  while ((pEnum = CertEnumCertificatesInStore(hStore, pEnum)) != NULL) {
    BYTE storeHash[20];
    DWORD storeHashLen = sizeof(storeHash);
    if (CertGetCertificateContextProperty(pEnum, CERT_HASH_PROP_ID, storeHash, &storeHashLen)) {
      if (hashLen == storeHashLen && memcmp(hash, storeHash, hashLen) == 0) {
        // Duplicate found
        return true;
      }
    }
  }
  return false;
}

void OdAlternativeCertificateStoreWinImpl::loadCertificates() const
{
  m_certObjects.clear();

  m_certFileNames.clear();
  m_privateKeyFileNames.clear();

  findFiles(m_certsDir, m_certFileNames, OD_T("pem"));
  findFiles(m_privateKeysDir, m_privateKeyFileNames, OD_T("pem"));

  for (OdUInt32 i = 0; i < m_privateKeyFileNames.size(); i++)
  {
    BIO* keyBIO = BIO_new(BIO_s_file());

    OdAnsiString ansiPrivateKeyFileName = (const char*)(m_privateKeyFileNames[i]);

    if (1 == BIO_read_filename(keyBIO, ansiPrivateKeyFileName.c_str()))
    {
      EVP_PKEY* privKey = NULL;
      PEM_read_bio_PrivateKey(keyBIO, &privKey, NULL, NULL);
      if (privKey)
      {
        for (OdUInt32 j = 0; j < m_certFileNames.size(); j++)
        {
          BIO* certBIO = BIO_new(BIO_s_file());

          OdAnsiString ansiCertFileName = (const char*)(m_certFileNames[j]);

          if (1 == BIO_read_filename(certBIO, ansiCertFileName.c_str()))
          {
            X509* currentCert = NULL;
            currentCert = PEM_read_bio_X509(certBIO, NULL, 0, NULL);
            if (currentCert)
            {
              if (1 == X509_check_private_key(currentCert, privKey))
              {
                PKCS12* p12 = PKCS12_create(
                  "",
                  ansiPrivateKeyFileName.c_str(),
                  privKey,
                  currentCert,
                  NULL, 0, 0, 0, 0, 0
                );
                
                BIO* mem = BIO_new(BIO_s_mem());
                i2d_PKCS12_bio(mem, p12);

                PKCS12_free(p12);

                char* pfx_data = nullptr;
                long pfx_size = BIO_get_mem_data(mem, &pfx_data);

                DATA_BLOB pfxBlob;
                pfxBlob.pbData = (BYTE*)pfx_data;
                pfxBlob.cbData = (DWORD)pfx_size;

                HCERTSTORE hMemStore = PFXImportCertStore(
                  &pfxBlob,
                  NULL,
                  CRYPT_EXPORTABLE
                );

                BIO_free(mem);

                if (hMemStore)
                {
                  PCCERT_CONTEXT pCertContext = CertEnumCertificatesInStore(hMemStore, NULL);
                  if (pCertContext) {
                    if (!certificateExistsInStore(m_hMemoryStore, pCertContext)) {
                      CertAddCertificateContextToStore(m_hMemoryStore, pCertContext, CERT_STORE_ADD_ALWAYS, NULL);
                      OdCertificateObjectPtr obj = new OdCertificateObjectWinImpl(pCertContext);
                      m_certObjects.push_back(obj);
                    }
                    CertFreeCertificateContext(pCertContext);
                  }
                  CertCloseStore(hMemStore, 0);
                }
                
                X509_free(currentCert);
                BIO_free(certBIO);
                break;
              }

              X509_free(currentCert);
            }
          }

          if (certBIO)
            BIO_free(certBIO);
        }

        EVP_PKEY_free(privKey);
      }
    }

    if (keyBIO)
      BIO_free(keyBIO);
  }
}

void OdAlternativeCertificateStoreWinImpl::loadCACertificates() const
{
  OdArray<OdString> caFileNames;
  findFiles(caDirectory(), caFileNames, OD_T("pem"));

  for (OdUInt32 i = 0; i < caFileNames.size(); ++i)
  {
    OdAnsiString ansiCAFName = (const char*)(caFileNames[i]);
    BIO* caBIO = BIO_new_file(ansiCAFName.c_str(), "rb");
    if (!caBIO)
      continue;

    while (true)
    {
      X509* x509ca = PEM_read_bio_X509(caBIO, NULL, 0, NULL);
      if (!x509ca)
        break;

      int certLen = i2d_X509(x509ca, NULL);
      if (certLen > 0)
      {
        OdBinaryData encodedCert;
        encodedCert.resize(certLen);

        unsigned char* p = (unsigned char*)encodedCert.asArrayPtr();
        i2d_X509(x509ca, &p);

        PCCERT_CONTEXT pCertContext = CertCreateCertificateContext(
          X509_ASN_ENCODING,
          encodedCert.asArrayPtr(),
          encodedCert.size()
        );
        if (pCertContext)
        {
          if (!certificateExistsInStore(m_hMemoryStore, pCertContext)) {
            CertAddCertificateContextToStore(m_hMemoryStore, pCertContext, CERT_STORE_ADD_ALWAYS, NULL);
          }
          CertFreeCertificateContext(pCertContext);
        }
      }
      X509_free(x509ca);
    }
    BIO_free(caBIO);
  }
}

void OdAlternativeCertificateStoreWinImpl::loadCABundle() const
{
  OdString caBundlePath = caBundleFile();
  if (caBundlePath.isEmpty())
    return;

  BIO* bio = BIO_new_file((const char*)caBundlePath, "rb");
  if (!bio)
    return;

  while (true)
  {
    X509* x509ca = PEM_read_bio_X509(bio, NULL, 0, NULL);
    if (!x509ca)
      break;

    int certLen = i2d_X509(x509ca, NULL);
    if (certLen > 0)
    {
      OdBinaryData encodedCert;
      encodedCert.resize(certLen);

      unsigned char* p = (unsigned char*)encodedCert.asArrayPtr();
      i2d_X509(x509ca, &p);

      PCCERT_CONTEXT pCertContext = CertCreateCertificateContext(
        X509_ASN_ENCODING,
        encodedCert.asArrayPtr(),
        encodedCert.size()
      );
      if (pCertContext)
      {
        if (!certificateExistsInStore(m_hMemoryStore, pCertContext)) {
          CertAddCertificateContextToStore(m_hMemoryStore, pCertContext, CERT_STORE_ADD_ALWAYS, NULL);
        }
        CertFreeCertificateContext(pCertContext);
      }
    }
    X509_free(x509ca);
  }
  BIO_free(bio);
}

void OdAlternativeCertificateStoreWinImpl::setAllowSelfSignedCerts(bool allow)
{
  m_allowSelfSignedCerts = allow;
}

void OdAlternativeCertificateStoreWinImpl::getAllPrivateKeyFileNames(OdStringArray& privateKeyFileNames)
{
  ensureLoaded();

  privateKeyFileNames = m_privateKeyFileNames;
}