--- a/src/ch/ethz/ssh2/crypto/dh/DhExchange.java	Wed Jul 30 14:21:50 2014 -0700
+++ b/src/ch/ethz/ssh2/crypto/dh/DhExchange.java	Wed Jul 30 14:29:39 2014 -0700
@@ -1,193 +1,140 @@
-/*
- * Copyright (c) 2006-2013 Christian Plattner. All rights reserved.
- * Please refer to the LICENSE.txt for licensing details.
+/**
+ *
  */
 package ch.ethz.ssh2.crypto.dh;
 
 import java.io.IOException;
 import java.math.BigInteger;
-import java.security.SecureRandom;
+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 ch.ethz.ssh2.crypto.digest.HashForSSH2Types;
-import ch.ethz.ssh2.log.Logger;
-import ch.ethz.ssh2.util.StringEncoder;
+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 Christian Plattner
- * @version $Id: DhExchange.java 152 2014-04-28 11:02:23Z dkocher@sudo.ch $
+ * @author kenny
+ *
  */
-public class DhExchange {
-    private static final Logger log = Logger.getLogger(DhExchange.class);
+public class DhExchange extends GenericDhExchange {
 
     /* Given by the standard */
 
-    static final BigInteger p1, p14;
-    static final BigInteger g;
+    private static final BigInteger P1 = new BigInteger(
+        "FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"
+        + "29024E088A67CC74020BBEA63B139B22514A08798E3404DD"
+        + "EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"
+        + "E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"
+        + "EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE65381"
+        + "FFFFFFFFFFFFFFFF", 16);
 
-    BigInteger p;
+    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 */
 
-    BigInteger e;
-    BigInteger x;
-
-    /* Server public and private */
+    private DHPrivateKey clientPrivate;
+    private DHPublicKey clientPublic;
 
-    BigInteger f;
-    BigInteger y;
+    /* Server public */
 
-    /* Shared secret */
-
-    BigInteger k;
+    private DHPublicKey serverPublic;
 
-    static {
-        final String p1_string = "17976931348623159077083915679378745319786029604875"
-                                 + "60117064444236841971802161585193689478337958649255415021805654859805036464"
-                                 + "40548199239100050792877003355816639229553136239076508735759914822574862575"
-                                 + "00742530207744771258955095793777842444242661733472762929938766870920560605"
-                                 + "0270810842907692932019128194467627007";
-        final String p14_string = "FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129"
-                                  + "024E088A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B302B0"
-                                  + "A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9A637ED6B0BFF5CB"
-                                  + "6F406B7EDEE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3DC2007CB8A"
-                                  + "163BF0598DA48361C55D39A69163FA8FD24CF5F83655D23DCA3AD961C62F356208"
-                                  + "552BB9ED529077096966D670C354E4ABC9804F1746C08CA18217C32905E462E36C"
-                                  + "E3BE39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9DE2BCBF69558171"
-                                  + "83995497CEA956AE515D2261898FA051015728E5A8AACAA68FFFFFFFFFFFFFFFF";
-        p1 = new BigInteger(p1_string);
-        p14 = new BigInteger(p14_string, 16);
-        g = new BigInteger("2");
-    }
+    @Override
+    public void init(String name) throws IOException {
+        final DHParameterSpec spec;
 
-    public DhExchange() {
-    }
-
-    public void clientInit(int group, SecureRandom rnd) {
-        k = null;
-
-        if (group == 1) {
-            p = p1;
+        if ("diffie-hellman-group1-sha1".equals(name)) {
+            spec = new DHParameterSpec(P1, G);
         }
-        else if (group == 14) {
-            p = p14;
+        else if ("diffie-hellman-group14-sha1".equals(name)) {
+            spec = new DHParameterSpec(P14, G);
         }
         else {
-            throw new IllegalArgumentException("Unknown DH group " + group);
-        }
-
-        while (true) {
-            x = new BigInteger(p.bitLength() - 1, rnd);
-
-            if (x.compareTo(BigInteger.ONE) > 0) {
-                break;
-            }
+            throw new IllegalArgumentException("Unknown DH group " + name);
         }
 
-        e = g.modPow(x, p);
-    }
-
-    public void serverInit(int group, SecureRandom rnd) {
-        k = null;
-
-        if (group == 1) {
-            p = p1;
+        try {
+            KeyPairGenerator kpg = KeyPairGenerator.getInstance("DH");
+            kpg.initialize(spec);
+            KeyPair pair = kpg.generateKeyPair();
+            clientPrivate = (DHPrivateKey) pair.getPrivate();
+            clientPublic = (DHPublicKey) pair.getPublic();
         }
-        else if (group == 14) {
-            p = p14;
+        catch (NoSuchAlgorithmException e) {
+            throw(IOException) new IOException("No DH keypair generator").initCause(e);
         }
-        else {
-            throw new IllegalArgumentException("Unknown DH group " + group);
+        catch (InvalidAlgorithmParameterException e) {
+            throw(IOException) new IOException("Invalid DH parameters").initCause(e);
         }
-
-        y = new BigInteger(p.bitLength() - 1, rnd);
-        f = g.modPow(y, p);
     }
 
-    /**
-     * @return Returns the e.
-     * @throws IllegalStateException
-     */
-    public BigInteger getE() {
-        if (e == null) {
-            throw new IllegalStateException("DhDsaExchange not initialized!");
-        }
+    @Override
+    public byte[] getE() {
+        if (clientPublic == null)
+            throw new IllegalStateException("DhExchange not initialized!");
 
-        return e;
+        return clientPublic.getY().toByteArray();
+    }
+
+    @Override
+    protected byte[] getServerE() {
+        if (serverPublic == null)
+            throw new IllegalStateException("DhExchange not initialized!");
+
+        return serverPublic.getY().toByteArray();
     }
 
-    /**
-     * @return Returns the f.
-     * @throws IllegalStateException
-     */
-    public BigInteger getF() {
-        if (f == null) {
-            throw new IllegalStateException("DhDsaExchange not initialized!");
-        }
+    @Override
+    public void setF(byte[] f) throws IOException {
+        if (clientPublic == null)
+            throw new IllegalStateException("DhExchange not initialized!");
 
-        return f;
-    }
+        final KeyAgreement ka;
 
-    /**
-     * @return Returns the shared secret k.
-     * @throws IllegalStateException
-     */
-    public BigInteger getK() {
-        if (k == null) {
-            throw new IllegalStateException("Shared secret not yet known, need f first!");
+        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);
         }
-
-        return k;
-    }
-
-    /**
-     * @param f
-     */
-    public void setF(BigInteger f) {
-        if (e == null) {
-            throw new IllegalStateException("DhDsaExchange not initialized!");
+        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);
         }
 
-        if (BigInteger.ZERO.compareTo(f) >= 0 || p.compareTo(f) <= 0) {
-            throw new IllegalArgumentException("Invalid f specified!");
-        }
-
-        this.f = f;
-        this.k = f.modPow(x, p);
+        sharedSecret = new BigInteger(ka.generateSecret());
     }
 
-    /**
-     * @param e
-     */
-    public void setE(BigInteger e) {
-        if (f == null) {
-            throw new IllegalStateException("DhDsaExchange not initialized!");
-        }
-
-        if (BigInteger.ZERO.compareTo(e) >= 0 || p.compareTo(e) <= 0) {
-            throw new IllegalArgumentException("Invalid e specified!");
-        }
-
-        this.e = e;
-        this.k = e.modPow(y, p);
-    }
-
-    public byte[] calculateH(byte[] clientversion, byte[] serverversion, byte[] clientKexPayload,
-                             byte[] serverKexPayload, byte[] hostKey) throws IOException {
-        HashForSSH2Types hash = new HashForSSH2Types("SHA1");
-
-        if (log.isInfoEnabled()) {
-            log.info("Client: '" + StringEncoder.GetString(clientversion) + "'");
-            log.info("Server: '" + StringEncoder.GetString(serverversion) + "'");
-        }
-
-        hash.updateByteString(clientversion);
-        hash.updateByteString(serverversion);
-        hash.updateByteString(clientKexPayload);
-        hash.updateByteString(serverKexPayload);
-        hash.updateByteString(hostKey);
-        hash.updateBigInt(e);
-        hash.updateBigInt(f);
-        hash.updateBigInt(k);
-        return hash.getDigest();
+    @Override
+    public String getHashAlgo() {
+        return "SHA1";
     }
 }