diff --git a/strix/tools/reporting/tool.py b/strix/tools/reporting/tool.py index c2a6abf7..164dd3af 100644 --- a/strix/tools/reporting/tool.py +++ b/strix/tools/reporting/tool.py @@ -2539,6 +2539,54 @@ def _sandbox_session(ctx: RunContextWrapper) -> Any: return inner.get("sandbox_session") +def _entry_kind(entry: Any) -> str: + kind = getattr(entry, "kind", None) + if kind is None: + return "file" + value = getattr(kind, "value", kind) + return str(value).lower() + + +def _entry_matches_path(entry: Any, path: str) -> bool: + entry_path = str(getattr(entry, "path", "") or "") + if not entry_path: + return False + wanted = PurePosixPath(path) + listed = PurePosixPath(entry_path) + return entry_path in {path, wanted.name} or listed.name == wanted.name + + +async def _inspect_workspace_file(session: Any, path: str) -> str | None: + """Reject symlinks and oversized files before the sandbox read follows them.""" + ls = getattr(session, "ls", None) + if not callable(ls): + return None + parent = PurePosixPath(path).parent.as_posix() + try: + entries = await ls(Path(parent)) + except Exception as exc: + logger.exception("export_markdown_report ls failed") + return f"Failed to inspect '{path}' in the sandbox: {exc!s}" + entry = next((item for item in entries or [] if _entry_matches_path(item, path)), None) + if entry is None: + return f"Markdown file not found: {path}" + return _reject_workspace_entry(path, entry) + + +def _reject_workspace_entry(path: str, entry: Any) -> str | None: + kind = _entry_kind(entry) + if kind == "symlink": + return f"'{path}' is a symlink; export a regular markdown file inside /workspace" + if kind != "file": + return f"'{path}' is not a regular markdown file" + size = getattr(entry, "size", None) + if isinstance(size, int) and size <= 0: + return f"Markdown file is empty: {path}" + if isinstance(size, int) and size > _MAX_EXPORTED_REPORT_BYTES: + return f"Markdown file exceeds {_MAX_EXPORTED_REPORT_BYTES} bytes ({size} bytes)" + return None + + def _export_destination( ctx: RunContextWrapper, path: str ) -> tuple[Any, Path, str, str] | dict[str, Any]: @@ -2563,7 +2611,33 @@ def _export_destination( return session, report_state.get_run_dir(), workspace_path, name +def _read_limited_bytes(stream: Any) -> tuple[bytes | None, str | None]: + """Copy at most one extra byte past the export limit, then close ``stream``.""" + limit = _MAX_EXPORTED_REPORT_BYTES + 1 + try: + chunks: list[bytes] = [] + remaining = limit + while remaining > 0: + chunk = stream.read(remaining) + if not chunk: + break + if isinstance(chunk, str): + chunk = chunk.encode("utf-8") + if not isinstance(chunk, bytes | bytearray): + return None, "sandbox read did not return file bytes" + chunks.append(bytes(chunk)) + remaining -= len(chunk) + return b"".join(chunks), None + finally: + close = getattr(stream, "close", None) + if callable(close): + close() + + async def _read_workspace_markdown(session: Any, path: str) -> tuple[bytes | None, str | None]: + inspect_error = await _inspect_workspace_file(session, path) + if inspect_error is not None: + return None, inspect_error try: stream = await session.read(Path(path)) except FileNotFoundError: @@ -2571,17 +2645,14 @@ async def _read_workspace_markdown(session: Any, path: str) -> tuple[bytes | Non except Exception as exc: logger.exception("export_markdown_report read failed") return None, f"Failed to read '{path}' from the sandbox: {exc!s}" - try: - data = stream.read() - finally: - close = getattr(stream, "close", None) - if callable(close): - close() - if isinstance(data, str): - return data.encode("utf-8"), None - if isinstance(data, bytes | bytearray): - return bytes(data), None - return None, "sandbox read did not return file bytes" + data, read_error = _read_limited_bytes(stream) + if data is None: + return None, read_error + if len(data) > _MAX_EXPORTED_REPORT_BYTES: + return None, ( + f"Markdown file exceeds {_MAX_EXPORTED_REPORT_BYTES} bytes ({len(data)} bytes)" + ) + return data, None def _markdown_text(path: str, payload: bytes) -> tuple[str | None, str | None]: @@ -2605,6 +2676,15 @@ async def _export_markdown_report(ctx: RunContextWrapper, path: str) -> dict[str if isinstance(ready, dict): return ready session, run_dir, workspace_path, name = ready + destination = run_dir / name + if destination.exists(): + return { + "success": False, + "error": ( + f"'{name}' already exists in the host run directory; " + "choose another markdown filename" + ), + } payload, read_error = await _read_workspace_markdown(session, workspace_path) if payload is None: @@ -2613,7 +2693,6 @@ async def _export_markdown_report(ctx: RunContextWrapper, path: str) -> dict[str if text is None: return {"success": False, "error": text_error} - destination = run_dir / name try: atomic_write_text(destination, text) except OSError as exc: diff --git a/tests/test_export_markdown_report.py b/tests/test_export_markdown_report.py index ae16cf89..e600caf1 100644 --- a/tests/test_export_markdown_report.py +++ b/tests/test_export_markdown_report.py @@ -4,31 +4,83 @@ from __future__ import annotations import io import json -from collections.abc import Iterator from typing import TYPE_CHECKING, Any import pytest +from agents.sandbox.files import EntryKind, FileEntry +from agents.sandbox.types import Permissions from agents.tool_context import ToolContext from strix.report.state import ReportState, set_global_report_state -from strix.tools.reporting.tool import export_markdown_report +from strix.tools.reporting.tool import _MAX_EXPORTED_REPORT_BYTES, export_markdown_report if TYPE_CHECKING: + from collections.abc import Iterator from pathlib import Path +def _file_entry(path: str, payload: bytes, *, kind: EntryKind = EntryKind.FILE) -> FileEntry: + return FileEntry( + path=path, + permissions=Permissions.from_mode(0o100644), + owner="root", + group="root", + size=len(payload), + kind=kind, + ) + + +class _SimpleEntry: + def __init__(self, path: str, *, kind: str = "file", size: int | None = None) -> None: + self.path = path + self.kind = kind + if size is not None: + self.size = size + + class _Sandbox: - def __init__(self, files: dict[str, bytes]) -> None: + def __init__( + self, + files: dict[str, bytes], + *, + kinds: dict[str, EntryKind] | None = None, + listed: dict[str, object] | None = None, + ) -> None: self.files = files + self.kinds = kinds or {} + self.listed = listed or {} self.reads: list[str] = [] + self.bytes_read = 0 + + async def ls(self, path: Path) -> list[object]: + parent = path.as_posix().rstrip("/") or "/" + entries: list[object] = [] + for file_path, payload in self.files.items(): + if file_path.rsplit("/", 1)[0] == parent: + entries.append( + self.listed.get(file_path) + or _file_entry( + file_path, payload, kind=self.kinds.get(file_path, EntryKind.FILE) + ) + ) + return entries async def read(self, path: Path) -> io.BytesIO: self.reads.append(path.as_posix()) payload = self.files.get(path.as_posix()) if payload is None: raise FileNotFoundError(path.as_posix()) - return io.BytesIO(payload) + + sandbox = self + + class _CountedBytes(io.BytesIO): + def read(self, size: int | None = -1) -> bytes: + chunk = super().read(-1 if size is None else size) + sandbox.bytes_read += len(chunk) + return chunk + + return _CountedBytes(payload) @pytest.fixture @@ -97,3 +149,66 @@ async def test_rejects_paths_outside_a_named_workspace_markdown_file( assert sandbox.reads == [] assert not (report_state.get_run_dir() / "report.md").exists() assert not (report_state.get_run_dir() / "penetration_test_report.md").exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "kind", "match"), + [ + (b"", EntryKind.FILE, "empty"), + (b"# Findings\n\x00hidden\n", EntryKind.FILE, "markdown text file"), + (b"\xff\xfe# Findings\n", EntryKind.FILE, "UTF-8"), + (b"x" * (_MAX_EXPORTED_REPORT_BYTES + 1), EntryKind.FILE, "exceeds"), + (b"# Findings\n", EntryKind.SYMLINK, "symlink"), + ], +) +async def test_rejects_empty_invalid_oversized_and_symlink_reports( + report_state: ReportState, + payload: bytes, + kind: EntryKind, + match: str, +) -> None: + sandbox = _Sandbox({"/workspace/report.md": payload}, kinds={"/workspace/report.md": kind}) + + result = await _export({"sandbox_session": sandbox}, "/workspace/report.md") + + assert result["success"] is False + assert match in result["error"] + assert not (report_state.get_run_dir() / "report.md").exists() + if kind is EntryKind.SYMLINK or match in {"empty", "exceeds"}: + assert sandbox.reads == [] + if match == "exceeds": + assert sandbox.bytes_read == 0 + + +@pytest.mark.asyncio +async def test_size_limit_stops_the_sandbox_read_after_one_extra_byte( + report_state: ReportState, +) -> None: + oversized = b"x" * (_MAX_EXPORTED_REPORT_BYTES + 64) + sandbox = _Sandbox( + {"/workspace/report.md": oversized}, + listed={"/workspace/report.md": _SimpleEntry("/workspace/report.md")}, + ) + + result = await _export({"sandbox_session": sandbox}, "/workspace/report.md") + + assert result["success"] is False + assert "exceeds" in result["error"] + assert sandbox.bytes_read == _MAX_EXPORTED_REPORT_BYTES + 1 + assert not (report_state.get_run_dir() / "report.md").exists() + + +@pytest.mark.asyncio +async def test_same_named_exports_do_not_overwrite_an_existing_report( + report_state: ReportState, +) -> None: + run_dir = report_state.get_run_dir() + (run_dir / "report.md").write_text("# first\n", encoding="utf-8") + sandbox = _Sandbox({"/workspace/notes/report.md": b"# second\n"}) + + result = await _export({"sandbox_session": sandbox}, "/workspace/notes/report.md") + + assert result["success"] is False + assert "already exists" in result["error"] + assert (run_dir / "report.md").read_text(encoding="utf-8") == "# first\n"