/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package org.mozilla.jss.provider.javax.crypto;

import java.io.CharConversionException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;

import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.SecretKeyFactorySpi;
import javax.crypto.spec.DESKeySpec;
import javax.crypto.spec.DESedeKeySpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

import org.mozilla.jss.CryptoManager;
import org.mozilla.jss.crypto.CryptoToken;
import org.mozilla.jss.crypto.EncryptionAlgorithm;
import org.mozilla.jss.crypto.KeyGenAlgorithm;
import org.mozilla.jss.crypto.KeyWrapAlgorithm;
import org.mozilla.jss.crypto.KeyWrapper;
import org.mozilla.jss.crypto.PBEAlgorithm;
import org.mozilla.jss.crypto.PBEKeyGenParams;
import org.mozilla.jss.crypto.SecretKeyFacade;
import org.mozilla.jss.crypto.SymmetricKey;
import org.mozilla.jss.crypto.TokenException;
import org.mozilla.jss.crypto.TokenRuntimeException;
import org.mozilla.jss.crypto.TokenSupplierManager;
import org.mozilla.jss.util.Password;

public class JSSSecretKeyFactorySpi extends SecretKeyFactorySpi {

    private KeyGenAlgorithm alg = null;
    private CryptoToken token = null;

    protected JSSSecretKeyFactorySpi(KeyGenAlgorithm alg) {
        this.alg = alg;
        token = TokenSupplierManager.getTokenSupplier().getThreadToken();
    }

    private SecretKey
    generateKeyFromBits(byte[] bits, SymmetricKey.Type keyType)
        throws NoSuchAlgorithmException, TokenException,
        InvalidKeySpecException, InvalidAlgorithmParameterException
    {
      try {
        KeyWrapper wrapper = token.getKeyWrapper(KeyWrapAlgorithm.PLAINTEXT);

        wrapper.initUnwrap();

        SymmetricKey symk = wrapper.unwrapSymmetric(bits, keyType, 0);

        return new SecretKeyFacade(symk);
      } catch(InvalidKeyException e) {
            throw new InvalidKeySpecException(e.getMessage());
      }
    }

    private static PBEKeyGenParams
    makePBEKeyGenParams(PBEKeySpec spec)
        throws InvalidKeySpecException
    {
        // The PBEKeySpec in JCE 1.2.1 does not contain salt or iteration.
        // The PBEKeySpec in JDK 1.4 does contain salt and iteration.
        // If this is a JCE 1.2.1 PBEKeySpec, we don't have enough
        // information and we have to throw an exception. If it's a JDK 1.4
        // PBEKeySpec, we can get the information. The only way I know of
        // to find this out at runtime and be compatible with both
        // versions is to use the reflection API.
        Class<?> specClass = spec.getClass();
        try {
            Method getSaltMethod = specClass.getMethod("getSalt",
                                       (java.lang.Class[]) null);
            Method getIterationMethod =
                specClass.getMethod("getIterationCount",
                                    (java.lang.Class[]) null);

            byte[] salt = (byte[]) getSaltMethod.invoke(spec,
                                       (java.lang.Object[]) null);

            Integer itCountObj =
                (Integer) getIterationMethod.invoke(spec,
                                                    (java.lang.Object[]) null);
            int iterationCount = itCountObj.intValue();

            Password pass = new Password(spec.getPassword());

            PBEKeyGenParams params =
                new PBEKeyGenParams(pass, salt, iterationCount);
            pass.clear();
            return params;
        } catch(NoSuchMethodException nsme) {
            // fall through
        } catch(SecurityException se) {
            throw new InvalidKeySpecException(
                "SecurityException calling getMethod() on the key " +
                "spec's class: " + se.getMessage());
        } catch(IllegalAccessException iae) {
            throw new InvalidKeySpecException(
                "IllegalAccessException invoking method on PBEKeySpec: " +
                iae.getMessage());
        } catch(InvocationTargetException ite) {
            String message="";
            Throwable t = ite.getTargetException();
            if( t != null ) {
                message = t.getMessage();
            }
            throw new InvalidKeySpecException(
                "InvocationTargetException invoking method on PBEKeySpec: "+
                message);
        }

        throw new InvalidKeySpecException(
            "This version of PBEKeySpec is unsupported. It must " +
            "implement getSalt() and getIterationCount(). The PBEKeySpec in " +
            "JDK 1.4 works, as does org.mozilla.jss.crypto.PBEKeyGenParams. " +
            "The PBEKeySpec in JCE 1.2.1 and earlier does NOT work.");
    }

    @Override
    public SecretKey
    engineGenerateSecret(KeySpec spec) throws InvalidKeySpecException
    {
      try {
        if( spec instanceof PBEKeySpec ||
            spec instanceof PBEKeyGenParams) {

            PBEKeyGenParams params;
            if( spec instanceof PBEKeySpec ) {
                params = makePBEKeyGenParams((PBEKeySpec)spec);
            } else {
                params = (org.mozilla.jss.crypto.PBEKeyGenParams) spec;
            }
            org.mozilla.jss.crypto.KeyGenerator gen =token.getKeyGenerator(alg);
            gen.initialize(params);
            SymmetricKey symk = gen.generate();
            params.clear();
            return new SecretKeyFacade(symk);
        } else if (spec instanceof DESedeKeySpec) {
            if( alg != KeyGenAlgorithm.DES3 ) {
                throw new InvalidKeySpecException(
                    "Incorrect KeySpec type (" + spec.getClass().getName() +
                    ") for algorithm (" + alg.toString() + ")");
            }
            return generateKeyFromBits(
                ((DESedeKeySpec)spec).getKey(), SymmetricKey.Type.DES3 );
        } else if (spec instanceof DESKeySpec) {
            if( alg != KeyGenAlgorithm.DES ) {
                throw new InvalidKeySpecException(
                    "Incorrect KeySpec type (" + spec.getClass().getName() +
                    ") for algorithm (" + alg.toString() + ")");
            }
            return generateKeyFromBits(
                ((DESKeySpec)spec).getKey(), SymmetricKey.Type.DES );
        } else if( spec instanceof SecretKeySpec ) {
            SecretKeySpec kspec = (SecretKeySpec) spec;
            SymmetricKey.Type type =
                SymmetricKey.Type.fromName( kspec.getAlgorithm());
            return generateKeyFromBits( kspec.getEncoded(), type);
        } else {
            throw new InvalidKeySpecException(
                "Unsupported KeySpec: " + spec.getClass().getName());
        }
      } catch(TokenException te) {
            throw new TokenRuntimeException(te.getMessage());
      } catch(InvalidAlgorithmParameterException iape) {
            throw new InvalidKeySpecException(
                "InvalidAlgorithmParameterException: " + iape.getMessage());
      } catch(IllegalStateException e) {
            throw (TokenRuntimeException) new TokenRuntimeException("IllegalStateException: " +
                e.getMessage()).initCause(e);
      } catch(CharConversionException e) {
            throw new InvalidKeySpecException("CharConversionException: " +
                e.getMessage(), e);
      } catch(NoSuchAlgorithmException e) {
            throw new InvalidKeySpecException("NoSuchAlgorithmException: " +
                e.getMessage(), e);
      }
    }

    @Override
    public KeySpec engineGetKeySpec(SecretKey key, Class<?> keySpec)
            throws InvalidKeySpecException
    {
      try {
        if( ! (key instanceof SecretKeyFacade) ) {
            throw new InvalidKeySpecException("key is not a JSS key");
        }
        SymmetricKey symkey = ((SecretKeyFacade)key).key;
        byte[] keyBits = symkey.getKeyData();
        SymmetricKey.Type keyType = symkey.getType();
        if( keySpec.equals(DESedeKeySpec.class) ) {
            if( keyType != SymmetricKey.Type.DES3 ) {
                throw new InvalidKeySpecException(
                    "key/spec mismatch: " + keyType + " key, DESede spec");
            }
            return new DESedeKeySpec(keyBits);
        } else if( keySpec.equals(DESKeySpec.class) ) {
            if( keyType != SymmetricKey.Type.DES ) {
                throw new InvalidKeySpecException(
                    "key/spec mismatch: " + keyType + " key, DES spec");
            }
            return new DESKeySpec(keyBits);
        } else if( keySpec.equals(SecretKeySpec.class) ) {
            return new SecretKeySpec(keyBits, keyType.toString());
        } else {
            throw new InvalidKeySpecException(
                "Unsupported key spec: " + keySpec.getName());
        }
      } catch(SymmetricKey.NotExtractableException e) {
          throw new InvalidKeySpecException("Key is not extractable: " + e.getMessage(), e);
      } catch(InvalidKeyException e) {
          // This gets thrown by the key spec constructor if there's something
          // wrong with the key bits. But since those key bits came from
          // a real key, this should never happen.
          throw new InvalidKeySpecException("Invalid key: " + e.getMessage(), e);
      }
    }

    @Override
    public SecretKey engineTranslateKey(SecretKey key)
        throws InvalidKeyException
    {
        if( key instanceof SecretKeyFacade ) {
            // try cloning the key
            try {
                SymmetricKey oldkey = ((SecretKeyFacade)key).key;
                CryptoToken owningToken = oldkey.getOwningToken();
                org.mozilla.jss.crypto.KeyGenerator keygen =
                    token.getKeyGenerator(oldkey.getType().getKeyGenAlg());

                SymmetricKey newkey = keygen.clone(oldkey);

                return new SecretKeyFacade(newkey);
            } catch(SymmetricKey.NotExtractableException nee) {
                // no way around this, we fail
                throw new InvalidKeyException("key is not extractable");
            } catch(TokenException te) {
                // fall through and try doing it the long way
            } catch(NoSuchAlgorithmException nsae) {
                throw new InvalidKeyException("Unsupported algorithm: " +
                    nsae.getMessage());
            }
        }

        // try extracting the key value and then creating a new key
        try {
            byte[] keyBits = key.getEncoded();
            if( keyBits == null ) {
                throw new InvalidKeyException("Key is not extractable");
            }
            SymmetricKey.Type keyType =
                SymmetricKey.Type.fromName( key.getAlgorithm() );
            return generateKeyFromBits( keyBits, keyType);
        } catch( NoSuchAlgorithmException nsae ) {
            throw new InvalidKeyException("Unsupported algorithm: "
                + key.getAlgorithm());
        } catch(TokenException te) {
            throw new InvalidKeyException("Token failed to process key: " +
                te.getMessage());
        } catch(InvalidKeySpecException ikse) {
            throw new InvalidKeyException("Invalid key spec: "
                + ikse.getMessage());
        } catch(InvalidAlgorithmParameterException iape) {
            throw new InvalidKeyException("Invalid algorithm parameters: " +
                iape.getMessage());
        }
    }

    public static void main(String args[]) {

      try {
        CryptoManager.initialize(".");

        CryptoManager cm = CryptoManager.getInstance();
        CryptoToken tok = cm.getInternalCryptoToken();
        cm.setThreadToken(tok);

        org.mozilla.jss.crypto.KeyGenerator keygen =
            tok.getKeyGenerator(KeyGenAlgorithm.DES3);

        SymmetricKey symk = keygen.generate();
        SecretKeyFacade origKey = new SecretKeyFacade(symk);

        JSSSecretKeyFactorySpi fact =
            new JSSSecretKeyFactorySpi(KeyGenAlgorithm.DES3);

        DESedeKeySpec kspec = (DESedeKeySpec)
                fact.engineGetKeySpec(origKey, DESedeKeySpec.class);

        SecretKeyFacade newKey = (SecretKeyFacade)
                fact.engineGenerateSecret(kspec);

        org.mozilla.jss.crypto.Cipher cipher =
            tok.getCipherContext(EncryptionAlgorithm.DES3_ECB);
        cipher.initEncrypt(origKey.key);
        String original = "Hello, World!!!!";
        byte[] cipherText = cipher.doFinal( original.getBytes("UTF-8") );
        System.out.println("ciphertext is " + cipherText.length + " bytes");

        cipher.initDecrypt(newKey.key);
        byte[] plainText = cipher.doFinal(cipherText);
        System.out.println("recovered plaintext is " + plainText.length +
            " bytes");

        String recovered = new String(plainText, "UTF-8");
        System.out.println("Recovered '" + recovered + "'");
        if( ! recovered.equals(original) ) {
            throw new Exception("recovered string is different from original");
        }

        char[] pw = "foobarpw".toCharArray();
        byte[] salt = new byte[] { 0, 1, 2, 3, 4, 5, 6, 7 };
        int iterationCount = 2;

        // generate a PBE key the old-fashioned way
        keygen = tok.getKeyGenerator(PBEAlgorithm.PBE_SHA1_DES3_CBC);
        PBEKeyGenParams jssKeySpec =
            new PBEKeyGenParams(pw, salt, iterationCount);
        keygen.initialize(jssKeySpec);
        symk = keygen.generate();
        byte[] keydata = symk.getKeyData();

        // generate a PBE key with the JCE
        SecretKeyFactory keyFact = SecretKeyFactory.getInstance("PBEWithSHA1AndDESede", "Mozilla-JSS");
        newKey = (SecretKeyFacade) keyFact.generateSecret(jssKeySpec);
        byte[] newkeydata = newKey.key.getKeyData();
        if( ! java.util.Arrays.equals(keydata, newkeydata) ) {
            throw new Exception("generated PBE keys are different");
        }
        System.out.println("generated PBE keys are the same");

/* XXX JDK 1.4 ONLY
        // now try with a JDK 1.4 PBEKeySpec
        PBEKeySpec keySpec = new PBEKeySpec(pw, salt, iterationCount);
        newKey = (SecretKeyFacade) keyFact.generateSecret(keySpec);
        if( ! java.util.Arrays.equals(keydata, newKey.key.getKeyData()) ) {
            throw new Exception("generated PBE keys are different");
        }
        System.out.println("generated PBE keys are the same");
*/

        System.exit(0);
      } catch(Throwable t) {
            t.printStackTrace();
            System.exit(-1);
      }
    }

    public static class DES extends JSSSecretKeyFactorySpi {
        public DES() {
            super(KeyGenAlgorithm.DES);
        }
    }
    public static class DESede extends JSSSecretKeyFactorySpi {
        public DESede() {
            super(KeyGenAlgorithm.DESede);
        }
    }
    public static class AES extends JSSSecretKeyFactorySpi {
        public AES() {
            super(KeyGenAlgorithm.AES);
        }
    }
    public static class RC4 extends JSSSecretKeyFactorySpi {
        public RC4() {
            super(KeyGenAlgorithm.RC4);
        }
    }
    public static class RC2 extends JSSSecretKeyFactorySpi {
        public RC2() {
            super(KeyGenAlgorithm.RC2);
        }
    }
    public static class PBE_MD5_DES_CBC extends JSSSecretKeyFactorySpi {
        public PBE_MD5_DES_CBC() {
            super(PBEAlgorithm.PBE_MD5_DES_CBC);
        }
    }
    public static class PBE_SHA1_DES_CBC extends JSSSecretKeyFactorySpi {
        public PBE_SHA1_DES_CBC() {
            super(PBEAlgorithm.PBE_SHA1_DES_CBC);
        }
    }
    public static class PBE_SHA1_DES3_CBC extends JSSSecretKeyFactorySpi {
        public PBE_SHA1_DES3_CBC() {
            super(PBEAlgorithm.PBE_SHA1_DES3_CBC);
        }
    }
    public static class PBE_SHA1_RC4_128 extends JSSSecretKeyFactorySpi {
        public PBE_SHA1_RC4_128() {
            super(PBEAlgorithm.PBE_SHA1_RC4_128);
        }
    }

    public static class HmacSHA1 extends JSSSecretKeyFactorySpi {
        public HmacSHA1() {
            super(KeyGenAlgorithm.SHA1_HMAC);
        }
    }

    public static class PBAHmacSHA1 extends JSSSecretKeyFactorySpi {
        public PBAHmacSHA1() {
            super(KeyGenAlgorithm.PBA_SHA1_HMAC);
        }
    }

    public static class HmacSHA256 extends JSSSecretKeyFactorySpi {
        public HmacSHA256() {
            super(KeyGenAlgorithm.SHA256_HMAC);
        }
    }

    public static class HmacSHA384 extends JSSSecretKeyFactorySpi {
        public HmacSHA384() {
            super(KeyGenAlgorithm.SHA384_HMAC);
        }
    }

    public static class HmacSHA512 extends JSSSecretKeyFactorySpi {
        public HmacSHA512() {
            super(KeyGenAlgorithm.SHA512_HMAC);
        }
    }

    public static class GenericSecret extends JSSSecretKeyFactorySpi {
        public GenericSecret() {
            super(KeyGenAlgorithm.GENERIC_SECRET);
        }
    }
}
