package eu.stork.documentservice.utils;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.security.AlgorithmParameters;
import java.security.PublicKey;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;

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

import com.sun.org.apache.xml.internal.security.utils.Base64;

import eu.stork.documentservice.exceptions.EncryptionException;

public class EncryptionHelper {
	
	private static String key;
	private static String iv;
	private static Cipher cipher;
	
	public EncryptionHelper() throws EncryptionException
	{
		this.generateKeys();
	}
	
	public EncryptionHelper(String inKey, String inIv) throws EncryptionException
	{
		this.initKeys(inKey, inIv);
	}
	
	/**
	 * Generate new symmetric keys
	 * @throws EncryptionException
	 */
	public void generateKeys() throws EncryptionException
	{
		try 
		{
			KeyGenerator keyGen = KeyGenerator.getInstance("AES");
			keyGen.init(256);
			SecretKey secretKey = keyGen.generateKey();
			cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
			cipher.init(Cipher.ENCRYPT_MODE, secretKey);
			AlgorithmParameters params = cipher.getParameters();
			key = Base64.encode(secretKey.getEncoded());
			iv = Base64.encode(params.getParameterSpec(IvParameterSpec.class).getIV());
		} 
		catch (Exception e) 
		{				
			e.printStackTrace();
			throw new EncryptionException("Unable to generate encryption key.", e);
		}		
	}
	
	/**
	 * Initialize keys with specified keys
	 * @param inKey the key to use
	 * @param inIv the IV to use
	 * @throws EncryptionException the exception thrown
	 */
	public void initKeys(String inKey, String inIv) throws EncryptionException
	{
		try 
		{
			key = inKey;
			iv = inIv;
			SecretKeySpec skeySpec = new SecretKeySpec(Base64.decode(inKey), "AES");
			cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
			cipher.init(Cipher.DECRYPT_MODE, skeySpec, new IvParameterSpec(Base64.decode(inIv)));
		} 
		catch (Exception e) 
		{				
			e.printStackTrace();
			throw new EncryptionException("Unable to initialize encryption key.", e);
		}		
	}
	
	/**
	 * Encrypt data with key
	 * @param clearData the clear data
	 * @return the encrypted data
	 * @throws EncryptionException the exception thrown
	 */
	public byte[] encrypt(byte[] clearData) throws EncryptionException
	{		
		if (clearData != null)
		{
			try
			{
				return cipher.doFinal(clearData);
			}
			catch (Exception ex)
			{
				throw new EncryptionException("Could not decrypt data.", ex);
			}
		}
		else
			throw new EncryptionException("Clear data is null.");
	}
	
	/**
	 * Decrypt data with keys
	 * @param encData the encrypted data
	 * @return decrypted data
	 * @throws EncryptionException the exception thrown
	 */
	public byte[] decrypt(byte[] encData) throws EncryptionException
	{
		if (encData != null)
		{
			try
			{
				return cipher.doFinal(encData);
			}
			catch (Exception ex)
			{
				throw new EncryptionException("Could not encrypt data.", ex);
			}
		}
		else
			throw new EncryptionException("Encrypted data is null.");
	}
	
	/**
	 * Get the key string
	 * @return the key
	 */
	public String getKey()
	{
		return key;
	}
	
	/**
	 * Get the IV string
	 * @return the iv
	 */
	public String getIv()
	{
		return iv; 
	}
	
	/**
	 * Encrypt string with certificate
	 * @param certString the PEM formated certificate
	 * @param input the string to encrypt
	 * @return encrypted string
	 * @throws EncryptionException the exception thrown
	 */
	public String encryptWithCert(String certString, String input) throws EncryptionException
	{
		if (certString != null && !certString.isEmpty())
		{
			if (input != null && !input.isEmpty())
			{
				try {					
					certString = certString.replace("-----BEGIN CERTIFICATE-----", "");
					certString = certString.replace("-----END CERTIFICATE-----", "");
					InputStream inStream = new ByteArrayInputStream(Base64.decode(certString));
					CertificateFactory cf = CertificateFactory.getInstance("X.509");
					X509Certificate cert = (X509Certificate)cf.generateCertificate(inStream);
					PublicKey pk = cert.getPublicKey();
					Cipher rsaCipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
					rsaCipher.init(Cipher.ENCRYPT_MODE, pk);
					byte[] encrypted = rsaCipher.doFinal(input.getBytes("UTF-8"));
					return Base64.encode(encrypted);
				} 
				catch (Exception e) {
					e.printStackTrace();
					throw new EncryptionException("Unabled to encrypt string.", e);
				}
			}
			else
				throw new EncryptionException("Input is null or empty.");
		}
		else
			throw new EncryptionException("Certificate is null or empty.");
	}
	
}