package org.bouncycastle.cms;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.security.AlgorithmParameters;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.bouncycastle.asn1.ASN1Null;
import org.bouncycastle.asn1.ASN1OctetString;
import org.bouncycastle.asn1.ASN1OutputStream;
import org.bouncycastle.asn1.DEREncodable;
import org.bouncycastle.asn1.DERNull;
import org.bouncycastle.asn1.cms.EncryptedContentInfo;
import org.bouncycastle.asn1.cms.IssuerAndSerialNumber;
import org.bouncycastle.asn1.cms.KeyTransRecipientInfo;
import org.bouncycastle.asn1.cms.RecipientIdentifier;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;


/**
 * the RecipientInfo class for a recipient who has been sent a secret
 * key encrypted using their public key that needs to be used to
 * extract the message.
 */
public class KeyTransRecipientInformation
    extends RecipientInformation
{
    private KeyTransRecipientInfo   info;
    private EncryptedContentInfo    data;
    private static ASN1Null         asn1Null = new DERNull();

    public KeyTransRecipientInformation(
        KeyTransRecipientInfo   info,
        EncryptedContentInfo    data)
    {
        this.info = info;
        this.rid = new RecipientId();
        this.data = data;

        RecipientIdentifier r = info.getRecipientIdentifier();

        try
        {
            if (r.isTagged())
            {
                ASN1OctetString octs = ASN1OctetString.getInstance(r.getId());

                rid.setSubjectKeyIdentifier(octs.getOctets());
            }
            else
            {
                IssuerAndSerialNumber   iAnds = IssuerAndSerialNumber.getInstance(r.getId());

                ByteArrayOutputStream   bOut = new ByteArrayOutputStream();
                ASN1OutputStream        aOut = new ASN1OutputStream(bOut);

                aOut.writeObject(iAnds.getName());

                rid.setIssuer(bOut.toByteArray());
                rid.setSerialNumber(iAnds.getSerialNumber().getValue());
            }
        }
        catch (IOException e)
        {
            throw new IllegalArgumentException("invalid rid in KeyTransRecipientInformation");
        }
    }

    /**
     * decrypt the content and return it as a byte array.
     */
	public byte[] getContent(
        Key      key,
        String   prov)
        throws CMSException, NoSuchProviderException
    {
        try
        {
            byte[]              encryptedKey = info.getEncryptedKey().getOctets();
            AlgorithmIdentifier keyEncAlg = AlgorithmIdentifier.getInstance(info.getKeyEncryptionAlgorithm());

            Cipher  keyCipher = Cipher.getInstance(keyEncAlg.getObjectId().getId(), prov);

            keyCipher.init(Cipher.DECRYPT_MODE, key);

            byte[]              keyBytes = keyCipher.doFinal(encryptedKey);
            byte[]              enc = data.getEncryptedContent().getOctets();
            AlgorithmIdentifier aid = data.getContentEncryptionAlgorithm();
            String              alg = aid.getObjectId().getId();
            SecretKey           sKey = new SecretKeySpec(keyBytes, alg);
            Cipher              cipher = Cipher.getInstance(alg, prov);
            DEREncodable        sParams = aid.getParameters();

            if (sParams != null && !asn1Null.equals(sParams))
            {
                AlgorithmParameters     params = AlgorithmParameters.getInstance(alg, prov);

                ByteArrayOutputStream   bOut = new ByteArrayOutputStream();
                ASN1OutputStream        aOut = new ASN1OutputStream(bOut);

                aOut.writeObject(aid.getParameters());

                params.init(bOut.toByteArray(), "ASN.1");

                cipher.init(Cipher.DECRYPT_MODE, sKey, params);
            }
            else
            {
                if (alg.equals(CMSEnvelopedDataGenerator.DES_EDE3_CBC)
                    || alg.equals(CMSEnvelopedDataGenerator.IDEA_CBC)
                    || alg.equals(CMSEnvelopedDataGenerator.CAST5_CBC))
                {
                    cipher.init(Cipher.DECRYPT_MODE, sKey, new IvParameterSpec(new byte[8]));
                }
                else
                {
                    cipher.init(Cipher.DECRYPT_MODE, sKey);
                }
            }

            return cipher.doFinal(enc);
        }
        catch (NoSuchAlgorithmException e)
        {
            throw new CMSException("can't find algorithm.", e);
        }
        catch (IllegalBlockSizeException e)
        {
            throw new CMSException("illegal blocksize in message.", e);
        }
        catch (InvalidKeyException e)
        {
            throw new CMSException("key invalid in message.", e);
        }
        catch (NoSuchPaddingException e)
        {
            throw new CMSException("required padding not supported.", e);
        }
        catch (BadPaddingException e)
        {
            throw new CMSException("bad padding in message.", e);
        }
        catch (InvalidAlgorithmParameterException e)
        {
            throw new CMSException("algorithm parameters invalid.", e);
        }
        catch (IOException e)
        {
            throw new CMSException("error decoding algorithm parameters.", e);
        }
    }
}
