package aalgTest;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.math.BigInteger;

import aalg.algebra.DlDomainParameters;
import aalg.algebra.IField;
import aalg.algebra.IGroup;
import aalg.algebra.PrimeField;
import aalg.algebra.PrimeFieldElement;
import aalg.algebra.SubGroupDefinition;
import aalg.crypto.AesCipher;
import aalg.crypto.ByteArraySymmetricKey;
import aalg.crypto.CbcStreamCipher;
import aalg.crypto.DHStreamCipher;
import aalg.crypto.DlKeyGenerator;
import aalg.crypto.HashElementKeyExtractor;
import aalg.crypto.IBlockCipher;
import aalg.crypto.IElementKeyExtractor;
import aalg.crypto.IHashFunction;
import aalg.crypto.IInputStreamFilter;
import aalg.crypto.IOutputStreamFilter;
import aalg.crypto.IPadding;
import aalg.crypto.IStreamDecryptor;
import aalg.crypto.IStreamEncryptor;
import aalg.crypto.Pkcs7Padding;
import aalg.crypto.Sha512;
import aalg.crypto.TestRandomGenerator;

public class DHStreamCipherTest  extends PipeTest {

	private final TestRandomGenerator rng;
	private final ByteArraySymmetricKey key;
	private final AesCipher blockCipher;
	private final IPadding padding;
	private final DHStreamCipher<PrimeFieldElement> dhCipher;

	public DHStreamCipherTest() {
		rng = new TestRandomGenerator();
		
		final byte[] keyBytes = new byte[AesCipher.BLOCK_LENGTH];
		rng.fillArray(keyBytes, 0, keyBytes.length);
		
		key = new ByteArraySymmetricKey(keyBytes);
		
		blockCipher = new AesCipher(AesCipher.BLOCK_LENGTH);
		padding = new Pkcs7Padding();
		
		final PrimeGroupCreator creator = new PrimeGroupCreator();
		final DlDomainParameters<PrimeFieldElement> domainParameters = creator.createGroup();
		final SubGroupDefinition<PrimeFieldElement> subGroupDefinition = domainParameters.getFullSubGroup();
		
		final DlKeyGenerator<PrimeFieldElement> keyGenerator = new DlKeyGenerator<PrimeFieldElement>(subGroupDefinition, rng);
		keyGenerator.generateKeys();
		
		final CbcStreamCipher cipher = new CbcStreamCipher(blockCipher, padding, rng);
		final IHashFunction hashFunction = new Sha512();
		
		final IElementKeyExtractor<PrimeFieldElement> keyExtractor = new HashElementKeyExtractor<PrimeFieldElement>(subGroupDefinition.getGroup(), hashFunction);
		
		dhCipher = new DHStreamCipher<PrimeFieldElement>(subGroupDefinition, rng, keyExtractor, cipher);
		dhCipher.setPublicKey(keyGenerator.getPublicKey());
		dhCipher.setPrivateKey(keyGenerator.getPrivateKey());
	}
	
	@Override
	protected IInputStreamFilter getInputFilter(final ByteArrayInputStream inputStream) throws IOException {
		final IStreamDecryptor decryptor = dhCipher.getDecryptor();
		decryptor.initialize(inputStream);
		
		return decryptor;
	}

	@Override
	protected void flushInputFilter(final IInputStreamFilter filter) throws IOException {
		((IStreamDecryptor) filter).finish();
	}

	@Override
	protected IOutputStreamFilter getOutputFilter(final ByteArrayOutputStream outputStream) throws IOException {
		final IStreamEncryptor encryptor = dhCipher.getEncryptor();
		encryptor.initialize(outputStream);
		
		return encryptor;
	}

	@Override
	protected void flushOutputFilter(final IOutputStreamFilter filter) throws IOException {
		((IStreamEncryptor) filter).finish();
	}

}
