package aalg.crypto;

import java.io.IOException;
import java.io.InputStream;

import aalg.algebra.Utils;

public class CbcStreamDecryptor implements ISymmetricStreamDecryptor {
	
	private final IBlockCipherTransformation blockDecryptor;
	private final IPaddingReader paddingReader;
	private InputStream cipherTextStream;
	private final int blockLength;
	private final byte[] plainTextBlock;
	private int plainTextOffset;
	private final byte[] currentCipherTextBlock;
	private final byte[] lastCipherTextBlock;

	
	private final InputStream blockDecryptingStream = new InputStream() {
		@Override
		public int read() throws IOException {
			if (plainTextOffset >= plainTextBlock.length) {
				if (Utils.readByteArray (cipherTextStream, currentCipherTextBlock, 0, blockLength) < blockLength)
					return -1;
					
				blockDecryptor.transformBlock (currentCipherTextBlock, 0, plainTextBlock, 0);
				
				for (int i = 0; i < plainTextBlock.length; i++)
					plainTextBlock[i] ^= lastCipherTextBlock[i];
				
				System.arraycopy(currentCipherTextBlock, 0, lastCipherTextBlock, 0, blockLength);
				plainTextOffset = 0;
			}
			
			final byte b = plainTextBlock[plainTextOffset];
			plainTextOffset++;
			
			return b & 0xFF;
		}
	};

	public CbcStreamDecryptor(final IBlockCipherTransformation blockDecryptor, final IPaddingReader paddingReader) {
		this.blockDecryptor = blockDecryptor;
		this.paddingReader = paddingReader;
		
		this.blockLength = blockDecryptor.getCipher().getBlockLength();
		this.currentCipherTextBlock = new byte[blockLength];
		this.lastCipherTextBlock = new byte[blockLength];
		this.plainTextBlock = new byte[blockLength];
	}

	@Override
	public void initialize(final InputStream cipherTextStream) throws IOException {
		this.cipherTextStream = cipherTextStream;
		
		paddingReader.initialize(blockDecryptingStream, plainTextBlock.length);
		
		if (Utils.readByteArray (cipherTextStream, lastCipherTextBlock, 0, blockLength) < blockLength)
			throw new RuntimeException("Cannot read initialization vector");
		
		plainTextOffset = blockLength;
	}

	@Override
	public InputStream getTargetStream() throws IOException {
		return paddingReader.getTargetStream();
	}

	@Override
	public void finish() {
	}

	@Override
	public InputStream getSourceStream() throws IOException {
		return cipherTextStream;
	}
	
	@Override
	public void setKey(final ISymmetricKey key) {
		blockDecryptor.setKey(key);
	}

}
