//////////////////////////////////////////////////////////////////////////// 
// 
// 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.Arrays;
import java.util.Set;
import java.util.HashSet;
import java.util.Iterator;

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

import javax.mail.*;
import javax.mail.internet.*;
import javax.activation.*;
import java.awt.datatransfer.*;

import com.dstc.security.provider.OID;
import com.dstc.security.cms.Signer;
import com.dstc.security.cms.SignedMessage;
import com.dstc.security.x509.Attribute;
import com.dstc.security.x509.AlgorithmId;
import com.dstc.security.asn1.Asn1Exception;

/**
 * <p>Implementation of a Part for a multipart/signed MIME message.
 *
 * <p>For a receiver of a multipart/signed MIME message, this class
 *    is transparently instantiated using the JavaMail framework (once
 *    the mailcap is properly set up). To get a MimeMultipartSigned
 *    object, first create a MimeMessage object from an InputStream 
 *    encapsulating the MIME message, and then call getContent() on it,
 *    with a cast to MimeMultipartSigned. The signed message can then
 *    be verified with verify() after calling setTrustedCAs() to 
 *    specify the Set of CA certificates to trust.
 *
 * <p>For a sender of a multipart/signed MIME message, a MimeBodyPart
 *    object is first constructed from the headers and message text. This
 *    is used to construct a MimeMultipartSigned object. Next, the digest
 *    and signature algorithms are specified by calling setDigestAlgorithm()
 *    and setSignatureAlgorithm(). The signing Private Key and associated
 *    X509Certificate are then set by calling setKeys(). Any extra 
 *    certificates to faciliate cert path processing for verifiers can be 
 *    added by calling setCertificates(). Next, a call to sign() will
 *    sign the message (resulting in a CMS SignedData which is base-64
 *    encoded and included as the second part of this (two-part)
 *    multipart/signed message) 
 *
 * <p>NB. Only text messages are supported currently.
 * <p>NB. Signing for only single signer is supported.
 *        Verification for multiple signers is supported.
 *
 * @version 0.98, 98/11/01
 * @author Ming Yung
 */
public class MimeMultipartSigned extends MimeMultipart
{
  private byte[] msg;
  private MimeBodyPart firstPart;
  private MimeBodyPart sigPart;
  private String digestAlgoName;
  private String sigAlgoName;
  private PrivateKey privateKey;
  private Set certificates;
  private X509Certificate signerCert;
  private SignedMessage signedMessage;

  /**
   * Constructs a MimeMultipartSigned from a DataSource
   */
  public MimeMultipartSigned(DataSource ds) 
    throws IOException, MessagingException
  {
    super(ds);

    this.firstPart = (MimeBodyPart)this.getBodyPart(0);

    this.sigPart = (MimeBodyPart)this.getBodyPart(1);

    if (!this.sigPart.isMimeType("application/x-pkcs7-signature"))
      throw new MessagingException("Bad mime part");

    this.signedMessage 
      = (SignedMessage)sigPart.getDataHandler().getContent();
  }

  /**
   * Constructs a MimeMultipartSigned from a MimeBodyPart
   * (the part to be signed)
   */
  public MimeMultipartSigned(MimeBodyPart msgPart)
  {
    super("signed; protocol=\"application/x-pkcs7-signature\"");

    try
    {
      this.firstPart = msgPart;
      addBodyPart(msgPart);
    }
    catch (Exception e)
    {
      e.printStackTrace();
    }
  }

  /**
   *  Constructs an MimeMultipartSigned from a MimeMultipart.
   */
  public MimeMultipartSigned(Multipart mp)
  {
    super("signed; protocol=\"application/x-pkcs7-signature\"");

    try
    {
      this.firstPart = new MimeBodyPart();
      this.firstPart.setContent(mp, mp.getContentType());
      this.firstPart.addHeader("Content-Type", mp.getContentType());

      addBodyPart(this.firstPart);
    }
    catch (MessagingException me)
    {
      me.printStackTrace();
    }
  }

  /**
   * Sets the signature algorithm to be used
   */
  public void setSignatureAlgorithm(String name)
  {
    this.sigAlgoName = name;
  }

  /**
   * Sets the digest algorithm to use. "SHA-1" and "MD5" are supported.
   */
  public void setDigestAlgorithm(String name)
  {
    this.digestAlgoName = name;

    try
    {
      if (name.equals("SHA-1"))
        setSubType("signed; micalg=sha1");
      else if (name.equals("MD5"))
        setSubType("signed; micalg=md5");
    }
    catch (MessagingException e)
    {
      e.printStackTrace();
    }
  }

  /**
   * Sets the private key and corresponding certificate for signing
   */
  public void setKeys(PrivateKey key, X509Certificate cert)
  {
    this.privateKey = key;
    this.signerCert = cert;
  }

  /**
   * Sets the X509Certificate chain (as a Set of X509Certificate)
   * for the SignedData.
   */
  public void setCertificates(Set certs) 
    throws Asn1Exception, CertificateException
  {
    this.certificates = certs;
  }

  /**
   * Sets the trusted CAs
   */
  public void setTrustedCAs(Set certs)
  {
    this.signedMessage.setTrustedCAs(certs);
  }

  /**
   * Signs the message and finalizes this MimeMultipartSigned 
   */
  public void sign() throws SMIMEException
  {
    if (this.sigAlgoName == null)
      throw new SMIMEException("Signature algorithm not yet set");

    if (this.digestAlgoName == null)
      throw new SMIMEException("Digest algorithm not yet set");

    if (this.privateKey == null)
      throw new SMIMEException("Private Key not yet set");

    try
    {
      TextCanonicalizer can = new TextCanonicalizer(this.firstPart);
      this.msg = can.getDigest();

      Set atts = new HashSet();
      Attribute att1 
        = new com.dstc.security.cms.atts.ContentType(OID.id_data);
      atts.add(att1);
 
      Signer signer 
        = new Signer(this.signerCert, this.certificates, this.privateKey);
      signer.setDigestAlgorithm(this.digestAlgoName);
      signer.setSignatureAlgorithm(this.sigAlgoName);
      signer.setSignedAttributes(atts);
  
      this.signedMessage = new SignedMessage(signer, this.msg);

      MimeBodyPart signedPart = new MimeBodyPart();
      signedPart.setContent(this.signedMessage, 
        "application/x-pkcs7-signature");

      signedPart.addHeader("Content-Type", 
                     "application/x-pkcs7-signature; name=\"smime.p7s\"");
      signedPart.addHeader("Content-Transfer-Encoding", "base64");
      signedPart.addHeader("Content-Disposition", 
                           "attachment; filename=\"smime.p7s\"");
      signedPart.addHeader("Content-Description",
                           "S/MIME Cryptographic Signature");

      addBodyPart(signedPart);
    }
    catch (Exception e)
    {
      e.printStackTrace();
    }
  }
  
  /**
   * Verifies this MimeMultipartSigned. This first calculates the message 
   * digest for the first part (the part to be signed). This message digest
   * is checked for each SignerInfo in the second part. The signature in
   * each SignerInfo is also checked and the signer certificate is checked
   * against the Set of trusted certs.
   */
  public void verify() throws SMIMEException
  {
    try
    {
      String digestAlgName = null;

      StringTokenizer tknz = new StringTokenizer(this.getContentType(), ";");
      while (tknz.hasMoreTokens())
      {
        String token = tknz.nextToken().trim();
        if (token.startsWith("micalg"))
        {
          if (token.equals("micalg=sha1"))
            digestAlgName = "SHA-1";
          else if (token.equals("micalg=md5"))
            digestAlgName = "MD5";
          else
            throw new SMIMEException("Unknown micalg");

          break;
        }
      }

      TextCanonicalizer can = new TextCanonicalizer(this.firstPart);
      this.msg = can.getDigest();

      Set atts = new HashSet();
      this.signedMessage.verify(this.msg);
    }
    catch (Exception e)
    {
      throw new SMIMEException(e.getMessage());
    }
  }

  /**
   * Returns the content of this MimeMultipartSigned (the part that is signed)
   */
  public String getMessage() throws IOException, MessagingException
  {
    return (String)this.firstPart.getContent();
  }

  /**
   * Returns the Set of signer X509Certificate for this 
   * MimeMultipartSigned
   */
  public Set getSignerCertificates()
  {
    Set retval = new HashSet();

    Iterator signers =  this.signedMessage.getSigners().iterator();
    while (signers.hasNext())
    {
      retval.add(((Signer)signers.next()).getCertPath().getSignerCert());
    }
    return retval;
  }

  protected static void debug(String msg, byte bytes[])
  {
    System.out.print(msg);
    for (int i=0; i<bytes.length; i++)
    {
      System.out.print(java.lang.Integer.toHexString(bytes[i] & 0xff) + " ");
    }
    System.out.println(" ");
  }
}
