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

import java.util.Collection;
import java.util.Set;
import java.util.HashSet;
import java.net.*;
import java.util.Vector;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.BufferedInputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.InputStreamReader;
import java.io.BufferedReader;
import java.io.IOException;
import java.math.BigInteger;
import java.security.PublicKey;
import java.security.PrivateKey;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.Security;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.AlgorithmParameters;
import java.security.AlgorithmParameterGenerator;
import java.security.cert.X509Certificate;
import java.security.cert.CertificateFactory;
import java.security.spec.AlgorithmParameterSpec;
import java.security.interfaces.RSAPublicKey;
import java.security.interfaces.RSAPrivateCrtKey;

import javax.crypto.spec.DHParameterSpec;

import com.dstc.security.provider.OID;
import com.dstc.security.x509.Attribute;
import com.dstc.security.x509.X500Name;
import com.dstc.security.x509.X509CertImpl;
import com.dstc.security.x509.AlgorithmId;
import com.dstc.security.x509.SubjectPublicKeyInfo;
import com.dstc.security.pki.pkcs10.CertificationRequest;
import com.dstc.security.pki.pkcs10.CertificationRequestInfo;
import com.dstc.security.pkcs8.PKCS8EncryptedPrivateKey;
import com.dstc.security.asn1.Asn1;
import com.dstc.security.asn1.Asn1Exception;
import com.dstc.security.asn1.IA5String;
import com.dstc.security.asn1.Oid;

public class KeyGenerator
{
  private String keyAlg;
  private String sigAlg;
  private KeyPair kp;
  private SecureRandom ran;
  private String dn;
  private String email;
  private AlgorithmParameterSpec paramSpec;
  private int strength = 0;

  /**
   * Constructs a KeyGenerator from a SecureRandom
   */
  public KeyGenerator(SecureRandom ran)
  {
    this.ran = ran;
  }

  /**
   * Sets the user's distinguished name
   */
  public void setDistinguishedName(String dn)
  {
    this.dn = dn;
  }

  /**
   * Sets the user's email address
   */
  public void setEmailAddress(String email)
  {
    this.email = email;
  }

  /**
   * Sets the required Public Key algorithm
   */
  public void setKeyAlgorithm(String alg)
  {
    this.keyAlg = alg;

    if (alg.equals("RSA"))
    {
      this.sigAlg = OID.md5WithRSAEncryption;
    }
    else if (alg.equals("DSA"))
    {
      this.sigAlg = OID.dsaWithSHA;
    }
    else if (alg.equals("DH"))
    {
      this.sigAlg = OID.dhe;     //Diffie-Hellman Ephemeral POP alg
    }
  }

  /**
   * Sets the required key length in bits
   */
  public void setKeyLength(int len)
  {
    this.strength = len;
  }

  /**
   * Generates the key pair 
   */
  public void generateKeyPair() throws ToolException
  {
    if (keyAlg == null)
      throw new ToolException("Key algorithm not yet set");

    try
    {
      KeyPairGenerator kpg = KeyPairGenerator.getInstance(keyAlg);

      if(keyAlg.equalsIgnoreCase("DH") || keyAlg.equalsIgnoreCase("DSA"))
      {
if (paramSpec == null)
          throw new ToolException("Algorithm parameters unspecified");

        kpg.initialize(paramSpec);
      }
      else
      {
        if (strength == 0)
          throw new ToolException("Key length not yet set");

        kpg.initialize(strength, ran);
      }

      kp = kpg.generateKeyPair();
    }
    catch (Exception e)
    {
      e.printStackTrace();
      throw new ToolException(e.getMessage());
    }
  }

  /**
   * Returns the Public Key most recently generated
   */
  public PublicKey getPublic()
  {
    return kp.getPublic();
  }

  /**
   * Returns the Private Key most recently generated
   */
  public PrivateKey getPrivate()
  {
    return kp.getPrivate();
  }

  /**
   * Returns the Distinguished Name for the user
   */
  public String getDN()
  {
    return this.dn;
  }

  /**
   * Returns the email address for the user
   */
  public String getEmail()
  {
    return this.email;
  }

  /**
   * Returns the key algorithm
   */
  public String getKeyAlgorithm()
  {
    return this.keyAlg;
  }

  /**
   * Constructs a PKCS#10 Certification request for the generated public key
   */
  public CertificationRequest constructRequest() 
    throws Asn1Exception, ToolException, IOException
  {
    if (dn == null)
      throw new ToolException("Distinguished Name not yet set");

    if (email == null)
      throw new ToolException("Email address not yet set");

    if (kp == null)
      throw new ToolException("Key pair not yet generated");

    CertificationRequest retval;
  
    Set atts = new HashSet();
    atts.add(new Attribute(OID.emailAddress, new IA5String(email)));

    SubjectPublicKeyInfo spki 
      = new SubjectPublicKeyInfo(kp.getPublic().getEncoded());

    CertificationRequestInfo info 
      = new CertificationRequestInfo(new X500Name(dn), spki, atts);
  
    retval = new CertificationRequest(info, sigAlg, kp.getPrivate());
    return retval;
  }

  /**
   * Generates algorithm parameters
   */
  public void generateAlgParameters() throws ToolException
  {
    if (keyAlg == null)
      throw new ToolException("Key algorithm not yet set");

    if (!keyAlg.equals("DSA") && !keyAlg.equals("DH"))
      throw new ToolException("No parameters for selected algorithm");

    if (strength == 0)
      throw new ToolException("Key length not yet set");

    try 
    {
      AlgorithmParameters algParams;

      AlgorithmParameterGenerator paramGen = 
        AlgorithmParameterGenerator.getInstance(keyAlg);
      paramGen.init(strength, ran);
      algParams = paramGen.generateParameters();

      if (keyAlg.equals("DSA"))
      {
        paramSpec 
         = algParams.getParameterSpec(Class.forName(
             "java.security.spec.DSAParameterSpec"));
      }
      else if (keyAlg.equals("DH"))
      {
        paramSpec 
         = algParams.getParameterSpec(Class.forName(
             "javax.crypto.spec.DHParameterSpec"));
      }
    }
    catch (Exception e)
    {
      e.printStackTrace();
      throw new ToolException(e.getMessage());
    }
  }

  /**
   * Sets the algorithm parameters to be those of the subject public key
   * in the supplied certificate
   */
  public void setAlgParameters(X509Certificate cert) throws ToolException
  {
    if (keyAlg == null)
      throw new ToolException("Key algorithm not yet set");

    if (!keyAlg.equals("DSA") && !keyAlg.equals("DH"))
      throw new ToolException("No parameters for selected algorithm");

    try 
    {
      AlgorithmParameters algParams;

      algParams 
        = (new SubjectPublicKeyInfo(cert.getPublicKey().
              getEncoded())).getAlgorithmId().getParams();

      if (keyAlg.equals("DSA"))
      {
        paramSpec 
         = algParams.getParameterSpec(Class.forName(
             "java.security.spec.DSAParameterSpec"));
      }
      else if (keyAlg.equals("DH"))
      {
        paramSpec 
         = algParams.getParameterSpec(Class.forName(
             "javax.crypto.spec.DHParameterSpec"));
      }
    }
    catch(Exception e)
    {
      e.printStackTrace();
      throw new ToolException(e.getMessage());
    } 
  }

  /**
   * KeyGenerator-specific Exception class
   */
  protected class ToolException extends Exception
  {
    protected ToolException(String msg)
    {
      super(msg);
    }
  }
}
