From 27c4d2a6b02eac05350bad9cdb1dc5ffeb7f5719 Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Thu, 9 Jul 2026 10:58:33 +0200
Subject: [PATCH] error handling cube attack, less round than model precomputed
 rounds

---
 src/tea3/cube_attack.py | 26 ++++++++++++++++++--------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/tea3/cube_attack.py b/src/tea3/cube_attack.py
index ad40d96..d0e6860 100644
--- a/src/tea3/cube_attack.py
+++ b/src/tea3/cube_attack.py
@@ -52,7 +52,10 @@ def build_target_poly(
     if fixed_bits:
         model = specialize_model(model, fixed_bits)
 
-    for _ in range((rounds - base_model.step_count) if base_model is not None else rounds):
+    start = model.step_count
+    if rounds < start:
+        raise ValueError(f"rounds={rounds} is smaller than current step_count={start}")
+    for _ in range(rounds - start):
         model.step(skip_abstract=True)
 
     return model.R_bits[target_reg][target_bit]
@@ -203,13 +206,20 @@ def run_cube_attack(
     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,
-    )
+    try:
+        poly = build_target_poly(
+            rounds=rounds,
+            target_reg=target_reg,
+            target_bit=target_bit,
+            fixed_bits=fixed_bits,
+            base_model=model,
+        )
+    except ValueError as e:
+        print("=" * 50)
+        print("Cube attack aborted.")
+        print(f"Error: {e}")
+        print("=" * 50)
+        return
 
     print("=" * 50)
     if model is not None: