problem 24, part 1

3 files changed, 33 insertions, 0 deletions

diff --git a/src/crack.rs b/src/crack.rs
index 92950c4..b351249 100644
--- a/src/crack.rs
+++ b/src/crack.rs
@@ -386,6 +386,21 @@ pub fn clone_mersenne_twister_from_output (outputs: &[u32]) -> MersenneTwister {
     return MersenneTwister::from_seed((state, 0));
 }
 
+pub fn recover_16_bit_mt19937_key (ciphertext: &[u8], suffix: &[u8]) -> Option<u16> {
+    for _key in 0..65536u32 {
+        let key = _key as u16;
+        let plaintext = ::random::mt19937_stream_cipher(
+            ciphertext,
+            key as u32
+        );
+        if &plaintext[(ciphertext.len() - suffix.len())..] == suffix {
+            return Some(key);
+        }
+    }
+
+    return None;
+}
+
 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/lib.rs b/src/lib.rs
index 089250c..bbaf93f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -37,3 +37,4 @@ pub use crack::crack_single_byte_xor;
 pub use crack::crack_repeating_key_xor;
 pub use crack::crack_fixed_nonce_ctr_statistically;
 pub use crack::clone_mersenne_twister_from_output;
+pub use crack::recover_16_bit_mt19937_key;
diff --git a/tests/lib.rs b/tests/lib.rs
index 71c5e32..f498e1f 100644
--- a/tests/lib.rs
+++ b/tests/lib.rs
@@ -483,3 +483,20 @@ fn problem_22 () {
         assert_eq!(mt.gen::<u32>(), mt2.gen::<u32>());
     }
 }
+
+#[test]
+fn problem_23 () {
+    let key: u16 = rand::thread_rng().gen();
+    let fixed_suffix = b"AAAAAAAAAAAAAA";
+    let plaintext: Vec<u8> = rand::thread_rng()
+        .gen_iter()
+        .take(rand::thread_rng().gen_range(0, 32))
+        .chain(fixed_suffix.iter().map(|x| *x))
+        .collect();
+    let ciphertext = matasano::mt19937_stream_cipher(&plaintext[..], key as u32);
+    let got = matasano::recover_16_bit_mt19937_key(
+        &ciphertext[..],
+        &fixed_suffix[..],
+    ).unwrap();
+    assert_eq!(got, key);
+}