//----------------------------------------------------------------------------
//
//  Copyright (C) 2007 Intel Corporation
//
//  File:       CertHash.cpp
//
//  Contents:   Certificate hash utility
//
//----------------------------------------------------------------------------

#include <windows.h>
#include <wincrypt.h>
#include <atlbase.h>
#include <fstream>
#include <stdio.h>
#include <string>
using namespace std;

#include "CertHash.h"

//Declaration of the wanted functions in libcrypto.dll
typedef unsigned char* (*SHA1funcPtr)(const unsigned char* data, size_t dataLength, unsigned char* hash);  
typedef unsigned char* (*SHA256funcPtr)(const unsigned char* data, size_t dataLength, unsigned char* hash);  
typedef unsigned char* (*SHA384funcPtr)(const unsigned char* data, size_t dataLength, unsigned char* hash);  

CertHash::CertHash(const CertHash& h) : _valid(false)
{
	_valid = h._valid;
	_dllFailed = h._dllFailed;
	_friendlyName = h._friendlyName;
	_hashAlgorithm = h._hashAlgorithm;
	_hashLength = h._hashLength;
		
	_hash = nullptr;
	if(h._hash)
	{
		_hash = new unsigned char[MAX_HASH_LENGTH];
		if (memcpy_s(_hash, MAX_HASH_LENGTH, h._hash, h._hashLength))
		{
			printf("Error: Can't Copy Memory");
		}
	}
}

CertHash::CertHash(const string& fileName, const string& name, const CertHashAlgorithm& hashAlgorithm) : _valid(false), _dllFailed(false)
{
	_hashLength = 0;

	_friendlyName = name;
	this->_hashAlgorithm = hashAlgorithm;
	_hash = new unsigned char[MAX_HASH_LENGTH];
	memset(_hash, 0 , MAX_HASH_LENGTH);

	HCERTSTORE store = NULL;
	PCCERT_CONTEXT pCertContext = NULL;
	CComBSTR str(fileName.c_str());

	do {		
		store = CertOpenStore(
			CERT_STORE_PROV_FILENAME,   // the store provider type
			X509_ASN_ENCODING,          // if needed, use the usual
			NULL,                       // use the default HCRYPTPROV
			0,                          // accept the default for all
			str);						// the name of an existing file

		if (!store)
			break;

		pCertContext = CertEnumCertificatesInStore(store, NULL);
		if (!pCertContext)
			break;

		// for sha2 algorithms (sha256 and sha384) use openssl
		if ((USB_CERT_HASH_ALGORITHM_SHA256 == hashAlgorithm) || (USB_CERT_HASH_ALGORITHM_SHA384 == hashAlgorithm))
		{
			HINSTANCE loadLib;
			SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_DEFAULT_DIRS);
			loadLib = LoadLibrary(L"libcrypto.dll");
			// Check to see if the library was loaded successfully 
			if (0 != loadLib)
			{
				if (USB_CERT_HASH_ALGORITHM_SHA256 == hashAlgorithm)
				{
					_hashLength = CERT_HASH_SHA256_LENGTH;

					SHA256funcPtr sha256funcPtr = (SHA256funcPtr)GetProcAddress(loadLib, "SHA256");

					sha256funcPtr(pCertContext->pbCertEncoded, pCertContext->cbCertEncoded, _hash);
				}
				else if (USB_CERT_HASH_ALGORITHM_SHA384 == hashAlgorithm)
				{
					_hashLength = CERT_HASH_SHA384_LENGTH;

					SHA384funcPtr sha384funcPtr = (SHA384funcPtr)GetProcAddress(loadLib, "SHA384");

					sha384funcPtr(pCertContext->pbCertEncoded, pCertContext->cbCertEncoded, _hash);
				}

				_valid = true;
				FreeLibrary(loadLib);
			}
			else
			{
				_dllFailed = true;
			}
		}
		// for sha1 algorithm windows
		else // sha1 or unsupported hash - use the default (sha1)
		{
			_hashLength = CERT_HASH_SHA1_LENGTH;

			DWORD cbData = CERT_HASH_SHA1_LENGTH;
			if (!CertGetCertificateContextProperty(pCertContext, CERT_SHA1_HASH_PROP_ID, _hash, &cbData))
				break;

			// do nothing on error...caller must check valid
			_valid = true;
		}

	} while(0);

	if (pCertContext)
		CertFreeCertificateContext(pCertContext);
	if (store)
		CertCloseStore(store, 0);
}

CertHash& CertHash::operator=(const CertHash& h)
{
	_valid = h._valid;
	_dllFailed = h._dllFailed;
	_friendlyName = h._friendlyName;
	_hashAlgorithm = h._hashAlgorithm;
	_hashLength = h._hashLength;

	_hash = nullptr;
	if (h._hash)
	{
		_hash = new unsigned char[MAX_HASH_LENGTH];
		if (memcpy_s(_hash, MAX_HASH_LENGTH, h._hash, h._hashLength))
		{
			printf("Error: Can't Copy Memory");
		}
	}
	return *this;
}

CertHash::CertHash(const string& hashstring, const string& friendlyName) :_valid(false), _dllFailed(false)
{
	
		_hash = new unsigned char[MAX_HASH_LENGTH];
		_friendlyName = friendlyName;

		_hashLength = hashstring.length()/2;		

		// check the length of the string
		if(_hashLength == CERT_HASH_SHA1_LENGTH)	
			_hashAlgorithm = USB_CERT_HASH_ALGORITHM_SHA1;
		
		else if(_hashLength == CERT_HASH_SHA256_LENGTH)	
			_hashAlgorithm = USB_CERT_HASH_ALGORITHM_SHA256;
		
		else if(_hashLength == CERT_HASH_SHA384_LENGTH)
			_hashAlgorithm = USB_CERT_HASH_ALGORITHM_SHA384;
		
		else
		{
			_hashAlgorithm = INVALID_USB_CERT_HASH_ALGORITHM;
			return;
		}
			

		//check if all characters are hex 
		for (int i = 0 ; i < hashstring.length(); i++)
		{
			char test = hashstring[i];
			if (!isxdigit(test))
				return;
		}

		for (int i = 0, j=0; i < hashstring.length(); i=i+2,j++)
		{
			_hash[j] = strtoul((hashstring.substr(i,2)).c_str(), 0, 16);
		}
		
		_valid = true;


}

CertHash::~CertHash()
{
	if(_hash)
	{
		delete[] _hash;
		_hash = NULL;
	}
}