/*
 * Copyright (c) 2007 Kungliga Tekniska Högskolan
 * (Royal Institute of Technology, Stockholm, Sweden).
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the Institute nor the names of its contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE INSTITUTE AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE INSTITUTE OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include "hx_locl.h"

#ifdef HAVE_FRAMEWORK_SECURITY

#include <Security/Security.h>

/* Missing function decls in pre Leopard */
#ifdef NEED_SECKEYGETCSPHANDLE_PROTO
OSStatus SecKeyGetCSPHandle(SecKeyRef, CSSM_CSP_HANDLE *);
OSStatus SecKeyGetCredentials(SecKeyRef, CSSM_ACL_AUTHORIZATION_TAG,
			      int, const CSSM_ACCESS_CREDENTIALS **);
#define kSecCredentialTypeDefault 0
#define CSSM_SIZE uint32_t
#endif


static int
getAttribute(SecKeychainItemRef itemRef, SecItemAttr item,
	     SecKeychainAttributeList **attrs)
{	
    SecKeychainAttributeInfo attrInfo;
    UInt32 attrFormat = 0;
    OSStatus ret;

    *attrs = NULL;

    attrInfo.count = 1;
    attrInfo.tag = &item;
    attrInfo.format = &attrFormat;

    ret = SecKeychainItemCopyAttributesAndData(itemRef, &attrInfo, NULL,
					       attrs, NULL, NULL);
    if (ret)
	return EINVAL;
    return 0;
}


/*
 *
 */

struct kc_rsa {
    SecKeychainItemRef item;
    size_t keysize;
};


static int
kc_rsa_public_encrypt(int flen,
		      const unsigned char *from,
		      unsigned char *to,
		      RSA *rsa,
		      int padding)
{
    return -1;
}

static int
kc_rsa_public_decrypt(int flen,
		      const unsigned char *from,
		      unsigned char *to,
		      RSA *rsa,
		      int padding)
{
    return -1;
}


static int
kc_rsa_private_encrypt(int flen,
		       const unsigned char *from,
		       unsigned char *to,
		       RSA *rsa,
		       int padding)
{
    struct kc_rsa *kc = RSA_get_app_data(rsa);

    CSSM_RETURN cret;
    OSStatus ret;
    const CSSM_ACCESS_CREDENTIALS *creds;
    SecKeyRef privKeyRef = (SecKeyRef)kc->item;
    CSSM_CSP_HANDLE cspHandle;
    const CSSM_KEY *cssmKey;
    CSSM_CC_HANDLE sigHandle = 0;
    CSSM_DATA sig, in;
    int fret = 0;

    if (padding != RSA_PKCS1_PADDING)
	return -1;

    cret = SecKeyGetCSSMKey(privKeyRef, &cssmKey);
    if(cret) abort();

    cret = SecKeyGetCSPHandle(privKeyRef, &cspHandle);
    if(cret) abort();

    ret = SecKeyGetCredentials(privKeyRef, CSSM_ACL_AUTHORIZATION_SIGN,
			       kSecCredentialTypeDefault, &creds);
    if(ret) abort();

    ret = CSSM_CSP_CreateSignatureContext(cspHandle, CSSM_ALGID_RSA,
					  creds, cssmKey, &sigHandle);
    if(ret) abort();

    in.Data = (uint8 *)from;
    in.Length = flen;
	
    sig.Data = (uint8 *)to;
    sig.Length = kc->keysize;
	
    cret = CSSM_SignData(sigHandle, &in, 1, CSSM_ALGID_NONE, &sig);
    if(cret) {
	/* cssmErrorString(cret); */
	fret = -1;
    } else
	fret = sig.Length;

    if(sigHandle)
	CSSM_DeleteContext(sigHandle);

    return fret;
}

static int
kc_rsa_private_decrypt(int flen, const unsigned char *from, unsigned char *to,
		       RSA * rsa, int padding)
{
    struct kc_rsa *kc = RSA_get_app_data(rsa);

    CSSM_RETURN cret;
    OSStatus ret;
    const CSSM_ACCESS_CREDENTIALS *creds;
    SecKeyRef privKeyRef = (SecKeyRef)kc->item;
    CSSM_CSP_HANDLE cspHandle;
    const CSSM_KEY *cssmKey;
    CSSM_CC_HANDLE handle = 0;
    CSSM_DATA out, in, rem;
    int fret = 0;
    CSSM_SIZE outlen = 0;
    char remdata[1024];

    if (padding != RSA_PKCS1_PADDING)
	return -1;

    cret = SecKeyGetCSSMKey(privKeyRef, &cssmKey);
    if(cret) abort();

    cret = SecKeyGetCSPHandle(privKeyRef, &cspHandle);
    if(cret) abort();

    ret = SecKeyGetCredentials(privKeyRef, CSSM_ACL_AUTHORIZATION_DECRYPT,
			       kSecCredentialTypeDefault, &creds);
    if(ret) abort();


    ret = CSSM_CSP_CreateAsymmetricContext (cspHandle,
					    CSSM_ALGID_RSA,
					    creds,
					    cssmKey,
					    CSSM_PADDING_PKCS1,
					    &handle);
    if(ret) abort();

    in.Data = (uint8 *)from;
    in.Length = flen;
	
    out.Data = (uint8 *)to;
    out.Length = kc->keysize;
	
    rem.Data = (uint8 *)remdata;
    rem.Length = sizeof(remdata);

    cret = CSSM_DecryptData(handle, &in, 1, &out, 1, &outlen, &rem);
    if(cret) {
	/* cssmErrorString(cret); */
	fret = -1;
    } else
	fret = out.Length;

    if(handle)
	CSSM_DeleteContext(handle);

    return fret;
}

static int
kc_rsa_init(RSA *rsa)
{
    return 1;
}

static int
kc_rsa_finish(RSA *rsa)
{
    struct kc_rsa *kc_rsa = RSA_get_app_data(rsa);
    CFRelease(kc_rsa->item);
    memset(kc_rsa, 0, sizeof(*kc_rsa));
    free(kc_rsa);
    return 1;
}

static const RSA_METHOD kc_rsa_pkcs1_method = {
    "hx509 Keychain PKCS#1 RSA",
    kc_rsa_public_encrypt,
    kc_rsa_public_decrypt,
    kc_rsa_private_encrypt,
    kc_rsa_private_decrypt,
    NULL,
    NULL,
    kc_rsa_init,
    kc_rsa_finish,
    0,
    NULL,
    NULL,
    NULL
};

static int
set_private_key(hx509_context context,
		SecKeychainItemRef itemRef,
		hx509_cert cert)
{
    struct kc_rsa *kc;
    hx509_private_key key;
    RSA *rsa;
    int ret;

    ret = _hx509_private_key_init(&key, NULL, NULL);
    if (ret)
	return ret;

    kc = calloc(1, sizeof(*kc));
    if (kc == NULL)
	_hx509_abort("out of memory");

    kc->item = itemRef;

    rsa = RSA_new();
    if (rsa == NULL)
	_hx509_abort("out of memory");

    /* Argh, fake modulus since OpenSSL API is on crack */
    {
	SecKeychainAttributeList *attrs = NULL;
	uint32_t size;
	void *data;

	rsa->n = BN_new();
	if (rsa->n == NULL) abort();

	ret = getAttribute(itemRef, kSecKeyKeySizeInBits, &attrs);
	if (ret) abort();

	size = *(uint32_t *)attrs->attr[0].data;
	SecKeychainItemFreeAttributesAndData(attrs, NULL);

	kc->keysize = (size + 7) / 8;

	data = malloc(kc->keysize);
	memset(data, 0xe0, kc->keysize);
	BN_bin2bn(data, kc->keysize, rsa->n);
	free(data);
    }
    rsa->e = NULL;

    RSA_set_method(rsa, &kc_rsa_pkcs1_method);
    ret = RSA_set_app_data(rsa, kc);
    if (ret != 1)
	_hx509_abort("RSA_set_app_data");

    _hx509_private_key_assign_rsa(key, rsa);
    _hx509_cert_assign_key(cert, key);

    return 0;
}

/*
 *
 */

struct ks_keychain {
    int anchors;
    SecKeychainRef keychain;
};

static int
keychain_init(hx509_context context,
	      hx509_certs certs, void **data, int flags,
	      const char *residue, hx509_lock lock)
{
    struct ks_keychain *ctx;

    ctx = calloc(1, sizeof(*ctx));
    if (ctx == NULL) {
	hx509_clear_error_string(context);
	return ENOMEM;
    }

    if (residue) {
	if (strcasecmp(residue, "system-anchors") == 0) {
	    ctx->anchors = 1;
	} else if (strncasecmp(residue, "FILE:", 5) == 0) {
	    OSStatus ret;

	    ret = SecKeychainOpen(residue + 5, &ctx->keychain);
	    if (ret != noErr) {
		hx509_set_error_string(context, 0, ENOENT,
				       "Failed to open %s", residue);
		return ENOENT;
	    }
	} else {
	    hx509_set_error_string(context, 0, ENOENT,
				   "Unknown subtype %s", residue);
	    return ENOENT;
	}
    }

    *data = ctx;
    return 0;
}

/*
 *
 */

static int
keychain_free(hx509_certs certs, void *data)
{
    struct ks_keychain *ctx = data;
    if (ctx->keychain)
	CFRelease(ctx->keychain);
    memset(ctx, 0, sizeof(*ctx));
    free(ctx);
    return 0;
}

/*
 *
 */

struct iter {
    hx509_certs certs;
    void *cursor;
    SecKeychainSearchRef searchRef;
};

static int
keychain_iter_start(hx509_context context,
		    hx509_certs certs, void *data, void **cursor)
{
    struct ks_keychain *ctx = data;
    struct iter *iter;

    iter = calloc(1, sizeof(*iter));
    if (iter == NULL) {
	hx509_set_error_string(context, 0, ENOMEM, "out of memory");
	return ENOMEM;
    }

    if (ctx->anchors) {
        CFArrayRef anchors;
	int ret;
	int i;

	ret = hx509_certs_init(context, "MEMORY:ks-file-create",
			       0, NULL, &iter->certs);
	if (ret) {
	    free(iter);
	    return ret;
	}

	ret = SecTrustCopyAnchorCertificates(&anchors);
	if (ret != 0) {
	    hx509_certs_free(&iter->certs);
	    free(iter);
	    hx509_set_error_string(context, 0, ENOMEM,
				   "Can't get trust anchors from Keychain");
	    return ENOMEM;
	}
	for (i = 0; i < CFArrayGetCount(anchors); i++) {
	    SecCertificateRef cr;
	    hx509_cert cert;
	    CSSM_DATA cssm;

	    cr = (SecCertificateRef)CFArrayGetValueAtIndex(anchors, i);

	    SecCertificateGetData(cr, &cssm);

	    ret = hx509_cert_init_data(context, cssm.Data, cssm.Length, &cert);
	    if (ret)
		continue;

	    ret = hx509_certs_add(context, iter->certs, cert);
	    hx509_cert_free(cert);
	}
	CFRelease(anchors);
    }

    if (iter->certs) {
	int ret;
	ret = hx509_certs_start_seq(context, iter->certs, &iter->cursor);
	if (ret) {
	    hx509_certs_free(&iter->certs);
	    free(iter);
	    return ret;
	}
    } else {
	OSStatus ret;

	ret = SecKeychainSearchCreateFromAttributes(ctx->keychain,
						    kSecCertificateItemClass,
						    NULL,
						    &iter->searchRef);
	if (ret) {
	    free(iter);
	    hx509_set_error_string(context, 0, ret,
				   "Failed to start search for attributes");
	    return ENOMEM;
	}
    }

    *cursor = iter;
    return 0;
}

/*
 *
 */

static int
keychain_iter(hx509_context context,
	      hx509_certs certs, void *data, void *cursor, hx509_cert *cert)
{
    SecKeychainAttributeList *attrs = NULL;
    SecKeychainAttributeInfo attrInfo;
    UInt32 attrFormat[1] = { 0 };
    SecKeychainItemRef itemRef;
    SecItemAttr item[1];
    struct iter *iter = cursor;
    OSStatus ret;
    UInt32 len;
    void *ptr = NULL;

    if (iter->certs)
	return hx509_certs_next_cert(context, iter->certs, iter->cursor, cert);

    *cert = NULL;

    ret = SecKeychainSearchCopyNext(iter->searchRef, &itemRef);
    if (ret == errSecItemNotFound)
	return 0;
    else if (ret != 0)
	return EINVAL;
	
    /*
     * Pick out certificate and matching "keyid"
     */

    item[0] = kSecPublicKeyHashItemAttr;

    attrInfo.count = 1;
    attrInfo.tag = item;
    attrInfo.format = attrFormat;

    ret = SecKeychainItemCopyAttributesAndData(itemRef, &attrInfo, NULL,
					       &attrs, &len, &ptr);
    if (ret)
	return EINVAL;

    ret = hx509_cert_init_data(context, ptr, len, cert);
    if (ret)
	goto out;

    /*
     * Find related private key if there is one by looking at
     * kSecPublicKeyHashItemAttr == kSecKeyLabel
     */
    {
	SecKeychainSearchRef search;
	SecKeychainAttribute attrKeyid;
	SecKeychainAttributeList attrList;

	attrKeyid.tag = kSecKeyLabel;
	attrKeyid.length = attrs->attr[0].length;
	attrKeyid.data = attrs->attr[0].data;
	
	attrList.count = 1;
	attrList.attr = &attrKeyid;

	ret = SecKeychainSearchCreateFromAttributes(NULL,
						    CSSM_DL_DB_RECORD_PRIVATE_KEY,
						    &attrList,
						    &search);
	if (ret) {
	    ret = 0;
	    goto out;
	}

	ret = SecKeychainSearchCopyNext(search, &itemRef);
	CFRelease(search);
	if (ret == errSecItemNotFound) {
	    ret = 0;
	    goto out;
	} else if (ret) {
	    ret = EINVAL;
	    goto out;
	}
	set_private_key(context, itemRef, *cert);
    }

out:
    SecKeychainItemFreeAttributesAndData(attrs, ptr);

    return ret;
}

/*
 *
 */

static int
keychain_iter_end(hx509_context context,
		  hx509_certs certs,
		  void *data,
		  void *cursor)
{
    struct iter *iter = cursor;

    if (iter->certs) {
	hx509_certs_end_seq(context, iter->certs, iter->cursor);
	hx509_certs_free(&iter->certs);
    } else {
	CFRelease(iter->searchRef);
    }

    memset(iter, 0, sizeof(*iter));
    free(iter);
    return 0;
}

/*
 *
 */

struct hx509_keyset_ops keyset_keychain = {
    "KEYCHAIN",
    0,
    keychain_init,
    NULL,
    keychain_free,
    NULL,
    NULL,
    keychain_iter_start,
    keychain_iter,
    keychain_iter_end
};

#endif /* HAVE_FRAMEWORK_SECURITY */

/*
 *
 */

void
_hx509_ks_keychain_register(hx509_context context)
{
#ifdef HAVE_FRAMEWORK_SECURITY
    _hx509_ks_register(context, &keyset_keychain);
#endif
}
