#define CRYPT_SIGN_MESSAGE_PARA_HAS_CMS_FIELDS
#include <iostream>
#include <Windows.h>
#include <io.h>
#include <fcntl.h>
#include "PureExampleC0000225.h"

#define ERROR_BUFFER_SIZE 2048
#define MAX_SIZE_CERT_NAME 1024
#define MY_MSG "CryptoAPI is a good way to handle security"

static BYTE* msg = nullptr;
static LPWSTR errmsg = nullptr;
static LPCWSTR HandleError(DWORD err, LPCWSTR custMsg);
static PCCERT_CONTEXT
FindCertificate(
    const HCERTSTORE hStore,
    const WCHAR* CertSearchString);
static DWORD Sign(BYTE* message, DWORD cbmess_len,
    BYTE** signMessage, LPCWSTR cert, LPCWSTR* errMessage, BOOL Detached);

int main()
{
    auto a = _setmode(_fileno(stdout), _O_U16TEXT);
    a = _setmode(_fileno(stderr), _O_U16TEXT);
    a = _setmode(_fileno(stdin), _O_U16TEXT);

    BYTE* pbMessage = (BYTE*)MY_MSG;
    DWORD cbMessage = (strlen((CHAR*)pbMessage) + 1);
    BYTE* signMessage = nullptr;
    LPCWSTR errMessage = nullptr;
    DWORD result = Sign(pbMessage, cbMessage, &signMessage, L"Common name (CN) сертификата",
        &errMessage, TRUE);
    if (result != -1)
    {
        std::wcout << L"Все хорошо!\n";
        ClearSignMemory();
    }
    else
    {
        std::wcout << errMessage << "\n";
        ClearErrorMemory();
    }
}

static DWORD Sign(BYTE* message, DWORD cbmess_len,
    BYTE** signMessage, LPCWSTR cert, LPCWSTR* errMessage, BOOL Detached)
{
    DWORD cbsignMessage = 0;
    HCERTSTORE certStore = CertOpenStore(CERT_STORE_PROV_SYSTEM_W,
        X509_ASN_ENCODING | PKCS_7_ASN_ENCODING,
        NULL, CERT_STORE_OPEN_EXISTING_FLAG |
        CERT_SYSTEM_STORE_LOCAL_MACHINE | CERT_STORE_READONLY_FLAG, L"MY");
    if (certStore != nullptr)
    {
        PCCERT_CONTEXT certContext = FindCertificate(certStore, cert);
        if (certContext != nullptr)
        {
            CRYPT_OBJID_BLOB hashParam;
            hashParam.cbData = 0;
            hashParam.pbData = nullptr;
            CRYPT_ALGORITHM_IDENTIFIER hashAlg;
            hashAlg.pszObjId = (LPSTR)"1.2.643.7.1.1.2.2";
            hashAlg.Parameters = hashParam;
            CRYPT_ALGORITHM_IDENTIFIER hashEncryptionAlg;
            hashEncryptionAlg.pszObjId = (LPSTR)"1.2.643.7.1.1.1.1";
            hashEncryptionAlg.Parameters = hashParam;
            CRYPT_SIGN_MESSAGE_PARA signPara;
            signPara.cbSize = sizeof(signPara);
            signPara.dwMsgEncodingType = X509_ASN_ENCODING | PKCS_7_ASN_ENCODING;
            signPara.pSigningCert = certContext;
            signPara.HashAlgorithm = hashAlg;
            signPara.pvHashAuxInfo = nullptr;
            signPara.cMsgCert = 1;
            signPara.rgpMsgCert = &certContext;
            signPara.cMsgCrl = 0;
            signPara.cAuthAttr = 0;
            signPara.rgAuthAttr = NULL;
            signPara.cUnauthAttr = 0;
            signPara.dwFlags = 0;
            signPara.dwInnerContentType = 0;
            signPara.HashEncryptionAlgorithm = hashEncryptionAlg;
            signPara.pvHashEncryptionAuxInfo = 0;

            const BYTE* MessageArray[] = { message };
            DWORD_PTR MessageSizeArray[1];
            MessageSizeArray[0] = cbmess_len;

            // !!!!!!!!!! вот здесь бывает ошибка, 1 раз в 10 запусков примерно, случайным образом.
            if (!CryptSignMessage(
                &signPara,
                Detached,
                1,
                MessageArray,
                MessageSizeArray,
                NULL,
                &cbsignMessage))
            {
                *errMessage = HandleError(GetLastError(), L"Calculation size of buffer is failed!!!.");
                CertFreeCertificateContext(certContext);
                CertCloseStore(certStore, 0);
                return -1;
            }

            msg = new BYTE[cbsignMessage];

            if (!CryptSignMessage(
                &signPara,
                Detached,
                1,
                MessageArray,
                MessageSizeArray,
                msg,
                &cbsignMessage))
            {
                *errMessage = HandleError(GetLastError(), L"Sign message is failed.");
                delete[] msg;
                CertFreeCertificateContext(certContext);
                CertCloseStore(certStore, 0);
                return -1;
            }

            *signMessage = msg;
            CertFreeCertificateContext(certContext);
            CertCloseStore(certStore, 0);
            return cbsignMessage;
        }
        else
        {
            *errMessage = HandleError(GetLastError(), L"Find certificate is failed.");
            CertCloseStore(certStore, 0);
            return -1;
        }
    }
    else
    {
        *errMessage = HandleError(GetLastError(), L"Find certificate store is failed.");
        return -1;
    }
}

static PCCERT_CONTEXT
FindCertificate(
    const HCERTSTORE hStore,
    const WCHAR* CertSearchString)
{
    PCCERT_CONTEXT capiCertificate = NULL;
    WCHAR certname[MAX_SIZE_CERT_NAME] = { 0 };

    for (;;)
    {
        capiCertificate = CertEnumCertificatesInStore(hStore, capiCertificate);
        if (NULL == capiCertificate)
        {
            break;
        }

        if (FALSE ==
            CertGetNameStringW(capiCertificate, CERT_NAME_ATTR_TYPE,
                0, (LPSTR)szOID_COMMON_NAME, certname, MAX_SIZE_CERT_NAME))
        {
            CertFreeCertificateContext(capiCertificate);
            capiCertificate = NULL;
            break;
        }

        if ((0 == wcsncmp(certname, CertSearchString, MAX_SIZE_CERT_NAME)))
        {
            break;
        }
    }
    return capiCertificate;
}

static LPCWSTR HandleError(DWORD err, LPCWSTR custMsg)
{
    errmsg = new WCHAR[ERROR_BUFFER_SIZE];
    DWORD nSize = ERROR_BUFFER_SIZE;
    wcscpy_s(errmsg, ERROR_BUFFER_SIZE, custMsg);
    wcscat_s(errmsg, ERROR_BUFFER_SIZE, L"\n");
    wcscat_s(errmsg, ERROR_BUFFER_SIZE, L"GetLastError: ");
    _ultow_s(err, errmsg + wcslen(errmsg), nSize - wcslen(errmsg), 10);
    wcscat_s(errmsg, ERROR_BUFFER_SIZE, L"\n");
    FormatMessageW(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
        nullptr, err, 0, errmsg + wcslen(errmsg), nSize - wcslen(errmsg), nullptr);
    return errmsg;
}

void ClearSignMemory()
{
    if (msg != nullptr)
    {
        delete msg;
        msg = nullptr;
    }
}

void ClearErrorMemory()
{
    if (errmsg != nullptr)
    {
        delete errmsg;
        errmsg = nullptr;
    }
}