package com.eactive.eai.common.security;

import java.nio.ByteBuffer;
import java.security.Key;
import java.security.SecureRandom;
import java.security.Security;
import java.security.spec.AlgorithmParameterSpec;

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

import org.bouncycastle.jcajce.spec.FPEParameterSpec;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

public class ARIACryptoModuleExtension implements CryptoModuleExtension {

	private Cipher encryptEngine;
	private Cipher decryptEngine;

	@Override
	public void init(String alg, String mode, String pad, byte[] encKey, byte[] decKey) throws Exception {
		this.init("ARIA", mode, pad, (byte[]) null, encKey, decKey);
	}

	@Override
	public void init(String alg, String mode, String pad, byte[] iv, byte[] encKey, byte[] decKey) throws Exception {
		if (Security.getProvider("BC") == null) {
			Security.addProvider(new BouncyCastleProvider());
		}
		String m = mode != null ? mode : "";
		String p = pad != null ? pad : "";
		Key encKeySpec = new SecretKeySpec(encKey, "ARIA");
		Key decKeySpec = new SecretKeySpec(decKey, "ARIA");
		AlgorithmParameterSpec param = iv != null ? new IvParameterSpec(iv) : null;
		if (m.toLowerCase().contains("ff")) {
			param = new FPEParameterSpec(256, iv);
		}
		this.encryptEngine = Cipher.getInstance(String.format("ARIA/%s/%s", m, p));
		this.encryptEngine.init(Cipher.ENCRYPT_MODE, encKeySpec, param, (SecureRandom) null);
		this.decryptEngine = Cipher.getInstance(String.format("ARIA/%s/%s", m, p));
		this.decryptEngine.init(Cipher.DECRYPT_MODE, decKeySpec, param, (SecureRandom) null);
	}

	@Override
	public int encrypt(byte[] src, int sindex, int slength, byte[] dst, int dindex) throws Exception {
		return this.encryptEngine.doFinal(src, sindex, slength, dst, dindex);
	}

	@Override
	public byte[] encrypt(byte[] src, int sindex, int slength) throws Exception {
		int size = (slength + 17) / this.encryptEngine.getBlockSize() * this.encryptEngine.getBlockSize();
		size += (slength + 17) % this.encryptEngine.getBlockSize() == 0 ? 0 : this.encryptEngine.getBlockSize();
		byte[] dst = new byte[this.encryptEngine.getOutputSize(size)];
		size = this.encrypt(src, sindex, slength, dst, 0);
		return ByteBuffer.allocate(size).put(dst, 0, size).array();
	}

	@Override
	public int decrypt(byte[] src, int sindex, int slength, byte[] dst, int dindex) throws Exception {
		return this.decryptEngine.doFinal(src, sindex, slength, dst, dindex);
	}

	@Override
	public byte[] decrypt(byte[] src, int sindex, int slength) throws Exception {
		return this.decryptEngine.doFinal(src, sindex, slength);
	}
}
