import java.io.UnsupportedEncodingException;
import java.security.GeneralSecurityException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.util.Random;

import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;

public class CredentialMatcherTests
{
    static interface CredentialMatcher
    {
        public boolean matches(String plainText, byte[] storedCredentials)
            throws GeneralSecurityException, UnsupportedEncodingException;
        public byte[] mutate(String plainText)
            throws GeneralSecurityException, UnsupportedEncodingException;
    }

    public static void main(String[] args)
        throws Exception
    {
        CredentialMatcher cm;

        if("pbkdf2".equalsIgnoreCase(System.getProperty("matcher")))
            cm = new PBKDF2CredentialMatcher();
        else
            cm = new MessageDigestCredentialMatcher();

        if("create".equals(args[0]))
        {
            if(cm instanceof PBKDF2CredentialMatcher)
                ((PBKDF2CredentialMatcher)cm).setIterations(((PBKDF2CredentialMatcher)cm).benchmark(2000, 100));

            System.out.println(((BaseCredentialMatcher)cm).getAlgorithm() + ": " + toByteString(cm.mutate(args[1])));
        }
        else if("check".equals(args[0]))
            System.out.println("matches=" + cm.matches(args[1], fromByteString(args[2])));
    }

    static final char[] hex = new char[] { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f' };
    /**
     * Converts byte array into a String with hexadecimal encoding
     * (e.g. { 0xca, 0xfe } -> "cafe").
     *
     * @param a A byte array.
     *
     * @return A String representing the btye array as text.
     *
     * @see #fromByteString
     */
    public static String toByteString(final byte[] a)
    {
        final int len = a.length;
        final char[] chars = new char[len << 1];
        int j=0;
        for(int i=0; i<len; ++i)
        {
            byte b = a[i];
            chars[j++] = hex[(b >> 4) & 0x0f];
            chars[j++] = hex[(b     ) & 0x0f];
        }

        return new String(chars);
    }

    /**
     * Converts a String of "bytes" into a byte array.
     *
     * @param s A String of characters representing bytes.
     *
     * @return An array of bytes specified by the string.
     *
     * @see #toByteString
     */
    public static byte[] fromByteString(final String s)
    {
        if(null == s)
            return null;

        if(1 == (s.length() & 0x01))
            throw new IllegalArgumentException("String must have an even number of characters.");

        final int length = s.length() >>> 1;

        final byte[] bytes = new byte[length];

        for(int i=0; i < length; ++i)
        {
            int idx = i << 1; // i * 2
            bytes[i] = (byte)
                (
                 (byte)(getNibble(s.charAt(idx)) << 4) // move to upper nibble
                 |
                 getNibble(s.charAt(idx + 1))
                 )
                ;
        }

        return bytes;
    }

    private static byte getNibble(final char c)
    {
        if(c >= '0' && c <= '9')
            return (byte)(c - '0');
        if(c >= 'a' && c <= 'f')
            return (byte)(c - 'a' + 10);
        if(c >= 'A' && c <= 'F')
            return (byte)(c - 'A' + 10);

        throw new IllegalArgumentException("Invalid byte value: " + c);
    }

    static abstract class BaseCredentialMatcher
    {
        private String _algorithm = getDefaultAlgorithm();

        /**
         * The number of iterations through which the password will be
         * passed-through the password-derivation algorithm.
         */
        private int _iterations = getDefaultIterations();

        /**
         * The length of the salt, in bytes.
         */
        private int _saltLength = getDefaultSaltLength();

        private final Random _random = new SecureRandom();

        protected abstract String getDefaultAlgorithm();
        protected abstract int getDefaultIterations();
        protected abstract int getDefaultSaltLength();
        
        public void setAlgorithm(String algorithm) {
            _algorithm = algorithm;
        }
        public void setIterations(int iterations) {
            _iterations = iterations;
        }
        public void setSaltLength(int saltLength) {
            _saltLength = saltLength;
        }

        public String getAlgorithm() {
            return _algorithm;
        }
        public int getIterations() {
            return _iterations;
        }
        public int getSaltLength() {
            return _saltLength;
        }

        /**
         * Generates a new random salt of length {@link #getSaltLength()}.
         *
         * @return A new array filled with random bytes.
         */
        protected byte[] generateSalt() {
            return generateSalt(getSaltLength(), _random);
        }
        
        /**
         * Generates a new random salt of length <code>saltLength</code>
         * filled with random data provided by <code>random</code>.
         *
         * @return A new array filled with random bytes.
         */
        protected static byte[] generateSalt(int saltLength, Random random) {
            byte[] salt = new byte[saltLength];

            /*
              // un-comment if you want to have a consistent salt for testing
              for(int i=0; i<saltLength; ++i)
                  salt[i] = (byte)i;
             */
            random.nextBytes(salt);

            return salt;
        }

        protected byte[] pack(byte[] salt, int iterations, byte[] mutated) {
            int saltLength = salt.length;

            // The complete credential = salt + '$' + iteration count + '$' + encoded credential
            byte[] complete = new byte[saltLength + 1 + 4 + 1 + mutated.length];
            System.arraycopy(salt, 0, complete, 0, saltLength);
            complete[saltLength] = (byte)'$';
            complete[saltLength + 1] = (byte)((iterations >> 24) & 0xff);
            complete[saltLength + 2] = (byte)((iterations >> 16) & 0xff);
            complete[saltLength + 3] = (byte)((iterations >>  8) & 0xff);
            complete[saltLength + 4] = (byte)((iterations      ) & 0xff);
            complete[saltLength + 5] = (byte)'$';
            System.arraycopy(mutated, 0, complete, saltLength + 6, mutated.length);

            return complete;
        }
        /**
         * Compares two arrays byte-wise.
         *
         * @param a The first array.
         * @param apos The offset into the first array.
         * @param b The second array.
         * @param bpos The offset into the second array.
         * @param length The number of items to compare.
         *
         * @return <code>true</code> if <code>length</code> bytes of both
         *         arrays are identical starting at their respective offsets.
         */
        public boolean equals(byte[] a, int apos, byte[] b, int bpos, int length)
        {
            // Protect against array overflow
            if(a.length - apos - length < 0
               || b.length - bpos - length < 0)
                return false;

            for(int i=0; i<length; ++i)
                if(a[i + apos] != b[i + bpos])
                    return false;
            
            return true;
        }
    }
    static class MessageDigestCredentialMatcher
        extends BaseCredentialMatcher
        implements CredentialMatcher
    {
        public static final String DEFAULT_ALGORITHM = "SHA-512";
        public static final int DEFAULT_ITERATION_COUNT = 1;
        public static final int DEFAULT_SALT_LENGTH = 32;

        /**
         * The character encoding that will be used to extract bytes
         * from plain-text passwords.
         */
        private String _encoding = "UTF-8";

        public String getEncoding() {
            return _encoding;
        }

        @Override
        protected String getDefaultAlgorithm() {
            return DEFAULT_ALGORITHM;
        }

        @Override
        protected int getDefaultIterations()
        {
            return DEFAULT_ITERATION_COUNT;
        }

        @Override
        protected int getDefaultSaltLength() {
            return DEFAULT_SALT_LENGTH;
        }

        @Override
        public boolean matches(String plainText, byte[] storedCredentials)
            throws GeneralSecurityException, UnsupportedEncodingException
        {
            return equals(storedCredentials, 0,
                           mutate(plainText), 0,
                           storedCredentials.length);
        }

        @Override
        public byte[] mutate(String plainText)
            throws GeneralSecurityException, UnsupportedEncodingException
        {
            MessageDigest md = MessageDigest.getInstance(getAlgorithm());
            byte[] plainTextBytes = plainText.getBytes(getEncoding());

            int iterations = getIterations();

            byte[] salt = generateSalt();

            for(int i=0; i<iterations; ++i) {
                if (null != salt)
                    md.update(salt);

                md.update(plainTextBytes);
            }

            // Blank-out the plain-text copy of the byte array
            for(int i=plainTextBytes.length; i>=0; --i)
                plainTextBytes[0] = 0x00;

            return pack(salt, iterations, md.digest());
        }
    }

    static class PBKDF2CredentialMatcher
        extends BaseCredentialMatcher
        implements CredentialMatcher
    {
        public static final String DEFAULT_ALGORITHM = "PBKDF2WithHmacSHA1";
        public static final int DEFAULT_ITERATION_COUNT = 20000;
        public static final int DEFAULT_SALT_LENGTH = 32;

        @Override
        protected String getDefaultAlgorithm() {
            return DEFAULT_ALGORITHM;
        }

        @Override
        protected int getDefaultIterations()
        {
            return DEFAULT_ITERATION_COUNT;
        }

        @Override
        protected int getDefaultSaltLength() {
            return DEFAULT_SALT_LENGTH;
        }

        public int benchmark(long target, long epsilon)
            throws NoSuchAlgorithmException, InvalidKeySpecException
        {
            // SHA-1 generates 160 bit hashes, so that's what makes sense here
            //int derivedKeyLength = 160;
            int derivedKeyLength = 160;

            byte[] salt = generateSalt();

            // The time for iterations scales linearly. We'll take a sample
            // and try to hone-in as quickly as possible on what it would take
            // to reach the target time.
            int iterations = 1000;

            SecretKeyFactory f = SecretKeyFactory.getInstance(getAlgorithm());

            long elapsed = 0;

            boolean stop = false;
            char[] password = "testing".toCharArray();
            while(!stop)
            {
                elapsed = System.currentTimeMillis();

                KeySpec spec = new PBEKeySpec(password, salt, iterations, derivedKeyLength);
                f.generateSecret(spec).getEncoded();
                elapsed = System.currentTimeMillis() - elapsed;


                if(Math.abs(target - elapsed) < epsilon)
                    stop = true;
                else
                    iterations *= ((double)target / (double)elapsed);
            }

            return iterations;
        }

        @Override
        public byte[] mutate(String password)
            throws GeneralSecurityException
        {
            return mutate(password.toCharArray());
        }

        public byte[] mutate(char[] password)
            throws GeneralSecurityException
        {
            byte[] salt = generateSalt();
            int iterations = getIterations();

            return pack(salt, iterations, mutate(getAlgorithm(), password, salt, iterations));
        }

        protected static byte[] mutate(String algorithm, char[] password, byte[] salt, int iterations)
            throws GeneralSecurityException
        {
            // SHA-1 generates 160 bit hashes, so that's what makes sense here
            int derivedKeyLength = 160; // How can we determine this value at runtime?

            KeySpec spec = new PBEKeySpec(password, salt, iterations, derivedKeyLength);

            SecretKeyFactory f = SecretKeyFactory.getInstance(algorithm);

            return f.generateSecret(spec).getEncoded();
        }

        @Override
        public boolean matches(String plainText, byte[] storedCredential)
            throws GeneralSecurityException
        {
            // Unpack the salt and iteration count from the stored credentials
            byte[] salt;
            int iterations;

            int saltDivider = -1;
            int iterationDivider = -1;

            for(int i=0; i<storedCredential.length; ++i)
            {
                if('$' == storedCredential[i])
                {
                    if(-1 == saltDivider)
                        saltDivider = i;
                    else
                    {
                        iterationDivider = i;
                        break;
                    }
                }
            }

            if(-1 == saltDivider || -1 == iterationDivider)
                throw new IllegalArgumentException("Improperly encoded stored credential");

            salt = new byte[saltDivider];
            System.arraycopy(storedCredential, 0, salt, 0, saltDivider);
            iterations = (((int)storedCredential[saltDivider + 1]) << 24) & 0xff000000
                    | (((int)storedCredential[saltDivider + 2]) << 16) & 0x00ff0000
                    | (((int)storedCredential[saltDivider + 3]) <<  8) & 0x0000ff00
                    | (((int)storedCredential[saltDivider + 4]) <<  0) & 0x000000ff
                    ;

            byte[] recomputed = mutate(getAlgorithm(), plainText.toCharArray(), salt, iterations);

            return equals(storedCredential, saltDivider + 6, recomputed, 0, storedCredential.length - iterationDivider - 1);
        }
    }
}