diff --git a/pyrit/cli/api_client.py b/pyrit/cli/api_client.py index 14e7588855..258c3e99bc 100644 --- a/pyrit/cli/api_client.py +++ b/pyrit/cli/api_client.py @@ -171,17 +171,20 @@ async def health_check_async(self) -> bool: # Scenarios # ------------------------------------------------------------------ - async def list_scenarios_async(self, *, limit: int = 200) -> list[RegisteredScenario]: + async def list_scenarios_async(self, *, limit: int | None = None) -> list[RegisteredScenario]: """ List all available scenarios. + Args: + limit (int | None): Maximum number of scenarios to return. None returns all scenarios. + Returns: - list[RegisteredScenario]: All scenarios in the catalog. + list[RegisteredScenario]: Scenarios in catalog order. """ from pyrit.models.catalog import RegisteredScenario - payload = await self._get_json_async(path="/api/scenarios/catalog", params={"limit": limit}) - return [RegisteredScenario.model_validate(item) for item in payload.get("items", [])] + items = await self._get_catalog_items_async(path="/api/scenarios/catalog", limit=limit) + return [RegisteredScenario.model_validate(item) for item in items] async def get_scenario_async(self, *, scenario_name: str) -> RegisteredScenario | None: """ @@ -209,17 +212,20 @@ async def get_scenario_async(self, *, scenario_name: str) -> RegisteredScenario # Initializers # ------------------------------------------------------------------ - async def list_initializers_async(self, *, limit: int = 200) -> list[RegisteredInitializer]: + async def list_initializers_async(self, *, limit: int | None = None) -> list[RegisteredInitializer]: """ List all available initializers. + Args: + limit (int | None): Maximum number of initializers to return. None returns all initializers. + Returns: - list[RegisteredInitializer]: All initializers in the catalog. + list[RegisteredInitializer]: Initializers in catalog order. """ from pyrit.models.catalog import RegisteredInitializer - payload = await self._get_json_async(path="/api/initializers", params={"limit": limit}) - return [RegisteredInitializer.model_validate(item) for item in payload.get("items", [])] + items = await self._get_catalog_items_async(path="/api/initializers", limit=limit) + return [RegisteredInitializer.model_validate(item) for item in items] async def register_initializer_async(self, *, name: str, script_content: str) -> RegisteredInitializer: """ @@ -252,17 +258,20 @@ async def register_initializer_async(self, *, name: str, script_content: str) -> # Targets # ------------------------------------------------------------------ - async def list_targets_async(self, *, limit: int = 200) -> list[TargetInstance]: + async def list_targets_async(self, *, limit: int | None = None) -> list[TargetInstance]: """ List all available targets. + Args: + limit (int | None): Maximum number of targets to return. None returns all targets. + Returns: - list[TargetInstance]: All targets registered on the server. + list[TargetInstance]: Registered targets in catalog order. """ from pyrit.models.catalog import TargetInstance - payload = await self._get_json_async(path="/api/targets", params={"limit": limit}) - return [TargetInstance.model_validate(item) for item in payload.get("items", [])] + items = await self._get_catalog_items_async(path="/api/targets", limit=limit) + return [TargetInstance.model_validate(item) for item in items] # ------------------------------------------------------------------ # Datasets @@ -591,6 +600,43 @@ def _get_client(self) -> Any: ) return self._client + async def _get_catalog_items_async(self, *, path: str, limit: int | None) -> list[dict[str, Any]]: + """ + Read a cursor-paginated catalog up to the optional total limit. + + Args: + path (str): Catalog endpoint. + limit (int | None): Maximum number of items, or None to fetch all pages. + + Returns: + list[dict[str, Any]]: Catalog items in page order. + + Raises: + ValueError: If the limit is nonpositive or a required cursor is missing or repeated. + """ + if limit is not None and limit < 1: + raise ValueError("Catalog limit must be positive.") + + items: list[dict[str, Any]] = [] + seen_cursors: set[str] = set() + cursor: str | None = None + while True: + page_size = 200 if limit is None else min(200, limit - len(items)) + params: dict[str, int | str] = {"limit": page_size} + if cursor is not None: + params["cursor"] = cursor + payload = await self._get_json_async(path=path, params=params) + items.extend(payload.get("items", [])) + if limit is not None and len(items) >= limit: + return items[:limit] + pagination = payload.get("pagination", {}) + if not pagination.get("has_more"): + return items + cursor = pagination.get("next_cursor") + if not isinstance(cursor, str) or not cursor or cursor in seen_cursors: + raise ValueError(f"Catalog response from {path} has a missing or repeated next-page cursor.") + seen_cursors.add(cursor) + async def _get_json_async(self, *, path: str, params: dict[str, Any] | None = None) -> dict[str, Any]: """ GET a JSON endpoint and return the parsed response. diff --git a/tests/unit/cli/test_api_client.py b/tests/unit/cli/test_api_client.py index 3c8cb00c6b..55a3d46276 100644 --- a/tests/unit/cli/test_api_client.py +++ b/tests/unit/cli/test_api_client.py @@ -9,6 +9,7 @@ import threading from datetime import UTC, datetime from pathlib import Path +from typing import Any from unittest.mock import ANY, AsyncMock, MagicMock, call, patch import httpx @@ -737,6 +738,111 @@ async def test_list_targets_async(client, mock_httpx_client): mock_httpx_client.get.assert_awaited_once_with("/api/targets", params={"limit": 7}) +@pytest.mark.parametrize( + ("method", "path", "name_field", "payload"), + [ + ("list_targets_async", "/api/targets", "target_registry_name", _target_payload()), + ("list_scenarios_async", "/api/scenarios/catalog", "scenario_name", _scenario_payload()), + ("list_initializers_async", "/api/initializers", "initializer_name", _initializer_payload()), + ], +) +async def test_catalog_list_follows_next_cursor_async( + *, + client: PyRITApiClient, + mock_httpx_client: MagicMock, + method: str, + path: str, + name_field: str, + payload: dict[str, Any], +) -> None: + mock_httpx_client.get.side_effect = [ + _make_response( + json_data={ + "items": [payload | {name_field: "first"}], + "pagination": {"has_more": True, "next_cursor": "first"}, + } + ), + _make_response( + json_data={ + "items": [payload | {name_field: "second"}], + "pagination": {"has_more": False, "next_cursor": None}, + } + ), + ] + + result = await getattr(client, method)() + + assert [getattr(item, name_field) for item in result] == ["first", "second"] + assert mock_httpx_client.get.await_args_list == [ + call(path, params={"limit": 200}), + call(path, params={"limit": 200, "cursor": "first"}), + ] + + +@pytest.mark.parametrize( + ("limit", "second_request_limit", "expected_count"), + [(None, 200, 205), (201, 1, 201)], + ids=["all-205-targets", "limit-201"], +) +async def test_list_targets_crosses_page_size_boundary_async( + *, + client: PyRITApiClient, + mock_httpx_client: MagicMock, + limit: int | None, + second_request_limit: int, + expected_count: int, +) -> None: + items = [_target_payload(target_registry_name=f"target-{i}") for i in range(205)] + mock_httpx_client.get.side_effect = [ + _make_response( + json_data={ + "items": items[:200], + "pagination": {"has_more": True, "next_cursor": "first"}, + } + ), + _make_response( + json_data={ + "items": items[200:expected_count], + "pagination": { + "has_more": expected_count < len(items), + "next_cursor": "second" if expected_count < len(items) else None, + }, + } + ), + ] + + result = await client.list_targets_async(limit=limit) + + assert [item.target_registry_name for item in result] == [f"target-{i}" for i in range(expected_count)] + assert mock_httpx_client.get.await_args_list == [ + call("/api/targets", params={"limit": 200}), + call("/api/targets", params={"limit": second_request_limit, "cursor": "first"}), + ] + + +@pytest.mark.parametrize("next_cursor", [None, "first"]) +async def test_list_targets_rejects_invalid_next_cursor_async( + *, client: PyRITApiClient, mock_httpx_client: MagicMock, next_cursor: str | None +) -> None: + mock_httpx_client.get.side_effect = [ + _make_response( + json_data={ + "items": [_target_payload(target_registry_name="first")], + "pagination": {"has_more": True, "next_cursor": "first"}, + } + ), + _make_response( + json_data={ + "items": [_target_payload(target_registry_name="second")], + "pagination": {"has_more": True, "next_cursor": next_cursor}, + } + ), + ] + + with pytest.raises(ValueError, match="cursor"): + await client.list_targets_async() + + async def test_list_converters_async(client, mock_httpx_client): mock_httpx_client.get.return_value = _make_response(json_data={"items": []}) await client.list_converters_async() diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index 6cbdc48f35..909dc94146 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -14,11 +14,12 @@ from uuid import uuid4 import pytest +import respx from pyrit.cli import _config_reader as pyrit_scan_config_reader from pyrit.cli import pyrit_scan from pyrit.cli._cli_args import parse_run_arguments -from pyrit.models import Operation, Parameter +from pyrit.models import Operation, Parameter, TargetCapabilities from unit.mocks import make_scenario_result # A valid UTF-8 initializer whose text is not pure ASCII. Decoded with a Windows ANSI code page @@ -563,6 +564,33 @@ def test_main_list_targets(self, mock_client_class, mock_probe): assert result == 0 mock_client.list_targets_async.assert_awaited_once() + @respx.mock + def test_main_list_targets_prints_total_across_pages(self, capsys: pytest.CaptureFixture[str]) -> None: + base_url = "http://testserver:8000" + compatibility_id = "0.14.0+g" + "a" * 40 + capabilities = TargetCapabilities().model_dump(mode="json") + items = [ + { + "target_registry_name": f"target-{i}", + "identifier": {"class_name": "OpenAIChatTarget", "class_module": "pyrit.prompt_target"}, + "capabilities": capabilities, + } + for i in range(205) + ] + respx.get(f"{base_url}/api/health").respond(200, json={"status": "healthy", "service": "pyrit-backend"}) + respx.get(f"{base_url}/api/version").respond(200, json={"compatibility_id": compatibility_id}) + respx.get(f"{base_url}/api/targets", params__eq={"limit": "200"}).respond( + 200, json={"items": items[:200], "pagination": {"has_more": True, "next_cursor": "next-page"}} + ) + respx.get(f"{base_url}/api/targets", params__eq={"limit": "200", "cursor": "next-page"}).respond( + 200, json={"items": items[200:], "pagination": {"has_more": False, "next_cursor": None}} + ) + + with patch("pyrit._compatibility.get_compatibility_id", return_value=compatibility_id): + assert pyrit_scan.main(["list-targets", "--server-url", base_url, "--auth-mode", "none"]) == 0 + + assert "Total targets: 205" in capsys.readouterr().out.splitlines() + @patch("pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new_callable=AsyncMock, return_value=True) @patch("pyrit.cli.api_client.PyRITApiClient") def test_main_list_converters(self, mock_client_class, mock_probe):