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


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


/************************************************************************
* Получить имена контейнеров, не содержащих сертификат                  *
************************************************************************/
BOOL EnumContainers(IN LPWSTR lpszCSPs[],
                    IN DWORD pdwCSPTypes[],
                    IN DWORD dwSelectedCSP,
                    OUT LPWSTR lpszContainers[],
                    OUT DWORD* pdwContainersCount)
{
	BOOL bResult = TRUE;                       // Вспомогательная переменная для хранения результата выполнения функции

	HCRYPTPROV hCryptProv = 0;                 // Хэндл криптопровайдера

	LPSTR lpszCurContA = NULL;                 // Буфер для временного хранения имени контейнера в ANSI кодировке
	LPWSTR lpszCurContW = NULL;                // Буфер для временного хранения имени контейнера в Unicode кодировке
	DWORD dwCurContNameLen = 0;                // Размер буфера в байтах для имени контейнера
	DWORD dwIndex = 0;                         // Счетчик цикла
	BOOL bIsContNotNeeded = TRUE;              // Если TRUE, то найденный контейнер содержит сертификат

	DWORD dwError = ERROR_SUCCESS;             // Вспомогательная переменная для хранения кода возврата

	int iCharsNum = 0;

	*pdwContainersCount = 0;

	/************************************************************************
	* Получить хэндл криптопровайдера                                       *
	************************************************************************/
	bResult = CryptAcquireContextW(&hCryptProv,
	                               NULL,
	                               lpszCSPs[dwSelectedCSP],
	                               pdwCSPTypes[dwSelectedCSP],
	                               CRYPT_VERIFYCONTEXT);
	if (!bResult) {
		dwError = GetLastError();
		PrintErrorText(L"CryptAcquireContextW",
		               dwError);
		return bResult;
	}

	for (dwIndex = 0; dwIndex < C2C_MAX_ENUM_COUNT; ++dwIndex) {
		/************************************************************************
		* Получить имена всех контейнеров в криптопровайдере                    *
		* (перебор производится в цикле)                                        *
		************************************************************************/
		dwCurContNameLen = 0;

		bResult = CryptGetProvParam(hCryptProv,
		                            PP_ENUMCONTAINERS,
		                            NULL,
		                            &dwCurContNameLen,
		                            dwIndex == 0 ? CRYPT_FIRST : 0);
		if (!bResult) {
			dwError = GetLastError();
			if (dwError == ERROR_NO_MORE_ITEMS) {
				bResult = TRUE;
				break;
			}

			PrintErrorText(L"CryptGetProvParam",
			               dwError);
			break;
		}

		lpszCurContA = (LPSTR)LocalAlloc(LMEM_ZEROINIT,
		                                 dwCurContNameLen * sizeof(CHAR));

		bResult = CryptGetProvParam(hCryptProv,
		                            PP_ENUMCONTAINERS,
		                            (BYTE*)lpszCurContA,
		                            &dwCurContNameLen,
		                            dwIndex == 0 ? CRYPT_FIRST : 0);
		if (!bResult) {
			dwError = GetLastError();
			LocalFree(lpszCurContA);
			if (dwError == ERROR_NO_MORE_ITEMS) {
				bResult = TRUE;
				break;
			}

			PrintErrorText(L"CryptGetProvParam",
			               dwError);
			break;
		}

		iCharsNum = MultiByteToWideChar(CP_ACP, 0, lpszCurContA, -1, NULL, 0);
		if (iCharsNum == 0) {
			LocalFree(lpszCurContA);
			PrintErrorText(L"MultiByteToWideChar",
			               GetLastError());
			bResult = FALSE;
			break;
		}
		lpszCurContW = (LPWSTR)LocalAlloc(LMEM_ZEROINIT,
		                                  iCharsNum * sizeof(WCHAR));
		iCharsNum = MultiByteToWideChar(CP_UTF8, 0, lpszCurContA, -1, lpszCurContW, iCharsNum);
		dwError = GetLastError();
		LocalFree(lpszCurContA);

		if (iCharsNum == 0) {
			LocalFree(lpszCurContW);
			PrintErrorText(L"MultiByteToWideChar",
			               dwError);
			bResult = FALSE;
			break;
		}

		/************************************************************************
		* Проверить наличие сертификата в контейнере                            *
		************************************************************************/
		bResult = CheckCertificateInContainer(lpszCSPs[dwSelectedCSP],
		                                      pdwCSPTypes[dwSelectedCSP],
		                                      (LPWSTR)lpszCurContW,
		                                      &bIsContNotNeeded);
		if (!bResult) {
			LocalFree(lpszCurContW);
			break;
		}

		if (bIsContNotNeeded) {
			LocalFree(lpszCurContW);
			continue;
		} else {
			/************************************************************************
			* Поместить имя контейнера без сертификата в выходной массив            *
			************************************************************************/
			++*pdwContainersCount;
			lpszContainers[*pdwContainersCount - 1] = lpszCurContW;
			lpszCurContW = 0;
		}
	}

	/************************************************************************
	* Освободить хэндл криптопровайдера                                     *
	************************************************************************/
	CryptReleaseContext(hCryptProv,
	                    0);

	bResult &= *pdwContainersCount > 0;
	if (!bResult) {
		wprintf(L"Error: there are no available containers in CSP,\n\tor no containers without certificate.\n");
	}

	return bResult;
}

/************************************************************************
* Проверить наличие сертификата в контейнере                            *
************************************************************************/
BOOL CheckCertificateInContainer(IN LPCWSTR lpszCSP,
                                 IN DWORD dwCryptProvType,
                                 IN LPCWSTR lpszContainer,
                                 OUT BOOL* pbCertificateExists)
{
	BOOL bResult = TRUE;                       // Вспомогательная переменная для хранения результата выполнения функции
	HCRYPTPROV hCryptProv = 0;                 // Хэндл криптопровайдера
	DWORD dwError = ERROR_SUCCESS;             // Вспомогательная переменная для хранения кода возврата


	/************************************************************************
	* Получить хэндл криптопровайдера                                       *
	************************************************************************/
	bResult = CryptAcquireContextW(&hCryptProv,
	                               lpszContainer,
	                               lpszCSP,
	                               dwCryptProvType,
	                               0);
	if (!bResult) {
		dwError = GetLastError();
		PrintErrorText(L"CryptAcquireContextW",
		               dwError);
		return bResult;
	}

	/************************************************************************
	* Проверить наличие сертификата для пары  AT_KEYEXCHANGE                *
	************************************************************************/
	bResult = CheckKeyCertificate(hCryptProv,
	                              AT_KEYEXCHANGE,
	                              pbCertificateExists);
	if (!bResult
	    ||    *pbCertificateExists) {
		CryptReleaseContext(hCryptProv, 0);
		return bResult;
	}

	/************************************************************************
	* Проверить наличие сертификата для пары  AT_SIGNATURE                  *
	************************************************************************/
	bResult = CheckKeyCertificate(hCryptProv,
	                              AT_SIGNATURE,
	                              pbCertificateExists);

	/************************************************************************
	* Освободить хэндл криптопровайдера                                     *
	************************************************************************/
	CryptReleaseContext(hCryptProv,
	                    0);
	return bResult;
}

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

	DWORD dwError = ERROR_SUCCESS;             // Вспомогательная переменная для хранения кода возврата

	DWORD dwCertBufferLen = 0;                 // Вспомогательная переменная для хранения размера буфера в байтах
	PBYTE pbCert = NULL;                       // Вспомогательная переменная для хранения тела сертификата

	PCCERT_CONTEXT pCertCtx = NULL;            // Вспомогательная переменная для хранения контекста сертификата

	DWORD dwIssuerLen = 0;                     // Вспомогательная переменная для хранения размера имени издателя в байтах
	DWORD dwSubjectLen = 0;                    // Вспомогательная переменная для хранения размера имени субъекта в байтах


	*pbCertificateExists = FALSE;

	/************************************************************************
	* Получить хэндл ключевой пары                                          *
	************************************************************************/
	bResult = CryptGetUserKey(hCryptProv,
	                          dwKeySpec,
	                          &hCryptKey);

	bResult = (bResult || GetLastError() == NTE_NO_KEY);

	if (!bResult) {
		PrintErrorText(L"CryptGetUserKey",
		               GetLastError());
		return bResult;
	}

	if (GetLastError() == NTE_NO_KEY) {
		return bResult;
	}

	/************************************************************************
	* Проверить существование сертификата для ключевой пары                 *
	************************************************************************/
	bResult = CryptGetKeyParam(hCryptKey,
	                           KP_CERTIFICATE,
	                           NULL,
	                           &dwCertBufferLen,
	                           0);

	/************************************************************************
	* Освободить хэндл ключевой пары                                        *
	************************************************************************/
	if (!bResult) {
		dwError = GetLastError();
		CryptDestroyKey(hCryptKey);
		hCryptKey = 0;
		if (dwError == SCARD_E_NO_SUCH_CERTIFICATE) {
			bResult = TRUE;
			return bResult;
		}

		PrintErrorText(L"CryptGetKeyParam",
		               dwError);
		return bResult;
	}

	/************************************************************************
	* Запросить сертификат                                                  *
	************************************************************************/
	pbCert = (BYTE*)LocalAlloc(LMEM_ZEROINIT,
	                           dwCertBufferLen);

	bResult = CryptGetKeyParam(hCryptKey,
	                           KP_CERTIFICATE,
	                           pbCert,
	                           &dwCertBufferLen,
	                           0);

	dwError = GetLastError();
	/************************************************************************
	* Освободить хэндл ключевой пары                                        *
	************************************************************************/
	CryptDestroyKey(hCryptKey);
	hCryptKey = 0;

	if (!bResult) {
		LocalFree(pbCert);
		PrintErrorText(L"CryptGetKeyParam",
		               dwError);
		return bResult;
	}

	/************************************************************************
	* Создать контекст сертификата                                          *
	************************************************************************/
	pCertCtx = CertCreateCertificateContext(C2C_CERT_ENCODING,
	                                        pbCert,
	                                        dwCertBufferLen);

	LocalFree(pbCert);
	if (!pCertCtx) {
		return FALSE;
	}

	/************************************************************************
	* Проверить, что имена издателя и субъекта имеют ненулевую длину        *
	************************************************************************/
	dwIssuerLen = CertNameToStrW(X509_ASN_ENCODING,
	                             &pCertCtx->pCertInfo->Issuer,
	                             CERT_X500_NAME_STR,
	                             NULL,
	                             0);

	dwSubjectLen = CertNameToStrW(X509_ASN_ENCODING,
	                              &pCertCtx->pCertInfo->Subject,
	                              CERT_X500_NAME_STR,
	                              NULL,
	                              0);

	if (dwSubjectLen > 1 && dwIssuerLen > 1) {
		*pbCertificateExists = TRUE;
	}

	/************************************************************************
	* Освободить контекст сертификата                                       *
	************************************************************************/
	CertFreeCertificateContext(pCertCtx);
	pCertCtx = NULL;

	return bResult;
}


/************************************************************************
* Получить порядковый номер элемента в массиве имен контейнеров,        *
* для которого имя контейнера из массива совпадает с именем             *
* контейнера, переданным пользователем                                  *
************************************************************************/
DWORD SelectContainer(IN LPCWSTR lpszPreferredContainer,
                      IN LPWSTR lpszContainers[],
                      IN DWORD dwContainersCount)
{
	DWORD dwSelectedContainer = C2C_BAD_SELECT;   // Возвращаемое значение. Порядковый номер контейнера в массиве
	DWORD dwIndex = 0;                            // Счетчик цикла


	if (lpszPreferredContainer) {
		/************************************************************************
		* Если пользователь передал в командной строке имя контейнера, выполнить*
		* следующие действия:                                                   *
		*  - сравнить имя контейнера, переданное пользователем, с именами       *
		*    контейнеров, присутствующими в массиве имен;                       *
		*  - вернуть индекс элемента массива, совпадающего с тем, который       *
		*    передал пользователь.                                              *
		************************************************************************/
		for (dwIndex = 0; dwIndex < dwContainersCount; ++dwIndex) {
			if (wcscmp(lpszPreferredContainer,
			           lpszContainers[dwIndex])
			    == 0) {
				dwSelectedContainer = dwIndex;
				break;
			}
		}
	} else {
		/************************************************************************
		* Если пользователь не передал в командной строке имя контейнера,       *
		* выполнить следующие действия:                                         *
		*  - распечатать пронумерованный список имен контейнеров;               *
		*  - запросить у пользователя номер контейнера, с которым будет         *
		*    производиться дальнейшая работа.                                   *
		************************************************************************/
		wprintf(L"\nSelect container:\n");
		for (dwIndex = 0; dwIndex < dwContainersCount; ++dwIndex) {
			wprintf(L"[%d] - %s\n", dwIndex, lpszContainers[dwIndex]);
		}

		if (wscanf(L"%d", &dwSelectedContainer) != 1
		    ||  dwSelectedContainer + 1 > dwContainersCount) {
			wprintf(L"Input error.\n");
			dwSelectedContainer = C2C_BAD_SELECT;
			return dwSelectedContainer;
		}
	}

	if (dwSelectedContainer == C2C_BAD_SELECT) {
		wprintf(L"Error: no such container in CSP, or \n\tcertificate is in container already.\n");
	} else {
		wprintf(L"Selected container: \"%s\".\n", lpszContainers[dwSelectedContainer]);
	}

	return dwSelectedContainer;
}
