#include "ckcapi.h"

#include <memory.h>

#include <wincrypt.h>

#ifndef CERT_FIND_KEY_IDENTIFIER
#define CERT_FIND_KEY_IDENTIFIER 983040
#endif

#ifndef CERT_KEY_IDENTIFIER_PROP_ID
#define CERT_KEY_IDENTIFIER_PROP_ID	20
#endif 

#define USE_ENCODINGS	(X509_ASN_ENCODING | PKCS_7_ASN_ENCODING)

/* All the stores we look in for certificates, in this order */
static const char* CERT_STORES[] = {
	"My",
	"AddressBook",
	"CA", 
	"Root", 
	"Trust", 
	"TrustedPeople", 
	"AuthRoot", 
	NULL
};

typedef struct _CertObject
{
	CkCapiObject obj;
	const char* store;

	/* Together these can uniquely identify a certificate */
	CRYPT_INTEGER_BLOB serial;
	CERT_NAME_BLOB issuer;

	/* 
	 * This must stay together. It comprises a unique 
	 * key, together with the data that runs off the end.
	 */
	int otype;
	BYTE cert_data[1];
}
CertObject;

static CK_RV
cert_bool_attribute(void* obj, CK_ATTRIBUTE_TYPE type, 
					CK_VOID_PTR data, CK_ULONG_PTR len)
{
	PCCERT_CONTEXT cert = (PCCERT_CONTEXT)obj;
	CK_BBOOL val;

	ASSERT(obj);

	switch(type) 
	{
	/* 
	 * Resides on the token
	 * - Always true for CAPI objects.
	 */
	case CKA_TOKEN:
		val = CK_TRUE;
		break;

	/*
	 * Private vs. Public object.
	 * - Always false for certificates.
	 */
	case CKA_PRIVATE:
		val = CK_FALSE;
		break;

	/*
	 * If object can be modified.
	 * - Currently always false. In the future with additional 
	 *   functionality this may change.
	 */
	case CKA_MODIFIABLE:
		val = CK_FALSE;
		break;

	/*
	 * Whether the certificate can be trusted for the application
	 * in which it was created.
	 * - Use CertGetCertificateChain to build up a chain 
	 *   for this certificate, and then look in the CERT_CHAIN_CONTEXT 
	 *   TrustStatus field.
	 */
	case CKA_TRUSTED:
		/* TODO: Implement */

	default:
		return CKR_ATTRIBUTE_TYPE_INVALID;
	};

	return ckcapi_return_data(data, len, &val, sizeof(CK_BBOOL));
}

static CK_RV
cert_ulong_attribute(void* obj, CK_ATTRIBUTE_TYPE type, 
					 CK_VOID_PTR data, CK_ULONG_PTR len)
{
	PCCERT_CONTEXT cert = (PCCERT_CONTEXT)obj;
	CK_ULONG val;

	ASSERT(obj);

	switch(type)
	{

	/*
	 * Object class.
	 * - Always CKO_CERTIFICATE for certificates.
	 */
	case CKA_CLASS:
		val = CKO_CERTIFICATE;
		break;

	/*
	 * Type of certificate. 
	 * - Always X509.
	 */
	case CKA_CERTIFICATE_TYPE:
		val = CKC_X_509;
		break;

	/*
	 * Whether a CA, user certificate, other.
	 * - Get certificate szOID_ENHANCED_KEY_USAGE 
	 * extension or CERT_CTL_PROP_ID and look into CTL_USAGE structure.
	 */
	case CKA_CERTIFICATE_CATEGORY:
		/* TODO: Implement */

	default:
		return CKR_ATTRIBUTE_TYPE_INVALID;
	};

	if(*len < sizeof(CK_ULONG))
	{
		*len = sizeof(CK_ULONG);
		return CKR_BUFFER_TOO_SMALL;
	}

	return ckcapi_return_data(data, len, &val, sizeof(CK_ULONG));
}

CK_RV
ckcapi_cert_get_bytes_attribute(void* obj, CK_ATTRIBUTE_TYPE type, 
								CK_VOID_PTR data, CK_ULONG_PTR len)
{
	PCCERT_CONTEXT cert = (PCCERT_CONTEXT)obj;
	DWORD err;
			
	ASSERT(sizeof(CK_ULONG) == sizeof(DWORD));
	ASSERT(obj);

	switch(type)
	{

	/*
	 * Description of the object.
	 * - We use CAPI's CERT_FRIENDLY_NAME_PROP_ID property, 
	 *   converted into UTF8.
	 * - Yes this is slow, but this is not really a property
	 *   that's searched on or retrieved intensively.
	 */
	case CKA_LABEL:
		{
			WCHAR* utf16 = NULL;
			DWORD size;

			/* Get the UTF16 string, a worst case of twice as long as UTF8 */
			if(data)
			{
				size = *len * sizeof(WCHAR);
				utf16 = _alloca(size);
			}

			if(!CertGetCertificateContextProperty(cert, CERT_FRIENDLY_NAME_PROP_ID,
												  utf16, &size))
			{
				err = GetLastError();
				if(err == CRYPT_E_NOT_FOUND)
					utf16 = L"";
				else
					return ckcapi_winerr_to_ckr(err);
			}

			if(utf16)
			{
				/* Always have a default name */
				if(!utf16[0])
				{
					utf16 = L"Unnamed Certificate";
					size = wcslen(utf16) * 2;
				}

				/* Convert the data into the buffer */
				*len = WideCharToMultiByte(CP_UTF8, 0, utf16, size / sizeof(WCHAR), 
										   data, *len, NULL, NULL);
				if(!*len)
					return ckcapi_winerr_to_ckr(GetLastError());
			}
			else
			{
				/* Just return an appropriate allocation length */
				*len = (size / sizeof(WCHAR)) + sizeof(WCHAR);
			}

		}
		return CKR_OK;

	/*
	 * A byte array unique to this certificate. The CKA_ID of 
	 * matching certificates and private keys should match. 
	 * Should match the key identifier in an X.509v3 certificate.
	 * 
	 * We use CAPI's CERT_KEY_IDENTIFIER_PROP_ID property directly.
	 */
	case CKA_ID:
		if(!CertGetCertificateContextProperty(cert, CERT_KEY_IDENTIFIER_PROP_ID, 
												  data, (DWORD*)len))
		{
			err = GetLastError();
			if(err == CRYPT_E_NOT_FOUND)
				return CKR_ATTRIBUTE_TYPE_INVALID;
			return ckcapi_winerr_to_ckr(err);
		}
		return CKR_OK;
		

	/*
	 * DER-encoding of the certificate subject name.
	 * 
	 * We use CAPI's CERT_CONTEXT pCertInfo->Subject field
	 * directly. 
	 */
	case CKA_SUBJECT:
		return ckcapi_return_data(data, len, cert->pCertInfo->Subject.pbData,
								cert->pCertInfo->Subject.cbData);

	/*
	 * DER-encoding of the certificate issuer name.
	 * 
	 * We use CAPI's CERT_CONTEXT pCertInfo->Issuer field
	 * directly.
	 */
	case CKA_ISSUER:
		return ckcapi_return_data(data, len, cert->pCertInfo->Issuer.pbData,
								cert->pCertInfo->Issuer.cbData);

	/*
	 * DER-encoding of the certificate serial number.
	 * 
	 * TODO:
	 *  BOOL rc = CryptEncodeObject(X509_ASN_ENCODING,
     *                    X509_MULTI_BYTE_INTEGER,
     *                   &certContext->pCertInfo->SerialNumber,
	 *		 co->derSerial,
     *                    &size);
	 */
	case CKA_SERIAL_NUMBER:
		if(!CryptEncodeObject(X509_ASN_ENCODING, X509_MULTI_BYTE_INTEGER,
							  &cert->pCertInfo->SerialNumber, data, len))
		{
			err = GetLastError();
			if(err == ERROR_FILE_NOT_FOUND)
				return CKR_GENERAL_ERROR;
			return ckcapi_winerr_to_ckr(err);
		}
		return CKR_OK;

	/*
	 * BER-encoding of the full certificate.
	 *
	 * We use CAPI's CERT_CONTEXT pbCertEncoded field directly.
	 */
	case CKA_VALUE:
		return ckcapi_return_data(data, len, cert->pbCertEncoded,
								cert->cbCertEncoded);

	/*
	 * If CKA_VALUE not specified, this is where the full 
	 * certificate can be found. 
	 * 
	 * We don't support this. All our certificates are present
	 * in full.
	 */
	case CKA_URL:
		break;

	/*
	 * Checksum
	 * - TODO: Work out what to do here
	 */
	case CKA_CHECK_VALUE:
		break;

	/*
	 * TODO: Should we support these?
	 */
	case CKA_HASH_OF_SUBJECT_PUBLIC_KEY:
	case CKA_HASH_OF_ISSUER_PUBLIC_KEY:
		break;

	/* Not supported */

	default:
		break;
	};

	return CKR_ATTRIBUTE_TYPE_INVALID;
}

static CK_RV
cert_date_attribute(void* obj, CK_ATTRIBUTE_TYPE type, 
					CK_VOID_PTR data, CK_ULONG_PTR len)
{
	switch(type)
	{

	/*
	 * Start date for the certificate.
	 * - TODO: Work out where to get this.
	 *	pCertInfo->NotBefore;
	 */
	case CKA_START_DATE:
		break;

	/*
	 * End date for the certificate.
	 * - TODO: Work out where to get this.
	 * pCertInfo->NotBefore;
	 */
	case CKA_END_DATE:
		break;
	};

	return CKR_ATTRIBUTE_TYPE_INVALID;
}

static void
cert_release(void* data)
{
	PCCERT_CONTEXT cert = (PCCERT_CONTEXT)data;
	ASSERT(cert);
	CertFreeCertificateContext(cert);
}

static const CkCapiObjectDataVtable cert_objdata_vtable = {
	cert_bool_attribute,
	cert_ulong_attribute,
	ckcapi_cert_get_bytes_attribute,
	cert_date_attribute,
	cert_release,
};

static CK_RV 
cert_load_data(CkCapiSession* sess, CkCapiObject* obj, CkCapiObjectData* objdata)
{
	CertObject* cobj = (CertObject*)obj;
	HCERTSTORE store;
	CERT_INFO info;
	PCCERT_CONTEXT cert;

	ASSERT(cobj);
	ASSERT(objdata);

	ASSERT(cobj->store);
	store = CertOpenSystemStore((HCRYPTPROV)NULL, cobj->store);
	if(!store)
		return ckcapi_winerr_to_ckr(GetLastError());

	ASSERT(cobj->issuer.pbData);
	ASSERT(cobj->issuer.cbData);
	ASSERT(cobj->serial.pbData);
	ASSERT(cobj->serial.cbData);

	/* Setup our search */
	memset(&info, 0, sizeof(info));
	memcpy(&info.SerialNumber, &cobj->serial, sizeof(info.SerialNumber));
	memcpy(&info.Issuer, &cobj->issuer, sizeof(info.Issuer));

	cert = CertGetSubjectCertificateFromStore(store, USE_ENCODINGS, &info);

	CertCloseStore(store, 0);

	if(!cert)
	{
		DWORD err = GetLastError();

		/* TODO: Is this right for a deleted certificate? */
		ASSERT(err != E_INVALIDARG);
		if(err == CRYPT_E_NOT_FOUND)
			return CKR_OBJECT_HANDLE_INVALID;
		else
			return ckcapi_winerr_to_ckr(GetLastError());
	}
	
	/* 
	 * Important: The trust stuff depends on the objdata value of a 
	 * certificate object is a CERT_CONTEXT pointer.
	 */
	objdata->data = (void*)cert;
	objdata->data_funcs = &cert_objdata_vtable;
	return CKR_OK;
}


static void 
cert_object_release(void* data)
{
	CertObject* cobj = (CertObject*)data;
	ASSERT(cobj);
	free(cobj);
}

static const CkCapiObjectVtable cert_object_vtable = {
	cert_load_data,
	cert_object_release,
};

static CK_RV
register_cert_object(CkCapiSession* sess, const char* store, PCCERT_CONTEXT cert, 
					 CkCapiObject** obj)
{
	CertObject* cobj;
	CK_RV ret;
	size_t len;

	/* We save the Issuer and SerialNumber for identification later */
	len = cert->pCertInfo->SerialNumber.cbData + 
		  cert->pCertInfo->Issuer.cbData;

	/* Add one in case null termination is needed */
	len++;

	cobj = calloc(sizeof(CertObject) + len, 1);
	if(!cobj)
		return CKR_HOST_MEMORY;

	cobj->otype = OBJECT_CERT;
	cobj->store = store;

	cobj->obj.id = 0;
	cobj->obj.unique_key = UNIQUE_KEY_AT(cobj, otype);
	cobj->obj.unique_len = UNIQUE_KEY_VAR_LEN(cobj, otype, cert_data, len);
	cobj->obj.obj_funcs = &cert_object_vtable;

	/* Copy Issuer data in */
	cobj->issuer.cbData = cert->pCertInfo->Issuer.cbData;
	cobj->issuer.pbData = cobj->cert_data;
	memcpy(cobj->issuer.pbData, cert->pCertInfo->Issuer.pbData,
		   cobj->issuer.cbData);

	/* Copy Serial Number data in */
	cobj->serial.cbData = cert->pCertInfo->SerialNumber.cbData;
	cobj->serial.pbData = cobj->cert_data + cobj->issuer.cbData;
	memcpy(cobj->serial.pbData, cert->pCertInfo->SerialNumber.pbData,
		   cobj->serial.cbData);
	
	ret = ckcapi_object_register(sess, &(cobj->obj));
	if(ret != CKR_OK)
	{
		free(cobj);
		return ret;
	}

	ASSERT(cobj->obj.id != 0);
	*obj = &cobj->obj;
	return CKR_OK;
}

static void
clear_object_data_for_store(CkCapiSession* sess, CkCapiObject* obj, 
							CkCapiObjectData* data, void* arg)
{
	const char* store = (const char*) arg;
	CertObject *cobj = (CertObject*)obj;

	// Is it one of ours? 
	if(obj->obj_funcs != &cert_object_vtable)
		return;

	if(strcmp(cobj->store, store) == 0)
		ckcapi_session_clear_object_data(sess, obj);
}

static CK_RV
find_in_store(CkCapiSession* sess, const char* store_name, 
			  CK_ATTRIBUTE_PTR match, CK_ULONG count, CkCapiArray* arr)
{
	PCCERT_CONTEXT cert = NULL;
	CkCapiObject* obj;
	HCERTSTORE store;
	CkCapiObjectData objdata;
	DWORD err;
	CK_RV ret = CKR_OK;

	
	/* Clear any loaded data for objects in this store */
	ckcapi_session_enum_object_data(sess, clear_object_data_for_store, (void*)store_name);

	store = CertOpenSystemStore((HCRYPTPROV)NULL, store_name);
	if(store == NULL)
	{
		err = GetLastError();

		/* Store not found, we don't care */
		if(err == ERROR_FILE_NOT_FOUND)
			return CKR_OK;

		else
			return ckcapi_winerr_to_ckr(err);
	}

	/* Match each certificate */
	while((cert = CertEnumCertificatesInStore(store, cert)) != NULL)
	{
		objdata.data = (void*)cert;
		objdata.data_funcs = &cert_objdata_vtable;

		if(ckcapi_object_data_match(&objdata, match, count))
		{
			ret = register_cert_object(sess, store_name, cert, &obj);
			if(ret != CKR_OK)
				break;

			ASSERT(obj);

			/* Store away the object data for performance reasons */
			objdata.data = (void*)CertDuplicateCertificateContext(cert);
			if(objdata.data) {
				if(ckcapi_session_set_object_data(sess, obj, &objdata) != CKR_OK)
					CertFreeCertificateContext((PCCERT_CONTEXT)objdata.data);
			}

			ckcapi_array_append(arr, obj);
		}
	}

	ASSERT(store);
	CertCloseStore(store, 0);

	return ret;
}

static CK_RV
match_in_store(CkCapiSession* sess, const char* store_name, PCERT_INFO info,
			   CK_ATTRIBUTE_PTR match, CK_ULONG count, CkCapiArray* arr)
{
	PCCERT_CONTEXT cert = NULL;
	CkCapiObject* obj;
	HCERTSTORE store;
	CkCapiObjectData objdata;
	DWORD err;
	CK_RV ret = CKR_OK;

	store = CertOpenSystemStore((HCRYPTPROV)NULL, store_name);
	if(store == NULL)
	{
		err = GetLastError();

		/* Store not found, we don't care */
		if(err == ERROR_FILE_NOT_FOUND)
			return CKR_OK;
		else
			return ckcapi_winerr_to_ckr(err);
	}

	cert = CertGetSubjectCertificateFromStore(store, USE_ENCODINGS, info);
	if(cert == NULL)
	{
		err = GetLastError();

		/* Certificate not found, we don't care */
		if(err == CRYPT_E_NOT_FOUND)
			return CKR_OK;
		else
			return ckcapi_winerr_to_ckr(err);
	}

	/* Match the certificate */
	objdata.data = (void*)cert;
	objdata.data_funcs = &cert_objdata_vtable;

	if(ckcapi_object_data_match(&objdata, match, count))
	{
		ret = register_cert_object(sess, store_name, cert, &obj);
		if(ret == CKR_OK)
		{
			ASSERT(obj);

			/* Store away the object data for performance reasons */
			ret = ckcapi_session_set_object_data(sess, obj, &objdata);
			if(ret == CKR_OK)
				ckcapi_array_append(arr, obj);
		}
	}

	if(ret != CKR_OK && cert)
		CertFreeCertificateContext(cert);


	ASSERT(store);
	CertCloseStore(store, 0);

	return ret;
}

CK_RV
ckcapi_cert_find(CkCapiSession* sess, CK_OBJECT_CLASS cls, CK_ATTRIBUTE_PTR match, 
				 CK_ULONG count, CkCapiArray* arr)
{
	CRYPT_INTEGER_BLOB* serial = NULL; 
	CERT_INFO info;
	CK_RV ret;
	BOOL specific;
	CK_ULONG i;
	DWORD size;

	/* We only have certificates here */
	if(cls != CKO_CERTIFICATE && cls != CKO_ANY)
		return CKR_OK;

	/* 
	 * See if we have a issuer and serial number for a 
	 * specific certificate to find.
	 */
	memset(&info, 0, sizeof(info));

	for(i = 0; i < count; ++i)
	{
		if(!match[i].pValue || !match[i].ulValueLen)
			continue;

		if(match[i].type == CKA_ISSUER)
		{
			info.Issuer.cbData = match[i].ulValueLen;
			info.Issuer.pbData = match[i].pValue;
		}
		else if(!serial && match[i].type == CKA_SERIAL_NUMBER)
		{
			if(!CryptDecodeObject(USE_ENCODINGS, X509_MULTI_BYTE_INTEGER,
								  match[i].pValue, match[i].ulValueLen, 0, NULL, &size))
			{
				continue;
			}

			serial = calloc(1, size);
			if(!CryptDecodeObject(USE_ENCODINGS, X509_MULTI_BYTE_INTEGER,
								  match[i].pValue, match[i].ulValueLen, 0, serial, &size))
				continue;

			ASSERT(serial->cbData);
			ASSERT(serial->pbData);

			info.SerialNumber.cbData = serial->cbData;
			info.SerialNumber.pbData = serial->pbData;
		}
	}

	specific = info.SerialNumber.cbData && info.Issuer.cbData;

	for(i = 0; CERT_STORES[i]; ++i)
	{
		/* Match a specific certificate */
		if(specific)
			ret = match_in_store(sess, CERT_STORES[i], &info, match, count, arr);

		/* Match any ol certificate */
		else 
			ret = find_in_store(sess, CERT_STORES[i], match, count, arr);

		if(ret != CKR_OK)
			break;
	}

	if(serial)
		free(serial);

	return ret;
}