add another variable change

src/tea3/cli.py

@@ -1,7 +1,7 @@
 from tea3.pretty_print import pretty_print
 from tea3.cliutils import prompt_int, prompt_choice, prompt_list
 from tea3.tea3model import Tea3Model
-from tea3.variable_search import run_exhaustive, run_exhaustive_staircase
+from tea3.variable_search import run_exhaustive, run_exhaustive_staircase, run_exhaustive_staircase2
 from tea3.sbox import run_sbox
 from tea3.variable_xor import run_variable_xor, run_exhaustive_xor
 from tea3.f31f32 import run_f31f32
@@ -40,8 +40,9 @@ def run_exhaustive_cli():
     print("\nChoose the variable-change family:")
     print("  1) First-row matrix")
     print("  2) Staircase matrix")
+    print("  3) Staircase 2 matrix")
 
-    family = prompt_choice("Your choice (1 or 2): ", {1, 2})
+    family = prompt_choice("Your choice (1, 2 or 3): ", {1, 2, 3})
 
     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.")
@@ -54,8 +55,10 @@ def run_exhaustive_cli():
 
     if family == 1:
         run_exhaustive(steps, target_reg, target_bit)
-    else:
+    elif family == 2:
         run_exhaustive_staircase(steps, target_reg, target_bit)
+    else:
+        run_exhaustive_staircase2(steps, target_reg, target_bit)
 
     print("\n" + "=" * 50)
     print("Done.")

src/tea3/variable_search.py

@@ -74,6 +74,30 @@ def make_staircase_matrix(coeffs):
     M[7, 0] = GF(2)(coeffs[7])
     return M
 
+def make_staircase2_matrix(coeffs):
+    """
+    y = M x with:
+
+    | 1   0   0   0   0   0   0   a1 |
+    | a2  1   0   0   0   0   0   0  |
+    | 0   a3  1   0   0   0   0   0  |
+    | 0   0   a4  1   0   0   0   0  |
+    | 0   0   0   a5  1   0   0   0  |
+    | 0   0   0   0   a6  1   0   0  |
+    | 0   0   0   0   0   a7  1   0  |
+    | 0   0   0   0   0   0   a8  1  |
+    """
+    M = identity_matrix(GF(2), 8)
+    M[0, 7] = GF(2)(coeffs[0])
+    M[1, 0] = GF(2)(coeffs[1])
+    M[2, 1] = GF(2)(coeffs[2])
+    M[3, 2] = GF(2)(coeffs[3])
+    M[4, 3] = GF(2)(coeffs[4])
+    M[5, 4] = GF(2)(coeffs[5])
+    M[6, 5] = GF(2)(coeffs[6])
+    M[7, 6] = GF(2)(coeffs[7])
+    return M
+
 # wrappers
 
 def run_exhaustive_generic(
@@ -142,3 +166,9 @@ def run_exhaustive_staircase(steps: int, target_reg: int = 0, target_bit: int =
         steps, make_staircase_matrix, n_params=8,
         target_reg=target_reg, target_bit=target_bit,
     )
+
+def run_exhaustive_staircase2(steps: int, target_reg: int = 0, target_bit: int = -1):
+    return run_exhaustive_generic(
+        steps, make_staircase2_matrix, n_params=8,
+        target_reg=target_reg, target_bit=target_bit,
+    )