// key_import_tester.cpp : Defines the entry point for the console application.
//

// Usage: 
//
//DWORD ImportNullKey(
//        HCRYPTPROV			hCryptProv,
//        HCRYPTKEY*			phKey
//        )
//{
//    BYTE nullKey[16]={0};
//    KeyImporter<16> importKey;
//    return importKey.Import(hCryptProv,CALG_AES_128,nullKey,sizeof(nullKey),phKey);
//} 

#include <stdio.h>
#include <tchar.h>
#include <windows.h>
#include <wincrypt.h>
#include <assert.h>


#pragma pack (push, 1)
template <int KeySize>
struct PlainTextKeyBlob
{
	BLOBHEADER _hdr;
	DWORD      _cbKey;
	BYTE _key[KeySize];
};
#pragma pack (pop)

template <int  KeySize> 
class KeyImporter:PlainTextKeyBlob<KeySize>
{
	public: 
		DWORD Import( const HCRYPTPROV hProv, const ALG_ID alg, VOID const * const  pKey, const DWORD cbKey, HCRYPTKEY* const phKey)
		{
			DWORD ret = Init(pKey,cbKey,alg);
			if (ret != NO_ERROR)
			{
				return ret;
			}

			if (!CryptImportKey(hProv, (PBYTE)this, cbToImport(), NULL, 0, phKey))
			{
				return GetLastError();
			}

			return NO_ERROR;
		}

	private:

		DWORD Init(VOID const * const pKey, const DWORD cbKey, const ALG_ID alg)
		{
			if (cbKey != KeySize)
			{
				return ERROR_INVALID_PARAMETER;
			}  

			_hdr.bType = PLAINTEXTKEYBLOB;
			_hdr.bVersion =  CUR_BLOB_VERSION;
			_hdr.reserved = 0;
			_hdr.aiKeyAlg = alg;
			_cbKey = KeySize;
			memcpy (_key,pKey,sizeof(_key));
			return NO_ERROR;
		}

		DWORD cbToImport() const 
		{
			return sizeof(_hdr) + sizeof(_cbKey) +  sizeof(_key);
		}
};

C_ASSERT(sizeof(PlainTextKeyBlob<17>)==sizeof(KeyImporter<17>));

// asserts on failure. 

void ut_KeyImporter()
{
	// Acquire Contexts
	HCRYPTKEY hCryptProv1,hCryptProv2;
	HCRYPTPROV hGeneratedKey,hImportedKey;
	BOOL bRet =
		CryptAcquireContext(&hCryptProv1,NULL,NULL,PROV_RSA_AES,CRYPT_NEWKEYSET
		|CRYPT_VERIFYCONTEXT);
	assert(bRet);

	bRet =
		CryptAcquireContext(&hCryptProv2,NULL,NULL,PROV_RSA_AES,CRYPT_NEWKEYSET
		|CRYPT_VERIFYCONTEXT);
	assert(bRet);


	// Generate Key
	bRet = CryptGenKey(hCryptProv1,CALG_AES_128,CRYPT_EXPORTABLE,&hGeneratedKey);
	assert(bRet);

	// Export Key 
	PlainTextKeyBlob<16> exportedKey;
	DWORD cbExportedKey = sizeof(exportedKey);
	bRet = CryptExportKey(hGeneratedKey,NULL,PLAINTEXTKEYBLOB,0,reinterpret_cast<PBYTE>(&exportedKey),&cbExportedKey);
	assert(bRet && cbExportedKey == sizeof(exportedKey));

	// Import that key into a new key container. 
	KeyImporter<16> keyImporter;
	DWORD ret = keyImporter.Import(hCryptProv2,exportedKey._hdr.aiKeyAlg,exportedKey._key,sizeof(exportedKey._key),&hImportedKey);
	assert(ret==NO_ERROR);

	// Create Plain Text
	BYTE plainText[16];
	const DWORD cbPlainText=sizeof(plainText);
	memset(plainText,'p',cbPlainText);

	BYTE scratch[MAX_PATH];
	C_ASSERT(cbPlainText < sizeof(scratch));
	memcpy(scratch,plainText,cbPlainText);

	// Encypt plain text with generated key 
	DWORD cbEncrypted=cbPlainText;
	bRet = CryptEncrypt(hGeneratedKey,NULL,TRUE,0,scratch,&cbEncrypted,sizeof(scratch));
	assert(bRet);

	// Decrypt ciphertext with imported key 
	DWORD cbDecrypted = cbEncrypted;
	bRet = CryptDecrypt(hImportedKey,NULL,TRUE,0,scratch,&cbDecrypted);
	assert(bRet);

	// ensure a match
	assert (cbDecrypted == cbPlainText);
	assert (memcmp(scratch,plainText,cbPlainText) == 0);

	// cleanup 
	CryptDestroyKey(hGeneratedKey);
	CryptDestroyKey(hImportedKey);
	CryptReleaseContext(hCryptProv1,0);
	CryptReleaseContext(hCryptProv2,0);
}

int _tmain(int argc, _TCHAR* argv[])
{
	
	ut_KeyImporter();
	return 0;
}