

package gov.va.med.cds.util;


import gov.va.med.cds.exception.EncryptionException;
import gov.va.med.cds.exception.ErrorCodeEnum;

import org.apache.commons.codec.binary.Base64;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.UnrecoverableEntryException;
import java.security.cert.CertificateException;
import java.util.Enumeration;
import java.util.Properties;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;


public class CipherUtilities
{
    public static final String ENCRYPTED_PREFIX = "{AES}";

    private static final String ENCRYPTION_ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/CBC/PKCS5Padding";
    
    // build the initialization vector.
    // Note: When using an initialization vector, the decryption must use the same initialization vector as the encryption
    private static final SecureRandom random = new SecureRandom();
    private static final byte[] iv = random.generateSeed( 16 );
    // Fortify is flagging the below line initialization vector with all 0's for weak encryption
//    private static final byte[] iv = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
//    private static final byte[] iv = { 66, 101, 99, 97, 117, 115, 101, 79, 102, 70, 111, 114, 116, 105, 102, 121 };
    private static final IvParameterSpec ivspec = new IvParameterSpec( iv );
    

    public static SecretKey generateKey( )
    {
        KeyGenerator kgen = null;

        try
        {
            kgen = KeyGenerator.getInstance( ENCRYPTION_ALGORITHM );
            kgen.init( 128 );
        }
        catch ( NoSuchAlgorithmException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_MISSING_KEYSTORE_FILEPATH_EXCEPTION );
        }

        return kgen.generateKey();
    }


    public static SecretKeySpec generateKeySpec( )
    {
        SecretKey skey = generateKey();
        byte[] raw = skey.getEncoded();

        return new SecretKeySpec( raw, ENCRYPTION_ALGORITHM );
    }


    public static SecretKeySpec generateKeySpec( SecretKey aSecretKey )

    {
        byte[] raw = aSecretKey.getEncoded();
        return new SecretKeySpec( raw, ENCRYPTION_ALGORITHM );
    }


    public static SecretKeySpec configureKeyStore( String aKeyStoreFileLocation, String aKeyStorePassword )
        throws IOException
    {
        FileInputStream fileInStream = null;
        KeyStore kstore = null;

        if ( aKeyStorePassword == null )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_MISSING_KEYSTORE_PASSWORD_EXCEPTION );
        }

        if ( aKeyStoreFileLocation == null )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_MISSING_KEYSTORE_FILEPATH_EXCEPTION );
        }

        try
        {
            kstore = KeyStore.getInstance( "JCEKS" );
            // get user password and file input stream
            char[] password = aKeyStorePassword.toCharArray();
            fileInStream = new FileInputStream( aKeyStoreFileLocation );
            kstore.load( fileInStream, password );
        }
        catch ( KeyStoreException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( NoSuchAlgorithmException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( CertificateException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        finally
        {
            if ( fileInStream != null )
            {
                fileInStream.close();
            }
        }

        return retrieveKeyEntry( aKeyStoreFileLocation, aKeyStorePassword, kstore );

    }


    private static SecretKeySpec retrieveKeyEntry( String aKeyStoreFileLocation, String aKeyStorePassword, KeyStore aKeyStore )
        throws IOException
    {
        FileOutputStream fos = null;
        KeyStore.PasswordProtection kp = new KeyStore.PasswordProtection( aKeyStorePassword.toCharArray() );

        // get secret key
        SecretKey secretKey = null;
        KeyStore.SecretKeyEntry secretEntry = null;
        try
        {
            secretEntry = ( KeyStore.SecretKeyEntry )aKeyStore.getEntry( "secretKeyAlias", kp );

            if ( secretEntry == null )
            {
                // save secret key
                secretKey = generateKey();
                KeyStore.SecretKeyEntry skEntry = new KeyStore.SecretKeyEntry( secretKey );
                aKeyStore.setEntry( "secretKeyAlias", skEntry, kp );

                // store away the keystore
                fos = new java.io.FileOutputStream( aKeyStoreFileLocation );
                aKeyStore.store( fos, aKeyStorePassword.toCharArray() );
            }
            else
            {
                secretKey = secretEntry.getSecretKey();
            }
        }
        catch ( NoSuchAlgorithmException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( UnrecoverableEntryException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( KeyStoreException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( CertificateException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        finally
        {
            if ( fos != null )
            {
                fos.close();
            }
        }

        return generateKeySpec( secretKey );
    }


    public static String encrypt( SecretKeySpec aSecretKeySpec, String aStringToEncrypt )
    {
        return encrypt( aSecretKeySpec, aStringToEncrypt, true );
    }


    /**
     * 
     * @param aSecretKeySpec
     * @param aStringToEncrypt
     * @return
     */
    public static String encrypt( SecretKeySpec aSecretKeySpec, String aStringToEncrypt, boolean aAddPrefixFlag )
    {
        byte[] encrypted = null;
        try
        {
            Cipher cipher = Cipher.getInstance( ENCRYPTION_ALGORITHM );
            cipher.init( Cipher.ENCRYPT_MODE, aSecretKeySpec );

            byte[] in = aStringToEncrypt.getBytes();
            encrypted = cipher.doFinal( in );
        }
        catch ( InvalidKeyException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( NoSuchAlgorithmException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( NoSuchPaddingException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( IllegalBlockSizeException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( BadPaddingException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }

        //use 64 base encoding to preserve padding, when convert bytes to string
        String encryptedString = Base64.encodeBase64String( encrypted );

        if ( aAddPrefixFlag )
        {
            encryptedString = ENCRYPTED_PREFIX + encryptedString;
        }
        
        return encryptedString;
    }


    public static void encryptAndSaveProperties( Properties aProperties, SecretKeySpec aSecretKeySpec, String aPropertyFilePath )
        throws IOException
    {
        FileOutputStream fos = new FileOutputStream( new File( aPropertyFilePath ) );
        Properties propsToBeSaved = ( Properties )aProperties.clone();

        Enumeration<?> e = propsToBeSaved.propertyNames();
        String propKey = ( String )e.nextElement();
        if ( propKey.endsWith( ".password" ) )
        {
            String propValue = propsToBeSaved.getProperty( propKey );
            propsToBeSaved.setProperty( propKey, encrypt( aSecretKeySpec, propValue ) );
        }

        try
        {
            propsToBeSaved.store( fos, "File encrypted" );
        }
        finally
        {
            if ( fos != null )
            {
                fos.close();
            }
        }
    }


    /**
     * Used for Cache calls
     * @param data
     * @param key
     * @return
     * @throws NoSuchAlgorithmException
     * @throws NoSuchPaddingException
     * @throws IllegalBlockSizeException
     * @throws BadPaddingException
     * @throws InvalidKeyException
     * @throws InvalidAlgorithmParameterException
     */
    public static String encrypt( String aStringToEncrypt, byte[] aKey )
        throws NoSuchAlgorithmException,
            NoSuchPaddingException,
            IllegalBlockSizeException,
            BadPaddingException,
            InvalidKeyException,
            InvalidAlgorithmParameterException
    {
        // use the class level initialization vector 
        return encrypt( aStringToEncrypt, aKey, iv );
    }


    public static String encrypt( String aStringToEncrypt, byte[] aKey, byte[] anInitializationVector )
        throws NoSuchAlgorithmException,
            NoSuchPaddingException,
            IllegalBlockSizeException,
            BadPaddingException,
            InvalidKeyException,
            InvalidAlgorithmParameterException
    {
        String encodedResult = null;

        SecretKeySpec secretKeySpec = new SecretKeySpec( aKey, ENCRYPTION_ALGORITHM );
        // When using an initialization vector both encryption and decryption must use the same initialization vector, so 
        IvParameterSpec ivParameterSpec = new IvParameterSpec( anInitializationVector );

        // initialize the cipher for encrypt mode
        Cipher cipher = Cipher.getInstance( TRANSFORMATION );
        cipher.init( Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec );

        // encrypt the message
        byte[] encryptedData = cipher.doFinal( aStringToEncrypt.getBytes() );
        encodedResult = Base64.encodeBase64String( encryptedData );

        return encodedResult;
    }


    /**
     * 
     * @param key
     * @param encryptedString
     * @return
     */
    public static String decrypt( SecretKeySpec aSecretKeySpec, String aStringToDecrypt )
    {
        String decryptedString = null;
        if ( isEncrypted( aStringToDecrypt ) )
        {
            aStringToDecrypt = removePrefix( aStringToDecrypt );

            try
            {
                Cipher cipher = Cipher.getInstance( ENCRYPTION_ALGORITHM );
                cipher.init( Cipher.DECRYPT_MODE, aSecretKeySpec );

                //encrypted string was encoded with padded cipher, get bytes back using the same method
                byte[] in = Base64.decodeBase64( aStringToDecrypt );
                byte[] bitesOut = cipher.doFinal( in );
                decryptedString = new String( bitesOut );
            }
            catch ( NoSuchAlgorithmException e )
            {
                throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
            }
            catch ( NoSuchPaddingException e )
            {
                throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
            }
            catch ( InvalidKeyException e )
            {
                throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
            }
            catch ( IllegalBlockSizeException e )
            {
                throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
            }
            catch ( BadPaddingException e )
            {
                throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
            }

            return decryptedString;
        }
        else
        {
            return aStringToDecrypt;
        }
    }


    /**
     * 
     * @param key
     * @param encryptedString
     * @return
     */
    public static String decrypt( byte[] aKey, String aStringToDecrypt )
    {
        // use the class level initialization vector 
        return decrypt( aKey, iv, aStringToDecrypt );
    }


    public static String decrypt( byte[] aKey, byte[] anInitializationVector, String aStringToDecrypt )
    {
        String decryptedString = null;

        try
        {
            SecretKeySpec secretKeySpec = new SecretKeySpec( aKey, ENCRYPTION_ALGORITHM );
            IvParameterSpec ivParameterSpec = new IvParameterSpec( anInitializationVector );

            // initialize the cipher for decrypt mode
//            Cipher cipher = Cipher.getInstance( TRANSFORMATION );
//            Cipher cipher = Cipher.getInstance( ENCRYPTION_ALGORITHM );
//            cipher.init( Cipher.DECRYPT_MODE, keySpec );
            // When using an initialization vector both encryption and decryption must use the same initialization vector, so 
            // use the class level IvParameterSpec 
            Cipher cipher = Cipher.getInstance( TRANSFORMATION );
            cipher.init( Cipher.DECRYPT_MODE, secretKeySpec, ivParameterSpec );

            //encrypted string was encoded with padded cipher, get bytes back using the same method
            byte[] in = Base64.decodeBase64( aStringToDecrypt );
            byte[] bitesOut = cipher.doFinal( in );
            decryptedString = new String( bitesOut );
        }
        catch ( NoSuchAlgorithmException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( NoSuchPaddingException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( InvalidKeyException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( IllegalBlockSizeException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( BadPaddingException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }
        catch ( InvalidAlgorithmParameterException e )
        {
            throw new EncryptionException( ErrorCodeEnum.ENCRYPTION_EXCEPTION, e );
        }

        return decryptedString;
    }


    public static boolean isEncrypted( String anEncryptedString )
    {
        return anEncryptedString.startsWith( ENCRYPTED_PREFIX );
    }


    private static String removePrefix( String anEncryptedString )
    {
        return anEncryptedString.replace( ENCRYPTED_PREFIX, "" );
    }

}
