package aalg.crypto;

import aalg.algebra.Utils;

public class CtrStreamBase {

	private final IBlockCipher dataBlockCipher;
	private final IBlockCipher ivBlockCipher;
	private final IBlockCipherTransformation dataEncryptor;
	private final int[] counter;
	private final int[] increment;
	private final byte[] counterBytes;
	private final byte[] keyStreamBlock;
	private int keyStreamIndex;
	protected final int blockLength;
	protected ISymmetricKey symmetricKey;
	
	
	protected CtrStreamBase (final IBlockCipher dataBlockCipher, final IBlockCipher ivBlockCipher, final int[] increment) {
		this.dataBlockCipher = dataBlockCipher;
		this.ivBlockCipher = ivBlockCipher;
		
		this.blockLength = dataBlockCipher.getBlockLength();
		
		if (ivBlockCipher.getBlockLength() != blockLength)
			throw new RuntimeException("Different block lengths of data block cipher and IV block cipher");
		
		if (blockLength % Utils.INT_BYTE_COUNT != 0)
			throw new RuntimeException("Block length is not divisible by 4");
		
		this.counter = new int[blockLength / Utils.INT_BYTE_COUNT];
		this.counterBytes = new byte[blockLength];
		this.keyStreamBlock = new byte[blockLength];
		this.dataEncryptor = dataBlockCipher.getEncryptor();
		this.increment = increment;
	}
	
	protected void initializeCounter(final byte[] nonce) {
		final int ivKeyLength = ivBlockCipher.getKeySize();
		final int dataKeyLength = dataBlockCipher.getKeySize();
		final ISymmetricKey ivKey = symmetricKey.getSubKey(0, ivKeyLength);
		final ISymmetricKey dataKey = symmetricKey.getSubKey(ivKeyLength, dataKeyLength);
		
		final IBlockCipherTransformation ivEncryptor = ivBlockCipher.getEncryptor();
		ivEncryptor.setKey(ivKey);
		
		final byte[] iv = new byte[blockLength];
		ivEncryptor.transformBlock(nonce, 0, iv, 0);
		
		Utils.byteBlockToIntBlock(iv, counter, 0, 0, blockLength / Utils.INT_BYTE_COUNT);
		dataEncryptor.setKey(dataKey);
		
		keyStreamIndex = 0;
	}
	
	private void incrementCounter() {
		long carry = 0;
		
		for (int i = 0; i < counter.length; i++) {
			carry += counter[i] & Utils.INT_MASK;
			carry += increment[i] & Utils.INT_MASK;
			
			counter[i] = (int) carry;
			carry >>= Integer.SIZE;
		}
	}
	
	private void generateKeyStreamBlock() {
		Utils.intBlockToByteBlock(counter, counterBytes, 0, 0, counter.length);
		dataEncryptor.transformBlock(counterBytes, 0, keyStreamBlock, 0);
		
		incrementCounter();
	}
	
	protected byte getKeyStreamByte() {
		if (keyStreamIndex >= blockLength) { 
			generateKeyStreamBlock();
			keyStreamIndex = 0;
		}
		
		final byte result = keyStreamBlock[keyStreamIndex];
		keyStreamIndex++;
		
		return result;
	}
	
	public void setKey(final ISymmetricKey key) {
		this.symmetricKey = key;
	}
}
