package ssl;

import java.io.UnsupportedEncodingException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;

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

/**
 * AES encryption stream
 *
 * @author phoenix
 *
 */
public class AESStream {

	/** ???? */
	private String IV = "AAAAAAAAAAAAAAAA";

	/** Encryption key */
	private byte[] key = null;

	/** Used cipher suide */
	private Cipher cipher;

	public AESStream() {
		this.key = null;
		try {
			cipher = Cipher.getInstance("AES/CBC/NoPadding", "SunJCE");
		} catch (NoSuchAlgorithmException e) {
			throw new RuntimeException("Encryption algorithm not supported", e);
		} catch (NoSuchProviderException e) {
			throw new RuntimeException("Encryption provider not supported", e);
		} catch (NoSuchPaddingException e) {
			throw new RuntimeException("Padding algoritm not supported", e);
		}
	}

	/**
	 *
	 * @param key
	 *            UTF-8 encoded encryption key
	 */
	public AESStream(final String key) {
		this();
		try {
			this.key = key.getBytes("UTF-8");
		} catch (UnsupportedEncodingException e) {
			throw new RuntimeException("Unsupported UTF-8 encoding", e);
		}
	}

	/**
	 *
	 * @param key
	 *            UTF-8 encoded encryption key
	 */
	public AESStream(final byte[] key) {
		this();
		this.key = key;
	}

	/**
	 * Do byte paddding, based on 16 bytes
	 *
	 * @param input
	 *            to be padded
	 * @return extended byte array, that is a multiple of 16 long
	 */
	private static byte[] padding(final byte[] input) {
		return padding(input, 16);
	}

	/**
	 * Do byte paddding, based on n bytes (n = size)
	 *
	 * @param input
	 *            to be padded
	 * @param size
	 *            return block will be a multiple of this number
	 * @return extended byte array, that is a multiple of size
	 */
	private static byte[] padding(final byte[] input, final int size) {
		if (input.length % size == 0)
			return input;

		final int newSize = (int) (Math.ceil(input.length / (float) size)) * size;
		final byte[] ret = new byte[newSize];
		for (int i = 0; i < input.length; i++)
			ret[i] = input[i];
		for (int i = input.length; i < newSize; i++)
			ret[i] = '\0';
		return ret;

	}

	/**
	 * Remove tailing '\0' strings
	 *
	 * @param input
	 *            to be de-padded
	 * @return de-padded byte array
	 */
	private static byte[] depadding(final byte[] input) {
		int len = input.length;
		while (len > 0 && input[len - 1] == '\0')
			len--;

		if (len == input.length)
			return input;
		else {
			final byte[] ret = new byte[len];
			for (int i = 0; i < len; i++)
				ret[i] = input[i];
			return ret;
		}
	}

	public byte[] getKey() {
		return key;
	}

	public void setKey(byte[] key) {
		this.key = padding(key);
	}

	public void setKey(final String key) {
		try {
			setKey(padding(key.getBytes("UTF-8")));
		} catch (UnsupportedEncodingException e) {
			throw new RuntimeException("Unsupported UTF-8 encoding", e);
		}
	}

	public byte[] encrypt(final String plainText) {
		try {
			return encrypt(padding(plainText.getBytes("UTF-8")));
		} catch (UnsupportedEncodingException e) {
			throw new RuntimeException("Unsupported encoding", e);
		}
	}

	public byte[] encrypt(final byte[] buffer) {
		final SecretKeySpec key = new SecretKeySpec(this.key, "AES");
		synchronized (this.cipher) {
			try {
				cipher.init(Cipher.ENCRYPT_MODE, key, new IvParameterSpec(IV.getBytes("UTF-8")));
				return cipher.doFinal(buffer);
			} catch (InvalidKeyException e) {
				throw new RuntimeException("Invalid key", e);
			} catch (InvalidAlgorithmParameterException e) {
				throw new RuntimeException("Invalid algorithm", e);
			} catch (UnsupportedEncodingException e) {
				throw new RuntimeException("Unsupported encoding", e);
			} catch (IllegalBlockSizeException e) {
				throw new RuntimeException("Illegal block size", e);
			} catch (BadPaddingException e) {
				throw new RuntimeException("Bad padding", e);
			}

		}
	}

	public String decryptStr(byte[] cipherText) {
		final byte[] buf = decrypt(cipherText);
		try {
			return new String(buf, "UTF-8");
		} catch (UnsupportedEncodingException e) {
			throw new RuntimeException("Unsupported encoding", e);
		}
	}

	public String decryptStr(final String cipherText) {
		try {
			return decryptStr(cipherText.getBytes("UTF-8"));
		} catch (UnsupportedEncodingException e) {
			throw new RuntimeException("Unsupported encoding", e);
		}
	}

	public byte[] decrypt(final String cipherText) {
		try {
			return decrypt(cipherText.getBytes("UTF-8"));
		} catch (UnsupportedEncodingException e) {
			throw new RuntimeException("Unsupported encoding", e);
		}
	}

	public byte[] decrypt(byte[] cipherText) {
		if (cipherText.length % 16 != 0)
			throw new RuntimeException("Input block size wrong");

		synchronized (this.cipher) {
			final SecretKeySpec key = new SecretKeySpec(this.key, "AES");
			try {
				cipher.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(IV.getBytes("UTF-8")));
				return depadding(cipher.doFinal(cipherText));
			} catch (InvalidKeyException e) {
				throw new RuntimeException("Invalid key", e);
			} catch (InvalidAlgorithmParameterException e) {
				throw new RuntimeException("Invalid algorithm", e);
			} catch (UnsupportedEncodingException e) {
				throw new RuntimeException("Unsupported encoding", e);
			} catch (IllegalBlockSizeException e) {
				throw new RuntimeException("Illegal block size", e);
			} catch (BadPaddingException e) {
				throw new RuntimeException("Bad padding", e);
			}
		}
	}

}