//////////////////////////////////////////////////////////////////////////// 
// 
// Copyright (C) DSTC Pty Ltd (ACN 052 372 577) 1993, 1994, 1995.
// Unpublished work.  All Rights Reserved.
// 
// The software contained on this media is the property of the
// DSTC Pty Ltd.  Use of this software is strictly in accordance
// with the license agreement in the accompanying LICENSE.DOC 
// file. If your distribution of this software does not contain 
// a LICENSE.DOC file then you have no rights to use this 
// software in any manner and should contact DSTC at the address 
// below to determine an appropriate licensing arrangement.
// 
//      DSTC Pty Ltd
//      Level 7, GP South
//      University of Queensland
//      St Lucia, 4072
//      Australia
//      Tel: +61 7 3365 4310
//      Fax: +61 7 3365 4311
//      Email: jcsi@dstc.qut.edu.au
// 
// This software is being provided "AS IS" without warranty of
// any kind.  In no event shall DSTC Pty Ltd be liable for
// damage of any kind arising out of or in connection with
// the use or performance of this software.
// 
//////////////////////////////////////////////////////////////////////////// 

package com.dstc.security.smime;

import java.util.Set;

import java.io.*;
import java.util.Vector;
import java.util.Enumeration;
import java.security.SecureRandom;
import java.security.PrivateKey;
import java.security.cert.X509Certificate;

import javax.activation.DataHandler;

import javax.mail.*;
import javax.mail.internet.*;

import com.dstc.security.cms.EnvelopedMessage;
import com.dstc.security.cms.CMSException;
import com.dstc.security.cms.Recipient;
import com.dstc.security.provider.OID;

/**
 * <p>An extension of MimeBodyPart to handle an S/MIME defined mime body
 * part with mime type application/x-pkcs7-mime.
 *
 * Usage: <pre>
 *
 *    SMIMEPart smp = new SMIMEPart(mbp);
 *    smp.setEncryptionAlgorithm("DESede");
 *    smp.addRecipient(certificate);
 *    smp.encrypt();
 *    ...
 *    smp.decrypt(cert, privKey); </pre>
 *
 * @version 0.98, 98/11/01
 * @author Ming Yung
 * @author Daniel Bradley
 */
public class SMIMEPart extends MimeBodyPart
{
  private static final String CONTENT_TYPE = "application/x-pkcs7-mime";
  private static final String CONTENT_TRANSFER_ENCODING = "base64";
  private static final String CONTENT_DISPOSITION
      = "attachment; filename=\"smime.p7m\"";
  private static final String CONTENT_DESCRIPTION
      = "S/MIME Encrypted Message";
  
  private EnvelopedMessage envelopedMessage = null;
  private SecureRandom rand = null;

  public SMIMEPart() {}

  /**
   *  Constructs an SMIMEPart encapsulating a MimeMultipart.
   */
  public SMIMEPart(Multipart mp, SecureRandom rand)
  {
    try
    {
      this.rand = rand;

      MimeBodyPart mbp = new MimeBodyPart();
      mbp.setContent(mp, mp.getContentType());
      mbp.addHeader("Content-Type", mp.getContentType());
      mbp.addHeader("Content-Transfer-Encoding", "7bit");

      this.initialize(mbp);
    } 
    catch (Exception me) 
    {
      me.printStackTrace();
    }
  }
    
  /**
   * Constructs an SMIMEPart from a MimeBodyPart encapsulating the
   * plainText part (to be encrypted)
   */
  public SMIMEPart(MimeBodyPart part, SecureRandom rand) 
  {
    this.rand = rand;

    initialize(part);
  }

  /**
   *  Initializes the SMIMEPart to clone the passed MimeBodyPart.
   */
  public void initialize(MimeBodyPart part)
  {
    try
    {    
      this.envelopedMessage = new EnvelopedMessage(rand);

      setContent(part.getContent(), part.getContentType());

      setHeader("Content-Type", part.getContentType());
      setHeader("Content-Transfer-Encoding", part.getEncoding());
      setDisposition(part.getDisposition());
      setDescription(part.getDescription());
      if (part.getFileName() != null)
        setFileName(part.getFileName());
    } 
    catch (MessagingException me) 
    {
      me.printStackTrace();
    } 
    catch (IOException ioe) 
    {
      ioe.printStackTrace();
    } 
  }

  public void setContent(Object o, String type) throws MessagingException
  {
    if ((o instanceof EnvelopedMessage) && 
       (type.startsWith("application/x-pkcs7-mime")))
    {
      this.envelopedMessage = (EnvelopedMessage)o;
    }
    super.setContent(o, type);
  }

  /**
   * Sets the encryption algorithm to use 
   */
  public void setEncryptionAlgorithm(String alg) throws SMIMEException
  {
    if (alg == null)
      throw new SMIMEException("Algorithm not set");

    try
    {
      this.envelopedMessage.setEncryptionAlgorithm(alg);
    }
    catch (CMSException e)
    {
      e.printStackTrace();
      throw new SMIMEException(e.getMessage());
    }
  }

  /**
   * Adds a recipient with given X509Certificate 
   */
  public void addRecipient(X509Certificate cert) throws SMIMEException
  {
    try
    {
      this.envelopedMessage.addRecipient(new Recipient(cert));
    }
    catch (CMSException e)
    {
      throw new SMIMEException(e.getMessage());
    }
  }

  /**
   * Encrypts the message and finalizes this SMIMEPart
   */
  public void encrypt() throws SMIMEException
  {
    try
    {
      ByteArrayOutputStream baos = new ByteArrayOutputStream();
      super.writeTo(baos);
      this.envelopedMessage.setMessage(baos.toByteArray());
      this.envelopedMessage.encrypt();

      super.setContent(this.envelopedMessage, "application/x-pkcs7-mime");

      setHeader("Content-Type", CONTENT_TYPE);
      setHeader("Content-Transfer-Encoding", CONTENT_TRANSFER_ENCODING);
      setHeader("Content-Disposition", CONTENT_DISPOSITION);
      setHeader("Content-Description", CONTENT_DESCRIPTION);
    }  
    catch (Exception e)
    {
      e.printStackTrace();
      throw new SMIMEException(e.toString());
    }
  }

  /**
   * Decrypts this EnvelopedMessage for a Recipient with supplied
   * X509Certificate and PrivateKey
   */
  public void decrypt(X509Certificate cert, PrivateKey privKey) 
    throws SMIMEException
  {
    try
    {
      this.envelopedMessage.decrypt(cert, privKey);
    } 
    catch (Exception e) 
    {
      e.printStackTrace();
      throw new SMIMEException(e.getMessage());
    }
  }

  /**
   * Returns the Set of certificates for the Originator of this message
   */
  public Set getOriginatorCerts()
  {
    return this.envelopedMessage.getOriginatorCerts();
  }

  /**
   * Returns the decrypted message
   */
  public byte[] getDecrypted() throws SMIMEException
  {
    try
    {
      return this.envelopedMessage.getMessage();
    }
    catch (CMSException e)
    {
      throw new SMIMEException(e.getMessage());
    }
  }
}
