package org.bouncycastle.mail.smime;

import java.io.IOException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.cert.X509Certificate;
import java.util.Iterator;

import javax.mail.MessagingException;
import javax.mail.internet.InternetHeaders;
import javax.mail.internet.MimeBodyPart;
import javax.mail.internet.MimeMessage;
import javax.mail.internet.MimeMultipart;

import org.bouncycastle.asn1.cms.AttributeTable;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.cms.CMSSignedData;
import org.bouncycastle.cms.CMSSignedDataGenerator;
import org.bouncycastle.cms.SignerInformation;
import org.bouncycastle.jce.cert.CertStore;
import org.bouncycastle.jce.cert.CertStoreException;
import org.bouncycastle.util.encoders.Base64;

/**
 * general class for generating a pkcs7-signature message.
 * <p>
 * A simple example of usage.
 *
 * <pre>
 *      CertStore           certs...
 *      SMIMESignedGenerator  fact = new SMIMESignedGenerator();
 *
 *      fact.addSigner(privKey, cert, SMIMESignedGenerator.DIGEST_SHA1);
 *      fact.addCertificatesAndCRLs(certs);
 *
 *      MimeMultipart       smime = fact.generate(content, "BC");
 * </pre>
 */
public class SMIMESignedGenerator
    extends SMIMEGenerator
{
    public static final String  DIGEST_SHA1 = "1.3.14.3.2.26";
    public static final String  DIGEST_MD5 = "1.2.840.113549.2.5";

    public static final String  ENCRYPTION_RSA = "1.2.840.113549.1.1.1";
    public static final String  ENCRYPTION_DSA = "1.2.840.10040.4.3";

    private CMSSignedDataGenerator      fact;

    /**
     * base constructor
     */
	public SMIMESignedGenerator()
    {
        fact = new CMSSignedDataGenerator();
    }

    /**
     * add a signer - no attributes other than the default ones will be
     * provided here.
     */
    public void addSigner(
        PrivateKey      key,
        X509Certificate cert,
        String          digestOID)
        throws IllegalArgumentException
    {
        fact.addSigner(key, cert, digestOID);
    }

    /**
     * add a signer with extra signed/unsigned attributes.
     */
    public void addSigner(
        PrivateKey      key,
        X509Certificate cert,
        String          digestOID,
        AttributeTable  signedAttr,
        AttributeTable  unsignedAttr)
        throws IllegalArgumentException
    {
        fact.addSigner(key, cert, digestOID, signedAttr, unsignedAttr);
    }

    /**
     * add the certificates and CRLs contained in the given CertStore
     * to the pool that will be included in the encoded signature block.
     * <p>
     * Note: this assumes the CertStore will support null in the get
     * methods.
     */
    public void addCertificatesAndCRLs(
        CertStore               certStore)
        throws CertStoreException, SMIMEException
    {
        try
        {
            fact.addCertificatesAndCRLs(certStore);
        }
        catch (CMSException e)
        {
            throw new SMIMEException(e.getMessage(), e.getUnderlyingException());
        }
    }

    /**
     * at this point we expect our body part to be well defined.
     */
    private MimeMultipart make(
        MimeBodyPart    content,
        String          sigProvider)
        throws NoSuchAlgorithmException, NoSuchProviderException, SMIMEException
    {
        StringBuffer        header = new StringBuffer(
                "signed; protocol=\"application/pkcs7-signature\"");
        int                 count = 0;
        CMSSignedData       signedData;

        try
        {
            signedData = fact.generate(
                            new CMSProcessableBodyPart(content), sigProvider);
        }
        catch (CMSException e)
        {
            throw new SMIMEException(e.getMessage(), e.getUnderlyingException());
        }

        //
        // build the hash header
        //
        Iterator it = signedData.getSignerInfos().getSigners().iterator();

        while (it.hasNext())
        {
            SignerInformation       signer = (SignerInformation)it.next();
            AlgorithmIdentifier     digAlgId, encAlgId;

            if (count == 0)
            {
                header.append("; micalg=");
            }
            else
            {
                header.append(",");
            }

            if (signer.getDigestAlgOID().equals(DIGEST_SHA1))
            {
                header.append("sha1");
            }
            else if (signer.getDigestAlgOID().equals(DIGEST_MD5))
            {
                header.append("md5");
            }
            else
            {
                header.append("unknown");
            }

            count++;
        }

        InternetHeaders sigHeader = new InternetHeaders();

        sigHeader.addHeader("Content-Type", "application/pkcs7-signature; name=smime.p7s; smime-type=signed-data");
        sigHeader.addHeader("Content-Disposition", "attachment; filename=\"smime.p7s\"");
        sigHeader.addHeader("Content-Description", "S/MIME Cryptographic Signature");

        try
        {
            MimeBodyPart    sig;
            
            if (useBase64)
            {
                sigHeader.addHeader("Content-Transfer-Encoding", "base64");

                sig = new MimeBodyPart(sigHeader, Base64.encode(signedData.getEncoded()));
            }
            else
            {
                sigHeader.addHeader("Content-Transfer-Encoding", encoding);

                sig = new MimeBodyPart(sigHeader, signedData.getEncoded());
            }

            MimeMultipart   mm = new MimeMultipart(header.toString());

            mm.addBodyPart(content);
            mm.addBodyPart(sig);

            return mm;
        }
        catch (IOException e)
        {
            throw new SMIMEException("exception encoding signature.", e);
        }
        catch (MessagingException e)
        {
            throw new SMIMEException("exception putting multi-part together.", e);
        }
    }

    /**
     * generate a signed object that contains an SMIME Signed Multipart
     * object using the given provider.
     */
    public MimeMultipart generate(
        MimeBodyPart    content,
        String          sigProvider)
        throws NoSuchAlgorithmException, NoSuchProviderException, SMIMEException
    {
        return make(makeContentBodyPart(content), sigProvider);
    }

    /**
     * generate a signed object that contains an SMIME Signed Multipart
     * object using the given provider from the given MimeMessage
     */
    public MimeMultipart generate(
        MimeMessage     message,
        String          sigProvider)
        throws NoSuchAlgorithmException, NoSuchProviderException, SMIMEException
    {
        try
        {
            message.saveChanges();      // make sure we're up to date.
        }
        catch (MessagingException e)
        {
            throw new SMIMEException("unable to save message", e);
        }
                        
        return make(makeContentBodyPart(message), sigProvider);
    }
}
