diff --git a/src/tea3/cli.py b/src/tea3/cli.py index a4b965a..2a4c519 100644 --- a/src/tea3/cli.py +++ b/src/tea3/cli.py @@ -23,19 +23,37 @@ def main() -> None: steps = prompt_int("\nHow many steps do you want to run? (1–100): ", 1, 100) 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) - bit = prompt_int("Which bit of that register? (0–7): ", 0, 7) + print("Enter -1 to print all registers, or all bits.") + + 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) + + if reg == -1 and bit == -1: + print(f"\nRunning {steps} step(s), watching all R registers and all bits") + elif reg == -1: + print(f"\nRunning {steps} step(s), watching all R registers, bit [{bit}]") + elif bit == -1: + print(f"\nRunning {steps} step(s), watching R[{reg}] and all bits") + else: + print(f"\nRunning {steps} step(s), watching R[{reg}][{bit}]") - print(f"\nRunning {steps} step(s), watching R[{reg}][{bit}]") print("-" * 50) model = Tea3Model() for i in range(steps): model.step() - poly = model.R_bits[reg][bit] - print(f"\n[Step {i+1}] R_bits[{reg}][{bit}] =") - print(pretty_print(poly)) + print(f"\n[Step {i + 1}]") + + 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.")