fix(legacy): handle cancelled source selection

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)

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()