From 147fe38e8f70d1482ee162600ae60994145d14b2 Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Tue, 7 Jul 2026 13:46:28 +0200
Subject: [PATCH] cube attack, use loaded model

---
 src/tea3/cli.py         |  6 ++++--
 src/tea3/cube_attack.py | 23 +++++++++++++++++------
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/src/tea3/cli.py b/src/tea3/cli.py
index 75f05a7..6ee5425 100644
--- a/src/tea3/cli.py
+++ b/src/tea3/cli.py
@@ -150,7 +150,7 @@ def run_exhaustive_xor_cli():
     print("Done.")
 
 
-def run_cube_attack_cli():
+def run_cube_attack_cli(model = None):
     print("\nCube attack search on TEA3 model to search for low-degree superpolys.")
 
     rounds = prompt_int("How many rounds? (1–100): ", 1, 100)
@@ -193,6 +193,7 @@ def run_cube_attack_cli():
         max_degree=max_degree,
         fixed_zero_bits=fixed_zero_bits,
         fixed_one_bits=fixed_one_bits,
+        model=model,
     )
 
     print("\n" + "=" * 50)
@@ -201,6 +202,7 @@ def run_cube_attack_cli():
 
 
 def main():
+    model = None
     while True:
         print("=" * 50)
         print(" Tea3 Model ")
@@ -242,7 +244,7 @@ def main():
         elif mode == 7:
             run_exhaustive_xor_cli()
         elif mode == 8:
-            run_cube_attack_cli()
+            run_cube_attack_cli(model)
         elif mode == 9:
             run_precompute_cli()
         elif mode == 10:
diff --git a/src/tea3/cube_attack.py b/src/tea3/cube_attack.py
index 02ef5c5..ad40d96 100644
--- a/src/tea3/cube_attack.py
+++ b/src/tea3/cube_attack.py
@@ -1,4 +1,5 @@
 from __future__ import annotations
+import copy
 
 from dataclasses import dataclass
 from itertools import combinations
@@ -36,18 +37,23 @@ def build_target_poly(
     target_reg: int = 7,
     target_bit: int = 0,
     fixed_bits: dict[object, int] | None = None,
+    base_model: Tea3Model | None = None,
 ):
     """
     Return one Boolean polynomial from the model after `rounds` steps.
-    The model is treated as an oracle. The returned polynomial is the chosen bit of the register state after `rounds` iterations.
+    If `base_model` is provided, continue from its state.
+    The model is treated as an oracle.
     """
-    model = Tea3Model()
+    if base_model is not None:
+        model = copy.deepcopy(base_model)
+    else:
+        model = Tea3Model()
 
     if fixed_bits:
         model = specialize_model(model, fixed_bits)
 
-    for _ in range(rounds):
-        model.step(skip_abstract = True)
+    for _ in range((rounds - base_model.step_count) if base_model is not None else rounds):
+        model.step(skip_abstract=True)
 
     return model.R_bits[target_reg][target_bit]
 
@@ -183,6 +189,7 @@ def run_cube_attack(
     max_degree: int = 1,
     fixed_zero_bits: Sequence[str] | None = None,
     fixed_one_bits: Sequence[str] | None = None,
+    model: Tea3Model | None = None,
 ):
     fixed_bits: dict[str, int] = {}
 
@@ -193,16 +200,20 @@ def run_cube_attack(
         for name in fixed_one_bits:
             fixed_bits[name] = 1
 
-    model = Tea3Model()
-    public_vars = pick_public_vars(model, fixed_bits=fixed_bits)
+    tmp_model = Tea3Model()
+    public_vars = pick_public_vars(tmp_model, fixed_bits=fixed_bits)
+
     poly = build_target_poly(
         rounds=rounds,
         target_reg=target_reg,
         target_bit=target_bit,
         fixed_bits=fixed_bits,
+        base_model=model,
     )
 
     print("=" * 50)
+    if model is not None:
+        print(f"Starting from precomputed model ({model.step_count} steps already done).")
     print(f"Target: {rounds}-round output bit R{target_reg}[{target_bit}]")
     print(f"Public variables: {len(public_vars)} R bits")
     print(f"Cube size: {cube_size}")