package org.bouncycastle.cms;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;

import org.bouncycastle.asn1.ASN1InputStream;
import org.bouncycastle.asn1.ASN1OutputStream;
import org.bouncycastle.asn1.ASN1Set;
import org.bouncycastle.asn1.cms.ContentInfo;
import org.bouncycastle.asn1.cms.EnvelopedData;
import org.bouncycastle.asn1.cms.KEKRecipientInfo;
import org.bouncycastle.asn1.cms.KeyTransRecipientInfo;
import org.bouncycastle.asn1.cms.RecipientInfo;

/**
 * containing class for an CMS Enveloped Data object
 */
public class CMSEnvelopedData
{
    RecipientInformationStore   recipientInfoStore;
    ContentInfo                 contentInfo;

    private static ContentInfo readContentInfo(
        InputStream envelopedData)
        throws CMSException
    {
        try
        {
            ASN1InputStream in = new ASN1InputStream(envelopedData);

            return ContentInfo.getInstance(in.readObject());
        }
        catch (IOException e)
        {
            throw new CMSException("IOException reading content.", e);
        }
    }

	public CMSEnvelopedData(
        byte[]    envelopedData) 
        throws CMSException
    {
        this(readContentInfo(new ByteArrayInputStream(envelopedData)));
    }

	public CMSEnvelopedData(
        InputStream    envelopedData) 
        throws CMSException
    {
        this(readContentInfo(envelopedData));
    }

	public CMSEnvelopedData(
        ContentInfo contentInfo)
        throws CMSException
    {
        this.contentInfo = contentInfo;

        EnvelopedData  envData = EnvelopedData.getInstance(contentInfo.getContent());

        //
        // load the RecepientInfoStore
        //
        ASN1Set     s = envData.getRecipientInfos();
        ArrayList   infos = new ArrayList();

        for (int i = 0; i != s.size(); i++)
        {
            RecipientInfo   info = RecipientInfo.getInstance(s.getObjectAt(i));

            if (info.getInfo() instanceof KeyTransRecipientInfo)
            {
                infos.add(new KeyTransRecipientInformation(
                            (KeyTransRecipientInfo)info.getInfo(), envData.getEncryptedContentInfo()));
            }
            else if (info.getInfo() instanceof KEKRecipientInfo)
            {
                infos.add(new KEKRecipientInformation(
                            (KEKRecipientInfo)info.getInfo(), envData.getEncryptedContentInfo()));
            }
        }

        this.recipientInfoStore = new RecipientInformationStore(infos);
    }

    /**
     * return a store of the intended recipients for this message
     */
    public RecipientInformationStore getRecipientInfos()
    {
        return recipientInfoStore;
    }

    /**
     * return the ASN.1 encoded representation of this object.
     */
    public byte[] getEncoded()
        throws IOException
    {
        ByteArrayOutputStream   bOut = new ByteArrayOutputStream();
        ASN1OutputStream        aOut = new ASN1OutputStream(bOut);

        aOut.writeObject(contentInfo);

        return bOut.toByteArray();
    }
}
