diff --git a/src/tea3/cli.py b/src/tea3/cli.py index 2a4c519..4f05682 100644 --- a/src/tea3/cli.py +++ b/src/tea3/cli.py @@ -1,5 +1,6 @@ from tea3.pretty_print import pretty_print from tea3.tea3model import Tea3Model +from tea3.variable_search import run_exhaustive def prompt_int(message: str, lo: int, hi: int) -> int: @@ -15,28 +16,28 @@ def prompt_int(message: str, lo: int, hi: int) -> int: print(f"Value must be between {lo} and {hi} (inclusive).") -def main() -> None: - print("=" * 50) - print(" Tea3 Model ") - print("=" * 50) +def prompt_choice(message: str, choices: set[int]) -> int: + choices_str = ", ".join(map(str, sorted(choices))) + while True: + raw = input(message).strip() + try: + value = int(raw) + except ValueError: + print("Please enter a number.") + continue + if value in choices: + return value + print(f"Value must be one of: {choices_str}") - steps = prompt_int("\nHow many steps do you want to run? (1–100): ", 1, 100) +def run_classic_cli(): print("\nR registers are indexed 0–7; bits within each register are 0–7.") print("Enter -1 to print all registers, or all bits.") + 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) 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("-" * 50) model = Tea3Model() @@ -59,4 +60,28 @@ def main() -> None: print("Done.") +def main(): + print("=" * 50) + print(" Tea3 Model ") + print("=" * 50) + + print("\nChoose a mode:") + print(" 1) Classic inspection") + print(" 2) Exhaustive variable-change search") + + mode = prompt_choice("Your choice (1 or 2): ", {1, 2}) + + if mode == 1: + run_classic_cli() + else: + steps = prompt_int("How many steps? (1–100): ", 1, 100) + target_reg = prompt_int("Target register (0–7): ", 0, 7) + + print("-" * 50) + run_exhaustive(steps, target_reg) + + print("\n" + "=" * 50) + print("Done.") + + main() diff --git a/src/tea3/tea3model.py b/src/tea3/tea3model.py index f15368a..2355e63 100644 --- a/src/tea3/tea3model.py +++ b/src/tea3/tea3model.py @@ -12,24 +12,24 @@ class Tea3Model: self.step_count = 0 names = ( - [f"x{i}{j}" for i in range(5) for j in range(8)] + - [f"r{i}{j}" for i in range(5) for j in range(8)] + - [f"R{i}{j}" for i in range(8) for j in range(8)] + - ["g"] + [f"x{i}{j}" for i in range(5) for j in range(8)] + # 0–39 + [f"y{i}{j}" for i in range(5) for j in range(8)] + # 40–79 + [f"r{i}{j}" for i in range(5) for j in range(8)] + # 80–119 + [f"R{i}{j}" for i in range(8) for j in range(8)] + # 120–183 + ["g"] # 184 ) name_string = ",".join(names) self.S = BooleanPolynomialRing(len(names), name_string) self.v = self.S.gens() - self.x_bits = [list(self.v[i*8:(i+1)*8]) for i in range(5)] - self.r_bits = [list(self.v[40 + i*8 : 40 + (i+1)*8]) for i in range(5)] - self.R_bits = [list(self.v[80 + i*8 : 80 + (i+1)*8]) for i in range(8)] - + self.x_bits = [list(self.v[i*8 : (i+1)*8 ]) for i in range(5)] + self.y_bits = [list(self.v[40 + i*8 : 40 + (i+1)*8]) for i in range(5)] + self.r_bits = [list(self.v[80 + i*8 : 80 + (i+1)*8]) for i in range(5)] + self.R_bits = [list(self.v[120 + i*8 : 120 + (i+1)*8]) for i in range(8)] self.g = self.v[-1] def _abstract_R(self): - s = self.step_count one = self.S.one() zero = self.S.zero() @@ -38,28 +38,28 @@ class Tea3Model: poly = self.R_bits[i][j] groups = {} - pure_xr = zero - const = one if bool(poly.constant_coefficient()) else zero + pure_xyr = zero + const = one if bool(poly.constant_coefficient()) else zero for monom in poly: term_vars = monom.variables() if not term_vars: continue - xr_vars = [v for v in term_vars if str(v)[0] in ('x', 'r')] - Rf_vars = [v for v in term_vars if str(v)[0] in ('R', 'f', 'g')] + xyr_vars = [v for v in term_vars if str(v)[0] in ('x', 'y', 'r')] + Rg_vars = [v for v in term_vars if str(v)[0] in ('R', 'g')] - xr_mono = reduce(mul, (self.S(v) for v in xr_vars), one) - xr_key = frozenset(str(v) for v in xr_vars) + xyr_mono = reduce(mul, (self.S(v) for v in xyr_vars), one) + xyr_key = frozenset(str(v) for v in xyr_vars) - if not Rf_vars: - pure_xr += xr_mono + if not Rg_vars: + pure_xyr += xyr_mono else: - groups[xr_key] = xr_mono + groups[xyr_key] = xyr_mono - result = pure_xr + const - for xr_key, xr_mono in groups.items(): - result += xr_mono * self.g + result = pure_xyr + const + for xyr_key, xyr_mono in groups.items(): + result += xyr_mono * self.g self.R_bits[i][j] = result @@ -101,6 +101,7 @@ class Tea3Model: return R7 + def S(r): # placeholder return r diff --git a/src/tea3/variable_search.py b/src/tea3/variable_search.py new file mode 100644 index 0000000..0c0c8b0 --- /dev/null +++ b/src/tea3/variable_search.py @@ -0,0 +1,40 @@ +from itertools import product as iproduct +from tea3.tea3model import Tea3Model +from tea3.pretty_print import pretty_print + +def apply_variable_change(model, coeffs): + orig_x = [list(model.v[i*8 : (i+1)*8]) for i in range(5)] + orig_y = model.y_bits + + subs = {} + + for i in range(1, 5): + for j in range(8): + subs[orig_x[i][j]] = orig_y[i][j] + + for j in range(8): + new_expr = orig_y[0][j] + for i, ai in enumerate(coeffs, start=1): + if ai: + new_expr = new_expr + orig_y[i][j] + subs[orig_x[0][j]] = new_expr + + new_R = [] + for i in range(8): + row = [poly.subs(subs) for poly in model.R_bits[i]] + new_R.append(row) + return new_R + +def run_exhaustive(steps: int, target_reg: int = 0): + for idx, coeffs in enumerate(iproduct([0, 1], repeat=4)): + model = Tea3Model() + for _ in range(steps): + model.step() + + new_R = apply_variable_change(model, coeffs) + + label = "".join(map(str, coeffs)) + print(f"\n[{idx:02d}] (a1,a2,a3,a4) = {label}") + for j in range(8): + poly = new_R[target_reg][j] + print(f" R[{target_reg}][{j}] = {pretty_print(poly)}")