package org.bouncycastle.crypto.agreement.kdf;

import java.io.IOException;

import org.bouncycastle.asn1.ASN1Encoding;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.ASN1OctetString;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.DERNull;
import org.bouncycastle.asn1.DEROctetString;
import org.bouncycastle.asn1.DERSequence;
import org.bouncycastle.asn1.DERTaggedObject;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.crypto.DataLengthException;
import org.bouncycastle.crypto.DerivationParameters;
import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.DigestDerivationFunction;
import org.bouncycastle.crypto.generators.KDF2BytesGenerator;
import org.bouncycastle.crypto.params.KDFParameters;
import org.bouncycastle.util.Exceptions;
import org.bouncycastle.util.Pack;

/**
 * X9.63 based key derivation function for ECDH CMS.
 */
public class ECDHKEKGenerator
    implements DigestDerivationFunction
{
    private DigestDerivationFunction kdf;

    private ASN1ObjectIdentifier algorithm;
    private int                 keySize;
    private byte[]              z;

    public ECDHKEKGenerator(
        Digest digest)
    {
        this.kdf = new KDF2BytesGenerator(digest);
    }

    public void init(DerivationParameters param)
    {
        DHKDFParameters params = (DHKDFParameters)param;

        this.algorithm = params.getAlgorithm();
        this.keySize = params.getKeySize();
        this.z = params.getZ();
    }

    public Digest getDigest()
    {
        return kdf.getDigest();
    }

    public int generateBytes(byte[] out, int outOff, int len)
        throws DataLengthException, IllegalArgumentException
    {
        if (outOff + len > out.length)
        {
            throw new DataLengthException("output buffer too small");
        }

        AlgorithmIdentifier keyInfo = new AlgorithmIdentifier(algorithm, DERNull.INSTANCE);
        ASN1OctetString suppPubInfo = DEROctetString.withContents(Pack.intToBigEndian(keySize));

        // TODO org.bouncycastle.asn1.cms.ecc.ECCCMSSharedInfo exists, but is located in the 'util' jar.
        // TODO Should the optional DHKDFParameters.getExtraInfo be used for ECCCMSSharedInfo.entityUInfo?
        ASN1Sequence eccCMSSharedInfo = new DERSequence(keyInfo, new DERTaggedObject(2, suppPubInfo));

        try
        {
            byte[] iv = eccCMSSharedInfo.getEncoded(ASN1Encoding.DER);

            kdf.init(new KDFParameters(z, iv));
        }
        catch (IOException e)
        {
            throw Exceptions.illegalArgumentException("unable to initialise kdf", e);
        }

        return kdf.generateBytes(out, outOff, len);
    }
}
