From 173426c0d1297d472e8148f5bfbf5b3d32518e71 Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Thu, 2 Jul 2026 09:23:06 +0200
Subject: [PATCH] clean exit

---
 src/tea3/cli.py | 62 ++++++++++++++++++++++++++++---------------------
 1 file changed, 36 insertions(+), 26 deletions(-)

diff --git a/src/tea3/cli.py b/src/tea3/cli.py
index a05b5b2..f81fa11 100644
--- a/src/tea3/cli.py
+++ b/src/tea3/cli.py
@@ -190,35 +190,45 @@ def run_exhaustive_xor_cli():
 
 
 def main():
-    print("=" * 50)
-    print(" Tea3 Model ")
-    print("=" * 50)
+    while True:
+        print("=" * 50)
+        print(" Tea3 Model ")
+        print("=" * 50)
 
-    print("\nChoose a mode:")
-    print("  1) Classic inspection")
-    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")
-    print("  6) F31, F32 analysis")
-    print("  7) Exhaustive XOR search")
+        print("\nChoose a mode:")
+        print("  1) Classic inspection")
+        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")
+        print("  6) F31, F32 analysis")
+        print("  7) Exhaustive XOR search")
+        print("  0) Exit")
 
-    mode = prompt_choice("Your choice (1, 2, 3, 4, 5, 6 or 7): ", {1, 2, 3, 4, 5, 6, 7})
+        mode = prompt_choice(
+            "Your choice (0-7): ",
+            {0, 1, 2, 3, 4, 5, 6, 7}
+        )
 
-    if mode == 1:
-        run_classic_cli()
-    elif mode == 2:
-        run_advanced_cli()
-    elif mode == 3:
-        run_exhaustive_cli()
-    elif mode == 4:
-        run_sbox()
-    elif mode == 5:
-        run_variable_xor_cli()
-    elif mode == 6:
-        run_f31f32()
-    elif mode == 7:
-        run_exhaustive_xor_cli()
+        if mode == 0:
+            print("\nGoodbye")
+            break
+        elif mode == 1:
+            run_classic_cli()
+        elif mode == 2:
+            run_advanced_cli()
+        elif mode == 3:
+            run_exhaustive_cli()
+        elif mode == 4:
+            run_sbox()
+        elif mode == 5:
+            run_variable_xor_cli()
+        elif mode == 6:
+            run_f31f32()
+        elif mode == 7:
+            run_exhaustive_xor_cli()
 
+        input("\nPress Enter to return to the main menu...")
+        print()
 
 main()