#include "ckcapi.h"
#include "x509-usages.h"

#include "pkcs11/pkcs11n.h"

#include <wincrypt.h>

/*
 * These are the attributes expected by NSS on a trust object:
 *
 * CKA_CLASS
 * CKA_TOKEN
 * CKA_LABEL
 * CKA_CERT_SHA1_HASH
 * CKA_CERT_MD5_HASH
 * CKA_ISSUER
 * CKA_SUBJECT
 * CKA_TRUST_SERVER_AUTH
 * CKA_TRUST_CLIENT_AUTH
 * CKA_TRUST_EMAIL_PROTECTION
 * CKA_TRUST_CODE_SIGNING
 * CKA_SERIAL_NUMBER
 */

typedef struct _TrustObject
{
	CkCapiObject obj;

	/* Together these form the unique key. Must be contiguous */
	unsigned int otype;
	CK_OBJECT_HANDLE cert_obj;
}
TrustObject;

typedef struct _TrustData
{
	PCCERT_CONTEXT cert;
	CTL_USAGE* usage;
}
TrustData;

static CK_ULONG
has_usage(TrustData* trust_data, const char* oid)
{
	CTL_USAGE* usage = trust_data->usage;
	DWORD i;

	/* No usages, means anything goes */
	if(usage == NULL)
		return CKT_NETSCAPE_TRUSTED_DELEGATOR;

	for(i = 0; i < usage->cUsageIdentifier; ++i)
	{
		if(usage->rgpszUsageIdentifier[i] && 
		   strcmp(oid, usage->rgpszUsageIdentifier[i]) == 0)
			return CKT_NETSCAPE_TRUSTED_DELEGATOR;
	}

	/* TODO: What's the correct thing to return here? */
	return CKT_NETSCAPE_VALID_DELEGATOR;
}

static CK_RV
trust_bool_attribute(void* obj, CK_ATTRIBUTE_TYPE type, 
					 CK_VOID_PTR data, CK_ULONG_PTR len)
{
	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.
	 * TODO: This needs to be set to yes once we get change support.
	 */
	case CKA_MODIFIABLE:
		val = CK_FALSE;
		break;

	/* 
	 * TODO: Figure out what this is.
	 */
	case CKA_TRUST_STEP_UP_APPROVED:
		val = CK_FALSE;
		break;

	default:
		return CKR_ATTRIBUTE_TYPE_INVALID;
	};

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

static CK_RV
trust_ulong_attribute(void* obj, CK_ATTRIBUTE_TYPE type, 
					  CK_VOID_PTR data, CK_ULONG_PTR len)
{
	TrustData* trust_data = (TrustData*)obj;
	CK_ULONG val;

	ASSERT(obj);

	switch(type)
	{

	/*
	 * Object class.
	 * - Always CKO_NETSCAPE_TRUST for netscape trust
	 */
	case CKA_CLASS:
		val = CKO_NETSCAPE_TRUST;
		break;

	/* 
	 * Various trust flags 
	 */
	case CKA_TRUST_SERVER_AUTH:
		val = has_usage(trust_data, X509_USAGE_SERVER_AUTH);
		break;
	case CKA_TRUST_CLIENT_AUTH:
		val = has_usage(trust_data, X509_USAGE_CLIENT_AUTH);
		break;
	case CKA_TRUST_CODE_SIGNING:
		val = has_usage(trust_data, X509_USAGE_CODE_SIGNING);
		break;
	case CKA_TRUST_EMAIL_PROTECTION:
		val = has_usage(trust_data, X509_USAGE_EMAIL);
		break;
	case CKA_TRUST_IPSEC_END_SYSTEM:
		val = has_usage(trust_data, X509_USAGE_IPSEC_ENDPOINT);
		break;
	case CKA_TRUST_IPSEC_TUNNEL:
		val = has_usage(trust_data, X509_USAGE_IPSEC_TUNNEL);
		break;
	case CKA_TRUST_IPSEC_USER:
		val = has_usage(trust_data, X509_USAGE_IPSEC_USER);
		break;
	case CKA_TRUST_TIME_STAMPING:
		val = has_usage(trust_data, X509_USAGE_TIME_STAMPING);
		break;

	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));
}

static CK_RV
trust_bytes_attribute(void* obj, CK_ATTRIBUTE_TYPE type, 
					 CK_VOID_PTR data, CK_ULONG_PTR len)
{
	TrustData* trust_data = (TrustData*)obj;		
	ASSERT(obj);

	switch(type)
	{
	/*
	 * Forward these through to the certificate itself.
	 */
	case CKA_SUBJECT:
	case CKA_ISSUER:
	case CKA_SERIAL_NUMBER:
	case CKA_LABEL:
		ASSERT(trust_data->cert);
		return ckcapi_cert_get_bytes_attribute((void*)(trust_data->cert), 
											   type, data, len);

	/* 
	 * The hash of the DER encoded certificate.
	 */
	case CKA_CERT_MD5_HASH:
	case CKA_CERT_SHA1_HASH:
		if(!CryptHashCertificate(0, type == CKA_CERT_MD5_HASH ? CALG_MD5 : CALG_SHA1,
								 0, trust_data->cert->pbCertEncoded, 
								 trust_data->cert->cbCertEncoded, data, (DWORD*)len))
			return ckcapi_winerr_to_ckr(GetLastError());
		return CKR_OK;
	};

	return CKR_ATTRIBUTE_TYPE_INVALID;
}

static CK_RV
trust_date_attribute(void* obj, CK_ATTRIBUTE_TYPE type, 
					 CK_VOID_PTR data, CK_ULONG_PTR len)
{
	return CKR_ATTRIBUTE_TYPE_INVALID;
}

static void
trust_release(void* data)
{
	TrustData* trust_data = (TrustData*)data;
	ASSERT(trust_data);

	ASSERT(trust_data->cert);
	CertFreeCertificateContext(trust_data->cert);

	if(trust_data->usage)
		free(trust_data->usage);

	free(trust_data);
}

static const CkCapiObjectDataVtable trust_objdata_vtable = {
	trust_bool_attribute,
	trust_ulong_attribute,
	trust_bytes_attribute,
	trust_date_attribute,
	trust_release,
};

static CK_RV
parse_usage(TrustData* trust_data)
{
	DWORD size, usize, err;
	CTL_USAGE* usage;
	void* buf;

	ASSERT(!trust_data->usage);

	/* Get the size of the usage property */
	if(!CertGetCertificateContextProperty(trust_data->cert, CERT_CTL_USAGE_PROP_ID, 
										  NULL, &size))
	{
		err = GetLastError();

		/* No usage data is not an error */
		if(err == CRYPT_E_NOT_FOUND)
			return CKR_OK; 
		return ckcapi_winerr_to_ckr(err);
	}

	/* Now get the actual usage property */
	buf = _alloca(size);
	if(!CertGetCertificateContextProperty(trust_data->cert, CERT_CTL_USAGE_PROP_ID, 
										  buf, &size))
	{
		err = GetLastError();
		if(err == CRYPT_E_NOT_FOUND)
			return CKR_OK; 
		return ckcapi_winerr_to_ckr(err);
	}

	/* Get the decoded size of the usage property */
	if(!CryptDecodeObject(X509_ASN_ENCODING, X509_ENHANCED_KEY_USAGE,
						  buf, size, 0, NULL, &usize))
		return ckcapi_winerr_to_ckr(GetLastError());

	/* Allocate and decode it */
	usage = (CTL_USAGE*)calloc(1, usize);
	if(!usage) 
		return CKR_HOST_MEMORY;
	if(!CryptDecodeObject(X509_ASN_ENCODING, X509_ENHANCED_KEY_USAGE,
						  buf, size, 0, buf, &usize))
	{
		free(usage);
		return ckcapi_winerr_to_ckr(GetLastError());
	}

	trust_data->usage = usage;
	return CKR_OK;
}

static CK_RV 
trust_load_data(CkCapiSession* sess, CkCapiObject* obj, CkCapiObjectData* objdata)
{
	TrustObject* tobj = (TrustObject*)obj;
	TrustData* trust_data;
	CkCapiObjectData* certdata;
	CK_RV ret;

	ASSERT(tobj);
	ASSERT(objdata);

	/* Get the raw data for the certificate */
	ret = ckcapi_session_get_object_data_for(sess, tobj->cert_obj, &certdata);
	if(ret != CKR_OK)
		return ret;

	trust_data = (TrustData*)calloc(1, sizeof(TrustData));
	if(!trust_data)
		return CKR_HOST_MEMORY;

	ASSERT(certdata->data);
	trust_data->cert = certdata->data;

	/* Dig up the usage data */
	ret = parse_usage(trust_data);
	if(ret != CKR_OK)
	{
		free(trust_data);
		return ret;
	}

	/* And keep a reference to the certificate */
	trust_data->cert = CertDuplicateCertificateContext((PCCERT_CONTEXT)(certdata->data));

	objdata->data = trust_data;
	objdata->data_funcs = &trust_objdata_vtable;

	return CKR_OK;
}


static void 
trust_object_release(void* data)
{
	TrustObject* tobj = (TrustObject*)data;
	ASSERT(tobj);
	free(tobj);
}

static const CkCapiObjectVtable trust_object_vtable = {
	trust_load_data,
	trust_object_release,
};

static CK_RV
register_trust_object(CkCapiSession* sess, CkCapiObject* cert, CkCapiObject** obj)
{
	TrustObject* tobj;
	CK_RV ret;

	tobj = calloc(sizeof(TrustObject), 1);
	if(!tobj)
		return CKR_HOST_MEMORY;

	tobj->otype = OBJECT_TRUST;
	tobj->cert_obj = cert->id;

	tobj->obj.id = 0;
	tobj->obj.obj_funcs = &trust_object_vtable;
	tobj->obj.unique_key = UNIQUE_KEY_AT(tobj, otype);
	tobj->obj.unique_len = UNIQUE_KEY_LEN(tobj, otype, cert_obj);

	ret = ckcapi_object_register(sess, &(tobj->obj));
	if(ret != CKR_OK)
	{
		free(tobj);
		return ret;
	}

	ASSERT(tobj->obj.id != 0);
	*obj = &(tobj->obj);

	return CKR_OK;
}

static CK_RV
list_matching_certificates(CkCapiSession* sess, CK_ATTRIBUTE_PTR match, 
						   CK_ULONG count, CkCapiArray* arr)
{
	CK_OBJECT_CLASS cert_class = CKO_CERTIFICATE;
	CK_ATTRIBUTE search[3];
	CK_ULONG n_search = 0;
	CK_ULONG i;

	/* The class */
	search[0].type = CKA_CLASS;
	search[0].pValue = &cert_class;
	search[0].ulValueLen = sizeof(CK_OBJECT_CLASS);
	++n_search;

	for(i = 0; i < count && n_search < 3; ++i)
	{
		/* 
		 * These are the attributes that tie a certificate 
		 * to trust object, so try match certs with these
		 */
		if(match[i].type == CKA_ISSUER || 
		   match[i].type == CKA_SERIAL_NUMBER)
		{
			search[n_search].type = match[i].type;
			search[n_search].pValue = match[i].pValue;
			search[n_search].ulValueLen = match[i].ulValueLen;
			++n_search;
		}
	}

	/* Do the certificate search */
	return ckcapi_cert_find(sess, CKO_CERTIFICATE, search, n_search, arr);
}

CK_RV
ckcapi_trust_find(CkCapiSession* sess, CK_OBJECT_CLASS cls, 
				  CK_ATTRIBUTE_PTR match, CK_ULONG count, CkCapiArray* arr)
{
	CkCapiObject* obj;
	CkCapiObject* certobj;
	CkCapiObjectData* objdata;
	CkCapiArray* certarr;
	CK_RV ret = CKR_OK;
	CK_ULONG i;

	if(cls != CKO_NETSCAPE_TRUST)
		return CKR_OK;

	/* Get a list of all certificates */
	certarr = ckcapi_array_new(0, 1, sizeof(CkCapiObject*));
	if(!certarr)
		return CKR_HOST_MEMORY;
	ret = list_matching_certificates(sess, match, count, certarr);

	/* Now match each of them against our criteria */
	if(ret == CKR_OK)
	{
		for(i = 0; i < certarr->len; ++i)
		{
			certobj = ckcapi_array_index(certarr, CkCapiObject*, i);
			ASSERT(certobj);

			/* We'll register a trust object for any loaded certificate */
			ret = register_trust_object(sess, certobj, &obj);
			if(ret != CKR_OK)
				break;

			ASSERT(obj);

			ret = ckcapi_session_get_object_data(sess, obj, &objdata);
			if(ret != CKR_OK)
				break;

			/* Only return new object if it matches */
			if(ckcapi_object_data_match(objdata, match, count))
				ckcapi_array_append(arr, obj);
		}
	}

	ckcapi_array_free(certarr, TRUE);
	return ret;
}