abstraction, XOR multiple bits

src/tea3/cli.py

@@ -3,7 +3,7 @@ from tea3.cliutils import prompt_int, prompt_choice, prompt_list
 from tea3.tea3model import Tea3Model
 from tea3.variable_search import run_exhaustive
 from tea3.sbox import run_sbox
-from tea3.variable_xor import run_variable_xor, run_exhaustive_two_bit_xor
+from tea3.variable_xor import run_variable_xor, run_exhaustive_xor
 from tea3.f31f32 import run_f31f32
 
 
@@ -70,15 +70,17 @@ def run_variable_xor_cli():
     print("\n" + "=" * 50)
     print("Done.")
 
-def run_exhaustive_two_bit_xor_cli():
+def run_exhaustive_xor_cli():
     print("\nR registers are indexed 0–7; bits within each register are 0–7.")
-    print("This mode checks all 2-bit XOR pairs and prints the best one(s) at each step.")
+    print("This mode searches all XOR combinations of a given size.")
 
     steps = prompt_int("How many steps? (1–100): ", 1, 100)
     target_reg = prompt_int("Target register (0–7): ", 0, 7)
+    xor_size = prompt_int("Number of bits to XOR (2–8): ", 2, 8)
 
     print("-" * 50)
-    run_exhaustive_two_bit_xor(steps, target_reg)
+
+    run_exhaustive_xor(steps=steps, target_reg=target_reg, xor_size=xor_size)
 
     print("\n" + "=" * 50)
     print("Done.")
@@ -96,7 +98,7 @@ def main():
     print("  3) S box analysis")
     print("  4) variable XOR")
     print("  5) F31, F32 analysis")
-    print("  6) Exhaustive 2-bit XOR search")
+    print("  6) Exhaustive XOR search")
 
     mode = prompt_choice("Your choice (1, 2, 3, 4, 5 or 6): ", {1, 2, 3, 4, 5, 6})
 
@@ -111,7 +113,7 @@ def main():
     elif mode == 5:
         run_f31f32()
     elif mode == 6:
-        run_exhaustive_two_bit_xor_cli()
+        run_exhaustive_xor_cli()
 
 
 main()

src/tea3/variable_xor.py

@@ -6,6 +6,12 @@ from tea3.pretty_print import pretty_print
 def monomial_count(poly):
     return len(poly.monomials())
 
+def xor_poly_from_bits(model, target_reg, bits):
+    poly = model.S.zero()
+    for bit in bits:
+        poly += model.R_bits[target_reg][bit]
+    return poly
+
 def run_variable_xor(steps, target_reg, bits_to_xor):
     model = Tea3Model()
 
@@ -20,9 +26,8 @@ def run_variable_xor(steps, target_reg, bits_to_xor):
     for i in range(steps):
         model.step()
 
-        xor_poly = model.S.zero()
-        for bit in bits_to_xor:
-            xor_poly += model.R_bits[target_reg][bit]
+        xor_poly = xor_poly_from_bits(model, target_reg, bits_to_xor)
+
 
         print(f"\n[Step {i + 1}]")
         print(f"XOR of R_bits[{target_reg}][{bits_to_xor}] =")
@@ -30,34 +35,42 @@ def run_variable_xor(steps, target_reg, bits_to_xor):
         print()
 
 
-def run_exhaustive_two_bit_xor(steps, target_reg, bit_indices=range(8)):
-    model = Tea3Model()
+def run_exhaustive_xor(steps, target_reg, xor_size, bit_indices=range(8)):
+    if xor_size < 2:
+        print("xor_size must be at least 2.")
+        return
 
-    bit_pairs = list(combinations(bit_indices, 2))
+    bit_indices = list(bit_indices)
+    if xor_size > len(bit_indices):
+        print(f"xor_size must be at most {len(bit_indices)}.")
+        return
+
+    model = Tea3Model()
+    all_combinations = list(combinations(bit_indices, xor_size))
 
     print(f"Target register: R{target_reg}")
-    print(f"Searching all 2-bit XORs among bits: {list(bit_indices)}")
+    print(f"Searching all {xor_size}-bit XORs among bits: {bit_indices}")
     print("-" * 50)
 
     for i in range(steps):
         model.step()
 
         best_count = None
-        best_pairs = []
+        best_combinations = []
         best_poly = None
 
-        for b1, b2 in bit_pairs:
-            xor_poly = model.R_bits[target_reg][b1] + model.R_bits[target_reg][b2]
+        for bits in all_combinations:
+            xor_poly = xor_poly_from_bits(model, target_reg, bits)
             count = monomial_count(xor_poly)
 
             if best_count is None or count < best_count:
                 best_count = count
-                best_pairs = [(b1, b2)]
+                best_combinations = [bits]
                 best_poly = xor_poly
             elif count == best_count:
-                best_pairs.append((b1, b2))
+                best_combinations.append(bits)
 
         print(f"\n[Step {i + 1}]")
         print(f"Best XOR(s) with {best_count} total monomials:")
-        for b1, b2 in best_pairs:
-            print(f"  bits XORed: ({b1}, {b2})")
+        for bits in best_combinations:
+            print(f"  bits XORed: {bits}")