forcing bits to 1

src/tea3/cli.py

@@ -37,8 +37,11 @@ def run_classic_cli():
     print("\n" + "=" * 50)
     print("Done.")
 
-def _set_bits_to_zero(model, bit_names):
-    zero = model.S.zero()
+def _set_bits(model, bit_names, value):
+    if value not in (0, 1):
+        raise ValueError("value must be 0 or 1")
+
+    const = model.S.one() if value else model.S.zero()
 
     for raw in bit_names:
         name = raw.strip()
@@ -57,17 +60,17 @@ def _set_bits_to_zero(model, bit_names):
         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
+            model.R_bits[reg][bit] = const
 
         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
+            model.x_bits[reg][bit] = const
 
         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
+            model.r_bits[reg][bit] = const
 
         else:
             raise ValueError(
@@ -77,9 +80,7 @@ def _set_bits_to_zero(model, bit_names):
 
 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")
+    print("x and r registers are indexed 0–4; bits within each register are 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)
@@ -88,15 +89,20 @@ def run_advanced_cli():
     raw = input("Bits to set to 0 (space-separated): ").strip()
     zero_bits = raw.split() if raw else []
 
+    raw = input("Bits to set to 1 (space-separated): ").strip()
+    one_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
+
+    try:
+        # 1 overrides 0
+        _set_bits(model, zero_bits, 0)
+        _set_bits(model, one_bits, 1)
+    except ValueError as e:
+        print(f"Error: {e}")
+        return
 
     for i in range(steps):
         model.step()
@@ -190,7 +196,7 @@ def main():
 
     print("\nChoose a mode:")
     print("  1) Classic inspection")
-    print("  2) Advanced inspection with forced zero bits")
+    print("  2) Advanced inspection with forced 0 or 1 bits")
     print("  3) Exhaustive variable-change search")
     print("  4) S box analysis")
     print("  5) variable XOR")