fix: contain legacy session saves under test_history

pentestgpt_legacy/utils/pentest_gpt.py

@@ -41,6 +41,32 @@
 _GOOGLE_PLACEHOLDER = "Google search results:\nstill under development."
 
 
+def resolve_session_save_path(save_root: str, save_name: str) -> str:
+    """Return a file path that is a direct child of ``save_root``.
+
+    Session names must be plain filenames. Absolute paths, nested components,
+    ``.`` / ``..``, and any resolved path whose parent is not ``save_root`` are
+    rejected so transcripts cannot escape the session directory.
+
+    Raises:
+        ValueError: if ``save_name`` is not a safe basename under ``save_root``.
+    """
+    if not save_name or not save_name.strip():
+        raise ValueError("session name must be a non-empty filename")
+    if save_name in {".", ".."}:
+        raise ValueError("session name must be a plain filename")
+    if os.path.isabs(save_name):
+        raise ValueError("session name must not be an absolute path")
+    if os.path.basename(save_name) != save_name:
+        raise ValueError("session name must be a plain filename")
+
+    root = os.path.realpath(save_root)
+    candidate = os.path.realpath(os.path.join(root, save_name))
+    if os.path.dirname(candidate) != root:
+        raise ValueError("session path escapes the save directory")
+    return candidate
+
+
 class pentestGPT:
     postfix_options: ClassVar[dict[str, str]] = {
         "tool": "The input content is from a security testing tool. You need to list down all the points that are interesting to you; you should summarize it as if you are reporting to a senior penetration tester for further guidance.\n",
@@ -401,17 +427,32 @@ def save_session(self) -> None:
             "Before you quit, you may want to save the current session.",
             style="bold green",
         )
-        save_name = prompt_ask(
-            "Please enter the name of the current session. (Default with current timestamp)\n> ",
-            multiline=False,
-        )
-        if save_name == "":
-            save_name = str(time.time())
         save_root = os.path.join(os.getcwd(), self.save_dir)
         os.makedirs(save_root, exist_ok=True)
-        with open(os.path.join(save_root, save_name), "w") as f:
-            json.dump({"task_log": self.task_log, "history": self.history}, f)
-        self.console.print(f"The current session is saved as {save_name}", style="bold green")
+
+        while True:
+            save_name = prompt_ask(
+                "Please enter the name of the current session. (Default with current timestamp)\n> ",
+                multiline=False,
+            )
+            if save_name == "":
+                save_name = str(time.time())
+            try:
+                save_path = resolve_session_save_path(save_root, save_name)
+            except ValueError as exc:
+                self.console.print(
+                    f"Invalid session name: {exc}. Please enter a plain filename.",
+                    style="bold red",
+                )
+                continue
+
+            with open(save_path, "w", encoding="utf-8") as f:
+                json.dump({"task_log": self.task_log, "history": self.history}, f)
+            self.console.print(
+                f"The current session is saved as {os.path.basename(save_path)}",
+                style="bold green",
+            )
+            return
 
     def main(self) -> None:
         """Run the interactive session."""

tests/legacy/test_session_save.py

@@ -0,0 +1,99 @@
+"""Regression tests for legacy session path containment (#502)."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+import pytest
+from rich.console import Console
+
+from pentestgpt_legacy.utils import pentest_gpt as pentest_gpt_mod
+from pentestgpt_legacy.utils.pentest_gpt import pentestGPT, resolve_session_save_path
+
+pytestmark = pytest.mark.unit
+
+
+def test_resolve_rejects_traversal(tmp_path) -> None:
+    root = tmp_path / "test_history"
+    root.mkdir()
+    with pytest.raises(ValueError):
+        resolve_session_save_path(str(root), "../outside.json")
+
+
+def test_resolve_rejects_nested(tmp_path) -> None:
+    root = tmp_path / "test_history"
+    root.mkdir()
+    with pytest.raises(ValueError):
+        resolve_session_save_path(str(root), "nested/session.json")
+
+
+def test_resolve_rejects_absolute(tmp_path) -> None:
+    root = tmp_path / "test_history"
+    root.mkdir()
+    absolute = str(tmp_path / "elsewhere.json")
+    with pytest.raises(ValueError):
+        resolve_session_save_path(str(root), absolute)
+
+
+def test_resolve_accepts_plain_filename(tmp_path) -> None:
+    root = tmp_path / "test_history"
+    root.mkdir()
+    path = resolve_session_save_path(str(root), "session.json")
+    assert path == os.path.realpath(str(root / "session.json"))
+    assert os.path.dirname(path) == os.path.realpath(str(root))
+
+
+def test_resolve_rejects_dot_and_dotdot(tmp_path) -> None:
+    root = tmp_path / "test_history"
+    root.mkdir()
+    with pytest.raises(ValueError):
+        resolve_session_save_path(str(root), ".")
+    with pytest.raises(ValueError):
+        resolve_session_save_path(str(root), "..")
+
+
+def _session_stub(monkeypatch: pytest.MonkeyPatch, cwd: Any, answers: list[str]) -> pentestGPT:
+    """Build a minimal pentestGPT without LLM clients."""
+    session = object.__new__(pentestGPT)
+    session.save_dir = "test_history"
+    session.task_log = {"step": 1}
+    session.history = {"user": [(1.0, "hi")]}
+    session.console = Console(force_terminal=False)
+    monkeypatch.chdir(cwd)
+    queue = list(answers)
+
+    def fake_prompt_ask(_prompt: str, multiline: bool = False) -> str:
+        del _prompt, multiline
+        assert queue, "prompt_ask called with no remaining answers"
+        return queue.pop(0)
+
+    monkeypatch.setattr(pentest_gpt_mod, "prompt_ask", fake_prompt_ask)
+    return session
+
+
+def test_save_session_reprompts_and_stays_in_test_history(
+    tmp_path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    outside = tmp_path / "outside.json"
+    session = _session_stub(
+        monkeypatch,
+        tmp_path,
+        answers=["../outside.json", "nested/session.json", "good-session.json"],
+    )
+    session.save_session()
+
+    assert not outside.exists()
+    saved = tmp_path / "test_history" / "good-session.json"
+    assert saved.is_file()
+    payload = json.loads(saved.read_text(encoding="utf-8"))
+    assert payload == {"task_log": {"step": 1}, "history": {"user": [[1.0, "hi"]]}}
+
+
+def test_save_session_empty_name_uses_timestamp(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(pentest_gpt_mod.time, "time", lambda: 12345.0)
+    session = _session_stub(monkeypatch, tmp_path, answers=[""])
+    session.save_session()
+    saved = tmp_path / "test_history" / "12345.0"
+    assert saved.is_file()