/*
 * The Initial Developer of the Original Code is International
 * Business Machines Corporation. Portions created by IBM
 * Corporation are Copyright (C) 2005 International Business
 * Machines Corporation. All Rights Reserved.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the Common Public License as published by
 * IBM Corporation; either version 1 of the License, or (at your option)
 * any later version.
 *
 * This program 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
 * Common Public License for more details.
 *
 * You should have received a copy of the Common Public License
 * along with this program; if not, a copy can be viewed at
 * http://www.opensource.org/licenses/cpl1.0.php.
 */
/* (C) COPYRIGHT International Business Machines Corp. 2001, 2002, 2005 */
/*
 * Copyright 2009 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */

#include "tpmtok_int.h"

static CK_BBOOL true = TRUE, false = FALSE;

static CK_RV
key_mgr_get_private_key_type(
        CK_BYTE *keydata,
        CK_ULONG        keylen,
        CK_KEY_TYPE *keytype)
{
        CK_BYTE  *alg = NULL;
        CK_BYTE  *priv_key = NULL;
        CK_ULONG  alg_len;
        CK_RV    rc;

        rc = ber_decode_PrivateKeyInfo(keydata, keylen, &alg,
            &alg_len, &priv_key);
        if (rc != CKR_OK) {
                return (rc);
        }
        if (alg_len >= ber_rsaEncryptionLen) {
                if (memcmp(alg, ber_rsaEncryption,
                    ber_rsaEncryptionLen) == 0) {
                        *keytype = CKK_RSA;
                        return (CKR_OK);
                }
        }

        return (CKR_TEMPLATE_INCOMPLETE);
}

CK_RV
key_mgr_generate_key_pair(SESSION          * sess,
        CK_MECHANISM    * mech,
        CK_ATTRIBUTE    * publ_tmpl,
        CK_ULONG            publ_count,
        CK_ATTRIBUTE    * priv_tmpl,
        CK_ULONG            priv_count,
        CK_OBJECT_HANDLE  * publ_key_handle,
        CK_OBJECT_HANDLE  * priv_key_handle)
{
        OBJECT  * publ_key_obj = NULL;
        OBJECT  * priv_key_obj = NULL;
        CK_ATTRIBUTE  * attr     = NULL;
        CK_ATTRIBUTE  * new_attr        = NULL;
        CK_ULONG        i, keyclass, subclass = 0;
        CK_BBOOL        flag;
        CK_RV      rc;

        if (! sess || ! mech || ! publ_key_handle || ! priv_key_handle) {
                return (CKR_FUNCTION_FAILED);
        }
        if (! publ_tmpl && (publ_count != 0)) {
                return (CKR_FUNCTION_FAILED);
        }
        if (! priv_tmpl && (priv_count != 0)) {
                return (CKR_FUNCTION_FAILED);
        }

        for (i = 0; i < publ_count; i++) {
                if (publ_tmpl[i].type == CKA_CLASS) {
                        keyclass = *(CK_OBJECT_CLASS *)publ_tmpl[i].pValue;
                        if (keyclass != CKO_PUBLIC_KEY) {
                                return (CKR_TEMPLATE_INCONSISTENT);
                        }
                }

                if (publ_tmpl[i].type == CKA_KEY_TYPE)
                        subclass = *(CK_ULONG *)publ_tmpl[i].pValue;
        }


        for (i = 0; i < priv_count; i++) {
                if (priv_tmpl[i].type == CKA_CLASS) {
                        keyclass = *(CK_OBJECT_CLASS *)priv_tmpl[i].pValue;
                        if (keyclass != CKO_PRIVATE_KEY) {
                                return (CKR_TEMPLATE_INCONSISTENT);
                        }
                }

                if (priv_tmpl[i].type == CKA_KEY_TYPE) {
                        CK_ULONG temp = *(CK_ULONG *)priv_tmpl[i].pValue;
                        if (temp != subclass) {
                                return (CKR_TEMPLATE_INCONSISTENT);
                        }
                }
        }


        switch (mech->mechanism) {
                case CKM_RSA_PKCS_KEY_PAIR_GEN:
                        if (subclass != 0 && subclass != CKK_RSA) {
                                return (CKR_TEMPLATE_INCONSISTENT);
                        }

                        subclass = CKK_RSA;
                break;

                default:
                        return (CKR_MECHANISM_INVALID);
        }


        rc = object_mgr_create_skel(sess,
            publ_tmpl,  publ_count, MODE_KEYGEN,
            CKO_PUBLIC_KEY,  subclass, &publ_key_obj);

        if (rc != CKR_OK) {
                goto error;
        }
        rc = object_mgr_create_skel(sess,
            priv_tmpl,  priv_count, MODE_KEYGEN,
            CKO_PRIVATE_KEY, subclass, &priv_key_obj);

        if (rc != CKR_OK) {
                goto error;
        }

        switch (mech->mechanism) {
                case CKM_RSA_PKCS_KEY_PAIR_GEN:
                        rc = ckm_rsa_key_pair_gen(
                            sess->hContext,
                            publ_key_obj->template,
                            priv_key_obj->template);
                break;

                default:
                        rc = CKR_MECHANISM_INVALID;
                break;
        }

        if (rc != CKR_OK) {
                goto error;
        }

        /*
         * we can now set CKA_ALWAYS_SENSITIVE and CKA_NEVER_EXTRACTABLE
         * to their appropriate values.  this only applies to CKO_SECRET_KEY
         * and CKO_PRIVATE_KEY objects
         */
        flag = template_attribute_find(priv_key_obj->template,
            CKA_SENSITIVE, &attr);
        if (flag == TRUE) {
                flag = *(CK_BBOOL *)attr->pValue;

                rc = build_attribute(CKA_ALWAYS_SENSITIVE, &flag,
                    sizeof (CK_BBOOL), &new_attr);
                if (rc != CKR_OK) {
                        goto error;
                }
                (void) template_update_attribute(priv_key_obj->template,
                    new_attr);

        } else {
                rc = CKR_FUNCTION_FAILED;
                goto error;
        }


        flag = template_attribute_find(priv_key_obj->template,
            CKA_EXTRACTABLE, &attr);
        if (flag == TRUE) {
                flag = *(CK_BBOOL *)attr->pValue;

                rc = build_attribute(CKA_NEVER_EXTRACTABLE, &true,
                    sizeof (CK_BBOOL), &new_attr);
                if (rc != CKR_OK) {
                        goto error;
                }
                if (flag == TRUE)
                        *(CK_BBOOL *)new_attr->pValue = false;

                (void) template_update_attribute(priv_key_obj->template,
                    new_attr);

        } else {
                rc = CKR_FUNCTION_FAILED;
                goto error;
        }

        rc = object_mgr_create_final(sess, publ_key_obj, publ_key_handle);
        if (rc != CKR_OK) {
                goto error;
        }
        rc = object_mgr_create_final(sess, priv_key_obj, priv_key_handle);
        if (rc != CKR_OK) {
                (void) object_mgr_destroy_object(sess, *publ_key_handle);
                publ_key_obj = NULL;
                goto error;
        }
        return (rc);

error:
        if (publ_key_obj)
                (void) object_free(publ_key_obj);
        if (priv_key_obj)
                (void) object_free(priv_key_obj);

        *publ_key_handle = 0;
        *priv_key_handle = 0;

        return (rc);
}

CK_RV
key_mgr_wrap_key(SESSION           * sess,
        CK_BBOOL        length_only,
        CK_MECHANISM    * mech,
        CK_OBJECT_HANDLE    h_wrapping_key,
        CK_OBJECT_HANDLE    h_key,
        CK_BYTE   * wrapped_key,
        CK_ULONG  * wrapped_key_len) {
        ENCR_DECR_CONTEXT * ctx = NULL;
        OBJECT    * key1_obj  = NULL;
        OBJECT    * key2_obj  = NULL;
        CK_ATTRIBUTE    * attr  = NULL;
        CK_BYTE   * data        = NULL;
        CK_ULONG    data_len;
        CK_OBJECT_CLASS class;
        CK_KEY_TYPE      keytype;
        CK_BBOOL    flag;
        CK_RV   rc;

        if (! sess || ! wrapped_key_len) {
                return (CKR_FUNCTION_FAILED);
        }

        rc = object_mgr_find_in_map1(sess->hContext, h_wrapping_key, &key1_obj);
        if (rc != CKR_OK) {
                return (CKR_WRAPPING_KEY_HANDLE_INVALID);
        }
        rc = object_mgr_find_in_map1(sess->hContext, h_key, &key2_obj);
        if (rc != CKR_OK) {
                return (CKR_KEY_HANDLE_INVALID);
        }

        rc = template_attribute_find(key2_obj->template,
            CKA_EXTRACTABLE, &attr);
        if (rc == FALSE) {
                return (CKR_KEY_NOT_WRAPPABLE);
        } else {
                flag = *(CK_BBOOL *)attr->pValue;
                if (flag == FALSE) {
                        return (CKR_KEY_NOT_WRAPPABLE);
                }
        }

        rc = template_attribute_find(key2_obj->template, CKA_CLASS, &attr);
        if (rc == FALSE) {
                return (CKR_KEY_NOT_WRAPPABLE);
        } else
                class = *(CK_OBJECT_CLASS *)attr->pValue;

        switch (mech->mechanism) {
                case CKM_RSA_PKCS:
                if (class != CKO_SECRET_KEY) {
                        return (CKR_KEY_NOT_WRAPPABLE);
                }
                break;

                default:
                return (CKR_KEY_NOT_WRAPPABLE);
        }

        rc = template_attribute_find(key2_obj->template,
            CKA_KEY_TYPE, &attr);
        if (rc == FALSE)
                return (CKR_KEY_NOT_WRAPPABLE);
        else
                keytype = *(CK_KEY_TYPE *)attr->pValue;

        switch (keytype) {
                case CKK_RSA:
                rc = rsa_priv_wrap_get_data(key2_obj->template, length_only,
                    &data, &data_len);
                if (rc != CKR_OK) {
                        return (rc);
                }
                break;

                case CKK_GENERIC_SECRET:
                rc = generic_secret_wrap_get_data(key2_obj->template,
                    length_only, &data, &data_len);
                if (rc != CKR_OK) {
                        return (rc);
                }
                break;
                default:
                return (CKR_KEY_NOT_WRAPPABLE);
        }

        switch (mech->mechanism) {
                case CKM_RSA_PKCS:
                break;

                default:
                return (CKR_KEY_NOT_WRAPPABLE);
        }

        ctx = (ENCR_DECR_CONTEXT *)malloc(sizeof (ENCR_DECR_CONTEXT));
        if (! ctx) {
                return (CKR_HOST_MEMORY);
        }
        (void) memset(ctx, 0x0, sizeof (ENCR_DECR_CONTEXT));

        rc = encr_mgr_init(sess, ctx, OP_WRAP, mech, h_wrapping_key);
        if (rc != CKR_OK) {
                return (rc);
        }
        rc = encr_mgr_encrypt(sess,     length_only,
            ctx, data,  data_len, wrapped_key, wrapped_key_len);

        if (data != NULL) {
                free(data);
        }
        (void) encr_mgr_cleanup(ctx);
        free(ctx);

        return (rc);
}

CK_RV
key_mgr_unwrap_key(SESSION         * sess,
        CK_MECHANISM    * mech,
        CK_ATTRIBUTE    * attributes,
        CK_ULONG            attrib_count,
        CK_BYTE    * wrapped_key,
        CK_ULONG            wrapped_key_len,
        CK_OBJECT_HANDLE    h_unwrapping_key,
        CK_OBJECT_HANDLE  * h_unwrapped_key)
{
        ENCR_DECR_CONTEXT * ctx = NULL;
        OBJECT      * key_obj = NULL;
        CK_BYTE    * data = NULL;
        CK_ULONG            data_len;
        CK_ULONG            keyclass, keytype;
        CK_ULONG            i;
        CK_BBOOL            found_class, found_type, fromend;
        CK_RV           rc;


        if (! sess || ! wrapped_key || ! h_unwrapped_key) {
                return (CKR_FUNCTION_FAILED);
        }

        rc = object_mgr_find_in_map1(sess->hContext, h_unwrapping_key,
            &key_obj);
        if (rc != CKR_OK) {
                return (CKR_WRAPPING_KEY_HANDLE_INVALID);
        }

        found_class    = FALSE;
        found_type      = FALSE;

        switch (mech->mechanism) {
                case CKM_RSA_PKCS:
                        keyclass = CKO_SECRET_KEY;
                        found_class = TRUE;
                break;
        }

        for (i = 0; i < attrib_count; i++) {
                switch (attributes[i].type) {
                        case CKA_CLASS:
                        keyclass = *(CK_OBJECT_CLASS *)attributes[i].pValue;
                        found_class = TRUE;
                        break;

                        case CKA_KEY_TYPE:
                        keytype = *(CK_KEY_TYPE *)attributes[i].pValue;
                        found_type = TRUE;
                        break;
                }
        }

        if (found_class == FALSE || (found_type == FALSE && keyclass !=
            CKO_PRIVATE_KEY)) {
                return (CKR_TEMPLATE_INCOMPLETE);
        }

        switch (mech->mechanism) {
                case CKM_RSA_PKCS:
                if (keyclass != CKO_SECRET_KEY) {
                        return (CKR_TEMPLATE_INCONSISTENT);
                }
                break;
                default:
                        return (CKR_MECHANISM_INVALID);
        }


        ctx = (ENCR_DECR_CONTEXT *)malloc(sizeof (ENCR_DECR_CONTEXT));
        if (! ctx) {
                return (CKR_HOST_MEMORY);
        }
        (void) memset(ctx, 0x0, sizeof (ENCR_DECR_CONTEXT));

        rc = decr_mgr_init(sess, ctx, OP_UNWRAP, mech, h_unwrapping_key);
        if (rc != CKR_OK)
                return (rc);

        rc = decr_mgr_decrypt(sess,
            TRUE, ctx, wrapped_key, wrapped_key_len,
            data,       &data_len);
        if (rc != CKR_OK) {
                goto error;
        }
        data = (CK_BYTE *)malloc(data_len);
        if (! data) {
                rc = CKR_HOST_MEMORY;
                goto error;
        }

        rc = decr_mgr_decrypt(sess,
            FALSE, ctx, wrapped_key, wrapped_key_len,
            data,       &data_len);

        (void) decr_mgr_cleanup(ctx);
        free(ctx);

        if (rc != CKR_OK) {
                goto error;
        }
        /*
         * if we use X.509, the data will be padded from the front with zeros.
         * PKCS #11 specifies that for this mechanism, CK_VALUE is to be read
         * from the end of the data.
         *
         * Note: the PKCS #11 reference implementation gets this wrong.
         */
        if (mech->mechanism == CKM_RSA_X_509)
                fromend = TRUE;
        else
        fromend = FALSE;

        if (keyclass == CKO_PRIVATE_KEY) {
                rc = key_mgr_get_private_key_type(data, data_len, &keytype);
                if (rc != CKR_OK) {
                        goto error;
                }
        }

        rc = object_mgr_create_skel(sess,
            attributes,    attrib_count,
            MODE_UNWRAP, keyclass,      keytype,
            &key_obj);
        if (rc != CKR_OK) {
                goto error;
        }
        switch (keyclass) {
                case CKO_SECRET_KEY:
                rc = secret_key_unwrap(key_obj->template, keytype, data,
                    data_len, fromend);
                break;

                case CKO_PRIVATE_KEY:
                rc = priv_key_unwrap(key_obj->template, keytype,
                    data, data_len);
                break;

                default:
                rc = CKR_WRAPPED_KEY_INVALID;
                break;
        }

        if (rc != CKR_OK) {
                goto error;
        }
        rc = object_mgr_create_final(sess, key_obj, h_unwrapped_key);
        if (rc != CKR_OK) {
                goto error;
        }
        if (data) free(data);
        return (rc);

error:
        if (key_obj) (void) object_free(key_obj);
        if (data)    free(data);

        return (rc);
}