package aalg.crypto;

import java.io.IOException;
import java.io.OutputStream;

import aalg.algebra.IElement;
import aalg.algebra.IGroup;
import aalg.algebra.SubGroupDefinition;

public class DHStreamEncryptor<T extends IElement> implements IStreamEncryptor {
	
	private final OutputStream plainTextStream = new OutputStream() {
		@Override
		public void write(final int b) throws IOException {
			encryptor.getSourceStream().write(b);
		}
	};
	
	private OutputStream cipherTextStream;
	private final DlKeyGenerator<T> keyGenerator;
	private final SubGroupDefinition<T> subGroupDefinition;
	private final DHKeyExchange<T> keyExchange;
	private final DlPublicKey<T> publicKey;
	private final IElementKeyExtractor<T> keyExtractor;
	private final ISymmetricStreamCipher cipher;
	private ISymmetricStreamEncryptor encryptor;
		
	public DHStreamEncryptor(final SubGroupDefinition<T> subGroupDefinition, final IRandomGenerator rng, final IElementKeyExtractor<T> keyExtractor, final ISymmetricStreamCipher cipher, final DlPublicKey<T> publicKey) {
		this.keyGenerator = new DlKeyGenerator<T>(subGroupDefinition, rng);
		this.subGroupDefinition = subGroupDefinition;
		this.keyExchange = new DHKeyExchange<T>(subGroupDefinition);
		this.publicKey = publicKey;
		this.keyExtractor = keyExtractor;
		this.cipher = cipher;
	}
	
	@Override
	public OutputStream getSourceStream() throws IOException {
		return plainTextStream;
	}

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

	@Override
	public void initialize(final OutputStream cipherTextStream) throws IOException {
		this.cipherTextStream = cipherTextStream;
		
		keyGenerator.generateKeys();
		final DlPrivateKey randomPrivateKey = keyGenerator.getPrivateKey();
		final DlPublicKey<T> randomPublicKey = keyGenerator.getPublicKey();
		
		final IGroup<T> group = subGroupDefinition.getGroup();
		randomPublicKey.writeToStream(group, cipherTextStream);
		
		final T sharedElement = keyExchange.calculateSharedKey(publicKey, randomPrivateKey);
		final byte[] sharedKey = keyExtractor.extractKey(sharedElement, 0, cipher.getKeySize());

		encryptor = cipher.getEncryptor();
		encryptor.setKey(new ByteArraySymmetricKey(sharedKey));
		encryptor.initialize(cipherTextStream);
	}

	@Override
	public void finish() throws IOException {
		encryptor.finish();
	}

}
