4 files changed, 93 insertions, 12 deletions

diff --git a/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/KeyAgreementTestable.java b/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/KeyAgreementTestable.java
index 6566a9c..f3786de 100644
--- a/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/KeyAgreementTestable.java
+++ b/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/KeyAgreementTestable.java
@@ -140,6 +140,27 @@ public class KeyAgreementTestable extends StandaloneTestable<KeyAgreementTestabl
         }
     }
 
+    @Override
+    public KeyAgreementTestable clone() {
+        try {
+            KeyAgreementTestable kat = builder()
+                    .ka(KeyAgreement.getInstance(ka.getAlgorithm(), ka.getProvider()))
+                    .privateKey(privateKey)
+                    .publicKey(publicKey)
+                    .privateKgt(kgtPrivate) // Do not clone the kgts?
+                    .publicKgt(kgtPublic) // Do not clone the kgts?
+                    .spec(spec)
+                    .keyAlgo(keyAlgo)
+                    .random(random)
+                    .build();
+            kat.secret = secret;
+            kat.derived = derived;
+            return kat;
+        } catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     public static Builder builder() {
         return new Builder();
     }
diff --git a/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/KeyGeneratorTestable.java b/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/KeyGeneratorTestable.java
index f9c84e1..cca9dc5 100644
--- a/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/KeyGeneratorTestable.java
+++ b/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/KeyGeneratorTestable.java
@@ -1,9 +1,7 @@
 package cz.crcs.ectester.standalone.test.base;
 
-import java.security.InvalidAlgorithmParameterException;
-import java.security.KeyPair;
-import java.security.KeyPairGenerator;
-import java.security.SecureRandom;
+import javax.crypto.KeyAgreement;
+import java.security.*;
 import java.security.spec.AlgorithmParameterSpec;
 import java.security.spec.ECGenParameterSpec;
 import java.security.spec.ECParameterSpec;
@@ -12,12 +10,13 @@ import java.security.spec.ECParameterSpec;
  * @author Jan Jancar johny@neuromancer.sk
  */
 public class KeyGeneratorTestable extends StandaloneTestable<KeyGeneratorTestable.KeyGeneratorStage> {
-    private KeyPair kp;
-    private final KeyPairGenerator kpg;
+    private KeyPairGenerator kpg;
     private int keysize = 0;
     private AlgorithmParameterSpec spec = null;
     private SecureRandom random;
 
+    private KeyPair kp;
+
     public KeyGeneratorTestable(KeyPairGenerator kpg) {
         this.kpg = kpg;
     }
@@ -80,6 +79,32 @@ public class KeyGeneratorTestable extends StandaloneTestable<KeyGeneratorTestabl
         hasRun = true;
     }
 
+    @Override
+    public void reset() {
+        super.reset();
+        try {
+            kpg = KeyPairGenerator.getInstance(kpg.getAlgorithm(), kpg.getProvider());
+        } catch (NoSuchAlgorithmException e) {
+        }
+        kp = null;
+    }
+
+    @Override
+    public KeyGeneratorTestable clone() {
+        try {
+            KeyGeneratorTestable kgt = builder()
+                    .keyPairGenerator(KeyPairGenerator.getInstance(kpg.getAlgorithm(), kpg.getProvider()))
+                    .keysize(keysize)
+                    .spec(spec)
+                    .random(random)
+                    .build();
+            kgt.kp = kp;
+            return kgt;
+        } catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     public static Builder builder() {
         return new Builder();
     }
@@ -107,6 +132,11 @@ public class KeyGeneratorTestable extends StandaloneTestable<KeyGeneratorTestabl
             return this;
         }
 
+        public Builder spec(AlgorithmParameterSpec spec) {
+            this.spec = spec;
+            return this;
+        }
+
         public Builder spec(ECGenParameterSpec spec) {
             this.spec = spec;
             return this;
diff --git a/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/SignatureTestable.java b/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/SignatureTestable.java
index 5839497..c12f3c7 100644
--- a/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/SignatureTestable.java
+++ b/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/SignatureTestable.java
@@ -1,9 +1,6 @@
 package cz.crcs.ectester.standalone.test.base;
 
-import java.security.InvalidKeyException;
-import java.security.SecureRandom;
-import java.security.Signature;
-import java.security.SignatureException;
+import java.security.*;
 import java.security.interfaces.ECPrivateKey;
 import java.security.interfaces.ECPublicKey;
 
@@ -11,12 +8,13 @@ import java.security.interfaces.ECPublicKey;
  * @author Jan Jancar johny@neuromancer.sk
  */
 public class SignatureTestable extends StandaloneTestable<SignatureTestable.SignatureStage> {
-    private final Signature sig;
+    private Signature sig;
     private ECPrivateKey signKey;
     private ECPublicKey verifyKey;
     private KeyGeneratorTestable kgt;
     private SecureRandom random;
     private byte[] data;
+
     private byte[] signature;
     private boolean verified;
 
@@ -59,6 +57,31 @@ public class SignatureTestable extends StandaloneTestable<SignatureTestable.Sign
         return verified;
     }
 
+    
+    @Override
+    public void reset() {
+        super.reset();
+        try {
+            sig = Signature.getInstance(sig.getAlgorithm(), sig.getProvider());
+        } catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException(e);
+        }
+        signature = null;
+        verified = false;
+    }
+
+    @Override
+    public SignatureTestable clone() throws CloneNotSupportedException {
+        SignatureTestable sig = (SignatureTestable) super.clone();
+        try {
+            sig.sig = Signature.getInstance(this.sig.getAlgorithm(), this.sig.getProvider());
+        } catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException(e);
+        }
+        sig.kgt = kgt; // Do not clone the kgts?
+        return sig;
+    }
+
     @Override
     public void run() {
         try {
@@ -68,7 +91,7 @@ public class SignatureTestable extends StandaloneTestable<SignatureTestable.Sign
                 verifyKey = (ECPublicKey) kgt.getKeyPair().getPublic();
             }
 
-            if(signKey != null) {
+            if (signKey != null) {
                 stage = SignatureStage.InitSign;
                 try {
                     if (random != null) {
diff --git a/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/StandaloneTestable.java b/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/StandaloneTestable.java
index 47bffc1..3ae6fa7 100644
--- a/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/StandaloneTestable.java
+++ b/standalone/src/main/java/cz/crcs/ectester/standalone/test/base/StandaloneTestable.java
@@ -22,4 +22,11 @@ public abstract class StandaloneTestable<T extends Enum<T>> extends BaseTestable
         hasRun = true;
         exception = ex;
     }
+
+    @Override
+    public void reset() {
+        super.reset();
+        stage = null;
+        exception = null;
+    }
 }