/* 
 * Copyright (C) 2007 Stef Walter
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.	 See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 59 Temple Place - Suite 330,
 * Boston, MA 02111-1307, USA.
 */

#include "ckcapi.h"
#include "ckcapi-cert.h"
#include "ckcapi-object.h"
#include "ckcapi-session.h"
#include "ckcapi-token.h"

#include <memory.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;

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

typedef struct _CertObjectData
{
	CkCapiObjectData base;
	PCCERT_CONTEXT cert;
}
CertObjectData;

static CK_RV
cert_bool_attribute(CkCapiObjectData* objdata, CK_ATTRIBUTE_PTR attr)
{
	CertObjectData* cdata = (CertObjectData*)objdata;
	CK_BBOOL val;

	ASSERT(cdata);

	switch(attr->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(attr->pValue, &attr->ulValueLen, &val, sizeof(CK_BBOOL));
}

static CK_RV
cert_ulong_attribute(CkCapiObjectData* objdata, CK_ATTRIBUTE_PTR attr)
{
	CertObjectData* cdata = (CertObjectData*)objdata;
	CK_ULONG val;

	ASSERT(objdata);

	switch(attr->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;
	};

	return ckcapi_return_data(attr->pValue, &attr->ulValueLen, &val, sizeof(CK_ULONG));
}

static CK_RV
cert_bytes_attribute(CkCapiObjectData* objdata, CK_ATTRIBUTE_PTR attr)
{
	CertObjectData* cdata = (CertObjectData*)objdata;
	PCCERT_CONTEXT cert = cdata->cert;
			
	ASSERT(sizeof(CK_ULONG) == sizeof(DWORD));
	ASSERT(cdata);

	return ckcapi_cert_certificate_get_bytes(cdata->cert, attr);
}

static CK_RV
cert_date_attribute(CkCapiObjectData* objdata, CK_ATTRIBUTE_PTR attr)
{
	switch(attr->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_data_release(void* data)
{
	CertObjectData* cdata = (CertObjectData*)data;
	ASSERT(cdata && cdata->cert);
	CertFreeCertificateContext(cdata->cert);
	free(cdata);
}

static const CkCapiObjectDataVtable cert_objdata_vtable = {
	cert_bool_attribute,
	cert_ulong_attribute,
	cert_bytes_attribute,
	cert_date_attribute,
	cert_data_release,
};

static CkCapiObjectData*
cert_alloc_data(CkCapiSession* sess, CkCapiObject* obj, PCCERT_CONTEXT cert)
{
	CertObjectData* cdata;

	cdata = (CertObjectData*)calloc(sizeof(CertObjectData), 1);
	if(!cdata)
		return NULL;
	
	cdata->cert = cert;

	cdata->base.object = obj->id;
	cdata->base.data_funcs = &cert_objdata_vtable;

	return &(cdata->base);
}

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

	ASSERT(cobj);
	ASSERT(objdata);

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

	/* No store should mean no objects were loaded */
	ASSERT(sess->store);

	/* 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(sess->store, CKCAPI_ENCODINGS, &info);

	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 = cert_alloc_data(sess, obj, cert);
	if(!(*objdata))
	{
		CertFreeCertificateContext(cert);
		return CKR_HOST_MEMORY;
	}
	
	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,
};

CK_RV 
ckcapi_cert_certificate_get_bytes(PCCERT_CONTEXT cert, CK_ATTRIBUTE_PTR attr)
{
	CK_VOID_PTR data = attr->pValue;
	CK_ULONG_PTR len = &attr->ulValueLen;
	DWORD err;

	ASSERT(cert);
	ASSERT(len);
	ASSERT(attr);

	switch(attr->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;
}

PCCERT_CONTEXT	
ckcapi_cert_object_data_get_certificate(CkCapiObjectData* objdata)
{
	CertObjectData* cdata;

	ASSERT(objdata);
	ASSERT(objdata->data_funcs == &cert_objdata_vtable);

	cdata = (CertObjectData*)objdata;
	return cdata->cert;
}

static CK_RV
register_cert_object(CkCapiSession* sess, 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->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_token_register_object(sess->slot, &(cobj->obj));
	if(ret != CKR_OK)
	{
		free(cobj);
		return ret;
	}

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

static CK_RV
find_in_store(CkCapiSession* sess, CK_ATTRIBUTE_PTR match, 
			  CK_ULONG count, CkCapiArray* arr)
{
	PCCERT_CONTEXT cert = NULL;
	CkCapiObject* obj;
	CertObjectData cdata;
	CkCapiObjectData* objdata;
	CK_RV ret = CKR_OK;

	/* No store, no objects */
	if(!sess->store)
		return CKR_OK;

	/* Match each certificate */
	while((cert = CertEnumCertificatesInStore(sess->store, cert)) != NULL)
	{
		cdata.cert = cert;
		cdata.base.object = 0;
		cdata.base.data_funcs = &cert_objdata_vtable;

		if(ckcapi_object_data_match(&cdata.base, match, count))
		{
			ret = register_cert_object(sess, cert, &obj);
			if(ret != CKR_OK)
				break;

			ASSERT(obj);

			/* Store away the object data for performance reasons */
			objdata = cert_alloc_data(sess, obj, cert);
			if(objdata)
			{
				ckcapi_session_take_object_data(sess, obj, objdata);

				/* For continuing the enumeration */
				cert = CertDuplicateCertificateContext(cert);
			}

			ckcapi_array_append(arr, obj->id);
		}
	}

	return ret;
}

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

	/* No store, no objects */
	if(!sess->store)
		return CKR_OK;

	cert = CertGetSubjectCertificateFromStore(sess->store, CKCAPI_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 */
	cdata.cert = cert;
	cdata.base.object = 0;
	cdata.base.data_funcs = &cert_objdata_vtable;

	if(ckcapi_object_data_match(&cdata.base, match, count))
	{
		ret = register_cert_object(sess, cert, &obj);
		if(ret == CKR_OK)
		{
			ASSERT(obj);

			/* Store away the object data for performance reasons */
			objdata = cert_alloc_data(sess, obj, cert);
			if(objdata)
			   ckcapi_session_take_object_data(sess, obj, objdata);
			ckcapi_array_append(arr, obj->id);
		}
	}

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

	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;
	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(CKCAPI_ENCODINGS, X509_MULTI_BYTE_INTEGER,
								  match[i].pValue, match[i].ulValueLen, 0, NULL, &size))
			{
				continue;
			}

			serial = calloc(1, size);
			if(!CryptDecodeObject(CKCAPI_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;
		}
	}

	/* Match a specific certificate */
	if(info.SerialNumber.cbData && info.Issuer.cbData)
		ret = match_in_store(sess, &info, match, count, arr);

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

	if(serial)
		free(serial);

	return ret;
}