/*
 * Copyright (c) 1997 - 2008 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 "krb5_locl.h"

#include <pkinit_asn1.h>

KRB5_LIB_FUNCTION krb5_error_code KRB5_LIB_CALL
_krb5_pk_octetstring2key(krb5_context context,
			 krb5_enctype type,
			 const void *dhdata,
			 size_t dhsize,
			 const heim_octet_string *c_n,
			 const heim_octet_string *k_n,
			 krb5_keyblock *key)
{
    struct _krb5_encryption_type *et = _krb5_find_enctype(type);
    krb5_error_code ret;
    size_t keylen, offset;
    void *keydata;
    unsigned char counter;
    unsigned char shaoutput[SHA_DIGEST_LENGTH];
    EVP_MD_CTX *m;

    if(et == NULL) {
	krb5_set_error_message(context, KRB5_PROG_ETYPE_NOSUPP,
			       N_("encryption type %d not supported", ""),
			       type);
	return KRB5_PROG_ETYPE_NOSUPP;
    }
    keylen = (et->keytype->bits + 7) / 8;

    keydata = malloc(keylen);
    if (keydata == NULL)
	return krb5_enomem(context);

    m = EVP_MD_CTX_create();
    if (m == NULL) {
	free(keydata);
	return krb5_enomem(context);
    }

    counter = 0;
    offset = 0;
    do {

	EVP_DigestInit_ex(m, EVP_sha1(), NULL);
	EVP_DigestUpdate(m, &counter, 1);
	EVP_DigestUpdate(m, dhdata, dhsize);

	if (c_n)
	    EVP_DigestUpdate(m, c_n->data, c_n->length);
	if (k_n)
	    EVP_DigestUpdate(m, k_n->data, k_n->length);

	EVP_DigestFinal_ex(m, shaoutput, NULL);

	memcpy((unsigned char *)keydata + offset,
	       shaoutput,
	       min(keylen - offset, sizeof(shaoutput)));

	offset += sizeof(shaoutput);
	counter++;
    } while(offset < keylen);
    memset_s(shaoutput, sizeof(shaoutput), 0, sizeof(shaoutput));

    EVP_MD_CTX_destroy(m);

    ret = krb5_random_to_key(context, type, keydata, keylen, key);
    memset_s(keydata, sizeof(keylen), 0, sizeof(keylen));
    free(keydata);
    return ret;
}

static krb5_error_code
encode_uvinfo(krb5_context context, krb5_const_principal p, krb5_data *data)
{
    KRB5PrincipalName pn;
    krb5_error_code ret;
    size_t size = 0;

    pn.principalName = p->name;
    pn.realm = p->realm;

    ASN1_MALLOC_ENCODE(KRB5PrincipalName, data->data, data->length,
		       &pn, &size, ret);
    if (ret) {
	krb5_data_zero(data);
	krb5_set_error_message(context, ret,
			       N_("Failed to encode KRB5PrincipalName", ""));
	return ret;
    }
    if (data->length != size)
	krb5_abortx(context, "asn1 compiler internal error");
    return 0;
}

static krb5_error_code
encode_otherinfo(krb5_context context,
		 const AlgorithmIdentifier *ai,
		 krb5_const_principal client,
		 krb5_const_principal server,
		 krb5_enctype enctype,
		 const krb5_data *as_req,
		 const krb5_data *pk_as_rep,
		 const Ticket *ticket,
		 krb5_data *other)
{
    PkinitSP80056AOtherInfo otherinfo;
    PkinitSuppPubInfo pubinfo;
    krb5_error_code ret;
    krb5_data pub;
    size_t size = 0;

    krb5_data_zero(other);
    memset(&otherinfo, 0, sizeof(otherinfo));
    memset(&pubinfo, 0, sizeof(pubinfo));

    pubinfo.enctype = enctype;
    pubinfo.as_REQ = *as_req;
    pubinfo.pk_as_rep = *pk_as_rep;
    pubinfo.ticket = *ticket;
    ASN1_MALLOC_ENCODE(PkinitSuppPubInfo, pub.data, pub.length,
		       &pubinfo, &size, ret);
    if (ret) {
	krb5_set_error_message(context, ret, N_("malloc: out of memory", ""));
	return ret;
    }
    if (pub.length != size)
	krb5_abortx(context, "asn1 compiler internal error");

    ret = encode_uvinfo(context, client, &otherinfo.partyUInfo);
    if (ret) {
	free(pub.data);
	return ret;
    }
    ret = encode_uvinfo(context, server, &otherinfo.partyVInfo);
    if (ret) {
	free(otherinfo.partyUInfo.data);
	free(pub.data);
	return ret;
    }

    otherinfo.algorithmID = *ai;
    otherinfo.suppPubInfo = &pub;

    ASN1_MALLOC_ENCODE(PkinitSP80056AOtherInfo, other->data, other->length,
		       &otherinfo, &size, ret);
    free(otherinfo.partyUInfo.data);
    free(otherinfo.partyVInfo.data);
    free(pub.data);
    if (ret) {
	krb5_set_error_message(context, ret, N_("malloc: out of memory", ""));
	return ret;
    }
    if (other->length != size)
	krb5_abortx(context, "asn1 compiler internal error");

    return 0;
}



KRB5_LIB_FUNCTION krb5_error_code KRB5_LIB_CALL
_krb5_pk_kdf(krb5_context context,
	     const struct AlgorithmIdentifier *ai,
	     const void *dhdata,
	     size_t dhsize,
	     krb5_const_principal client,
	     krb5_const_principal server,
	     krb5_enctype enctype,
	     const krb5_data *as_req,
	     const krb5_data *pk_as_rep,
	     const Ticket *ticket,
	     krb5_keyblock *key)
{
    struct _krb5_encryption_type *et;
    krb5_error_code ret;
    krb5_data other;
    size_t keylen, offset;
    uint32_t counter;
    unsigned char *keydata;
    unsigned char shaoutput[SHA512_DIGEST_LENGTH];
    const EVP_MD *md;
    EVP_MD_CTX *m;

    if (der_heim_oid_cmp(&asn1_oid_id_pkinit_kdf_ah_sha1, &ai->algorithm) == 0) {
        md = EVP_sha1();
    } else if (der_heim_oid_cmp(&asn1_oid_id_pkinit_kdf_ah_sha256, &ai->algorithm) == 0) {
        md = EVP_sha256();
    } else if (der_heim_oid_cmp(&asn1_oid_id_pkinit_kdf_ah_sha512, &ai->algorithm) == 0) {
        md = EVP_sha512();
    } else {
	krb5_set_error_message(context, KRB5_PROG_ETYPE_NOSUPP,
			       N_("KDF not supported", ""));
	return KRB5_PROG_ETYPE_NOSUPP;
    }
    if (ai->parameters != NULL &&
	(ai->parameters->length != 2 ||
	 memcmp(ai->parameters->data, "\x05\x00", 2) != 0))
	{
	    krb5_set_error_message(context, KRB5_PROG_ETYPE_NOSUPP,
				   N_("kdf params not NULL or the NULL-type",
				      ""));
	    return KRB5_PROG_ETYPE_NOSUPP;
	}

    et = _krb5_find_enctype(enctype);
    if(et == NULL) {
	krb5_set_error_message(context, KRB5_PROG_ETYPE_NOSUPP,
			       N_("encryption type %d not supported", ""),
			       enctype);
	return KRB5_PROG_ETYPE_NOSUPP;
    }
    keylen = (et->keytype->bits + 7) / 8;

    keydata = malloc(keylen);
    if (keydata == NULL)
	return krb5_enomem(context);

    ret = encode_otherinfo(context, ai, client, server,
			   enctype, as_req, pk_as_rep, ticket, &other);
    if (ret) {
	free(keydata);
	return ret;
    }

    m = EVP_MD_CTX_create();
    if (m == NULL) {
	free(keydata);
	free(other.data);
	return krb5_enomem(context);
    }

    offset = 0;
    counter = 1;
    do {
	unsigned char cdata[4];

	EVP_DigestInit_ex(m, md, NULL);
	_krb5_put_int(cdata, counter, 4);
	EVP_DigestUpdate(m, cdata, 4);
	EVP_DigestUpdate(m, dhdata, dhsize);
	EVP_DigestUpdate(m, other.data, other.length);

	EVP_DigestFinal_ex(m, shaoutput, NULL);

	memcpy((unsigned char *)keydata + offset,
	       shaoutput,
	       min(keylen - offset, EVP_MD_CTX_size(m)));

	offset += EVP_MD_CTX_size(m);
	counter++;
    } while(offset < keylen);
    memset_s(shaoutput, sizeof(shaoutput), 0, sizeof(shaoutput));

    EVP_MD_CTX_destroy(m);
    free(other.data);

    ret = krb5_random_to_key(context, enctype, keydata, keylen, key);
    memset_s(keydata, sizeof(keylen), 0, sizeof(keylen));
    free(keydata);

    return ret;
}