package aalgTest;

import java.io.IOException;
import java.math.BigInteger;
import org.junit.Assert;
import org.junit.Test;

import aalg.algebra.DlDomainParameters;
import aalg.algebra.IGroup;
import aalg.algebra.IElement;
import aalg.algebra.SubGroupDefinition;
import aalg.crypto.DHKeyExchange;
import aalg.crypto.DlKeyGenerator;
import aalg.crypto.DlPrivateKey;
import aalg.crypto.DlPublicKey;
import aalg.crypto.IRandomGenerator;
import aalg.crypto.TestRandomGenerator;

public class DHTest {

	private <T extends IElement> void dhTest (final IGroupCreator<T> creator) throws IOException {
		final DlDomainParameters<T> domainParameters = creator.createGroup();
		final SubGroupDefinition<T> subGroupDefinition = domainParameters.getFullSubGroup();
		
		final IRandomGenerator rng = new TestRandomGenerator();
		
		final DlKeyGenerator<T> keyGenerator = new DlKeyGenerator<T>(subGroupDefinition, rng);
		
		keyGenerator.generateKeys();
		
		final DlPrivateKey privA = keyGenerator.getPrivateKey();
		final DlPublicKey<T> pubA = keyGenerator.getPublicKey();

		keyGenerator.generateKeys();
		
		final DlPrivateKey privB = keyGenerator.getPrivateKey();
		final DlPublicKey<T> pubB = keyGenerator.getPublicKey();
		
		final DHKeyExchange<T> dh = new DHKeyExchange<T>(subGroupDefinition);
		final T keyA = dh.calculateSharedKey(pubB, privA);
		final T keyB = dh.calculateSharedKey(pubA, privB);
				
		Assert.assertEquals(keyA, keyB);
	}
	
	@Test
	public void testDHPrimeField() throws IOException {
		final PrimeGroupCreator creator = new PrimeGroupCreator();
		
		dhTest (creator);
	}

	@Test
	public void testDHMultiGroup() throws IOException {
		final MultiGroupCreator creator = new MultiGroupCreator(BigInteger.valueOf(65537));
		
		dhTest (creator);
	}

}
