//-------------------------------------------------------------------------------------------------
// <copyright file="CryptoProcessor.cpp" company="LiquidRobotics">
//  Copyright (c) Liquid Robotics Corporation, confidential and proprietary.  All rights reserved. 
// </copyright>
//
// <summary>
// 	Handles the encryption and decryption of messages using the Mozilla
//  NSS Library.
// </summary>
//
// <owner>Arash Ushani</owner>
//-------------------------------------------------------------------------------------------------

#include "CryptoProcessor.h"
#include <string.h>
#include <nss/nss.h>
#include <nss/pk11pub.h>
#include <prerror.h>
#include <iostream>
#include <stdexcept>

using namespace std;

namespace LiquidRobotics {
namespace CryptoLib {

bool CryptoProcessor::nss_initialized = false;

CryptoProcessor::CryptoProcessor()
{
	init_nss();

	/* choose mechanism: CKM_DES_CBC_PAD, CKM_DES3_ECB, CKM_DES3_CBC.....
	* Note that some mechanisms (*_PAD) imply the padding is handled for you
	* by NSS. If you choose something else, then data padding is the
	* application's responsibility
	*/
	cipherMech = CKM_AES_CBC_PAD;

	/* slot = PK11_GetInternalKeySlot(); is a simpler alternative but in
	* theory, it *may not* return the optimal slot for the operation. For
	* DES ops, Internal slot is typically the best slot
	*/
	slot = PK11_GetBestSlot(cipherMech, NULL);

	SymKey = NULL;
}

CryptoProcessor::~CryptoProcessor()
{
}

void CryptoProcessor::init_nss()
{
	if (!nss_initialized)
	{
		nss_initialized = true;

		SECStatus rv = NSS_NoDB_Init(".");
		if (rv != SECSuccess)
		{
			cout << "NSS initialization failed (err " << PR_GetError() << ")" << endl;
		}

		// Init the random number generator
		srand(time(NULL));
	}
}

void CryptoProcessor::SetKey(unsigned char *newKey, int len)
{
#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Setting key of length=%i\n", len);
#endif
	/* NSS passes blobs around as SECItems. These contain a pointer to
	* data and a length. Turn the raw key into a SECItem. */
	SECItem keyItem;

	keyItem.type = siBuffer;
	keyItem.data = newKey;
	keyItem.len = len;

#ifdef CRYTO_DEBUG
	fprintf(stderr, "AES Key:\n");
	for (int i=0; i<len; i++)
	{
		fprintf(stderr, "0x%02x ", keyItem.data[i]);
	}
	fprintf(stderr, "\n");
#endif
	/* Turn the raw key into a key object. We use PK11_OriginUnwrap
	* to indicate the key was unwrapped - which is what should be done
	* normally anyway - using raw keys isn't a good idea */

	//if (SymKey)
	//{
	//	cerr << "Freeing key... " << endl;
	//	PK11_FreeSymKey(SymKey);
	//}

	SymKey = PK11_ImportSymKey(slot, cipherMech, PK11_OriginUnwrap, CKA_ENCRYPT, &keyItem, NULL);
	if (SymKey == NULL)
	{
		cerr << "Failure to import key into NSS (err " << PR_GetError() << ")" << endl;
	}
}

//unsigned char* CryptoProcessor::Encrypt(unsigned char *data, int dataLen, unsigned char *iv, int ivLen, int *result_len)
bool CryptoProcessor::Encrypt(unsigned char *data, int dataLen, unsigned char *iv, int ivLen, unsigned char *result, int *result_len)
{
	if (!SymKey)
	{
		fprintf(stderr, "Encrypt Error: Have no key to encrypt with.\n");
		return false;
	}

#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Starting encrypt with dataLen=%i and ivLen=%i\n", dataLen, ivLen);
#endif

	// TEST TO SEE IF THIS FIXES THE PROBLEM
	unsigned char clearTextData[dataLen];
	memcpy(clearTextData, data, dataLen);

	unsigned char theIV[ivLen];
	memcpy(theIV, iv, ivLen);

#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Clear text data:\n");
	for (int i=0; i<dataLen; i++)
	{
		fprintf(stderr, "0x%02x ", clearTextData[i]);
	}
	fprintf(stderr, "\n");

	fprintf(stderr, "IV:\n");
	for (int i=0; i<ivLen; i++)
	{
		fprintf(stderr, "0x%02x ", theIV[i]);
	}
	fprintf(stderr, "\n");
#endif

	/* set up the PKCS11 encryption paramters.
	* when not using CBC mode, ivItem.data and ivItem.len can be 0, or you
	* can simply pass NULL for the iv parameter in PK11_ParamFromIV func
	*/
	SECItem ivItem;
#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Encrypt: Setting up iv...\n");
#endif
	ivItem.type = siBuffer;
	ivItem.data = theIV;
	ivItem.len = ivLen;
	SECItem* SecParam = PK11_ParamFromIV(cipherMech, &ivItem);
	if (SecParam == NULL)
	{
		fprintf(stderr, "Failure to set up PKCS11 param (err %d)", PR_GetError());
		return false;
	}

#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Encrypt: Have iv...\n");
#endif

	//fprintf(stderr, "IV:");
	//for (unsigned int i=0; i<ivItem.len; i++)
	//	fprintf(stderr, "%02x ", ivItem.data[i]);
	//fprintf(stderr, "\n");

	/* ENCRYPT data into buf1.*/
	unsigned char buf1[dataLen+16];
	memset(buf1, 0, dataLen+16);

	int tmp1_outlen = 0;
	int tmp2_outlen = 0;

#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Encrypting...\n");
#endif

	/* Create cipher context */
	PK11Context* EncContext = PK11_CreateContextBySymKey(cipherMech, CKA_ENCRYPT, SymKey, SecParam);
	SECStatus rv1 = PK11_CipherOp(EncContext, buf1, &tmp1_outlen, dataLen+16, clearTextData, dataLen);
	SECStatus rv2 = PK11_DigestFinal(EncContext, buf1+tmp1_outlen, (unsigned int*) &tmp2_outlen, (unsigned int) (sizeof(buf1)-tmp1_outlen));
	PK11_DestroyContext(EncContext, PR_TRUE);
	*result_len = tmp1_outlen + tmp2_outlen;
	if (rv1 != SECSuccess || rv2 != SECSuccess)
	{
		fprintf(stderr, "Failed encrypt, rv1=%i, rv2=%i\n", rv1, rv2);
		return false;
	}

	//fprintf(stderr, "Encrypted Data: ");
	//for (int i=0; i<*result_len; i++)
	//	fprintf(stderr, "%02x ", buf1[i]);
	//fprintf(stderr, "\n");

	//fprintf(stderr, "---------------------\n");

	//if (SecParam)
	//	SECITEM_FreeItem(SecParam, PR_TRUE);
	//if (slot)
	//	PK11_FreeSlot(slot);

	memcpy(result, buf1, *result_len);

#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Encrypt: Freeing params... not!\n");

	fprintf(stderr, "Encrypted data:\n");
	for (int i=0; i<*result_len; i++)
	{
		fprintf(stderr, "0x%02x ", result[i]);
	}
	fprintf(stderr, "\n");
#endif

	return true;
}

//unsigned char* CryptoProcessor::Decrypt(unsigned char *data, int dataLen, unsigned char *iv, int ivLen, int *result_len)
bool CryptoProcessor::Decrypt(unsigned char *data, int dataLen, unsigned char *iv, int ivLen, unsigned char *result, int *result_len)
{
	if (!SymKey)
	{
		fprintf(stderr, "Decrypt Error: Have no key to decrypt with.\n");
		return false;
	}

#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Starting decrypt with dataLen=%i and ivLen=%i\n", dataLen, ivLen);
#endif
	// TEST TO SEE IF THIS FIXES THE PROBLEM
	unsigned char encryptedData[dataLen];
	memcpy(encryptedData, data, dataLen);

	unsigned char theIV[ivLen];
	memcpy(theIV, iv, ivLen);

	/* DECRYPT into buf2.*/
	unsigned char buf2[dataLen+8];
	memset(buf2, 0, dataLen+8);
	int tmp1_outlen = 0;
	int tmp2_outlen = 0;

	/* set up the PKCS11 encryption paramaters.
	* when not using CBC mode, ivItem.data and ivItem.len can be 0, or you
	* can simply pass NULL for the iv parameter in PK11_ParamFromIV func
	*/
	SECItem ivItem2;

	ivItem2.type = siBuffer;
	ivItem2.data = theIV;
	ivItem2.len = ivLen;

	SECItem* SecParam2 = PK11_ParamFromIV(cipherMech, &ivItem2);
	if (SecParam2 == NULL)
	{
		fprintf(stderr, "Failure to set up PKCS11 param (err %d)", PR_GetError());
		return false;
	}

#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Decrypted data:\n");
	for (int i=0; i<dataLen; i++)
	{
		fprintf(stderr, "0x%02x ", encryptedData[i]);
	}
	fprintf(stderr, "\n");

	fprintf(stderr, "IV:\n");
	for (int i=0; i<ivLen; i++)
	{
		fprintf(stderr, "0x%02x ", theIV[i]);
	}
	fprintf(stderr, "\n");
#endif

	/* Create cipher context */
	PK11Context* EncContext2 = PK11_CreateContextBySymKey(cipherMech, CKA_DECRYPT, SymKey, SecParam2);
	SECStatus rv12 = PK11_CipherOp(EncContext2, buf2, &tmp1_outlen, dataLen+8, encryptedData, dataLen);
	SECStatus rv22 = PK11_DigestFinal(EncContext2, buf2+tmp1_outlen, (unsigned int*) &tmp2_outlen, (unsigned int) (dataLen-tmp1_outlen));
	PK11_DestroyContext(EncContext2, PR_TRUE);
	*result_len = tmp1_outlen + tmp2_outlen;
	if (rv12 != SECSuccess || rv22 != SECSuccess)
	{
		fprintf(stderr, "Failed decrypt, rv1=%i, rv2=%i\n", rv12, rv22);
		return false;
	}

	memcpy(result, buf2, *result_len);

#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Decrypt: freeing params... not!\n");
#endif

	//if (SecParam2)
	//	SECITEM_FreeItem(SecParam2, PR_TRUE);
	//if (slot)
	//	PK11_FreeSlot(slot);

	return true;
}

void CryptoProcessor::GenerateRandomIV(int ivLen, unsigned char *iv)
{

#ifdef CRYPTO_DEBUG
	fprintf(stderr, "Generating random IV of len=%i\n", ivLen);
#endif

	for (int i=0; i<ivLen; i++)
		iv[i] = rand()%256;
}

void CryptoProcessor::test()
{
	fprintf(stderr, "\nTesting Crypto...\n");

	unsigned char gKey[] = {0xe8, 0xa7, 0x7c, 0xe2, 0x05, 0x63, 0x6a, 0x31, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18};
	unsigned char gIV[] = {0xe4, 0xbb, 0x3b, 0xd3, 0xc3, 0x71, 0x2e, 0x58, 0xe4, 0xbb, 0x3b, 0xd3, 0xc3, 0x71, 0x2e, 0x58};

	std::string message("Here's a sample sentence: The quick brown fox jumped over the lazy fence. Let's see if this works.");

	unsigned char data[message.length()+1], encryptedData[message.length()+16], decryptedData[message.length()+16];

	strcpy((char*) data, message.c_str());

	unsigned int dataLen = strlen((const char*) data)+1;

	CryptoProcessor test;
	test.SetKey(gKey, sizeof(gKey));

	fprintf(stderr, (const char*) "\nOriginal (%i bytes):\n", dataLen);
	unsigned int i = 0;
	while (i<dataLen)
		fprintf(stderr, "%02x ", *(data+i++));
	fprintf(stderr, "\n");

	int result_len=0;

	test.Encrypt(data, dataLen, gIV, sizeof(gIV), encryptedData, &result_len);
	fprintf(stderr, "\nEncrypted (%i bytes):\n", result_len);
	for (int i=0; i<result_len; i++)
		fprintf(stderr, "%02x ", encryptedData[i]);
	fprintf(stderr, "\n");

	/*
	unsigned char *encrypted = test.Encrypt(data, strlen(data)+1, gIV, sizeof(gIV), &result_len);

	fprintf(stderr, "\nEncrypted (%i bytes):\n", result_len);
	i = 0;
	while (i<result_len)
		fprintf(stderr, "%02x ", *(encrypted+i++));
	fprintf(stderr, "\n");

	memcpy(encryptedData, encrypted, result_len);
	*/

	int decrypted_len=0;
	test.Decrypt(encryptedData, result_len, gIV, sizeof(gIV), decryptedData, &decrypted_len);
	fprintf(stderr, "\nDecrypted (%i bytes):\n", decrypted_len);
	for (int i=0; i<decrypted_len; i++)
		fprintf(stderr, "%02x ", decryptedData[i]);
	fprintf(stderr, "\n");

	/*
	unsigned char *decrypted = test.Decrypt(encrypted, result_len, gIV, sizeof(gIV), &decrypted_len);

	fprintf(stderr, "\nDecrypted (%i bytes):\n", decrypted_len);
	i = 0;
	while (i<decrypted_len)
		fprintf(stderr, "%02x ", *(encrypted+i++));
	fprintf(stderr, "\n");

	memcpy(decryptedData, decrypted, result_len);
	*/

	std::string orig((const char*) data);
	std::string final((const char*) decryptedData);

	fprintf(stderr, "\n\nOrig= <%s>", data);
	fprintf(stderr, "\nRecv= <%s>", decryptedData);
	if (orig.compare(final)==0)
		fprintf(stderr, "\nMatched!\n");
	else
		fprintf(stderr, "\nDid not match...\n");
}

}
}
