diff --git a/src/tea3/cli.py b/src/tea3/cli.py index 69e2938..a05b5b2 100644 --- a/src/tea3/cli.py +++ b/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")