package org.jivesoftware.util;

import java.nio.charset.StandardCharsets;
import java.security.Key;
import java.security.Security;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Utility class providing symmetric AES encryption/decryption. To strengthen
 * the encrypted result, use the {@link #setKey} method to provide a custom
 * key prior to invoking the {@link #encrypt} or {@link #decrypt} methods.
 *
 * @author Tom Evans
 */
public class AesEncryptor implements Encryptor {

	private static final Logger log = LoggerFactory.getLogger(AesEncryptor.class);
	private static final String ALGORITHM = "AES/CBC/PKCS7Padding";

	private static final byte[] INIT_PARM =
	{
		(byte)0xcd, (byte)0x91, (byte)0xa7, (byte)0xc5,
		(byte)0x27, (byte)0x8b, (byte)0x39, (byte)0xe0,
		(byte)0xfa, (byte)0x72, (byte)0xd0, (byte)0x29,
		(byte)0x83, (byte)0x65, (byte)0x9d, (byte)0x74
	};

	private static final byte[] DEFAULT_KEY =
	{
		(byte)0xf2, (byte)0x46, (byte)0x5d, (byte)0x2a,
		(byte)0xd1, (byte)0x73, (byte)0x0b, (byte)0x18,
		(byte)0xcb, (byte)0x86, (byte)0x95, (byte)0xa3,
		(byte)0xb1, (byte)0xe5, (byte)0x89, (byte)0x27
	};

    private static boolean isInitialized = false;

	private byte[] cipherKey = null;

    /** Default constructor */
    public AesEncryptor() { initialize(); }

    /** Custom key constructor */
    public AesEncryptor(String key) { 
    	initialize();
    	setKey(key);
    }

	/* (non-Javadoc)
	 * @see org.jivesoftware.util.Encryptor#encrypt(java.lang.String)
	 */
	@Override
	public String encrypt(String value)
	{
		if (value == null) { return null; }
		byte [] bytes = value.getBytes(StandardCharsets.UTF_8);
		return Base64.encodeBytes( cipher(bytes, getKey(), Cipher.ENCRYPT_MODE) );
	}

	/* (non-Javadoc)
	 * @see org.jivesoftware.util.Encryptor#decrypt(java.lang.String)
	 */
	@Override
	public String decrypt(String value)
	{
		if (value == null) { return null; }
		byte [] bytes = cipher(Base64.decode(value), getKey(), Cipher.DECRYPT_MODE);
		if (bytes == null) { return null; }
		return new String(bytes, StandardCharsets.UTF_8);
	}

	/**
	 * Symmetric encrypt/decrypt routine.
	 *
	 * @param attribute The value to be converted
	 * @param key The encryption key
	 * @param mode The cipher mode (encrypt or decrypt)
	 * @return The converted attribute, or null if conversion fails
	 */
	private byte [] cipher(byte [] attribute, byte [] key, int mode)
	{
		byte [] result = null;
		try
		{
		    // Create AES encryption key
		    Key aesKey = new SecretKeySpec(key, "AES");

		    // Create AES Cipher
		    Cipher aesCipher = Cipher.getInstance(ALGORITHM);

		    // Initialize AES Cipher and convert
		    aesCipher.init(mode, aesKey, new IvParameterSpec(INIT_PARM));
		    result = aesCipher.doFinal(attribute);
		}
		catch (Exception e)
		{
			log.error("AES cipher failed", e);
		}
		return result;
	}

	/**
	 * Return the encryption key. This will return the user-defined
	 * key (if available) or a default encryption key.
	 *
	 * @return The encryption key
	 */
	private byte [] getKey()
	{
		return cipherKey == null ? DEFAULT_KEY : cipherKey;
	}

	/**
	 * Set the encryption key. This will apply the user-defined key,
	 * truncated or filled (via the default key) as needed  to meet
	 * the key length specifications.
	 *
	 * @param key The encryption key
	 */
	private void setKey(byte [] key)
	{
		cipherKey = editKey(key);
	}

	/* (non-Javadoc)
	 * @see org.jivesoftware.util.Encryptor#setKey(java.lang.String)
	 */
	@Override
	public void setKey(String key)
	{
		if (key == null) { 
			cipherKey = null; 
			return;
		}
		byte [] bytes = key.getBytes(StandardCharsets.UTF_8);
		setKey(editKey(bytes));
	}

	/**
	 * Validates an optional user-defined encryption key. Only the
	 * first sixteen bytes of the input array will be used for the key.
	 * It will be filled (if necessary) to a minimum length of sixteen.
	 *
	 * @param key The user-defined encryption key
	 * @return A valid encryption key, or null
	 */
	private byte [] editKey(byte [] key)
	{
		if (key == null) { return null; }
		byte [] result = new byte [DEFAULT_KEY.length];
		for (int x=0; x<DEFAULT_KEY.length; x++)
		{
			result[x] = x < key.length ? key[x] : DEFAULT_KEY[x];
		}
		return result;
	}

	/** Installs the required security provider(s) */
	private synchronized void initialize()
	{
		if (!isInitialized)
		{
			try
			{
		        Security.addProvider(new BouncyCastleProvider());
				isInitialized = true;
			}
			catch (Throwable t)
			{
				log.warn("JCE provider failure; unable to load BC", t);
			}
		}
	}

/* */
	
}
