/*
 * Copyright (c) 2011 Jan Dittberner
 *
 * Permission is hereby granted, free of charge, to any person obtaining
 * a copy of this software and associated documentation files (the
 * "Software"), to deal in the Software without restriction, including
 * without limitation the rights to use, copy, modify, merge, publish,
 * distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to
 * the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */
package de.communardo.jdi.bcsmime_demo;

import java.security.KeyStore;
import java.security.cert.CertStore;
import java.security.cert.CertStoreParameters;
import java.security.cert.Certificate;
import java.security.cert.CollectionCertStoreParameters;
import java.security.cert.X509CertSelector;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Enumeration;
import java.util.List;

import javax.mail.Address;
import javax.mail.internet.InternetAddress;
import javax.mail.internet.MimeBodyPart;
import javax.mail.internet.MimeMessage;

import org.bouncycastle.asn1.x509.GeneralName;
import org.bouncycastle.cms.CMSAlgorithm;
import org.bouncycastle.cms.RecipientInfoGenerator;
import org.bouncycastle.cms.jcajce.JceCMSContentEncryptorBuilder;
import org.bouncycastle.cms.jcajce.JceKeyTransRecipientInfoGenerator;
import org.bouncycastle.mail.smime.SMIMEEnvelopedGenerator;
import org.bouncycastle.operator.OutputEncryptor;

/**
 * S/MIME encryption using the new BouncyCastle 1.46 APIs.
 *
 * @author Jan Dittberner
 */
public class SMIMEEncrypt {
    private final CertStore certs;

    /**
     * Creates a new SMIMEEncrypt instance.
     * 
     * @param keystore
     *            key store to use for recipient certificates
     */
    public SMIMEEncrypt(KeyStore keystore) {
        List<Certificate> certificates = new ArrayList<Certificate>();

        try {
            Enumeration<String> aliases = keystore.aliases();
            while (aliases.hasMoreElements()) {
                Certificate cert = keystore.getCertificate(aliases
                        .nextElement());
                if (cert != null) {
                    certificates.add(cert);
                }
            }

            CertStoreParameters params = new CollectionCertStoreParameters(
                    certificates);
            this.certs = CertStore.getInstance("Collection", params);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Encrypts a MimeMessage to all its recipients.
     *
     * @param message
     *            MIME message to encrypt
     * @return encrypted S/MIME message
     * @throws Exception
     *             if an error occurs
     */
    public MimeMessage encryptMessage(MimeMessage message) throws Exception {
        SMIMEEnvelopedGenerator smeg = new SMIMEEnvelopedGenerator();
        for (Address recipient : message.getAllRecipients()) {
            Collection<? extends Certificate> certificates = getCertificates((InternetAddress) recipient);
            for (Certificate cert : certificates) {
                RecipientInfoGenerator recipientInfoGen = new JceKeyTransRecipientInfoGenerator(
                        (X509Certificate) cert);
                smeg.addRecipientInfoGenerator(recipientInfoGen);
            }
        }
        OutputEncryptor encryptor = new JceCMSContentEncryptorBuilder(
                CMSAlgorithm.AES256_CBC).build();
        MimeBodyPart encryptedContent = smeg.generate(message, encryptor);
        MimeMessage result = new MimeMessage(message);
        result.setContent(encryptedContent.getContent(), encryptedContent
                .getContentType());
        result.saveChanges();
        return result;
    }

    /**
     * Helper method for getting certificates from a keystore.
     * 
     * @param recipient
     *            recipient address
     * @return X.509 certificate for recipient
     * @throws Exception
     *             if an error occurs
     */
    private Collection<? extends Certificate> getCertificates(
            InternetAddress recipient) throws Exception {
        X509CertSelector selector = new X509CertSelector();
        selector.setMatchAllSubjectAltNames(false);
        selector.addSubjectAlternativeName(GeneralName.rfc822Name, recipient
                .getAddress());

        return this.certs.getCertificates(selector);
    }
}