package aalg.crypto;

import java.io.IOException;
import java.io.OutputStream;
import java.util.Arrays;

public class CbcStreamEncryptor implements ISymmetricStreamEncryptor {
	
	private final IBlockCipherTransformation blockEncryptor;
	private final IPaddingWriter paddingWriter;
	private final IRandomGenerator rng;
	private final byte[] plainTextBlock;
	private int plainTextByteCount;
	private final byte[] lastCipherTextBlock;
	private OutputStream cipherTextStream;
	
	private final OutputStream blockEncryptingStream = new OutputStream() {
		@Override
		public void write(final int b) throws IOException {
			plainTextBlock[plainTextByteCount] = (byte) (b ^ lastCipherTextBlock[plainTextByteCount]);
			plainTextByteCount++;
			
			if (plainTextByteCount >= plainTextBlock.length) {
				blockEncryptor.transformBlock(plainTextBlock, 0, lastCipherTextBlock, 0);
				cipherTextStream.write(lastCipherTextBlock);
				
				plainTextByteCount = 0;
			}
		}
		
	};

	public CbcStreamEncryptor(final IBlockCipherTransformation blockEncryptor, final IPaddingWriter paddingWriter, final IRandomGenerator rng) {
		this.blockEncryptor = blockEncryptor;
		this.paddingWriter = paddingWriter;
		this.rng = rng;
		
		final int blockLength = blockEncryptor.getCipher().getBlockLength();
		this.lastCipherTextBlock = new byte[blockLength];
		this.plainTextBlock = new byte[blockLength];
	}

	@Override
	public OutputStream getSourceStream() throws IOException {
		return paddingWriter.getSourceStream();
	}

	@Override
	public OutputStream getTargetStream() throws IOException {
		return cipherTextStream;
	}

	@Override
	public void initialize(final OutputStream cipherTextStream) throws IOException {
		this.cipherTextStream = cipherTextStream;
		this.paddingWriter.initialize(blockEncryptingStream, lastCipherTextBlock.length);
		
		// Write IV
		rng.fillArray(lastCipherTextBlock, 0, lastCipherTextBlock.length);
		cipherTextStream.write(lastCipherTextBlock);
		
		plainTextByteCount = 0;
	}

	@Override
	public void finish() throws IOException {
		paddingWriter.finish();
		
		// Clear sensitive data
		Arrays.fill(lastCipherTextBlock, (byte) 0);
	}

	@Override
	public void setKey(final ISymmetricKey key) {
		blockEncryptor.setKey(key);
	}

}
