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

/**
 * <p>A class supporting key pair generation and PKCS#10 requests.
 *
 * <p>A pseudo-interactive key tool is provided, enabling generation 
 * of key pairs for RSA, DSA and Diffie-Hellman keys, and construction 
 * of PKCS#10 Certification requests for the generated public key.
 * The PKCS#10 request can be saved to a file to be communicated to a
 * CA, or sent directly to the online CA supplied with JCSI.
 *
 * <p>To use the tool, simply do
 * <pre>
 *      java com.dstc.security.pki.KeyTool
 * </pre>
 *
 * @version 0.98, 98/11/01
 * @author Ming Yung, Nishanth Chandran
 */
public class ConsoleKeyTool extends KeyGenerator
{
  private BufferedReader inr;
  private String resp;

  public static void main(String args[])
  {
    try
    {
      InputStreamReader isr = new InputStreamReader(System.in);
      BufferedReader in = new BufferedReader(isr);
      String resp;
  
      System.out.println("Seeding random number generator. Please wait...");
      SecureRandom ran = new SecureRandom();
      ran.nextInt();

      ConsoleKeyTool keyTool = new ConsoleKeyTool(ran);

      while (true)
      {
        System.out.print("\nGenerate (K)eyPair or (Q)uit? ");
        resp = in.readLine();
        if (resp.equalsIgnoreCase("Q"))
        {
          break;
        }
        else if (resp.equalsIgnoreCase("K"))
        {
          keyTool.setKeyAlgorithm();
  if (keyTool.getKeyAlgorithm().equals("RSA"))
            keyTool.setKeyLength();
  else
            keyTool.setAlgParameters();

          keyTool.generateKeyPair();

          PrivateKey privKey = keyTool.getPrivate();
  PublicKey subjectPub = keyTool.getPublic();
 
          System.out.print("\nFile to save private key? ");
          FileOutputStream keyFileOs = new FileOutputStream(in.readLine());

          PKCS8EncryptedPrivateKey pkcs8 
            = new PKCS8EncryptedPrivateKey(privKey);
          System.out.print("\nEnter password to lock private key? ");
          pkcs8.encrypt(in.readLine());
          keyFileOs.write(pkcs8.getEncoded());
          keyFileOs.close();

          System.out.print("\n(G)enerate PKCS#10 request or (Q)uit? ");
          resp = in.readLine();
          if (resp.equalsIgnoreCase("Q"))
          {
            break;
          }
          else if (resp.equalsIgnoreCase("G"))
          {
            keyTool.setDistinguishedName();
  
            keyTool.setEmailAddress();

            CertificationRequest req = keyTool.constructRequest();

    System.out.print("\nPKCS#10 Request to (F)ile or (S)end to CA: ");
            resp = in.readLine();
            if (resp.equalsIgnoreCase("F"))
            {            
      System.out.print("\nFile to save PKCS#10 request? ");
      keyFileOs = new FileOutputStream(in.readLine());
      keyFileOs.write(req.encode());
      keyFileOs.close();
            }
            else if (resp.equalsIgnoreCase("S"))
            {
      byte[] certBytes = keyTool.sendRequest(req.encode());
        
      CertificateFactory certFact = 
        CertificateFactory.getInstance("X509");

      ByteArrayInputStream bais = new ByteArrayInputStream(certBytes);
      X509Certificate cert = 
        (X509Certificate)certFact.generateCertificate(bais);

      System.out.println("\nThe returned certificate's details: ");
              System.out.println("######################################");
              System.out.println(cert.toString());
              System.out.println("######################################");

              System.out.print("\nFile to save certificate? ");
              keyFileOs = new FileOutputStream(in.readLine());
              keyFileOs.write(cert.getEncoded());
              keyFileOs.close();
    }
          }
        }
      }
    }
    catch (Exception e)
    {
      e.printStackTrace();
    }
  }

  /**
   * Constructs a ConsoleKeyTool from a SecureRandom
   */
  public ConsoleKeyTool(SecureRandom ran)
  {
    super(ran);

    InputStreamReader isr = new InputStreamReader(System.in);
    inr = new BufferedReader(isr);
  }

  /**
   * Prompts the user to enter her distinguished name
   */
  protected void setDistinguishedName() throws IOException
  {
    System.err.println("Enter Subject Distinguished Name: ");
    StringBuffer sb = new StringBuffer();

    System.err.print("Common Name: ");
    sb.append("2.5.4.3=");
    sb.append(inr.readLine());

    System.err.print("Organizational Unit: ");
    sb.append(", 2.5.4.11=");
    sb.append(inr.readLine());

    System.err.print("Organization: ");
    sb.append(", 2.5.4.10=");
    sb.append(inr.readLine());

    System.err.print("Country: ");
    sb.append(", 2.5.4.6=");
    sb.append(inr.readLine());

    setDistinguishedName(sb.toString());
  }

  /**
   * Prompts the user to enter her email address
   */
  protected void setEmailAddress() throws IOException
  {
    System.err.print("Email: ");
    setEmailAddress(inr.readLine());
  }

  /**
   * Prompts the user to enter the required Public Key algorithm
   */
  protected void setKeyAlgorithm() throws IOException
  {
    System.out.print("\nKey Algorithm: (R)SA, (D)SA " +
                     "or Diffie-(H)ellman? ");
    resp = inr.readLine();
 
    if (resp.equalsIgnoreCase("R"))
    {
      setKeyAlgorithm("RSA");
    }
    else if (resp.equalsIgnoreCase("D"))
    {
      setKeyAlgorithm("DSA");
    }
    else if (resp.equalsIgnoreCase("H"))
    {
      setKeyAlgorithm("DH");
    }
  }

  /**
   * Prompts the user to enter the required key length
   */
  protected void setKeyLength() throws IOException
  {
    System.out.print("\nKey length: (1) 1024, (2) 768 " +
                   "or (3) 512? ");
    resp = inr.readLine();

    if (resp.equalsIgnoreCase("1"))
    {
      setKeyLength(1024);
    }
    else if (resp.equalsIgnoreCase("2"))
    {
      setKeyLength(768);
    }
    if (resp.equalsIgnoreCase("3"))
    {
      setKeyLength(512);
    }
  }

  /**
   * Sends the supplied DER-encoding of the PKCS#10 request to the
   * JCSI online CA and returns with the DER-encoding of an 
   * X.509Certificate
   *
   * <p>The URL must be of the form
   * <pre>
   *       http://path.to.host:port/servlet/pkcs10
   * </pre>
   * where "path.to.host" and "port" depends on where the JCSI online
   * CA is listening
   */
  protected byte[] sendRequest(byte[] encoded) throws ToolException
  {
    try
    {
      HttpURLConnection connection;
      byte[] data=null;
  
      System.out.print("\nEnter URL for online CA: ");
  
      URL url = new URL(inr.readLine());
      connection = (HttpURLConnection) url.openConnection();
      connection.setDoOutput(true);
      connection.setDoInput(true);
      connection.setRequestMethod("POST");
      OutputStream toServer = connection.getOutputStream();
  
      toServer.write(encoded);
      toServer.flush();
      toServer.close();

      System.out.println("Request sent. Waiting for reply...");
  
      if (connection.getResponseCode() == HttpURLConnection.HTTP_OK)
      {
        InputStream is = connection.getInputStream();
        while (is.available() == 0) ;
        data = new byte[is.available()];
        int datalen=is.read(data);
        is.close();
      }    

      return data;
    }
    catch (Exception e)
    {
      e.printStackTrace();
      throw new ToolException(e.getMessage());
    }
  }

  /**
   * Prompts the user to either supply parameters through an X.509Certificate
   * or have them generated
   */
  protected void setAlgParameters() throws ToolException
  {
    if (getKeyAlgorithm().equals("RSA"))
      return;

    try
    {
      while(true)
      {
        System.out.print("\n(G)enerate or (P)rovide Algorithm parameters? ");
        resp = inr.readLine();
      
        if(resp.equalsIgnoreCase("G")) 
        {
  setKeyLength();
          System.out.println("Generating parameters. Please wait...");
  generateAlgParameters();
          break;
        }
        else if(resp.equalsIgnoreCase("P")) 
        {
          System.out.print("\nCertificate containing required parameters? ");
          FileInputStream fis = new FileInputStream(inr.readLine());
          byte[] encoded = new byte[fis.available()];
          fis.read(encoded);
          fis.close();
    
          CertificateFactory certFact =
            CertificateFactory.getInstance("X509");
    
          ByteArrayInputStream bais = new ByteArrayInputStream(encoded);
          X509Certificate cert = 
            (X509Certificate)certFact.generateCertificate(bais);

  setAlgParameters(cert);
          break;
        }
      }
    }
    catch(Exception e)
    {
      e.printStackTrace();
      throw new ToolException(e.getMessage());
    } 
  }
}
