refactor to allow for multiple client implementations

4 files changed, 94 insertions, 61 deletions

diff --git a/src/crack.rs b/src/crack.rs
index e65ea35..731cfbd 100644
--- a/src/crack.rs
+++ b/src/crack.rs
@@ -1,3 +1,4 @@
+use num_bigint::RandBigInt;
 use rand::Rng;
 use rayon::prelude::*;
 use std::borrow::ToOwned;
@@ -932,6 +933,89 @@ impl DiffieHellmanMessageExchanger
     }
 }
 
+pub trait SRPClient {
+    fn server(&mut self) -> &mut crate::dh::SRPServer;
+
+    fn key_exchange_impl(
+        &mut self,
+        user: &str,
+        pass: &str,
+    ) -> (Vec<u8>, Vec<u8>, num_bigint::BigUint);
+
+    fn register(&mut self, user: &str, pass: &str) {
+        let n = &self.server().n.clone();
+        let g = &self.server().g.clone();
+
+        let mut salt = [0; 16];
+        rand::thread_rng().fill(&mut salt);
+        let input = [&salt[..], pass.as_bytes()].concat();
+        let xh = crate::sha1::sha1(&input);
+        let x = num_bigint::BigUint::from_bytes_le(&xh[..]);
+        let v = g.modpow(&x, n);
+        self.server().register(user, &salt, &v);
+    }
+
+    fn key_exchange(
+        &mut self,
+        user: &str,
+        pass: &str,
+    ) -> Option<num_bigint::BigUint> {
+        let (session, salt, s) = self.key_exchange_impl(user, pass);
+        let k = crate::sha1::sha1(&s.to_bytes_le());
+        let hmac = crate::sha1::sha1_hmac(&k, &salt);
+
+        if !self.server().verify(session, hmac.to_vec()) {
+            return None;
+        }
+
+        Some(s)
+    }
+}
+
+#[derive(Debug)]
+pub struct CorrectSRPClient<'a> {
+    server: &'a mut crate::dh::SRPServer,
+}
+
+impl<'a> CorrectSRPClient<'a> {
+    pub fn new(server: &'a mut crate::dh::SRPServer) -> CorrectSRPClient<'a> {
+        CorrectSRPClient { server }
+    }
+}
+
+impl<'a> SRPClient for CorrectSRPClient<'a> {
+    fn server(&mut self) -> &mut crate::dh::SRPServer {
+        self.server
+    }
+
+    fn key_exchange_impl(
+        &mut self,
+        user: &str,
+        pass: &str,
+    ) -> (Vec<u8>, Vec<u8>, num_bigint::BigUint) {
+        let n = &self.server.n.clone();
+        let g = &self.server.g.clone();
+        let k = &self.server.k.clone();
+
+        let a_priv = rand::thread_rng().gen_biguint_below(n);
+        let a_pub = g.modpow(&a_priv, n);
+        let (session, salt, b_pub) =
+            self.server.exchange_pubkeys(user, &a_pub);
+
+        let uinput = [a_pub.to_bytes_le(), b_pub.to_bytes_le()].concat();
+        let uh = crate::sha1::sha1(&uinput);
+        let u = num_bigint::BigUint::from_bytes_le(&uh[..]);
+
+        let xinput = [salt.clone(), pass.as_bytes().to_vec()].concat();
+        let xh = crate::sha1::sha1(&xinput);
+        let x = num_bigint::BigUint::from_bytes_le(&xh[..]);
+
+        let s = (b_pub - k * g.modpow(&x, n)).modpow(&(a_priv + u * x), n);
+
+        (session, salt, s)
+    }
+}
+
 fn crack_single_byte_xor_with_confidence(input: &[u8]) -> (u8, f64) {
     let mut min_diff = 100.0;
     let mut best_key = 0;
diff --git a/src/dh.rs b/src/dh.rs
index 5317b29..74f5f54 100644
--- a/src/dh.rs
+++ b/src/dh.rs
@@ -31,11 +31,12 @@ impl DHKeyPair {
 
 #[derive(Debug)]
 pub struct SRPServer {
+    pub n: num_bigint::BigUint,
+    pub g: num_bigint::BigUint,
+    pub k: num_bigint::BigUint,
+
     users: std::collections::HashMap<String, SRPUser>,
     sessions: std::collections::HashMap<Vec<u8>, SRPSession>,
-    n: num_bigint::BigUint,
-    g: num_bigint::BigUint,
-    k: num_bigint::BigUint,
 }
 
 impl SRPServer {
@@ -45,11 +46,11 @@ impl SRPServer {
         k: num_bigint::BigUint,
     ) -> SRPServer {
         SRPServer {
-            users: std::collections::HashMap::new(),
-            sessions: std::collections::HashMap::new(),
             n,
             g,
             k,
+            users: std::collections::HashMap::new(),
+            sessions: std::collections::HashMap::new(),
         }
     }
 
@@ -131,57 +132,3 @@ pub struct SRPSession {
     v: num_bigint::BigUint,
     salt: Vec<u8>,
 }
-
-#[derive(Debug)]
-pub struct SRPClient<'a> {
-    server: &'a mut SRPServer,
-}
-
-impl<'a> SRPClient<'a> {
-    pub fn new(server: &'a mut SRPServer) -> SRPClient<'a> {
-        SRPClient { server }
-    }
-
-    pub fn register(&mut self, user: &str, pass: &str) {
-        let mut salt = [0; 16];
-        rand::thread_rng().fill(&mut salt);
-        let input = [&salt[..], pass.as_bytes()].concat();
-        let xh = crate::sha1::sha1(&input);
-        let x = num_bigint::BigUint::from_bytes_le(&xh[..]);
-        let v = self.server.g.modpow(&x, &self.server.n);
-        self.server.register(user, &salt, &v);
-    }
-
-    pub fn key_exchange(
-        &mut self,
-        user: &str,
-        pass: &str,
-    ) -> Option<num_bigint::BigUint> {
-        let n = &self.server.n.clone();
-        let g = &self.server.g.clone();
-        let k = &self.server.k.clone();
-
-        let a_priv = rand::thread_rng().gen_biguint_below(n);
-        let a_pub = g.modpow(&a_priv, n);
-        let (session, salt, b_pub) =
-            self.server.exchange_pubkeys(user, &a_pub);
-
-        let uinput = [a_pub.to_bytes_le(), b_pub.to_bytes_le()].concat();
-        let uh = crate::sha1::sha1(&uinput);
-        let u = num_bigint::BigUint::from_bytes_le(&uh[..]);
-
-        let xinput = [salt.clone(), pass.as_bytes().to_vec()].concat();
-        let xh = crate::sha1::sha1(&xinput);
-        let x = num_bigint::BigUint::from_bytes_le(&xh[..]);
-
-        let s = (b_pub - k * g.modpow(&x, n)).modpow(&(a_priv + u * x), n);
-        let k = crate::sha1::sha1(&s.to_bytes_le());
-        let hmac = crate::sha1::sha1_hmac(&k, &salt);
-
-        if !self.server.verify(session, hmac.to_vec()) {
-            return None
-        }
-
-        Some(s)
-    }
-}
diff --git a/src/lib.rs b/src/lib.rs
index d7566e8..97ce460 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -37,11 +37,12 @@ pub use crack::recover_16_bit_mt19937_key;
 pub use crack::recover_mersenne_twister_seed_from_time;
 pub use crack::recover_mt19937_key_from_time;
 pub use crack::BlockCipherMode;
+pub use crack::CorrectSRPClient;
 pub use crack::DiffieHellmanMessageExchanger;
 pub use crack::NullDiffieHellmanMessageExchanger;
 pub use crack::ParameterInjectionDiffieHellmanMessageExchanger;
+pub use crack::SRPClient;
 pub use dh::DHKeyPair;
-pub use dh::SRPClient;
 pub use dh::SRPServer;
 pub use dh::SRPUser;
 pub use http::create_query_string;
diff --git a/tests/set5.rs b/tests/set5.rs
index 6aa089a..66d2420 100644
--- a/tests/set5.rs
+++ b/tests/set5.rs
@@ -1,3 +1,4 @@
+use matasano::SRPClient;
 use rand::Rng;
 
 #[test]
@@ -251,7 +252,7 @@ fn problem_36() {
     let pass = "supersecret";
 
     let mut server = matasano::SRPServer::new(n, g, k);
-    let mut client = matasano::SRPClient::new(&mut server);
+    let mut client = matasano::CorrectSRPClient::new(&mut server);
 
     client.register(user, pass);