//////////////////////////////////////////////////////////////////////////// 
// 
// 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.cms;

import java.util.Set;
import java.util.HashSet;
import java.util.Iterator;

import java.security.SecureRandom;
import java.security.AlgorithmParameters;
import java.security.PrivateKey;
import java.security.cert.X509Certificate;
import java.security.spec.AlgorithmParameterSpec;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.SecretKeySpec;
import javax.crypto.spec.RC2ParameterSpec;
import javax.crypto.spec.IvParameterSpec;

import com.dstc.security.provider.OID;
import com.dstc.security.x509.AlgorithmId;
import com.dstc.security.cms.v1.ContentInfo;
import com.dstc.security.cms.v1.EnvelopedData;
import com.dstc.security.cms.v1.RecipientInfo;
import com.dstc.security.cms.v1.EncryptedContentInfo;
import com.dstc.security.cms.v1.IssuerAndSerialNumber;
import com.dstc.security.cms.v1.OriginatorInfo;
import com.dstc.security.asn1.Asn1Exception;

/**
 * <p>A class representing a CMS enveloped message, (used, eg. to
 * carry encrypted messages in S/MIME).
 *
 * <p>To encrypt a message according to CMS, use the message
 * and optionally the Originator to construct an EnvelopedMessage
 * object. Next, specify the encryption algorithm by calling
 * setEncryptionAlgorithm(). Next, call addRecipient() for each
 * Recipient of the message. Next, call encrypt() to encrypt the
 * message. Thereafter, the DER encoding for the ContentInfo-encapsulated
 * EnvelopedData structure for this encrypted message is obtained by
 * calling getEncoded().
 *
 * <p>To decrypt an encrypted message formatted as a ContentInfo-encapsulated
 * EnvelopedData structure accoring to CMS, first use its DER encoding
 * to construct an EnvelopedMessage object. Next pass the recipient
 * Private Key and associated X509Certificate in a a call to decrypt().
 * If this process succeeds, the decrypted message can be obtained by
 * calling getMessage().
 *
 * NB. This implementation currently supports only the Key Transport
 * method (ie. RSA encryption of session keys).
 *
 * @version 0.99, 99/01/26
 * @author Ming Yung
 *
 */
public class EnvelopedMessage
{
  private Set recipients;
  private Set recipientInfos;

  private String encryptionAlg;
  private String encryptionAlgMode;
  private String encryptionAlgModePad;
  private String encryptionAlgOid;
  private String paramSpecClassName;

  private Originator originator = null;
  private byte[] message = null;
  private SecretKey sessionKey = null;
  private SecureRandom rand = new SecureRandom(new byte[8]);;
  private Cipher cipher;
  private AlgorithmParameters params;
  private byte[] encryptedContent;
  private ContentInfo contentInfo;

  /**
   * Constructs an EnvelopedMessage from a SecureRandom
   */
  public EnvelopedMessage(SecureRandom rand)
  {
    this.recipients = new HashSet();
    this.recipientInfos = new HashSet();
    this.rand = rand;
  }

  /**
   * Constructs an EnvelopedMessage for an Originator and a message
   */
  public EnvelopedMessage(Originator originator, byte[] message,
                          SecureRandom rand)
  {
    this.recipients = new HashSet();
    this.originator = originator;
    this.message = message;
    this.rand = rand;
  }

  /**
   * Constructs an EnvelopedMessage from a DER encoding of a
   * ContentInfo with an EnvelopedData
   */
  public EnvelopedMessage(byte[] encoded) throws Asn1Exception, CMSException
  {
    this.contentInfo = new ContentInfo(encoded);
    EnvelopedData envData 
      = new EnvelopedData(contentInfo.getContent().encode());

    OriginatorInfo origInfo = envData.getOriginatorInfo();
    if (origInfo != null)
      this.originator = new Originator(origInfo);

    this.recipientInfos = envData.getRecipientInfos();

    this.recipients = new HashSet();
    Iterator it = this.recipientInfos.iterator();
    while(it.hasNext())
    {
      this.recipients.add(
        new Recipient((RecipientInfo)it.next()));
    }

    EncryptedContentInfo eci = envData.getEncryptedContentInfo();
    this.encryptedContent = eci.getEncryptedContent();

    String encryptionAlgOid 
      = eci.getContentEncryptionAlgorithm().getOid();
    this.params = eci.getContentEncryptionAlgorithm().getParams();

    if (encryptionAlgOid.equals(OID.rc2_cbc))
    {
      this.encryptionAlgModePad = "RC2/CBC/PKCS5Padding";
      this.encryptionAlgMode = "RC2/CBC";
      this.encryptionAlg = "RC2";
      this.paramSpecClassName = "javax.crypto.spec.RC2ParameterSpec";
    }
    else if (encryptionAlgOid.equals(OID.des_ede3_cbc))
    {
      this.encryptionAlgModePad = "DESede/CBC/PKCS5Padding";
      this.encryptionAlgMode = "DESede/CBC";
      this.encryptionAlg = "DESede";
      this.paramSpecClassName = "javax.crypto.spec.IvParameterSpec";
    }
  }

  /**
   * Sets the message to be enveloped
   */
  public void setMessage(byte[] message)
  {
    this.message = message;
  }

  /**
   * Returns the message in this EnvelopedMessage
   */
  public byte[] getMessage() throws CMSException
  {
    if (this.message == null)
      throw new CMSException("Message not yet decrypted");

    return this.message;
  }

  /**
   * Returns the Set of certificates for the Originator of this 
   * EnvelopedMessage
   */
  public Set getOriginatorCerts()
  {
    if (this.originator == null)
      return null;

    return this.originator.getCerts();
  }

  /**
   * Sets the encryption algorithm for this EnvelopedMessage
   *
   * (RC2/40 and DESede are supported as per S/MIME)
   */
  public void setEncryptionAlgorithm(String alg) throws CMSException
  {
    if (alg.equals("RC2/40"))
    {
      this.encryptionAlgModePad = "RC2/CBC/PKCS5Padding";
      this.encryptionAlgMode = "RC2/CBC";
      this.encryptionAlg = "RC2";
      this.encryptionAlgOid = OID.rc2_cbc;
      this.paramSpecClassName = "javax.crypto.spec.RC2ParameterSpec";
      this.sessionKey = this.generateSessionKey(40);
    }
    else if (alg.equals("DESede"))
    {
      this.encryptionAlgModePad = "DESede/CBC/PKCS5Padding";
      this.encryptionAlgMode = "DESede/CBC";
      this.encryptionAlg = "DESede";
      this.encryptionAlgOid = OID.des_ede3_cbc;
      this.paramSpecClassName = "javax.crypto.spec.IvParameterSpec";
      this.sessionKey = this.generateSessionKey(192);
    }
  }

  /**
   * Returns the encryption algorithm for this EnvelopedMessage
   */
  public String getCipherAlgorithm()
  {
    return this.encryptionAlgModePad;
  }

  /**
   * Adds a recipient to this EnvelopedMessage
   */
  public void addRecipient(Recipient recipient) throws CMSException
  {
    if (this.sessionKey == null)
      throw new CMSException("Encryption algorithm not yet set");

    this.recipients.add(recipient);
    this.recipientInfos.add(recipient.encryptKey(
      this.sessionKey.getEncoded()));
  }

  /**
   * Returns the Set of Recipients for this EnvelopedMessage
   */
  public Set getRecipients()
  {
    return this.recipients;
  }

  /**
   * Returns the CMS (PKCS#7) encoding for this EnvelopedMessage
   */
  public byte[] getEncoded()
  {
    return this.contentInfo.encode();
  }

  /**
   * Encrypts and finalizes this EnvelopedMessage
   */
  public void encrypt() throws CMSException
  {
    EnvelopedData envelopedData 
      = new EnvelopedData(0, null, this.recipientInfos, 
          this.encryptContent());
    this.contentInfo =
       new ContentInfo(OID.id_envelopedData, envelopedData);
  }

  /**
   * Decrypts this EnvelopedMessage for a Recipient with given
   * X509Certificate and PrivateKey
   */
  public void decrypt(X509Certificate cert, PrivateKey privKey)
    throws CMSException
  {
    Recipient recipient = this.getRecipient(cert);

    SecretKeySpec keySpec 
      = new SecretKeySpec(recipient.decryptKey(privKey), 
                          this.encryptionAlg);

    try
    {
      SecretKeyFactory keyFact 
        = SecretKeyFactory.getInstance(this.encryptionAlg);
      this.sessionKey = keyFact.translateKey(keySpec);

      this.cipher = Cipher.getInstance(getCipherAlgorithm());

      this.cipher.init(Cipher.DECRYPT_MODE, this.sessionKey,
        this.params.getParameterSpec(Class.forName(
          paramSpecClassName)), (SecureRandom)null);

      this.message = this.cipher.doFinal(this.encryptedContent);
    }
    catch (Exception e)
    {
      throw new CMSException(e.toString());
    }
  }

  /**
   * Returns the EncryptedContentInfo for this EnvelopedMessage
   */
  private EncryptedContentInfo encryptContent()
    throws CMSException
  {
    if (this.sessionKey == null)
      throw new CMSException("Encryption algorithm not yet set");

    if (this.message == null)
      throw new CMSException("Message not yet set");

    this.encryptedContent = this.encrypt(this.message);

    try
    {
      return new EncryptedContentInfo(OID.id_data, 
        new AlgorithmId(this.encryptionAlgOid, this.params), encryptedContent);
    }
    catch (Exception e)
    {
      throw new CMSException(e.toString());
    }
  }

  /**
   * Encrypts and returns the given data with the session key
   */
  private byte[] encrypt(byte[] data) throws CMSException
  {
    if (this.sessionKey == null)
      throw new CMSException("Encryption algorithm not yet set");

    try
    { 
      this.cipher = Cipher.getInstance(getCipherAlgorithm());

      byte iv[] = new byte[8];
      this.rand.nextBytes(iv);

      this.params 
        = AlgorithmParameters.getInstance(this.encryptionAlgMode);

      AlgorithmParameterSpec paramSpec = null;
      if (this.encryptionAlg.equals("RC2"))
      {
        paramSpec = new RC2ParameterSpec(128, iv);
      }
      else if (this.encryptionAlg.equals("DESede"))
      {
        paramSpec = new IvParameterSpec(iv);
      }

      this.params.init(paramSpec);

      this.cipher.init(Cipher.ENCRYPT_MODE, this.sessionKey, 
                       paramSpec, this.rand);
      return this.cipher.doFinal(data);
    }
    catch (Exception e)
    {
      throw new CMSException(e.toString());
    }
  }

  /**
   * Generates and returns a session key of given length
   * to encrypt this EnvelopedMessage
   */
  private SecretKey generateSessionKey(int keyLength) throws CMSException
  {
    if (this.encryptionAlg == null)
      throw new CMSException("Encryption algorithm not yet set");

    try
    {
      KeyGenerator keyGen 
        = KeyGenerator.getInstance(this.encryptionAlg);
      keyGen.init(keyLength, this.rand);
      return keyGen.generateKey();
    }
    catch (Exception e)
    {
      e.printStackTrace();
      throw new CMSException(e.toString());
    }
  }

  /**
   * Returns the Recipient corresponding to a given X509Certificate
   */
  private Recipient getRecipient(X509Certificate cert) 
    throws CMSException
  {
    IssuerAndSerialNumber iss = new IssuerAndSerialNumber(cert);

    Iterator it = this.getRecipients().iterator();
    while(it.hasNext())
    {
      Recipient rcpt = (Recipient)it.next();
      if(iss.equals(rcpt.getIssuerAndSerialNumber()))
        return rcpt;
    }
    throw new CMSException("Wrong certificate for this Enveloped message");
  }
}
