CLI changes

.gitignore

@@ -3,4 +3,4 @@
 __pycache__/
 .pytest_cache/
 *.egg-info/
-
+*.txt

src/tea3/cli.py

@@ -1,5 +1,6 @@
 from tea3.pretty_print import pretty_print
 from tea3.tea3model import Tea3Model
+from tea3.variable_search import run_exhaustive
 
 
 def prompt_int(message: str, lo: int, hi: int) -> int:
@@ -15,30 +16,80 @@ def prompt_int(message: str, lo: int, hi: int) -> int:
         print(f"Value must be between {lo} and {hi} (inclusive).")
 
 
-def main() -> None:
+def prompt_choice(message: str, choices: set[int]) -> int:
-    print("=" * 50)
+    choices_str = ", ".join(map(str, sorted(choices)))
-    print(" Tea3 Model ")
+    while True:
-    print("=" * 50)
+        raw = input(message).strip()
+        try:
+            value = int(raw)
+        except ValueError:
+            print("Please enter a number.")
+            continue
+        if value in choices:
+            return value
+        print(f"Value must be one of: {choices_str}")
 
-    steps = prompt_int("\nHow many steps do you want to run? (1–100): ", 1, 100)
 
+def run_classic_cli():
     print("\nR registers are indexed 0–7; bits within each register are 0–7.")
-    reg = prompt_int("Which R register do you want to inspect? (0–7): ", 0, 7)
+    print("Enter -1 to print all registers, or all bits.")
-    bit = prompt_int("Which bit of that register?              (0–7): ", 0, 7)
+
+    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)
 
-    print(f"\nRunning {steps} step(s), watching R[{reg}][{bit}]")
     print("-" * 50)
 
-    model = Tea3Model(max_steps=steps)
+    model = Tea3Model()
 
     for i in range(steps):
         model.step()
-        poly = model.R_bits[reg][bit]
+        print(f"\n[Step {i + 1}]")
-        print(f"\n[Step {i+1}]  R_bits[{reg}][{bit}] =")
+
+        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("\nR registers are indexed 0–7; bits within each register are 0–7.")
+    print("Enter -1 to print all bits in the chosen register.")
+
+    steps = prompt_int("How many steps? (1–100): ", 1, 100)
+    target_reg = prompt_int("Target register (0–7): ", 0, 7)
+    target_bit = prompt_int("Target bit (-1 or 0–7): ", -1, 7)
+
+    print("-" * 50)
+    run_exhaustive(steps, target_reg, target_bit)
 
     print("\n" + "=" * 50)
     print("Done.")
 
 
+
+def main():
+    print("=" * 50)
+    print(" Tea3 Model ")
+    print("=" * 50)
+
+    print("\nChoose a mode:")
+    print("  1) Classic inspection")
+    print("  2) Exhaustive variable-change search")
+
+    mode = prompt_choice("Your choice (1 or 2): ", {1, 2})
+
+    if mode == 1:
+        run_classic_cli()
+    else:
+        run_exhaustive_cli()
+
+
 main()

src/tea3/pretty_print.py

@@ -1,11 +1,12 @@
 def pretty_print(poly):
-    ring = poly.parent()
-
     R_terms = []
     r_terms = []
     x_terms = []
     mixed_terms = []
 
+    if isinstance(poly, int):
+        return str(poly)
+
     has_const = bool(poly.constant_coefficient())
 
     for monom in poly:
@@ -28,11 +29,12 @@ def pretty_print(poly):
 
     parts = []
 
-    if R_terms or (has_const and not r_terms):
+    if has_const:
-        parts.append("f(Ri)")
-    elif has_const:
         parts.append("1")
 
+    if R_terms:
+        parts.append("f(Ri)")
+
     if r_terms:
         parts.append("f(ri)")
 

src/tea3/tea3model.py

@@ -1,19 +1,22 @@
 from sage.all import GF, BooleanPolynomialRing
+from functools import reduce
+from operator import mul
 
 from tea3.constants import TEA3_SBOX, T_F1, T_F2
 from tea3.pretty_print import pretty_print, pretty_print_vec
 
 
 class Tea3Model:
-    def __init__(self, max_steps=20):
+    def __init__(self):
         self.F = GF(2)
         self.step_count = 0
 
         names = (
-            [f"x{i}{j}" for i in range(5) for j in range(8)] +
+            [f"x{i}{j}" for i in range(5) for j in range(8)] +  # 0–39
-            [f"r{i}{j}" for i in range(5) for j in range(8)] +
+            [f"y{i}{j}" for i in range(5) for j in range(8)] +  # 40–79
-            [f"R{i}{j}" for i in range(8) for j in range(8)] +
+            [f"r{i}{j}" for i in range(5) for j in range(8)] +  # 80–119
-            [f"f{s}_{i}{j}" for s in range(max_steps) for i in range(8) for j in range(8)]
+            [f"R{i}{j}" for i in range(8) for j in range(8)] +  # 120–183
+            ["g"]                                               # 184
         )
 
         name_string = ",".join(names)
@@ -21,63 +24,44 @@ class Tea3Model:
         self.v = self.S.gens()
 
         self.x_bits = [list(self.v[i*8       : (i+1)*8      ]) for i in range(5)]
-        self.r_bits = [list(self.v[40 + i*8 : 40 + (i+1)*8]) for i in range(5)]
+        self.y_bits = [list(self.v[40 + i*8  : 40 + (i+1)*8]) for i in range(5)]
-        self.R_bits = [list(self.v[80 + i*8 : 80 + (i+1)*8]) for i in range(8)]
+        self.r_bits = [list(self.v[80 + i*8  : 80 + (i+1)*8]) for i in range(5)]
-
+        self.R_bits = [list(self.v[120 + i*8 : 120 + (i+1)*8]) for i in range(8)]
-        # Abstract variables
+        self.g = self.v[-1]
-        base = 80 + 64
-        self.fR_bits = [
-            [list(self.v[base + s*64 + i*8 : base + s*64 + i*8 + 8])
-             for i in range(8)]
-            for s in range(max_steps)
-        ]
-
-    def _split_poly(self, poly):
-        """
-        Split a polynomial into:
-          - R_f_part: monomials involving only 'R' or 'f' (abstract) variables
-          - xr_part:  monomials involving 'x' or 'r' variables
-
-        constant term is grouped with R_f_part when R_f_part is non-zero
-        """
-        zero = self.S.zero()
-        R_f_part = zero
-        xr_part = zero
-
-        has_const = bool(poly.constant_coefficient())
-
-        for monom in poly:
-            vars_in_term = monom.variables()
-            if not vars_in_term:
-                continue
-            families = {str(v)[0] for v in vars_in_term}
-            monom_poly = self.S(monom)
-            if families <= {'R', 'f'}:
-                R_f_part += monom_poly
-            else:
-                xr_part += monom_poly
-
-        if has_const:
-            if R_f_part != zero:
-                R_f_part += self.S.one()
-            else:
-                xr_part += self.S.one()
-
-        return R_f_part, xr_part
 
     def _abstract_R(self):
-        """
+        one  = self.S.one()
-        Replace the R/f-dependent part of every R_bits[i][j] with an
+        zero = self.S.zero()
-        abstract variable f{step}_{i}{j}, leaving only x and r terms explicit.
+
-        """
-        s = self.step_count
         for i in range(8):
             for j in range(8):
-                R_f_part, xr_part = self._split_poly(self.R_bits[i][j])
+                poly = self.R_bits[i][j]
-                if R_f_part != self.S.zero():
+
-                    self.R_bits[i][j] = self.fR_bits[s][i][j] + xr_part
+                groups  = {}
+                pure_xyr = zero
+                const    = one if bool(poly.constant_coefficient()) else zero
+
+                for monom in poly:
+                    term_vars = monom.variables()
+                    if not term_vars:
+                        continue
+
+                    xyr_vars = [v for v in term_vars if str(v)[0] in ('x', 'y', 'r')]
+                    Rg_vars  = [v for v in term_vars if str(v)[0] in ('R', 'g')]
+
+                    xyr_mono = reduce(mul, (self.S(v) for v in xyr_vars), one)
+                    xyr_key  = frozenset(str(v) for v in xyr_vars)
+
+                    if not Rg_vars:
+                        pure_xyr += xyr_mono
                     else:
-                    self.R_bits[i][j] = xr_part
+                        groups[xyr_key] = xyr_mono
+
+                result = pure_xyr + const
+                for xyr_key, xyr_mono in groups.items():
+                    result += xyr_mono * self.g
+
+                self.R_bits[i][j] = result
 
     def step(self):
         R = self.R_bits.copy()
@@ -117,6 +101,7 @@ class Tea3Model:
 
         return R7
 
+
 def S(r):
     # placeholder
     return r

src/tea3/variable_search.py

@@ -0,0 +1,50 @@
+from itertools import product as iproduct
+from tea3.tea3model import Tea3Model
+from tea3.pretty_print import pretty_print
+
+def apply_variable_change(model, coeffs):
+    orig_x = [list(model.v[i*8 : (i+1)*8]) for i in range(5)]
+    orig_y = model.y_bits
+
+    subs = {}
+
+    for i in range(1, 5):
+        for j in range(8):
+            subs[orig_x[i][j]] = orig_y[i][j]
+
+    for j in range(8):
+        new_expr = orig_y[0][j]
+        for i, ai in enumerate(coeffs, start=1):
+            if ai:
+                new_expr = new_expr + orig_y[i][j]
+        subs[orig_x[0][j]] = new_expr
+
+    new_R = []
+    for i in range(8):
+        row = [poly.subs(subs) for poly in model.R_bits[i]]
+        new_R.append(row)
+    return new_R
+
+def run_exhaustive(steps: int, target_reg: int = 0, target_bit: int = -1):
+    model = Tea3Model()
+    for _ in range(steps):
+        model.step()
+
+    snapshot = model
+
+    for idx, coeffs in enumerate(iproduct([0, 1], repeat=4)):
+        new_R = apply_variable_change(snapshot, coeffs)
+
+        label = "".join(map(str, coeffs))
+        print(f"\n[{idx:02d}] (a1,a2,a3,a4) = {label}")
+
+        if target_bit == -1:
+            bits = range(8)
+        else:
+            if not (0 <= target_bit < 8):
+                raise ValueError("target_bit must be in [0, 7] or -1")
+            bits = [target_bit]
+
+        for j in bits:
+            poly = new_R[target_reg][j]
+            print(f"  R[{target_reg}][{j}] = {pretty_print(poly)}")