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

//From JDK
import java.lang.Math;
import java.math.BigInteger;
import java.io.ByteArrayOutputStream;
import java.security.Key;
import java.security.MessageDigest;
import java.security.spec.AlgorithmParameterSpec;
import java.security.InvalidKeyException;
import java.security.InvalidAlgorithmParameterException;
//From Javax
import javax.crypto.spec.PBEParameterSpec;
//From JCSI
import com.dstc.security.provider.PBEKeyDerivation;
import com.dstc.security.provider.PBEKey;
import com.dstc.security.provider.PBEParameters;

/**
 * <p>An Implementation of KeyDerivation using PKCS#12 described method
 *
 * @see com.dstc.security.provider.KeyDerivation
 * @see com.dstc.security.PKCS5.KeyDerivation
 *
 * @version 0.98, 99/05/26
 * @author Eddy Cheung
 */

/* Log Message 
 * $Log: PKCS12KeyDerivation.java,v $
 * Revision 1.2  1999/07/05 23:59:46  cheung
 * Added/Fixed some PKCS#12 doc. info
 *
 */

public class PKCS12KeyDerivation implements PBEKeyDerivation
{
  public byte[] generateKey(int opmode, MessageDigest md, 
                            Key key, AlgorithmParameterSpec params, 
                            int keyLen)
    throws InvalidKeyException, InvalidAlgorithmParameterException
  {
    //Assertion
    if (opmode < GENERATE_KEY_MODE || opmode > GENERATE_INTEGRITY_KEY_MODE)
      throw new IllegalArgumentException("Unrecognise or unsupported Mode");

    if (!(key instanceof PBEKey)) 
    {
      throw new InvalidKeyException ("Not a PBE Key");
    }

    if (!(params instanceof PBEParameterSpec)) {
      throw new InvalidAlgorithmParameterException("Not PBE parameters");
    }

    byte password[] = ((PBEKey)key).getEncoded ();
    byte salt[] = ((PBEParameterSpec)params).getSalt ();
    int iteration = ((PBEParameterSpec)params).getIterationCount ();


    //find out the block size in bytes
    int v = getMessageDigestBlockSize(md);
    //get the digest size in bytes
    int u = getMessageDigestOutputSize(md);

    ByteArrayOutputStream I_stream = new ByteArrayOutputStream();

    //Step 1. Construct the string D
    byte[] D = new byte[v];
    for (int i = 0; i < D.length; i++)
      D[i] = (byte)opmode;

    try 
    {
      //Step 2. Construct String S
      if (salt.length!=0) 
      {
        int Slen = 
          (v * (int)Math.ceil((double)salt.length / (double)v));

        for (int i = 0; i < (Slen / salt.length); i++)
          I_stream.write(salt);

        I_stream.write(salt, 0, Slen % salt.length);
      }
      //Step 3. Construct String P
      if (password.length != 0) 
      {
        int Plen = 
          (v * (int)Math.ceil((double)password.length / (double)v));

        for (int i = 0; i < (Plen / password.length); i++)
          I_stream.write(password);

        I_stream.write(password, 0, Plen % password.length);
      }
    } 
    catch (Exception err) 
    {
      err.printStackTrace();
    }

    //Step 4. Combine S & P into I
    byte[] I = I_stream.toByteArray();

    //Step 5. set c = ceiling ( n / u )
    int c = (int)Math.ceil((double)keyLen / (double)u);

    //Step 6
    byte[] A = new byte[v * u];
    for (int i = 0; i < c; i++) 
    {
      //Step 6a
      md.reset();
      md.update(D);
      md.update(I);

      for (int n = 1; n < iteration; n++) 
      {
        byte[] hash_result = md.digest();
        md.reset();
        md.update(hash_result);
      }

      byte[] result = md.digest();
  
      byte[] B = new byte[v];
      for (int n = 1; n <= (v / u); n++)
        System.arraycopy(result, 0, B, u*(n-1), result.length);
  
      System.arraycopy(result, 0, B, (v / u)*u, v % u);
      System.arraycopy(B, 0, A, i * u, u);

      //Modify Ij
      for (int j = 1; j < I.length-1; j=j+v) 
      {
        byte[] Ij_block = new byte[v];
        System.arraycopy(I, j-1, Ij_block, 0, v);
        BigInteger Ij = new BigInteger(B);
        Ij = Ij.add(BigInteger.ONE).add(new BigInteger(Ij_block));
        Ij_block = Ij.toByteArray();
        System.arraycopy(Ij_block, 0, I, j-1, v);
      }
    }

    byte[] result = new byte[keyLen];
    System.arraycopy(A, 0, result, 0, keyLen);
    return result;
  }
  
  // Return the MessageDigest Block Size
  private static int getMessageDigestBlockSize(MessageDigest md) 
  {
    if (md.getAlgorithm().equals("SHA"))
      return 64;   //512 bit block size
    else if (md.getAlgorithm().equals("MD5"))
      return 64;
    else if (md.getAlgorithm().equals("MD2"))
      return 64;
    else if (md.getAlgorithm().equals("RIPEMD160"))
      return 64;
    else 
      return 0;
  }
  
  // Return the MessageDigest Block Size
  private static int getMessageDigestOutputSize(MessageDigest md) 
  {
    if (md.getAlgorithm().equals("SHA"))
      return 20;
    else if (md.getAlgorithm().equals("MD5"))
      return 16;
    else if (md.getAlgorithm().equals("MD2"))
      return 16;
    else if (md.getAlgorithm().equals("RIPMD160"))
      return 20;
    else 
      return 0;
  }
}

