/**
 *
 */
package com.trilead.ssh2.crypto.dh;

import java.io.IOException;
import java.math.BigInteger;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;

import javax.crypto.KeyAgreement;
import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;
import javax.crypto.spec.DHPublicKeySpec;

/**
 * @author kenny
 *
 */
public class DhExchange extends GenericDhExchange {

	/* Given by the standard */

	private static final BigInteger P1 = new BigInteger(
			  "FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"
			+ "29024E088A67CC74020BBEA63B139B22514A08798E3404DD"
			+ "EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"
			+ "E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"
			+ "EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE65381"
			+ "FFFFFFFFFFFFFFFF", 16);

	private static final BigInteger P14 = new BigInteger(
			  "FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"
			+ "29024E088A67CC74020BBEA63B139B22514A08798E3404DD"
			+ "EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"
			+ "E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"
			+ "EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3D"
			+ "C2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F"
			+ "83655D23DCA3AD961C62F356208552BB9ED529077096966D"
			+ "670C354E4ABC9804F1746C08CA18217C32905E462E36CE3B"
			+ "E39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9"
			+ "DE2BCBF6955817183995497CEA956AE515D2261898FA0510"
			+ "15728E5A8AACAA68FFFFFFFFFFFFFFFF", 16);

	private static final BigInteger G = BigInteger.valueOf(2);

	/* Client public and private */

	private DHPrivateKey clientPrivate;
	private DHPublicKey clientPublic;

	/* Server public */

	private DHPublicKey serverPublic;

	@Override
	public void init(String name) throws IOException {
		final DHParameterSpec spec;
		if ("diffie-hellman-group1-sha1".equals(name)) {
			spec = new DHParameterSpec(P1, G);
		} else if ("diffie-hellman-group14-sha1".equals(name)) {
			spec = new DHParameterSpec(P14, G);
		} else {
			throw new IllegalArgumentException("Unknown DH group " + name);
		}

		try {
			KeyPairGenerator kpg = KeyPairGenerator.getInstance("DH");
			kpg.initialize(spec);
			KeyPair pair = kpg.generateKeyPair();
			clientPrivate = (DHPrivateKey) pair.getPrivate();
			clientPublic = (DHPublicKey) pair.getPublic();
		} catch (NoSuchAlgorithmException e) {
			throw (IOException) new IOException("No DH keypair generator").initCause(e);
		} catch (InvalidAlgorithmParameterException e) {
			throw (IOException) new IOException("Invalid DH parameters").initCause(e);
		}
	}

	@Override
	public byte[] getE() {
		if (clientPublic == null)
			throw new IllegalStateException("DhExchange not initialized!");

		return clientPublic.getY().toByteArray();
	}

	@Override
	protected byte[] getServerE() {
		if (serverPublic == null)
			throw new IllegalStateException("DhExchange not initialized!");

		return serverPublic.getY().toByteArray();
	}

	@Override
	public void setF(byte[] f) throws IOException {
		if (clientPublic == null)
			throw new IllegalStateException("DhExchange not initialized!");

		final KeyAgreement ka;
		try {
			KeyFactory kf = KeyFactory.getInstance("DH");
			DHParameterSpec params = clientPublic.getParams();
			this.serverPublic = (DHPublicKey) kf.generatePublic(new DHPublicKeySpec(
					new BigInteger(f), params.getP(), params.getG()));

			ka = KeyAgreement.getInstance("DH");
			ka.init(clientPrivate);
			ka.doPhase(serverPublic, true);
		} catch (NoSuchAlgorithmException e) {
			throw (IOException) new IOException("No DH key agreement method").initCause(e);
		} catch (InvalidKeyException e) {
			throw (IOException) new IOException("Invalid DH key").initCause(e);
		} catch (InvalidKeySpecException e) {
			throw (IOException) new IOException("Invalid DH key").initCause(e);
		}

		sharedSecret = new BigInteger(ka.generateSecret());
	}

	@Override
	public String getHashAlgo() {
		return "SHA1";
	}
}