From f0a8f1f5a183686f143b3eff9010029e67b5bfcf Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Thu, 25 Jun 2026 09:47:46 +0200
Subject: [PATCH] advanced inspection mode

---
 src/tea3/cli.py | 101 +++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 91 insertions(+), 10 deletions(-)

diff --git a/src/tea3/cli.py b/src/tea3/cli.py
index d5eb9f2..a977625 100644
--- a/src/tea3/cli.py
+++ b/src/tea3/cli.py
@@ -36,6 +36,84 @@ def run_classic_cli():
     print("\n" + "=" * 50)
     print("Done.")
 
+def _set_bits_to_zero(model, bit_names):
+    zero = model.S.zero()
+
+    for raw in bit_names:
+        name = raw.strip()
+        if not name:
+            continue
+
+        if len(name) != 3 or not name[1:].isdigit():
+            raise ValueError(
+                f"Invalid bit name '{name}'. Expected formats like R00, x00, r10."
+            )
+
+        prefix = name[0]
+        reg = int(name[1])
+        bit = int(name[2])
+
+        if prefix == "R":
+            if not (0 <= reg <= 7 and 0 <= bit <= 7):
+                raise ValueError(f"Invalid R bit '{name}'. Use R00..R77.")
+            model.R_bits[reg][bit] = zero
+
+        elif prefix == "x":
+            if not (0 <= reg <= 4 and 0 <= bit <= 7):
+                raise ValueError(f"Invalid x bit '{name}'. Use x00..x47.")
+            model.x_bits[reg][bit] = zero
+
+        elif prefix == "r":
+            if not (0 <= reg <= 4 and 0 <= bit <= 7):
+                raise ValueError(f"Invalid r bit '{name}'. Use r00..r47.")
+            model.r_bits[reg][bit] = zero
+
+        else:
+            raise ValueError(
+                f"Invalid bit name '{name}'. Expected R, x, or r."
+            )
+
+
+def run_advanced_cli():
+    print("\nR registers are indexed 0–7; bits within each register are 0–7.")
+    print("x, y, r registers are indexed 0–4; bits within each register are 0–7.")
+    print("Enter the bits you want to force to 0, for example:")
+    print("  R00 x00 r00 r01 r10")
+
+    steps = prompt_int("How many steps do you want to run? (1–100): ", 1, 100)
+    reg = prompt_int("Which R register do you want to inspect? (-1 or 0–7): ", -1, 7)
+    bit = prompt_int("Which bit of that register?              (-1 or 0–7): ", -1, 7)
+
+    raw = input("Bits to set to 0 (space-separated): ").strip()
+    zero_bits = raw.split() if raw else []
+
+    print("-" * 50)
+
+    model = Tea3Model()
+    if zero_bits:
+        try:
+            _set_bits_to_zero(model, zero_bits)
+        except ValueError as e:
+            print(f"Error: {e}")
+            return
+
+    for i in range(steps):
+        model.step()
+        print(f"\n[Step {i + 1}]")
+
+        regs = range(8) if reg == -1 else [reg]
+        bits = range(8) if bit == -1 else [bit]
+
+        for r in regs:
+            for b in bits:
+                poly = model.R_bits[r][b]
+                print(f"R_bits[{r}][{b}] =")
+                print(pretty_print(poly))
+                print()
+
+    print("\n" + "=" * 50)
+    print("Done.")
+
 def run_exhaustive_cli():
     print("\nChoose the variable-change family:")
     print("  1) First-row matrix")
@@ -110,25 +188,28 @@ def main():
 
     print("\nChoose a mode:")
     print("  1) Classic inspection")
-    print("  2) Exhaustive variable-change search")
-    print("  3) S box analysis")
-    print("  4) variable XOR")
-    print("  5) F31, F32 analysis")
-    print("  6) Exhaustive XOR search")
+    print("  2) Advanced inspection with forced zero bits")
+    print("  3) Exhaustive variable-change search")
+    print("  4) S box analysis")
+    print("  5) variable XOR")
+    print("  6) F31, F32 analysis")
+    print("  7) Exhaustive XOR search")
 
-    mode = prompt_choice("Your choice (1, 2, 3, 4, 5 or 6): ", {1, 2, 3, 4, 5, 6})
+    mode = prompt_choice("Your choice (1, 2, 3, 4, 5, 6 or 7): ", {1, 2, 3, 4, 5, 6, 7})
 
     if mode == 1:
         run_classic_cli()
     elif mode == 2:
-        run_exhaustive_cli()
+        run_advanced_cli()
     elif mode == 3:
-        run_sbox()
+        run_exhaustive_cli()
     elif mode == 4:
-        run_variable_xor_cli()
+        run_sbox()
     elif mode == 5:
-        run_f31f32()
+        run_variable_xor_cli()
     elif mode == 6:
+        run_f31f32()
+    elif mode == 7:
         run_exhaustive_xor_cli()