package aalg.crypto;

import java.util.Arrays;

import aalg.algebra.Utils;

/**
 * Encryptor of AES block cipher.
 * @author Ing. Petr Ležák
 */
public class AesEncryptor extends AesBase {
	
	private static final int[] ROUND_TABLE = createRoundTable();
	
	private final AesCipher cipher;
	private final int[] roundKeys = new int[MAX_ROUNDS * COLUMN_COUNT];

	private static int[] createRoundTable() {
		final int[] table = new int[BYTE_LAST * ROW_COUNT];
		final int[][] mixColumnMatrix = {
			{2, 3, 1, 1},
			{1, 2, 3, 1},
			{1, 1, 2, 3},
			{3, 1, 1, 2}
		};
		
		for (int row = 0; row < ROW_COUNT; row++) {
			for (int value = BYTE_FIRST; value < BYTE_LAST; value++) {
				final int sub = subByteForward(value);
				int result = 0;
				
				for (int targetRow = 0; targetRow < ROW_COUNT; targetRow++) {
					final int resultByte = multiplyBytes(sub, mixColumnMatrix[targetRow][row]);
					
					result |= resultByte << (Byte.SIZE * targetRow);
				}
				
				final int tableIndex = calculateRoundTableIndex(row, value);
				table[tableIndex] = result;
			}
		}
		
		return table;
	}
	
	public AesEncryptor(final AesCipher cipher) {
		this.cipher = cipher;
	}
	
	private void subBytes(final int[] block) {
		for (int i = 0; i < COLUMN_COUNT; i++) {
			final int origColumn = block[i];
			int transformedColumn = 0;
			
			for (int j = (ROW_COUNT - 1) * Byte.SIZE; j >= 0; j -= Byte.SIZE) {
				transformedColumn |= subByteForward((origColumn >>> j) & 0xFF) << j;
			}
			
			block[i] = transformedColumn;
		}
	}

	private static int calculateRoundTableIndex (final int row, final int value) {
		return row * BYTE_LAST + value;
	}
	
	private void shiftRows(final int[] block) {
		for (int i = 1; i < ROW_COUNT; i++) {
			final int shift = i * Byte.SIZE;
			final int mask = ~(0xFF << shift);
			int data = 0;
			
			for (int j = COLUMN_COUNT - 1; j >= 0; j--)
				data = (data << Byte.SIZE) | ((block[j] >>> shift) & 0xFF);
			
			data = Integer.rotateRight(data, i * Byte.SIZE);
			
			for (int j = 0; j < COLUMN_COUNT; j++) {
				block[j] = (block[j] & mask) | ((data & 0xFF) << shift);
				data >>>= Byte.SIZE;
			}
		}
	}

	private void addRoundKey(final int round, final int[] block) {
		final int roundKeyAddress = ROW_COUNT * round;
		
		for (int i = 0; i < ROW_COUNT; i++) {
			block[i] ^= roundKeys[roundKeyAddress + i];
		}
	}
	
	@Override
	public void setRoundKey (final int round, final int[] key, final int offset) {
		System.arraycopy(key, offset, roundKeys, round * COLUMN_COUNT, COLUMN_COUNT);
	}

	public void oneRound(final int round, final int[] prevBlock, final int[] nextBlock) {
		Arrays.fill(nextBlock, 0);

		for (int columnIndex = 0; columnIndex < COLUMN_COUNT; columnIndex++) {
			int column = prevBlock[columnIndex];
			
			for (int rowIndex = 0; rowIndex < ROW_COUNT; rowIndex++) {
				final int value = column & 0xFF;
				final int targetColumnIndex = (columnIndex - rowIndex) & 0x03;
				final int tableIndex = calculateRoundTableIndex(rowIndex, value);
					
				nextBlock[targetColumnIndex] ^= ROUND_TABLE[tableIndex];
				column >>>= Byte.SIZE;
			}
		}
		
		addRoundKey(round, nextBlock);
	}

	@Override
	public IBlockCipher getCipher() {
		return cipher;
	}

	@Override
	public void transformBlock (final byte[] sourceBlock, final int sourceOffset, final byte[] targetBlock, final int targetOffset) {
		final int[] prevBlock = new int[ROW_COUNT];
		final int[] nextBlock = new int[ROW_COUNT];
		
		Utils.byteBlockToIntBlock(sourceBlock, prevBlock, sourceOffset, 0, COLUMN_COUNT);
		addRoundKey(0, prevBlock);
		
		for (int round = 1; round < roundCount; round++) {
			oneRound(round, prevBlock, nextBlock);
			System.arraycopy(nextBlock, 0, prevBlock, 0, ROW_COUNT);
		}
		
		subBytes(prevBlock);
		shiftRows(prevBlock);
		addRoundKey(roundCount, prevBlock);	
		
		Utils.intBlockToByteBlock(prevBlock, targetBlock, 0, targetOffset, COLUMN_COUNT);
	}

}
