From 4af940ac1dfa0136fdc79c54ea6fcdaa429d1fe4 Mon Sep 17 00:00:00 2001 From: jsong468 Date: Thu, 27 Aug 2026 11:39:03 -0700 Subject: [PATCH] phase 4 draft --- pyrit/cli/_cli_args.py | 46 +++++++ pyrit/cli/_output.py | 55 ++++++++- pyrit/cli/_results.py | 189 +++++++++++++++++++++++++++-- pyrit/cli/pyrit_scan.py | 34 ++---- pyrit/cli/pyrit_shell.py | 46 +++---- tests/unit/cli/test_output.py | 132 ++++++++++++++++++++ tests/unit/cli/test_pyrit_scan.py | 29 +++++ tests/unit/cli/test_pyrit_shell.py | 10 ++ tests/unit/cli/test_results.py | 186 +++++++++++++++++++++++++++- 9 files changed, 662 insertions(+), 65 deletions(-) diff --git a/pyrit/cli/_cli_args.py b/pyrit/cli/_cli_args.py index 8c941cb479..1fcae6697a 100644 --- a/pyrit/cli/_cli_args.py +++ b/pyrit/cli/_cli_args.py @@ -91,6 +91,44 @@ def parse_scenario_result_view(raw: str) -> ScenarioResultView: raise argparse.ArgumentTypeError(f"invalid view '{raw}' (choose from {valid})") from None +class OutputFormat(str, Enum): + """ + Serialization format for a ``scenario-results`` render. + + Like ``ScenarioResultView``, this lives in this parse-time-safe module so the + argument parsers can reference it without importing ``pydantic``. ``html`` is + intentionally omitted until a later phase adds its renderer. + """ + + #: Human-readable console text (the default). + CONSOLE = "console" + #: Machine-readable JSON via the payload's ``model_dump_json``. + JSON = "json" + + +def parse_output_format(raw: str) -> OutputFormat: + """ + Parse an ``--output`` token into an ``OutputFormat``. + + Used as an argparse ``type=`` so an invalid value produces an error that + lists the valid format names (mirroring ``parse_scenario_result_view``). + + Args: + raw (str): The raw ``--output`` token. + + Returns: + OutputFormat: The matching format. + + Raises: + argparse.ArgumentTypeError: If *raw* is not a valid format name. + """ + try: + return OutputFormat(raw) + except ValueError: + valid = ", ".join(fmt.value for fmt in OutputFormat) + raise argparse.ArgumentTypeError(f"invalid output '{raw}' (choose from {valid})") from None + + # --------------------------------------------------------------------------- # Pure validators # --------------------------------------------------------------------------- @@ -465,6 +503,14 @@ def add_results_arguments(*, parser: argparse.ArgumentParser) -> None: help="Show at most N attacks (ignored for --view overview; defaults to 5 for " "--view conversations/full when no --attack-result-ids is given)", ) + group.add_argument( + "--output", + type=parse_output_format, + default=OutputFormat.CONSOLE, + metavar="{" + ",".join(fmt.value for fmt in OutputFormat) + "}", + help="Output format: 'console' (human-readable, default) or 'json' " + "(machine-readable; informational notes go to stderr so stdout stays parseable)", + ) def build_scenario_results_parser() -> argparse.ArgumentParser: diff --git a/pyrit/cli/_output.py b/pyrit/cli/_output.py index 5b04d2611f..e368f6b2ae 100644 --- a/pyrit/cli/_output.py +++ b/pyrit/cli/_output.py @@ -16,7 +16,12 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from pyrit.cli._results import AttacksTablePayload, ConversationsPayload, TranscriptMessage + from pyrit.cli._results import ( + AttacksTablePayload, + ConversationsPayload, + ResultsPayload, + TranscriptMessage, + ) from pyrit.models import ScenarioResult from pyrit.models.catalog import ( RegisteredInitializer, @@ -457,7 +462,7 @@ def print_attacks_table(*, payload: AttacksTablePayload) -> None: print(f" technique: {row.atomic_attack_name}") print(f" objective: {row.objective}") - shown = len(payload.rows) + shown = payload.shown if shown < payload.total: print(f"\nShowing {shown} of {payload.total} attacks (use --limit to change).") else: @@ -487,7 +492,7 @@ def print_conversations(*, payload: ConversationsPayload) -> None: print(f" objective: {convo.objective}") _print_transcript(messages=convo.messages) - shown = len(payload.conversations) + shown = payload.shown if shown < payload.total: print(f"\nShowing {shown} of {payload.total} attacks (use --limit or --attack-result-ids to change).") else: @@ -514,6 +519,50 @@ def _print_transcript(*, messages: list[TranscriptMessage]) -> None: print(_wrap(text=f"rationale: {message.score.rationale}", indent=" ")) +def print_results_json(*, payload: ResultsPayload) -> None: + """ + Serialize a results payload as JSON to stdout. + + Emits the payload's own ``model_dump_json`` so the JSON always matches its + fields; informational notes are routed to stderr elsewhere so this stdout + document stays parseable. + + Args: + payload (ResultsPayload): The view payload to serialize. + """ + print(payload.model_dump_json(indent=2)) + + +async def print_results_console_async(*, result: ScenarioResult, payload: ResultsPayload) -> None: + """ + Render a results payload to the console, dispatching on its concrete type. + + The ``overview`` payload keeps the rich framework pretty printer (which reads + *result* directly), so its typed payload is used only by the JSON format; the + other views render from *payload*. + + Args: + result (ScenarioResult): The scenario result, used by the overview printer. + payload (ResultsPayload): The payload built for the requested view. + """ + from pyrit.cli._results import ( + AttacksTablePayload, + ConversationsPayload, + FullPayload, + ScenarioOverviewPayload, + ) + + if isinstance(payload, ScenarioOverviewPayload): + await print_scenario_result_async(result=result) + elif isinstance(payload, AttacksTablePayload): + print_attacks_table(payload=payload) + elif isinstance(payload, ConversationsPayload): + print_conversations(payload=payload) + elif isinstance(payload, FullPayload): + print_attacks_table(payload=payload.attacks) + print_conversations(payload=payload.conversations) + + # --------------------------------------------------------------------------- # Scenario run history # --------------------------------------------------------------------------- diff --git a/pyrit/cli/_results.py b/pyrit/cli/_results.py index b4831b6aa8..de3af3733b 100644 --- a/pyrit/cli/_results.py +++ b/pyrit/cli/_results.py @@ -16,9 +16,10 @@ from __future__ import annotations +import sys from typing import TYPE_CHECKING, Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, computed_field from pyrit.cli._cli_args import ScenarioResultView @@ -33,6 +34,36 @@ _DEFAULT_HEAVY_VIEW_LIMIT = 5 +class GroupSummary(BaseModel): + """One display group's totals and success rate in the ``overview`` view.""" + + name: str + total_results: int + success_rate: int + + +class ScenarioOverviewPayload(BaseModel): + """ + The ``overview`` view: scenario-level aggregates and per-group success rates. + + Captures the numbers the console pretty printer computes inline, but as a + typed object so ``--output json`` can emit the same aggregates the console + shows. The console path still renders via the framework pretty printer while + JSON serializes this payload, so the two compute independently; the aggregates + are kept in lockstep by ``test_overview_payload_matches_framework_printer``. + """ + + scenario_result_id: str + scenario_name: str + scenario_version: int + pyrit_version: str + total_techniques: int + total_attack_results: int + overall_success_rate: int + unique_objectives: int + groups: list[GroupSummary] = Field(default_factory=list) + + class AttackRow(BaseModel): """A single attack result rendered as one row of the attacks table.""" @@ -49,14 +80,21 @@ class AttacksTablePayload(BaseModel): The ``attacks`` view: one row per attack result in a scenario run. ``total`` is the number of attacks that matched the selection before - ``--limit`` was applied; ``len(rows)`` is how many are actually included. - Exposing both lets any renderer show a "showing N of M" note. + ``--limit`` was applied; ``shown`` (== ``len(rows)``) is how many are actually + included. Exposing both lets any renderer — or a JSON consumer — detect when + the list was truncated. """ scenario_result_id: str rows: list[AttackRow] = Field(default_factory=list) total: int = 0 + @computed_field # type: ignore[prop-decorator] + @property + def shown(self) -> int: + """Rows actually included, i.e. ``total`` minus any ``--limit`` truncation.""" + return len(self.rows) + class TranscriptScore(BaseModel): """ @@ -98,14 +136,40 @@ class ConversationsPayload(BaseModel): The ``conversations`` view: the main-conversation transcript per attack. ``total`` is the number of attacks that matched the selection before - ``--limit`` was applied; ``len(conversations)`` is how many are actually - included. Exposing both lets any renderer show a "showing N of M" note. + ``--limit`` was applied; ``shown`` (== ``len(conversations)``) is how many are + actually included. Exposing both lets any renderer — or a JSON consumer — + detect when the list was truncated. """ scenario_result_id: str conversations: list[AttackConversation] = Field(default_factory=list) total: int = 0 + @computed_field # type: ignore[prop-decorator] + @property + def shown(self) -> int: + """Transcripts actually included, i.e. ``total`` minus any ``--limit`` truncation.""" + return len(self.conversations) + + +class FullPayload(BaseModel): + """ + The ``full`` view: the ``attacks`` table together with the ``conversations``. + + Bundling both sub-payloads into one model lets ``--output json`` emit a + single well-formed document (rather than two concatenated ones) while the + console renderer prints the table then the transcripts. + """ + + scenario_result_id: str + attacks: AttacksTablePayload + conversations: ConversationsPayload + + +#: Any of the per-view payloads a ``scenario-results`` render can produce. Named +#: so the builder, the JSON renderer, and both front-ends share one type. +ResultsPayload = ScenarioOverviewPayload | AttacksTablePayload | ConversationsPayload | FullPayload + def resolve_view(*, view: ScenarioResultView | None) -> ScenarioResultView: """ @@ -155,20 +219,61 @@ def apply_view_limit_policy( """ if view is ScenarioResultView.OVERVIEW: if limit is not None: - print("Note: --limit has no effect with --view overview; ignoring it.") + print("Note: --limit has no effect with --view overview; ignoring it.", file=sys.stderr) return None if view in (ScenarioResultView.CONVERSATIONS, ScenarioResultView.FULL): if limit is None and not attack_result_ids: print( f"Note: no --attack-result-ids or --limit given; showing at most " f"{_DEFAULT_HEAVY_VIEW_LIMIT} conversations. Pass --limit or " - "--attack-result-ids to see more." + "--attack-result-ids to see more.", + file=sys.stderr, ) return _DEFAULT_HEAVY_VIEW_LIMIT return limit return limit +def build_overview_payload(*, result: ScenarioResult, scenario_result_id: str) -> ScenarioOverviewPayload: + """ + Build the ``overview`` payload from an already-fetched scenario result. + + Reproduces the console pretty printer's "Overall Statistics" and "Per-Group + Breakdown" by calling the same ``ScenarioResult`` aggregate methods. The + console printer computes the same numbers independently, so + ``test_overview_payload_matches_framework_printer`` locks the two in sync and + fails if the printer's math changes. + + Args: + result (ScenarioResult): The scenario result to aggregate. + scenario_result_id (str): The run id, echoed back on the payload. + + Returns: + ScenarioOverviewPayload: The scenario-level aggregates and per-group rates. + """ + from pyrit.models import AttackOutcome + + groups: list[GroupSummary] = [] + for group_name, group_results in result.get_display_groups().items(): + total_group = len(group_results) + successful = sum(1 for attack_result in group_results if attack_result.outcome == AttackOutcome.SUCCESS) + success_rate = int((successful / total_group) * 100) if total_group else 0 + groups.append(GroupSummary(name=group_name, total_results=total_group, success_rate=success_rate)) + + total_attack_results = sum(len(results) for results in result.attack_results.values()) + return ScenarioOverviewPayload( + scenario_result_id=scenario_result_id, + scenario_name=result.scenario_name, + scenario_version=result.scenario_version, + pyrit_version=result.pyrit_version, + total_techniques=len(result.get_techniques_used()), + total_attack_results=total_attack_results, + overall_success_rate=result.objective_achieved_rate(), + unique_objectives=len(result.get_objectives()), + groups=groups, + ) + + def build_attacks_table_payload( *, result: ScenarioResult, @@ -279,6 +384,76 @@ async def build_conversations_payload_async( ) +async def build_results_payload_async( + *, + view: ScenarioResultView, + result: ScenarioResult, + client: PyRITApiClient, + scenario_result_id: str, + attack_result_ids: list[str] | None = None, + limit: int | None = None, +) -> ResultsPayload: + """ + Build the payload for *view*, dispatching to the per-view builder. + + A single entry point so both front-ends ask for "the payload for this view" + without special-casing which one: the ``overview`` / ``attacks`` builders are + pure and local, while ``conversations`` / ``full`` fetch each attack's + transcript, and this hides that async/sync split behind one ``await``. + + Args: + view (ScenarioResultView): The resolved view to build. + result (ScenarioResult): The already-fetched scenario result. + client (PyRITApiClient): Client used by the transcript-fetching views. + scenario_result_id (str): The run id, echoed back on the payload. + attack_result_ids (list[str] | None): Restrict to these attack ids. + Defaults to None (all attacks). + limit (int | None): Effective attack cap from ``apply_view_limit_policy``. + Defaults to None. + + Returns: + ResultsPayload: The payload matching *view*. + """ + if view is ScenarioResultView.OVERVIEW: + return build_overview_payload(result=result, scenario_result_id=scenario_result_id) + + if view is ScenarioResultView.ATTACKS: + return build_attacks_table_payload( + result=result, + scenario_result_id=scenario_result_id, + attack_result_ids=attack_result_ids, + limit=limit, + ) + + if view is ScenarioResultView.CONVERSATIONS: + return await build_conversations_payload_async( + result=result, + client=client, + scenario_result_id=scenario_result_id, + attack_result_ids=attack_result_ids, + limit=limit, + ) + + attacks = build_attacks_table_payload( + result=result, + scenario_result_id=scenario_result_id, + attack_result_ids=attack_result_ids, + limit=limit, + ) + conversations = await build_conversations_payload_async( + result=result, + client=client, + scenario_result_id=scenario_result_id, + attack_result_ids=attack_result_ids, + limit=limit, + ) + return FullPayload( + scenario_result_id=scenario_result_id, + attacks=attacks, + conversations=conversations, + ) + + def _select_attacks(*, result: ScenarioResult, attack_result_ids: list[str] | None) -> list[tuple[str, AttackResult]]: """ Return ``(atomic_attack_name, attack_result)`` pairs, optionally id-filtered. diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index e14e0e9bce..3c0bafa21c 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -791,11 +791,10 @@ async def _handle_results_async(*, client: Any, parsed_args: Namespace) -> int: int: Exit code (``0`` on success, ``1`` on error). """ from pyrit.cli import _output - from pyrit.cli._cli_args import ScenarioResultView + from pyrit.cli._cli_args import OutputFormat from pyrit.cli._results import ( apply_view_limit_policy, - build_attacks_table_payload, - build_conversations_payload_async, + build_results_payload_async, resolve_view, ) @@ -805,27 +804,8 @@ async def _handle_results_async(*, client: Any, parsed_args: Namespace) -> int: try: result = await client.get_scenario_run_results_async(scenario_result_id=scenario_result_id) - except Exception as exc: - _print_cli_exception(exc=exc) - return 1 - - if view is ScenarioResultView.OVERVIEW: - await _output.print_scenario_result_async(result=result) - return 0 - - if view in (ScenarioResultView.ATTACKS, ScenarioResultView.FULL): - attacks_payload = build_attacks_table_payload( - result=result, - scenario_result_id=scenario_result_id, - attack_result_ids=parsed_args.attack_result_ids, - limit=limit, - ) - _output.print_attacks_table(payload=attacks_payload) - if view is ScenarioResultView.ATTACKS: - return 0 - - try: - conversations_payload = await build_conversations_payload_async( + payload = await build_results_payload_async( + view=view, result=result, client=client, scenario_result_id=scenario_result_id, @@ -835,7 +815,11 @@ async def _handle_results_async(*, client: Any, parsed_args: Namespace) -> int: except Exception as exc: _print_cli_exception(exc=exc) return 1 - _output.print_conversations(payload=conversations_payload) + + if parsed_args.output is OutputFormat.JSON: + _output.print_results_json(payload=payload) + else: + await _output.print_results_console_async(result=result, payload=payload) return 0 diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py index be4d634cef..2239bbe2c7 100644 --- a/pyrit/cli/pyrit_shell.py +++ b/pyrit/cli/pyrit_shell.py @@ -549,6 +549,7 @@ def do_scenario_results(self, arg: str) -> None: scenario-results [--view overview|attacks|conversations|full] [--attack-result-ids ...] [--limit N] + [--output console|json] Views: overview Scenario-level aggregate: totals and per-group success @@ -559,6 +560,10 @@ def do_scenario_results(self, arg: str) -> None: (messages plus their scores and full rationale). full The attacks table followed by the transcripts. + Output: + --output console (default) renders human-readable text; --output json + emits the same data as JSON on stdout (informational notes go to stderr). + For conversations/full, when neither --attack-result-ids nor --limit is given, at most 5 attacks are shown to avoid dumping a whole run. """ @@ -567,12 +572,11 @@ def do_scenario_results(self, arg: str) -> None: import shlex - from pyrit.cli._cli_args import ScenarioResultView, build_scenario_results_parser - from pyrit.cli._output import print_attacks_table, print_conversations, print_scenario_result_async + from pyrit.cli._cli_args import OutputFormat, build_scenario_results_parser + from pyrit.cli._output import print_results_console_async, print_results_json from pyrit.cli._results import ( apply_view_limit_policy, - build_attacks_table_payload, - build_conversations_payload_async, + build_results_payload_async, resolve_view, ) @@ -584,7 +588,8 @@ def do_scenario_results(self, arg: str) -> None: if not tokens: print( "Usage: scenario-results " - "[--view overview|attacks|conversations|full] [--attack-result-ids ...] [--limit N]" + "[--view overview|attacks|conversations|full] [--attack-result-ids ...] " + "[--limit N] [--output console|json]" ) print("Use 'scenario-history' to see available run IDs.") return @@ -602,28 +607,9 @@ def do_scenario_results(self, arg: str) -> None: result = self._run_async( self._api_client.get_scenario_run_results_async(scenario_result_id=parsed.scenario_result_id) ) - except Exception as exc: - print(f"Error: {exc}") - return - - if view is ScenarioResultView.OVERVIEW: - self._run_async(print_scenario_result_async(result=result)) - return - - if view in (ScenarioResultView.ATTACKS, ScenarioResultView.FULL): - attacks_payload = build_attacks_table_payload( - result=result, - scenario_result_id=parsed.scenario_result_id, - attack_result_ids=parsed.attack_result_ids, - limit=limit, - ) - print_attacks_table(payload=attacks_payload) - if view is ScenarioResultView.ATTACKS: - return - - try: - conversations_payload = self._run_async( - build_conversations_payload_async( + payload = self._run_async( + build_results_payload_async( + view=view, result=result, client=self._api_client, scenario_result_id=parsed.scenario_result_id, @@ -634,7 +620,11 @@ def do_scenario_results(self, arg: str) -> None: except Exception as exc: print(f"Error: {exc}") return - print_conversations(payload=conversations_payload) + + if parsed.output is OutputFormat.JSON: + print_results_json(payload=payload) + else: + self._run_async(print_results_console_async(result=result, payload=payload)) def do_print_scenario(self, arg: str) -> None: """ diff --git a/tests/unit/cli/test_output.py b/tests/unit/cli/test_output.py index 8880c95466..2a88788061 100644 --- a/tests/unit/cli/test_output.py +++ b/tests/unit/cli/test_output.py @@ -823,6 +823,138 @@ def test_print_conversations_shows_truncation_note(capsys): assert "(no messages)" in out +# --------------------------------------------------------------------------- +# print_results_json +# --------------------------------------------------------------------------- + + +def test_print_results_json_emits_parseable_payload(capsys): + import json + + payload = _attacks_payload( + rows=[ + { + "attack_result_id": "aid-1", + "atomic_attack_name": "tech_a", + "objective": "extract secrets", + "outcome": "success", + "executed_turns": 3, + "score_value": "0.9", + } + ], + total=1, + ) + _output.print_results_json(payload=payload) + out = capsys.readouterr().out + # stdout is a single parseable document matching the payload's own serialization. + assert json.loads(out) == json.loads(payload.model_dump_json()) + assert json.loads(out)["rows"][0]["attack_result_id"] == "aid-1" + # The computed `shown` sibling is serialized so JSON consumers see truncation. + assert json.loads(out)["shown"] == 1 + + +# --------------------------------------------------------------------------- +# print_results_console_async (payload -> renderer dispatch) +# --------------------------------------------------------------------------- + + +async def test_render_console_overview_uses_framework_printer(): + from pyrit.cli._results import ScenarioOverviewPayload + + payload = ScenarioOverviewPayload( + scenario_result_id="SID", + scenario_name="TestScenario", + scenario_version=1, + pyrit_version="1.0.0", + total_techniques=1, + total_attack_results=1, + overall_success_rate=100, + unique_objectives=1, + ) + result = make_scenario_result(scenario_name="TestScenario", attack_results={}) + with patch("pyrit.cli._output.print_scenario_result_async", new_callable=AsyncMock) as mock_print: + await _output.print_results_console_async(result=result, payload=payload) + # overview keeps the rich framework printer (reads result, ignores payload). + mock_print.assert_awaited_once() + + +async def test_render_console_attacks_prints_table(capsys): + result = make_scenario_result(scenario_name="TestScenario", attack_results={}) + payload = _attacks_payload( + rows=[ + { + "attack_result_id": "aid-1", + "atomic_attack_name": "tech_a", + "objective": "extract secrets", + "outcome": "success", + "executed_turns": 1, + "score_value": None, + } + ], + total=1, + ) + await _output.print_results_console_async(result=result, payload=payload) + assert "extract secrets" in capsys.readouterr().out + + +async def test_render_console_conversations_prints_transcripts(capsys): + result = make_scenario_result(scenario_name="TestScenario", attack_results={}) + payload = _conversations_payload( + conversations=[ + { + "attack_result_id": "aid-1", + "atomic_attack_name": "tech_a", + "objective": "extract secrets", + "outcome": "success", + "conversation_id": "conv-1", + "messages": [{"role": "user", "turn": 0, "text": "hello there", "score": None}], + } + ], + total=1, + ) + await _output.print_results_console_async(result=result, payload=payload) + out = capsys.readouterr().out + assert "Conversations" in out + assert "hello there" in out + + +async def test_render_console_full_prints_table_then_transcripts(capsys): + from pyrit.cli._results import FullPayload + + result = make_scenario_result(scenario_name="TestScenario", attack_results={}) + attacks = _attacks_payload( + rows=[ + { + "attack_result_id": "aid-1", + "atomic_attack_name": "tech_a", + "objective": "extract secrets", + "outcome": "success", + "executed_turns": 1, + "score_value": None, + } + ], + total=1, + ) + conversations = _conversations_payload( + conversations=[ + { + "attack_result_id": "aid-1", + "atomic_attack_name": "tech_a", + "objective": "extract secrets", + "outcome": "success", + "conversation_id": "conv-1", + "messages": [], + } + ], + total=1, + ) + payload = FullPayload(scenario_result_id="SID", attacks=attacks, conversations=conversations) + await _output.print_results_console_async(result=result, payload=payload) + out = capsys.readouterr().out + assert "Attack Results" in out + assert "Conversations" in out + + # --------------------------------------------------------------------------- # print_scenario_runs_list # --------------------------------------------------------------------------- diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index 911d629cd4..1dc2003317 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -1486,6 +1486,35 @@ def test_handle_results_reports_fetch_error(self, capsys): assert rc == 1 assert "boom" in capsys.readouterr().out + def test_handle_results_attacks_json_emits_parseable_document(self, capsys): + import asyncio + import json + + client = AsyncMock() + client.get_scenario_run_results_async.return_value = _make_scenario_result() + parsed = pyrit_scan.parse_args(["scenario-results", "SID", "--view", "attacks", "--output", "json"]) + rc = asyncio.run(pyrit_scan._handle_results_async(client=client, parsed_args=parsed)) + assert rc == 0 + document = json.loads(capsys.readouterr().out) + assert document["scenario_result_id"] == "SID" + assert document["rows"][0]["objective"] == "extract data" + + def test_handle_results_overview_json_emits_aggregates(self, capsys): + import asyncio + import json + + client = AsyncMock() + client.get_scenario_run_results_async.return_value = _make_scenario_result() + parsed = pyrit_scan.parse_args(["scenario-results", "SID", "--output", "json"]) + # overview JSON goes through the payload, not the framework console printer. + with patch("pyrit.cli._output.print_scenario_result_async", new_callable=AsyncMock) as mock_print: + rc = asyncio.run(pyrit_scan._handle_results_async(client=client, parsed_args=parsed)) + assert rc == 0 + mock_print.assert_not_awaited() + document = json.loads(capsys.readouterr().out) + assert document["scenario_result_id"] == "SID" + assert "overall_success_rate" in document + class TestScenarioHistory: """Tests for the ``scenario-history`` verb and its handler.""" diff --git a/tests/unit/cli/test_pyrit_shell.py b/tests/unit/cli/test_pyrit_shell.py index 4168df61d6..f08a073025 100644 --- a/tests/unit/cli/test_pyrit_shell.py +++ b/tests/unit/cli/test_pyrit_shell.py @@ -950,6 +950,16 @@ def test_print_scenario_alias_warns_and_delegates(self, shell, capsys): assert "deprecated" in capsys.readouterr().out.lower() mock_print.assert_awaited_once() + def test_attacks_view_json_emits_parseable_document(self, shell, capsys): + import json + + s, client = shell + client.get_scenario_run_results_async = AsyncMock(return_value=_attacks_scenario_result()) + s.do_scenario_results("rid-1 --view attacks --output json") + document = json.loads(capsys.readouterr().out) + assert document["scenario_result_id"] == "rid-1" + assert {row["objective"] for row in document["rows"]} == {"obj-alpha", "obj-beta"} + def test_stop_server_close_client_swallows_errors(self, shell): s, client = shell launcher = MagicMock() diff --git a/tests/unit/cli/test_results.py b/tests/unit/cli/test_results.py index fcd5a980ce..67f8cb5c2d 100644 --- a/tests/unit/cli/test_results.py +++ b/tests/unit/cli/test_results.py @@ -11,14 +11,22 @@ import pytest from pyrit.cli._cli_args import ( + OutputFormat, ScenarioResultView, add_results_arguments, build_scenario_results_parser, + parse_output_format, ) from pyrit.cli._results import ( + AttacksTablePayload, + ConversationsPayload, + FullPayload, + ScenarioOverviewPayload, apply_view_limit_policy, build_attacks_table_payload, build_conversations_payload_async, + build_overview_payload, + build_results_payload_async, resolve_view, ) from pyrit.models import AttackOutcome, AttackResult, ComponentIdentifier, Score @@ -92,7 +100,7 @@ def test_resolve_view_passes_through_explicit_value(): def test_limit_policy_drops_and_warns_for_overview(capsys): effective = apply_view_limit_policy(view=ScenarioResultView.OVERVIEW, limit=5) assert effective is None - assert "no effect" in capsys.readouterr().out + assert "no effect" in capsys.readouterr().err def test_limit_policy_keeps_limit_for_attacks(capsys): @@ -161,6 +169,10 @@ def test_builder_limit_caps_rows_but_total_is_pre_limit(): payload = build_attacks_table_payload(result=result, scenario_result_id="SID", limit=2) assert payload.total == 5 assert len(payload.rows) == 2 + # `shown` tracks the rendered rows; `total` stays the pre-limit count. + assert payload.shown == 2 + assert payload.shown < payload.total + assert payload.model_dump()["shown"] == 2 def test_builder_handles_no_attacks(): @@ -243,7 +255,7 @@ def test_resolve_view_passes_through_conversations(): def test_limit_policy_defaults_heavy_view_when_unscoped(capsys): effective = apply_view_limit_policy(view=ScenarioResultView.CONVERSATIONS, limit=None) assert effective == 5 - assert "at most 5" in capsys.readouterr().out + assert "at most 5" in capsys.readouterr().err def test_limit_policy_heavy_view_respects_explicit_limit(capsys): @@ -394,3 +406,173 @@ async def test_build_conversations_payload_limit_gates_fetch(): assert len(payload.conversations) == 2 # --limit caps the number of message fetches, not just the rendered rows. assert len(client.calls) == 2 + + +# --------------------------------------------------------------------------- +# OutputFormat / --output +# --------------------------------------------------------------------------- + + +def test_output_format_values(): + assert OutputFormat.CONSOLE.value == "console" + assert OutputFormat.JSON.value == "json" + + +def test_parse_output_format_valid_and_invalid(): + import argparse + + assert parse_output_format("json") is OutputFormat.JSON + with pytest.raises(argparse.ArgumentTypeError, match="choose from console, json"): + parse_output_format("xml") + + +def test_add_results_arguments_registers_output_defaulting_to_console(): + import argparse + + parser = argparse.ArgumentParser() + add_results_arguments(parser=parser) + assert parser.parse_args([]).output is OutputFormat.CONSOLE + assert parser.parse_args(["--output", "json"]).output is OutputFormat.JSON + + +def test_shell_parser_rejects_unknown_output(capsys): + parser = build_scenario_results_parser() + with pytest.raises(SystemExit): + parser.parse_args(["SID", "--output", "xml"]) + assert "choose from console, json" in capsys.readouterr().err + + +# --------------------------------------------------------------------------- +# build_overview_payload +# --------------------------------------------------------------------------- + + +def test_build_overview_payload_computes_aggregates(): + result = _result( + { + "tech_a": [ + _attack(objective="a1", outcome=AttackOutcome.SUCCESS), + _attack(objective="a2", outcome=AttackOutcome.FAILURE), + ], + "tech_b": [_attack(objective="b1", outcome=AttackOutcome.SUCCESS)], + } + ) + + payload = build_overview_payload(result=result, scenario_result_id="SID") + + assert payload.scenario_result_id == "SID" + assert payload.scenario_name == "TestScenario" + assert payload.total_attack_results == 3 + assert payload.total_techniques == 2 + # 2 of 3 attacks succeeded -> int(2/3 * 100). + assert payload.overall_success_rate == 66 + assert payload.unique_objectives == 3 + + +def test_build_overview_payload_per_group_rates(): + result = _result( + { + "tech_a": [ + _attack(objective="a1", outcome=AttackOutcome.SUCCESS), + _attack(objective="a2", outcome=AttackOutcome.FAILURE), + ], + "tech_b": [_attack(objective="b1", outcome=AttackOutcome.SUCCESS)], + } + ) + + payload = build_overview_payload(result=result, scenario_result_id="SID") + + groups = {group.name: group for group in payload.groups} + assert groups["tech_a"].total_results == 2 + assert groups["tech_a"].success_rate == 50 + assert groups["tech_b"].total_results == 1 + assert groups["tech_b"].success_rate == 100 + + +async def test_overview_payload_matches_framework_printer(): + # The console overview (framework printer) and the JSON overview + # (build_overview_payload) compute their numbers independently. This locks + # the two together: if the printer's aggregate math changes, its rendered + # numbers stop matching the payload and this fails, flagging the drift. + from pyrit.output.scenario_result.pretty import PrettyScenarioResultPrinter + + result = _result( + { + "tech_a": [ + _attack(objective="a1", outcome=AttackOutcome.SUCCESS), + _attack(objective="a2", outcome=AttackOutcome.FAILURE), + ], + "tech_b": [_attack(objective="b1", outcome=AttackOutcome.SUCCESS)], + } + ) + payload = build_overview_payload(result=result, scenario_result_id="SID") + + rendered = await PrettyScenarioResultPrinter(enable_colors=False).render_async(result) + + assert f"Total Techniques: {payload.total_techniques}" in rendered + assert f"Total Attack Results: {payload.total_attack_results}" in rendered + assert f"Overall Success Rate: {payload.overall_success_rate}%" in rendered + assert f"Unique Objectives: {payload.unique_objectives}" in rendered + for group in payload.groups: + assert f"Group: {group.name}" in rendered + assert f"Success Rate: {group.success_rate}%" in rendered + + +# --------------------------------------------------------------------------- +# build_results_payload_async (view -> payload dispatch) +# --------------------------------------------------------------------------- + + +async def test_build_results_payload_overview_returns_overview_view(): + result = _result({"tech_a": [_attack(objective="a1")]}) + client = _FakeMessagesClient() + + payload = await build_results_payload_async( + view=ScenarioResultView.OVERVIEW, result=result, client=client, scenario_result_id="SID" + ) + + assert isinstance(payload, ScenarioOverviewPayload) + assert client.calls == [] + + +async def test_build_results_payload_attacks_returns_table_without_fetch(): + result = _result({"tech_a": [_attack(objective="a1")]}) + client = _FakeMessagesClient() + + payload = await build_results_payload_async( + view=ScenarioResultView.ATTACKS, result=result, client=client, scenario_result_id="SID" + ) + + assert isinstance(payload, AttacksTablePayload) + assert client.calls == [] + + +async def test_build_results_payload_conversations_returns_transcripts(): + attack = _attack(objective="a1") + result = _result({"tech_a": [attack]}) + client = _FakeMessagesClient() + + payload = await build_results_payload_async( + view=ScenarioResultView.CONVERSATIONS, result=result, client=client, scenario_result_id="SID" + ) + + assert isinstance(payload, ConversationsPayload) + assert client.calls == [(attack.attack_result_id, attack.conversation_id)] + + +async def test_build_results_payload_full_composes_both_with_shared_limit(): + attacks = [_attack(objective=f"o{i}") for i in range(3)] + result = _result({"tech_a": attacks}) + client = _FakeMessagesClient() + + payload = await build_results_payload_async( + view=ScenarioResultView.FULL, result=result, client=client, scenario_result_id="SID", limit=2 + ) + + assert isinstance(payload, FullPayload) + # The shared limit gates both the table rows and the fetched transcripts. + assert len(payload.attacks.rows) == 2 + assert len(payload.conversations.conversations) == 2 + assert len(client.calls) == 2 + assert payload.attacks.total == 3 + assert payload.conversations.total == 3