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

import java.io.UnsupportedEncodingException;
import java.security.PrivateKey;
import java.security.AlgorithmParameters;
import java.security.spec.DSAParameterSpec;
import java.security.spec.InvalidKeySpecException;

import javax.crypto.Cipher;
import javax.crypto.spec.PBEParameterSpec;
import javax.crypto.spec.DHParameterSpec;

import com.dstc.security.provider.OID;
import com.dstc.security.x509.AlgorithmId;
import com.dstc.security.provider.PBEKey;
import com.dstc.security.provider.RSAPrivateCrtKey;
import com.dstc.security.provider.DSAPrivateKey;
import com.dstc.security.provider.DHPrivateKey;
import com.dstc.security.provider.PBEKeyDerivation;
import com.dstc.security.provider.PBECipher;
import com.dstc.security.pkcs12.PKCS12KeyDerivation;
import com.dstc.security.asn1.Asn1Exception;

/**
 * This class encrypts/decrypts a PrivateKey to/from PKCS#8 format.
 * Password-based encryption based on PKCS#5 and PKCS#12 are supported.
 *
 * @author Ming Yung, Eddy Cheung
 */ 
public class PKCS8EncryptedPrivateKey
{
  private EncryptedPrivateKeyInfo epki = null;
  private PrivateKey priv = null;

  public PKCS8EncryptedPrivateKey(PrivateKey priv)
  {
    this.priv = priv;
  }

  public PKCS8EncryptedPrivateKey(byte[] encoded) throws Asn1Exception
  {
    this.epki = new EncryptedPrivateKeyInfo(encoded);
  }

  public void encrypt(String password)
  {
    encrypt(password.getBytes(), new PKCS5KeyDerivation());
  }

  public void encrypt(byte[] passwordBytes)
  {
    encrypt(passwordBytes, new PKCS5KeyDerivation());
  }

  public void encrypt(byte[] passwordBytes, PBEKeyDerivation deriveStyle)
  {
    try
    {
      //Encrypt it with a PBECipher
      //Generate salt;
      byte preSalt[] = "Hello there!".getBytes ();
      byte salt[] = new byte[8];
      System.arraycopy(preSalt, 0, salt, 0, 8);
      PBEKey pbeKey = new PBEKey(passwordBytes);
      PBEParameterSpec paramsSpec = new PBEParameterSpec(salt, 5);
      AlgorithmParameters algParams = null;
      Cipher cipher = null;

      //XXX TODO, There should be a choice of algorithm within each method
      if (deriveStyle instanceof PKCS5KeyDerivation)
      {
        algParams = AlgorithmParameters.getInstance("PBEwithMD5andDES-CBC");
        cipher = Cipher.getInstance("PBEwithMD5andDES-CBC");
      }
      else if (deriveStyle instanceof PKCS12KeyDerivation)
      {
        algParams 
          = AlgorithmParameters.getInstance("PBEwithSHAand3-KeyTripleDES-CBC");
        cipher = Cipher.getInstance("PBEwithSHAand3-KeyTripleDES-CBC");
      }
      else
      {
        throw new Exception ("Unsupported PBEKeyDerivation style");
      }
      algParams.init(paramsSpec);
      cipher.init (Cipher.ENCRYPT_MODE, pbeKey, paramsSpec);

      byte pkiBytes[] = priv.getEncoded();

      byte encrypted[] = cipher.doFinal(pkiBytes);

      this.epki = new EncryptedPrivateKeyInfo(
                    new AlgorithmId(OID.pbeWithMD5AndDES_CBC,
                                    algParams), encrypted);
    }
    catch (Exception e)
    {
      e.printStackTrace();
    }
  }

  /**
     * Encrypt using the password and key derivation style
     */
  public void encrypt(String password, PBEKeyDerivation keyStyle)
  {
    if (keyStyle instanceof PKCS5KeyDerivation)
    {
      encrypt(password.getBytes(), keyStyle);
    }
    else if (keyStyle instanceof PKCS12KeyDerivation)
    {
      try
      {
        byte[] tmp = password.getBytes("UnicodeBigUnmarked");
        byte[] passwordBytes = new byte[tmp.length + 2];
        System.arraycopy(tmp, 0, passwordBytes, 0, tmp.length);
        passwordBytes[tmp.length] = 0x00;
        passwordBytes[tmp.length + 1] = 0x00;
        encrypt(passwordBytes, keyStyle);
      }
      catch (UnsupportedEncodingException err)
      {
        err.printStackTrace();
      }
    }
    else
    {
      throw new IllegalArgumentException("Unsupported PBEKeyDerivation");
    }
  }

  public void decrypt(String password) throws InvalidKeySpecException
  {
    decrypt(password.getBytes());
  }

  public void decrypt(byte[] passwordBytes) throws InvalidKeySpecException
  {
    try
    {
      //Get encrypted key and encryption algorithm
      byte encrypted[] = epki.getEncryptedKey();
      AlgorithmId algoId = epki.getEncryptionAlgorithm();

      //Get the Cipher object
      Cipher cipher = null;
      if (algoId.getOid().equals(OID.pbeWithMD5AndDES_CBC))
      {
        cipher = Cipher.getInstance("PBEwithMD5andDES-CBC");
      }
      else if ( algoId.getOid().equals(
            OID.pbeWithSHAAnd3_KeyTripleDES_CBC))
      {
        cipher = Cipher.getInstance("PBEwithSHAand3-KeyTripleDES-CBC");
      }
      else
      {
        throw new InvalidKeySpecException("Unknown encryption algorithm!!");
      }
      //Get encryption algorithm parameters
      PBEParameterSpec paramsSpec = (PBEParameterSpec)
          algoId.getParams().getParameterSpec(
            Class.forName("javax.crypto.spec.PBEParameterSpec"));

      PBEKey pbeKey = new PBEKey(passwordBytes);

      //Decrypt encoded private key
      cipher.init (Cipher.DECRYPT_MODE, pbeKey, paramsSpec, null);
      byte encodedKey[] = cipher.doFinal(encrypted);

      PrivateKeyInfo pki = new PrivateKeyInfo(encodedKey);

      //PrivateKeyInfo pki = new PrivateKeyInfo(epki.decrypt(passwordBytes));
      AlgorithmId algo = pki.getAlgorithmId();

      if (algo.getOid().equals(OID.rsaEncryption))
      {
        priv = new RSAPrivateCrtKey(pki.getEncodedKey());
      }
      else if (algo.getOid().equals(OID.dsa))
      {
        DSAParameterSpec dsaParams = (DSAParameterSpec)
            algo.getParams().getParameterSpec(
              Class.forName("java.security.spec.DSAParameterSpec"));

        priv = new DSAPrivateKey(pki.getEncodedKey(), dsaParams);
      }
      else if (algo.getOid().equals(OID.dhKeyAgreement))
      {
        DHParameterSpec dhParams = (DHParameterSpec)
            algo.getParams().getParameterSpec(
              Class.forName("javax.crypto.spec.DHParameterSpec"));

        priv = new DHPrivateKey(pki.getEncodedKey(), dhParams);
      }
    }
    catch (Exception e)
    {
      e.printStackTrace();
    }
  }

  /**
   * Deprecated - Please use decrypt directly
   */
  public void decrypt(String password,
                      PBEKeyDerivation keyStyle) throws InvalidKeySpecException
  {
    if (keyStyle instanceof PKCS5KeyDerivation)
    {
      decrypt(password.getBytes());
    }
    else if (keyStyle instanceof PKCS12KeyDerivation)
    {
      try
      {
        byte[] tmp = password.getBytes("UnicodeBigUnmarked");
        byte[] passwordBytes = new byte[tmp.length + 2];
        System.arraycopy(tmp, 0, passwordBytes, 0, tmp.length);
        passwordBytes[tmp.length] = 0x00;
        passwordBytes[tmp.length + 1] = 0x00;
        decrypt(passwordBytes);
      }
      catch (UnsupportedEncodingException err)
      {
        err.printStackTrace();
      }
    }
    else
    {
      throw new IllegalArgumentException("Unsupported PBEKeyDerivation");
    }
  }

  public byte[] getEncoded()
  {
    return this.epki.encode();
  }

  public PrivateKey getPrivateKey()
  {
    return this.priv;
  }
}
