Add Parallel web search to interactive PentestGPT

pentestgpt_legacy/utils/pentest_gpt.py

@@ -35,11 +35,10 @@
     main_task_entry,
     mainTaskCompleter,
 )
+from pentestgpt_legacy.utils.web_search import search_web
 
 logger = loguru.logger
 
-_GOOGLE_PLACEHOLDER = "Google search results:\nstill under development."
-
 
 class pentestGPT:
     postfix_options: ClassVar[dict[str, str]] = {
@@ -204,8 +203,23 @@ def test_generation_handler(self, text: str) -> str:
         self.log_conversation("generation", response)
         return response
 
+    def web_search_handler(self) -> str:
+        query = prompt_ask("Web search query (sent to Parallel): ", multiline=False)
+        self.log_conversation("user", f"Web search: {query}")
+        try:
+            with self.console.status("[bold green] Searching the web..."):
+                response = search_web(query)
+        except ValueError as exc:
+            response = str(exc)
+        except Exception as exc:
+            logger.warning("Web search failed: {}", type(exc).__name__)
+            response = "Web search failed. Please try again."
+        self.console.print(response + "\n", style="yellow")
+        self.log_conversation("pentestGPT", response)
+        return response
+
     def local_input_handler(self) -> str:
-        """Handle the sub-task ('more') loop: discuss / brainstorm / google / continue."""
+        """Handle the sub-task ('more') loop: discuss / brainstorm / search / continue."""
         local_task_response = ""
         self.chat_count += 1
         local_request_option = local_task_entry()
@@ -246,11 +260,8 @@ def local_input_handler(self) -> str:
             self.console.print(local_task_response + "\n", style="yellow")
             self.log_conversation("pentestGPT", local_task_response)
 
-        elif local_request_option == "google":
-            self.console.print("Google integration is still under development.", style="bold green")
-            self.log_conversation("pentestGPT", _GOOGLE_PLACEHOLDER)
-            self.console.print(_GOOGLE_PLACEHOLDER + "\n", style="yellow")
-            return _GOOGLE_PLACEHOLDER
+        elif local_request_option in {"search", "google"}:
+            return self.web_search_handler()
 
         elif local_request_option == "continue":
             self.console.print("Exit the local task and continue the main task.")
@@ -260,7 +271,7 @@ def local_input_handler(self) -> str:
         return local_task_response
 
     def input_handler(self) -> Any:
-        """Main REPL dispatch: next / more / todo / discuss / google / help / quit."""
+        """Main REPL dispatch: next / more / todo / discuss / search / help / quit."""
         self.chat_count += 1
 
         request_option = main_task_entry()
@@ -378,11 +389,8 @@ def input_handler(self) -> Any:
             self.console.print(response + "\n", style="yellow")
             self.log_conversation("pentestGPT", response)
 
-        elif request_option == "google":
-            self.console.print("Google integration is still under development.", style="bold green")
-            self.log_conversation("pentestGPT", _GOOGLE_PLACEHOLDER)
-            self.console.print(_GOOGLE_PLACEHOLDER + "\n", style="yellow")
-            return _GOOGLE_PLACEHOLDER
+        elif request_option in {"search", "google"}:
+            return self.web_search_handler()
 
         elif request_option == "quit":
             response = False

pentestgpt_legacy/utils/task_handler.py

@@ -18,7 +18,8 @@ class localTaskCompleter(Completer):
         "discuss",  # discuss with pentestGPT on the local task
         "brainstorm",  # let pentestGPT brainstorm on the local task
         "help",  # show the help page (for this local task)
-        "google",  # search on Google
+        "search",  # search the public web
+        "google",  # compatibility alias for the old placeholder command
         "continue",  # quit the local task  (for this local task)
     ]
 
@@ -28,7 +29,8 @@ class localTaskCompleter(Completer):
             "Let <b>PentestGPT</b> brainstorm on the local task for all the possible solutions."
         ),
         "help": HTML("Show the help page for this local task."),
-        "google": HTML("Search on Google."),
+        "search": HTML("Search the web with Parallel."),
+        "google": HTML("Alias for search (Parallel)."),
         "continue": HTML("Quit the local task and continue the previous testing."),
     }
 
@@ -37,7 +39,9 @@ class localTaskCompleter(Completer):
     - discuss: Discuss with PentestGPT about this local task.
     - brainstorm: Let PentestGPT brainstorm on the local task for all the possible solutions.
     - help: Show the help page for this local task.
-    - google: Search on Google.
+    - search: Search the web with Parallel (https://search.parallel.ai/mcp).
+      Your query is sent to Parallel. No key is needed; free access is rate limited.
+    - google: Alias for search; uses Parallel.
     - quit: Quit the local task and continue the testing."""
 
     def get_completions(self, document, complete_event):
@@ -58,6 +62,7 @@ class mainTaskCompleter(Completer):
         "more",
         "todo",
         "discuss",
+        "search",
         "google",
         "help",
         "quit",
@@ -68,7 +73,8 @@ class mainTaskCompleter(Completer):
         "more": HTML("Explain the task with more details."),
         "todo": HTML("Ask <b>PentestGPT</b> for todos."),
         "discuss": HTML("Discuss with <b>PentestGPT</b>."),
-        "google": HTML("Search on Google."),
+        "search": HTML("Search the web with Parallel."),
+        "google": HTML("Alias for search (Parallel)."),
         "help": HTML("Show the help page."),
         "quit": HTML("End the current session."),
     }
@@ -79,7 +85,9 @@ class mainTaskCompleter(Completer):
  - more: Explain the previous given task with more details.
  - todo: Ask PentestGPT for the task list and what to do next.
  - discuss: Discuss with PentestGPT. You can ask for help, discuss the task, or give any feedbacks.
- - google: Search your question on Google. The results are automatically parsed by Google.
+ - search: Search the web with Parallel (https://search.parallel.ai/mcp).
+   Your query is sent to Parallel. No key is needed; free access is rate limited.
+ - google: Alias for search; uses Parallel.
  - help: Show this help page.
  - quit: End the current session."""
 

pentestgpt_legacy/utils/web_search.py

@@ -0,0 +1,70 @@
+"""Interactive web search through Parallel's free Search MCP server."""
+
+from __future__ import annotations
+
+import json
+from importlib.metadata import version
+from typing import Any
+
+import anyio
+import httpx
+from mcp import ClientSession
+from mcp.client.streamable_http import streamable_http_client
+
+ENDPOINT = "https://search.parallel.ai/mcp"
+
+
+def _format_results(payload: dict[str, Any]) -> str:
+    results = payload.get("results")
+    if not isinstance(results, list):
+        raise ValueError("Parallel returned an unexpected search response.")
+    if not results:
+        return "No web search results found."
+
+    lines = ["Web search results from Parallel:"]
+    for result in results[:5]:
+        if not isinstance(result, dict) or not isinstance(result.get("url"), str):
+            raise ValueError("Parallel returned an unexpected search result.")
+        title = result.get("title") or result["url"]
+        lines.append(f"\n{title}\n{result['url']}")
+        excerpts = result.get("excerpts")
+        if isinstance(excerpts, list):
+            for excerpt in excerpts[:2]:
+                if isinstance(excerpt, str) and excerpt.strip():
+                    lines.append(excerpt.strip()[:800])
+    return "\n".join(lines)
+
+
+async def _search(
+    query: str, *, endpoint: str = ENDPOINT, transport: httpx.AsyncBaseTransport | None = None
+) -> str:
+    if not query.strip() or len(query) > 200:
+        raise ValueError("Enter a search query of 1 to 200 characters.")
+
+    # Identify this project for aggregate free MCP usage; do not identify a user or session.
+    user_agent = f"PentestGPT/{version('pentestgpt')} mcp-python"
+    async with (
+        httpx.AsyncClient(
+            headers={"User-Agent": user_agent}, timeout=30, transport=transport
+        ) as http_client,
+        streamable_http_client(endpoint, http_client=http_client) as (read, write, _),
+        ClientSession(read, write) as session,
+    ):
+        await session.initialize()
+        result = await session.call_tool(
+            "web_search", {"objective": query, "search_queries": [query]}
+        )
+        if result.isError:
+            raise RuntimeError("Parallel could not complete the web search.")
+        payload = result.structuredContent
+        if payload is None:
+            text = next((item.text for item in result.content if item.type == "text"), None)
+            payload = json.loads(text) if text else None
+        if not isinstance(payload, dict):
+            raise ValueError("Parallel returned an unexpected search response.")
+        return _format_results(payload)
+
+
+def search_web(query: str) -> str:
+    """Search the public web when an interactive user requests it."""
+    return anyio.run(_search, query)

tests/legacy/test_web_search.py

@@ -0,0 +1,119 @@
+"""Interactive search reaches the actual MCP client and displays cited results."""
+
+from __future__ import annotations
+
+import json
+from io import StringIO
+
+import httpx
+import pytest
+from rich.console import Console
+
+from pentestgpt_legacy.utils import pentest_gpt, web_search
+
+
+@pytest.mark.asyncio
+async def test_search_uses_parallel_mcp_with_project_identity() -> None:
+    requests: list[httpx.Request] = []
+
+    def respond(request: httpx.Request) -> httpx.Response:
+        requests.append(request)
+        message = json.loads(request.content)
+        method = message["method"]
+        if method == "initialize":
+            result = {
+                "protocolVersion": "2025-03-26",
+                "capabilities": {"tools": {}},
+                "serverInfo": {"name": "fixture", "version": "1"},
+            }
+        elif method == "tools/list":
+            result = {
+                "tools": [
+                    {
+                        "name": "web_search",
+                        "description": "Search the web",
+                        "inputSchema": {
+                            "type": "object",
+                            "properties": {
+                                "objective": {"type": "string"},
+                                "search_queries": {"type": "array", "items": {"type": "string"}},
+                            },
+                            "required": ["objective", "search_queries"],
+                        },
+                        "outputSchema": {
+                            "type": "object",
+                            "properties": {"results": {"type": "array"}},
+                        },
+                    }
+                ]
+            }
+        elif method == "tools/call":
+            result = {
+                "content": [],
+                "structuredContent": {
+                    "results": [
+                        {
+                            "title": "Official reference",
+                            "url": "https://example.org/reference",
+                            "excerpts": ["A useful finding."],
+                        }
+                    ]
+                },
+                "isError": False,
+            }
+        else:
+            return httpx.Response(202)
+        return httpx.Response(
+            200,
+            json={"jsonrpc": "2.0", "id": message["id"], "result": result},
+            headers={"Content-Type": "application/json"},
+        )
+
+    output = await web_search._search(
+        "public reference",
+        endpoint="https://example.org/mcp",
+        transport=httpx.MockTransport(respond),
+    )
+
+    calls = [json.loads(request.content) for request in requests]
+    assert [call["method"] for call in calls] == [
+        "initialize",
+        "notifications/initialized",
+        "tools/call",
+        "tools/list",
+    ]
+    assert calls[2]["params"] == {
+        "name": "web_search",
+        "arguments": {"objective": "public reference", "search_queries": ["public reference"]},
+    }
+    assert all(
+        request.headers["user-agent"] == "PentestGPT/1.0.0 mcp-python" for request in requests
+    )
+    assert all("authorization" not in request.headers for request in requests)
+    assert "https://example.org/reference" in output
+    assert "A useful finding." in output
+
+
+@pytest.mark.parametrize("local", [False, True])
+def test_interactive_search_command_shows_results(
+    monkeypatch: pytest.MonkeyPatch, local: bool
+) -> None:
+    agent = object.__new__(pentest_gpt.pentestGPT)
+    output = StringIO()
+    agent.console = Console(file=output)
+    agent.chat_count = 0
+    agent.history = {"user": [], "pentestGPT": []}
+    monkeypatch.setattr(pentest_gpt, "prompt_ask", lambda *args, **kwargs: "public reference")
+    monkeypatch.setattr(
+        pentest_gpt, "search_web", lambda query: f"Result: {query}\nhttps://example.org/reference"
+    )
+    if local:
+        monkeypatch.setattr(pentest_gpt, "local_task_entry", lambda: "search")
+        result = agent.local_input_handler()
+    else:
+        monkeypatch.setattr(pentest_gpt, "main_task_entry", lambda: "search")
+        result = agent.input_handler()
+
+    assert "https://example.org/reference" in result
+    assert "https://example.org/reference" in output.getvalue()
+    assert agent.history["user"][-1][1] == "Web search: public reference"