From fb5049b0d8aa9642a6c9077752096b7012f81f53 Mon Sep 17 00:00:00 2001
From: Iams4kura <126048986+Iams4kura@users.noreply.github.com>
Date: Wed, 19 Aug 2026 11:27:49 +0800
Subject: [PATCH] fix(legacy): handle cancelled source selection

---
 pentestgpt_legacy/utils/pentest_gpt.py |  5 +++++
 tests/legacy/test_pentest_gpt.py       | 22 ++++++++++++++++++++++
 2 files changed, 27 insertions(+)
 create mode 100644 tests/legacy/test_pentest_gpt.py

diff --git a/pentestgpt_legacy/utils/pentest_gpt.py b/pentestgpt_legacy/utils/pentest_gpt.py
index 3573d147..4994ac69 100644
--- a/pentestgpt_legacy/utils/pentest_gpt.py
+++ b/pentestgpt_legacy/utils/pentest_gpt.py
@@ -285,6 +285,11 @@ def input_handler(self) -> Any:
             source = prompt_select(
                 title="Please choose the source of the information.", values=value_list
             )
+            if source is None:
+                message = "Source selection cancelled."
+                self.console.print(message, style="yellow")
+                self.log_conversation("pentestGPT", message)
+                return message
             self.console.print("Your input: (End with <shift + right-arrow>)", style="bold green")
             user_input = prompt_ask("> ", multiline=True)
             self.log_conversation("user", f"Source: {options[int(source)]}" + "\n" + user_input)
diff --git a/tests/legacy/test_pentest_gpt.py b/tests/legacy/test_pentest_gpt.py
new file mode 100644
index 00000000..3840dc98
--- /dev/null
+++ b/tests/legacy/test_pentest_gpt.py
@@ -0,0 +1,22 @@
+from unittest.mock import Mock
+
+import pentestgpt_legacy.utils.pentest_gpt as pentest_gpt_module
+
+
+def test_input_handler_returns_to_menu_when_source_selection_is_cancelled(
+    monkeypatch,
+) -> None:
+    app = object.__new__(pentest_gpt_module.pentestGPT)
+    app.chat_count = 0
+    app.console = Mock()
+    app.log_conversation = Mock()
+
+    monkeypatch.setattr(pentest_gpt_module, "main_task_entry", lambda: "next")
+    monkeypatch.setattr(pentest_gpt_module, "prompt_select", lambda **_kwargs: None)
+    prompt_ask = Mock(return_value="unused input")
+    monkeypatch.setattr(pentest_gpt_module, "prompt_ask", prompt_ask)
+
+    result = app.input_handler()
+
+    assert result == "Source selection cancelled."
+    prompt_ask.assert_not_called()