/*************************************************************************
* Rutoken                                                                *
* Copyright (c) 2003-2017, CJSC Aktiv-Soft. All rights reserved.         *
* Подробная информация:  http://www.rutoken.ru                           *
*------------------------------------------------------------------------*
* В данном файле содержится реализация следующих вспомогательных         *
* функций для работы с Рутокен при помощи программы Cert2Cont:           *
*  - получения хэндла ключевой пары;                                     *
*  - получения структуры данных типа CERT_PUBLIC_KEY_INFO.               *
*************************************************************************/


#include "stdafx.h"
#include "Keys.h"
#include "Cert2Cont.h"
#include "Util_Funcs.h"


/************************************************************************
* Получить хэндл ключевой пары                                          *
************************************************************************/
BOOL GetCryptKey(IN HCRYPTPROV hCryptProv,
                 IN DWORD dwPreferredKeySpec,
                 OUT HCRYPTKEY* hCryptKey,
                 OUT DWORD* dwSelectedKeySpec)
{
	BOOL bResult = TRUE;                        // Вспомогательная переменная для хранения результата выполнения функции

	HCRYPTKEY hCryptKeyExchange = 0;            // Хэндл ключевой пары с типом AT_KEYEXCHANGE
	HCRYPTKEY hCryptKeySignature = 0;           // Хэндл ключевой пары с типом AT_SIGNATURE

	BOOL bNeedToSelect = FALSE;                 // Если TRUE, то Пользователь должен самостоятельно задать тип ключевой пары
	DWORD dwSelectedKey = C2C_BAD_SELECT;       // Возвращаемое значение. Порядковый номер криптопровайдера в массиве


	/************************************************************************
	* Получить временный хэндл ключевой пары                                *
	* с типом AT_KEYEXCHANGE                                                *
	************************************************************************/
	bResult = CryptGetUserKey(hCryptProv,
	                          AT_KEYEXCHANGE,
	                          &hCryptKeyExchange);
	if (!bResult) {
		if (GetLastError() == NTE_NO_KEY) {
			hCryptKeyExchange = 0;
			bResult = TRUE;
		} else {
			PrintErrorText(L"CryptGetUserKey",
			               GetLastError());
			return bResult;
		}
	}

	/************************************************************************
	* Получить временный хэндл ключевой пары                                *
	* с типом AT_SIGNATURE                                                  *
	************************************************************************/
	bResult = CryptGetUserKey(hCryptProv,
	                          AT_SIGNATURE,
	                          &hCryptKeySignature);
	if (!bResult) {
		if (GetLastError() == NTE_NO_KEY) {
			hCryptKeySignature = 0;
			bResult = TRUE;
		} else {
			PrintErrorText(L"CryptGetUserKey",
			               GetLastError());
			return bResult;
		}
	}

	bNeedToSelect = hCryptKeyExchange != 0
	                && hCryptKeySignature != 0
	                && dwPreferredKeySpec == 0;

	if (bNeedToSelect) {
		/************************************************************************
		* Распечатать приглашение для ввода типа ключевой пары                  *
		************************************************************************/
		wprintf(L"Select key:\n");
		wprintf(L"[0] - AT_KEYEXCHANGE\n");
		wprintf(L"[1] - AT_SIGNATURE\n");

		if (wscanf(L"%d", &dwSelectedKey) != 1) {
			wprintf(L"Input error.\n");
			bResult = FALSE;
		} else {
			switch (dwSelectedKey) {
			case 0:
				wprintf(L"You've selected AT_KEYEXCHANGE key.\n");
				break;
			case 1:
				wprintf(L"You've selected AT_SIGNATURE key.\n");
				break;
			default:
				wprintf(L"Input error.\n");
				bResult = FALSE;
				break;
			}
		}
	} else {
		/************************************************************************
		* Определить требуемый тип ключевой пары                                *
		* на основании типа ключевой пары, переданного пользователем            *
		************************************************************************/
		if (dwPreferredKeySpec) {
			dwSelectedKey = dwPreferredKeySpec;
		} else {
			if (hCryptKeyExchange != 0
			    && hCryptKeySignature == 0) {
				dwSelectedKey = AT_KEYEXCHANGE;
			}
			if (hCryptKeyExchange == 0
			    && hCryptKeySignature != 0) {
				dwSelectedKey = AT_SIGNATURE;
			}
		}
	}

	/************************************************************************
	* Освободить временные хэндлы ключевых пар с типами                     *
	* AT_KEYEXCHANGE и AT_SIGNATURE                                         *
	************************************************************************/
	if (hCryptKeyExchange != 0) {
		CryptDestroyKey(hCryptKeyExchange);
	}
	if (hCryptKeySignature != 0) {
		CryptDestroyKey(hCryptKeySignature);
	}

	if (!bResult) {
		return bResult;
	}


	/************************************************************************
	* Получить хэндл ключевой пары с заданным типом                         *
	************************************************************************/
	bResult = CryptGetUserKey(hCryptProv,
	                          dwSelectedKey,
	                          hCryptKey);
	if (!bResult) {
		hCryptKey = 0;
		PrintErrorText(L"CryptGetUserKey",
		               GetLastError());
	}
	*dwSelectedKeySpec = dwSelectedKey;

	return bResult;
}

/************************************************************************
* Получить структуру данных типа CERT_PUBLIC_KEY_INFO;                  *
************************************************************************/
BOOL GetPublicKeyBlobFromCryptProv(IN HCRYPTPROV hCryptProv,
                                   IN DWORD dwKeySpec,
                                   OUT PCERT_PUBLIC_KEY_INFO* pInfo,
                                   OUT DWORD* dwInfoLen)
{
	BOOL bResult = TRUE;           // Вспомогательная переменная для хранения результата выполнения функции
	DWORD dwError = ERROR_SUCCESS; // Вспомогательная переменная для хранения кода возврата


	/************************************************************************
	* Получить размер структуры данных типа CERT_PUBLIC_KEY_INFO;           *
	************************************************************************/
	bResult = CryptExportPublicKeyInfo(hCryptProv,
	                                   dwKeySpec,
	                                   C2C_CERT_ENCODING,
	                                   NULL,
	                                   dwInfoLen);

	dwError = GetLastError();
	bResult = bResult
	          || dwError == ERROR_MORE_DATA;
	if (!bResult) {
		PrintErrorText(L"CryptExportPublicKeyInfo",
		               dwError);
		return bResult;
	}

	*pInfo = (PCERT_PUBLIC_KEY_INFO)LocalAlloc(LMEM_ZEROINIT, *dwInfoLen);

	/************************************************************************
	* Получить структуру данных типа CERT_PUBLIC_KEY_INFO;                  *
	************************************************************************/
	bResult = CryptExportPublicKeyInfo(hCryptProv,
	                                   dwKeySpec,
	                                   C2C_CERT_ENCODING,
	                                   *pInfo,
	                                   dwInfoLen);
	if (!bResult) {
		dwError = GetLastError();
		LocalFree(*pInfo);
		*pInfo = NULL;
		PrintErrorText(L"CryptExportPublicKeyInfo",
		               dwError);
		return bResult;
	}

	return bResult;
}

