//package eLecture;

import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.math.BigInteger;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;

public class CounterMode {

	public static byte[] encrypt(SecretKey sk, boolean encryption, byte[] msg){
      if (encryption)
          return _encrypt(sk,msg);
      else
          return _decrypt(sk,msg);

	}
	
	private static byte[] _encrypt(SecretKey sk,byte[] msg) {
        byte[] finalCipherTxt = new byte[0];
        try{
            Cipher cipher = Cipher.getInstance("AES/ECB/NoPadding");
            cipher.init(Cipher.ENCRYPT_MODE, sk);


            //generate IV
            SecureRandom tempSR = SecureRandom.getInstance("SHA1PRNG");
            byte[] byteHelper = new byte[cipher.getBlockSize()];
            tempSR.nextBytes(byteHelper);
            IvParameterSpec ivSpec = new IvParameterSpec(byteHelper);
            ivSpec.getIV();
            byte[] ctrIV = ivSpec.getIV();
            BigInteger bIV = new BigInteger(ctrIV);

            //initial blocks
            int numBlocks = msg.length / cipher.getBlockSize();
            int numExtraBytes = msg.length % cipher.getBlockSize();
            byte[][] cipherTxt = new byte[numBlocks + (numExtraBytes > 0 ? 2 : 1)][cipher.getBlockSize()];

            for (int i = 0; i < numBlocks; i++) {
                byte tempMsg[] = new byte[cipher.getBlockSize()];
                int k = 0;
                for (int j = i * cipher.getBlockSize(); j < (i + 1) * cipher.getBlockSize(); j++) {
                    tempMsg[k++] = msg[j];
                }
                byte[] cipherTemp = cipher.doFinal(bIV.toByteArray());
                cipherTxt[i] = (new BigInteger(tempMsg)).xor(new BigInteger(cipherTemp)).toByteArray();
                bIV.add(BigInteger.ONE);
            }
            //last block
            if (numExtraBytes != 0) {
                byte lastBlock[] = new byte[cipher.getBlockSize()];
                for (int i = 0; i < lastBlock.length; i++) {
                    lastBlock[i] = 0;
                }
                int k = 0;
                for (int j = numBlocks * cipher.getBlockSize(); j < msg.length; j++) {
                    lastBlock[k++] = msg[j];
                }
                byte[] cipherTemp = cipher.doFinal(bIV.toByteArray());
                cipherTxt[numBlocks] = (new BigInteger(lastBlock)).xor(new BigInteger(cipherTemp)).toByteArray();
            }
            cipherTxt[cipherTxt.length-1][cipher.getBlockSize()-1] = (byte)numExtraBytes;

            //put things together
            finalCipherTxt = new byte[cipher.getBlockSize() + cipherTxt.length * cipher.getBlockSize()];
            for (int i = 0; i < ctrIV.length; i++){
                finalCipherTxt[i] = ctrIV[i];
            }
            int k = cipher.getBlockSize();
            for (int i = 0; i < cipherTxt.length; i++){
                for (int j = 0; j < cipher.getBlockSize(); j++ ){
                    finalCipherTxt[k++] = cipherTxt[i][j];
                }
            }


        }
        catch(Exception e){
            System.out.println(e);
        }
        System.out.println(finalCipherTxt.toString());
        return finalCipherTxt;

	}
    private static byte[] _decrypt(SecretKey sk, byte[] cipherTxt){
        byte[] finalPlainTxt = new byte[0];
        try{
            Cipher cipher = Cipher.getInstance("AES/ECB/NoPadding");
            cipher.init(Cipher.ENCRYPT_MODE, sk);

            byte[] iv = new byte[cipher.getBlockSize()];
            for (int i = 0; i < cipher.getBlockSize(); i++){
                iv[i] = cipherTxt[i];
            }
            BigInteger bIV = new BigInteger(iv);
            //initial blocks
            int numBlocks = cipherTxt.length / cipher.getBlockSize() - 2;
            int numExtraBytes = (int)cipherTxt[cipherTxt.length - 1];
            byte[][] plainTxt = new byte[numBlocks][cipher.getBlockSize()];

            for (int i = 0; i < numBlocks; i++) {
                byte tempMsg[] = new byte[cipher.getBlockSize()];
                int k = 0;
                for (int j = (i+1) * cipher.getBlockSize(); j < (i + 2) * cipher.getBlockSize(); j++) {
                    tempMsg[k++] = cipherTxt[j];
                }
                byte[] cipherTemp = cipher.doFinal(bIV.toByteArray());
                plainTxt[i] = (new BigInteger(tempMsg)).xor(new BigInteger(cipherTemp)).toByteArray();
                bIV.add(BigInteger.ONE);
            }
            finalPlainTxt = new byte[numBlocks * cipher.getBlockSize() + numExtraBytes];
            int k = 0;
            for (int i = 0; i < plainTxt.length; i++){
                for (int j = 0; j < cipher.getBlockSize() && k < finalPlainTxt.length; j++ ){
                    finalPlainTxt[k++] = plainTxt[i][j];
                }
            }
        }
        catch (Exception e){
            e.printStackTrace();
        }
        return finalPlainTxt;
    }

}
