Added support for ecdh-sha2-nistp256 and ecdh-sha2-nistp384 key exchange algorithms

src/main/java/net/schmizz/sshj/DefaultConfig.java

@@ -31,10 +31,7 @@ import net.schmizz.sshj.transport.cipher.BlowfishCBC;
 import net.schmizz.sshj.transport.cipher.Cipher;
 import net.schmizz.sshj.transport.cipher.TripleDESCBC;
 import net.schmizz.sshj.transport.compression.NoneCompression;
-import net.schmizz.sshj.transport.kex.DHG1;
-import net.schmizz.sshj.transport.kex.DHG14;
-import net.schmizz.sshj.transport.kex.DHGexSHA1;
-import net.schmizz.sshj.transport.kex.DHGexSHA256;
+import net.schmizz.sshj.transport.kex.*;
 import net.schmizz.sshj.transport.mac.HMACMD5;
 import net.schmizz.sshj.transport.mac.HMACMD596;
 import net.schmizz.sshj.transport.mac.HMACSHA1;
@@ -100,7 +97,13 @@ public class DefaultConfig
 
     protected void initKeyExchangeFactories(boolean bouncyCastleRegistered) {
         if (bouncyCastleRegistered)
-            setKeyExchangeFactories(new DHG14.Factory(), new DHG1.Factory(), new DHGexSHA1.Factory(), new DHGexSHA256.Factory());
+            setKeyExchangeFactories(new DHG14.Factory(),
+                    new DHG1.Factory(),
+                    new DHGexSHA1.Factory(),
+                    new DHGexSHA256.Factory(),
+                    new ECDHNistP.Factory256(),
+                    new ECDHNistP.Factory384());
+                    // TODO 521 fails sometimes with key verification errors new ECDHNistP.Factory521());
         else
             setKeyExchangeFactories(new DHG1.Factory(), new DHGexSHA1.Factory());
     }

src/main/java/net/schmizz/sshj/common/Message.java

@@ -30,7 +30,7 @@ public enum Message {
 
     KEXDH_INIT(30),
 
-    /** { KEXDH_REPLY, KEXDH_GEX_GROUP } */
+    /** { KEXDH_REPLY, KEXDH_GEX_GROUP, SSH_MSG_KEX_ECDH_REPLY } */
     KEXDH_31(31),
 
     KEX_DH_GEX_INIT(32),

src/main/java/net/schmizz/sshj/transport/digest/SHA384.java

@@ -0,0 +1,23 @@
+package net.schmizz.sshj.transport.digest;
+
+public class SHA384 extends BaseDigest {
+/** Named factory for SHA384 digest */
+public static class Factory
+        implements net.schmizz.sshj.common.Factory.Named<Digest> {
+
+    @Override
+    public Digest create() {
+        return new SHA384();
+    }
+
+    @Override
+    public String getName() {
+        return "sha384";
+    }
+}
+
+    /** Create a new instance of a SHA384 digest */
+    public SHA384() {
+        super("SHA-384", 48);
+    }
+}

src/main/java/net/schmizz/sshj/transport/digest/SHA512.java

@@ -0,0 +1,23 @@
+package net.schmizz.sshj.transport.digest;
+
+public class SHA512 extends BaseDigest {
+/** Named factory for SHA384 digest */
+public static class Factory
+        implements net.schmizz.sshj.common.Factory.Named<Digest> {
+
+    @Override
+    public Digest create() {
+        return new SHA512();
+    }
+
+    @Override
+    public String getName() {
+        return "sha512";
+    }
+}
+
+    /** Create a new instance of a SHA384 digest */
+    public SHA512() {
+        super("SHA-512", 64);
+    }
+}

src/main/java/net/schmizz/sshj/transport/kex/AbstractDH.java

@@ -0,0 +1,41 @@
+package net.schmizz.sshj.transport.kex;
+
+import net.schmizz.sshj.transport.digest.Digest;
+
+import java.math.BigInteger;
+import java.security.PublicKey;
+import java.util.Arrays;
+
+public abstract class AbstractDH extends KeyExchangeBase {
+    protected final Digest digest;
+    protected final DHBase dh;
+
+    protected byte[] H;
+    protected PublicKey hostKey;
+
+    public AbstractDH(DHBase dh, Digest digest) {
+        this.dh = dh;
+        this.digest = digest;
+    }
+
+    @Override
+    public byte[] getH() {
+        return Arrays.copyOf(H, H.length);
+    }
+
+    @Override
+    public BigInteger getK() {
+        return dh.getK();
+    }
+
+    @Override
+    public Digest getHash() {
+        return digest;
+    }
+
+    @Override
+    public PublicKey getHostKey() {
+        return hostKey;
+    }
+
+}

src/main/java/net/schmizz/sshj/transport/kex/AbstractDHG.java

@@ -38,46 +38,24 @@ import java.util.Arrays;
  * Base class for DHG key exchange algorithms. Implementations will only have to configure the required data on the
  * {@link DH} class in the
  */
-public abstract class AbstractDHG extends KeyExchangeBase
+public abstract class AbstractDHG extends AbstractDH
         implements KeyExchange {
 
     private final Logger log = LoggerFactory.getLogger(getClass());
 
-    private final Digest sha1 = new SHA1();
-    private final DH dh = new DH();
-
-    private byte[] H;
-    private PublicKey hostKey;
-
-    @Override
-    public byte[] getH() {
-        return Arrays.copyOf(H, H.length);
-    }
-
-    @Override
-    public BigInteger getK() {
-        return dh.getK();
-    }
-
-    @Override
-    public Digest getHash() {
-        return sha1;
-    }
-
-    @Override
-    public PublicKey getHostKey() {
-        return hostKey;
+    public AbstractDHG(DHBase dhBase, Digest digest) {
+        super(dhBase, digest);
     }
 
     @Override
     public void init(Transport trans, String V_S, String V_C, byte[] I_S, byte[] I_C)
             throws GeneralSecurityException, TransportException {
         super.init(trans, V_S, V_C, I_S, I_C);
-        sha1.init();
+        digest.init();
         initDH(dh);
 
         log.debug("Sending SSH_MSG_KEXDH_INIT");
-        trans.write(new SSHPacket(Message.KEXDH_INIT).putMPInt(dh.getE()));
+        trans.write(new SSHPacket(Message.KEXDH_INIT).putBytes(dh.getE()));
     }
 
     @Override
@@ -88,11 +66,11 @@ public abstract class AbstractDHG extends KeyExchangeBase
 
         log.debug("Received SSH_MSG_KEXDH_REPLY");
         final byte[] K_S;
-        final BigInteger f;
+        final byte[] f;
         final byte[] sig; // signature sent by server
         try {
             K_S = packet.readBytes();
-            f = packet.readMPInt();
+            f = packet.readBytes();
             sig = packet.readBytes();
             hostKey = new Buffer.PlainBuffer(K_S).readPublicKey();
         } catch (Buffer.BufferException be) {
@@ -103,11 +81,11 @@ public abstract class AbstractDHG extends KeyExchangeBase
 
         final Buffer.PlainBuffer buf = initializedBuffer()
                 .putString(K_S)
-                .putMPInt(dh.getE())
-                .putMPInt(f)
+                .putBytes(dh.getE())
+                .putBytes(f)
                 .putMPInt(dh.getK());
-        sha1.update(buf.array(), buf.rpos(), buf.available());
-        H = sha1.digest();
+        digest.update(buf.array(), buf.rpos(), buf.available());
+        H = digest.digest();
 
         Signature signature = Factory.Named.Util.create(trans.getConfig().getSignatureFactories(),
                                                         KeyType.fromKey(hostKey).toString());
@@ -119,7 +97,7 @@ public abstract class AbstractDHG extends KeyExchangeBase
         return true;
     }
 
-    protected abstract void initDH(DH dh)
+    protected abstract void initDH(DHBase dh)
             throws GeneralSecurityException;
 
 }

src/main/java/net/schmizz/sshj/transport/kex/AbstractDHGex.java

@@ -8,58 +8,30 @@ import net.schmizz.sshj.transport.digest.Digest;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.crypto.spec.DHParameterSpec;
 import java.math.BigInteger;
 import java.security.GeneralSecurityException;
-import java.security.PublicKey;
-import java.util.Arrays;
 
-public abstract class AbstractDHGex extends KeyExchangeBase {
+public abstract class AbstractDHGex extends AbstractDH {
     private final Logger log = LoggerFactory.getLogger(getClass());
 
-    private Digest digest;
-
     private int minBits = 1024;
     private int maxBits = 8192;
     private int preferredBits = 2048;
 
-    private DH dh;
-    private PublicKey hostKey;
-    private byte[] H;
-
     public AbstractDHGex(Digest digest) {
-        this.digest = digest;
+        super(new DH(), digest);
     }
 
     @Override
     public void init(Transport trans, String V_S, String V_C, byte[] I_S, byte[] I_C) throws GeneralSecurityException, TransportException {
         super.init(trans, V_S, V_C, I_S, I_C);
-        dh = new DH();
         digest.init();
 
         log.debug("Sending {}", Message.KEX_DH_GEX_REQUEST);
         trans.write(new SSHPacket(Message.KEX_DH_GEX_REQUEST).putUInt32(minBits).putUInt32(preferredBits).putUInt32(maxBits));
     }
 
-    @Override
-    public byte[] getH() {
-        return Arrays.copyOf(H, H.length);
-    }
-
-    @Override
-    public BigInteger getK() {
-        return dh.getK();
-    }
-
-    @Override
-    public Digest getHash() {
-        return digest;
-    }
-
-    @Override
-    public PublicKey getHostKey() {
-        return hostKey;
-    }
-
     @Override
     public boolean next(Message msg, SSHPacket buffer) throws GeneralSecurityException, TransportException {
         log.debug("Got message {}", msg);
@@ -78,7 +50,7 @@ public abstract class AbstractDHGex extends KeyExchangeBase {
 
     private boolean parseGexReply(SSHPacket buffer) throws Buffer.BufferException, GeneralSecurityException, TransportException {
         byte[] K_S = buffer.readBytes();
-        BigInteger f = buffer.readMPInt();
+        byte[] f = buffer.readBytes();
         byte[] sig = buffer.readBytes();
         hostKey = new Buffer.PlainBuffer(K_S).readPublicKey();
 
@@ -90,10 +62,10 @@ public abstract class AbstractDHGex extends KeyExchangeBase {
                 .putUInt32(minBits)
                 .putUInt32(preferredBits)
                 .putUInt32(maxBits)
-                .putMPInt(dh.getP())
-                .putMPInt(dh.getG())
-                .putMPInt(dh.getE())
-                .putMPInt(f)
+                .putMPInt(((DH) dh).getP())
+                .putMPInt(((DH) dh).getG())
+                .putBytes(dh.getE())
+                .putBytes(f)
                 .putMPInt(k);
         digest.update(buf.array(), buf.rpos(), buf.available());
         H = digest.digest();
@@ -116,9 +88,9 @@ public abstract class AbstractDHGex extends KeyExchangeBase {
             throw new GeneralSecurityException("Server generated gex p is out of range (" + bitLength + " bits)");
         }
         log.debug("Received server p bitlength {}", bitLength);
-        dh.init(p, g);
+        dh.init(new DHParameterSpec(p, g));
         log.debug("Sending {}", Message.KEX_DH_GEX_INIT);
-        trans.write(new SSHPacket(Message.KEX_DH_GEX_INIT).putMPInt(dh.getE()));
+        trans.write(new SSHPacket(Message.KEX_DH_GEX_INIT).putBytes(dh.getE()));
         return false;
     }
 }

src/main/java/net/schmizz/sshj/transport/kex/DH.java

@@ -18,59 +18,44 @@ package net.schmizz.sshj.transport.kex;
 import net.schmizz.sshj.common.SSHRuntimeException;
 import net.schmizz.sshj.common.SecurityUtils;
 
-import javax.crypto.KeyAgreement;
 import javax.crypto.spec.DHParameterSpec;
 import javax.crypto.spec.DHPublicKeySpec;
 import java.math.BigInteger;
 import java.security.GeneralSecurityException;
 import java.security.KeyFactory;
 import java.security.KeyPair;
-import java.security.KeyPairGenerator;
 import java.security.PublicKey;
+import java.security.spec.AlgorithmParameterSpec;
 
 /** Diffie-Hellman key generator. */
-public class DH {
+public class DH extends DHBase {
 
     private BigInteger p;
     private BigInteger g;
-    private BigInteger e; // my public key
-    private BigInteger K; // shared secret key
-    private final KeyPairGenerator generator;
-    private final KeyAgreement agreement;
 
     public DH() {
-        try {
-            generator = SecurityUtils.getKeyPairGenerator("DH");
-            agreement = SecurityUtils.getKeyAgreement("DH");
-        } catch (GeneralSecurityException e) {
-            throw new SSHRuntimeException(e);
-        }
+        super("DH", "DH");
     }
 
-    public void init(BigInteger p, BigInteger g)
-            throws GeneralSecurityException {
-        this.p = p;
-        this.g = g;
-        generator.initialize(new DHParameterSpec(p, g));
+    @Override
+    protected void init(AlgorithmParameterSpec params) throws GeneralSecurityException {
+        if (!(params instanceof DHParameterSpec)) {
+            throw new SSHRuntimeException("Wrong algorithm parameters for Diffie Hellman");
+        }
+        this.p = ((DHParameterSpec) params).getP();
+        this.g = ((DHParameterSpec) params).getG();
+        generator.initialize(params);
         final KeyPair kp = generator.generateKeyPair();
         agreement.init(kp.getPrivate());
-        e = ((javax.crypto.interfaces.DHPublicKey) kp.getPublic()).getY();
+        setE(((javax.crypto.interfaces.DHPublicKey) kp.getPublic()).getY().toByteArray());
     }
 
-    public void computeK(BigInteger f)
-            throws GeneralSecurityException {
+    @Override
+    void computeK(byte[] f) throws GeneralSecurityException {
         final KeyFactory keyFactory = SecurityUtils.getKeyFactory("DH");
-        final PublicKey yourPubKey = keyFactory.generatePublic(new DHPublicKeySpec(f, p, g));
+        final PublicKey yourPubKey = keyFactory.generatePublic(new DHPublicKeySpec(new BigInteger(f), p, g));
         agreement.doPhase(yourPubKey, true);
-        K = new BigInteger(1, agreement.generateSecret());
-    }
-
-    public BigInteger getE() {
-        return e;
-    }
-
-    public BigInteger getK() {
-        return K;
+        setK(new BigInteger(1, agreement.generateSecret()));
     }
 
     public BigInteger getP() {

src/main/java/net/schmizz/sshj/transport/kex/DHBase.java

@@ -0,0 +1,47 @@
+package net.schmizz.sshj.transport.kex;
+
+import net.schmizz.sshj.common.SSHRuntimeException;
+import net.schmizz.sshj.common.SecurityUtils;
+
+import javax.crypto.KeyAgreement;
+import java.math.BigInteger;
+import java.security.GeneralSecurityException;
+import java.security.KeyPairGenerator;
+import java.security.spec.AlgorithmParameterSpec;
+
+abstract class DHBase {
+    protected final KeyPairGenerator generator;
+    protected final KeyAgreement agreement;
+
+    private byte[] e; // my public key
+    private BigInteger K; // shared secret key
+
+    public DHBase(String generator, String agreement) {
+        try {
+            this.generator = SecurityUtils.getKeyPairGenerator(generator);
+            this.agreement = SecurityUtils.getKeyAgreement(agreement);
+        } catch (GeneralSecurityException e) {
+            throw new SSHRuntimeException(e);
+        }
+    }
+
+    abstract void computeK(byte[] f) throws GeneralSecurityException;
+
+    protected abstract void init(AlgorithmParameterSpec params) throws GeneralSecurityException;
+
+    void setE(byte[] e) {
+        this.e = e;
+    }
+
+    void setK(BigInteger k) {
+        K = k;
+    }
+
+    public byte[] getE() {
+        return e;
+    }
+
+    public BigInteger getK() {
+        return K;
+    }
+}

src/main/java/net/schmizz/sshj/transport/kex/DHG1.java

@@ -15,12 +15,17 @@
  */
 package net.schmizz.sshj.transport.kex;
 
+import net.schmizz.sshj.transport.digest.SHA1;
+
+import javax.crypto.spec.DHParameterSpec;
 import java.security.GeneralSecurityException;
 
 /**
  * Diffie-Hellman key exchange with SHA-1 and Oakley Group 2 [RFC2409] (1024-bit MODP Group).
  *
  * @see <a href="http://www.ietf.org/rfc/rfc4253.txt">RFC 4253</a>
+ *
+ * TODO refactor away the (unneeded) class
  */
 public class DHG1
         extends AbstractDHG {
@@ -38,13 +43,14 @@ public class DHG1
         public String getName() {
             return "diffie-hellman-group1-sha1";
         }
+    }
 
+    public DHG1() {
+        super(new DH(), new SHA1());
     }
 
     @Override
-    protected void initDH(DH dh)
-            throws GeneralSecurityException {
-        dh.init(DHGroupData.P1, DHGroupData.G);
+    protected void initDH(DHBase dh) throws GeneralSecurityException {
+        dh.init(new DHParameterSpec(DHGroupData.P1, DHGroupData.G));
     }
-
 }

src/main/java/net/schmizz/sshj/transport/kex/DHG14.java

@@ -15,6 +15,9 @@
  */
 package net.schmizz.sshj.transport.kex;
 
+import net.schmizz.sshj.transport.digest.SHA1;
+
+import javax.crypto.spec.DHParameterSpec;
 import java.security.GeneralSecurityException;
 
 /**
@@ -42,10 +45,12 @@ public class DHG14
 
     }
 
-    @Override
-    protected void initDH(DH dh)
-            throws GeneralSecurityException {
-        dh.init(DHGroupData.P14, DHGroupData.G);
+    public DHG14() {
+        super(new DH(), new SHA1());
     }
 
+    @Override
+    protected void initDH(DHBase dh) throws GeneralSecurityException {
+        dh.init(new DHParameterSpec(DHGroupData.P14, DHGroupData.G));
+    }
 }

src/main/java/net/schmizz/sshj/transport/kex/ECDH.java

@@ -0,0 +1,88 @@
+package net.schmizz.sshj.transport.kex;
+
+import net.schmizz.sshj.common.SSHRuntimeException;
+import net.schmizz.sshj.common.SecurityUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.math.BigInteger;
+import java.security.GeneralSecurityException;
+import java.security.KeyFactory;
+import java.security.KeyPair;
+import java.security.PublicKey;
+import java.security.interfaces.ECPublicKey;
+import java.security.spec.*;
+import java.util.Arrays;
+
+public class ECDH extends DHBase {
+
+    private ECParameterSpec ecParameterSpec;
+
+    public ECDH() {
+        super("EC", "ECDH");
+    }
+
+    protected void init(AlgorithmParameterSpec params) throws GeneralSecurityException {
+        generator.initialize(params);
+        KeyPair keyPair = generator.generateKeyPair();
+        agreement.init(keyPair.getPrivate());
+        ECPublicKey ecPublicKey = (ECPublicKey) keyPair.getPublic();
+        this.ecParameterSpec = ecPublicKey.getParams();
+        ECPoint w = ecPublicKey.getW();
+        byte[] encoded = getEncoded(w, ecParameterSpec.getCurve());
+        setE(encoded);
+    }
+
+    @Override
+    public void computeK(byte[] f) throws GeneralSecurityException {
+        KeyFactory keyFactory = SecurityUtils.getKeyFactory("EC");
+        ECPublicKeySpec keySpec = new ECPublicKeySpec(getDecoded(f, ecParameterSpec.getCurve()), ecParameterSpec);
+        PublicKey yourPubKey = keyFactory.generatePublic(keySpec);
+        agreement.doPhase(yourPubKey, true);
+        setK(new BigInteger(1, agreement.generateSecret()));
+    }
+
+    /**
+     * SECG 2.3.4 Octet String to ECPoint
+     */
+    private static ECPoint getDecoded(byte[] M, EllipticCurve curve) {
+        int elementSize = getElementSize(curve);
+        if (M.length != 2 * elementSize + 1 || M[0] != 0x04) {
+            throw new SSHRuntimeException("Invalid 'f' for Elliptic Curve " + curve.toString());
+        }
+        byte[] xBytes = new byte[elementSize];
+        byte[] yBytes = new byte[elementSize];
+        System.arraycopy(M, 1, xBytes, 0, elementSize);
+        System.arraycopy(M, 1 + elementSize, yBytes, 0, elementSize);
+        return new ECPoint(new BigInteger(1, xBytes), new BigInteger(1, yBytes));
+    }
+
+    /**
+     * SECG 2.3.3 ECPoint to Octet String
+     */
+    private static byte[] getEncoded(ECPoint point, EllipticCurve curve) {
+        int elementSize = getElementSize(curve);
+        byte[] M = new byte[2 * elementSize + 1];
+        M[0] = 0x04;
+
+        byte[] xBytes = stripLeadingZeroes(point.getAffineX().toByteArray());
+        byte[] yBytes = stripLeadingZeroes(point.getAffineY().toByteArray());
+        System.arraycopy(xBytes, 0, M, 1 + elementSize - xBytes.length, xBytes.length);
+        System.arraycopy(yBytes, 0, M, 1 + 2 * elementSize - yBytes.length, yBytes.length);
+        return M;
+    }
+
+    private static byte[] stripLeadingZeroes(byte[] bytes) {
+        int start = 0;
+        while (bytes[start] == 0x0) {
+            start++;
+        }
+
+        return Arrays.copyOfRange(bytes, start, bytes.length);
+    }
+
+    private static int getElementSize(EllipticCurve curve) {
+        int fieldSize = curve.getField().getFieldSize();
+        return (fieldSize + 7) / 8;
+    }
+}

src/main/java/net/schmizz/sshj/transport/kex/ECDHNistP.java

@@ -0,0 +1,71 @@
+package net.schmizz.sshj.transport.kex;
+
+import net.schmizz.sshj.transport.digest.Digest;
+import net.schmizz.sshj.transport.digest.SHA256;
+import net.schmizz.sshj.transport.digest.SHA384;
+import net.schmizz.sshj.transport.digest.SHA512;
+import org.bouncycastle.jce.spec.ECNamedCurveGenParameterSpec;
+
+import java.security.GeneralSecurityException;
+import java.security.spec.ECGenParameterSpec;
+
+public class ECDHNistP extends AbstractDHG {
+
+    private String curve;
+
+    /** Named factory for ECDHNistP key exchange */
+    public static class Factory521
+            implements net.schmizz.sshj.common.Factory.Named<KeyExchange> {
+
+        @Override
+        public KeyExchange create() {
+            return new ECDHNistP("P-521", new SHA512());
+        }
+
+        @Override
+        public String getName() {
+            return "ecdh-sha2-nistp521";
+        }
+    }
+
+    /** Named factory for ECDHNistP key exchange */
+    public static class Factory384
+            implements net.schmizz.sshj.common.Factory.Named<KeyExchange> {
+
+        @Override
+        public KeyExchange create() {
+            return new ECDHNistP("P-384", new SHA384());
+        }
+
+        @Override
+        public String getName() {
+            return "ecdh-sha2-nistp384";
+        }
+    }
+
+    /** Named factory for ECDHNistP key exchange */
+    public static class Factory256
+            implements net.schmizz.sshj.common.Factory.Named<KeyExchange> {
+
+        @Override
+        public KeyExchange create() {
+            return new ECDHNistP("P-256", new SHA256());
+        }
+
+        @Override
+        public String getName() {
+            return "ecdh-sha2-nistp256";
+        }
+    }
+
+    public ECDHNistP(String curve, Digest digest) {
+        super(new ECDH(), digest);
+        this.curve = curve;
+    }
+
+    @Override
+    protected void initDH(DHBase dh) throws GeneralSecurityException {
+        dh.init(new ECNamedCurveGenParameterSpec(curve));
+    }
+
+}

src/main/java/net/schmizz/sshj/transport/kex/KeyExchangeBase.java

@@ -7,9 +7,6 @@ import net.schmizz.sshj.transport.TransportException;
 import java.security.GeneralSecurityException;
 import java.util.Arrays;
 
-/**
- * Created by ajvanerp on 29/10/15.
- */
 public abstract class KeyExchangeBase implements KeyExchange {
     protected Transport trans;
 

src/test/java/com/hierynomus/sshj/transport/kex/DiffieHellmanGroupExchangeTest.java

@@ -1,45 +0,0 @@
-package com.hierynomus.sshj.transport.kex;
-
-import com.hierynomus.sshj.test.SshFixture;
-import net.schmizz.sshj.SSHClient;
-import org.apache.sshd.common.KeyExchange;
-import org.apache.sshd.common.NamedFactory;
-import org.apache.sshd.server.kex.DHGEX;
-import org.apache.sshd.server.kex.DHGEX256;
-import org.junit.After;
-import org.junit.Rule;
-import org.junit.Test;
-
-import java.io.IOException;
-import java.util.Arrays;
-import java.util.Collections;
-
-import static org.hamcrest.MatcherAssert.assertThat;
-
-public class DiffieHellmanGroupExchangeTest {
-    @Rule
-    public SshFixture fixture = new SshFixture(false);
-
-    @After
-    public void stopServer() {
-        fixture.stopServer();
-    }
-
-    @Test
-    public void shouldKexWithGroupExchangeSha1() throws IOException {
-        setupAndCheckKex(new DHGEX.Factory());
-    }
-
-    @Test
-    public void shouldKexWithGroupExchangeSha256() throws IOException {
-        setupAndCheckKex(new DHGEX256.Factory());
-    }
-
-    private void setupAndCheckKex(NamedFactory<KeyExchange> factory) throws IOException {
-        fixture.getServer().setKeyExchangeFactories(Collections.singletonList(factory));
-        fixture.start();
-        SSHClient sshClient = fixture.setupConnectedDefaultClient();
-        assertThat("should be connected", sshClient.isConnected());
-        sshClient.disconnect();
-    }
-}

src/test/java/com/hierynomus/sshj/transport/kex/KeyExchangeTest.java

@@ -0,0 +1,69 @@
+package com.hierynomus.sshj.transport.kex;
+
+import com.hierynomus.sshj.test.KnownFailingTests;
+import com.hierynomus.sshj.test.SshFixture;
+import net.schmizz.sshj.DefaultConfig;
+import net.schmizz.sshj.SSHClient;
+import net.schmizz.sshj.common.Factory;
+import net.schmizz.sshj.transport.kex.DHGexSHA1;
+import net.schmizz.sshj.transport.kex.DHGexSHA256;
+import net.schmizz.sshj.transport.kex.ECDHNistP;
+import org.apache.sshd.common.KeyExchange;
+import org.apache.sshd.common.NamedFactory;
+import org.apache.sshd.server.kex.*;
+import org.junit.After;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+
+import java.io.IOException;
+import java.util.Collections;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+
+public class KeyExchangeTest {
+    @Rule
+    public SshFixture fixture = new SshFixture(false);
+
+    @After
+    public void stopServer() {
+        fixture.stopServer();
+    }
+
+    @Test
+    public void shouldKexWithDiffieHellmanGroupExchangeSha1() throws IOException {
+        setupAndCheckKex(new DHGEX.Factory(), new DHGexSHA1.Factory());
+    }
+
+    @Test
+    public void shouldKexWithDiffieHellmanGroupExchangeSha256() throws IOException {
+        setupAndCheckKex(new DHGEX256.Factory(), new DHGexSHA256.Factory());
+    }
+
+    @Test
+    public void shouldKexWithEllipticCurveDiffieHellmanNistP256() throws IOException {
+        setupAndCheckKex(new ECDHP256.Factory(), new ECDHNistP.Factory256());
+    }
+
+    @Test
+    public void shouldKexWithEllipticCurveDiffieHellmanNistP384() throws IOException {
+        setupAndCheckKex(new ECDHP384.Factory(), new ECDHNistP.Factory384());
+    }
+
+    @Test
+    @Category({KnownFailingTests.class})
+    public void shouldKexWithEllipticCurveDiffieHellmanNistP521() throws IOException {
+        setupAndCheckKex(new ECDHP521.Factory(), new ECDHNistP.Factory521());
+    }
+
+    private void setupAndCheckKex(NamedFactory<KeyExchange> serverFactory,
+                                  Factory.Named<net.schmizz.sshj.transport.kex.KeyExchange> clientFactory) throws IOException {
+        fixture.getServer().setKeyExchangeFactories(Collections.singletonList(serverFactory));
+        fixture.start();
+        DefaultConfig config = new DefaultConfig();
+        config.setKeyExchangeFactories(Collections.singletonList(clientFactory));
+        SSHClient sshClient = fixture.connectClient(fixture.setupClient(config));
+        assertThat("should be connected", sshClient.isConnected());
+        sshClient.disconnect();
+    }
+}