///////////////////////////////////////////////////////////////////////////////
// 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.
///////////////////////////////////////////////////////////////////////////////

#include "OdaCommon.h"
#include "StaticRxObject.h"
#include "ExSystemServices.h"
#include "DynamicLinker.h"
#include "diagnostics.h"
#include "RxDynamicModule.h"
#include "ExPrintConsole.h"
#include "RxObjectImpl.h"

#define STL_USING_ALL
#include "OdaSTL.h"

#include "OdCryptoServices/OdCryptoServices.h"
#include "OdCryptoServices/OdAlternativeCertificateStore.h"

#if defined(OD_WINDOWS_DESKTOP)
#include "OdCryptoServices/OdCertificateObjectWinImpl.h"
#include "OdAlternativeCertificateStoreWinImpl.h"
#else
#include "OdCryptoServices/OdCertificateObjectImpl.h"
#include "OdAlternativeCertificateStoreOpenSSLImpl.h"
#endif

#include <openssl/pem.h>
#include <openssl/x509.h>

#define STD(a) std::a

#ifdef OD_HAVE_CONSOLE_H_FILE
#include <console.h>
#endif

#define MAX_PATH_LENGTH 1024

#if defined(__GNUC__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif

#ifndef _TOOLKIT_IN_DLL_
ODRX_BEGIN_STATIC_MODULE_MAP()
ODRX_END_STATIC_MODULE_MAP()
#endif

void PrintUsage()
{
  odPrintConsoleString(OD_T("Usage:\n"));
  odPrintConsoleString(OD_T("  DigitalSignatureBasicSample [options] inputFile SIGN CertSubject CertIssuer CertSerialNumber signatureFile\n"));
  odPrintConsoleString(OD_T("  DigitalSignatureBasicSample [options] inputFile VALIDATE signatureFile\n"));
  odPrintConsoleString(OD_T("\nOptions:\n"));
  odPrintConsoleString(OD_T("  -selfsigned\n"));
  odPrintConsoleString(OD_T("      Allow self-signed certificates.\n"));
  odPrintConsoleString(OD_T("  -altstore -certsdir dir -privdir dir (-cadir dir | -cabundle file)\n"));
  odPrintConsoleString(OD_T("      Use alternative certificate store. Both -certsdir and -privdir are required. Exactly one of -cadir or -cabundle must be specified.\n"));
  odPrintConsoleString(OD_T("  -rawsig\n"));
  odPrintConsoleString(OD_T("      Use raw (non-containerized) signature and validation. When used:\n"));
  odPrintConsoleString(OD_T("        * -hash [sha1|sha256|sha384|sha512]   Hash algorithm for signing/validation (required).\n"));
  odPrintConsoleString(OD_T("        * For VALIDATE, -certfile path        Certificate file used for validation (required).\n"));
  odPrintConsoleString(OD_T("  -hash [sha1|sha256|sha384|sha512]\n"));
  odPrintConsoleString(OD_T("      Hash algorithm for raw signature (required with -rawsig).\n"));
  odPrintConsoleString(OD_T("  -certfile path\n"));
  odPrintConsoleString(OD_T("      Certificate file for validation with -rawsig (required with VALIDATE and -rawsig).\n"));
  odPrintConsoleString(OD_T("\nNotes:\n"));
  odPrintConsoleString(OD_T("  * For SIGN, CertSubject, CertIssuer, and CertSerialNumber are always required.\n"));
  odPrintConsoleString(OD_T("  * When using -rawsig, -hash must be specified.\n"));
  odPrintConsoleString(OD_T("  * For validation with -rawsig, -certfile must be specified.\n"));
  odPrintConsoleString(OD_T("  * For -altstore, always specify both -certsdir and -privdir, and one of -cadir or -cabundle.\n"));
}

OdSignatureHashAlgorithm parseHashAlgorithm(const OdString& value)
{
  OdString v(value);
  v.makeLower();
  if (v == OD_T("sha1"))
  {
    return kSHA1;
  }
  if (v == OD_T("sha256"))
  {
    return kSHA256;
  }
  if (v == OD_T("sha384"))
  {
    return kSHA384;
  }
  if (v == OD_T("sha512"))
  {
    return kSHA512;
  }
  return kSHA256;
}

OdCertificateObjectPtr loadCertificateFromFile(const OdString& certFilePath)
{
  OdAnsiString ansiCertFilePath = (const char*)certFilePath;
  BIO* certBIO = BIO_new_file(ansiCertFilePath.c_str(), "rb");
  if (!certBIO)
  {
    return OdCertificateObjectPtr();
  }
  X509* cert = PEM_read_bio_X509(certBIO, NULL, 0, NULL);
  BIO_free(certBIO);
  if (!cert)
  {
    return OdCertificateObjectPtr();
  }

#if defined(OD_WINDOWS_DESKTOP)
  int len = i2d_X509(cert, NULL);
  if (len <= 0)
  {
    X509_free(cert);
    return OdCertificateObjectPtr();
  }
  unsigned char* der = (unsigned char*)OPENSSL_malloc(len);
  if (!der)
  {
    X509_free(cert);
    return OdCertificateObjectPtr();
  }
  unsigned char* p = der;
  i2d_X509(cert, &p);
  PCCERT_CONTEXT pCertContext = CertCreateCertificateContext(
    X509_ASN_ENCODING, der, len);
  OPENSSL_free(der);
  X509_free(cert);
  if (!pCertContext)
  {
    return OdCertificateObjectPtr();
  }
  OdCertificateObjectPtr certObj = new OdCertificateObjectWinImpl(pCertContext);
  CertFreeCertificateContext(pCertContext);
  return certObj;
#else
  OdCertificateObjectPtr certObj = new OdCertificateObjectImpl(cert);
  X509_free(cert);
  return certObj;
#endif
}

void loadFileChunks(const OdString& inputFilePath, OdLinkedArray<OdBinaryData>& chunkArray)
{
  OdStreamBufPtr pFile = odrxSystemServices()->createFile(inputFilePath);
  OdUInt64 fileLen = pFile->length(), currentLen = fileLen;
  const OdUInt32 MAX_LEN_OF_CHUNK = 0xFFFFFFFF;
  while (currentLen >= MAX_LEN_OF_CHUNK)
  {
    OdBinaryData chunk;
    chunk.resize(MAX_LEN_OF_CHUNK);
    pFile->getBytes(chunk.asArrayPtr(), MAX_LEN_OF_CHUNK);
    chunkArray.append(chunk);
    currentLen -= MAX_LEN_OF_CHUNK;
  }
  if (currentLen > 0)
  {
    OdBinaryData chunk;
    chunk.resize((OdUInt32)currentLen);
    pFile->getBytes(chunk.asArrayPtr(), (OdUInt32)currentLen);
    chunkArray.append(chunk);
  }
}

OdString formatCertInfo(const OdCertificateDescription& certDesc)
{
  return L"Subject\t\t=\t" + certDesc.m_CertSubject +
    L"\nIssuer\t\t=\t" + certDesc.m_CertIssuer +
    L"\nSerial number\t=\t" + certDesc.m_CertSerialNum +
    L"\nValid from\t\t=\t" + certDesc.m_CertValidFrom +
    L"\nValid to\t\t=\t" + certDesc.m_CertValidTo;
}

struct CryptoEnv
{
  OdCryptoServicesPtr pCryptoServices;
  OdAlternativeCertificateStorePtr pAltStore;

  CryptoEnv() {}

  void init(
    bool allowSelfSigned,
    bool useAltStore,
    const OdString& altCertsDir,
    const OdString& altPrivDir,
    const OdString& altCaDir,
    const OdString& altCaBundle,
    bool altStoreCadirChosen,
    bool altStoreCabundleChosen)
  {
    OdRxClassPtr pService = odrxServiceDictionary()->getAt(OD_T("OdCryptoServices"));
    if (pService.isNull())
    {
      throw OdError(OD_T("Can't init CryptoServices!"));
    }
    pCryptoServices = pService->create();
    if (allowSelfSigned)
    {
      pCryptoServices->setAllowSelfSignedCerts(true);
    }

    if (useAltStore)
    {
#if defined(OD_WINDOWS_DESKTOP)
      pAltStore = OdRxObjectImpl<OdAlternativeCertificateStoreWinImpl>::createObject();
#else
      pAltStore = OdRxObjectImpl<OdAlternativeCertificateStoreOpenSSLImpl>::createObject();
#endif
      pAltStore->setCertsDirectory(altCertsDir);
      pAltStore->setPrivateKeysDirectory(altPrivDir);
      if (altStoreCadirChosen)
      {
        pAltStore->setCaDirectory(altCaDir);
      }
      else if (altStoreCabundleChosen)
      {
        pAltStore->setCaBundleFile(altCaBundle);
      }
      pAltStore->setAllowSelfSignedCerts(allowSelfSigned);
      pCryptoServices->setAlternativeCertificateStore(pAltStore);
    }
  }

  ~CryptoEnv()
  {
  }
};

void validate(
  OdCryptoServices* pCryptoServices,
  const OdString& inputFilePath,
  const OdString& signaturePath,
  bool rawSignature,
  OdSignatureHashAlgorithm hashAlg,
  const OdString& certFilePath)
{
  OdLinkedArray<OdBinaryData> chunkArray;
  loadFileChunks(inputFilePath, chunkArray);

  OdBinaryData signatureBlock;
  OdStreamBufPtr pSignatureFile = odrxSystemServices()->createFile(signaturePath);
  OdUInt64 signatureLen = pSignatureFile->length();
  signatureBlock.resize((OdUInt32)signatureLen);
  pSignatureFile->getBytes(signatureBlock.asArrayPtr(), (OdUInt32)signatureLen);

  OdCertificateObjectPtr pCertObj;

  if (!rawSignature)
  {
    OdCryptoServices::OdSignatureVerificationResult verificationResult = OdCryptoServices::kHasNoSignature;
    pCryptoServices->verifyDetachedSignature(chunkArray, signatureBlock, verificationResult);

    OdString resultMessage;
    switch (verificationResult)
    {
    case OdCryptoServices::kSuccess:
      resultMessage = L"The CMS (containerized) digital signature is valid.";
      break;
    case OdCryptoServices::kBadSignature:
      resultMessage = L"Bad signature!";
      break;
    case OdCryptoServices::kCertificateChainProblem:
      resultMessage = L"Problem with one of the certificates in the path!";
      break;
    default:
      resultMessage = L"Error during verification!";
      break;
    }
    pCertObj = pCryptoServices->getCertFromDetachedSignature(signatureBlock);
    odPrintConsoleString(OD_T("\n%ls\n"), resultMessage.c_str());
    odPrintConsoleString(OD_T("Digital ID (Certificate) info:\n"));
    if (!pCertObj.isNull())
    {
      odPrintConsoleString(OD_T("%ls\n"), formatCertInfo(pCertObj->getCertDescription()).c_str());
    }
    else
    {
      odPrintConsoleString(OD_T("Certificate not found.\n"));
    }
  }
  else // raw signature
  {
    OdBinaryData mergedData;
    for (auto it = chunkArray.begin(); it != chunkArray.end(); ++it)
    {
      mergedData.insert(mergedData.end(), it->begin(), it->end());
    }

    if (!certFilePath.isEmpty())
    {
      pCertObj = loadCertificateFromFile(certFilePath);
    }

    if (!pCertObj.isNull())
    {
      OdCryptoServices::OdSignatureVerificationResult verificationResult = pCryptoServices->verifySignature(
        pCertObj, hashAlg, mergedData, signatureBlock);

      OdString resultMessage;
      switch (verificationResult)
      {
      case OdCryptoServices::kSuccess:
        resultMessage = L"The non-containerized digital signature is valid.";
        break;
      case OdCryptoServices::kBadSignature:
        resultMessage = L"Bad signature (raw)!";
        break;
      case OdCryptoServices::kCertificateChainProblem:
        resultMessage = L"Certificate chain problem!";
        break;
      case OdCryptoServices::kNoSigner:
        resultMessage = L"No signer!";
        break;
      case OdCryptoServices::kBadAlgId:
        resultMessage = L"Bad hash algorithm for verification!";
        break;
      default:
        resultMessage = L"Error during verification!";
        break;
      }
      odPrintConsoleString(OD_T("\n%ls\n"), resultMessage.c_str());
      odPrintConsoleString(OD_T("Digital ID (Certificate) info:\n"));
      odPrintConsoleString(OD_T("%ls\n"), formatCertInfo(pCertObj->getCertDescription()).c_str());
    }
    else
    {
      odPrintConsoleString(OD_T("\nCould not find certificate for verification!\n"));
    }
  }
}

void sign(
  OdCryptoServices* pCryptoServices,
  const OdString& inputFilePath,
  const OdCertificateShortDesc& certShortDesc,
  const OdString& signaturePath,
  bool rawSignature,
  OdSignatureHashAlgorithm hashAlg)
{
  OdCertificateObjectPtr pCertObj = pCryptoServices->getCertObjByShortDesc(certShortDesc);

  if (pCertObj.isNull())
  {
    odPrintConsoleString(OD_T("\nCannot find certificate for specified subject/issuer/serial!\n"));
    throw OdError(OD_T("No certificate found!"));
  }

  OdLinkedArray<OdBinaryData> chunkArrayToSign;
  loadFileChunks(inputFilePath, chunkArrayToSign);

  OdBinaryData signatureBlock;

  if (!rawSignature)
  {
    OdSubjectKeyIdAttribPtr pAttrib;
    OdCryptSignMessageParaPtr pSignPara = pCryptoServices->newCryptSignMessagePara(pCertObj, pAttrib);

    if (!pCryptoServices->generateDetachedSignature(pSignPara, chunkArrayToSign, signatureBlock))
    {
      throw OdError(OD_T("Error while CMS signature generation!"));
    }
    else
    {
      OdStreamBufPtr pSignatureFile = odrxSystemServices()->createFile(signaturePath, Oda::kFileWrite, Oda::kShareDenyNo, Oda::kCreateAlways);
      pSignatureFile->putBytes(signatureBlock.asArrayPtr(), signatureBlock.size());
      odPrintConsoleString(OD_T("\nGenerated CMS (containerized) signature saved to %ls\n"), signaturePath.c_str());
    }
  }
  else
  {
    OdBinaryData mergedData;
    for (auto it = chunkArrayToSign.begin(); it != chunkArrayToSign.end(); ++it)
    {
      mergedData.insert(mergedData.end(), it->begin(), it->end());
    }

    OdSignDataStatus signDataStatus = pCryptoServices->signData(pCertObj, hashAlg, mergedData, signatureBlock);
    if (signDataStatus.result != kSignData_Success)
    {
      throw OdError(signDataStatus.description);
    }
    else
    {
      OdStreamBufPtr pSignatureFile = odrxSystemServices()->createFile(signaturePath, Oda::kFileWrite, Oda::kShareDenyNo, Oda::kCreateAlways);
      pSignatureFile->putBytes(signatureBlock.asArrayPtr(), signatureBlock.size());
      odPrintConsoleString(OD_T("\nGenerated non-containerized signature saved to %ls\n"), signaturePath.c_str());
    }
  }
}

void uninitHelperFunction()
{
#if defined(OD_WINDOWS_DESKTOP)
  OdAlternativeCertificateStoreWinImpl::rxUninit();
#else
  OdAlternativeCertificateStoreOpenSSLImpl::rxUninit();
#endif
  ::odrxUninitialize();
}

#if defined(OD_USE_WMAIN)
int wmain(int argc, wchar_t* argv[])
#else
int main(int argc, char* argv[])
#endif
{
#ifdef OD_HAVE_CCOMMAND_FUNC
  argc = ccommand(&argv);
#endif

#ifndef _TOOLKIT_IN_DLL_
  ODRX_INIT_STATIC_MODULE_MAP();
#endif
  OdStaticRxObject<ExSystemServices> svcs;
  odrxInitialize(&svcs);

#if defined(OD_WINDOWS_DESKTOP)
  OdAlternativeCertificateStoreWinImpl::rxInit();
#else
  OdAlternativeCertificateStoreOpenSSLImpl::rxInit();
#endif

  bool allowSelfSigned = false;
  bool useAltStore = false;
  bool rawSignature = false;
  bool altStoreCadirChosen = false;
  bool altStoreCabundleChosen = false;
  OdString altCertsDir, altPrivDir, altCaDir, altCaBundle, certFilePath;
  OdSignatureHashAlgorithm hashAlg = kSHA256;
  int argIdx = 1;

  while (argIdx < argc)
  {
    OdString curArg(argv[argIdx]);
    curArg.makeLower();
    if (curArg == OD_T("-selfsigned"))
    {
      allowSelfSigned = true;
      ++argIdx;
    }
    else if (curArg == OD_T("-altstore"))
    {
      useAltStore = true;
      ++argIdx;
      while (argIdx < argc)
      {
        OdString storeArg(argv[argIdx]);
        storeArg.makeLower();
        if (storeArg == OD_T("-certsdir"))
        {
          if (argIdx + 1 < argc)
          {
            altCertsDir = OdString(argv[argIdx + 1]);
            argIdx += 2;
          }
          else
          {
            PrintUsage(); uninitHelperFunction();
            return 10;
          }
        }
        else if (storeArg == OD_T("-privdir"))
        {
          if (argIdx + 1 < argc)
          {
            altPrivDir = OdString(argv[argIdx + 1]);
            argIdx += 2;
          }
          else
          {
            PrintUsage(); uninitHelperFunction();
            return 10;
          }
        }
        else if (storeArg == OD_T("-cadir"))
        {
          if (altStoreCabundleChosen)
          {
            odPrintConsoleString(OD_T("Error: only one of -cadir or -cabundle may be specified for -altstore.\n"));
            PrintUsage(); uninitHelperFunction();
            return 10;
          }
          if (argIdx + 1 < argc)
          {
            altCaDir = OdString(argv[argIdx + 1]);
            altStoreCadirChosen = true;
            argIdx += 2;
          }
          else
          {
            PrintUsage(); uninitHelperFunction();
            return 10;
          }
        }
        else if (storeArg == OD_T("-cabundle"))
        {
          if (altStoreCadirChosen)
          {
            odPrintConsoleString(OD_T("Error: only one of -cadir or -cabundle may be specified for -altstore.\n"));
            PrintUsage(); uninitHelperFunction();
            return 10;
          }
          if (argIdx + 1 < argc)
          {
            altCaBundle = OdString(argv[argIdx + 1]);
            altStoreCabundleChosen = true;
            argIdx += 2;
          }
          else
          {
            PrintUsage(); uninitHelperFunction();
            return 10;
          }
        }
        else
        {
          break;
        }
      }
    }
    else if (curArg == OD_T("-rawsig"))
    {
      rawSignature = true;
      ++argIdx;
    }
    else if (curArg == OD_T("-hash"))
    {
      if (argIdx + 1 < argc)
      {
        hashAlg = parseHashAlgorithm(OdString(argv[argIdx + 1]));
        argIdx += 2;
      }
      else
      {
        PrintUsage(); uninitHelperFunction();
        return 10;
      }
    }
    else if (curArg == OD_T("-certfile"))
    {
      if (argIdx + 1 < argc)
      {
        certFilePath = OdString(argv[argIdx + 1]);
        argIdx += 2;
      }
      else
      {
        PrintUsage(); uninitHelperFunction();
        return 10;
      }
    }
    else
    {
      break;
    }
  }

  if (useAltStore)
  {
    if (altCertsDir.isEmpty() || altPrivDir.isEmpty() ||
      (!altStoreCadirChosen && !altStoreCabundleChosen) ||
      (altStoreCadirChosen && altStoreCabundleChosen))
    {
      odPrintConsoleString(OD_T("Error: -altstore requires both -certsdir and -privdir, and exactly one of -cadir or -cabundle.\n"));
      PrintUsage(); uninitHelperFunction();
      return 10;
    }
  }

  if (rawSignature && (hashAlg != kSHA1 && hashAlg != kSHA256 && hashAlg != kSHA384 && hashAlg != kSHA512))
  {
    odPrintConsoleString(OD_T("Error: -rawsig requires -hash argument (sha1|sha256|sha384|sha512).\n"));
    PrintUsage(); uninitHelperFunction();
    return 10;
  }

  int cmdArgs = argc - argIdx;
  if (cmdArgs < 3)
  {
    PrintUsage(); uninitHelperFunction();
    return 10;
  }

  OdString filePath(argv[argIdx]);
  OdString operation(argv[argIdx + 1]);
  int status = 0;

  try
  {
    CryptoEnv env;
    env.init(
      allowSelfSigned, useAltStore, altCertsDir, altPrivDir,
      altCaDir, altCaBundle, altStoreCadirChosen, altStoreCabundleChosen
    );
    
    if (operation.compare(OD_T("SIGN")) == 0)
    {
      if (cmdArgs != 6)
      {
        PrintUsage(); uninitHelperFunction();
        return 10;
      }
      OdCertificateShortDesc certShortDesc;
      certShortDesc.m_CertSubject = OdString(argv[argIdx + 2]);
      certShortDesc.m_CertIssuer = OdString(argv[argIdx + 3]);
      certShortDesc.m_CertSerialNum = OdString(argv[argIdx + 4]);
      OdString signaturePath(argv[argIdx + 5]);
      sign(env.pCryptoServices, filePath, certShortDesc, signaturePath, rawSignature, hashAlg);
    }
    else if (operation.compare(OD_T("VALIDATE")) == 0)
    {
      bool validRaw = rawSignature && !certFilePath.isEmpty() && (cmdArgs == 4);
      if ((cmdArgs != 3) && !validRaw)
      {
        PrintUsage(); uninitHelperFunction();
        return 10;
      }
      OdString signaturePath(argv[argIdx + cmdArgs - 1]);
      if (rawSignature && certFilePath.isEmpty())
      {
        odPrintConsoleString(OD_T("Error: -rawsig with VALIDATE requires -certfile.\n"));
        PrintUsage(); uninitHelperFunction();
        return 10;
      }
      validate(env.pCryptoServices, filePath, signaturePath, rawSignature, hashAlg, certFilePath);
    }
    else
    {
      PrintUsage(); uninitHelperFunction();
      return 10;
    }
  }
  catch (OdError& e)
  {
    odPrintConsoleString(OD_T("Exception (%ls) during the file processing!\n"), e.description().c_str());
    status = 10;
  }
  catch (...)
  {
    odPrintConsoleString(OD_T("Unknown Exception during the file processing!\n"));
    status = 10;
  }

  uninitHelperFunction();

  return status;
}

#if defined(__GNUC__)
#pragma GCC diagnostic pop
#endif