/*
* Copyright(C) 2000-2010 Проект ИОК
*
* Этот файл содержит информацию, являющуюся
* собственностью компании Крипто Про.
*
* Любая часть этого файла не может быть скопирована,
* исправлена, переведена на другие языки,
* локализована или модифицирована любым способом,
* откомпилирована, передана по сети с или на
* любую компьютерную систему без предварительного
* заключения соглашения с компанией Крипто Про.
*
* Программный код, содержащийся в этом файле, предназначен
* исключительно для целей обучения и не может быть использован
* для защиты информации.
*
* Компания Крипто-Про не несет никакой
* ответственности за функционирование этого кода.
*/

#include <stdio.h>
#ifdef _WIN32
#   include <windows.h>
#   include <wincrypt.h>
#   include <tchar.h>
#else
#   include <stdlib.h>
#   include "reader/tchar.h"
#endif
#include <string.h>
#include <WinCryptEx.h>

// Начало примера (не следует удалять данный комментарий, он используется 
// для автоматической сборки документации)
//--------------------------------------------------------------------
// Пример кода для зашифрования данных и создания зашифрованного 
// сообщения при помощи поточного режима функций CryptMsgXXX.

// Для функционирования данного кода необходимы:
// - контейнер с ключом AT_KEYEXCHANGE в провайдере PROV_GOST_2012_256
// - сертификат этого ключа, установленный в хранилище пользователя ("MY")
//--------------------------------------------------------------------

#define MY_ENCODING_TYPE  (PKCS_7_ASN_ENCODING | X509_ASN_ENCODING)

static void CleanUp(void);
static void HandleError(const char *s);
static PCCERT_CONTEXT GetRecipientCert(HCERTSTORE hCertStore);
static void DecryptMessage(BYTE *pbEncryptedBlob, DWORD cbEncryptedBlob, PCCERT_CONTEXT pRecipientCert);
static void GetCertDName(PCERT_NAME_BLOB pNameBlob, char **pszName);

static HCRYPTPROV hCryptProv = 0;        // дескриптор CSP 
static HCERTSTORE hStoreHandle = 0;      // дескриптор хранилища сертификатов
static PCCERT_CONTEXT pRecipientCert = NULL;
static char *szDName = NULL;		  // DName сертификата

typedef struct _MEMBUF
{
    BYTE* pbData;
    DWORD cbData;
    DWORD cbSize;
} MEMBUF;

static BOOL WINAPI StreamEncryptDecryptCallback(
    IN const void *pvArg,  //in
    IN BYTE *pbData,       //in
    IN DWORD cbData,       //in
    IN BOOL fFinal)
{
    MEMBUF *pBuf = (MEMBUF*)pvArg;
    (void)(fFinal);
    if (pBuf->cbSize < pBuf->cbData + cbData)
    {
	pBuf->cbSize += 2 * cbData;
	pBuf->pbData = realloc(pBuf->pbData, pBuf->cbSize);
    }
    memcpy(pBuf->pbData + pBuf->cbData, pbData, cbData);
    pBuf->cbData += cbData;
    return TRUE;
}

#undef MIN
#define MIN(x,y) ((x)<(y)?(x):(y))

int main(void)
{
    // Объявление и инициализация переменных. Они получают указатель на  
    // сообщение, которое будет зашифровано. В данном коде создается сообщение,
    // получается указатель на него.

    BYTE pbContent[0x10000] = "Security is our business"; // Сообщение
    DWORD cbContent = sizeof(pbContent);	   // Длина сообщения, включая конечный 0

    CRYPT_ALGORITHM_IDENTIFIER EncryptAlgorithm;
   
    printf("source message: %s\n", pbContent);
    printf("message length: %d bytes \n", cbContent);

    // Получение дескриптора криптографического провайдера.
    if (!CryptAcquireContext(
	&hCryptProv,         // Адрес возврашаемого дескриптора.
	0,                // Используется имя текущего зарегестрированного пользователя.
	NULL,                // Используется провайдер по умолчанию.
	PROV_GOST_2012_256,   // Необходимо для зашифрования и подписи.
	CRYPT_VERIFYCONTEXT))		     // Никакие флаги не нужны.
    {
	HandleError("Cryptographic context could not be acquired.");
    }
    printf("CSP has been acquired. \n");

    // Открытие системного хранилища сертификатов.
    hStoreHandle = CertOpenSystemStore(hCryptProv, _TEXT("MY"));

    if (!hStoreHandle)
    {
	HandleError("Error getting store handle.");
    }
    printf("The MY store is open. \n");

    // Получение указателя на сертификат получателя с помощью
    // функции GetRecipientCert. 
    pRecipientCert = GetRecipientCert(hStoreHandle);

    if (!pRecipientCert)
    {
	printf("No certificate with a CERT_KEY_CONTEXT_PROP_ID \n");
	printf("property and an AT_KEYEXCHANGE private key available. \n");
	printf("While the message could be encrypted, in this case, \n");
	printf("it could not be decrypted in this program. \n");
	printf("For more information, see the documentation for \n");
	printf("CryptEncryptMessage and CryptDecryptMessage.\n\n");
	HandleError("No Certificate with AT_KEYEXCHANGE key in store.");
    }
    GetCertDName(&pRecipientCert->pCertInfo->Subject, &szDName);
    printf("A recipient's certificate has been acquired: %s\n", szDName);

    // Инициализация структуры с нулем.
    memset(&EncryptAlgorithm, 0, sizeof(CRYPT_ALGORITHM_IDENTIFIER));
    EncryptAlgorithm.pszObjId = szOID_CP_GOST_28147;

    CMSG_ENVELOPED_ENCODE_INFO envelopedInfo = { sizeof(envelopedInfo) };

    PCERT_INFO pCertInfo = pRecipientCert->pCertInfo;

    envelopedInfo.ContentEncryptionAlgorithm = EncryptAlgorithm;
    envelopedInfo.cRecipients = 1;
    envelopedInfo.rgpRecipients = &pCertInfo;
    envelopedInfo.pvEncryptionAuxInfo = 0;
    envelopedInfo.hCryptProv = hCryptProv;

    MEMBUF EcryptedMessage = { 0 };

    // Fill the CMSG_STREAM_INFO structure.
    CMSG_STREAM_INFO MsgStreamInfo;
    memset(&MsgStreamInfo, 0, sizeof(CMSG_STREAM_INFO));
    MsgStreamInfo.cbContent = CMSG_INDEFINITE_LENGTH;
    MsgStreamInfo.pfnStreamOutput = StreamEncryptDecryptCallback;
    MsgStreamInfo.pvArg = &EcryptedMessage;

    HCRYPTMSG hMsg = CryptMsgOpenToEncode(
	PKCS_7_ASN_ENCODING | X509_ASN_ENCODING,
	0,
	CMSG_ENVELOPED,
	&envelopedInfo,
	NULL,
	&MsgStreamInfo);

    if (!hMsg) {
	HandleError("Error CryptMsgOpenToEncode.");
    }

    BOOL bFin = FALSE;

#define STREAM_BLOCK_SIZE 0x1000

    PBYTE tbEnc = pbContent;
    DWORD cbEnc = 0;
    DWORD cbCnt = 0;
    while (bFin == FALSE)
    {	
	cbCnt = MIN(STREAM_BLOCK_SIZE, cbContent - cbEnc);

	if (cbCnt < STREAM_BLOCK_SIZE) {
	    bFin = TRUE;
	}	
	/*-------------------------------------------------------------*/
	/* В созданное сообщение поместим данные для шифрования*/
	 BOOL ret = CryptMsgUpdate(
	    hMsg,		/* Handle to the message*/
	    tbEnc,		/* Pointer to the content*/
	    cbCnt,	/* Size of the content*/
	    bFin);		/* Last call*/
	if (!ret) {	    
	    HandleError("Error CryptMsgUpdate.");
	}
	tbEnc += cbCnt;
	cbEnc += cbCnt;
    }

    printf("The encrypted message is %d bytes. \n", EcryptedMessage.cbData);

    printf("Encryption succeeded. \n");

    CryptMsgClose(hMsg);    

    // Вызов функции DecryptMessage, код которой описан после main, для расшифрования сообщения.
    DecryptMessage(EcryptedMessage.pbData, EcryptedMessage.cbData, pRecipientCert);

    free(EcryptedMessage.pbData);

    CleanUp();
    return 0;
}

//  Определение функции DecryptMessage.
// Пример функции для расшифрования зашифрованного сообщения с 
// использованием функции CryptDecryptMessage. ЕЕ параметрами являются
// pbEncryptedBlob, зашифрованное сообщение; cbEncryptedBlob, длина
// этого сообщения
void DecryptMessage(BYTE *pbEncryptedBlob, DWORD cbEncryptedBlob, PCCERT_CONTEXT pRecipientCert)
{
    char * ep = getenv("COLUMNS");
    int brk;
    int i;
    brk = ep ? atoi(ep) : 80;
    brk = ((brk <= 3) ? 80 : brk) / 3;

    printf("The encrypted string is: \n");
    for (i = 0; i < (int)cbEncryptedBlob; i++)
	printf("%02x%c", pbEncryptedBlob[i], (i % brk == (brk - 1)) ? '\n' : ' ');
    printf("\n");

    DWORD dwKeySpec;
    BOOL fCallerFree;
    HCRYPTPROV hProv = 0;
    if (!CryptAcquireCertificatePrivateKey(
	pRecipientCert, 0, 0, &hProv, &dwKeySpec, &fCallerFree))
    {
	HandleError("Error CryptAcquireCertificatePrivateKey.");	
    }

    MEMBUF DecryptedMessage = { 0 };
    // Fill the CMSG_STREAM_INFO structure.
    CMSG_STREAM_INFO MsgStreamInfo;
    memset(&MsgStreamInfo, 0, sizeof(CMSG_STREAM_INFO));
    MsgStreamInfo.cbContent = CMSG_INDEFINITE_LENGTH;
    MsgStreamInfo.pfnStreamOutput = StreamEncryptDecryptCallback;
    MsgStreamInfo.pvArg = &DecryptedMessage;

    HCRYPTMSG hMsg = CryptMsgOpenToDecode(
	PKCS_7_ASN_ENCODING | X509_ASN_ENCODING,
	0, 0, 0, 0,
	&MsgStreamInfo);

    if (!hMsg) {
	HandleError("Error CryptMsgOpenToEncode.");
    }

    BOOL bDecrypt_started = FALSE;
    BOOL bFin = FALSE;
    PBYTE tbEnc = pbEncryptedBlob;
    DWORD cbEnc = 0;
    DWORD cbCnt = 0;
    while (bFin == FALSE)
    {
	cbCnt = MIN(STREAM_BLOCK_SIZE, cbEncryptedBlob - cbEnc);

	if (cbCnt < STREAM_BLOCK_SIZE) {
	    bFin = TRUE;
	}
	/*-------------------------------------------------------------*/
	/* В созданное сообщение поместим данные для шифрования*/
	BOOL ret = CryptMsgUpdate(
	    hMsg,		/* Handle to the message*/
	    tbEnc,		/* Pointer to the content*/
	    cbCnt,	/* Size of the content*/
	    bFin);		/* Last call*/
	if (!ret) {
	    HandleError("Error CryptMsgUpdate.");
	}
	tbEnc += cbCnt;
	cbEnc += cbCnt;

	// При первом вызове нужен вызов CryptMsgControl()
	if (!bDecrypt_started) {
            int i;
	    CMSG_CTRL_DECRYPT_PARA MsgCtrlDecryptPara;
	    memset(&MsgCtrlDecryptPara, 0, sizeof(CMSG_CTRL_DECRYPT_PARA));
	    MsgCtrlDecryptPara.cbSize = sizeof(CMSG_CTRL_DECRYPT_PARA);
	    MsgCtrlDecryptPara.dwKeySpec = dwKeySpec;
	    MsgCtrlDecryptPara.hCryptProv = hProv;
	    MsgCtrlDecryptPara.dwRecipientIndex = (DWORD)-1;

	    /* Определим индекс в списке получателей сообщения, соответствующий */
	    /* сертификату, заданному для расшифрования*/
	    DWORD dwBytes = sizeof(DWORD);
	    DWORD ret = 0;
	    DWORD recip_count = 0;
	    ret = CryptMsgGetParam(hMsg,
		CMSG_RECIPIENT_COUNT_PARAM,
		0,
		&recip_count,
		&dwBytes);
	    if (!ret)
	    {
		HandleError("Error CryptMsgGetParam.");
	    }

	    for (i = 0; i < (int)recip_count; i++) {
		ret = CryptMsgGetParam(hMsg,
		    CMSG_RECIPIENT_INFO_PARAM,
		    i,
		    NULL,
		    &dwBytes);
		if (!ret)
		{
		    HandleError("Error CryptMsgGetParam.");
		}

		CERT_INFO *recip_info = NULL; /* переменная для возврата ссылки из списка получателей*/
		recip_info = (CERT_INFO*)malloc(dwBytes);
		if (!recip_info)
		{
		    HandleError("Error malloc");
		}

		ret = CryptMsgGetParam(hMsg,
		    CMSG_RECIPIENT_INFO_PARAM,
		    i,
		    recip_info,
		    &dwBytes);
		if (!ret)
		{
		    HandleError("Error CryptMsgGetParam.");
		}

		/* Сравненим серийные номера сертификатов и имена издателей*/
		if (CertCompareCertificateName(
		    X509_ASN_ENCODING | PKCS_7_ASN_ENCODING,
		    &(pRecipientCert->pCertInfo->Issuer),
		    &(recip_info->Issuer)) &&
		    CertCompareIntegerBlob(&(pRecipientCert->pCertInfo->SerialNumber),
			&(recip_info->SerialNumber))) {
		    free(recip_info);
		    recip_info = NULL;
		    MsgCtrlDecryptPara.dwRecipientIndex = i;
		    break;
		}
		free(recip_info);
		recip_info = NULL;
	    }
	    if (MsgCtrlDecryptPara.dwRecipientIndex == (DWORD)-1)
	    {
		HandleError("Error dwRecipientIndex. ");
	    }

	    if (!CryptMsgControl(hMsg,
		0,
		CMSG_CTRL_DECRYPT,
		&MsgCtrlDecryptPara))
	    {
		HandleError("Error CryptMsgControl");
	    }
	    bDecrypt_started = TRUE;
	}
    }

    printf("Message Decrypted Successfully. \n");
    printf("The decrypted string is: %s\n", (LPSTR)DecryptedMessage.pbData);
    
    if (fCallerFree)
	CryptReleaseContext(hProv, 0);
    free(DecryptedMessage.pbData);
}

// Проверка типа провайдера 
static BOOL isGostType(DWORD dwProvType) {
    return IS_GOST_PROV(dwProvType);
}

// GetRecipientCert перечисляет сертификаты в хранилище и находит
// первый сертификат, обладающий ключем AT_EXCHANGE. Если сертификат
// сертификат найден, то возвращается указатель на него.  
PCCERT_CONTEXT GetRecipientCert(HCERTSTORE hCertStore)
{
    PCCERT_CONTEXT pCertContext = NULL;
    BOOL bCertNotFind = TRUE;
    DWORD dwSize = 0;
    CRYPT_KEY_PROV_INFO* pKeyInfo = NULL;
    DWORD PropId = CERT_KEY_PROV_INFO_PROP_ID;
    HCRYPTPROV hProv = 0;
    DWORD dwKeySpec = 0;
    BOOL  fFreeProv = FALSE;

    if (!hCertStore) return NULL;

    do
    {
	// Поиск сертификатов в хранилище до тех пор, пока не будет достигнут 
	// конец хранилища, или сертификат с ключем AT_KEYEXCHANGE не будет найден. 
	pCertContext = CertFindCertificateInStore(
	    hCertStore, // Дескриптор хранилища, в котором будет осуществлен поиск. 
	    MY_ENCODING_TYPE,
	    0,
	    CERT_FIND_PROPERTY,
	    &PropId,
	    pCertContext);
	if (!pCertContext)
	    break;

	// Для простоты в этом коде реализован только поиск первого 
	// вхождения ключа AT_KEYEXCHANGE. Во многих случаях, помимо 
	// поиска типа ключа, осуществляется также поиск определенного 
	// имени субъекта. 

	// Однократный вызов функции CertGetCertificateContextProperty  
	// для получения возврашенного размера структуры. 
	if (!(CertGetCertificateContextProperty(
	    pCertContext,
	    CERT_KEY_PROV_INFO_PROP_ID,
	    NULL,
	    &dwSize)))
	{
	    printf("Error getting key property.\n");
	    return NULL;
	}

	//-------------------------------------------------------------- 
	// распределение памяти под возвращенную структуру. 

	free(pKeyInfo);

	pKeyInfo = (CRYPT_KEY_PROV_INFO*)malloc(dwSize);

	if (!pKeyInfo)
	{
	    HandleError("Error allocating memory for pKeyInfo.");
	}

	//-------------------------------------------------------------- 
	// Получение структуры информации о ключе. 

	if (!(CertGetCertificateContextProperty(
	    pCertContext,
	    CERT_KEY_PROV_INFO_PROP_ID,
	    pKeyInfo,
	    &dwSize)))
	{
	    HandleError("The second call to the function failed.");
	}

	//------------------------------------------- 
	// Проверка члена dwKeySpec на расширенный ключ и типа провайдера
	if (pKeyInfo->dwKeySpec == AT_KEYEXCHANGE && isGostType(pKeyInfo->dwProvType))
	{
	    //-------------------------------------------
	    //попробуем открыть провайдер
	    fFreeProv = FALSE;
	    if (CryptAcquireCertificatePrivateKey(pCertContext, CRYPT_ACQUIRE_COMPARE_KEY_FLAG | CRYPT_ACQUIRE_SILENT_FLAG, NULL, &hProv, &dwKeySpec, &fFreeProv))
	    {
		HCRYPTKEY hKey = 0;
		if (CryptGetUserKey(hProv, dwKeySpec, &hKey))
		{
		    bCertNotFind = FALSE;
		    CryptDestroyKey(hKey);
		}
		if (fFreeProv)
		    CryptReleaseContext(hProv, 0);
	    }
	}
    } while (bCertNotFind && pCertContext);

    free(pKeyInfo);

    if (bCertNotFind)
	return NULL;
    else
	return (pCertContext);
} // Конец определения GetRecipientCert 

//----------------------------------------------------------------------------
// Получение имени из CERT_NAME_BLOB
void GetCertDName(PCERT_NAME_BLOB pNameBlob, char **pszName) {
    DWORD	cbName;

    cbName = CertNameToStr(
	X509_ASN_ENCODING, pNameBlob,
	CERT_X500_NAME_STR | CERT_NAME_STR_NO_PLUS_FLAG,
	NULL, 0);
    if (cbName == 1)
	HandleError("CertNameToStr(NULL)");

    *pszName = (char *)malloc(cbName * sizeof(char));
    if (!*pszName)
	HandleError("Out of memory.");

    cbName = CertNameToStrA(
	X509_ASN_ENCODING, pNameBlob,
	CERT_X500_NAME_STR | CERT_NAME_STR_NO_PLUS_FLAG,
	*pszName, cbName);
    if (cbName == 1)
	HandleError("CertNameToStr(pbData)");
}

// Конец примера 
// (не следует удалять данный комментарий, он используется 
//  для автоматической сборки документации)

//  В этом примере используется функция HandleError, функция обработки
//  простых ошибок, для печати сообщения об ошибке в стандартный файл 
//  ошибок (stderr) и выхода из программы. 
//  В большинстве приложений эта функция заменяется другой функцией, 
//  которая выводит более полное сообщение об ошибке.
void HandleError(const char *s)
{
    DWORD err = GetLastError();
    printf("Error number     : 0x%x\n", err);
    printf("Error description: %s\n", s);
    CleanUp();
    if (!err) err = 1;
    exit(err);
}

void CleanUp(void)
{
    CertFreeCertificateContext(pRecipientCert);
    if (hStoreHandle)
    {
	/* !! удалит hCryptProv */
	if (CertCloseStore(hStoreHandle, CERT_CLOSE_STORE_CHECK_FLAG)) {
	    printf("The MY store was closed without incident. \n");
	}
	else {
	    printf("Store closed after encryption -- \n"
		"but not all certificates or CRLs were freed. \n");
	}
    }
    free(szDName);
}
