package aalgTest;


import java.util.Arrays;

import org.junit.Assert;
import org.junit.Test;

import aalg.crypto.AesCipher;
import aalg.crypto.AesDecryptor;
import aalg.crypto.AesEncryptor;
import aalg.crypto.ByteArraySymmetricKey;
import aalg.crypto.IBlockCipherTransformation;
import aalg.crypto.IRandomGenerator;
import aalg.crypto.TestRandomGenerator;

public class AesTest {

	private static class RoundTestCase {
		public final String input;
		public final String key;
		public final String output;
		
		public RoundTestCase (final String input, final String key, final String output) {
			this.input = input;
			this.key = key;
			this.output = output;
		}
	};
	
	private static final RoundTestCase[] ROUND_TEST_CASES = {
		new RoundTestCase("000000000000000000F8000000000000", "00000000000000000000000000000000", "63636363052741416363636363636363"),
		new RoundTestCase("89d810e8855ace682d1843d8cb128fe4", "b692cf0b643dbdf1be9bc5006830b3fe", "4915598f55e5d7a0daca94fa1f0a63f7")
	};
	
	private static final RoundTestCase[] INV_ROUND_TEST_CASES = {
		new RoundTestCase("6363636363636363637C636363636363", "00000000000000000000000000000000", "0000000000000000000000000b0e090d"),
		new RoundTestCase("000000000000000000F8000000000000", "00000000000000000000000000000000", "525252525252525252525252733A0EF4"),
		new RoundTestCase("a7be1a6997ad739bd8c9ca451f618b61", "d6aa74fdd2af72fadaa678f1d6ab76fe", "6353e08c0960e104cd70b751bacad0e7")
	};

	
	private static final RoundTestCase[] ENCRYPT_TEST_CASES = {
		new RoundTestCase("00112233445566778899aabbccddeeff", "000102030405060708090a0b0c0d0e0f", "69c4e0d86a7b0430d8cdb78070b4c55a"),
		new RoundTestCase("3243f6a8885a308d313198a2e0370734", "2b7e151628aed2a6abf7158809cf4f3c", "3925841d02dc09fbdc118597196a0b32"),
		new RoundTestCase("00112233445566778899aabbccddeeff", "000102030405060708090a0b0c0d0e0f1011121314151617", "dda97ca4864cdfe06eaf70a0ec0d7191"),
		new RoundTestCase("00112233445566778899aabbccddeeff", "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f", "8ea2b7ca516745bfeafc49904b496089")
	};
	
	private int[] parseHexInts(final String input) {
		final int digitsPerItem = 8;
		final int[] result = new int[input.length() / digitsPerItem];
		
		for (int i = 0; i < result.length; i++) {
			final String numStr = input.substring(digitsPerItem * i, digitsPerItem * (i+1));
			final int number = (int) Long.parseLong(numStr, 16);
			
			result[i] = Integer.reverseBytes(number);
		}
		
		return result;
	}
	
	private byte[] parseHexBytes(final String input) {
		final int digitsPerItem = 2;
		final byte[] result = new byte[input.length() / digitsPerItem];
		
		for (int i = 0; i < result.length; i++) {
			final String numStr = input.substring(digitsPerItem * i, digitsPerItem * (i+1));
			final int number = Integer.parseInt(numStr, 16);
			
			result[i] = (byte) number;
		}
		
		return result;
	}
	
	@Test
	public void testRound() {
		final AesEncryptor encryptor = new AesEncryptor(null);
		
		for (RoundTestCase testCase: ROUND_TEST_CASES) {
			final int[] input = parseHexInts (testCase.input);
			final int[] roundKey = parseHexInts (testCase.key);
			final int[] expectedOutput = parseHexInts (testCase.output);
			final int[] output = new int[expectedOutput.length];
			
			encryptor.setRoundKey(0, roundKey, 0);
			encryptor.oneRound(0, input, output);
			
			Assert.assertTrue(Arrays.equals(expectedOutput, output));
		}
	}
	
	@Test
	public void testEncrypt() {
		final AesEncryptor encryptor = new AesEncryptor(null);
		
		for (RoundTestCase testCase: ENCRYPT_TEST_CASES) {
			final byte[] input = parseHexBytes (testCase.input);
			final byte[] key = parseHexBytes (testCase.key);
			final byte[] expectedOutput = parseHexBytes (testCase.output);
			final byte[] output = new byte[expectedOutput.length];
			
			encryptor.setKey(new ByteArraySymmetricKey(key));
			encryptor.transformBlock(input, 0, output, 0);
			
			Assert.assertTrue(Arrays.equals(expectedOutput, output));
		}
	}
	
	@Test
	public void testInvRound() {
		final AesDecryptor decryptor = new AesDecryptor(null);
		
		for (RoundTestCase testCase: INV_ROUND_TEST_CASES) {
			final int[] input = parseHexInts (testCase.input);
			final int[] roundKey = parseHexInts (testCase.key);
			final int[] expectedOutput = parseHexInts (testCase.output);
			final int[] output = new int[expectedOutput.length];
			
			decryptor.setRoundKey(1, roundKey, 0);
			decryptor.oneRound(1, input, output);
			
			Assert.assertTrue(Arrays.equals(expectedOutput, output));
		}
	}

	@Test
	public void testDecrypt() {
		final AesDecryptor decryptor = new AesDecryptor(null);
		
		for (RoundTestCase testCase: ENCRYPT_TEST_CASES) {
			final byte[] input = parseHexBytes (testCase.output);
			final byte[] key = parseHexBytes (testCase.key);
			final byte[] expectedOutput = parseHexBytes (testCase.input);
			final byte[] output = new byte[expectedOutput.length];
			
			decryptor.setKey(new ByteArraySymmetricKey(key));
			decryptor.transformBlock(input, 0, output, 0);
			
			Assert.assertTrue(Arrays.equals(expectedOutput, output));
		}
	}
	
	private static final int SPEED_TEST_COUNT = 2000000;
	
	@Test
	public void measureEncryptSpeed() {
		final AesCipher cipher = new AesCipher(AesCipher.AES_256_KEY_SIZE);
		
		final IRandomGenerator rng = new TestRandomGenerator();
		
		final byte[] key = new byte[cipher.getKeySize()];
		rng.fillArray(key, 0, key.length);
		
		final int blockLength = cipher.getBlockLength();
		final byte[] input = new byte[blockLength];
		final byte[] output = new byte[blockLength];
		
		final IBlockCipherTransformation encryptor = cipher.getEncryptor();
		encryptor.setKey(new ByteArraySymmetricKey(key));

		encryptLoop(encryptor, input, output);
		
		final long t1 = System.currentTimeMillis();
		encryptLoop(encryptor, input, output);
		final long t2 = System.currentTimeMillis();
		
		System.out.println ((1000L * SPEED_TEST_COUNT * blockLength / (t2 - t1) / (1024*1024)) + "MB/s");
	}

	private void encryptLoop(final IBlockCipherTransformation encryptor, final byte[] input, final byte[] output) {
		for (int i = 0; i < SPEED_TEST_COUNT; i++) {
			encryptor.transformBlock(input, 0, output, 0);
			System.arraycopy(output, 0, input, 0, input.length);
		}
	}


}
