#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 

typedef struct _CertObject
{
	CkCapiObject obj;
	const char* store;
	BYTE* key_id;
	DWORD key_id_len;
}
CertObject;

static CK_RV
copy_static_data(CK_VOID_PTR val, CK_ULONG_PTR len, 
				 CK_VOID_PTR data, DWORD cb)
{
	if(cb == 0)
		return CKR_ATTRIBUTE_TYPE_INVALID;

	if(!val)
	{
		*len = cb;
		return CKR_OK;
	}

	if(cb > *len)
	{
		*len = cb;
		return CKR_BUFFER_TOO_SMALL;
	}

	*len = cb;
	memcpy(val, data, cb);
	return CKR_OK;
}

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 copy_static_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 copy_static_data(data, len, &val, sizeof(CK_ULONG));
}

static CK_RV
cert_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 copy_static_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 copy_static_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 copy_static_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,
	cert_bytes_attribute,
	cert_date_attribute,
	cert_release,
};

static CK_RV 
cert_load(CkCapiObject* obj, CkCapiObjectData* objdata)
{
	CertObject* cobj = (CertObject*)obj;
	HCERTSTORE store;
	PCCERT_CONTEXT cert;
	CRYPT_HASH_BLOB blob;

	ASSERT(cobj);
	ASSERT(objdata);

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

	ASSERT(cobj->key_id);
	ASSERT(cobj->key_id_len);
	blob.pbData = cobj->key_id;
	blob.cbData = cobj->key_id_len;

	cert = CertFindCertificateInStore(store, X509_ASN_ENCODING | PKCS_7_ASN_ENCODING,
									  0, CERT_FIND_KEY_IDENTIFIER, &blob, NULL);

	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());
	}
	
	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,
	cert_object_release,
};

static CK_RV
register_cert_object(CkCapiSession* sess, const char* store, PCCERT_CONTEXT cert, 
					 CK_OBJECT_HANDLE_PTR id)
{
	CertObject* cobj;
	CK_RV ret;
	DWORD len;
	
	if(!CertGetCertificateContextProperty(cert, CERT_KEY_IDENTIFIER_PROP_ID, NULL, &len))
	{
		DBG(("cannot get certificate key identifier: %d", GetLastError()));
		return CKR_ATTRIBUTE_TYPE_INVALID;
	}

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

	/* Store keyid in allocated area after CertObject */
	cobj->key_id = (BYTE*)(cobj + 1);
	cobj->key_id_len = len;

	if(!CertGetCertificateContextProperty(cert, CERT_KEY_IDENTIFIER_PROP_ID, 
										  cobj->key_id, &(cobj->key_id_len)))
	{
		DBG(("cannot read certificate key identifier: %d", GetLastError()));
		free(cobj);
		return CKR_ATTRIBUTE_TYPE_INVALID;
	}

	cobj->store = store;
	cobj->obj.id = 0;
	cobj->obj.unique_key = cobj->key_id;
	cobj->obj.unique_len = cobj->key_id_len;
	cobj->obj.obj_funcs = cert_object_vtable;
	cobj->obj.data_funcs = cert_objdata_vtable;
	
	ret = ckcapi_object_register(sess, &(cobj->obj));
	if(ret != CKR_OK)
	{
		free(cobj);
		return ret;
	}

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

CK_RV
ckcapi_cert_find_in_store(CkCapiSession* sess, const char* store_name, 
						  CK_ATTRIBUTE_PTR match, CK_ULONG count, CkCapiArray* arr)
{
	PCCERT_CONTEXT cert = NULL;
	CK_OBJECT_HANDLE 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);
	}

	/* 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;

			ckcapi_array_append(arr, obj);
		}
	}

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

	return ret;
}

CK_RV
ckcapi_cert_find_all(CkCapiSession* sess, CK_ATTRIBUTE_PTR match, 
					 CK_ULONG count, CkCapiArray* arr)
{
	CK_RV ret;

	ret = ckcapi_cert_find_in_store(sess, "My", match, count, arr);
	if(ret == CKR_OK)
		ret = ckcapi_cert_find_in_store(sess, "AddressBook", match, count, arr);
	if(ret == CKR_OK)
		ret = ckcapi_cert_find_in_store(sess, "CA", match, count, arr);
	if(ret == CKR_OK)
		ret = ckcapi_cert_find_in_store(sess, "Root", match, count, arr);
	if(ret == CKR_OK)
		ret = ckcapi_cert_find_in_store(sess, "Trust", match, count, arr);
	if(ret == CKR_OK)
		ret = ckcapi_cert_find_in_store(sess, "TrustedPeople", match, count, arr);
	if(ret == CKR_OK)
		ret = ckcapi_cert_find_in_store(sess, "AuthRoot", match, count, arr);
	return ret;
}