From f36db2825f6f6ad6a758f124bb40907e5c5c3433 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 26 Oct 2026 20:05:57 +0000 Subject: [PATCH 01/38] FIX: Reject path parameters by default when building components from API input External construction now rejects a parameter whose type includes Path unless the calling service passes the name in allowed_path_params after making the value safe. The converter service allows its upload parameters once they are stored, and the target service allows the server-set upload folder and the GitHubCopilotTarget working directory. The scorer-only check is removed, and the "Valid parameters" list leaves out path parameters callers cannot set. --- doc/code/registry/0_registry.md | 7 +++ pyrit/backend/services/converter_service.py | 53 ++++++++++++++----- pyrit/backend/services/scorer_service.py | 38 +------------ pyrit/backend/services/target_service.py | 35 +++++++++++- pyrit/models/parameter.py | 5 ++ pyrit/registry/registry.py | 40 +++++++++++--- pyrit/registry/resolution.py | 52 ++++++++++++++---- tests/unit/backend/test_converter_service.py | 48 +++++++++++++++++ tests/unit/backend/test_scorer_service.py | 12 +++++ tests/unit/backend/test_target_service.py | 41 +++++++++++++- .../unit/backend/test_temporary_components.py | 4 +- .../test_lsb_image_stego_converter.py | 7 ++- tests/unit/models/test_parameter.py | 20 +++++++ .../unit/registry/test_converter_registry.py | 31 +++++++++++ tests/unit/registry/test_resolution.py | 51 ++++++++++++++++++ 15 files changed, 368 insertions(+), 76 deletions(-) diff --git a/doc/code/registry/0_registry.md b/doc/code/registry/0_registry.md index 0be69bf192..fd47a19918 100644 --- a/doc/code/registry/0_registry.md +++ b/doc/code/registry/0_registry.md @@ -161,6 +161,13 @@ orders have the same metadata, `type_name: "Path | str"`, including after a JSON round-trip. Optional forms accept `None` in Python; the display type omits `None`, as it does for other optional parameters. +External construction rejects a parameter whose type includes `Path`, because its +value can name a file on the server. A caller that has made the value safe, as the +backend does by storing an uploaded file first, passes the name in +`allowed_path_params` to `create_instance_from_external_input`, or to +`create_named_instance` or `recreate_instance` with `external_input=True`. +In-process `create_instance` still accepts paths. + The backend exposes scorer types at `GET /api/scorers/types`, lists registered instances at `GET /api/scorers`, retrieves one at `GET /api/scorers/{name}`, and creates one with `POST /api/scorers` (`name`, `type`, and optional `params`). The diff --git a/pyrit/backend/services/converter_service.py b/pyrit/backend/services/converter_service.py index 0086bac1f4..48779da411 100644 --- a/pyrit/backend/services/converter_service.py +++ b/pyrit/backend/services/converter_service.py @@ -50,6 +50,7 @@ from pyrit.common.azure_storage import is_azure_blob_uri from pyrit.memory import data_serializer_factory from pyrit.models import ConverterIdentifier, MessagePiece, PromptDataType +from pyrit.models.parameter import Parameter from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer from pyrit.registry.components import ConverterRegistry from pyrit.registry.registry import Reconstructable @@ -229,6 +230,8 @@ async def create_converter_async( converter_type=request.type, params=request.source.params if request.source else request.params, ) + # Every upload parameter now holds a stored upload or an accepted blob URL. + upload_names = frozenset(self._upload_params(converter_type=request.type)) try: # Uploads may have yielded to another request that took the name. if request.register: @@ -237,11 +240,18 @@ async def create_converter_async( self._registry.instances.validate_name_available(request.name) converter_obj = ( await construct_component_async( - self._registry.recreate_instance, source=source, params=params, external_input=True + self._registry.recreate_instance, + source=source, + params=params, + external_input=True, + allowed_path_params=upload_names, ) if source is not None else await construct_component_async( - self._registry.create_instance_from_external_input, request.type, params=params + self._registry.create_instance_from_external_input, + request.type, + params=params, + allowed_path_params=upload_names, ) ) if not request.register: @@ -311,12 +321,17 @@ async def preview_conversion_async(self, *, request: ConverterPreviewRequest) -> if spec.source_name != converter_id: raise ValueError("A converter specification must match its pipeline source") source = self._registry.resolve_source(name=spec.source_name, identifier_hash=spec.source_hash) + converter_type = type(source).__name__ params, paths = await self._persist_data_uri_params_async( - converter_type=type(source).__name__, params=spec.params + converter_type=converter_type, params=spec.params ) owned_paths.extend(paths) obj = await construct_component_async( - self._registry.recreate_instance, source=source, params=params, external_input=True + self._registry.recreate_instance, + source=source, + params=params, + external_input=True, + allowed_path_params=frozenset(self._upload_params(converter_type=converter_type)), ) if spec.effective_hash and obj.get_identifier().hash != spec.effective_hash: raise ValueError("Temporary converter configuration has changed") @@ -385,6 +400,25 @@ def get_converter_objects_for_ids(self, *, converter_ids: list[str]) -> list[Any # Private Helper Methods # ======================================================================== + def _upload_params(self, *, converter_type: str) -> dict[str, Parameter]: + """ + Return the converter's ``Path`` and ``Path | str`` parameters, which callers send as uploads. + + Args: + converter_type (str): The registered converter class name. + + Returns: + dict[str, Parameter]: The upload parameters by name. + """ + metadata = self._registry.get_registered_class_metadata(converter_type) + if metadata is None: + return {} + return { + parameter.name: parameter + for parameter in metadata.parameters + if parameter.is_path or parameter.is_path_or_str + } + async def _persist_data_uri_params_async( self, *, @@ -424,16 +458,7 @@ async def _persist_data_uri_params_async( ValueError: If a ``Path`` value is not a valid data URI, or an Azure Blob URL is outside the results container's media subdirectories. """ - metadata = self._registry.get_registered_class_metadata(converter_type) - path_params = ( - { - parameter.name: parameter - for parameter in metadata.parameters - if parameter.is_path or parameter.is_path_or_str - } - if metadata - else {} - ) + path_params = self._upload_params(converter_type=converter_type) result = dict(params) owned_paths: list[Path] = [] diff --git a/pyrit/backend/services/scorer_service.py b/pyrit/backend/services/scorer_service.py index d30568909c..889194beb1 100644 --- a/pyrit/backend/services/scorer_service.py +++ b/pyrit/backend/services/scorer_service.py @@ -19,25 +19,11 @@ from pyrit.models.catalog.scorer import ScorerInstance from pyrit.models.component_spec import RegisteredComponentReference from pyrit.models.identifiers.scorer_identifier import ScorerIdentifier -from pyrit.models.parameter import Parameter from pyrit.output.scorer.json import format_scorer_type, scorer_metrics_to_dict from pyrit.registry import ScorerRegistry from pyrit.score import Scorer -def _names_server_file(parameter: Parameter) -> bool: - """ - Return whether a scorer parameter names a file on the server. - - Args: - parameter (Parameter): The scorer constructor parameter. - - Returns: - bool: True for ``Path`` and ``Path | str`` parameters. - """ - return parameter.is_path or parameter.is_path_or_str - - class ScorerService: """Expose ScorerRegistry metadata and instances without duplicating construction logic.""" @@ -100,25 +86,6 @@ def find_registered_scorer(self, identifier_hash: str) -> Scorer | None: None, ) - def _reject_server_file_params(self, *, type_name: str, params: dict[str, Any]) -> None: - """ - Reject parameters that name a file on the server, before any scorer is built. - - Args: - type_name (str): The registered scorer type. - params (dict[str, Any]): The constructor parameters from the API caller. - - Raises: - ValueError: If a parameter names a file on the server. - """ - metadata = self._registry.get_registered_class_metadata(type_name) - for parameter in metadata.parameters if metadata else []: - if parameter.name in params and _names_server_file(parameter): - raise ValueError( - f"Parameter '{parameter.name}' of '{type_name}' names a server file and cannot be set " - "through the API; pass it from Python or an initializer instead." - ) - async def list_scorer_types_async(self) -> ScorerTypeResponse: """ List the scorer types external callers can build, without constructing scorers. @@ -140,14 +107,14 @@ def list_types() -> ScorerTypeResponse: parameters=[ parameter.for_external_catalog() for parameter in metadata.parameters - if parameter.is_external_input and not _names_server_file(parameter) + if parameter.is_external_input and not parameter.mentions_path ], is_llm_based=metadata.is_llm_based, description=metadata.class_description or None, ) for metadata in self._registry.get_all_registered_class_metadata() if all( - parameter.is_external_input and not _names_server_file(parameter) + parameter.is_external_input and not parameter.mentions_path for parameter in metadata.parameters if parameter.required ) @@ -214,7 +181,6 @@ async def create_scorer_async(self, *, request: CreateScorerRequest) -> ScorerIn def create() -> ScorerInstance | UnregisteredScorer: if request.type not in self._registry: raise ValueError(f"Scorer type '{request.type}' not found") - self._reject_server_file_params(type_name=request.type, params=request.params) if not request.register: scorer = self._registry.create_instance_from_external_input(request.type, params=request.params) return UnregisteredScorer( diff --git a/pyrit/backend/services/target_service.py b/pyrit/backend/services/target_service.py index 5fdefe022a..2cdce7f38b 100644 --- a/pyrit/backend/services/target_service.py +++ b/pyrit/backend/services/target_service.py @@ -48,6 +48,11 @@ "PromptShieldTarget": frozenset({"endpoint"}), } +# Path parameters API callers may set. Any other path parameter is rejected when a target is built from API input. +_CALLER_PATH_PARAMETERS: dict[str, frozenset[str]] = { + "GitHubCopilotTarget": frozenset({"working_directory"}), +} + # Constructor parameter through which a target accepts an explicit auth mode. _AUTH_MODE_PARAM = "auth_mode" @@ -167,6 +172,27 @@ def _get_supported_auth_modes(auth_modes: tuple[str, ...]) -> list[AuthMode]: supported_auth_modes.append(auth_mode) return supported_auth_modes + @staticmethod + def _allowed_path_params(*, target_cls: type[PromptTarget]) -> frozenset[str]: + """ + Return the path parameters a target built from API input may receive. + + These are the path parameters callers may set, plus the upload folder of a + target that uploads local files, which only the server sets. + + Args: + target_cls (type[PromptTarget]): Target class being described or built. + + Returns: + frozenset[str]: Path parameter names the target may receive. + """ + names: set[str] = { + name for base in target_cls.__mro__ for name in _CALLER_PATH_PARAMETERS.get(base.__name__, frozenset()) + } + if target_cls.upload_directory_parameter is not None: + names.add(target_cls.upload_directory_parameter) + return frozenset(names) + def _project_target_parameters(self, *, target_type: str, parameters: tuple[Parameter, ...]) -> list[Parameter]: """ Project registry parameters into the API contract. @@ -175,7 +201,8 @@ def _project_target_parameters(self, *, target_type: str, parameters: tuple[Para can resolve them from dotenv configuration. The GUI must still collect them explicitly, so the API marks those values required without changing the target constructor signatures. The upload folder of a target that uploads - local files is set by the server, so it is left out. + local files is set by the server, so it is left out, and so are other path + parameters callers may not set. Args: target_type (str): Registered target class name. @@ -188,10 +215,12 @@ def _project_target_parameters(self, *, target_type: str, parameters: tuple[Para required_names = frozenset().union( *(_ENV_BACKED_REQUIRED_PARAMETERS.get(base.__name__, frozenset()) for base in target_cls.__mro__) ) + path_names = self._allowed_path_params(target_cls=target_cls) return [ parameter.model_copy(update={"default": REQUIRED_VALUE}) if parameter.name in required_names else parameter for parameter in parameters if parameter.name != target_cls.upload_directory_parameter + and (not parameter.mentions_path or parameter.name in path_names) ] async def list_target_types_async(self, *, upload_directory: str | None = None) -> TargetTypeResponse: @@ -434,7 +463,9 @@ async def create_target_async( target_obj = ( await self.build_from_source_async(request.source) if request.source - else self._registry.create_instance_from_external_input(request.type, params=params) + else self._registry.create_instance_from_external_input( + request.type, params=params, allowed_path_params=self._allowed_path_params(target_cls=target_cls) + ) ) if not request.register: try: diff --git a/pyrit/models/parameter.py b/pyrit/models/parameter.py index bfb71c634a..c81f076c12 100644 --- a/pyrit/models/parameter.py +++ b/pyrit/models/parameter.py @@ -246,6 +246,11 @@ def is_path_or_str(self) -> bool: """Whether this parameter accepts paths or strings without URL normalization.""" return self.reference is None and _is_path_or_str(self.param_type) + @property + def mentions_path(self) -> bool: + """Whether a value of this parameter can name a file or folder on the server.""" + return self.reference is None and _mentions_path(self.param_type) + @computed_field @property def reference_type(self) -> str | None: diff --git a/pyrit/registry/registry.py b/pyrit/registry/registry.py index 40d7c6a593..f78fc378cd 100644 --- a/pyrit/registry/registry.py +++ b/pyrit/registry/registry.py @@ -40,7 +40,7 @@ ) if TYPE_CHECKING: - from collections.abc import Iterator, Mapping + from collections.abc import Collection, Iterator, Mapping from types import ModuleType from typing import Self @@ -763,18 +763,24 @@ def create_instance(self, name: str, **kwargs: object) -> T: ) return cls(**resolved) - def create_instance_from_external_input(self, name: str, *, params: Mapping[str, object]) -> T: + def create_instance_from_external_input( + self, name: str, *, params: Mapping[str, object], allowed_path_params: Collection[str] = frozenset() + ) -> T: """ Build a configured instance from external input such as a REST or CLI request. Unlike ``create_instance``, which serves in-process callers that may pass any Python object, this accepts only parameters with ``Parameter.is_external_input`` - and registry references given by name. Anything else is rejected before - construction; the constructor still validates the values it receives. + and registry references given by name. A parameter that can name a server file + or folder is accepted only when the caller lists it in ``allowed_path_params`` + after making it safe. Anything else is rejected before construction; the + constructor still validates the values it receives. Args: name (str): The catalog name to build. params (Mapping[str, object]): Constructor arguments from the external caller. + allowed_path_params (Collection[str]): Parameters that can name a server file + or folder but that the caller has made safe. Defaults to none. Returns: T: The constructed instance. @@ -782,8 +788,8 @@ def create_instance_from_external_input(self, name: str, *, params: Mapping[str, Raises: KeyError: If the name is not registered. ValueError: If an argument is not a valid constructor parameter or not an - external input, a registry reference cannot be resolved, or a value - cannot be coerced. + external input, names a server file or folder without being allowed, a + registry reference cannot be resolved, or a value cannot be coerced. """ cls = self.get_class(name) resolved = resolve_constructor_args( @@ -792,6 +798,7 @@ def create_instance_from_external_input(self, name: str, *, params: Mapping[str, identifier_type=self._identifier_type(), external_input=True, seam_base_type=self._seam_base_type(), + allowed_path_params=allowed_path_params, ) return cls(**resolved) @@ -870,6 +877,7 @@ def create_named_instance( params: Mapping[str, object] | None = None, registry_metadata: dict[str, Any] | None = None, external_input: bool = False, + allowed_path_params: Collection[str] = frozenset(), ) -> InstanceT: """ Build and store a configured instance under an explicit name. @@ -883,6 +891,8 @@ def create_named_instance( external_input (bool): Whether ``params`` come from an external caller, in which case the instance is built with ``create_instance_from_external_input``. Defaults to False. + allowed_path_params (Collection[str]): External parameters that can name a + server file or folder but that the caller has made safe. Defaults to none. Returns: InstanceT: The constructed and registered instance. @@ -890,7 +900,7 @@ def create_named_instance( self.instances.validate_name_available(name) args = dict(params) if params is not None else {} instance = ( - self.create_instance_from_external_input(type_name, params=args) + self.create_instance_from_external_input(type_name, params=args, allowed_path_params=allowed_path_params) if external_input else self.create_instance(type_name, **args) ) @@ -952,11 +962,24 @@ def get_reconstruction_parameters(self, source: InstanceT) -> dict[str, object]: } def recreate_instance( - self, *, source: InstanceT, params: Mapping[str, object], external_input: bool = False + self, + *, + source: InstanceT, + params: Mapping[str, object], + external_input: bool = False, + allowed_path_params: Collection[str] = frozenset(), ) -> InstanceT: """ Build an independent component; never change or register the source. + Args: + source (InstanceT): The registered component to copy. + params (Mapping[str, object]): Constructor overrides. + external_input (bool): Whether ``params`` come from an external caller. + Defaults to False. + allowed_path_params (Collection[str]): External overrides that can name a + server file or folder but that the caller has made safe. Defaults to none. + Returns: InstanceT: The new component, owned by the caller. @@ -988,6 +1011,7 @@ def recreate_instance( identifier_type=self._identifier_type(), external_input=True, seam_base_type=self._seam_base_type(), + allowed_path_params=allowed_path_params, ) with self._catalog_lock: self._ensure_discovered() diff --git a/pyrit/registry/resolution.py b/pyrit/registry/resolution.py index 8b4b60a5e2..cd2729d8bc 100644 --- a/pyrit/registry/resolution.py +++ b/pyrit/registry/resolution.py @@ -22,7 +22,9 @@ ``Parameter.coerce_value`` and resolving registry-reference parameters by name from the owning domain's registry. Defaults are left to the constructor. Callers building from external input (REST, CLI) select the external path, - which accepts only parameters with ``Parameter.is_external_input``. + which accepts only parameters with ``Parameter.is_external_input`` and rejects + parameters that can name a server file or folder unless the caller has made + them safe. - **Resolve from a declared list** (``resolve_declared_params``): the sibling for a component that declares an explicit ``list[Parameter]`` (e.g. a scenario's ``supported_parameters()``). It has no references, coerces every supplied @@ -621,6 +623,7 @@ def resolve_constructor_args( external_input: bool = False, strict_json: bool = False, seam_base_type: type | None = None, + allowed_path_params: Collection[str] = frozenset(), ) -> dict[str, Any]: """ Resolve a flat argument dict into constructor-ready keyword arguments. @@ -648,6 +651,9 @@ def resolve_constructor_args( contracts; ordinary CLI and Python resolution remain unchanged. seam_base_type (type | None): The domain base type whose non-required declared references are construction seams (see ``derive_parameters``). + allowed_path_params (Collection[str]): External-input parameters that can name a + server file or folder but that the caller has made safe, such as an uploaded + file the backend stored itself. Defaults to none. Returns: dict[str, Any]: Arguments ready to pass to ``cls(**resolved)``. @@ -655,14 +661,20 @@ def resolve_constructor_args( Raises: ValueError: If an argument is not a declared parameter, a registry reference cannot be resolved, a simple value cannot be coerced, or - external input sets a parameter that is not an external input. + external input sets a parameter that is not an external input or that + names a server file or folder without being allowed. """ by_name = { param.name: param for param in derive_parameters(cls=cls, identifier_type=identifier_type, seam_base_type=seam_base_type) } if external_input: - reject_non_external_params(params=raw_args, declared=list(by_name.values()), owner=cls.__name__) + reject_non_external_params( + params=raw_args, + declared=list(by_name.values()), + owner=cls.__name__, + allowed_path_params=allowed_path_params, + ) resolved: dict[str, Any] = {} for name, value in raw_args.items(): @@ -711,34 +723,52 @@ def resolve_constructor_args( return resolved -def reject_non_external_params(*, params: Mapping[str, Any], declared: Sequence[Parameter], owner: str) -> None: +def reject_non_external_params( + *, + params: Mapping[str, Any], + declared: Sequence[Parameter], + owner: str, + allowed_path_params: Collection[str] = frozenset(), +) -> None: """ Reject external input that is not an explicitly supported external input. Every name must be declared and have ``Parameter.is_external_input``, and registry - references must be given by name. Values are otherwise left to the component. + references must be given by name. A parameter that can name a server file or + folder is rejected unless it is in ``allowed_path_params``. Values are otherwise + left to the component. Args: params (Mapping[str, Any]): The parameter values supplied by an external caller. declared (Sequence[Parameter]): The parameters the component declares. owner (str): The owning class or scenario name, for error messages. + allowed_path_params (Collection[str]): Parameters that can name a server file or + folder but that the caller has made safe. Defaults to none. Raises: - ValueError: If ``params`` sets an undeclared parameter or one that is not an external - input, or gives a registry reference as anything but a name. + ValueError: If ``params`` sets an undeclared parameter, one that is not an external + input, or one that names a server file or folder without being allowed, or + gives a registry reference as anything but a name. """ declared_by_name = {parameter.name: parameter for parameter in declared} + settable_names = { + name + for name, parameter in declared_by_name.items() + if parameter.is_external_input and (not parameter.mentions_path or name in allowed_path_params) + } for name, value in params.items(): parameter = declared_by_name.get(name) if parameter is None: - external_names = sorted( - known for known, candidate in declared_by_name.items() if candidate.is_external_input - ) - raise ValueError(f"Unknown parameter '{name}' for '{owner}'. Valid parameters: {external_names}") + raise ValueError(f"Unknown parameter '{name}' for '{owner}'. Valid parameters: {sorted(settable_names)}") if not parameter.is_external_input: raise ValueError( f"Parameter '{name}' of '{owner}' cannot be set through the API; pass it from Python instead." ) + if name not in settable_names: + raise ValueError( + f"Parameter '{name}' of '{owner}' names a server file or folder and cannot be set through the API; " + "pass it from Python or an initializer instead." + ) if parameter.reference is not None: _require_reference_names(value=value, owner=owner, name=name) diff --git a/tests/unit/backend/test_converter_service.py b/tests/unit/backend/test_converter_service.py index 040670abd8..82d3827e07 100644 --- a/tests/unit/backend/test_converter_service.py +++ b/tests/unit/backend/test_converter_service.py @@ -35,6 +35,7 @@ Base64Converter, BinaryConverter, CaesarConverter, + ColloquialWordswapConverter, ImageCompressionConverter, QRCodeConverter, RepeatTokenConverter, @@ -47,6 +48,7 @@ from pyrit.converter.converter import get_converter_modalities from pyrit.memory import AzureBlobStorageIO, CentralMemory, MemoryInterface from pyrit.models import ComponentIdentifier, PromptDataType +from pyrit.models.component_spec import SourceInstanceSpec from pyrit.prompt_normalizer import PromptNormalizer from pyrit.registry.components import ConverterRegistry from unit.mocks import MockPromptTarget @@ -695,6 +697,52 @@ async def test_create_with_path_or_str_url( assert list(upload_service._upload_path.iterdir()) == [] assert await upload_service.delete_converter_async(converter_id=response.converter_id) + @pytest.mark.usefixtures("patch_central_database") + async def test_variants_and_previews_accept_uploaded_path_overrides(self, upload_service: ConverterService) -> None: + source = ColloquialWordswapConverter(deterministic=True) + upload_service._registry.instances.register(source, name="wordswap") + spec = SourceInstanceSpec( + source_name="wordswap", + source_hash=source.get_identifier().hash, + params={"wordswap_path": _make_data_uri(mime_type="application/yaml", content=b"hello:\n - wassup\n")}, + ) + + variant = await upload_service.create_converter_async( + request=CreateConverterRequest(type="ColloquialWordswapConverter", register=False, source=spec) + ) + preview = await upload_service.preview_conversion_async( + request=ConverterPreviewRequest( + original_value="hello there", converter_ids=["wordswap"], converter_specs=[spec] + ) + ) + + assert variant.identifier.params["substitution_keys"] == ["hello"] + assert preview.converted_value == "wassup there" + assert list(upload_service._upload_path.iterdir()) == [] + + @pytest.mark.usefixtures("patch_central_database") + async def test_variants_and_previews_reject_server_path_overrides( + self, upload_service: ConverterService, tmp_path: Path + ) -> None: + source = ColloquialWordswapConverter(deterministic=True) + upload_service._registry.instances.register(source, name="wordswap") + server_file = tmp_path / "server.yaml" + server_file.write_text("hello:\n - leaked\n", encoding="utf-8") + spec = SourceInstanceSpec( + source_name="wordswap", source_hash=source.get_identifier().hash, params={"wordswap_path": str(server_file)} + ) + + with pytest.raises(ValueError, match="'wordswap_path' must be uploaded as a data URI"): + await upload_service.create_converter_async( + request=CreateConverterRequest(type="ColloquialWordswapConverter", register=False, source=spec) + ) + with pytest.raises(ValueError, match="'wordswap_path' must be uploaded as a data URI"): + await upload_service.preview_conversion_async( + request=ConverterPreviewRequest( + original_value="hello", converter_ids=["wordswap"], converter_specs=[spec] + ) + ) + @pytest.mark.parametrize("value", [r"C:\server\input.mp4", "input.mp4", "https://example.org/input.mp4", 123]) async def test_path_or_str_rest_rejects_non_upload_non_blob_values( self, upload_service: ConverterService, value: object diff --git a/tests/unit/backend/test_scorer_service.py b/tests/unit/backend/test_scorer_service.py index 36f9c1c4f0..c949db10b7 100644 --- a/tests/unit/backend/test_scorer_service.py +++ b/tests/unit/backend/test_scorer_service.py @@ -466,3 +466,15 @@ async def test_decision_scorer_builds_through_the_api_with_its_bundled_question( assert created.scorer_registry_name == "decision" assert ScorerRegistry.get_registry_singleton().instances.get("decision") is not None + + +async def test_decision_scorer_unknown_parameter_does_not_list_its_question_path(chat_target: str) -> None: + with pytest.raises(ValueError, match="Unknown parameter 'nope' for 'DecisionScorer'") as error: + await ScorerService().create_scorer_async( + request=CreateScorerRequest( + name="decision", type="DecisionScorer", params={"chat_target": chat_target, "nope": 1} + ) + ) + + assert "'chat_target'" in str(error.value) + assert "decision_question_path" not in str(error.value) diff --git a/tests/unit/backend/test_target_service.py b/tests/unit/backend/test_target_service.py index 5a17eb04f7..67727f2967 100644 --- a/tests/unit/backend/test_target_service.py +++ b/tests/unit/backend/test_target_service.py @@ -529,6 +529,45 @@ async def test_create_github_copilot_target_requires_token_async( assert service.get_target_object(target_registry_name="copilot") is None +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("register", [True, False]) +async def test_create_github_copilot_target_accepts_working_directory_async(tmp_path: Path, register: bool) -> None: + params = {"model_name": "gpt-5-mini", "github_token": "test-github-token", "working_directory": str(tmp_path)} + + with patch.dict("sys.modules", {"copilot": ModuleType("copilot")}): + target = await TargetService().create_target_async( + request=CreateTargetRequest( + name="copilot" if register else None, type="GitHubCopilotTarget", register=register, params=params + ) + ) + + assert target.identifier.params["working_directory"] == str(tmp_path.resolve()) + + +class _FileTarget(MockPromptTarget): + """A target with a parameter that names a file on the server.""" + + def __init__(self, *, data_file: Path | None = None) -> None: + super().__init__() + self.data_file = data_file + + +@pytest.mark.usefixtures("patch_central_database") +async def test_target_path_parameters_are_left_out_and_rejected() -> None: + service = TargetService() + service._registry.register_class(_FileTarget) + + types = await service.list_target_types_async() + with pytest.raises(ValueError, match="'data_file' of '_FileTarget' names a server file"): + await service.create_target_async( + request=CreateTargetRequest(name="files", type="_FileTarget", params={"data_file": "/etc/hosts"}) + ) + + entry = next(item for item in types.items if item.target_type == "_FileTarget") + assert "data_file" not in [parameter.name for parameter in entry.parameters] + assert service.get_target_object(target_registry_name="files") is None + + @pytest.mark.usefixtures("patch_central_database") class TestCreateTargetCodeExecution: """Admin gate for targets whose construction runs model repository code.""" @@ -900,7 +939,7 @@ async def test_create_target_delegates_construction_to_registry(self, sqlite_ins ) as create: await service.create_target_async(request=CreateTargetRequest(type="TextTarget", params={})) - create.assert_called_once_with("TextTarget", params={}) + create.assert_called_once_with("TextTarget", params={}, allowed_path_params=frozenset()) async def test_create_target_rejects_invalid_enum_value(self, sqlite_instance) -> None: """An invalid enum value from the JSON request raises a parameter error before target construction.""" diff --git a/tests/unit/backend/test_temporary_components.py b/tests/unit/backend/test_temporary_components.py index 5b214d8c50..81fb1993fc 100644 --- a/tests/unit/backend/test_temporary_components.py +++ b/tests/unit/backend/test_temporary_components.py @@ -108,8 +108,8 @@ async def test_hugging_face_construction_uses_request_loop_async(self, register: registry = TargetRegistry.get_registry_singleton() factory = registry.create_instance_from_external_input - def create(name: str, *, params: Any) -> HuggingFaceChatTarget: - target = factory(name, params=params) + def create(name: str, *, params: Any, allowed_path_params: Any) -> HuggingFaceChatTarget: + target = factory(name, params=params, allowed_path_params=allowed_path_params) assert isinstance(target, HuggingFaceChatTarget) constructed.append(target) return target diff --git a/tests/unit/converter/test_lsb_image_stego_converter.py b/tests/unit/converter/test_lsb_image_stego_converter.py index 7fec17ad43..f72c793984 100644 --- a/tests/unit/converter/test_lsb_image_stego_converter.py +++ b/tests/unit/converter/test_lsb_image_stego_converter.py @@ -324,9 +324,12 @@ def test_lsb_converter_registry_accepts_carrier_image_path(tmp_path: Path, *, us assert parameter.is_external_input assert parameter.is_path_or_str + params = {"carrier_image_path": carrier_path if use_path else str(carrier_path)} + with pytest.raises(ValueError, match="names a server file"): + registry.create_instance_from_external_input("LSBImageStegoConverter", params=params) + converter = registry.create_instance_from_external_input( - "LSBImageStegoConverter", - params={"carrier_image_path": carrier_path if use_path else str(carrier_path)}, + "LSBImageStegoConverter", params=params, allowed_path_params={"carrier_image_path"} ) assert isinstance(converter, LSBImageStegoConverter) assert converter.get_identifier().params["carrier_image_path"] == str(carrier_path) diff --git a/tests/unit/models/test_parameter.py b/tests/unit/models/test_parameter.py index 55ff16b846..5a39761449 100644 --- a/tests/unit/models/test_parameter.py +++ b/tests/unit/models/test_parameter.py @@ -339,6 +339,26 @@ def test_list_path_or_str_contract_round_trip() -> None: assert candidate.coerce_value(values) == values +@pytest.mark.parametrize( + ("annotation", "expected"), + [ + (Path, True), + (Path | None, True), + (Path | str, True), + (str | Path | None, True), + (Path | int, True), + (list[Path | str], True), + (Sequence[Path], True), + (str, False), + (str | None, False), + (list[str], False), + (int, False), + ], +) +def test_mentions_path(annotation: object, expected: bool) -> None: + assert Parameter(name="source", description="d", param_type=annotation).mentions_path is expected + + class TestIsScalarParamType: """``_is_scalar_param_type`` recognizes plain and constrained scalars.""" diff --git a/tests/unit/registry/test_converter_registry.py b/tests/unit/registry/test_converter_registry.py index 80dbfe2ef5..252907d35b 100644 --- a/tests/unit/registry/test_converter_registry.py +++ b/tests/unit/registry/test_converter_registry.py @@ -228,6 +228,37 @@ def test_create_named_instance_selects_external_input_explicitly(self, registry: ) assert registry.instances.get("caesar") is converter + def test_create_named_instance_rejects_path_parameters_unless_allowed(self, registry: ConverterRegistry): + params = {"base_image": "input.png"} + with pytest.raises(ValueError, match="'base_image' of 'ImageOverlayConverter' names a server file"): + registry.create_named_instance( + name="overlay", type_name="ImageOverlayConverter", params=params, external_input=True + ) + assert registry.instances.get("overlay") is None + + converter = registry.create_named_instance( + name="overlay", + type_name="ImageOverlayConverter", + params=params, + external_input=True, + allowed_path_params={"base_image"}, + ) + assert registry.instances.get("overlay") is converter + + def test_recreate_instance_rejects_path_overrides_unless_allowed(self, registry: ConverterRegistry): + source = registry.create_named_instance( + name="overlay", type_name="ImageOverlayConverter", params={"base_image": "input.png"} + ) + params = {"base_image": "other.png"} + with pytest.raises(ValueError, match="'base_image' of 'ImageOverlayConverter' names a server file"): + registry.recreate_instance(source=source, params=params, external_input=True) + + variant = registry.recreate_instance( + source=source, params=params, external_input=True, allowed_path_params={"base_image"} + ) + assert variant.get_identifier().params["base_image"] == "other.png" + assert source.get_identifier().params["base_image"] == "input.png" + @pytest.mark.parametrize("name", ["preview", "types"]) def test_create_named_instance_rejects_reserved_name(self, registry: ConverterRegistry, name: str): with pytest.raises(ValueError, match="reserved"): diff --git a/tests/unit/registry/test_resolution.py b/tests/unit/registry/test_resolution.py index 52b1c563fb..4ecb8909c2 100644 --- a/tests/unit/registry/test_resolution.py +++ b/tests/unit/registry/test_resolution.py @@ -5,8 +5,10 @@ Tests for the shared registry constructor-argument resolution primitive. """ +import re from collections.abc import Callable from enum import Enum +from pathlib import Path from typing import Any, Literal import pytest @@ -459,6 +461,15 @@ class _Handle: """A Python object no external caller can supply.""" +class _Files: + """Helper whose constructor takes paths that name files on the server.""" + + def __init__(self, *, label: str = "", data_file: Path | None = None, source: Path | str | None = None) -> None: + self.label = label + self.data_file = data_file + self.source = source + + class _Mixed: """Helper whose constructor mixes external inputs with Python-object parameters.""" @@ -498,6 +509,46 @@ def test_external_still_rejects_unknown_parameters(self) -> None: with pytest.raises(ValueError, match="Unknown parameter 'nope'"): resolve_constructor_args(cls=_Mixed, raw_args={"nope": 1}, external_input=True) + @pytest.mark.parametrize( + "raw_args", [{"data_file": "/etc/hosts"}, {"source": "/etc/hosts"}, {"data_file": None}, {"source": ""}] + ) + def test_external_rejects_path_parameters_by_default(self, raw_args: dict[str, object]) -> None: + with pytest.raises( + ValueError, match="of '_Files' names a server file or folder and cannot be set through the API" + ): + resolve_constructor_args(cls=_Files, raw_args=raw_args, external_input=True) + + def test_external_accepts_only_the_allowed_path_parameters(self) -> None: + resolved = resolve_constructor_args( + cls=_Files, + raw_args={"label": "x", "data_file": "/uploads/data.bin"}, + external_input=True, + allowed_path_params={"data_file"}, + ) + + assert resolved == {"label": "x", "data_file": Path("/uploads/data.bin")} + with pytest.raises(ValueError, match="'source' of '_Files' names a server file"): + resolve_constructor_args( + cls=_Files, raw_args={"source": "/etc/hosts"}, external_input=True, allowed_path_params={"data_file"} + ) + + def test_unknown_parameter_lists_only_parameters_callers_may_set(self) -> None: + with pytest.raises(ValueError, match=re.escape("Valid parameters: ['label']")): + resolve_constructor_args(cls=_Files, raw_args={"nope": 1}, external_input=True) + with pytest.raises(ValueError, match=re.escape("Valid parameters: ['data_file', 'label']")): + resolve_constructor_args( + cls=_Files, raw_args={"nope": 1}, external_input=True, allowed_path_params={"data_file"} + ) + + def test_allowing_a_path_does_not_allow_python_objects(self) -> None: + with pytest.raises(ValueError, match="'handle' of '_Mixed' cannot be set through the API"): + resolve_constructor_args( + cls=_Mixed, raw_args={"handle": None}, external_input=True, allowed_path_params={"handle"} + ) + + def test_in_process_callers_keep_passing_paths(self) -> None: + assert _resolve(_Files, {"data_file": "/etc/hosts"}) == {"data_file": Path("/etc/hosts")} + def test_in_process_callers_keep_passing_python_objects(self) -> None: handle = _Handle() options = {"dtype": object()} From 291c462908f4ddb454659354435cf9f20078ceff Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 26 Oct 2026 20:51:04 +0000 Subject: [PATCH 02/38] Leave out target types that require a path callers may not set The target catalog now applies the path rule to required parameters too, so a target that needs a disallowed path is left out instead of listed without the field, matching the scorer catalog. allowed_path_params is typed frozenset[str] so a bare string cannot turn the allow-list into a substring match. --- pyrit/backend/services/target_service.py | 26 ++++++++++++++++--- pyrit/registry/registry.py | 14 +++++----- pyrit/registry/resolution.py | 8 +++--- tests/unit/backend/test_target_service.py | 14 ++++++++-- .../test_lsb_image_stego_converter.py | 2 +- .../unit/registry/test_converter_registry.py | 4 +-- tests/unit/registry/test_resolution.py | 11 +++++--- 7 files changed, 55 insertions(+), 24 deletions(-) diff --git a/pyrit/backend/services/target_service.py b/pyrit/backend/services/target_service.py index 2cdce7f38b..413c91df90 100644 --- a/pyrit/backend/services/target_service.py +++ b/pyrit/backend/services/target_service.py @@ -230,9 +230,9 @@ async def list_target_types_async(self, *, upload_directory: str | None = None) Returns every target that external callers can build, with the constructor parameters they may supply, each described in the form callers send it, and the auth modes it supports, all projected from the registry's - ``TargetMetadata``; targets that need a Python object for a required - parameter are left out, and so are targets that upload local files when the - server has no upload folder. Deciding which entries to surface to a + ``TargetMetadata``; targets that need a Python object or a path callers may not + set for a required parameter are left out, and so are targets that upload local + files when the server has no upload folder. Deciding which entries to surface to a user is a presentation concern owned by the caller (e.g. the frontend), not this service. @@ -259,7 +259,7 @@ async def list_target_types_async(self, *, upload_directory: str | None = None) description=metadata.class_description or None, ) for metadata in metadata_items - if all(parameter.is_external_input for parameter in metadata.parameters if parameter.required) + if self._has_required_params(target_type=metadata.class_name, parameters=metadata.parameters) and self._has_upload_directory(target_type=metadata.class_name, upload_directory=upload_directory) ] return TargetTypeResponse(items=items) @@ -287,6 +287,24 @@ async def resolve_binding_async(self, binding: TargetBinding) -> PromptTarget: """ return await self.build_from_source_async(binding.to_spec()) + def _has_required_params(self, *, target_type: str, parameters: tuple[Parameter, ...]) -> bool: + """ + Return whether a target type can receive every required parameter when built from API input. + + Args: + target_type (str): Registered target class name. + parameters (tuple[Parameter, ...]): Constructor parameters derived by the registry. + + Returns: + bool: False when a required parameter needs a Python object or a path callers may not set. + """ + path_names = self._allowed_path_params(target_cls=self._registry.get_class(target_type)) + return all( + parameter.is_external_input and (not parameter.mentions_path or parameter.name in path_names) + for parameter in parameters + if parameter.required + ) + def _has_upload_directory(self, *, target_type: str, upload_directory: str | None) -> bool: """ Return whether a target type uploads no local files or the server gives it an upload folder. diff --git a/pyrit/registry/registry.py b/pyrit/registry/registry.py index f78fc378cd..7ee36475e5 100644 --- a/pyrit/registry/registry.py +++ b/pyrit/registry/registry.py @@ -40,7 +40,7 @@ ) if TYPE_CHECKING: - from collections.abc import Collection, Iterator, Mapping + from collections.abc import Iterator, Mapping from types import ModuleType from typing import Self @@ -764,7 +764,7 @@ def create_instance(self, name: str, **kwargs: object) -> T: return cls(**resolved) def create_instance_from_external_input( - self, name: str, *, params: Mapping[str, object], allowed_path_params: Collection[str] = frozenset() + self, name: str, *, params: Mapping[str, object], allowed_path_params: frozenset[str] = frozenset() ) -> T: """ Build a configured instance from external input such as a REST or CLI request. @@ -779,7 +779,7 @@ def create_instance_from_external_input( Args: name (str): The catalog name to build. params (Mapping[str, object]): Constructor arguments from the external caller. - allowed_path_params (Collection[str]): Parameters that can name a server file + allowed_path_params (frozenset[str]): Parameters that can name a server file or folder but that the caller has made safe. Defaults to none. Returns: @@ -877,7 +877,7 @@ def create_named_instance( params: Mapping[str, object] | None = None, registry_metadata: dict[str, Any] | None = None, external_input: bool = False, - allowed_path_params: Collection[str] = frozenset(), + allowed_path_params: frozenset[str] = frozenset(), ) -> InstanceT: """ Build and store a configured instance under an explicit name. @@ -891,7 +891,7 @@ def create_named_instance( external_input (bool): Whether ``params`` come from an external caller, in which case the instance is built with ``create_instance_from_external_input``. Defaults to False. - allowed_path_params (Collection[str]): External parameters that can name a + allowed_path_params (frozenset[str]): External parameters that can name a server file or folder but that the caller has made safe. Defaults to none. Returns: @@ -967,7 +967,7 @@ def recreate_instance( source: InstanceT, params: Mapping[str, object], external_input: bool = False, - allowed_path_params: Collection[str] = frozenset(), + allowed_path_params: frozenset[str] = frozenset(), ) -> InstanceT: """ Build an independent component; never change or register the source. @@ -977,7 +977,7 @@ def recreate_instance( params (Mapping[str, object]): Constructor overrides. external_input (bool): Whether ``params`` come from an external caller. Defaults to False. - allowed_path_params (Collection[str]): External overrides that can name a + allowed_path_params (frozenset[str]): External overrides that can name a server file or folder but that the caller has made safe. Defaults to none. Returns: diff --git a/pyrit/registry/resolution.py b/pyrit/registry/resolution.py index cd2729d8bc..f641f436cd 100644 --- a/pyrit/registry/resolution.py +++ b/pyrit/registry/resolution.py @@ -623,7 +623,7 @@ def resolve_constructor_args( external_input: bool = False, strict_json: bool = False, seam_base_type: type | None = None, - allowed_path_params: Collection[str] = frozenset(), + allowed_path_params: frozenset[str] = frozenset(), ) -> dict[str, Any]: """ Resolve a flat argument dict into constructor-ready keyword arguments. @@ -651,7 +651,7 @@ def resolve_constructor_args( contracts; ordinary CLI and Python resolution remain unchanged. seam_base_type (type | None): The domain base type whose non-required declared references are construction seams (see ``derive_parameters``). - allowed_path_params (Collection[str]): External-input parameters that can name a + allowed_path_params (frozenset[str]): External-input parameters that can name a server file or folder but that the caller has made safe, such as an uploaded file the backend stored itself. Defaults to none. @@ -728,7 +728,7 @@ def reject_non_external_params( params: Mapping[str, Any], declared: Sequence[Parameter], owner: str, - allowed_path_params: Collection[str] = frozenset(), + allowed_path_params: frozenset[str] = frozenset(), ) -> None: """ Reject external input that is not an explicitly supported external input. @@ -742,7 +742,7 @@ def reject_non_external_params( params (Mapping[str, Any]): The parameter values supplied by an external caller. declared (Sequence[Parameter]): The parameters the component declares. owner (str): The owning class or scenario name, for error messages. - allowed_path_params (Collection[str]): Parameters that can name a server file or + allowed_path_params (frozenset[str]): Parameters that can name a server file or folder but that the caller has made safe. Defaults to none. Raises: diff --git a/tests/unit/backend/test_target_service.py b/tests/unit/backend/test_target_service.py index 67727f2967..50a7b75464 100644 --- a/tests/unit/backend/test_target_service.py +++ b/tests/unit/backend/test_target_service.py @@ -552,10 +552,19 @@ def __init__(self, *, data_file: Path | None = None) -> None: self.data_file = data_file +class _RequiredFileTarget(MockPromptTarget): + """A target that cannot be built without a file on the server.""" + + def __init__(self, *, data_file: Path) -> None: + super().__init__() + self.data_file = data_file + + @pytest.mark.usefixtures("patch_central_database") async def test_target_path_parameters_are_left_out_and_rejected() -> None: service = TargetService() service._registry.register_class(_FileTarget) + service._registry.register_class(_RequiredFileTarget) types = await service.list_target_types_async() with pytest.raises(ValueError, match="'data_file' of '_FileTarget' names a server file"): @@ -563,8 +572,9 @@ async def test_target_path_parameters_are_left_out_and_rejected() -> None: request=CreateTargetRequest(name="files", type="_FileTarget", params={"data_file": "/etc/hosts"}) ) - entry = next(item for item in types.items if item.target_type == "_FileTarget") - assert "data_file" not in [parameter.name for parameter in entry.parameters] + entries = {item.target_type: item for item in types.items} + assert "data_file" not in [parameter.name for parameter in entries["_FileTarget"].parameters] + assert "_RequiredFileTarget" not in entries assert service.get_target_object(target_registry_name="files") is None diff --git a/tests/unit/converter/test_lsb_image_stego_converter.py b/tests/unit/converter/test_lsb_image_stego_converter.py index f72c793984..b70ab997e6 100644 --- a/tests/unit/converter/test_lsb_image_stego_converter.py +++ b/tests/unit/converter/test_lsb_image_stego_converter.py @@ -329,7 +329,7 @@ def test_lsb_converter_registry_accepts_carrier_image_path(tmp_path: Path, *, us registry.create_instance_from_external_input("LSBImageStegoConverter", params=params) converter = registry.create_instance_from_external_input( - "LSBImageStegoConverter", params=params, allowed_path_params={"carrier_image_path"} + "LSBImageStegoConverter", params=params, allowed_path_params=frozenset({"carrier_image_path"}) ) assert isinstance(converter, LSBImageStegoConverter) assert converter.get_identifier().params["carrier_image_path"] == str(carrier_path) diff --git a/tests/unit/registry/test_converter_registry.py b/tests/unit/registry/test_converter_registry.py index 252907d35b..7710ca5e16 100644 --- a/tests/unit/registry/test_converter_registry.py +++ b/tests/unit/registry/test_converter_registry.py @@ -241,7 +241,7 @@ def test_create_named_instance_rejects_path_parameters_unless_allowed(self, regi type_name="ImageOverlayConverter", params=params, external_input=True, - allowed_path_params={"base_image"}, + allowed_path_params=frozenset({"base_image"}), ) assert registry.instances.get("overlay") is converter @@ -254,7 +254,7 @@ def test_recreate_instance_rejects_path_overrides_unless_allowed(self, registry: registry.recreate_instance(source=source, params=params, external_input=True) variant = registry.recreate_instance( - source=source, params=params, external_input=True, allowed_path_params={"base_image"} + source=source, params=params, external_input=True, allowed_path_params=frozenset({"base_image"}) ) assert variant.get_identifier().params["base_image"] == "other.png" assert source.get_identifier().params["base_image"] == "input.png" diff --git a/tests/unit/registry/test_resolution.py b/tests/unit/registry/test_resolution.py index 4ecb8909c2..aeaf4bff51 100644 --- a/tests/unit/registry/test_resolution.py +++ b/tests/unit/registry/test_resolution.py @@ -523,13 +523,16 @@ def test_external_accepts_only_the_allowed_path_parameters(self) -> None: cls=_Files, raw_args={"label": "x", "data_file": "/uploads/data.bin"}, external_input=True, - allowed_path_params={"data_file"}, + allowed_path_params=frozenset({"data_file"}), ) assert resolved == {"label": "x", "data_file": Path("/uploads/data.bin")} with pytest.raises(ValueError, match="'source' of '_Files' names a server file"): resolve_constructor_args( - cls=_Files, raw_args={"source": "/etc/hosts"}, external_input=True, allowed_path_params={"data_file"} + cls=_Files, + raw_args={"source": "/etc/hosts"}, + external_input=True, + allowed_path_params=frozenset({"data_file"}), ) def test_unknown_parameter_lists_only_parameters_callers_may_set(self) -> None: @@ -537,13 +540,13 @@ def test_unknown_parameter_lists_only_parameters_callers_may_set(self) -> None: resolve_constructor_args(cls=_Files, raw_args={"nope": 1}, external_input=True) with pytest.raises(ValueError, match=re.escape("Valid parameters: ['data_file', 'label']")): resolve_constructor_args( - cls=_Files, raw_args={"nope": 1}, external_input=True, allowed_path_params={"data_file"} + cls=_Files, raw_args={"nope": 1}, external_input=True, allowed_path_params=frozenset({"data_file"}) ) def test_allowing_a_path_does_not_allow_python_objects(self) -> None: with pytest.raises(ValueError, match="'handle' of '_Mixed' cannot be set through the API"): resolve_constructor_args( - cls=_Mixed, raw_args={"handle": None}, external_input=True, allowed_path_params={"handle"} + cls=_Mixed, raw_args={"handle": None}, external_input=True, allowed_path_params=frozenset({"handle"}) ) def test_in_process_callers_keep_passing_paths(self) -> None: From 797831b0ea227d251b0ef9184c57f2363f0c33f4 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 26 Oct 2026 21:33:14 +0000 Subject: [PATCH 03/38] Test a target whose required path is the server upload folder The target catalog still lists a target whose required path parameter is its server-set upload folder when the server has one, and leaves it out otherwise. --- tests/unit/backend/test_target_service.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/unit/backend/test_target_service.py b/tests/unit/backend/test_target_service.py index 50a7b75464..4e2849de8c 100644 --- a/tests/unit/backend/test_target_service.py +++ b/tests/unit/backend/test_target_service.py @@ -560,6 +560,16 @@ def __init__(self, *, data_file: Path) -> None: self.data_file = data_file +class _RequiredUploadTarget(MockPromptTarget): + """A target that uploads local files from a folder only the server sets.""" + + upload_directory_parameter = "upload_dir" + + def __init__(self, *, upload_dir: Path) -> None: + super().__init__() + self.upload_dir = upload_dir + + @pytest.mark.usefixtures("patch_central_database") async def test_target_path_parameters_are_left_out_and_rejected() -> None: service = TargetService() @@ -578,6 +588,19 @@ async def test_target_path_parameters_are_left_out_and_rejected() -> None: assert service.get_target_object(target_registry_name="files") is None +@pytest.mark.usefixtures("patch_central_database") +async def test_types_list_a_target_whose_required_path_is_the_server_upload_folder(tmp_path: Path) -> None: + service = TargetService() + service._registry.register_class(_RequiredUploadTarget) + + with_folder = await service.list_target_types_async(upload_directory=str(tmp_path)) + without_folder = await service.list_target_types_async() + + entry = next(item for item in with_folder.items if item.target_type == "_RequiredUploadTarget") + assert "upload_dir" not in [parameter.name for parameter in entry.parameters] + assert all(item.target_type != "_RequiredUploadTarget" for item in without_folder.items) + + @pytest.mark.usefixtures("patch_central_database") class TestCreateTargetCodeExecution: """Admin gate for targets whose construction runs model repository code.""" From 521ec7e5a0f1482a0ab09aeac2faeccf514aac8c Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 18:33:18 +0000 Subject: [PATCH 04/38] TEST: Cover saved converter copies in the server path check --- tests/unit/backend/test_converter_service.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/unit/backend/test_converter_service.py b/tests/unit/backend/test_converter_service.py index 46a829c051..cf33c9bd06 100644 --- a/tests/unit/backend/test_converter_service.py +++ b/tests/unit/backend/test_converter_service.py @@ -960,6 +960,13 @@ async def test_variants_and_previews_reject_server_path_overrides( await upload_service.create_converter_async( request=CreateConverterRequest(type="ColloquialWordswapConverter", register=False, source=spec) ) + with pytest.raises(ValueError, match="'wordswap_path' must be uploaded as a data URI"): + await upload_service.create_converter_async( + request=CreateConverterRequest( + type="ColloquialWordswapConverter", register=True, name="saved_wordswap", source=spec + ) + ) + assert upload_service._registry.instances.get("saved_wordswap") is None with pytest.raises(ValueError, match="'wordswap_path' must be uploaded as a data URI"): await upload_service.preview_conversion_async( request=ConverterPreviewRequest( From 66fd155e705676037e31f0cfcfeaac8a80557978 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 19:29:26 +0000 Subject: [PATCH 05/38] FEAT: Report the documents a storage listing cannot use --- pyrit/registry/file_document_storage.py | 125 ++++++++++++++ .../registry/test_file_document_storage.py | 158 ++++++++++++++++++ 2 files changed, 283 insertions(+) create mode 100644 tests/unit/registry/test_file_document_storage.py diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index 56432d80b7..bd16601de0 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -12,6 +12,7 @@ from __future__ import annotations +import fnmatch import hashlib import logging import os @@ -19,6 +20,7 @@ import tempfile import time from contextlib import contextmanager, suppress +from dataclasses import dataclass, field from pathlib import Path, PurePosixPath from typing import TYPE_CHECKING from urllib.parse import unquote, urlparse @@ -33,6 +35,32 @@ logger = logging.getLogger(__name__) + +@dataclass(frozen=True) +class UnreadableDocument: + """A stored document with a valid name whose content could not be read.""" + + name: str + message: str + + +@dataclass(frozen=True) +class IgnoredDocument: + """A stored file whose name breaks the naming rule, so no single-document call can address it.""" + + file_name: str + message: str + + +@dataclass(frozen=True) +class DocumentListing: + """Every document a listing found, including the ones it could not use.""" + + documents: dict[str, bytes] = field(default_factory=dict) + unreadable: tuple[UnreadableDocument, ...] = () + ignored: tuple[IgnoredDocument, ...] = () + + # Both platforms lock a single byte at offset zero, which they allow past the end of an # empty file. The lock belongs to the open file handle, so it excludes other threads in # this process as well as other processes, and the kernel drops it if the holder exits. @@ -199,6 +227,103 @@ def _list_documents(self) -> dict[str, bytes]: """ return self._list_blob_documents() if self._is_blob else self._list_local_documents() + def _list_documents_with_issues(self) -> DocumentListing: + """ + Read every stored document and report the ones that can't be used. + + Unlike ``_list_documents``, nothing is skipped quietly: a document that can't be read + is reported as unreadable, and a file whose name breaks the naming rule is reported as + ignored without being read, because no single-document call could address it. A + storage location that can't be listed at all raises instead of looking empty. + + Messages carry only a short reason, never a path or URL, so they can be shown to users + without exposing server paths or a SAS token. + + Returns: + DocumentListing: Readable documents keyed by name, plus the unreadable and ignored ones. + """ + return self._list_blob_documents_with_issues() if self._is_blob else self._list_local_documents_with_issues() + + def _list_local_documents_with_issues(self) -> DocumentListing: + """ + Read the local directory, reporting unreadable and badly named documents. + + Returns: + DocumentListing: The documents and issues found in the directory. + """ + directory = self._local_directory(create=True) + # ``Path.glob`` returns nothing for a directory it can't read; listing it directly raises instead. + # ``fnmatch`` keeps glob's per-platform case rule, so both listings see the same files. + with os.scandir(directory) as entries: + file_names = sorted(entry.name for entry in entries if fnmatch.fnmatch(entry.name, f"*{self._extension}")) + documents: dict[str, bytes] = {} + unreadable: list[UnreadableDocument] = [] + ignored: list[IgnoredDocument] = [] + for file_name in file_names: + name = Path(file_name).stem + problem = self._name_problem(name) + if problem is not None: + ignored.append(IgnoredDocument(file_name=file_name, message=problem)) + continue + try: + documents[name] = (directory / file_name).read_bytes() + except FileNotFoundError: + # Deleted after the directory was listed, so it is no longer part of the listing. + continue + except OSError as error: + unreadable.append(UnreadableDocument(name=name, message=error.strerror or type(error).__name__)) + return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) + + def _list_blob_documents_with_issues(self) -> DocumentListing: + """ + Read the blob container, reporting unreadable and badly named documents. + + Returns: + DocumentListing: The documents and issues found under the configured prefix. + """ + from azure.core.exceptions import AzureError, ResourceNotFoundError + + documents: dict[str, bytes] = {} + unreadable: list[UnreadableDocument] = [] + ignored: list[IgnoredDocument] = [] + with self._open_container_client() as client: + prefix = f"{self._blob_prefix}/" if self._blob_prefix else None + blobs = client.list_blobs(name_starts_with=prefix) if prefix else client.list_blobs() + blob_names = sorted( + blob.name for blob in blobs if self._is_direct_document_blob(blob_name=blob.name, prefix=prefix) + ) + for blob_name in blob_names: + file_name = blob_name.removeprefix(prefix or "") + name = PurePosixPath(file_name).stem + problem = self._name_problem(name) + if problem is not None: + ignored.append(IgnoredDocument(file_name=file_name, message=problem)) + continue + try: + documents[name] = client.download_blob(blob_name).readall() + except ResourceNotFoundError: + # Deleted after the container was listed, so it is no longer part of the listing. + continue + except AzureError as error: + status = getattr(error, "status_code", None) + reason = type(error).__name__ if status is None else f"{type(error).__name__} ({status})" + unreadable.append(UnreadableDocument(name=name, message=reason)) + return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) + + @staticmethod + def _name_problem(name: str) -> str | None: + """ + Return why a discovered document name can't be addressed, or None if it can. + + Returns: + str | None: The naming-rule error, or ``None`` for a legal name. + """ + try: + validate_registry_name(name) + except ValueError as error: + return str(error) + return None + def _list_local_documents(self) -> dict[str, bytes]: """ Read addressable documents from the configured local directory. diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py new file mode 100644 index 0000000000..e6d67bec77 --- /dev/null +++ b/tests/unit/registry/test_file_document_storage.py @@ -0,0 +1,158 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from azure.core.exceptions import HttpResponseError, ResourceNotFoundError +from azure.storage.blob import ContainerClient + +from pyrit.registry.file_document_storage import DocumentListing, FileDocumentStorage, UnreadableDocument + +_SAS_SOURCE = "https://account.blob.core.windows.net/documents/saved?sp=rl&sig=secret-signature" + + +def _local_storage(directory: Path) -> FileDocumentStorage: + return FileDocumentStorage(source=str(directory), extension=".json", source_label="Test documents") + + +def _blob_storage() -> FileDocumentStorage: + return FileDocumentStorage(source=_SAS_SOURCE, extension=".json", source_label="Test documents") + + +def _blob_client(*, names: list[str]) -> MagicMock: + client = MagicMock(spec=ContainerClient) + client.__enter__.return_value = client + client.list_blobs.return_value = [SimpleNamespace(name=name) for name in names] + return client + + +class TestLocalListingWithIssues: + def test_reports_unreadable_document_and_keeps_the_others(self, tmp_path: Path) -> None: + (tmp_path / "first.json").write_bytes(b"one") + (tmp_path / "second.json").write_bytes(b"two") + original_read_bytes = Path.read_bytes + + def read_bytes(path: Path) -> bytes: + if path.name == "second.json": + raise PermissionError(13, "Permission denied", str(path)) + return original_read_bytes(path) + + with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + listing = _local_storage(tmp_path)._list_documents_with_issues() + legacy = _local_storage(tmp_path)._list_documents() + + assert listing == DocumentListing( + documents={"first": b"one"}, + unreadable=(UnreadableDocument(name="second", message="Permission denied"),), + ) + assert legacy == {"first": b"one"} + + def test_reports_badly_named_files_as_ignored_without_reading_them(self, tmp_path: Path) -> None: + (tmp_path / "valid_name.json").write_bytes(b"ok") + (tmp_path / "Bad-Name.json").write_bytes(b"never read") + (tmp_path / "notes.txt").write_bytes(b"other extension") + (tmp_path / "nested").mkdir() + (tmp_path / "nested" / "inner.json").write_bytes(b"not a direct child") + read_names: list[str] = [] + original_read_bytes = Path.read_bytes + + def read_bytes(path: Path) -> bytes: + read_names.append(path.name) + return original_read_bytes(path) + + with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing.documents == {"valid_name": b"ok"} + assert listing.unreadable == () + assert [issue.file_name for issue in listing.ignored] == ["Bad-Name.json"] + assert "Bad-Name" in listing.ignored[0].message + assert read_names == ["valid_name.json"] + assert _local_storage(tmp_path)._list_documents() == listing.documents + + def test_skips_a_document_deleted_after_listing(self, tmp_path: Path) -> None: + (tmp_path / "kept.json").write_bytes(b"kept") + (tmp_path / "deleted.json").write_bytes(b"gone") + original_read_bytes = Path.read_bytes + + def read_bytes(path: Path) -> bytes: + if path.name == "deleted.json": + raise FileNotFoundError(2, "No such file or directory", str(path)) + return original_read_bytes(path) + + with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing(documents={"kept": b"kept"}) + + def test_creates_missing_directory_and_lists_nothing(self, tmp_path: Path) -> None: + listing = _local_storage(tmp_path / "missing")._list_documents_with_issues() + + assert listing == DocumentListing() + assert (tmp_path / "missing").is_dir() + + def test_raises_when_the_directory_cannot_be_listed(self, tmp_path: Path) -> None: + (tmp_path / "first.json").write_bytes(b"one") + + with patch("pyrit.registry.file_document_storage.os.scandir", side_effect=PermissionError(13, "denied")): + with pytest.raises(PermissionError): + _local_storage(tmp_path)._list_documents_with_issues() + + +class TestBlobListingWithIssues: + def test_reports_unreadable_and_ignored_blobs_without_exposing_the_url(self) -> None: + client = _blob_client( + names=[ + "saved/good.json", + "saved/broken.json", + "saved/Bad-Name.json", + "saved/notes.txt", + "saved/deeper/inner.json", + ] + ) + + def download_blob(blob_name: str) -> SimpleNamespace: + if blob_name == "saved/broken.json": + error = HttpResponseError(message=f"Forbidden: {_SAS_SOURCE}") + error.status_code = 403 + raise error + return SimpleNamespace(readall=lambda: b"content") + + client.download_blob.side_effect = download_blob + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + listing = _blob_storage()._list_documents_with_issues() + + assert listing.documents == {"good": b"content"} + assert listing.unreadable == (UnreadableDocument(name="broken", message="HttpResponseError (403)"),) + assert [issue.file_name for issue in listing.ignored] == ["Bad-Name.json"] + assert all("secret-signature" not in issue.message for issue in (*listing.unreadable, *listing.ignored)) + client.list_blobs.assert_called_once_with(name_starts_with="saved/") + downloaded = [call.args[0] for call in client.download_blob.call_args_list] + assert downloaded == ["saved/broken.json", "saved/good.json"] + + def test_skips_a_blob_deleted_after_listing(self) -> None: + client = _blob_client(names=["saved/kept.json", "saved/deleted.json"]) + + def download_blob(blob_name: str) -> SimpleNamespace: + if blob_name == "saved/deleted.json": + raise ResourceNotFoundError(message="The specified blob does not exist.") + return SimpleNamespace(readall=lambda: b"kept") + + client.download_blob.side_effect = download_blob + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + listing = _blob_storage()._list_documents_with_issues() + + assert listing == DocumentListing(documents={"kept": b"kept"}) + + def test_raises_when_the_container_cannot_be_listed(self) -> None: + client = _blob_client(names=[]) + client.list_blobs.side_effect = HttpResponseError(message="listing failed") + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + with pytest.raises(HttpResponseError): + _blob_storage()._list_documents_with_issues() From 2d8c465c9dfefb8efa0b1e47cf0d1ca1d8199927 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 20:12:11 +0000 Subject: [PATCH 06/38] Report non-files and missing containers instead of skipping them --- pyrit/registry/file_document_storage.py | 58 +++++++++---- .../registry/test_file_document_storage.py | 84 +++++++++++++++++-- 2 files changed, 116 insertions(+), 26 deletions(-) diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index bd16601de0..17f2219ebb 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -253,25 +253,30 @@ def _list_local_documents_with_issues(self) -> DocumentListing: """ directory = self._local_directory(create=True) # ``Path.glob`` returns nothing for a directory it can't read; listing it directly raises instead. - # ``fnmatch`` keeps glob's per-platform case rule, so both listings see the same files. + # ``fnmatch`` applies the platform's case rule, as ``glob`` does. with os.scandir(directory) as entries: - file_names = sorted(entry.name for entry in entries if fnmatch.fnmatch(entry.name, f"*{self._extension}")) + candidates = sorted( + (entry.name, entry.is_file()) for entry in entries if fnmatch.fnmatch(entry.name, f"*{self._extension}") + ) documents: dict[str, bytes] = {} unreadable: list[UnreadableDocument] = [] ignored: list[IgnoredDocument] = [] - for file_name in file_names: + for file_name, is_file in candidates: name = Path(file_name).stem problem = self._name_problem(name) if problem is not None: ignored.append(IgnoredDocument(file_name=file_name, message=problem)) - continue - try: - documents[name] = (directory / file_name).read_bytes() - except FileNotFoundError: - # Deleted after the directory was listed, so it is no longer part of the listing. - continue - except OSError as error: - unreadable.append(UnreadableDocument(name=name, message=error.strerror or type(error).__name__)) + elif not is_file: + # A broken link, a folder, or a pipe, which would block if read. + unreadable.append(UnreadableDocument(name=name, message="Not a regular file")) + else: + try: + documents[name] = (directory / file_name).read_bytes() + except FileNotFoundError: + # Deleted after the directory was listed, so it is no longer part of the listing. + continue + except OSError as error: + unreadable.append(UnreadableDocument(name=name, message=self._describe_read_error(error))) return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) def _list_blob_documents_with_issues(self) -> DocumentListing: @@ -281,7 +286,8 @@ def _list_blob_documents_with_issues(self) -> DocumentListing: Returns: DocumentListing: The documents and issues found under the configured prefix. """ - from azure.core.exceptions import AzureError, ResourceNotFoundError + from azure.core.exceptions import AzureError + from azure.storage.blob import StorageErrorCode documents: dict[str, bytes] = {} unreadable: list[UnreadableDocument] = [] @@ -301,15 +307,31 @@ def _list_blob_documents_with_issues(self) -> DocumentListing: continue try: documents[name] = client.download_blob(blob_name).readall() - except ResourceNotFoundError: - # Deleted after the container was listed, so it is no longer part of the listing. - continue except AzureError as error: - status = getattr(error, "status_code", None) - reason = type(error).__name__ if status is None else f"{type(error).__name__} ({status})" - unreadable.append(UnreadableDocument(name=name, message=reason)) + if getattr(error, "error_code", None) == StorageErrorCode.BLOB_NOT_FOUND: + # Deleted after the container was listed, so it is no longer part of the listing. + continue + unreadable.append(UnreadableDocument(name=name, message=self._describe_read_error(error))) return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) + @staticmethod + def _describe_read_error(error: Exception) -> str: + """ + Describe a storage error in a few words, leaving out any path or URL it carries. + + Returns: + str: For example ``Permission denied`` or ``HttpResponseError (403, AuthorizationPermissionMismatch)``. + """ + if isinstance(error, OSError) and error.strerror: + return error.strerror + error_code = getattr(error, "error_code", None) + details = [ + str(detail) + for detail in (getattr(error, "status_code", None), getattr(error_code, "value", error_code)) + if detail is not None + ] + return f"{type(error).__name__} ({', '.join(details)})" if details else type(error).__name__ + @staticmethod def _name_problem(name: str) -> str | None: """ diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index e6d67bec77..face54fe80 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -1,13 +1,14 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from collections.abc import Iterator from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest -from azure.core.exceptions import HttpResponseError, ResourceNotFoundError -from azure.storage.blob import ContainerClient +from azure.core.exceptions import HttpResponseError, ResourceNotFoundError, ServiceRequestError +from azure.storage.blob import ContainerClient, StorageErrorCode from pyrit.registry.file_document_storage import DocumentListing, FileDocumentStorage, UnreadableDocument @@ -88,6 +89,30 @@ def read_bytes(path: Path) -> bytes: assert listing == DocumentListing(documents={"kept": b"kept"}) + def test_reports_entries_that_are_not_regular_files_without_reading_them(self, tmp_path: Path) -> None: + (tmp_path / "folder.json").mkdir() + try: + (tmp_path / "broken_link.json").symlink_to(tmp_path / "missing_target.json") + except OSError: + pytest.skip("Creating symbolic links isn't permitted here") + + with patch.object(Path, "read_bytes", autospec=True) as read_bytes: + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing.unreadable == ( + UnreadableDocument(name="broken_link", message="Not a regular file"), + UnreadableDocument(name="folder", message="Not a regular file"), + ) + read_bytes.assert_not_called() + + def test_matches_extensions_like_the_legacy_listing(self, tmp_path: Path) -> None: + (tmp_path / "lower.json").write_bytes(b"one") + (tmp_path / "upper.JSON").write_bytes(b"two") + + storage = _local_storage(tmp_path) + + assert storage._list_documents_with_issues().documents == storage._list_documents() + def test_creates_missing_directory_and_lists_nothing(self, tmp_path: Path) -> None: listing = _local_storage(tmp_path / "missing")._list_documents_with_issues() @@ -134,12 +159,19 @@ def download_blob(blob_name: str) -> SimpleNamespace: downloaded = [call.args[0] for call in client.download_blob.call_args_list] assert downloaded == ["saved/broken.json", "saved/good.json"] - def test_skips_a_blob_deleted_after_listing(self) -> None: - client = _blob_client(names=["saved/kept.json", "saved/deleted.json"]) + def test_skips_only_a_blob_deleted_after_listing(self) -> None: + client = _blob_client(names=["saved/kept.json", "saved/deleted.json", "saved/container_gone.json"]) + error_codes = { + "saved/deleted.json": StorageErrorCode.BLOB_NOT_FOUND, + "saved/container_gone.json": StorageErrorCode.CONTAINER_NOT_FOUND, + } def download_blob(blob_name: str) -> SimpleNamespace: - if blob_name == "saved/deleted.json": - raise ResourceNotFoundError(message="The specified blob does not exist.") + if blob_name in error_codes: + error = ResourceNotFoundError(message="Not found.") + error.status_code = 404 + error.error_code = error_codes[blob_name] + raise error return SimpleNamespace(readall=lambda: b"kept") client.download_blob.side_effect = download_blob @@ -147,12 +179,48 @@ def download_blob(blob_name: str) -> SimpleNamespace: with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): listing = _blob_storage()._list_documents_with_issues() - assert listing == DocumentListing(documents={"kept": b"kept"}) + assert listing == DocumentListing( + documents={"kept": b"kept"}, + unreadable=( + UnreadableDocument(name="container_gone", message="ResourceNotFoundError (404, ContainerNotFound)"), + ), + ) def test_raises_when_the_container_cannot_be_listed(self) -> None: + def failing_pages() -> Iterator[SimpleNamespace]: + raise HttpResponseError(message="listing failed") + yield SimpleNamespace(name="saved/never.json") + client = _blob_client(names=[]) - client.list_blobs.side_effect = HttpResponseError(message="listing failed") + client.list_blobs.return_value = failing_pages() with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): with pytest.raises(HttpResponseError): _blob_storage()._list_documents_with_issues() + + +class TestDescribeReadError: + @pytest.mark.parametrize( + ("error", "expected"), + [ + (PermissionError(13, "Permission denied", "/server/path/secret.json"), "Permission denied"), + (OSError(), "OSError"), + (ServiceRequestError(message="https://account.blob.core.windows.net/x?sig=secret"), "ServiceRequestError"), + ], + ) + def test_describes_errors_without_their_details(self, error: Exception, expected: str) -> None: + assert FileDocumentStorage._describe_read_error(error) == expected + + @pytest.mark.parametrize( + ("error_code", "expected"), + [ + (StorageErrorCode.AUTHORIZATION_FAILURE, "HttpResponseError (403, AuthorizationFailure)"), + ("AuthorizationPermissionMismatch", "HttpResponseError (403, AuthorizationPermissionMismatch)"), + ], + ) + def test_includes_the_azure_status_and_error_code(self, error_code: str, expected: str) -> None: + error = HttpResponseError(message="https://account.blob.core.windows.net/x?sig=secret") + error.status_code = 403 + error.error_code = error_code + + assert FileDocumentStorage._describe_read_error(error) == expected From 5dec828431f959e3e1f68852f6b68680130954ea Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 21:10:09 +0000 Subject: [PATCH 07/38] Report an entry that can't be checked instead of failing the whole listing --- pyrit/registry/file_document_storage.py | 44 ++++++++++++------- .../registry/test_file_document_storage.py | 12 +++++ 2 files changed, 41 insertions(+), 15 deletions(-) diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index 17f2219ebb..4ba3f0949d 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -256,29 +256,43 @@ def _list_local_documents_with_issues(self) -> DocumentListing: # ``fnmatch`` applies the platform's case rule, as ``glob`` does. with os.scandir(directory) as entries: candidates = sorted( - (entry.name, entry.is_file()) for entry in entries if fnmatch.fnmatch(entry.name, f"*{self._extension}") + (entry for entry in entries if fnmatch.fnmatch(entry.name, f"*{self._extension}")), + key=lambda entry: entry.name, ) documents: dict[str, bytes] = {} unreadable: list[UnreadableDocument] = [] ignored: list[IgnoredDocument] = [] - for file_name, is_file in candidates: - name = Path(file_name).stem + for entry in candidates: + name = Path(entry.name).stem problem = self._name_problem(name) if problem is not None: - ignored.append(IgnoredDocument(file_name=file_name, message=problem)) - elif not is_file: - # A broken link, a folder, or a pipe, which would block if read. - unreadable.append(UnreadableDocument(name=name, message="Not a regular file")) - else: - try: - documents[name] = (directory / file_name).read_bytes() - except FileNotFoundError: - # Deleted after the directory was listed, so it is no longer part of the listing. - continue - except OSError as error: - unreadable.append(UnreadableDocument(name=name, message=self._describe_read_error(error))) + ignored.append(IgnoredDocument(file_name=entry.name, message=problem)) + continue + content = self._read_local_entry(entry=entry, name=name) + if isinstance(content, bytes): + documents[name] = content + elif content is not None: + unreadable.append(content) return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) + def _read_local_entry(self, *, entry: os.DirEntry[str], name: str) -> bytes | UnreadableDocument | None: + """ + Read one listed local document. + + Returns: + bytes | UnreadableDocument | None: The content, why it can't be read, or None if it was + deleted after the directory was listed. + """ + try: + if not entry.is_file(): + # A broken link, a folder, or a pipe, which would block if read. + return UnreadableDocument(name=name, message="Not a regular file") + return Path(entry.path).read_bytes() + except FileNotFoundError: + return None + except OSError as error: + return UnreadableDocument(name=name, message=self._describe_read_error(error)) + def _list_blob_documents_with_issues(self) -> DocumentListing: """ Read the blob container, reporting unreadable and badly named documents. diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index face54fe80..63bc7601dc 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -105,6 +105,18 @@ def test_reports_entries_that_are_not_regular_files_without_reading_them(self, t ) read_bytes.assert_not_called() + def test_reports_a_link_loop_and_keeps_the_others(self, tmp_path: Path) -> None: + (tmp_path / "good.json").write_bytes(b"ok") + try: + (tmp_path / "loop.json").symlink_to(tmp_path / "loop.json") + except OSError: + pytest.skip("Creating symbolic links isn't permitted here") + + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing.documents == {"good": b"ok"} + assert [issue.name for issue in listing.unreadable] == ["loop"] + def test_matches_extensions_like_the_legacy_listing(self, tmp_path: Path) -> None: (tmp_path / "lower.json").write_bytes(b"one") (tmp_path / "upper.JSON").write_bytes(b"two") From bf662b71e028aabb159cb7fa47051d3b0d1283db Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 21:50:58 +0000 Subject: [PATCH 08/38] FEAT: Read saved registry recipes --- pyrit/models/catalog/registry_recipe.py | 143 ++++++++++++++ pyrit/registry/registry_recipe_storage.py | 182 ++++++++++++++++++ .../registry/test_registry_recipe_storage.py | 150 +++++++++++++++ 3 files changed, 475 insertions(+) create mode 100644 pyrit/models/catalog/registry_recipe.py create mode 100644 pyrit/registry/registry_recipe_storage.py create mode 100644 tests/unit/registry/test_registry_recipe_storage.py diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py new file mode 100644 index 0000000000..ce2e544588 --- /dev/null +++ b/pyrit/models/catalog/registry_recipe.py @@ -0,0 +1,143 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Saved registry entries: the recipe kept for each one and what restoring them found.""" + +from __future__ import annotations + +from enum import Enum +from typing import Annotated, ClassVar + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator + +from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS + +_Identifier = Annotated[str, Field(max_length=MAX_IDENTIFIER_LENGTH)] +_TechniqueTag = Annotated[str, Field(pattern=r"^[A-Za-z][A-Za-z0-9_]{0,63}$")] + + +class RegistryRecipeKind(str, Enum): + """The registries whose entries can be saved, named after their storage folders.""" + + TARGETS = "targets" + CONVERTERS = "converters" + SCORERS = "scorers" + TECHNIQUES = "techniques" + + +class _RegistryRecipe(BaseModel): + """Fields every saved recipe shares. The entry's name is its document name, so it isn't repeated.""" + + CURRENT_SCHEMA_VERSION: ClassVar[int] = 1 + + model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False, frozen=True) + + schema_version: int = Field( + CURRENT_SCHEMA_VERSION, ge=1, le=CURRENT_SCHEMA_VERSION, description="Shape of this document" + ) + + +class ComponentRecipe(_RegistryRecipe): + """How to rebuild one saved target, converter, or scorer.""" + + class_module: _Identifier = Field( + ..., min_length=1, description="Module that defined the class when the entry was saved" + ) + class_name: _Identifier = Field(..., min_length=1, description="Registered class name") + params: dict[_Identifier, JsonValue] = Field( + default_factory=dict, + max_length=MAX_ITEMS, + description="Constructor parameters as the API accepted them; other entries are named, never embedded", + ) + + +class TechniqueRecipe(_RegistryRecipe): + """ + How to rebuild one saved technique: the fields of its create request, without the name. + + The fields and their limits mirror the technique API's create request, which this library + can't import. + """ + + type: str = Field(..., min_length=1, max_length=256, description="Registered attack type") + params: dict[_Identifier, JsonValue] = Field( + default_factory=dict, max_length=MAX_ITEMS, description="Attack constructor parameters" + ) + description: str | None = None + tags: list[_TechniqueTag] = Field(default_factory=list, max_length=MAX_ITEMS) + request_converters: list[_Identifier] | None = Field(default=None, max_length=MAX_ITEMS) + response_converters: list[_Identifier] | None = Field(default=None, max_length=MAX_ITEMS) + adversarial_chat: _Identifier | None = None + adversarial_system_prompt: str | None = None + adversarial_seed_prompt: str | None = None + adversarial_prompt_template: str | None = None + + @field_validator("tags") + @classmethod + def _validate_tags(cls, value: list[str]) -> list[str]: + folded = {tag.casefold() for tag in value} + if len(folded) != len(value): + raise ValueError("Tags must be unique, including letter case") + if {"all", "default"} & folded: + raise ValueError("all and default are reserved selectors") + return value + + +class SavedEntryRef(BaseModel): + """Points at one saved entry.""" + + model_config = ConfigDict(frozen=True) + + kind: RegistryRecipeKind + name: str + + +class StoredRecipe(BaseModel): + """A saved recipe and the version of the document it was read from.""" + + model_config = ConfigDict(frozen=True) + + ref: SavedEntryRef + recipe: ComponentRecipe | TechniqueRecipe + version: str = Field(..., description="Opaque version of the document this recipe was read from") + + +class RestoreFailureReason(str, Enum): + """Why a saved entry couldn't be restored.""" + + UNREADABLE_STORAGE = "unreadable_storage" + UNREADABLE_RECIPE = "unreadable_recipe" + INVALID_RECIPE = "invalid_recipe" + NEWER_SCHEMA_VERSION = "newer_schema_version" + UNKNOWN_CLASS = "unknown_class" + MODULE_MISMATCH = "module_mismatch" + UNSUPPORTED_CLASS = "unsupported_class" + UNSUPPORTED_CREDENTIAL = "unsupported_credential" + NAME_CLASH = "name_clash" + MISSING_DEPENDENCY = "missing_dependency" + FAILED_DEPENDENCY = "failed_dependency" + DEPENDENCY_LOOP = "dependency_loop" + BUILD_ERROR = "build_error" + + +class RestoreFailure(BaseModel): + """One saved entry, or one whole storage folder, that couldn't be restored.""" + + model_config = ConfigDict(frozen=True) + + kind: RegistryRecipeKind + name: str | None = Field(None, description="Entry name; None when the whole folder couldn't be read") + reason: RestoreFailureReason + message: str = Field(..., description="Why, in words safe to show users") + version: str | None = Field(None, description="Version of the stored document, when it could be read") + dependency: SavedEntryRef | None = Field(None, description="The entry this one needed that is missing or failed") + + +class IgnoredRecipeFile(BaseModel): + """A stored file whose name can't be an entry name, so it was never read.""" + + model_config = ConfigDict(frozen=True) + + kind: RegistryRecipeKind + file_name: str + message: str diff --git a/pyrit/registry/registry_recipe_storage.py b/pyrit/registry/registry_recipe_storage.py new file mode 100644 index 0000000000..0a956a19b8 --- /dev/null +++ b/pyrit/registry/registry_recipe_storage.py @@ -0,0 +1,182 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Storage for saved registry recipes, one folder per registry.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path +from urllib.parse import urlparse + +from pydantic import ValidationError + +from pyrit.common.azure_storage import is_azure_blob_uri +from pyrit.models.catalog.registry_recipe import ( + ComponentRecipe, + IgnoredRecipeFile, + RegistryRecipeKind, + RestoreFailure, + RestoreFailureReason, + SavedEntryRef, + StoredRecipe, + TechniqueRecipe, +) +from pyrit.registry.file_document_storage import FileDocumentStorage + + +@dataclass(frozen=True) +class RecipeListing: + """Every saved recipe of one kind, with the ones that couldn't be read or parsed.""" + + recipes: dict[str, StoredRecipe] = field(default_factory=dict) + failures: tuple[RestoreFailure, ...] = () + ignored: tuple[IgnoredRecipeFile, ...] = () + + +class RegistryRecipeStorage(FileDocumentStorage): + """ + Read the saved recipes of one registry as JSON documents. + + Each registry has its own folder under the configured source (``targets``, ``converters``, + ``scorers``, ``techniques``), so a target and a converter can share a name. Storage is read + through on every call, like scenario presets, because the folder can be shared between + processes. + """ + + def __init__(self, *, kind: RegistryRecipeKind, source: str | None = None) -> None: + """ + Initialize storage for one registry's folder. + + Args: + kind (RegistryRecipeKind): The registry whose recipes this storage reads. + source (str | None): Local directory or Azure Blob URI that holds the per-registry + folders. Defaults to ``registry_recipes`` under the PyRIT configuration directory. + + Raises: + ValueError: If the source has an unsupported URI scheme. + """ + root = source or str(self._get_default_storage_dir()) + super().__init__( + source=self._get_kind_source(root=root, kind=kind), + extension=".json", + source_label=f"Saved {kind.value}", + ) + self._kind = kind + + @staticmethod + def _get_default_storage_dir() -> Path: + """ + Get the default directory that holds the per-registry folders. + + Returns: + Path: Path to ``~/.pyrit/registry_recipes/``. + """ + # Deferred: importing pyrit.common.path triggers pyrit __init__.py + from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH + + return CONFIGURATION_DIRECTORY_PATH / "registry_recipes" + + @staticmethod + def _get_kind_source(*, root: str, kind: RegistryRecipeKind) -> str: + """ + Append the registry's folder to the configured source. + + For Azure Blob the folder goes on the URL path, so a SAS token in the query string is kept. + + Returns: + str: The local directory or Blob URI for the registry's folder. + """ + if is_azure_blob_uri(root): + parsed = urlparse(root) + return parsed._replace(path=f"{parsed.path.rstrip('/')}/{kind.value}").geturl() + return str(Path(root) / kind.value) + + def list_recipes(self) -> RecipeListing: + """ + Read every saved recipe of this kind and report the ones that can't be used. + + Returns: + RecipeListing: Parsed recipes keyed by name, plus the failures and ignored files. + + Raises: + OSError: If the local folder can't be listed. + AzureError: If the Blob container can't be listed. + """ + listing = self._list_documents_with_issues() + recipes: dict[str, StoredRecipe] = {} + failures = [ + self._failure(name=issue.name, reason=RestoreFailureReason.UNREADABLE_RECIPE, message=issue.message) + for issue in listing.unreadable + ] + for name, content in listing.documents.items(): + parsed = self._parse_recipe(name=name, content=content) + if isinstance(parsed, RestoreFailure): + failures.append(parsed) + else: + recipes[name] = parsed + ignored = tuple( + IgnoredRecipeFile(kind=self._kind, file_name=issue.file_name, message=issue.message) + for issue in listing.ignored + ) + return RecipeListing(recipes=recipes, failures=tuple(failures), ignored=ignored) + + def _parse_recipe(self, *, name: str, content: bytes) -> StoredRecipe | RestoreFailure: + """ + Parse one stored document, checking its schema version before anything else. + + Returns: + StoredRecipe | RestoreFailure: The recipe, or why it can't be used. + """ + version = self._compute_version(content) + invalid = RestoreFailureReason.INVALID_RECIPE + try: + payload = json.loads(content.decode("utf-8")) + except ValueError: + return self._failure(name=name, version=version, reason=invalid, message="Not valid UTF-8 JSON.") + except RecursionError: + return self._failure(name=name, version=version, reason=invalid, message="Nested too deeply to read.") + if not isinstance(payload, dict): + return self._failure(name=name, version=version, reason=invalid, message="Not a JSON object.") + schema_version = payload.get("schema_version") + current = ComponentRecipe.CURRENT_SCHEMA_VERSION + if isinstance(schema_version, int) and schema_version > current: + message = f"Saved by a newer PyRIT (format version {schema_version}; this one reads up to {current})." + return self._failure( + name=name, version=version, reason=RestoreFailureReason.NEWER_SCHEMA_VERSION, message=message + ) + model = TechniqueRecipe if self._kind is RegistryRecipeKind.TECHNIQUES else ComponentRecipe + try: + recipe = model.model_validate(payload) + except ValidationError as error: + return self._failure( + name=name, version=version, reason=invalid, message=self._describe_validation_error(error) + ) + except RecursionError: + return self._failure(name=name, version=version, reason=invalid, message="Nested too deeply to read.") + return StoredRecipe(ref=SavedEntryRef(kind=self._kind, name=name), recipe=recipe, version=version) + + def _failure( + self, *, name: str, reason: RestoreFailureReason, message: str, version: str | None = None + ) -> RestoreFailure: + """ + Build a failure for one entry in this registry's folder. + + Returns: + RestoreFailure: The failure. + """ + return RestoreFailure(kind=self._kind, name=name, reason=reason, message=message, version=version) + + @staticmethod + def _describe_validation_error(error: ValidationError) -> str: + """ + Summarize a validation error without echoing the stored values. + + Returns: + str: One line per problem, naming the field and what's wrong with it. + """ + problems = [ + f"{'.'.join(str(part) for part in issue['loc']) or 'document'}: {issue['msg']}" for issue in error.errors() + ] + return "Not a valid saved recipe: " + "; ".join(problems) diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py new file mode 100644 index 0000000000..58851e2b45 --- /dev/null +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -0,0 +1,150 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from azure.storage.blob import ContainerClient + +from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH +from pyrit.models.catalog.registry_recipe import ( + ComponentRecipe, + RegistryRecipeKind, + RestoreFailureReason, + TechniqueRecipe, +) +from pyrit.registry.registry_recipe_storage import RegistryRecipeStorage + +_TARGET = { + "schema_version": 1, + "class_module": "pyrit.prompt_target.openai.openai_chat_target", + "class_name": "OpenAIChatTarget", + "params": {"auth_mode": "identity", "endpoint": "https://contoso.openai.azure.com/openai/v1"}, +} + + +def _write(directory: Path, name: str, content: str | bytes) -> bytes: + directory.mkdir(parents=True, exist_ok=True) + data = content.encode("utf-8") if isinstance(content, str) else content + (directory / f"{name}.json").write_bytes(data) + return data + + +class TestRegistryRecipeStorageSource: + def test_local_source_uses_a_folder_per_registry(self, tmp_path: Path) -> None: + storage = RegistryRecipeStorage(kind=RegistryRecipeKind.CONVERTERS, source=str(tmp_path)) + + assert storage.display_source == str(tmp_path / "converters") + + def test_default_source_is_under_the_pyrit_configuration_directory(self) -> None: + storage = RegistryRecipeStorage(kind=RegistryRecipeKind.SCORERS) + + assert storage.display_source == str(CONFIGURATION_DIRECTORY_PATH / "registry_recipes" / "scorers") + + def test_blob_source_adds_the_folder_and_keeps_the_sas_token(self) -> None: + source = "https://account.blob.core.windows.net/recipes/team?sp=rl&sig=secret" + + kind_source = RegistryRecipeStorage._get_kind_source(root=source, kind=RegistryRecipeKind.TARGETS) + + assert kind_source == "https://account.blob.core.windows.net/recipes/team/targets?sp=rl&sig=secret" + + def test_blob_container_source_uses_the_folder_as_prefix(self) -> None: + client = MagicMock(spec=ContainerClient) + client.__enter__.return_value = client + client.list_blobs.return_value = [SimpleNamespace(name="techniques/formal.json")] + client.download_blob.return_value.readall.return_value = json.dumps({"type": "RedTeamingAttack"}).encode() + storage = RegistryRecipeStorage( + kind=RegistryRecipeKind.TECHNIQUES, source="https://account.blob.core.windows.net/recipes?sig=secret" + ) + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + listing = storage.list_recipes() + + client.list_blobs.assert_called_once_with(name_starts_with="techniques/") + assert list(listing.recipes) == ["formal"] + + +class TestRegistryRecipeStorageListing: + def test_reads_component_recipes_with_their_versions(self, tmp_path: Path) -> None: + content = _write(tmp_path / "targets", "contoso_gpt4o", json.dumps(_TARGET)) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + stored = listing.recipes["contoso_gpt4o"] + assert stored.ref.kind is RegistryRecipeKind.TARGETS + assert stored.ref.name == "contoso_gpt4o" + assert stored.recipe == ComponentRecipe.model_validate(_TARGET) + assert stored.version == hashlib.sha256(content).hexdigest() + assert listing.failures == () + + def test_reads_technique_recipes_as_technique_recipes(self, tmp_path: Path) -> None: + document = {"type": "RedTeamingAttack", "tags": ["custom"], "adversarial_chat": "contoso_gpt4o"} + _write(tmp_path / "techniques", "formal_red_team", json.dumps(document)) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TECHNIQUES, source=str(tmp_path)).list_recipes() + + assert listing.recipes["formal_red_team"].recipe == TechniqueRecipe.model_validate(document) + + def test_reports_documents_that_are_not_recipes(self, tmp_path: Path) -> None: + folder = tmp_path / "targets" + _write(folder, "not_json", "{oops") + _write(folder, "not_utf8", b"\xff\xfe") + _write(folder, "not_object", "[1, 2]") + _write(folder, "extra_field", json.dumps({**_TARGET, "credentials": {"api_key": "sk-secret-value"}})) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + assert listing.recipes == {} + assert {failure.name for failure in listing.failures} == {"not_json", "not_utf8", "not_object", "extra_field"} + assert all(failure.reason is RestoreFailureReason.INVALID_RECIPE for failure in listing.failures) + assert all(failure.version for failure in listing.failures) + extra = next(failure for failure in listing.failures if failure.name == "extra_field") + assert "credentials" in extra.message + assert "sk-secret-value" not in extra.message + + def test_reports_a_document_nested_too_deeply(self, tmp_path: Path) -> None: + _write(tmp_path / "targets", "deep", "[" * 100_000 + "]" * 100_000) + _write(tmp_path / "targets", "contoso_gpt4o", json.dumps(_TARGET)) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + (failure,) = listing.failures + assert (failure.name, failure.reason) == ("deep", RestoreFailureReason.INVALID_RECIPE) + assert list(listing.recipes) == ["contoso_gpt4o"] + + def test_reports_a_newer_schema_version_before_checking_fields(self, tmp_path: Path) -> None: + _write(tmp_path / "targets", "future", json.dumps({"schema_version": 2, "unknown_field": True})) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + (failure,) = listing.failures + assert failure.reason is RestoreFailureReason.NEWER_SCHEMA_VERSION + assert "newer PyRIT" in failure.message + + def test_reports_unreadable_and_ignored_files(self, tmp_path: Path) -> None: + folder = tmp_path / "scorers" + _write(folder, "locked", json.dumps(_TARGET)) + _write(folder, "Bad-Name", "{}") + original_read_bytes = Path.read_bytes + + def read_bytes(path: Path) -> bytes: + if path.name == "locked.json": + raise PermissionError(13, "Permission denied", str(path)) + return original_read_bytes(path) + + with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.SCORERS, source=str(tmp_path)).list_recipes() + + (failure,) = listing.failures + assert (failure.name, failure.reason, failure.message) == ( + "locked", + RestoreFailureReason.UNREADABLE_RECIPE, + "Permission denied", + ) + assert failure.version is None + (ignored,) = listing.ignored + assert (ignored.kind, ignored.file_name) == (RegistryRecipeKind.SCORERS, "Bad-Name.json") + assert "Bad-Name" in ignored.message From 9143ea64786ced1006c27e728d61e93c63839443 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 21:54:29 +0000 Subject: [PATCH 09/38] FEAT: Rebuild saved registry entries in dependency order --- pyrit/models/catalog/registry_recipe.py | 10 + pyrit/registry/registry_recipe_restore.py | 538 ++++++++++++++++++ .../registry/test_registry_recipe_restore.py | 483 ++++++++++++++++ 3 files changed, 1031 insertions(+) create mode 100644 pyrit/registry/registry_recipe_restore.py create mode 100644 tests/unit/registry/test_registry_recipe_restore.py diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index ce2e544588..a4e7349577 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -141,3 +141,13 @@ class IgnoredRecipeFile(BaseModel): kind: RegistryRecipeKind file_name: str message: str + + +class RegistryRestoreReport(BaseModel): + """What a restore rebuilt, what it couldn't, and which files it skipped.""" + + model_config = ConfigDict(frozen=True) + + restored: list[SavedEntryRef] = Field(default_factory=list) + failures: list[RestoreFailure] = Field(default_factory=list) + ignored: list[IgnoredRecipeFile] = Field(default_factory=list) diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py new file mode 100644 index 0000000000..0db95c36b7 --- /dev/null +++ b/pyrit/registry/registry_recipe_restore.py @@ -0,0 +1,538 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Rebuild saved registry entries in dependency order and report the ones that can't come back.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar +from urllib.parse import urlparse + +from pyrit.models.catalog.registry_recipe import ( + ComponentRecipe, + RegistryRecipeKind, + RegistryRestoreReport, + RestoreFailure, + RestoreFailureReason, + SavedEntryRef, + StoredRecipe, + TechniqueRecipe, +) +from pyrit.models.parameter import ComponentType +from pyrit.registry.registry_recipe_storage import RecipeListing, RegistryRecipeStorage + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from pyrit.models.parameter import Parameter + from pyrit.registry.registry import InstanceHoldingRegistry + + +@dataclass(frozen=True) +class _PlannedEntry: + """A saved entry whose class and references checked out, waiting to be built.""" + + stored: StoredRecipe + dependencies: tuple[SavedEntryRef, ...] + + +class RegistryRecipeRestorer: + """ + Rebuild saved targets, converters, scorers, and techniques into their registries. + + Each entry is built after the saved entries it names, with the same calls the API uses to + create one, so a restored entry passes the same checks. A file never chooses what gets + imported: classes and attack types are looked up among the ones already registered. Server + paths are refused as they are for API input, and a recipe never holds a value for a parameter + its class declares as a credential. Entries that can't come back are reported with the reason + instead of being skipped. + + Stored files are otherwise trusted the way configuration is. Rules for credentials a class + reads from server settings, and for fields that can carry a credential without being declared + one, need each class to declare them, so they arrive with saving and apply to restore too. + """ + + _RESERVED_TECHNIQUE_NAMES: ClassVar[frozenset[str]] = frozenset({"all", "default", "types"}) + _URL_PATTERN: ClassVar[re.Pattern[str]] = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*://[^\s'\"<>]+") + _KIND_BY_COMPONENT_TYPE: ClassVar[dict[ComponentType, RegistryRecipeKind]] = { + ComponentType.TARGET: RegistryRecipeKind.TARGETS, + ComponentType.CONVERTER: RegistryRecipeKind.CONVERTERS, + ComponentType.SCORER: RegistryRecipeKind.SCORERS, + } + _KIND_ORDER: ClassVar[tuple[RegistryRecipeKind, ...]] = tuple(RegistryRecipeKind) + + def __init__(self, *, source: str | None = None) -> None: + """ + Initialize the restorer. + + Args: + source (str | None): Local directory or Azure Blob URI holding the per-registry folders. + Defaults to ``registry_recipes`` under the PyRIT configuration directory. + """ + self._storages = {kind: RegistryRecipeStorage(kind=kind, source=source) for kind in RegistryRecipeKind} + + def restore(self) -> RegistryRestoreReport: + """ + Rebuild every saved entry that can be rebuilt and register it. + + Call it when the registries hold only startup entries, as they do at startup and after an + apply resets them; an entry restored earlier would otherwise clash with itself by name. If + any registry's folder can't be read, nothing is built, since entries in the other folders + may name entries in the unread one. + + Returns: + RegistryRestoreReport: What was restored, what failed and why, and which files were ignored. + """ + listings, failures = self._read_all() + ignored = [file for listing in listings for file in listing.ignored] + failures.extend(failure for listing in listings for failure in listing.failures) + if len(listings) < len(self._storages): + return RegistryRestoreReport(failures=self._sorted(failures), ignored=ignored) + stored = {recipe.ref: recipe for listing in listings for recipe in listing.recipes.values()} + planned = self._plan_all(stored=stored, failures=failures) + restored = self._build_in_order(planned=planned, failures=failures) + return RegistryRestoreReport(restored=restored, failures=self._sorted(failures), ignored=ignored) + + def _read_all(self) -> tuple[list[RecipeListing], list[RestoreFailure]]: + """ + Read every registry's folder, turning a folder that can't be listed into a failure. + + Returns: + tuple[list[RecipeListing], list[RestoreFailure]]: The folders that were read, and one + failure per folder that wasn't. + """ + from azure.core.exceptions import AzureError + + listings: list[RecipeListing] = [] + failures: list[RestoreFailure] = [] + for kind, storage in self._storages.items(): + try: + listings.append(storage.list_recipes()) + except (OSError, AzureError) as error: + message = f"The saved {kind.value} couldn't be listed: {storage._describe_read_error(error)}" + failures.append( + RestoreFailure(kind=kind, reason=RestoreFailureReason.UNREADABLE_STORAGE, message=message) + ) + return listings, failures + + def _plan_all( + self, *, stored: Mapping[SavedEntryRef, StoredRecipe], failures: list[RestoreFailure] + ) -> dict[SavedEntryRef, _PlannedEntry]: + """ + Check every readable recipe's class, credentials, name, and references. + + A reference resolves to a saved entry first, even one that failed to load, so an entry is + never silently attached to a startup entry that only shares its name. + + Returns: + dict[SavedEntryRef, _PlannedEntry]: The entries ready to build. Failures are appended. + """ + saved_refs = set(stored) | { + SavedEntryRef(kind=failure.kind, name=failure.name) for failure in failures if failure.name is not None + } + planned: dict[SavedEntryRef, _PlannedEntry] = {} + for ref in sorted(stored, key=self._sort_key): + outcome = self._plan(stored=stored[ref], saved_refs=saved_refs) + if isinstance(outcome, RestoreFailure): + failures.append(outcome) + else: + planned[ref] = outcome + return planned + + def _plan(self, *, stored: StoredRecipe, saved_refs: set[SavedEntryRef]) -> _PlannedEntry | RestoreFailure: + """ + Check one recipe and collect the saved entries it depends on. + + Returns: + _PlannedEntry | RestoreFailure: The entry ready to build, or why it can't be. + """ + recipe = stored.recipe + try: + references = ( + self._technique_references(ref=stored.ref, recipe=recipe) + if isinstance(recipe, TechniqueRecipe) + else self._component_references(ref=stored.ref, recipe=recipe) + ) + except _RecipeRejectedError as rejection: + return self._failure(stored=stored, reason=rejection.reason, message=str(rejection)) + except Exception as error: # Planning one entry must never stop the others from being checked. + return self._failure( + stored=stored, reason=RestoreFailureReason.INVALID_RECIPE, message=self._safe_message(error) + ) + dependencies: list[SavedEntryRef] = [] + for reference in references: + if reference in saved_refs: + dependencies.append(reference) + elif reference.name not in self._get_registry(reference.kind).instances: + message = f"Uses {self._describe_ref(reference)}, which doesn't exist." + return self._failure( + stored=stored, + reason=RestoreFailureReason.MISSING_DEPENDENCY, + message=message, + dependency=reference, + ) + return _PlannedEntry(stored=stored, dependencies=tuple(dict.fromkeys(dependencies))) + + def _component_references(self, *, ref: SavedEntryRef, recipe: ComponentRecipe) -> list[SavedEntryRef]: + """ + Check a target, converter, or scorer recipe and list the entries its parameters name. + + Returns: + list[SavedEntryRef]: The referenced entries, in parameter order. + + Raises: + _RecipeRejectedError: If the class is unknown, moved, or needs a server-chosen upload + folder, the recipe holds a credential, or the name is already taken. + """ + registry = self._get_registry(ref.kind) + try: + cls = registry.get_class(recipe.class_name) + except KeyError as error: + message = f"'{recipe.class_name}' is not a registered {self._singular(ref.kind)} class." + raise _RecipeRejectedError(reason=RestoreFailureReason.UNKNOWN_CLASS, message=message) from error + if cls.__module__ != recipe.class_module: + message = ( + f"'{recipe.class_name}' is registered from '{cls.__module__}', " + f"but the saved entry was made with the one in '{recipe.class_module}'." + ) + raise _RecipeRejectedError(reason=RestoreFailureReason.MODULE_MISMATCH, message=message) + if getattr(cls, "upload_directory_parameter", None) is not None: + message = ( + f"'{recipe.class_name}' uploads files from a folder the server chooses, which restore can't set yet." + ) + raise _RecipeRejectedError(reason=RestoreFailureReason.UNSUPPORTED_CLASS, message=message) + parameters = registry.get_class_metadata(cls).parameters + self._check_stored_credentials(parameters=parameters, params=recipe.params) + try: + registry.instances.validate_name_available(ref.name) + except ValueError as error: + raise _RecipeRejectedError(reason=RestoreFailureReason.NAME_CLASH, message=str(error)) from error + return self._parameter_references(parameters=parameters, params=recipe.params) + + def _technique_references(self, *, ref: SavedEntryRef, recipe: TechniqueRecipe) -> list[SavedEntryRef]: + """ + Check a technique recipe's name and attack type, and list the entries it names. + + Returns: + list[SavedEntryRef]: The adversarial target, converters, and attack parameter references. + + Raises: + _RecipeRejectedError: If the name is reserved or the attack type isn't registered. + """ + from pyrit.registry.components.attack_registry import AttackRegistry + + if ref.name.casefold() in self._RESERVED_TECHNIQUE_NAMES: + message = f"'{ref.name}' is a reserved technique name." + raise _RecipeRejectedError(reason=RestoreFailureReason.INVALID_RECIPE, message=message) + attack_registry = AttackRegistry.get_registry_singleton() + try: + attack_class = attack_registry.get_class(recipe.type) + except KeyError as error: + message = f"'{recipe.type}' is not a registered attack type." + raise _RecipeRejectedError(reason=RestoreFailureReason.UNKNOWN_CLASS, message=message) from error + converters = [*(recipe.request_converters or []), *(recipe.response_converters or [])] + references = [SavedEntryRef(kind=RegistryRecipeKind.CONVERTERS, name=name) for name in converters] + if recipe.adversarial_chat is not None: + references.append(SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name=recipe.adversarial_chat)) + parameters = attack_registry.get_class_metadata(attack_class).parameters + return references + self._parameter_references(parameters=parameters, params=recipe.params) + + @staticmethod + def _check_stored_credentials(*, parameters: Sequence[Parameter], params: Mapping[str, Any]) -> None: + """ + Refuse a recipe that stores a value for a parameter its class declares as a credential. + + Raises: + _RecipeRejectedError: If a declared credential parameter has a value. + """ + stored = sorted( + parameter.name + for parameter in parameters + if parameter.sensitive and params.get(parameter.name) not in (None, "") + ) + if stored: + raise _RecipeRejectedError( + reason=RestoreFailureReason.UNSUPPORTED_CREDENTIAL, + message=f"Saved entries can't hold credentials ('{stored[0]}').", + ) + + def _parameter_references( + self, *, parameters: Sequence[Parameter], params: Mapping[str, Any] + ) -> list[SavedEntryRef]: + """ + List the entries that supplied reference parameters name. + + An empty name leaves the reference unset, as it does when the API builds an entry. + + Returns: + list[SavedEntryRef]: One reference per named entry of a kind that can be saved. + """ + references: list[SavedEntryRef] = [] + for parameter in parameters: + if parameter.reference is None or params.get(parameter.name) in (None, ""): + continue + kind = self._KIND_BY_COMPONENT_TYPE.get(parameter.reference.component_type) + value = params[parameter.name] + names = value if isinstance(value, list) else [value] + if kind is not None: + references.extend(SavedEntryRef(kind=kind, name=name) for name in names if isinstance(name, str)) + return references + + def _build_in_order( + self, *, planned: dict[SavedEntryRef, _PlannedEntry], failures: list[RestoreFailure] + ) -> list[SavedEntryRef]: + """ + Build entries once everything they depend on has been settled. + + Returns: + list[SavedEntryRef]: The entries that were rebuilt and registered. Failures are appended. + """ + failed = {SavedEntryRef(kind=failure.kind, name=failure.name) for failure in failures if failure.name} + restored: list[SavedEntryRef] = [] + pending = dict(planned) + ready = self._ready(pending=pending, settled={*failed}) + while ready: + for ref in ready: + entry = pending.pop(ref) + failure = self._failed_dependency(entry=entry, failed=failed) or self._build(stored=entry.stored) + if failure is None: + restored.append(ref) + else: + failures.append(failure) + failed.add(ref) + ready = self._ready(pending=pending, settled={*failed, *restored}) + failures.extend(self._unresolved(pending=pending)) + return restored + + def _ready( + self, *, pending: Mapping[SavedEntryRef, _PlannedEntry], settled: set[SavedEntryRef] + ) -> list[SavedEntryRef]: + """ + List the pending entries whose dependencies have all been built or have failed. + + Returns: + list[SavedEntryRef]: The entries to handle next, in a stable order. + """ + ready = [ref for ref, entry in pending.items() if settled.issuperset(entry.dependencies)] + return sorted(ready, key=self._sort_key) + + def _failed_dependency(self, *, entry: _PlannedEntry, failed: set[SavedEntryRef]) -> RestoreFailure | None: + """ + Fail an entry if a saved entry it depends on couldn't be restored. + + Returns: + RestoreFailure | None: The failure, or None if every dependency was restored. + """ + dependency = next((ref for ref in entry.dependencies if ref in failed), None) + if dependency is None: + return None + return self._failure( + stored=entry.stored, + reason=RestoreFailureReason.FAILED_DEPENDENCY, + message=f"Uses {self._describe_ref(dependency)}, which couldn't be restored.", + dependency=dependency, + ) + + def _unresolved(self, *, pending: Mapping[SavedEntryRef, _PlannedEntry]) -> list[RestoreFailure]: + """ + Fail the entries left waiting: those on a dependency loop, and those waiting on one. + + Returns: + list[RestoreFailure]: One failure per entry that was never built. + """ + failures: list[RestoreFailure] = [] + for ref in sorted(pending, key=self._sort_key): + entry = pending[ref] + waiting_on = [dependency for dependency in entry.dependencies if dependency in pending] + on_loop = [ + dependency + for dependency in waiting_on + if dependency == ref or self._reaches(start=dependency, target=ref, pending=pending) + ] + blocker = (on_loop or waiting_on)[0] + if on_loop: + reason = RestoreFailureReason.DEPENDENCY_LOOP + message = f"Saved entries use each other in a loop, through {self._describe_ref(blocker)}." + else: + reason = RestoreFailureReason.FAILED_DEPENDENCY + message = f"Uses {self._describe_ref(blocker)}, which is waiting on a loop of saved entries." + failures.append(self._failure(stored=entry.stored, reason=reason, message=message, dependency=blocker)) + return failures + + @staticmethod + def _reaches( + *, start: SavedEntryRef, target: SavedEntryRef, pending: Mapping[SavedEntryRef, _PlannedEntry] + ) -> bool: + """ + Return whether following pending dependencies from ``start`` leads to ``target``. + + Returns: + bool: True if ``target`` can be reached. + """ + stack = [dependency for dependency in pending[start].dependencies if dependency in pending] + seen: set[SavedEntryRef] = set() + while stack: + current = stack.pop() + if current == target: + return True + if current in seen: + continue + seen.add(current) + stack.extend(dependency for dependency in pending[current].dependencies if dependency in pending) + return False + + def _build(self, *, stored: StoredRecipe) -> RestoreFailure | None: + """ + Build one entry and register it. + + Returns: + RestoreFailure | None: Why the entry couldn't be built or registered, or None on success. + """ + recipe = stored.recipe + try: + if isinstance(recipe, TechniqueRecipe): + self._build_technique(name=stored.ref.name, recipe=recipe) + else: + self._get_registry(stored.ref.kind).create_named_instance( + name=stored.ref.name, type_name=recipe.class_name, params=recipe.params, external_input=True + ) + except _RecipeRejectedError as rejection: + return self._failure(stored=stored, reason=rejection.reason, message=str(rejection)) + except Exception as error: # A constructor can raise anything; one bad entry must not stop the rest. + return self._failure( + stored=stored, reason=RestoreFailureReason.BUILD_ERROR, message=self._safe_message(error) + ) + return None + + @staticmethod + def _build_technique(*, name: str, recipe: TechniqueRecipe) -> None: + """ + Build a technique with the call the technique API uses, then register it. + + Raises: + _RecipeRejectedError: If the name or a tag conflicts with a registered technique. + """ + from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry + + registry = AttackTechniqueRegistry.get_registry_singleton() + args = recipe.model_dump(exclude_unset=True, exclude={"schema_version", "type", "tags"}) + factory = registry.create_factory( + name=name, attack_type=recipe.type, technique_tags=recipe.tags, strict_json=True, **args + ) + try: + registry.instances.register_runtime(factory) + except ValueError as error: + raise _RecipeRejectedError(reason=RestoreFailureReason.NAME_CLASH, message=str(error)) from error + + def _get_registry(self, kind: RegistryRecipeKind) -> InstanceHoldingRegistry[Any, Any]: + """ + Get the registry singleton that holds saved targets, converters, or scorers. + + Returns: + InstanceHoldingRegistry[Any, Any]: The registry. + """ + from pyrit.registry.components.converter_registry import ConverterRegistry + from pyrit.registry.components.scorer_registry import ScorerRegistry + from pyrit.registry.components.target_registry import TargetRegistry + + registries = { + RegistryRecipeKind.TARGETS: TargetRegistry, + RegistryRecipeKind.CONVERTERS: ConverterRegistry, + RegistryRecipeKind.SCORERS: ScorerRegistry, + } + return registries[kind].get_registry_singleton() + + def _failure( + self, + *, + stored: StoredRecipe, + reason: RestoreFailureReason, + message: str, + dependency: SavedEntryRef | None = None, + ) -> RestoreFailure: + """ + Build a failure for one saved entry. + + Returns: + RestoreFailure: The failure, carrying the stored document's version. + """ + return RestoreFailure( + kind=stored.ref.kind, + name=stored.ref.name, + reason=reason, + message=message, + version=stored.version, + dependency=dependency, + ) + + def _sort_key(self, ref: SavedEntryRef) -> tuple[int, str]: + """ + Order entries by registry, then by name, so restores are repeatable. + + Returns: + tuple[int, str]: The sort key. + """ + return self._KIND_ORDER.index(ref.kind), ref.name + + def _sorted(self, failures: list[RestoreFailure]) -> list[RestoreFailure]: + """ + Order failures by registry, then by name, with whole-folder failures first. + + Returns: + list[RestoreFailure]: The sorted failures. + """ + return sorted(failures, key=lambda failure: (self._KIND_ORDER.index(failure.kind), failure.name or "")) + + def _describe_ref(self, ref: SavedEntryRef) -> str: + """ + Describe an entry for a message, such as "target 'contoso_gpt4o'". + + Returns: + str: The description. + """ + return f"{self._singular(ref.kind)} '{ref.name}'" + + @staticmethod + def _singular(kind: RegistryRecipeKind) -> str: + """ + Name one entry of a kind, such as "target". + + Returns: + str: The singular name. + """ + return kind.value.removesuffix("s") + + @classmethod + def _safe_message(cls, error: Exception) -> str: + """ + Describe an error for the failure report, removing credentials from any URL it mentions. + + Constructor errors can quote a server setting, such as an endpoint with a SAS token. + + Returns: + str: The error's type and text, with every URL reduced to its scheme, host, and path. + """ + text = cls._URL_PATTERN.sub(lambda match: cls._redact_url(match.group(0)), str(error)) + return f"{type(error).__name__}: {text}" if text else type(error).__name__ + + @staticmethod + def _redact_url(url: str) -> str: + """ + Drop the user name, password, query, and fragment from a URL. + + Returns: + str: The URL without anything that can carry a credential. + """ + parsed = urlparse(url) + return parsed._replace(netloc=parsed.netloc.rpartition("@")[2], query="", fragment="").geturl() + + +class _RecipeRejectedError(Exception): + """A saved recipe failed a check before or while it was built.""" + + def __init__(self, *, reason: RestoreFailureReason, message: str) -> None: + """Initialize the error with the reason to report.""" + super().__init__(message) + self.reason = reason diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py new file mode 100644 index 0000000000..b05160c4e0 --- /dev/null +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -0,0 +1,483 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import json +import os +import sys +from collections.abc import Iterator +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import pytest +from pydantic import ValidationError + +from pyrit.backend.models.techniques import CreateTechniqueRequest +from pyrit.models.catalog.registry_recipe import ( + RegistryRecipeKind, + RestoreFailure, + RestoreFailureReason, + SavedEntryRef, + TechniqueRecipe, +) +from pyrit.prompt_target import OpenAIChatTarget, RoundRobinTarget +from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry +from pyrit.registry.components.converter_registry import ConverterRegistry +from pyrit.registry.components.scorer_registry import ScorerRegistry +from pyrit.registry.components.target_registry import TargetRegistry +from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer +from pyrit.registry.registry_recipe_storage import RegistryRecipeStorage + +_OPENAI_MODULE = "pyrit.prompt_target.openai.openai_chat_target" +_IDENTITY_TARGET = { + "class_module": _OPENAI_MODULE, + "class_name": "OpenAIChatTarget", + "params": { + "auth_mode": "identity", + "endpoint": "https://contoso.openai.azure.com/openai/v1", + "model_name": "gpt-4o", + }, +} +_ROUND_ROBIN_MODULE = "pyrit.prompt_target.round_robin_target" +_TECHNIQUE = { + "type": "RedTeamingAttack", + "params": {"max_turns": 5}, + "description": "Red teaming with a formal tone", + "tags": ["custom"], + "adversarial_chat": "contoso_gpt4o", + "adversarial_system_prompt": "You are a red teaming assistant.", + "request_converters": ["tone_formal"], + "response_converters": [], +} + +_REGISTRIES = (TargetRegistry, ConverterRegistry, ScorerRegistry, AttackTechniqueRegistry) + + +@pytest.fixture(autouse=True) +def isolated_registries() -> Iterator[None]: + for registry in _REGISTRIES: + registry.reset_registry_singleton() + yield + for registry in _REGISTRIES: + registry.reset_registry_singleton() + + +def _save(root: Path, kind: RegistryRecipeKind, name: str, document: dict[str, Any]) -> None: + _save_text(root, kind, name, json.dumps(document)) + + +def _save_text(root: Path, kind: RegistryRecipeKind, name: str, text: str) -> None: + folder = root / kind.value + folder.mkdir(parents=True, exist_ok=True) + (folder / f"{name}.json").write_text(text) + + +def _round_robin(*targets: str) -> dict[str, Any]: + return {"class_module": _ROUND_ROBIN_MODULE, "class_name": "RoundRobinTarget", "params": {"targets": list(targets)}} + + +def _failures_by_name(failures: list[RestoreFailure]) -> dict[str | None, RestoreFailure]: + return {failure.name: failure for failure in failures} + + +@pytest.mark.usefixtures("patch_central_database") +class TestRestoreOrderAndRebuild: + def test_restores_every_kind_after_the_entries_it_uses(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "formal_red_team", _TECHNIQUE) + _save( + tmp_path, + RegistryRecipeKind.SCORERS, + "refusal", + { + "class_module": "pyrit.score.true_false.self_ask_refusal_scorer", + "class_name": "SelfAskRefusalScorer", + "params": {"chat_target": "contoso_gpt4o"}, + }, + ) + _save( + tmp_path, + RegistryRecipeKind.CONVERTERS, + "tone_formal", + { + "class_module": "pyrit.converter.tone_converter", + "class_name": "ToneConverter", + "params": {"converter_target": "contoso_gpt4o", "tone": "formal"}, + }, + ) + _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + assert report.failures == [] + assert [(ref.kind, ref.name) for ref in report.restored] == [ + (RegistryRecipeKind.TARGETS, "contoso_gpt4o"), + (RegistryRecipeKind.CONVERTERS, "tone_formal"), + (RegistryRecipeKind.SCORERS, "refusal"), + (RegistryRecipeKind.TECHNIQUES, "formal_red_team"), + ] + target = TargetRegistry.get_registry_singleton().instances.get("contoso_gpt4o") + assert isinstance(target, OpenAIChatTarget) + converter = ConverterRegistry.get_registry_singleton().instances.get("tone_formal") + assert converter is not None + assert converter._converter_target is target + factory = AttackTechniqueRegistry.get_registry_singleton().instances.get("formal_red_team") + assert factory is not None + assert factory.technique_tags == ["custom"] + + def test_orders_entries_of_the_same_kind(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "a_pool", _round_robin("b_member", "c_member")) + _save(tmp_path, RegistryRecipeKind.TARGETS, "b_member", _IDENTITY_TARGET) + _save(tmp_path, RegistryRecipeKind.TARGETS, "c_member", _IDENTITY_TARGET) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + assert report.failures == [] + assert [ref.name for ref in report.restored] == ["b_member", "c_member", "a_pool"] + assert isinstance(TargetRegistry.get_registry_singleton().instances.get("a_pool"), RoundRobinTarget) + + def test_technique_recipe_matches_the_create_request(self) -> None: + recipe_fields = set(TechniqueRecipe.model_fields) - {"schema_version"} + request_fields = set(CreateTechniqueRequest.model_fields) - {"name"} + request = CreateTechniqueRequest.model_validate({"name": "formal_red_team", **_TECHNIQUE}) + recipe = TechniqueRecipe.model_validate(_TECHNIQUE) + + assert recipe_fields == request_fields + assert recipe.model_dump(exclude_unset=True, exclude={"schema_version", "type", "tags"}) == ( + request.model_dump(exclude_unset=True, exclude={"name", "type", "tags"}) + ) + + @pytest.mark.parametrize( + "changes", + [ + {"tags": ["all"]}, + {"tags": ["Custom", "custom"]}, + {"tags": ["bad-tag"]}, + {"type": ""}, + {"params": {"x" * 300: 1}}, + {"request_converters": ["c"] * 101}, + {"params": {"temperature": float("nan")}}, + {"unexpected": True}, + ], + ) + def test_technique_recipe_rejects_what_the_create_request_rejects(self, changes: dict[str, Any]) -> None: + document = {**_TECHNIQUE, **changes} + + with pytest.raises(ValidationError): + CreateTechniqueRequest.model_validate({"name": "formal_red_team", **document}) + with pytest.raises(ValidationError): + TechniqueRecipe.model_validate(document) + + +@pytest.mark.usefixtures("patch_central_database") +class TestRestoreFailures: + def test_a_technique_waits_for_its_saved_adversarial_target(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", {**_IDENTITY_TARGET, "class_name": "NoSuchTarget"}) + without_converters = {key: value for key, value in _TECHNIQUE.items() if key != "request_converters"} + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "formal_red_team", without_converters) + startup_target = OpenAIChatTarget( + endpoint="https://startup.openai.azure.com/openai/v1", model_name="gpt-4o", auth_mode="identity" + ) + TargetRegistry.get_registry_singleton().instances.register(startup_target, name="contoso_gpt4o") + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + failure = _failures_by_name(report.failures)["formal_red_team"] + assert failure.reason is RestoreFailureReason.FAILED_DEPENDENCY + assert failure.dependency == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="contoso_gpt4o") + + def test_technique_parameters_are_checked_as_strict_json(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) + without_converters = {key: value for key, value in _TECHNIQUE.items() if key != "request_converters"} + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "loose", {**without_converters, "params": {"max_turns": "5"}}) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert (failure.name, failure.reason) == ("loose", RestoreFailureReason.BUILD_ERROR) + + def test_reports_dependency_loops_and_entries_waiting_on_them(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "loop_a", _round_robin("loop_b")) + _save(tmp_path, RegistryRecipeKind.TARGETS, "loop_b", _round_robin("loop_a")) + _save(tmp_path, RegistryRecipeKind.TARGETS, "waiting", _round_robin("loop_a")) + _save(tmp_path, RegistryRecipeKind.TARGETS, "self_loop", _round_robin("self_loop")) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + failures = _failures_by_name(report.failures) + assert report.restored == [] + assert failures["loop_a"].reason is RestoreFailureReason.DEPENDENCY_LOOP + assert failures["loop_b"].reason is RestoreFailureReason.DEPENDENCY_LOOP + assert failures["self_loop"].reason is RestoreFailureReason.DEPENDENCY_LOOP + assert failures["waiting"].reason is RestoreFailureReason.FAILED_DEPENDENCY + assert failures["waiting"].dependency == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="loop_a") + + def test_names_a_loop_member_as_the_reason_for_a_loop(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "a", _round_robin("c", "b")) + _save(tmp_path, RegistryRecipeKind.TARGETS, "b", _round_robin("a")) + _save(tmp_path, RegistryRecipeKind.TARGETS, "c", _round_robin("d")) + _save(tmp_path, RegistryRecipeKind.TARGETS, "d", _round_robin("c")) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + failures = _failures_by_name(report.failures) + assert failures["a"].reason is RestoreFailureReason.DEPENDENCY_LOOP + assert failures["a"].dependency == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="b") + + def test_reports_missing_and_failed_dependencies(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "broken", {**_IDENTITY_TARGET, "class_name": "NoSuchTarget"}) + _save(tmp_path, RegistryRecipeKind.TARGETS, "uses_broken", _round_robin("broken")) + _save(tmp_path, RegistryRecipeKind.TARGETS, "uses_ghost", _round_robin("ghost")) + without_adversarial = {key: value for key, value in _TECHNIQUE.items() if key != "adversarial_chat"} + _save( + tmp_path, + RegistryRecipeKind.TECHNIQUES, + "uses_unparsable", + {**without_adversarial, "request_converters": ["bad"]}, + ) + _save_text(tmp_path, RegistryRecipeKind.CONVERTERS, "bad", "{not json") + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + failures = _failures_by_name(report.failures) + assert failures["broken"].reason is RestoreFailureReason.UNKNOWN_CLASS + assert failures["uses_broken"].reason is RestoreFailureReason.FAILED_DEPENDENCY + assert failures["uses_ghost"].reason is RestoreFailureReason.MISSING_DEPENDENCY + assert failures["uses_ghost"].dependency == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost") + assert failures["bad"].reason is RestoreFailureReason.INVALID_RECIPE + assert failures["uses_unparsable"].reason is RestoreFailureReason.FAILED_DEPENDENCY + assert failures["uses_unparsable"].dependency == SavedEntryRef(kind=RegistryRecipeKind.CONVERTERS, name="bad") + + def test_never_imports_a_module_named_by_a_file(self, tmp_path: Path) -> None: + module = "tests_registry_restore_never_imported" + _save(tmp_path, RegistryRecipeKind.TARGETS, "moved", {**_IDENTITY_TARGET, "class_module": module}) + _save(tmp_path, RegistryRecipeKind.TARGETS, "unknown", {"class_module": "os", "class_name": "system"}) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + failures = _failures_by_name(report.failures) + assert failures["moved"].reason is RestoreFailureReason.MODULE_MISMATCH + assert _OPENAI_MODULE in failures["moved"].message + assert failures["unknown"].reason is RestoreFailureReason.UNKNOWN_CLASS + assert module not in sys.modules + + def test_refuses_a_stored_credential(self, tmp_path: Path) -> None: + stored_key = {**_IDENTITY_TARGET, "params": {**_IDENTITY_TARGET["params"], "api_key": "sk-secret-value"}} + _save(tmp_path, RegistryRecipeKind.TARGETS, "stored_key", stored_key) + _save(tmp_path, RegistryRecipeKind.TARGETS, "good", _IDENTITY_TARGET) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert (failure.name, failure.reason) == ("stored_key", RestoreFailureReason.UNSUPPORTED_CREDENTIAL) + assert "sk-secret-value" not in failure.message + assert [ref.name for ref in report.restored] == ["good"] + + def test_refuses_a_class_that_needs_the_server_upload_folder(self, tmp_path: Path) -> None: + _save( + tmp_path, + RegistryRecipeKind.TARGETS, + "uploader", + { + "class_module": "pyrit.prompt_target.http_target.httpx_api_target", + "class_name": "HTTPXAPITarget", + "params": {"http_url": "https://example.test/upload", "file_path": "/etc/hosts"}, + }, + ) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert failure.reason is RestoreFailureReason.UNSUPPORTED_CLASS + assert "uploader" not in TargetRegistry.get_registry_singleton().instances + + @pytest.mark.parametrize("auth_mode", [[], {}, ["identity"]]) + def test_reports_malformed_values_without_stopping_the_restore(self, tmp_path: Path, auth_mode: object) -> None: + malformed = {**_IDENTITY_TARGET, "params": {**_IDENTITY_TARGET["params"], "auth_mode": auth_mode}} + _save(tmp_path, RegistryRecipeKind.TARGETS, "malformed", malformed) + _save(tmp_path, RegistryRecipeKind.TARGETS, "good", _IDENTITY_TARGET) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + assert [failure.name for failure in report.failures] == ["malformed"] + assert [ref.name for ref in report.restored] == ["good"] + + def test_a_planning_error_fails_only_that_entry(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "a_first", _IDENTITY_TARGET) + _save(tmp_path, RegistryRecipeKind.TARGETS, "b_second", _IDENTITY_TARGET) + + with patch.object( + RegistryRecipeRestorer, "_parameter_references", side_effect=[RuntimeError("unexpected"), []] + ): + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert (failure.name, failure.reason) == ("a_first", RestoreFailureReason.INVALID_RECIPE) + assert failure.message == "RuntimeError: unexpected" + assert [ref.name for ref in report.restored] == ["b_second"] + + def test_build_errors_drop_credentials_from_urls(self, tmp_path: Path) -> None: + endpoint = "https://user:hunter2@evil.example/openai/v1?sig=SYNTHETIC-SAS#frag" + no_endpoint = {**_IDENTITY_TARGET, "params": {"auth_mode": "identity", "model_name": "gpt-4o"}} + _save(tmp_path, RegistryRecipeKind.TARGETS, "server_endpoint", no_endpoint) + + with patch.dict(os.environ, {"OPENAI_CHAT_ENDPOINT": endpoint}): + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert failure.reason is RestoreFailureReason.BUILD_ERROR + assert "https://evil.example/openai/v1" in failure.message + for secret in ("hunter2", "SYNTHETIC-SAS", "frag"): + assert secret not in failure.message + + def test_an_empty_reference_is_left_unset(self, tmp_path: Path) -> None: + _save( + tmp_path, + RegistryRecipeKind.CONVERTERS, + "puzzled", + { + "class_module": "pyrit.converter.puzzled.puzzled_converter", + "class_name": "PuzzledConverter", + "params": {"converter_target": ""}, + }, + ) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + assert report.failures == [] + assert [ref.name for ref in report.restored] == ["puzzled"] + + def test_follows_converter_and_scorer_references(self, tmp_path: Path) -> None: + _save( + tmp_path, + RegistryRecipeKind.CONVERTERS, + "selective", + { + "class_module": "pyrit.converter.selective_text_converter", + "class_name": "SelectiveTextConverter", + "params": {"sub_converter": "ghost_converter"}, + }, + ) + _save( + tmp_path, + RegistryRecipeKind.SCORERS, + "composite", + { + "class_module": "pyrit.score.true_false.true_false_composite_scorer", + "class_name": "TrueFalseCompositeScorer", + "params": {"scorers": ["ghost_scorer"]}, + }, + ) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + failures = _failures_by_name(report.failures) + assert failures["selective"].dependency == SavedEntryRef( + kind=RegistryRecipeKind.CONVERTERS, name="ghost_converter" + ) + assert failures["composite"].dependency == SavedEntryRef(kind=RegistryRecipeKind.SCORERS, name="ghost_scorer") + assert {failure.reason for failure in report.failures} == {RestoreFailureReason.MISSING_DEPENDENCY} + + def test_refuses_reserved_technique_names_and_tags(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "types", {"type": "RedTeamingAttack"}) + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "tagged_all", {"type": "RedTeamingAttack", "tags": ["all"]}) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + failures = _failures_by_name(report.failures) + assert failures["types"].reason is RestoreFailureReason.INVALID_RECIPE + assert failures["tagged_all"].reason is RestoreFailureReason.INVALID_RECIPE + assert report.restored == [] + + def test_rejects_server_paths_the_way_the_api_does(self, tmp_path: Path) -> None: + _save( + tmp_path, + RegistryRecipeKind.CONVERTERS, + "wordswap", + { + "class_module": "pyrit.converter.colloquial_wordswap_converter", + "class_name": "ColloquialWordswapConverter", + "params": {"wordswap_path": "/etc/passwd"}, + }, + ) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert failure.reason is RestoreFailureReason.BUILD_ERROR + assert "wordswap_path" in failure.message + assert "wordswap" not in ConverterRegistry.get_registry_singleton().instances + + def test_reports_a_name_taken_at_startup_and_fails_its_dependents(self, tmp_path: Path) -> None: + startup_target = OpenAIChatTarget( + endpoint="https://startup.openai.azure.com/openai/v1", model_name="gpt-4o", auth_mode="identity" + ) + TargetRegistry.get_registry_singleton().instances.register(startup_target, name="contoso_gpt4o") + _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) + _save(tmp_path, RegistryRecipeKind.TARGETS, "pool", _round_robin("contoso_gpt4o")) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + failures = _failures_by_name(report.failures) + assert failures["contoso_gpt4o"].reason is RestoreFailureReason.NAME_CLASH + assert failures["pool"].reason is RestoreFailureReason.FAILED_DEPENDENCY + assert TargetRegistry.get_registry_singleton().instances.get("contoso_gpt4o") is startup_target + + def test_reports_a_technique_name_taken_in_another_case(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) + _save( + tmp_path, + RegistryRecipeKind.CONVERTERS, + "tone_formal", + { + "class_module": "pyrit.converter.tone_converter", + "class_name": "ToneConverter", + "params": {"converter_target": "contoso_gpt4o", "tone": "formal"}, + }, + ) + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "formal_red_team", _TECHNIQUE) + restorer = RegistryRecipeRestorer(source=str(tmp_path)) + assert restorer.restore().failures == [] + for registry in (TargetRegistry, ConverterRegistry): + registry.reset_registry_singleton() + (tmp_path / "techniques" / "formal_red_team.json").rename(tmp_path / "techniques" / "formal_red_team_copy.json") + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "custom", {**_TECHNIQUE, "tags": []}) + + report = restorer.restore() + + failures = _failures_by_name(report.failures) + assert failures["custom"].reason is RestoreFailureReason.NAME_CLASH + assert "custom" in failures["custom"].message + + def test_builds_nothing_when_a_folder_cannot_be_listed(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) + original_list_recipes = RegistryRecipeStorage.list_recipes + + def list_recipes(storage: RegistryRecipeStorage) -> Any: + if storage._kind is RegistryRecipeKind.CONVERTERS: + raise PermissionError(13, "Permission denied", str(tmp_path / "converters")) + return original_list_recipes(storage) + + with patch.object(RegistryRecipeStorage, "list_recipes", autospec=True, side_effect=list_recipes): + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert (failure.kind, failure.name, failure.reason) == ( + RegistryRecipeKind.CONVERTERS, + None, + RestoreFailureReason.UNREADABLE_STORAGE, + ) + assert str(tmp_path) not in failure.message + assert report.restored == [] + assert "contoso_gpt4o" not in TargetRegistry.get_registry_singleton().instances + + def test_reports_constructor_errors_and_keeps_going(self, tmp_path: Path) -> None: + bad_value = {**_IDENTITY_TARGET, "params": {**_IDENTITY_TARGET["params"], "temperature": "hot"}} + _save(tmp_path, RegistryRecipeKind.TARGETS, "bad_value", bad_value) + _save(tmp_path, RegistryRecipeKind.TARGETS, "good", _IDENTITY_TARGET) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert (failure.name, failure.reason) == ("bad_value", RestoreFailureReason.BUILD_ERROR) + assert failure.version + assert [ref.name for ref in report.restored] == ["good"] From c29d81279a020ffac74a5cc064b528338db69f94 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 20:37:35 +0000 Subject: [PATCH 10/38] FEAT: Restore saved registry entries and stay not ready if any fail --- doc/getting_started/pyrit_conf.md | 12 +++ pyrit/backend/services/runtime_lifecycle.py | 46 ++++++++- pyrit/setup/configuration_loader.py | 13 +++ tests/unit/backend/conftest.py | 8 ++ tests/unit/backend/test_runtime_lifecycle.py | 94 +++++++++++++++++++ tests/unit/setup/test_configuration_loader.py | 19 ++++ 6 files changed, 189 insertions(+), 3 deletions(-) diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index e8196da582..5d0306a53f 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -194,6 +194,18 @@ Suites store ordered references to scenario presets and resolve each preset's cu A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/scenario_suites`. Local writes use the same hidden lock-file and atomic replacement behavior as scenario presets. Startup scans the configured source so malformed suite documents are reported immediately without preventing the backend from starting. +### `registry_recipes_source` + +Stores saved targets, converters, scorers, and techniques in a local directory or Azure Blob container, so they come back when the backend restarts. An Azure URI may include a blob-name prefix: + +```yaml +registry_recipes_source: https://account.blob.core.windows.net/pyrit-storage/registry-recipes +``` + +Each kind has its own folder under the source (`targets/`, `converters/`, `scorers/`, `techniques/`), with one JSON file per saved entry named after the entry. A file holds what's needed to rebuild the entry, never a credential. A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/registry_recipes`. + +The backend restores saved entries at startup and after a live apply, once the configured initializers have registered theirs. If any saved entry can't be restored, the backend stays not ready (state `failed`, outcome `restore-failed`) and lists each failure with its reason in `GET /api/config/runtime`, so a scan can't quietly run without an entry it expects. Fix or delete the listed files, then restart. Files whose names aren't valid entry names (lowercase letters, digits, and underscores, starting with a letter) are listed as ignored and don't block startup. + ### `initialization_scripts` Local paths to custom Python scripts containing `PyRITInitializer` subclasses. Paths can be absolute or relative to the current working directory. diff --git a/pyrit/backend/services/runtime_lifecycle.py b/pyrit/backend/services/runtime_lifecycle.py index 9767c56749..8204c25b00 100644 --- a/pyrit/backend/services/runtime_lifecycle.py +++ b/pyrit/backend/services/runtime_lifecycle.py @@ -7,7 +7,7 @@ import logging import os import uuid -from typing import Any +from typing import TYPE_CHECKING, Any from fastapi import FastAPI @@ -26,10 +26,14 @@ from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH from pyrit.memory import CentralMemory from pyrit.registry import InitializerRegistry, ScenarioPresetRegistry, ScenarioSuiteRegistry +from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer from pyrit.setup.configuration_loader import ConfigurationLoader from pyrit.setup.environment_loading import resolve_environment_async from pyrit.setup.initialization import validate_reinitialization_memory +if TYPE_CHECKING: + from pyrit.models.catalog.registry_recipe import IgnoredRecipeFile, RestoreFailure + logger = logging.getLogger(__name__) @@ -46,6 +50,8 @@ def __init__(self, *, app: FastAPI, source: ConfigurationFileService) -> None: self.version: str | None = None self.outcome = "starting" self.message = "" + self.restore_failures: list[RestoreFailure] = [] + self.restore_ignored: list[IgnoredRecipeFile] = [] self.operations: set[asyncio.Task[None]] = set() self.management_operations: set[asyncio.Task[None]] = set() self.apply_task: asyncio.Task[None] | None = None @@ -75,6 +81,8 @@ def status(self) -> dict[str, Any]: "message": self.message, "enabled": self.topology_supported, "applying": self.apply_task is not None and not self.apply_task.done(), + "restore_failures": [failure.model_dump(mode="json") for failure in self.restore_failures], + "restore_ignored": [file.model_dump(mode="json") for file in self.restore_ignored], } async def _load_async(self) -> ConfigurationLoader: @@ -119,6 +127,36 @@ def _publish(self, config: ConfigurationLoader) -> None: self.generation = str(uuid.uuid4()) self.state, self.outcome, self.message = "ready", "success", "PyRIT is ready." + async def _restore_saved_entries_async(self, config: ConfigurationLoader) -> bool: + """ + Rebuild saved registry entries after the startup scripts have registered theirs. + + A failure keeps the runtime not ready, since a scan that would have used a missing entry + could quietly fall back to a different one. Fixing the stored file and restarting recovers. + + Args: + config (ConfigurationLoader): The loaded configuration naming the storage source. + + Returns: + bool: Whether every saved entry was restored. + """ + restorer = RegistryRecipeRestorer(source=config.registry_recipes_source) + report = await asyncio.to_thread(restorer.restore) + self.restore_failures, self.restore_ignored = report.failures, report.ignored + if not report.failures: + return True + self.state, self.outcome = "failed", "restore-failed" + problems = f"{len(report.failures)} problem{'s' if len(report.failures) > 1 else ''}" + self.message = ( + f"Saved registry entries couldn't be restored ({problems}). " + "Fix or delete the stored files listed in the runtime status, then restart the backend." + ) + logger.error( + "Saved registry entries couldn't be restored: %s", + [f"{failure.kind.value}/{failure.name or '*'}: {failure.reason.value}" for failure in report.failures], + ) + return False + async def startup_async(self) -> None: """Keep management alive even when parsing or initialization fails.""" self.app.state.allow_custom_initializers = False @@ -137,7 +175,8 @@ async def startup_async(self) -> None: await config.initialize_pyrit_async(raise_on_initializer_error=True) await get_scenario_run_service().reconcile_interrupted_runs_async() _, self.version = await self.source.read_with_version_async() - self._publish(config) + if await self._restore_saved_entries_async(config): + self._publish(config) except Exception: self.state, self.outcome = "restart-required", "initialization-failed" self.message = "PyRIT setup failed. Repair saved configuration, then restart the backend." @@ -219,7 +258,8 @@ async def _apply_async(self, *, version: str) -> None: await config.apply_prepared_reinitialization_async(prepared=prepared) await self._management_async(config) self.version = current_version - self._publish(config) + if await self._restore_saved_entries_async(config): + self._publish(config) except (Exception, asyncio.CancelledError) as error: logger.exception( "PyRIT live apply failed phase=%s generation=%s", diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index baf670ca3a..eb12e9b2e8 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -116,6 +116,8 @@ class ConfigurationLoader(YamlLoadable): followed by a blob prefix, used to persist scenario presets. scenario_suites_source: Local directory or Azure Blob container URI, optionally followed by a blob prefix, used to persist scenario suites. + registry_recipes_source: Local directory or Azure Blob container URI, optionally + followed by a blob prefix, used to persist saved targets, converters, scorers, and techniques. silent: Whether to suppress initialization messages. seed: Optional root seed for deterministic converter operations. operator: Name for the current operator, e.g. a team or username. @@ -171,6 +173,7 @@ class ConfigurationLoader(YamlLoadable): custom_initializers_source: str | None = None scenario_presets_source: str | None = None scenario_suites_source: str | None = None + registry_recipes_source: str | None = None target_upload_directory: str | None = None server: dict[str, Any] | None = None extensions: dict[str, Any] = field(default_factory=dict) @@ -200,6 +203,7 @@ def __post_init__(self) -> None: self._validate_custom_initializers_source() self._validate_scenario_presets_source() self._validate_scenario_suites_source() + self._validate_registry_recipes_source() self._validate_target_upload_directory() self._normalize_server() @@ -240,6 +244,15 @@ def _validate_scenario_suites_source(self) -> None: """ self._validate_document_source(name="scenario_suites_source", value=self.scenario_suites_source) + def _validate_registry_recipes_source(self) -> None: + """ + Validate the optional storage source for saved registry entries. + + Raises: + ValueError: If the source is not a non-empty string. + """ + self._validate_document_source(name="registry_recipes_source", value=self.registry_recipes_source) + @staticmethod def _validate_document_source(*, name: str, value: Any) -> None: """ diff --git a/tests/unit/backend/conftest.py b/tests/unit/backend/conftest.py index 6d0ce4bde3..d8a62cd66d 100644 --- a/tests/unit/backend/conftest.py +++ b/tests/unit/backend/conftest.py @@ -18,6 +18,7 @@ from pyrit.backend.services.manual_send_scheduler import get_manual_send_scheduler from pyrit.backend.services.message_send_service import get_message_send_service from pyrit.memory import SQLiteMemory +from pyrit.registry.registry_recipe_storage import RegistryRecipeStorage from test_support.fixtures import authentication_policy, compatibility_headers, compatibility_id # noqa: F401 from unit.async_utils import TEST_WATCHDOG_SECONDS @@ -57,3 +58,10 @@ def isolated_manual_message_services() -> Iterator[None]: yield for factory in (get_attack_service, get_message_send_service, get_manual_send_scheduler, get_dataset_service): factory.cache_clear() + + +@pytest.fixture(autouse=True) +def isolated_saved_registry_entries(tmp_path: Path) -> Iterator[None]: + """Keep backend startup and apply from restoring entries saved on the machine running the tests.""" + with patch.object(RegistryRecipeStorage, "_get_default_storage_dir", return_value=tmp_path / "registry_recipes"): + yield diff --git a/tests/unit/backend/test_runtime_lifecycle.py b/tests/unit/backend/test_runtime_lifecycle.py index 5bc4187327..4e9048ce78 100644 --- a/tests/unit/backend/test_runtime_lifecycle.py +++ b/tests/unit/backend/test_runtime_lifecycle.py @@ -4,6 +4,7 @@ """Admission and fail-fast runtime replacement invariants.""" import asyncio +import json import os from collections.abc import Generator from contextlib import ExitStack @@ -29,6 +30,7 @@ from pyrit.memory import CentralMemory, MemoryInterface from pyrit.models import ScenarioRunState from pyrit.models.catalog.scenario import RunScenarioRequest +from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer from pyrit.scenario import Scenario from pyrit.setup.configuration_loader import ConfigurationLoader from unit.mocks import make_scenario_result @@ -917,3 +919,95 @@ async def close_async() -> None: finally: release.set() await asyncio.gather(shutdown, return_exceptions=True) + + +def _save_recipe(root: Path, *, kind: str, name: str, document: dict[str, object]) -> Path: + folder = root / kind + folder.mkdir(parents=True, exist_ok=True) + path = folder / f"{name}.json" + path.write_text(json.dumps(document), encoding="utf-8") + return path + + +async def _start_with_recipes_async(runtime: RuntimeLifecycle, *, recipes: Path) -> None: + config = await runtime._load_async() + config.registry_recipes_source = str(recipes) + with ( + patch.object(config, "initialize_pyrit_async", AsyncMock()), + patch.object(lifecycle_module, "get_scenario_run_service", return_value=MagicMock(spec=ScenarioRunService)), + ): + await runtime.startup_async() + + +async def test_startup_with_an_unrestorable_entry_is_not_ready_async(runtime: RuntimeLifecycle, tmp_path: Path) -> None: + recipes = tmp_path / "recipes" + _save_recipe(recipes, kind="targets", name="broken", document={"class_module": "nowhere", "class_name": "Nope"}) + + await _start_with_recipes_async(runtime, recipes=recipes) + + assert (runtime.state, runtime.outcome) == ("failed", "restore-failed") + assert "1 problem" in runtime.message + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=runtime.app), base_url="http://test") as client: + health = await client.get("/api/health") + status = await client.get("/api/config/runtime") + assert health.json()["ready"] is False + (failure,) = status.json()["restore_failures"] + assert (failure["kind"], failure["name"], failure["reason"]) == ("targets", "broken", "unknown_class") + assert failure["version"] + + +async def test_startup_lists_ignored_files_without_blocking_readiness_async( + runtime: RuntimeLifecycle, tmp_path: Path +) -> None: + recipes = tmp_path / "recipes" + _save_recipe(recipes, kind="converters", name="Not-A-Name", document={}) + + await _start_with_recipes_async(runtime, recipes=recipes) + + assert runtime.state == "ready" + status = runtime.status() + assert status["restore_failures"] == [] + assert [(item["kind"], item["file_name"]) for item in status["restore_ignored"]] == [ + ("converters", "Not-A-Name.json") + ] + + +async def test_fixing_the_file_and_restarting_reaches_ready_async(runtime: RuntimeLifecycle, tmp_path: Path) -> None: + recipes = tmp_path / "recipes" + broken = _save_recipe(recipes, kind="scorers", name="broken", document={"schema_version": 99}) + await _start_with_recipes_async(runtime, recipes=recipes) + assert runtime.status()["restore_failures"][0]["reason"] == "newer_schema_version" + + broken.unlink() + await _start_with_recipes_async(runtime, recipes=recipes) + + assert (runtime.state, runtime.outcome) == ("ready", "success") + assert runtime.status()["restore_failures"] == [] + + +async def test_apply_with_an_unrestorable_entry_keeps_the_new_version_async( + runtime: RuntimeLifecycle, tmp_path: Path +) -> None: + recipes = tmp_path / "recipes" + _save_recipe(recipes, kind="techniques", name="broken", document={"type": "NoSuchAttack"}) + config = await runtime._load_async() + config.registry_recipes_source = str(recipes) + + await apply_async(runtime) + + _, current_version = await runtime.source.read_with_version_async() + assert (runtime.state, runtime.outcome) == ("failed", "restore-failed") + assert runtime.version == current_version + assert runtime.status()["restore_failures"][0]["reason"] == "unknown_class" + config.apply_prepared_reinitialization_async.assert_awaited_once() + assert runtime.begin_apply(version=current_version)["outcome"] == "restart-required" + + +async def test_restore_cancellation_propagates_from_startup_async(runtime: RuntimeLifecycle, tmp_path: Path) -> None: + runtime.state = "initializing" + + with patch.object(RegistryRecipeRestorer, "restore", side_effect=asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError): + await _start_with_recipes_async(runtime, recipes=tmp_path / "recipes") + + assert runtime.state == "initializing" diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 72b12fc6a7..5b86b1aa9c 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -85,6 +85,25 @@ def test_scenario_suites_source_rejects_invalid_value(self, invalid_value: objec with pytest.raises(ValueError, match="scenario_suites_source"): ConfigurationLoader(scenario_suites_source=invalid_value) # type: ignore[arg-type] + def test_registry_recipes_source_loads_from_yaml(self, tmp_path: pathlib.Path) -> None: + """Test loading an Azure Blob container URI for saved registry entries.""" + config_path = tmp_path / "config.yaml" + config_path.write_text( + "registry_recipes_source: https://account.blob.core.windows.net/pyrit-storage/registry-recipes\n", + encoding="utf-8", + ) + + config = ConfigurationLoader.from_yaml_file(config_path) + + assert config.registry_recipes_source == "https://account.blob.core.windows.net/pyrit-storage/registry-recipes" + assert ConfigurationLoader().registry_recipes_source is None + + @pytest.mark.parametrize("invalid_value", ["", " ", 42]) + def test_registry_recipes_source_rejects_invalid_value(self, invalid_value: object) -> None: + """Test rejecting empty or non-string sources for saved registry entries.""" + with pytest.raises(ValueError, match="registry_recipes_source"): + ConfigurationLoader(registry_recipes_source=invalid_value) # type: ignore[arg-type] + def test_target_upload_directory_loads_from_yaml(self, tmp_path: pathlib.Path) -> None: config_path = tmp_path / "config.yaml" config_path.write_text("target_upload_directory: /srv/uploads\n", encoding="utf-8") From cfa4fc03ab6acfd6ef5c2d90c304ce4ebb05345d Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 23:41:45 +0000 Subject: [PATCH 11/38] Limit technique descriptions like the create request and export the recipe models --- pyrit/models/catalog/__init__.py | 18 ++++++++++++++++++ pyrit/models/catalog/registry_recipe.py | 4 ++-- pyrit/models/request_limits.py | 1 + .../registry/test_registry_recipe_storage.py | 10 ++++++++++ 4 files changed, 31 insertions(+), 2 deletions(-) diff --git a/pyrit/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index 102788e2a8..94c1abb948 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -20,6 +20,16 @@ if TYPE_CHECKING: from pyrit.models.catalog.initializer import RegisteredInitializer + from pyrit.models.catalog.registry_recipe import ( + ComponentRecipe, + IgnoredRecipeFile, + RegistryRecipeKind, + RestoreFailure, + RestoreFailureReason, + SavedEntryRef, + StoredRecipe, + TechniqueRecipe, + ) from pyrit.models.catalog.scenario import ( AttackErrorSummary, AttackRetrySummary, @@ -60,9 +70,15 @@ _LAZY_EXPORTS: dict[str, str] = { "AttackErrorSummary": "pyrit.models.catalog.scenario", "AttackRetrySummary": "pyrit.models.catalog.scenario", + "ComponentRecipe": "pyrit.models.catalog.registry_recipe", + "IgnoredRecipeFile": "pyrit.models.catalog.registry_recipe", "PresetEntry": "pyrit.models.catalog.scenario_preset", "PresetIssue": "pyrit.models.catalog.scenario_preset", "RegisteredInitializer": "pyrit.models.catalog.initializer", + "RegistryRecipeKind": "pyrit.models.catalog.registry_recipe", + "RestoreFailure": "pyrit.models.catalog.registry_recipe", + "RestoreFailureReason": "pyrit.models.catalog.registry_recipe", + "SavedEntryRef": "pyrit.models.catalog.registry_recipe", "ScorerInstance": "pyrit.models.catalog.scorer", "RegisteredScenario": "pyrit.models.catalog.scenario", "RunScenarioRequest": "pyrit.models.catalog.scenario", @@ -87,8 +103,10 @@ "ScenarioRunSummary": "pyrit.models.catalog.scenario", "ScenarioTechniqueSummary": "pyrit.models.catalog.scenario", "StoredPreset": "pyrit.models.catalog.scenario_preset", + "StoredRecipe": "pyrit.models.catalog.registry_recipe", "StoredScenarioSuite": "pyrit.models.catalog.scenario_suite", "TargetInstance": "pyrit.models.catalog.target", + "TechniqueRecipe": "pyrit.models.catalog.registry_recipe", } __all__ = list(_LAZY_EXPORTS) diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index ce2e544588..5ac95be552 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -10,7 +10,7 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator -from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS +from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS, MAX_TEXT_LENGTH _Identifier = Annotated[str, Field(max_length=MAX_IDENTIFIER_LENGTH)] _TechniqueTag = Annotated[str, Field(pattern=r"^[A-Za-z][A-Za-z0-9_]{0,63}$")] @@ -63,7 +63,7 @@ class TechniqueRecipe(_RegistryRecipe): params: dict[_Identifier, JsonValue] = Field( default_factory=dict, max_length=MAX_ITEMS, description="Attack constructor parameters" ) - description: str | None = None + description: str | None = Field(default=None, max_length=MAX_TEXT_LENGTH) tags: list[_TechniqueTag] = Field(default_factory=list, max_length=MAX_ITEMS) request_converters: list[_Identifier] | None = Field(default=None, max_length=MAX_ITEMS) response_converters: list[_Identifier] | None = Field(default=None, max_length=MAX_ITEMS) diff --git a/pyrit/models/request_limits.py b/pyrit/models/request_limits.py index c9a20196ce..033c8ef9bd 100644 --- a/pyrit/models/request_limits.py +++ b/pyrit/models/request_limits.py @@ -7,3 +7,4 @@ MAX_IDENTIFIER_LENGTH = 256 MAX_LABEL_KEY_LENGTH = 128 MAX_LABEL_VALUE_LENGTH = 1_024 +MAX_TEXT_LENGTH = 100_000 diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index 58851e2b45..dfcecc970e 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -115,6 +115,16 @@ def test_reports_a_document_nested_too_deeply(self, tmp_path: Path) -> None: assert (failure.name, failure.reason) == ("deep", RestoreFailureReason.INVALID_RECIPE) assert list(listing.recipes) == ["contoso_gpt4o"] + def test_reports_a_technique_description_over_the_limit(self, tmp_path: Path) -> None: + document = {"type": "RedTeamingAttack", "description": "x" * 100_001} + _write(tmp_path / "techniques", "long_description", json.dumps(document)) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TECHNIQUES, source=str(tmp_path)).list_recipes() + + (failure,) = listing.failures + assert (failure.name, failure.reason) == ("long_description", RestoreFailureReason.INVALID_RECIPE) + assert "description" in failure.message + def test_reports_a_newer_schema_version_before_checking_fields(self, tmp_path: Path) -> None: _write(tmp_path / "targets", "future", json.dumps({"schema_version": 2, "unknown_field": True})) From 11d05190756733eac79315a7eb466daf9c388c48 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 23:45:37 +0000 Subject: [PATCH 12/38] Keep restore failure messages safe and complete, and check technique credentials --- pyrit/models/catalog/__init__.py | 2 + pyrit/registry/registry_recipe_restore.py | 27 ++-- .../registry/test_registry_recipe_restore.py | 118 ++++++++++++++---- 3 files changed, 112 insertions(+), 35 deletions(-) diff --git a/pyrit/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index 94c1abb948..a113a8bef0 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -24,6 +24,7 @@ ComponentRecipe, IgnoredRecipeFile, RegistryRecipeKind, + RegistryRestoreReport, RestoreFailure, RestoreFailureReason, SavedEntryRef, @@ -76,6 +77,7 @@ "PresetIssue": "pyrit.models.catalog.scenario_preset", "RegisteredInitializer": "pyrit.models.catalog.initializer", "RegistryRecipeKind": "pyrit.models.catalog.registry_recipe", + "RegistryRestoreReport": "pyrit.models.catalog.registry_recipe", "RestoreFailure": "pyrit.models.catalog.registry_recipe", "RestoreFailureReason": "pyrit.models.catalog.registry_recipe", "SavedEntryRef": "pyrit.models.catalog.registry_recipe", diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index 0db95c36b7..959e94565c 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -8,7 +8,6 @@ import re from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar -from urllib.parse import urlparse from pyrit.models.catalog.registry_recipe import ( ComponentRecipe, @@ -55,7 +54,7 @@ class RegistryRecipeRestorer: """ _RESERVED_TECHNIQUE_NAMES: ClassVar[frozenset[str]] = frozenset({"all", "default", "types"}) - _URL_PATTERN: ClassVar[re.Pattern[str]] = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*://[^\s'\"<>]+") + _URL_PATTERN: ClassVar[re.Pattern[str]] = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*://\S+") _KIND_BY_COMPONENT_TYPE: ClassVar[dict[ComponentType, RegistryRecipeKind]] = { ComponentType.TARGET: RegistryRecipeKind.TARGETS, ComponentType.CONVERTER: RegistryRecipeKind.CONVERTERS, @@ -219,7 +218,8 @@ def _technique_references(self, *, ref: SavedEntryRef, recipe: TechniqueRecipe) list[SavedEntryRef]: The adversarial target, converters, and attack parameter references. Raises: - _RecipeRejectedError: If the name is reserved or the attack type isn't registered. + _RecipeRejectedError: If the name is reserved, the attack type isn't registered, or the + recipe holds a credential. """ from pyrit.registry.components.attack_registry import AttackRegistry @@ -237,6 +237,7 @@ def _technique_references(self, *, ref: SavedEntryRef, recipe: TechniqueRecipe) if recipe.adversarial_chat is not None: references.append(SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name=recipe.adversarial_chat)) parameters = attack_registry.get_class_metadata(attack_class).parameters + self._check_stored_credentials(parameters=parameters, params=recipe.params) return references + self._parameter_references(parameters=parameters, params=recipe.params) @staticmethod @@ -509,24 +510,30 @@ def _safe_message(cls, error: Exception) -> str: """ Describe an error for the failure report, removing credentials from any URL it mentions. - Constructor errors can quote a server setting, such as an endpoint with a SAS token. + Constructor errors can quote a server setting, such as an endpoint with a SAS token. A URL + runs to the next whitespace, since quotes and other punctuation can appear inside one. Returns: - str: The error's type and text, with every URL reduced to its scheme, host, and path. + str: The error's type and text, with every URL cut before its query or fragment and + without a user name or password; just the type if the text can't be read. """ - text = cls._URL_PATTERN.sub(lambda match: cls._redact_url(match.group(0)), str(error)) + try: + text = cls._URL_PATTERN.sub(lambda match: cls._redact_url(match.group(0)), str(error)) + except Exception: + text = "" return f"{type(error).__name__}: {text}" if text else type(error).__name__ @staticmethod def _redact_url(url: str) -> str: """ - Drop the user name, password, query, and fragment from a URL. + Drop the user name, password, query, and fragment from a URL, without parsing it. Returns: - str: The URL without anything that can carry a credential. + str: The scheme, host, and path, which can't carry a credential. """ - parsed = urlparse(url) - return parsed._replace(netloc=parsed.netloc.rpartition("@")[2], query="", fragment="").geturl() + scheme, _, rest = url.partition("://") + rest = re.split(r"[?#]", rest, maxsplit=1)[0] + return f"{scheme}://{rest.rpartition('@')[2]}" class _RecipeRejectedError(Exception): diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index b05160c4e0..f56884a685 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import dataclasses import json import os import sys @@ -20,7 +21,9 @@ SavedEntryRef, TechniqueRecipe, ) +from pyrit.models.parameter import Parameter from pyrit.prompt_target import OpenAIChatTarget, RoundRobinTarget +from pyrit.registry.components.attack_registry import AttackRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry from pyrit.registry.components.converter_registry import ConverterRegistry from pyrit.registry.components.scorer_registry import ScorerRegistry @@ -146,6 +149,12 @@ def test_technique_recipe_matches_the_create_request(self) -> None: request.model_dump(exclude_unset=True, exclude={"name", "type", "tags"}) ) + def test_technique_recipe_defaults_match_the_create_request(self) -> None: + recipe = TechniqueRecipe.model_validate({"type": "RedTeamingAttack"}) + request = CreateTechniqueRequest.model_validate({"name": "formal_red_team", "type": "RedTeamingAttack"}) + + assert recipe.model_dump(exclude={"schema_version"}) == request.model_dump(exclude={"name"}) + @pytest.mark.parametrize( "changes", [ @@ -156,6 +165,7 @@ def test_technique_recipe_matches_the_create_request(self) -> None: {"params": {"x" * 300: 1}}, {"request_converters": ["c"] * 101}, {"params": {"temperature": float("nan")}}, + {"description": "x" * 100_001}, {"unexpected": True}, ], ) @@ -272,6 +282,29 @@ def test_refuses_a_stored_credential(self, tmp_path: Path) -> None: assert "sk-secret-value" not in failure.message assert [ref.name for ref in report.restored] == ["good"] + def test_restores_a_class_whose_optional_credentials_are_left_empty(self, tmp_path: Path) -> None: + _save( + tmp_path, + RegistryRecipeKind.TARGETS, + "anonymous_agent", + { + "class_module": "pyrit.prompt_target.a2a_target", + "class_name": "A2ATarget", + "params": {"endpoint": "https://agent.example.test"}, + }, + ) + + with patch.object(TargetRegistry, "create_named_instance") as build: + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + assert report.failures == [] + build.assert_called_once_with( + name="anonymous_agent", + type_name="A2ATarget", + params={"endpoint": "https://agent.example.test"}, + external_input=True, + ) + def test_refuses_a_class_that_needs_the_server_upload_folder(self, tmp_path: Path) -> None: _save( tmp_path, @@ -315,8 +348,14 @@ def test_a_planning_error_fails_only_that_entry(self, tmp_path: Path) -> None: assert failure.message == "RuntimeError: unexpected" assert [ref.name for ref in report.restored] == ["b_second"] - def test_build_errors_drop_credentials_from_urls(self, tmp_path: Path) -> None: - endpoint = "https://user:hunter2@evil.example/openai/v1?sig=SYNTHETIC-SAS#frag" + @pytest.mark.parametrize( + "endpoint", + [ + "https://user:hunter2@evil.example/openai/v1?sig=SYNTHETIC-SAS#frag", + "https://user:hunter2@evil.example/openai/v1'part?sig=SYNTHETIC-SAS#frag", + ], + ) + def test_build_errors_drop_credentials_from_urls(self, tmp_path: Path, endpoint: str) -> None: no_endpoint = {**_IDENTITY_TARGET, "params": {"auth_mode": "identity", "model_name": "gpt-4o"}} _save(tmp_path, RegistryRecipeKind.TARGETS, "server_endpoint", no_endpoint) @@ -325,10 +364,50 @@ def test_build_errors_drop_credentials_from_urls(self, tmp_path: Path) -> None: (failure,) = report.failures assert failure.reason is RestoreFailureReason.BUILD_ERROR - assert "https://evil.example/openai/v1" in failure.message + assert "evil.example/openai/v1" in failure.message for secret in ("hunter2", "SYNTHETIC-SAS", "frag"): assert secret not in failure.message + def test_error_text_that_is_not_a_valid_url_fails_only_that_entry(self, tmp_path: Path) -> None: + malformed = { + **_IDENTITY_TARGET, + "params": {**_IDENTITY_TARGET["params"], "temperature": "https://[x?sig=SECRET"}, + } + _save(tmp_path, RegistryRecipeKind.TARGETS, "malformed", malformed) + _save(tmp_path, RegistryRecipeKind.TARGETS, "good", _IDENTITY_TARGET) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert (failure.name, failure.reason) == ("malformed", RestoreFailureReason.BUILD_ERROR) + assert failure.version + assert "could not be coerced" in failure.message + assert "SECRET" not in failure.message + assert [ref.name for ref in report.restored] == ["good"] + + def test_refuses_a_stored_credential_in_attack_parameters(self, tmp_path: Path) -> None: + attack_registry = AttackRegistry.get_registry_singleton() + original_get_class_metadata = attack_registry.get_class_metadata + + def with_credential(cls: type) -> Any: + metadata = original_get_class_metadata(cls) + credential = Parameter(name="api_key", description="", sensitive=True) + return dataclasses.replace(metadata, parameters=(*metadata.parameters, credential)) + + _save( + tmp_path, + RegistryRecipeKind.TECHNIQUES, + "keyed", + {"type": "RedTeamingAttack", "params": {"api_key": "sk-x"}}, + ) + + with patch.object(attack_registry, "get_class_metadata", side_effect=with_credential): + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert (failure.name, failure.reason) == ("keyed", RestoreFailureReason.UNSUPPORTED_CREDENTIAL) + assert "sk-x" not in failure.message + def test_an_empty_reference_is_left_unset(self, tmp_path: Path) -> None: _save( tmp_path, @@ -422,31 +501,20 @@ def test_reports_a_name_taken_at_startup_and_fails_its_dependents(self, tmp_path assert failures["pool"].reason is RestoreFailureReason.FAILED_DEPENDENCY assert TargetRegistry.get_registry_singleton().instances.get("contoso_gpt4o") is startup_target - def test_reports_a_technique_name_taken_in_another_case(self, tmp_path: Path) -> None: + def test_reports_a_technique_name_that_matches_a_tag_in_another_case(self, tmp_path: Path) -> None: + without_converters = {key: value for key, value in _TECHNIQUE.items() if key != "request_converters"} _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) - _save( - tmp_path, - RegistryRecipeKind.CONVERTERS, - "tone_formal", - { - "class_module": "pyrit.converter.tone_converter", - "class_name": "ToneConverter", - "params": {"converter_target": "contoso_gpt4o", "tone": "formal"}, - }, - ) - _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "formal_red_team", _TECHNIQUE) - restorer = RegistryRecipeRestorer(source=str(tmp_path)) - assert restorer.restore().failures == [] - for registry in (TargetRegistry, ConverterRegistry): - registry.reset_registry_singleton() - (tmp_path / "techniques" / "formal_red_team.json").rename(tmp_path / "techniques" / "formal_red_team_copy.json") - _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "custom", {**_TECHNIQUE, "tags": []}) + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "formal_red_team", {**without_converters, "tags": ["Custom"]}) + assert RegistryRecipeRestorer(source=str(tmp_path)).restore().failures == [] + TargetRegistry.reset_registry_singleton() + (tmp_path / "techniques" / "formal_red_team.json").unlink() + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "custom", {**without_converters, "tags": []}) - report = restorer.restore() + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() - failures = _failures_by_name(report.failures) - assert failures["custom"].reason is RestoreFailureReason.NAME_CLASH - assert "custom" in failures["custom"].message + (failure,) = report.failures + assert (failure.name, failure.reason) == ("custom", RestoreFailureReason.NAME_CLASH) + assert "custom" in failure.message def test_builds_nothing_when_a_folder_cannot_be_listed(self, tmp_path: Path) -> None: _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) From b7bf0127eef438abd57b49fdcaa5db21306e8cb4 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 23:52:56 +0000 Subject: [PATCH 13/38] Document the recipes setting in the example config and pin the restore version check --- .pyrit_conf_example | 8 ++++++++ doc/getting_started/pyrit_conf.md | 2 +- test_support/analytics_benchmark.py | 1 + tests/unit/backend/test_runtime_lifecycle.py | 2 ++ 4 files changed, 12 insertions(+), 1 deletion(-) diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 51488ee0d8..5986c5bf3c 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -161,6 +161,14 @@ allow_custom_initializers: false # as scenario_presets_source. Defaults to ~/.pyrit/scenario_suites. # scenario_suites_source: https://account.blob.core.windows.net/container/scenario_suites +# Optional storage for saved targets, converters, scorers, and techniques, so they +# come back when the backend restarts. Each kind has its own folder under the source. +# This supports the same local directory and Azure Blob sources as +# scenario_presets_source. If any saved entry can't be restored, the backend stays +# not ready and lists the failures in its runtime status. Defaults to +# ~/.pyrit/registry_recipes. +# registry_recipes_source: https://account.blob.core.windows.net/container/registry_recipes + # Upload Folder for API-Created Targets # ------------------------------------- # Folder that targets created through the REST API may upload local files from diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 5d0306a53f..48a2e40516 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -202,7 +202,7 @@ Stores saved targets, converters, scorers, and techniques in a local directory o registry_recipes_source: https://account.blob.core.windows.net/pyrit-storage/registry-recipes ``` -Each kind has its own folder under the source (`targets/`, `converters/`, `scorers/`, `techniques/`), with one JSON file per saved entry named after the entry. A file holds what's needed to rebuild the entry, never a credential. A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/registry_recipes`. +Each kind has its own folder under the source (`targets/`, `converters/`, `scorers/`, `techniques/`), with one JSON file per saved entry named after the entry. A file holds what's needed to rebuild the entry; restore refuses a file that stores a value for a parameter its class marks as a credential, such as `api_key`. A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/registry_recipes`. The backend restores saved entries at startup and after a live apply, once the configured initializers have registered theirs. If any saved entry can't be restored, the backend stays not ready (state `failed`, outcome `restore-failed`) and lists each failure with its reason in `GET /api/config/runtime`, so a scan can't quietly run without an entry it expects. Fix or delete the listed files, then restart. Files whose names aren't valid entry names (lowercase letters, digits, and underscores, starting with a letter) are listed as ignored and don't block startup. diff --git a/test_support/analytics_benchmark.py b/test_support/analytics_benchmark.py index c8acfd455d..61c66e1df9 100644 --- a/test_support/analytics_benchmark.py +++ b/test_support/analytics_benchmark.py @@ -65,6 +65,7 @@ def benchmark_directory(tmp_path: Path) -> Iterator[Path]: "initializers": ["technique"], "initialization_scripts": [], "custom_initializers_source": str(directory / "initializers"), + "registry_recipes_source": str(directory / "registry_recipes"), "silent": True, } ), diff --git a/tests/unit/backend/test_runtime_lifecycle.py b/tests/unit/backend/test_runtime_lifecycle.py index 4e9048ce78..a01958004a 100644 --- a/tests/unit/backend/test_runtime_lifecycle.py +++ b/tests/unit/backend/test_runtime_lifecycle.py @@ -951,6 +951,8 @@ async def test_startup_with_an_unrestorable_entry_is_not_ready_async(runtime: Ru health = await client.get("/api/health") status = await client.get("/api/config/runtime") assert health.json()["ready"] is False + assert status.json()["version"] == runtime.version + assert runtime.version is not None (failure,) = status.json()["restore_failures"] assert (failure["kind"], failure["name"], failure["reason"]) == ("targets", "broken", "unknown_class") assert failure["version"] From e45865b3fbfddcc541ffdf17936753f9e39aa271 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 01:05:23 +0000 Subject: [PATCH 14/38] Test that a description at the limit is accepted --- tests/unit/registry/test_registry_recipe_storage.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index dfcecc970e..af55c36e64 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -115,14 +115,16 @@ def test_reports_a_document_nested_too_deeply(self, tmp_path: Path) -> None: assert (failure.name, failure.reason) == ("deep", RestoreFailureReason.INVALID_RECIPE) assert list(listing.recipes) == ["contoso_gpt4o"] - def test_reports_a_technique_description_over_the_limit(self, tmp_path: Path) -> None: - document = {"type": "RedTeamingAttack", "description": "x" * 100_001} - _write(tmp_path / "techniques", "long_description", json.dumps(document)) + def test_limits_a_technique_description_like_the_create_request(self, tmp_path: Path) -> None: + folder = tmp_path / "techniques" + _write(folder, "at_limit", json.dumps({"type": "RedTeamingAttack", "description": "x" * 100_000})) + _write(folder, "over_limit", json.dumps({"type": "RedTeamingAttack", "description": "x" * 100_001})) listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TECHNIQUES, source=str(tmp_path)).list_recipes() + assert list(listing.recipes) == ["at_limit"] (failure,) = listing.failures - assert (failure.name, failure.reason) == ("long_description", RestoreFailureReason.INVALID_RECIPE) + assert (failure.name, failure.reason) == ("over_limit", RestoreFailureReason.INVALID_RECIPE) assert "description" in failure.message def test_reports_a_newer_schema_version_before_checking_fields(self, tmp_path: Path) -> None: From c94ffbdd9163575a8c7c8b54ff38d632f4a25122 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 01:20:14 +0000 Subject: [PATCH 15/38] Show only the error type when a URL in it may carry credentials --- pyrit/registry/registry_recipe_restore.py | 48 ++++++++--------- .../registry/test_registry_recipe_restore.py | 52 +++++++++++++++++-- 2 files changed, 69 insertions(+), 31 deletions(-) diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index 959e94565c..39a9181768 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -54,7 +54,6 @@ class RegistryRecipeRestorer: """ _RESERVED_TECHNIQUE_NAMES: ClassVar[frozenset[str]] = frozenset({"all", "default", "types"}) - _URL_PATTERN: ClassVar[re.Pattern[str]] = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*://\S+") _KIND_BY_COMPONENT_TYPE: ClassVar[dict[ComponentType, RegistryRecipeKind]] = { ComponentType.TARGET: RegistryRecipeKind.TARGETS, ComponentType.CONVERTER: RegistryRecipeKind.CONVERTERS, @@ -102,14 +101,12 @@ def _read_all(self) -> tuple[list[RecipeListing], list[RestoreFailure]]: tuple[list[RecipeListing], list[RestoreFailure]]: The folders that were read, and one failure per folder that wasn't. """ - from azure.core.exceptions import AzureError - listings: list[RecipeListing] = [] failures: list[RestoreFailure] = [] for kind, storage in self._storages.items(): try: listings.append(storage.list_recipes()) - except (OSError, AzureError) as error: + except Exception as error: # Any listing error, not only storage ones, must become a failure. message = f"The saved {kind.value} couldn't be listed: {storage._describe_read_error(error)}" failures.append( RestoreFailure(kind=kind, reason=RestoreFailureReason.UNREADABLE_STORAGE, message=message) @@ -505,35 +502,34 @@ def _singular(kind: RegistryRecipeKind) -> str: """ return kind.value.removesuffix("s") - @classmethod - def _safe_message(cls, error: Exception) -> str: + @staticmethod + def _safe_message(error: Exception) -> str: """ - Describe an error for the failure report, removing credentials from any URL it mentions. + Describe an error for the failure report without any credential a URL in it may carry. - Constructor errors can quote a server setting, such as an endpoint with a SAS token. A URL - runs to the next whitespace, since quotes and other punctuation can appear inside one. + Constructor errors can quote a server setting, such as an endpoint with a SAS token. A word + with a ``:`` or ``//`` may hold a URL, so it's cut before any ``?`` or ``#``. If such a word + has an ``@`` (a user name or password), or a ``//`` URL can't be cut cleanly because it + contains a space, only the error's type is kept. Returns: - str: The error's type and text, with every URL cut before its query or fragment and - without a user name or password; just the type if the text can't be read. + str: The error's type and its text with URLs cut, or just the type. """ + name = type(error).__name__ try: - text = cls._URL_PATTERN.sub(lambda match: cls._redact_url(match.group(0)), str(error)) + text = str(error) except Exception: - text = "" - return f"{type(error).__name__}: {text}" if text else type(error).__name__ - - @staticmethod - def _redact_url(url: str) -> str: - """ - Drop the user name, password, query, and fragment from a URL, without parsing it. - - Returns: - str: The scheme, host, and path, which can't carry a credential. - """ - scheme, _, rest = url.partition("://") - rest = re.split(r"[?#]", rest, maxsplit=1)[0] - return f"{scheme}://{rest.rpartition('@')[2]}" + return name + words = re.split(r"([ \t\n\r\f\v]+)", text) + may_be_url = [":" in word or "//" in word for word in words] + if any("@" in word for word, url in zip(words, may_be_url, strict=True) if url): + return name + cut = "".join( + re.split(r"[?#]", word, maxsplit=1)[0] if url else word for word, url in zip(words, may_be_url, strict=True) + ) + if "//" in text and ("?" in cut or "#" in cut): + return name + return f"{name}: {cut}" if cut else name class _RecipeRejectedError(Exception): diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index f56884a685..0179874ec3 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -349,13 +349,16 @@ def test_a_planning_error_fails_only_that_entry(self, tmp_path: Path) -> None: assert [ref.name for ref in report.restored] == ["b_second"] @pytest.mark.parametrize( - "endpoint", + ("endpoint", "shown"), [ - "https://user:hunter2@evil.example/openai/v1?sig=SYNTHETIC-SAS#frag", - "https://user:hunter2@evil.example/openai/v1'part?sig=SYNTHETIC-SAS#frag", + ("https://evil.example/openai/v1?sig=SYNTHETIC-SAS#frag", "https://evil.example/openai/v1"), + ("https://evil.example/openai/v1'part?sig=SYNTHETIC-SAS#frag", "https://evil.example/openai/v1'part"), + ("https://evil.example/open\u00a0ai/v1?sig=SYNTHETIC-SAS", "https://evil.example/open\u00a0ai/v1"), + ("https://user:hunter2@evil.example/openai/v1?sig=SYNTHETIC-SAS#frag", None), + ("//user:hunter2@evil.example/openai/v1?sig=SYNTHETIC-SAS", None), ], ) - def test_build_errors_drop_credentials_from_urls(self, tmp_path: Path, endpoint: str) -> None: + def test_build_errors_drop_credentials_from_urls(self, tmp_path: Path, endpoint: str, shown: str | None) -> None: no_endpoint = {**_IDENTITY_TARGET, "params": {"auth_mode": "identity", "model_name": "gpt-4o"}} _save(tmp_path, RegistryRecipeKind.TARGETS, "server_endpoint", no_endpoint) @@ -364,7 +367,10 @@ def test_build_errors_drop_credentials_from_urls(self, tmp_path: Path, endpoint: (failure,) = report.failures assert failure.reason is RestoreFailureReason.BUILD_ERROR - assert "evil.example/openai/v1" in failure.message + if shown is None: + assert failure.message == "ValueError" + else: + assert f"'{shown}" in failure.message for secret in ("hunter2", "SYNTHETIC-SAS", "frag"): assert secret not in failure.message @@ -385,6 +391,25 @@ def test_error_text_that_is_not_a_valid_url_fails_only_that_entry(self, tmp_path assert "SECRET" not in failure.message assert [ref.name for ref in report.restored] == ["good"] + @pytest.mark.parametrize( + ("text", "expected"), + [ + ("got ://host/path?sig=SECRET", "ValueError: got ://host/path"), + ("got https://[x?sig=SECRET", "ValueError: got https://[x"), + ("got https:host/path?sig=SECRET#SECRET", "ValueError: got https:host/path"), + ("got https://host/pa\u00a0th?sig=SECRET#SECRET", "ValueError: got https://host/pa\u00a0th"), + ("keeps spacing https://host/?sig=SECRET intact", "ValueError: keeps spacing https://host/ intact"), + ("bad value: 'abc' (expected a number?)", "ValueError: bad value: 'abc' (expected a number?)"), + ("no address here", "ValueError: no address here"), + ("got //user:SECRET@host/path", "ValueError"), + ("got user:SECRET@https://host/path", "ValueError"), + ("got https://user\t:SECRET@host/path", "ValueError"), + ("got https://host/a b?sig=SECRET", "ValueError"), + ], + ) + def test_safe_messages_drop_credentials_from_anything_that_may_be_a_url(self, text: str, expected: str) -> None: + assert RegistryRecipeRestorer._safe_message(ValueError(text)) == expected + def test_refuses_a_stored_credential_in_attack_parameters(self, tmp_path: Path) -> None: attack_registry = AttackRegistry.get_registry_singleton() original_get_class_metadata = attack_registry.get_class_metadata @@ -538,6 +563,23 @@ def list_recipes(storage: RegistryRecipeStorage) -> Any: assert report.restored == [] assert "contoso_gpt4o" not in TargetRegistry.get_registry_singleton().instances + def test_builds_nothing_when_listing_a_folder_fails_unexpectedly(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) + original_list_recipes = RegistryRecipeStorage.list_recipes + + def list_recipes(storage: RegistryRecipeStorage) -> Any: + if storage._kind is RegistryRecipeKind.SCORERS: + raise ValueError("https://account.blob.core.windows.net/x?sig=SECRET") + return original_list_recipes(storage) + + with patch.object(RegistryRecipeStorage, "list_recipes", autospec=True, side_effect=list_recipes): + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert (failure.kind, failure.reason) == (RegistryRecipeKind.SCORERS, RestoreFailureReason.UNREADABLE_STORAGE) + assert "SECRET" not in failure.message + assert report.restored == [] + def test_reports_constructor_errors_and_keeps_going(self, tmp_path: Path) -> None: bad_value = {**_IDENTITY_TARGET, "params": {**_IDENTITY_TARGET["params"], "temperature": "hot"}} _save(tmp_path, RegistryRecipeKind.TARGETS, "bad_value", bad_value) From 26377ea234c35b18422cb6fe3aa86635a5c7ac6e Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 01:29:16 +0000 Subject: [PATCH 16/38] Point the restore message at where failures are listed and guard the benchmark's recipes folder --- doc/getting_started/pyrit_conf.md | 2 +- pyrit/backend/services/runtime_lifecycle.py | 2 +- tests/unit/backend/test_analytics_benchmark.py | 7 +++++++ 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 48a2e40516..b2fb7ed98f 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -204,7 +204,7 @@ registry_recipes_source: https://account.blob.core.windows.net/pyrit-storage/reg Each kind has its own folder under the source (`targets/`, `converters/`, `scorers/`, `techniques/`), with one JSON file per saved entry named after the entry. A file holds what's needed to rebuild the entry; restore refuses a file that stores a value for a parameter its class marks as a credential, such as `api_key`. A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/registry_recipes`. -The backend restores saved entries at startup and after a live apply, once the configured initializers have registered theirs. If any saved entry can't be restored, the backend stays not ready (state `failed`, outcome `restore-failed`) and lists each failure with its reason in `GET /api/config/runtime`, so a scan can't quietly run without an entry it expects. Fix or delete the listed files, then restart. Files whose names aren't valid entry names (lowercase letters, digits, and underscores, starting with a letter) are listed as ignored and don't block startup. +The backend restores saved entries at startup and after a live apply, once the configured initializers have registered theirs. If any saved entry can't be restored, the backend stays not ready (state `failed`, outcome `restore-failed`) and lists each failure with its reason in `GET /api/config/runtime`, so a scan can't quietly run without an entry it expects. Fix or delete the listed files, then restart. Files whose names aren't valid entry names (up to 64 lowercase letters, digits, and underscores, starting with a letter) are listed as ignored and don't block startup. ### `initialization_scripts` diff --git a/pyrit/backend/services/runtime_lifecycle.py b/pyrit/backend/services/runtime_lifecycle.py index 8204c25b00..a7e9f778ae 100644 --- a/pyrit/backend/services/runtime_lifecycle.py +++ b/pyrit/backend/services/runtime_lifecycle.py @@ -149,7 +149,7 @@ async def _restore_saved_entries_async(self, config: ConfigurationLoader) -> boo problems = f"{len(report.failures)} problem{'s' if len(report.failures) > 1 else ''}" self.message = ( f"Saved registry entries couldn't be restored ({problems}). " - "Fix or delete the stored files listed in the runtime status, then restart the backend." + "Fix or delete the files named in the server log or GET /api/config/runtime, then restart the backend." ) logger.error( "Saved registry entries couldn't be restored: %s", diff --git a/tests/unit/backend/test_analytics_benchmark.py b/tests/unit/backend/test_analytics_benchmark.py index 9e5b75096c..8dcbb30e59 100644 --- a/tests/unit/backend/test_analytics_benchmark.py +++ b/tests/unit/backend/test_analytics_benchmark.py @@ -4,6 +4,7 @@ """Local SQLite smoke coverage and unit tests of the integration workload's failure gates.""" import asyncio +import json import os import subprocess import sys @@ -60,6 +61,12 @@ def measurements() -> AnalyticsBenchmark: return benchmark +def test_benchmark_configuration_keeps_saved_entries_in_its_own_folder(benchmark_directory: Path) -> None: + config = json.loads((benchmark_directory / "config.yaml").read_text(encoding="utf-8")) + + assert config["registry_recipes_source"] == str(benchmark_directory / "registry_recipes") + + @pytest.mark.timeout(120) @pytest.mark.usefixtures("semantic_analytics_budgets") @pytest.mark.parametrize("elapsed_report_seconds", [0.0, 31.0], ids=["normal", "expired-report-deadline"]) From 8663852426b042b33fdeca5c9a153df7474f54e4 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 02:37:19 +0000 Subject: [PATCH 17/38] Cut failure messages at any URL query or fragment and hide any with user info --- pyrit/registry/registry_recipe_restore.py | 24 +++---- .../registry/test_registry_recipe_restore.py | 64 +++++++++++++------ 2 files changed, 52 insertions(+), 36 deletions(-) diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index 39a9181768..90181df568 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -5,7 +5,6 @@ from __future__ import annotations -import re from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar @@ -505,30 +504,23 @@ def _singular(kind: RegistryRecipeKind) -> str: @staticmethod def _safe_message(error: Exception) -> str: """ - Describe an error for the failure report without any credential a URL in it may carry. + Describe an error for the failure report without any credential it may quote. - Constructor errors can quote a server setting, such as an endpoint with a SAS token. A word - with a ``:`` or ``//`` may hold a URL, so it's cut before any ``?`` or ``#``. If such a word - has an ``@`` (a user name or password), or a ``//`` URL can't be cut cleanly because it - contains a space, only the error's type is kept. + Constructor errors can quote a server setting, such as an endpoint with a SAS token. The text + is cut at the first ``?`` or ``#``, where a URL's query or fragment starts. A user name and + password can't be told apart reliably, so text with an ``@`` isn't shown at all. Returns: - str: The error's type and its text with URLs cut, or just the type. + str: The error's type and its text before any ``?`` or ``#``, or the type alone. """ name = type(error).__name__ try: text = str(error) except Exception: return name - words = re.split(r"([ \t\n\r\f\v]+)", text) - may_be_url = [":" in word or "//" in word for word in words] - if any("@" in word for word, url in zip(words, may_be_url, strict=True) if url): - return name - cut = "".join( - re.split(r"[?#]", word, maxsplit=1)[0] if url else word for word, url in zip(words, may_be_url, strict=True) - ) - if "//" in text and ("?" in cut or "#" in cut): - return name + if "@" in text: + return f"{name} (message hidden because it may include a credential)" + cut = text.split("?", 1)[0].split("#", 1)[0] return f"{name}: {cut}" if cut else name diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index 0179874ec3..8a78e337d1 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -351,11 +351,14 @@ def test_a_planning_error_fails_only_that_entry(self, tmp_path: Path) -> None: @pytest.mark.parametrize( ("endpoint", "shown"), [ - ("https://evil.example/openai/v1?sig=SYNTHETIC-SAS#frag", "https://evil.example/openai/v1"), - ("https://evil.example/openai/v1'part?sig=SYNTHETIC-SAS#frag", "https://evil.example/openai/v1'part"), - ("https://evil.example/open\u00a0ai/v1?sig=SYNTHETIC-SAS", "https://evil.example/open\u00a0ai/v1"), - ("https://user:hunter2@evil.example/openai/v1?sig=SYNTHETIC-SAS#frag", None), - ("//user:hunter2@evil.example/openai/v1?sig=SYNTHETIC-SAS", None), + ("https://evil.example/openai/v1?sig=SYNTHETIC-SAS#frag", "'https://evil.example/openai/v1"), + ("https://evil.example/openai/v1'part?sig=SYNTHETIC-SAS", "'https://evil.example/openai/v1'part"), + ("https://evil.example/openai/v1?sig=SYNTH\tETIC-SAS", "'https://evil.example/openai/v1"), + ("evil.example/openai/v1?api-version=1 &sig=SYNTHETIC-SAS", "'evil.example/openai/v1"), + ("https://evil.example/openai/v1#SYNTHETIC-SAS", "'https://evil.example/openai/v1"), + ("https://user:hunter2@evil.example/openai/v1?sig=SYNTHETIC-SAS", None), + ("https://user:hun\nter2@evil.example/openai/v1", None), + ("//user:hunter2@evil.example/openai/v1", None), ], ) def test_build_errors_drop_credentials_from_urls(self, tmp_path: Path, endpoint: str, shown: str | None) -> None: @@ -368,12 +371,23 @@ def test_build_errors_drop_credentials_from_urls(self, tmp_path: Path, endpoint: (failure,) = report.failures assert failure.reason is RestoreFailureReason.BUILD_ERROR if shown is None: - assert failure.message == "ValueError" + assert failure.message == "ValueError (message hidden because it may include a credential)" else: - assert f"'{shown}" in failure.message - for secret in ("hunter2", "SYNTHETIC-SAS", "frag"): + assert failure.message.startswith("ValueError: Identity-based authentication requires") + assert shown in failure.message + for secret in ("hunter2", "hun", "ter2", "SYNTHETIC", "SYNTH", "ETIC", "frag"): assert secret not in failure.message + def test_build_errors_keep_text_without_a_credential(self, tmp_path: Path) -> None: + no_endpoint = {**_IDENTITY_TARGET, "params": {"auth_mode": "identity", "model_name": "gpt-4o"}} + _save(tmp_path, RegistryRecipeKind.TARGETS, "server_endpoint", no_endpoint) + + with patch.dict(os.environ, {"OPENAI_CHAT_ENDPOINT": "https://api.example.test/v1"}): + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert "'https://api.example.test/v1'. Pass auth_mode" in failure.message + def test_error_text_that_is_not_a_valid_url_fails_only_that_entry(self, tmp_path: Path) -> None: malformed = { **_IDENTITY_TARGET, @@ -394,22 +408,32 @@ def test_error_text_that_is_not_a_valid_url_fails_only_that_entry(self, tmp_path @pytest.mark.parametrize( ("text", "expected"), [ - ("got ://host/path?sig=SECRET", "ValueError: got ://host/path"), - ("got https://[x?sig=SECRET", "ValueError: got https://[x"), - ("got https:host/path?sig=SECRET#SECRET", "ValueError: got https:host/path"), - ("got https://host/pa\u00a0th?sig=SECRET#SECRET", "ValueError: got https://host/pa\u00a0th"), - ("keeps spacing https://host/?sig=SECRET intact", "ValueError: keeps spacing https://host/ intact"), - ("bad value: 'abc' (expected a number?)", "ValueError: bad value: 'abc' (expected a number?)"), - ("no address here", "ValueError: no address here"), - ("got //user:SECRET@host/path", "ValueError"), - ("got user:SECRET@https://host/path", "ValueError"), - ("got https://user\t:SECRET@host/path", "ValueError"), - ("got https://host/a b?sig=SECRET", "ValueError"), + ("got https://host/v1?sig=SECRET", "ValueError: got https://host/v1"), + ("got https:host/p#SECRET", "ValueError: got https:host/p"), + ("got https://host/v1#a?b=SECRET", "ValueError: got https://host/v1"), + ("got https://host/v1? sig=SEC\tRET tail", "ValueError: got https://host/v1"), + ("got host/v1?api-version=1 &sig=SECRET", "ValueError: got host/v1"), + ("could not be coerced to float: 'hot'", "ValueError: could not be coerced to float: 'hot'"), + ( + "Missing required parameters: ['objective_target']", + "ValueError: Missing required parameters: ['objective_target']", + ), + ("got https://host/v1", "ValueError: got https://host/v1"), + ("got //user:SECRET@host/v1", "ValueError (message hidden because it may include a credential)"), + ("got https://user:SEC\nRET@host/v1", "ValueError (message hidden because it may include a credential)"), + ("contact admin@example.com", "ValueError (message hidden because it may include a credential)"), ], ) - def test_safe_messages_drop_credentials_from_anything_that_may_be_a_url(self, text: str, expected: str) -> None: + def test_safe_messages_drop_queries_fragments_and_user_info(self, text: str, expected: str) -> None: assert RegistryRecipeRestorer._safe_message(ValueError(text)) == expected + def test_safe_message_survives_an_error_that_cannot_be_printed(self) -> None: + class UnprintableError(Exception): + def __str__(self) -> str: + raise RuntimeError("cannot print") + + assert RegistryRecipeRestorer._safe_message(UnprintableError()) == "UnprintableError" + def test_refuses_a_stored_credential_in_attack_parameters(self, tmp_path: Path) -> None: attack_registry = AttackRegistry.get_registry_singleton() original_get_class_metadata = attack_registry.get_class_metadata From 9a419b4ab11205dbd8483c0a8d30ddea782b7fdc Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 03:45:48 +0000 Subject: [PATCH 18/38] Escape text that isn't UTF-8 in restore reports so they can be sent as JSON --- pyrit/models/catalog/registry_recipe.py | 21 +++++++++--- .../registry/test_registry_recipe_storage.py | 34 +++++++++++++++++++ 2 files changed, 51 insertions(+), 4 deletions(-) diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index 5ac95be552..c069ccede5 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -8,7 +8,7 @@ from enum import Enum from typing import Annotated, ClassVar -from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, JsonValue, field_validator from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS, MAX_TEXT_LENGTH @@ -16,6 +16,19 @@ _TechniqueTag = Annotated[str, Field(pattern=r"^[A-Za-z][A-Za-z0-9_]{0,63}$")] +def _escape_invalid_text(value: str) -> str: + """ + Escape characters that can't be written as UTF-8, such as bytes from a file name that isn't UTF-8. + + Returns: + str: The text, safe to send in a JSON response. + """ + return value.encode("utf-8", "backslashreplace").decode("utf-8") + + +_DisplayText = Annotated[str, AfterValidator(_escape_invalid_text)] + + class RegistryRecipeKind(str, Enum): """The registries whose entries can be saved, named after their storage folders.""" @@ -128,7 +141,7 @@ class RestoreFailure(BaseModel): kind: RegistryRecipeKind name: str | None = Field(None, description="Entry name; None when the whole folder couldn't be read") reason: RestoreFailureReason - message: str = Field(..., description="Why, in words safe to show users") + message: _DisplayText = Field(..., description="Why, in words safe to show users") version: str | None = Field(None, description="Version of the stored document, when it could be read") dependency: SavedEntryRef | None = Field(None, description="The entry this one needed that is missing or failed") @@ -139,5 +152,5 @@ class IgnoredRecipeFile(BaseModel): model_config = ConfigDict(frozen=True) kind: RegistryRecipeKind - file_name: str - message: str + file_name: _DisplayText + message: _DisplayText diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index af55c36e64..e7ea4f0d91 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -3,16 +3,20 @@ import hashlib import json +import os from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, patch +import pytest from azure.storage.blob import ContainerClient from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH from pyrit.models.catalog.registry_recipe import ( ComponentRecipe, + IgnoredRecipeFile, RegistryRecipeKind, + RestoreFailure, RestoreFailureReason, TechniqueRecipe, ) @@ -67,6 +71,22 @@ def test_blob_container_source_uses_the_folder_as_prefix(self) -> None: assert list(listing.recipes) == ["formal"] +class TestRestoreReportText: + def test_text_that_is_not_utf8_is_escaped_so_it_can_be_sent_as_json(self) -> None: + ignored = IgnoredRecipeFile( + kind=RegistryRecipeKind.TARGETS, file_name="bad\udcffname.json", message="bad\udcffname" + ) + failure = RestoreFailure( + kind=RegistryRecipeKind.TARGETS, name="x", reason=RestoreFailureReason.BUILD_ERROR, message="got \udcff" + ) + + assert ignored.file_name == "bad\\udcffname.json" + assert failure.message == "got \\udcff" + json.dumps([ignored.model_dump(mode="json"), failure.model_dump(mode="json")], ensure_ascii=False).encode( + "utf-8" + ) + + class TestRegistryRecipeStorageListing: def test_reads_component_recipes_with_their_versions(self, tmp_path: Path) -> None: content = _write(tmp_path / "targets", "contoso_gpt4o", json.dumps(_TARGET)) @@ -136,6 +156,20 @@ def test_reports_a_newer_schema_version_before_checking_fields(self, tmp_path: P assert failure.reason is RestoreFailureReason.NEWER_SCHEMA_VERSION assert "newer PyRIT" in failure.message + def test_reports_a_file_name_that_is_not_utf8(self, tmp_path: Path) -> None: + folder = tmp_path / "targets" + folder.mkdir() + try: + (folder / os.fsdecode(b"bad\xffname.json")).write_text("{}") + except (OSError, UnicodeError): + pytest.skip("This file system only allows UTF-8 file names") + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + (ignored,) = listing.ignored + assert ignored.file_name == "bad\\udcffname.json" + json.dumps(ignored.model_dump(mode="json"), ensure_ascii=False).encode("utf-8") + def test_reports_unreadable_and_ignored_files(self, tmp_path: Path) -> None: folder = tmp_path / "scorers" _write(folder, "locked", json.dumps(_TARGET)) From 9a67367164b61fa4038848d34bcaf56661d50cf8 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 04:00:14 +0000 Subject: [PATCH 19/38] Read error text as plain normalized text and say when part of it is hidden --- pyrit/registry/registry_recipe_restore.py | 18 +++-- .../registry/test_registry_recipe_restore.py | 65 +++++++++++++++++-- 2 files changed, 71 insertions(+), 12 deletions(-) diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index 90181df568..cea24856eb 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -5,6 +5,7 @@ from __future__ import annotations +import unicodedata from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar @@ -504,24 +505,29 @@ def _singular(kind: RegistryRecipeKind) -> str: @staticmethod def _safe_message(error: Exception) -> str: """ - Describe an error for the failure report without any credential it may quote. + Describe an error for the failure report without any credential a URL in it may carry. Constructor errors can quote a server setting, such as an endpoint with a SAS token. The text is cut at the first ``?`` or ``#``, where a URL's query or fragment starts. A user name and - password can't be told apart reliably, so text with an ``@`` isn't shown at all. + password can't be told apart reliably, so text with an ``@`` isn't shown at all. Look-alike + characters, such as a full-width ``@``, count as the characters they look like. Returns: - str: The error's type and its text before any ``?`` or ``#``, or the type alone. + str: The error's type and its text, noting when the rest of the text or all of it is hidden. """ name = type(error).__name__ + hidden = "because it may include a credential" try: - text = str(error) + # A plain str, so no string method a subclass overrides runs on the text below. + text = unicodedata.normalize("NFKC", str.__str__(str(error))) except Exception: return name if "@" in text: - return f"{name} (message hidden because it may include a credential)" + return f"{name} (message hidden {hidden})" cut = text.split("?", 1)[0].split("#", 1)[0] - return f"{name}: {cut}" if cut else name + if cut == text: + return f"{name}: {text}" if text else name + return f"{name}: {cut} (rest hidden {hidden})" if cut.strip() else f"{name} (message hidden {hidden})" class _RecipeRejectedError(Exception): diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index 8a78e337d1..7a3b15ac9d 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -374,7 +374,7 @@ def test_build_errors_drop_credentials_from_urls(self, tmp_path: Path, endpoint: assert failure.message == "ValueError (message hidden because it may include a credential)" else: assert failure.message.startswith("ValueError: Identity-based authentication requires") - assert shown in failure.message + assert failure.message.endswith(f"{shown} (rest hidden because it may include a credential)") for secret in ("hunter2", "hun", "ter2", "SYNTHETIC", "SYNTH", "ETIC", "frag"): assert secret not in failure.message @@ -408,11 +408,31 @@ def test_error_text_that_is_not_a_valid_url_fails_only_that_entry(self, tmp_path @pytest.mark.parametrize( ("text", "expected"), [ - ("got https://host/v1?sig=SECRET", "ValueError: got https://host/v1"), - ("got https:host/p#SECRET", "ValueError: got https:host/p"), - ("got https://host/v1#a?b=SECRET", "ValueError: got https://host/v1"), - ("got https://host/v1? sig=SEC\tRET tail", "ValueError: got https://host/v1"), - ("got host/v1?api-version=1 &sig=SECRET", "ValueError: got host/v1"), + ( + "got https://host/v1?sig=SECRET", + "ValueError: got https://host/v1 (rest hidden because it may include a credential)", + ), + ( + "got https:host/p#SECRET", + "ValueError: got https:host/p (rest hidden because it may include a credential)", + ), + ( + "got https://host/v1#a?b=SECRET", + "ValueError: got https://host/v1 (rest hidden because it may include a credential)", + ), + ( + "got https://host/v1? sig=SEC\tRET tail", + "ValueError: got https://host/v1 (rest hidden because it may include a credential)", + ), + ( + "got host/v1?api-version=1 &sig=SECRET", + "ValueError: got host/v1 (rest hidden because it may include a credential)", + ), + ( + "got 'https://host/v1#SECRET' (again 'https://host/v1#SECRET')", + "ValueError: got 'https://host/v1 (rest hidden because it may include a credential)", + ), + ("?sig=SECRET", "ValueError (message hidden because it may include a credential)"), ("could not be coerced to float: 'hot'", "ValueError: could not be coerced to float: 'hot'"), ( "Missing required parameters: ['objective_target']", @@ -421,6 +441,12 @@ def test_error_text_that_is_not_a_valid_url_fails_only_that_entry(self, tmp_path ("got https://host/v1", "ValueError: got https://host/v1"), ("got //user:SECRET@host/v1", "ValueError (message hidden because it may include a credential)"), ("got https://user:SEC\nRET@host/v1", "ValueError (message hidden because it may include a credential)"), + ("got https://user:SE?CRET@host/v1", "ValueError (message hidden because it may include a credential)"), + ("got https://user:SECRET\uff20host/v1", "ValueError (message hidden because it may include a credential)"), + ( + "got https://host/v1\uff1fsig=SECRET", + "ValueError: got https://host/v1 (rest hidden because it may include a credential)", + ), ("contact admin@example.com", "ValueError (message hidden because it may include a credential)"), ], ) @@ -434,6 +460,33 @@ def __str__(self) -> str: assert RegistryRecipeRestorer._safe_message(UnprintableError()) == "UnprintableError" + def test_safe_message_ignores_string_methods_an_error_overrides(self) -> None: + class KeepEverything(str): + __slots__ = () + + def split(self, *args: Any, **kwargs: Any) -> list[str]: + return [self] + + class Exploding(str): + __slots__ = () + + def __contains__(self, item: object) -> bool: + raise RuntimeError("cannot search") + + class QuotingError(Exception): + def __init__(self, text: str) -> None: + super().__init__() + self.text = text + + def __str__(self) -> str: + return self.text + + keep = RegistryRecipeRestorer._safe_message(QuotingError(KeepEverything("got https://host/v1?sig=SECRET"))) + exploding = RegistryRecipeRestorer._safe_message(QuotingError(Exploding("got https://host/v1"))) + + assert keep == "QuotingError: got https://host/v1 (rest hidden because it may include a credential)" + assert exploding == "QuotingError: got https://host/v1" + def test_refuses_a_stored_credential_in_attack_parameters(self, tmp_path: Path) -> None: attack_registry = AttackRegistry.get_registry_singleton() original_get_class_metadata = attack_registry.get_class_metadata From 274b4ce57129e6a7d271ad3b274628df8f5a6da1 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 05:10:06 +0000 Subject: [PATCH 20/38] Escape entry names in restore reports too --- pyrit/models/catalog/registry_recipe.py | 4 ++-- .../unit/registry/test_registry_recipe_storage.py | 15 +++++++++++++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index c069ccede5..82be6a14e6 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -102,7 +102,7 @@ class SavedEntryRef(BaseModel): model_config = ConfigDict(frozen=True) kind: RegistryRecipeKind - name: str + name: _DisplayText class StoredRecipe(BaseModel): @@ -139,7 +139,7 @@ class RestoreFailure(BaseModel): model_config = ConfigDict(frozen=True) kind: RegistryRecipeKind - name: str | None = Field(None, description="Entry name; None when the whole folder couldn't be read") + name: _DisplayText | None = Field(None, description="Entry name; None when the whole folder couldn't be read") reason: RestoreFailureReason message: _DisplayText = Field(..., description="Why, in words safe to show users") version: str | None = Field(None, description="Version of the stored document, when it could be read") diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index e7ea4f0d91..1f49ac1487 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -18,6 +18,7 @@ RegistryRecipeKind, RestoreFailure, RestoreFailureReason, + SavedEntryRef, TechniqueRecipe, ) from pyrit.registry.registry_recipe_storage import RegistryRecipeStorage @@ -77,15 +78,25 @@ def test_text_that_is_not_utf8_is_escaped_so_it_can_be_sent_as_json(self) -> Non kind=RegistryRecipeKind.TARGETS, file_name="bad\udcffname.json", message="bad\udcffname" ) failure = RestoreFailure( - kind=RegistryRecipeKind.TARGETS, name="x", reason=RestoreFailureReason.BUILD_ERROR, message="got \udcff" + kind=RegistryRecipeKind.TARGETS, + name="x\udcff", + reason=RestoreFailureReason.MISSING_DEPENDENCY, + message="got \udcff", + dependency=SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800"), ) assert ignored.file_name == "bad\\udcffname.json" - assert failure.message == "got \\udcff" + assert (failure.name, failure.message) == ("x\\udcff", "got \\udcff") + assert failure.dependency == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\\ud800") json.dumps([ignored.model_dump(mode="json"), failure.model_dump(mode="json")], ensure_ascii=False).encode( "utf-8" ) + def test_valid_text_is_kept_as_written(self) -> None: + ref = SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="caf\u00e9 \uff10\uff0e\uff15 \U0001f600") + + assert ref.name == "caf\u00e9 \uff10\uff0e\uff15 \U0001f600" + class TestRegistryRecipeStorageListing: def test_reads_component_recipes_with_their_versions(self, tmp_path: Path) -> None: From a60323a9f448156558a0d74211f16488c5781814 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 05:19:22 +0000 Subject: [PATCH 21/38] Show error text with its original characters --- pyrit/registry/registry_recipe_restore.py | 24 +++++++++++++++---- .../registry/test_registry_recipe_restore.py | 13 ++++++++-- 2 files changed, 30 insertions(+), 7 deletions(-) diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index cea24856eb..7ce6288632 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -510,7 +510,8 @@ def _safe_message(error: Exception) -> str: Constructor errors can quote a server setting, such as an endpoint with a SAS token. The text is cut at the first ``?`` or ``#``, where a URL's query or fragment starts. A user name and password can't be told apart reliably, so text with an ``@`` isn't shown at all. Look-alike - characters, such as a full-width ``@``, count as the characters they look like. + characters, such as a full-width ``@``, count as the characters they look like, and the text + that is shown keeps its original characters. Returns: str: The error's type and its text, noting when the rest of the text or all of it is hidden. @@ -519,17 +520,30 @@ def _safe_message(error: Exception) -> str: hidden = "because it may include a credential" try: # A plain str, so no string method a subclass overrides runs on the text below. - text = unicodedata.normalize("NFKC", str.__str__(str(error))) + text = str.__str__(str(error)) + folded = unicodedata.normalize("NFKC", text) except Exception: return name - if "@" in text: + if "@" in folded: return f"{name} (message hidden {hidden})" - cut = text.split("?", 1)[0].split("#", 1)[0] - if cut == text: + if "?" not in folded and "#" not in folded: return f"{name}: {text}" if text else name + cut_at = next((index for index, char in enumerate(text) if _may_start_query(char)), 0) + cut = text[:cut_at] return f"{name}: {cut} (rest hidden {hidden})" if cut.strip() else f"{name} (message hidden {hidden})" +def _may_start_query(char: str) -> bool: + """ + Return whether a character is, or looks like, a ``?`` or ``#``. + + Returns: + bool: True if the character could start a URL's query or fragment. + """ + folded = unicodedata.normalize("NFKC", char) + return "?" in folded or "#" in folded + + class _RecipeRejectedError(Exception): """A saved recipe failed a check before or while it was built.""" diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index 7a3b15ac9d..8df0fdcee3 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -433,6 +433,7 @@ def test_error_text_that_is_not_a_valid_url_fails_only_that_entry(self, tmp_path "ValueError: got 'https://host/v1 (rest hidden because it may include a credential)", ), ("?sig=SECRET", "ValueError (message hidden because it may include a credential)"), + (" \t?sig=SECRET", "ValueError (message hidden because it may include a credential)"), ("could not be coerced to float: 'hot'", "ValueError: could not be coerced to float: 'hot'"), ( "Missing required parameters: ['objective_target']", @@ -448,6 +449,14 @@ def test_error_text_that_is_not_a_valid_url_fails_only_that_entry(self, tmp_path "ValueError: got https://host/v1 (rest hidden because it may include a credential)", ), ("contact admin@example.com", "ValueError (message hidden because it may include a credential)"), + ( + "could not convert '\uff10\uff0e\uff15' or 10\u00b3", + "ValueError: could not convert '\uff10\uff0e\uff15' or 10\u00b3", + ), + ( + "got '\uff15' at https://host/v1\ufe56sig=SECRET", + "ValueError: got '\uff15' at https://host/v1 (rest hidden because it may include a credential)", + ), ], ) def test_safe_messages_drop_queries_fragments_and_user_info(self, text: str, expected: str) -> None: @@ -464,8 +473,8 @@ def test_safe_message_ignores_string_methods_an_error_overrides(self) -> None: class KeepEverything(str): __slots__ = () - def split(self, *args: Any, **kwargs: Any) -> list[str]: - return [self] + def __getitem__(self, key: Any) -> str: + return str.__str__(self) class Exploding(str): __slots__ = () From e7558a41ed804339c4f6fb9cd80b241a5dbd3752 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 05:22:58 +0000 Subject: [PATCH 22/38] Test that a reference name that isn't valid text can still be reported --- .../registry/test_registry_recipe_restore.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index 8df0fdcee3..3683e0674f 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -496,6 +496,22 @@ def __str__(self) -> str: assert keep == "QuotingError: got https://host/v1 (rest hidden because it may include a credential)" assert exploding == "QuotingError: got https://host/v1" + def test_reports_a_reference_name_that_is_not_valid_text(self, tmp_path: Path) -> None: + _save_text( + tmp_path, + RegistryRecipeKind.TARGETS, + "pool", + '{"class_module": "pyrit.prompt_target.round_robin_target", "class_name": "RoundRobinTarget", ' + '"params": {"targets": ["ghost\\ud800"]}}', + ) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert failure.reason is RestoreFailureReason.MISSING_DEPENDENCY + assert failure.dependency == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\\ud800") + json.dumps(report.model_dump(mode="json"), ensure_ascii=False).encode("utf-8") + def test_refuses_a_stored_credential_in_attack_parameters(self, tmp_path: Path) -> None: attack_registry = AttackRegistry.get_registry_singleton() original_get_class_metadata = attack_registry.get_class_metadata From 686d47564d98f034abe0ac87614ff08759bf18d0 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 05:26:09 +0000 Subject: [PATCH 23/38] Test that the runtime status can report a reference name that isn't valid text --- tests/unit/backend/test_runtime_lifecycle.py | 24 ++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/unit/backend/test_runtime_lifecycle.py b/tests/unit/backend/test_runtime_lifecycle.py index a01958004a..78d7ca596f 100644 --- a/tests/unit/backend/test_runtime_lifecycle.py +++ b/tests/unit/backend/test_runtime_lifecycle.py @@ -958,6 +958,30 @@ async def test_startup_with_an_unrestorable_entry_is_not_ready_async(runtime: Ru assert failure["version"] +async def test_status_reports_a_reference_name_that_is_not_valid_text_async( + runtime: RuntimeLifecycle, tmp_path: Path +) -> None: + recipes = tmp_path / "recipes" + _save_recipe( + recipes, + kind="targets", + name="pool", + document={ + "class_module": "pyrit.prompt_target.round_robin_target", + "class_name": "RoundRobinTarget", + "params": {"targets": ["missing\udcff"]}, + }, + ) + + await _start_with_recipes_async(runtime, recipes=recipes) + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=runtime.app), base_url="http://test") as client: + status = await client.get("/api/config/runtime") + assert status.status_code == 200 + (failure,) = status.json()["restore_failures"] + assert (failure["reason"], failure["dependency"]["name"]) == ("missing_dependency", "missing\\udcff") + + async def test_startup_lists_ignored_files_without_blocking_readiness_async( runtime: RuntimeLifecycle, tmp_path: Path ) -> None: From c391b0d81124d64a9c44aa7968348742a1794e69 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 07:06:15 +0000 Subject: [PATCH 24/38] Escape report text only when it is sent as JSON --- pyrit/models/catalog/registry_recipe.py | 5 ++-- .../registry/test_registry_recipe_storage.py | 30 +++++++++++++------ 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index 82be6a14e6..43dd832ef1 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -8,7 +8,7 @@ from enum import Enum from typing import Annotated, ClassVar -from pydantic import AfterValidator, BaseModel, ConfigDict, Field, JsonValue, field_validator +from pydantic import BaseModel, ConfigDict, Field, JsonValue, PlainSerializer, field_validator from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS, MAX_TEXT_LENGTH @@ -26,7 +26,8 @@ def _escape_invalid_text(value: str) -> str: return value.encode("utf-8", "backslashreplace").decode("utf-8") -_DisplayText = Annotated[str, AfterValidator(_escape_invalid_text)] +# The model keeps the text it was given, so a name still matches the entry it names; only JSON output is escaped. +_DisplayText = Annotated[str, PlainSerializer(_escape_invalid_text, return_type=str, when_used="json")] class RegistryRecipeKind(str, Enum): diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index 1f49ac1487..49ef62c6e1 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -73,7 +73,7 @@ def test_blob_container_source_uses_the_folder_as_prefix(self) -> None: class TestRestoreReportText: - def test_text_that_is_not_utf8_is_escaped_so_it_can_be_sent_as_json(self) -> None: + def test_text_that_is_not_utf8_is_escaped_when_sent_as_json(self) -> None: ignored = IgnoredRecipeFile( kind=RegistryRecipeKind.TARGETS, file_name="bad\udcffname.json", message="bad\udcffname" ) @@ -85,17 +85,28 @@ def test_text_that_is_not_utf8_is_escaped_so_it_can_be_sent_as_json(self) -> Non dependency=SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800"), ) - assert ignored.file_name == "bad\\udcffname.json" - assert (failure.name, failure.message) == ("x\\udcff", "got \\udcff") - assert failure.dependency == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\\ud800") - json.dumps([ignored.model_dump(mode="json"), failure.model_dump(mode="json")], ensure_ascii=False).encode( - "utf-8" + sent_ignored, sent_failure = json.loads( + json.dumps([ignored.model_dump(mode="json"), failure.model_dump(mode="json")], ensure_ascii=False).encode( + "utf-8" + ) ) + assert (sent_ignored["file_name"], sent_ignored["message"]) == ("bad\\udcffname.json", "bad\\udcffname") + assert (sent_failure["name"], sent_failure["message"]) == ("x\\udcff", "got \\udcff") + assert sent_failure["dependency"]["name"] == "ghost\\ud800" + assert json.loads(failure.model_dump_json())["dependency"]["name"] == "ghost\\ud800" + + def test_names_keep_the_text_they_were_given(self) -> None: + ref = SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800") + + assert ref.name == "ghost\ud800" + assert ref == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800") + assert ref != SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\\ud800") + def test_valid_text_is_kept_as_written(self) -> None: ref = SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="caf\u00e9 \uff10\uff0e\uff15 \U0001f600") - assert ref.name == "caf\u00e9 \uff10\uff0e\uff15 \U0001f600" + assert ref.model_dump(mode="json")["name"] == "caf\u00e9 \uff10\uff0e\uff15 \U0001f600" class TestRegistryRecipeStorageListing: @@ -178,8 +189,9 @@ def test_reports_a_file_name_that_is_not_utf8(self, tmp_path: Path) -> None: listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() (ignored,) = listing.ignored - assert ignored.file_name == "bad\\udcffname.json" - json.dumps(ignored.model_dump(mode="json"), ensure_ascii=False).encode("utf-8") + assert ignored.file_name == os.fsdecode(b"bad\xffname.json") + sent = json.loads(json.dumps(ignored.model_dump(mode="json"), ensure_ascii=False).encode("utf-8")) + assert sent["file_name"] == "bad\\udcffname.json" def test_reports_unreadable_and_ignored_files(self, tmp_path: Path) -> None: folder = tmp_path / "scorers" From f986da5ceaf0fc04826112d99a23c971ee6e5189 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 07:13:17 +0000 Subject: [PATCH 25/38] Check each character on its own so long error messages stay fast --- pyrit/registry/registry_recipe_restore.py | 21 ++++++++++--------- .../registry/test_registry_recipe_restore.py | 20 ++++++++++++++++++ 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index 7ce6288632..e2e8ee46f7 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -509,8 +509,8 @@ def _safe_message(error: Exception) -> str: Constructor errors can quote a server setting, such as an endpoint with a SAS token. The text is cut at the first ``?`` or ``#``, where a URL's query or fragment starts. A user name and - password can't be told apart reliably, so text with an ``@`` isn't shown at all. Look-alike - characters, such as a full-width ``@``, count as the characters they look like, and the text + password can't be told apart reliably, so text with an ``@`` isn't shown at all. A look-alike + character, such as a full-width ``@``, counts as the character it looks like, and the text that is shown keeps its original characters. Returns: @@ -521,27 +521,28 @@ def _safe_message(error: Exception) -> str: try: # A plain str, so no string method a subclass overrides runs on the text below. text = str.__str__(str(error)) - folded = unicodedata.normalize("NFKC", text) except Exception: return name - if "@" in folded: + if any(_looks_like(char, "@") for char in text): return f"{name} (message hidden {hidden})" - if "?" not in folded and "#" not in folded: + cut_at = next((index for index, char in enumerate(text) if _looks_like(char, "?#")), None) + if cut_at is None: return f"{name}: {text}" if text else name - cut_at = next((index for index, char in enumerate(text) if _may_start_query(char)), 0) cut = text[:cut_at] return f"{name}: {cut} (rest hidden {hidden})" if cut.strip() else f"{name} (message hidden {hidden})" -def _may_start_query(char: str) -> bool: +def _looks_like(char: str, marks: str) -> bool: """ - Return whether a character is, or looks like, a ``?`` or ``#``. + Return whether a character is, or looks like, one of the given marks. + + Each character is folded on its own, which keeps the check linear in the length of the text. Returns: - bool: True if the character could start a URL's query or fragment. + bool: True if the character's compatibility form contains one of the marks. """ folded = unicodedata.normalize("NFKC", char) - return "?" in folded or "#" in folded + return any(mark in folded for mark in marks) class _RecipeRejectedError(Exception): diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index 81ad7488a1..3ded02f8a4 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -5,6 +5,7 @@ import json import os import sys +import unicodedata from collections.abc import Iterator from pathlib import Path from typing import Any @@ -476,6 +477,25 @@ def test_error_text_that_is_not_a_valid_url_fails_only_that_entry(self, tmp_path def test_safe_messages_drop_queries_fragments_and_user_info(self, text: str, expected: str) -> None: assert RegistryRecipeRestorer._safe_message(ValueError(text)) == expected + @pytest.mark.parametrize( + "mark", ["\u2047", "\u2048", "\u2049", "\ufe16", "\ufe56", "\ufe5f", "\uff03", "\uff1f", "\ufe6b", "\uff20"] + ) + def test_safe_message_treats_every_look_alike_as_the_real_mark(self, mark: str) -> None: + message = RegistryRecipeRestorer._safe_message(ValueError(f"got https://host/v1{mark}SECRET")) + + assert "SECRET" not in message + assert message.endswith("because it may include a credential)") + + def test_safe_message_folds_one_character_at_a_time(self) -> None: + marks = "\u0315" * 50 + "\u0300" * 50 + with patch( + "pyrit.registry.registry_recipe_restore.unicodedata.normalize", wraps=unicodedata.normalize + ) as normalize: + message = RegistryRecipeRestorer._safe_message(ValueError(f"a{marks}\uff1fsig=SECRET")) + + assert message == f"ValueError: a{marks} (rest hidden because it may include a credential)" + assert {len(call.args[1]) for call in normalize.call_args_list} == {1} + def test_safe_message_survives_an_error_that_cannot_be_printed(self) -> None: class UnprintableError(Exception): def __str__(self) -> str: From c86bfe620faf985454c95298bc1985b61897257f Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 09:05:02 +0000 Subject: [PATCH 26/38] Check the exact message for every form of each mark and bound the work per character --- .../registry/test_registry_recipe_restore.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index 3ded02f8a4..99dc130dbf 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -478,23 +478,30 @@ def test_safe_messages_drop_queries_fragments_and_user_info(self, text: str, exp assert RegistryRecipeRestorer._safe_message(ValueError(text)) == expected @pytest.mark.parametrize( - "mark", ["\u2047", "\u2048", "\u2049", "\ufe16", "\ufe56", "\ufe5f", "\uff03", "\uff1f", "\ufe6b", "\uff20"] + "mark", ["?", "#", "\u2047", "\u2048", "\u2049", "\ufe16", "\ufe56", "\ufe5f", "\uff03", "\uff1f"] ) - def test_safe_message_treats_every_look_alike_as_the_real_mark(self, mark: str) -> None: + def test_safe_message_cuts_at_any_form_of_a_query_or_fragment_mark(self, mark: str) -> None: message = RegistryRecipeRestorer._safe_message(ValueError(f"got https://host/v1{mark}SECRET")) - assert "SECRET" not in message - assert message.endswith("because it may include a credential)") + assert message == "ValueError: got https://host/v1 (rest hidden because it may include a credential)" + + @pytest.mark.parametrize("mark", ["@", "\ufe6b", "\uff20"]) + def test_safe_message_hides_text_with_any_form_of_an_at_sign(self, mark: str) -> None: + message = RegistryRecipeRestorer._safe_message(ValueError(f"got https://USER:SECRET{mark}host/v1")) + + assert message == "ValueError (message hidden because it may include a credential)" def test_safe_message_folds_one_character_at_a_time(self) -> None: marks = "\u0315" * 50 + "\u0300" * 50 + text = f"a{marks}\uff1fsig=SECRET" with patch( "pyrit.registry.registry_recipe_restore.unicodedata.normalize", wraps=unicodedata.normalize ) as normalize: - message = RegistryRecipeRestorer._safe_message(ValueError(f"a{marks}\uff1fsig=SECRET")) + message = RegistryRecipeRestorer._safe_message(ValueError(text)) assert message == f"ValueError: a{marks} (rest hidden because it may include a credential)" assert {len(call.args[1]) for call in normalize.call_args_list} == {1} + assert normalize.call_count <= 2 * len(text) def test_safe_message_survives_an_error_that_cannot_be_printed(self) -> None: class UnprintableError(Exception): From 9f9128278a3a276a989fc33499356781d53e3c5d Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 09:09:54 +0000 Subject: [PATCH 27/38] Test that the runtime status can list a file name that isn't valid text --- tests/unit/backend/test_runtime_lifecycle.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/unit/backend/test_runtime_lifecycle.py b/tests/unit/backend/test_runtime_lifecycle.py index 78d7ca596f..87a4a2959c 100644 --- a/tests/unit/backend/test_runtime_lifecycle.py +++ b/tests/unit/backend/test_runtime_lifecycle.py @@ -982,6 +982,22 @@ async def test_status_reports_a_reference_name_that_is_not_valid_text_async( assert (failure["reason"], failure["dependency"]["name"]) == ("missing_dependency", "missing\\udcff") +async def test_status_lists_a_file_name_that_is_not_valid_text_async(runtime: RuntimeLifecycle, tmp_path: Path) -> None: + folder = tmp_path / "recipes" / "converters" + folder.mkdir(parents=True) + try: + (folder / os.fsdecode(b"bad\xffname.json")).write_text("{}") + except (OSError, UnicodeError): + pytest.skip("This file system only allows UTF-8 file names") + + await _start_with_recipes_async(runtime, recipes=tmp_path / "recipes") + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=runtime.app), base_url="http://test") as client: + status = await client.get("/api/config/runtime") + assert status.status_code == 200 + assert [item["file_name"] for item in status.json()["restore_ignored"]] == ["bad\\udcffname.json"] + + async def test_startup_lists_ignored_files_without_blocking_readiness_async( runtime: RuntimeLifecycle, tmp_path: Path ) -> None: From a4554943caec42d4934c39500fa5c758affd2604 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 10:22:28 +0000 Subject: [PATCH 28/38] Pass the look-alike check's arguments by keyword --- pyrit/registry/registry_recipe_restore.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index e2e8ee46f7..0b147f4198 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -523,16 +523,16 @@ def _safe_message(error: Exception) -> str: text = str.__str__(str(error)) except Exception: return name - if any(_looks_like(char, "@") for char in text): + if any(_looks_like(char=char, marks="@") for char in text): return f"{name} (message hidden {hidden})" - cut_at = next((index for index, char in enumerate(text) if _looks_like(char, "?#")), None) + cut_at = next((index for index, char in enumerate(text) if _looks_like(char=char, marks="?#")), None) if cut_at is None: return f"{name}: {text}" if text else name cut = text[:cut_at] return f"{name}: {cut} (rest hidden {hidden})" if cut.strip() else f"{name} (message hidden {hidden})" -def _looks_like(char: str, marks: str) -> bool: +def _looks_like(*, char: str, marks: str) -> bool: """ Return whether a character is, or looks like, one of the given marks. From ce8b6843f9c08602a94c0578298702bf3fc396f4 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 10:28:49 +0000 Subject: [PATCH 29/38] Test that a technique waits for a saved converter in either list and reports an unknown attack type --- .../registry/test_registry_recipe_restore.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index 99dc130dbf..41fca00802 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -15,6 +15,7 @@ from pydantic import ValidationError from pyrit.backend.models.techniques import CreateTechniqueRequest +from pyrit.converter import Base64Converter from pyrit.models.catalog.registry_recipe import ( RegistryRecipeKind, RestoreFailure, @@ -210,6 +211,27 @@ def test_a_technique_waits_for_its_saved_adversarial_target(self, tmp_path: Path assert failure.reason is RestoreFailureReason.FAILED_DEPENDENCY assert failure.dependency == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="contoso_gpt4o") + @pytest.mark.parametrize("field", ["request_converters", "response_converters"]) + def test_a_technique_waits_for_a_saved_converter_in_either_list(self, tmp_path: Path, field: str) -> None: + ConverterRegistry.get_registry_singleton().instances.register(Base64Converter(), name="shared_name") + _save_text(tmp_path, RegistryRecipeKind.CONVERTERS, "shared_name", "{not json") + _save( + tmp_path, RegistryRecipeKind.TECHNIQUES, "uses_shared", {"type": "RedTeamingAttack", field: ["shared_name"]} + ) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + failure = _failures_by_name(report.failures)["uses_shared"] + assert failure.reason is RestoreFailureReason.FAILED_DEPENDENCY + assert failure.dependency == SavedEntryRef(kind=RegistryRecipeKind.CONVERTERS, name="shared_name") + + def test_reports_an_unknown_attack_type(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "unknown_attack", {"type": "NoSuchAttack"}) + + (failure,) = RegistryRecipeRestorer(source=str(tmp_path)).restore().failures + + assert (failure.name, failure.reason) == ("unknown_attack", RestoreFailureReason.UNKNOWN_CLASS) + def test_technique_parameters_are_checked_as_strict_json(self, tmp_path: Path) -> None: _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) without_converters = {key: value for key, value in _TECHNIQUE.items() if key != "request_converters"} From 74412fef5d26e534c6a7d945fdce855eb50e70f5 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 10:30:22 +0000 Subject: [PATCH 30/38] Test that restore runs after the startup scripts, at startup and on apply --- tests/unit/backend/test_runtime_lifecycle.py | 74 +++++++++++++++++++- 1 file changed, 71 insertions(+), 3 deletions(-) diff --git a/tests/unit/backend/test_runtime_lifecycle.py b/tests/unit/backend/test_runtime_lifecycle.py index 87a4a2959c..7ac560f20d 100644 --- a/tests/unit/backend/test_runtime_lifecycle.py +++ b/tests/unit/backend/test_runtime_lifecycle.py @@ -6,7 +6,7 @@ import asyncio import json import os -from collections.abc import Generator +from collections.abc import Callable, Generator from contextlib import ExitStack from datetime import UTC, datetime from pathlib import Path @@ -30,6 +30,9 @@ from pyrit.memory import CentralMemory, MemoryInterface from pyrit.models import ScenarioRunState from pyrit.models.catalog.scenario import RunScenarioRequest +from pyrit.prompt_target import OpenAIChatTarget +from pyrit.registry.components.converter_registry import ConverterRegistry +from pyrit.registry.components.target_registry import TargetRegistry from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer from pyrit.scenario import Scenario from pyrit.setup.configuration_loader import ConfigurationLoader @@ -929,16 +932,81 @@ def _save_recipe(root: Path, *, kind: str, name: str, document: dict[str, object return path -async def _start_with_recipes_async(runtime: RuntimeLifecycle, *, recipes: Path) -> None: +async def _start_with_recipes_async( + runtime: RuntimeLifecycle, *, recipes: Path, initialize: Callable[..., object] | None = None +) -> None: config = await runtime._load_async() config.registry_recipes_source = str(recipes) with ( - patch.object(config, "initialize_pyrit_async", AsyncMock()), + patch.object(config, "initialize_pyrit_async", AsyncMock(side_effect=initialize)), patch.object(lifecycle_module, "get_scenario_run_service", return_value=MagicMock(spec=ScenarioRunService)), ): await runtime.startup_async() +@pytest.fixture +def startup_chat() -> Generator[OpenAIChatTarget, None, None]: + TargetRegistry.reset_registry_singleton() + ConverterRegistry.reset_registry_singleton() + yield OpenAIChatTarget( + endpoint="https://startup.openai.azure.com/openai/v1", model_name="gpt-4o", auth_mode="identity" + ) + TargetRegistry.reset_registry_singleton() + ConverterRegistry.reset_registry_singleton() + + +def _save_converter_using_startup_chat(recipes: Path) -> None: + _save_recipe( + recipes, + kind="converters", + name="formal", + document={ + "class_module": "pyrit.converter.tone_converter", + "class_name": "ToneConverter", + "params": {"converter_target": "startup_chat", "tone": "formal"}, + }, + ) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_startup_restores_an_entry_that_uses_a_startup_target_async( + runtime: RuntimeLifecycle, tmp_path: Path, startup_chat: OpenAIChatTarget +) -> None: + recipes = tmp_path / "recipes" + _save_converter_using_startup_chat(recipes) + targets = TargetRegistry.get_registry_singleton() + + await _start_with_recipes_async( + runtime, recipes=recipes, initialize=lambda **_: targets.instances.register(startup_chat, name="startup_chat") + ) + + assert (runtime.state, runtime.outcome) == ("ready", "success") + converter = ConverterRegistry.get_registry_singleton().instances.get("formal") + assert converter is not None + assert converter._converter_target is startup_chat + + +@pytest.mark.usefixtures("patch_central_database") +async def test_apply_restores_an_entry_that_uses_a_startup_target_async( + runtime: RuntimeLifecycle, tmp_path: Path, startup_chat: OpenAIChatTarget +) -> None: + recipes = tmp_path / "recipes" + _save_converter_using_startup_chat(recipes) + targets = TargetRegistry.get_registry_singleton() + config = await runtime._load_async() + config.registry_recipes_source = str(recipes) + config.apply_prepared_reinitialization_async.side_effect = lambda **_: targets.instances.register( + startup_chat, name="startup_chat" + ) + + await apply_async(runtime) + + assert (runtime.state, runtime.outcome) == ("ready", "success") + converter = ConverterRegistry.get_registry_singleton().instances.get("formal") + assert converter is not None + assert converter._converter_target is startup_chat + + async def test_startup_with_an_unrestorable_entry_is_not_ready_async(runtime: RuntimeLifecycle, tmp_path: Path) -> None: recipes = tmp_path / "recipes" _save_recipe(recipes, kind="targets", name="broken", document={"class_module": "nowhere", "class_name": "Nope"}) From 784cd45bc916a1a34357cd218516cad71a785c39 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 10:31:23 +0000 Subject: [PATCH 31/38] Keep the offline browser fixture away from saved registry entries --- frontend/e2e/fixtures/manual_send_backend.py | 30 +++++++++++--------- pyrit/backend/README.md | 5 ++-- 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/frontend/e2e/fixtures/manual_send_backend.py b/frontend/e2e/fixtures/manual_send_backend.py index 03ae968b36..52c157703d 100644 --- a/frontend/e2e/fixtures/manual_send_backend.py +++ b/frontend/e2e/fixtures/manual_send_backend.py @@ -4,6 +4,7 @@ """Offline backend for manual-send browser tests, with no real configuration or providers.""" import os +import tempfile from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from unittest.mock import AsyncMock, patch @@ -28,19 +29,22 @@ async def _initialize_offline_async(self: ConfigurationLoader, *, raise_on_initi @asynccontextmanager async def _offline_lifespan_async(application: FastAPI) -> AsyncGenerator[None, None]: - config = ConfigurationLoader(memory_db_type="in_memory", env_files=[], env_akv_ref=[]) - with ( - patch.dict(os.environ, {"PYRIT_AUTH_MODE": "local", "PYRIT_DEV_MODE": "true"}, clear=True), - patch.object( - ConfigurationFileService, - "_read_source_async", - new=AsyncMock(return_value="memory_db_type: in_memory\nenv_files: []\ninitializers: []\n"), - ), - patch.object(ConfigurationLoader, "load_with_overrides", return_value=config), - patch.object(ConfigurationLoader, "initialize_pyrit_async", new=_initialize_offline_async), - ): - async with lifespan(application): - yield + with tempfile.TemporaryDirectory() as recipes: + config = ConfigurationLoader( + memory_db_type="in_memory", env_files=[], env_akv_ref=[], registry_recipes_source=recipes + ) + with ( + patch.dict(os.environ, {"PYRIT_AUTH_MODE": "local", "PYRIT_DEV_MODE": "true"}, clear=True), + patch.object( + ConfigurationFileService, + "_read_source_async", + new=AsyncMock(return_value="memory_db_type: in_memory\nenv_files: []\ninitializers: []\n"), + ), + patch.object(ConfigurationLoader, "load_with_overrides", return_value=config), + patch.object(ConfigurationLoader, "initialize_pyrit_async", new=_initialize_offline_async), + ): + async with lifespan(application): + yield app.router.lifespan_context = _offline_lifespan_async diff --git a/pyrit/backend/README.md b/pyrit/backend/README.md index c914f20503..52d458cb00 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -264,8 +264,9 @@ The chat preserves failed drafts across replacement, but requires converter choi to be reviewed again rather than restoring outputs tied to the previous registry. The offline browser fixture `frontend.e2e.fixtures.manual_send_backend:app` uses the real -backend lifecycle with isolated in-memory SQLite, no environment files, and no default -providers. The recovery tests register deterministic loopback targets. For example, start +backend lifecycle with isolated in-memory SQLite, no environment files, no saved registry +entries, and no default providers. The recovery tests register deterministic loopback +targets. For example, start it with `uv run python -m uvicorn frontend.e2e.fixtures.manual_send_backend:app --host 127.0.0.1 --port 18213` (set `PYTHONUTF8=1` on Windows). In another shell, set `PYRIT_BACKEND_URL=http://127.0.0.1:18213` and `E2E_FRONTEND_PORT=31213`, then run `npx playwright test chat-recovery --project seeded --workers 1` From 507567b7be04f37b4b88e7c7d3f18261aea1c10a Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 21:03:42 +0000 Subject: [PATCH 32/38] Check the opened file before reading it and skip links deleted since the listing --- pyrit/registry/file_document_storage.py | 39 +++++++- .../registry/test_file_document_storage.py | 99 +++++++++++++++---- 2 files changed, 115 insertions(+), 23 deletions(-) diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index 4ba3f0949d..c025c77eff 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -16,6 +16,7 @@ import hashlib import logging import os +import stat import sys import tempfile import time @@ -109,6 +110,27 @@ def _release_exclusive_lock(descriptor: int) -> None: fcntl.flock(descriptor, fcntl.LOCK_UN) +# Opening without blocking keeps a pipe from hanging the read; it changes nothing for a regular +# file. Windows has no O_NONBLOCK and needs O_BINARY so the bytes aren't translated. +_READ_FLAGS = os.O_RDONLY | getattr(os, "O_NONBLOCK", 0) | getattr(os, "O_BINARY", 0) + + +def _read_regular_file(path: str) -> bytes | None: + """ + Read a file only if what was opened is a regular file. + + The opened file is checked rather than the path, so the check still holds if the path + was replaced after it was listed. + + Returns: + bytes | None: The content, or None if the path names something other than a regular file. + """ + with open(os.open(path, _READ_FLAGS), "rb") as file: + if not stat.S_ISREG(os.fstat(file.fileno()).st_mode): + return None + return file.read() + + class DocumentConflictError(ValueError): """A stored document changed after the caller read it.""" @@ -279,19 +301,26 @@ def _read_local_entry(self, *, entry: os.DirEntry[str], name: str) -> bytes | Un """ Read one listed local document. + The directory can change after it is listed, so what is opened is checked again: an entry + deleted in between is skipped, whether or not it was a link, and anything that isn't a regular + file is reported without being read. + Returns: bytes | UnreadableDocument | None: The content, why it can't be read, or None if it was deleted after the directory was listed. """ try: - if not entry.is_file(): - # A broken link, a folder, or a pipe, which would block if read. - return UnreadableDocument(name=name, message="Not a regular file") - return Path(entry.path).read_bytes() + content = _read_regular_file(entry.path) if entry.is_file() else None except FileNotFoundError: - return None + content = None except OSError as error: return UnreadableDocument(name=name, message=self._describe_read_error(error)) + if content is not None: + return content + if not os.path.lexists(entry.path): + return None + # A broken link, a folder, or a pipe, which would block if read. + return UnreadableDocument(name=name, message="Not a regular file") def _list_blob_documents_with_issues(self) -> DocumentListing: """ diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index 63bc7601dc..a6e264a45f 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -1,7 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from collections.abc import Iterator +import os +from collections.abc import Callable, Generator, Iterator +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -10,7 +12,12 @@ from azure.core.exceptions import HttpResponseError, ResourceNotFoundError, ServiceRequestError from azure.storage.blob import ContainerClient, StorageErrorCode -from pyrit.registry.file_document_storage import DocumentListing, FileDocumentStorage, UnreadableDocument +from pyrit.registry.file_document_storage import ( + DocumentListing, + FileDocumentStorage, + UnreadableDocument, + _read_regular_file, +) _SAS_SOURCE = "https://account.blob.core.windows.net/documents/saved?sp=rl&sig=secret-signature" @@ -30,18 +37,46 @@ def _blob_client(*, names: list[str]) -> MagicMock: return client +@contextmanager +def _change_after_listing(*, change: Callable[[], None], cache_types: bool = False) -> Generator[None, None, None]: + # Runs `change` after the folder is listed and before its entries are read. With `cache_types`, + # each entry's type is looked up first, so it is out of date once `change` runs. + real_scandir = os.scandir + + @contextmanager + def scandir(path: Path) -> Generator[Iterator[os.DirEntry[str]], None, None]: + with real_scandir(path) as entries: + listed = list(entries) + if cache_types: + for entry in listed: + entry.is_file() + yield iter(listed) + change() + + with patch("pyrit.registry.file_document_storage.os.scandir", side_effect=scandir): + yield + + class TestLocalListingWithIssues: def test_reports_unreadable_document_and_keeps_the_others(self, tmp_path: Path) -> None: (tmp_path / "first.json").write_bytes(b"one") (tmp_path / "second.json").write_bytes(b"two") original_read_bytes = Path.read_bytes + def read(path: str) -> bytes | None: + if Path(path).name == "second.json": + raise PermissionError(13, "Permission denied", path) + return _read_regular_file(path) + def read_bytes(path: Path) -> bytes: if path.name == "second.json": raise PermissionError(13, "Permission denied", str(path)) return original_read_bytes(path) - with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + with ( + patch("pyrit.registry.file_document_storage._read_regular_file", side_effect=read), + patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes), + ): listing = _local_storage(tmp_path)._list_documents_with_issues() legacy = _local_storage(tmp_path)._list_documents() @@ -58,13 +93,12 @@ def test_reports_badly_named_files_as_ignored_without_reading_them(self, tmp_pat (tmp_path / "nested").mkdir() (tmp_path / "nested" / "inner.json").write_bytes(b"not a direct child") read_names: list[str] = [] - original_read_bytes = Path.read_bytes - def read_bytes(path: Path) -> bytes: - read_names.append(path.name) - return original_read_bytes(path) + def read(path: str) -> bytes | None: + read_names.append(Path(path).name) + return _read_regular_file(path) - with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + with patch("pyrit.registry.file_document_storage._read_regular_file", side_effect=read): listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing.documents == {"valid_name": b"ok"} @@ -75,20 +109,49 @@ def read_bytes(path: Path) -> bytes: assert _local_storage(tmp_path)._list_documents() == listing.documents def test_skips_a_document_deleted_after_listing(self, tmp_path: Path) -> None: - (tmp_path / "kept.json").write_bytes(b"kept") - (tmp_path / "deleted.json").write_bytes(b"gone") - original_read_bytes = Path.read_bytes + (tmp_path / "kept.json").write_bytes(b"{}\r\n") + (tmp_path / "empty.json").write_bytes(b"") + deleted = tmp_path / "deleted.json" + deleted.write_bytes(b"gone") - def read_bytes(path: Path) -> bytes: - if path.name == "deleted.json": - raise FileNotFoundError(2, "No such file or directory", str(path)) - return original_read_bytes(path) + with _change_after_listing(change=deleted.unlink, cache_types=True): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing(documents={"empty": b"", "kept": b"{}\r\n"}) - with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + def test_skips_a_link_deleted_after_listing(self, tmp_path: Path) -> None: + (tmp_path / "kept.json").write_bytes(b"kept") + (tmp_path / "target.txt").write_bytes(b"linked") + link = tmp_path / "linked.json" + try: + link.symlink_to(tmp_path / "target.txt") + except OSError: + pytest.skip("Creating symbolic links isn't permitted here") + + with _change_after_listing(change=link.unlink): listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing == DocumentListing(documents={"kept": b"kept"}) + @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + @pytest.mark.timeout(60) + def test_reports_a_pipe_put_in_place_of_a_listed_document_without_blocking(self, tmp_path: Path) -> None: + (tmp_path / "kept.json").write_bytes(b"kept") + swapped = tmp_path / "swapped.json" + swapped.write_bytes(b"replaced") + + def replace_with_pipe() -> None: + swapped.unlink() + os.mkfifo(swapped) + + with _change_after_listing(change=replace_with_pipe, cache_types=True): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing( + documents={"kept": b"kept"}, + unreadable=(UnreadableDocument(name="swapped", message="Not a regular file"),), + ) + def test_reports_entries_that_are_not_regular_files_without_reading_them(self, tmp_path: Path) -> None: (tmp_path / "folder.json").mkdir() try: @@ -96,14 +159,14 @@ def test_reports_entries_that_are_not_regular_files_without_reading_them(self, t except OSError: pytest.skip("Creating symbolic links isn't permitted here") - with patch.object(Path, "read_bytes", autospec=True) as read_bytes: + with patch("pyrit.registry.file_document_storage._read_regular_file") as read: listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing.unreadable == ( UnreadableDocument(name="broken_link", message="Not a regular file"), UnreadableDocument(name="folder", message="Not a regular file"), ) - read_bytes.assert_not_called() + read.assert_not_called() def test_reports_a_link_loop_and_keeps_the_others(self, tmp_path: Path) -> None: (tmp_path / "good.json").write_bytes(b"ok") From de2f4e3b6856ac819f45e89eb8d06190d86b01d4 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 21:35:07 +0000 Subject: [PATCH 33/38] Let restore set the path parameters a target lets callers set, as the API does --- pyrit/registry/registry_recipe_restore.py | 13 ++++++++-- .../registry/test_registry_recipe_restore.py | 25 +++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index 0b147f4198..2d6578d882 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -393,8 +393,17 @@ def _build(self, *, stored: StoredRecipe) -> RestoreFailure | None: if isinstance(recipe, TechniqueRecipe): self._build_technique(name=stored.ref.name, recipe=recipe) else: - self._get_registry(stored.ref.kind).create_named_instance( - name=stored.ref.name, type_name=recipe.class_name, params=recipe.params, external_input=True + registry = self._get_registry(stored.ref.kind) + # As in the API, a target may let callers set some path parameters, such as a working folder. + allowed_path_params = getattr( + registry.get_class(recipe.class_name), "caller_path_parameters", frozenset() + ) + registry.create_named_instance( + name=stored.ref.name, + type_name=recipe.class_name, + params=recipe.params, + external_input=True, + allowed_path_params=allowed_path_params, ) except _RecipeRejectedError as rejection: return self._failure(stored=stored, reason=rejection.reason, message=str(rejection)) diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index 41fca00802..ff58ac6d92 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -8,6 +8,7 @@ import unicodedata from collections.abc import Iterator from pathlib import Path +from types import ModuleType from typing import Any from unittest.mock import patch @@ -340,6 +341,7 @@ def test_restores_a_class_whose_optional_credentials_are_left_empty(self, tmp_pa type_name="A2ATarget", params={"endpoint": "https://agent.example.test"}, external_input=True, + allowed_path_params=frozenset(), ) def test_refuses_a_class_that_needs_the_server_upload_folder(self, tmp_path: Path) -> None: @@ -677,6 +679,29 @@ def test_rejects_server_paths_the_way_the_api_does(self, tmp_path: Path) -> None assert "wordswap_path" in failure.message assert "wordswap" not in ConverterRegistry.get_registry_singleton().instances + def test_lets_a_target_set_the_path_parameters_its_class_allows_callers(self, tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + recipes = tmp_path / "recipes" + _save( + recipes, + RegistryRecipeKind.TARGETS, + "copilot", + { + "class_module": "pyrit.prompt_target.github_copilot_target", + "class_name": "GitHubCopilotTarget", + "params": {"model_name": "gpt-5-mini", "working_directory": str(workspace)}, + }, + ) + + with patch.dict(sys.modules, {"copilot": ModuleType("copilot")}): + report = RegistryRecipeRestorer(source=str(recipes)).restore() + + assert report.failures == [] + target = TargetRegistry.get_registry_singleton().instances.get("copilot") + assert target is not None + assert target._working_directory == str(workspace.resolve()) + def test_reports_a_name_taken_at_startup_and_fails_its_dependents(self, tmp_path: Path) -> None: startup_target = OpenAIChatTarget( endpoint="https://startup.openai.azure.com/openai/v1", model_name="gpt-4o", auth_mode="identity" From a45ac239ba57f9a70b6805908b27c2a2cadf2996 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 22:46:18 +0000 Subject: [PATCH 34/38] Close a file that isn't regular before reporting it and report entries that can't be checked --- pyrit/registry/file_document_storage.py | 26 +++++-- .../registry/test_file_document_storage.py | 70 +++++++++++++++++-- 2 files changed, 86 insertions(+), 10 deletions(-) diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index c025c77eff..c75d398b59 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -125,10 +125,14 @@ def _read_regular_file(path: str) -> bytes | None: Returns: bytes | None: The content, or None if the path names something other than a regular file. """ - with open(os.open(path, _READ_FLAGS), "rb") as file: - if not stat.S_ISREG(os.fstat(file.fileno()).st_mode): + descriptor = os.open(path, _READ_FLAGS) + try: + if not stat.S_ISREG(os.fstat(descriptor).st_mode): return None - return file.read() + with open(descriptor, "rb", closefd=False) as file: + return file.read() + finally: + os.close(descriptor) class DocumentConflictError(ValueError): @@ -317,8 +321,22 @@ def _read_local_entry(self, *, entry: os.DirEntry[str], name: str) -> bytes | Un return UnreadableDocument(name=name, message=self._describe_read_error(error)) if content is not None: return content - if not os.path.lexists(entry.path): + return self._describe_non_file(path=entry.path, name=name) + + @classmethod + def _describe_non_file(cls, *, path: str, name: str) -> UnreadableDocument | None: + """ + Report a listed entry that isn't a regular file, unless it was deleted after the listing. + + Returns: + UnreadableDocument | None: Why it can't be read, or None if it no longer exists. + """ + try: + os.lstat(path) + except FileNotFoundError: return None + except OSError as error: + return UnreadableDocument(name=name, message=cls._describe_read_error(error)) # A broken link, a folder, or a pipe, which would block if read. return UnreadableDocument(name=name, message="Not a regular file") diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index a6e264a45f..2c52fdea3e 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. import os +import sys from collections.abc import Callable, Generator, Iterator from contextlib import contextmanager from pathlib import Path @@ -37,6 +38,12 @@ def _blob_client(*, names: list[str]) -> MagicMock: return client +def _count_open_files() -> int | None: + # Linux lists this process's open files here; elsewhere the count is skipped. + descriptors = Path("/proc/self/fd") + return len(list(descriptors.iterdir())) if descriptors.is_dir() else None + + @contextmanager def _change_after_listing(*, change: Callable[[], None], cache_types: bool = False) -> Generator[None, None, None]: # Runs `change` after the folder is listed and before its entries are read. With `cache_types`, @@ -133,24 +140,56 @@ def test_skips_a_link_deleted_after_listing(self, tmp_path: Path) -> None: assert listing == DocumentListing(documents={"kept": b"kept"}) - @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + @pytest.mark.parametrize("replacement", ["pipe", "folder", "dangling link"]) @pytest.mark.timeout(60) - def test_reports_a_pipe_put_in_place_of_a_listed_document_without_blocking(self, tmp_path: Path) -> None: + def test_reports_a_non_file_put_in_place_of_a_listed_document_without_reading_it( + self, tmp_path: Path, replacement: str + ) -> None: + if replacement == "pipe" and not hasattr(os, "mkfifo"): + pytest.skip("Named pipes need os.mkfifo") + if replacement == "folder" and sys.platform == "win32": + pytest.skip("Windows can't open a folder as a file") + if replacement == "dangling link" and sys.platform == "win32": + pytest.skip("Creating symbolic links may not be permitted on Windows") (tmp_path / "kept.json").write_bytes(b"kept") swapped = tmp_path / "swapped.json" swapped.write_bytes(b"replaced") - def replace_with_pipe() -> None: + def replace() -> None: swapped.unlink() - os.mkfifo(swapped) - - with _change_after_listing(change=replace_with_pipe, cache_types=True): + if replacement == "pipe": + os.mkfifo(swapped) + elif replacement == "folder": + swapped.mkdir() + else: + swapped.symlink_to(tmp_path / "missing_target.json") + + open_files = _count_open_files() + with _change_after_listing(change=replace, cache_types=True): listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing == DocumentListing( documents={"kept": b"kept"}, unreadable=(UnreadableDocument(name="swapped", message="Not a regular file"),), ) + assert _count_open_files() == open_files + + @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + def test_checks_the_opened_file_rather_than_the_path(self, tmp_path: Path) -> None: + document = tmp_path / "document.json" + document.write_bytes(b"original") + real_open = os.open + + def open_then_replace(path: str, flags: int, *args: int) -> int: + descriptor = real_open(path, flags, *args) + os.unlink(path) + os.mkfifo(path) + return descriptor + + with patch("pyrit.registry.file_document_storage.os.open", side_effect=open_then_replace): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing(documents={"document": b"original"}) def test_reports_entries_that_are_not_regular_files_without_reading_them(self, tmp_path: Path) -> None: (tmp_path / "folder.json").mkdir() @@ -168,6 +207,25 @@ def test_reports_entries_that_are_not_regular_files_without_reading_them(self, t ) read.assert_not_called() + @pytest.mark.skipif(sys.platform == "win32", reason="The null device is a POSIX character device here") + def test_reads_only_regular_files(self) -> None: + assert _read_regular_file(os.devnull) is None + + def test_reports_an_entry_that_cannot_be_checked_after_listing(self, tmp_path: Path) -> None: + deleted = tmp_path / "deleted.json" + deleted.write_bytes(b"gone") + + with ( + _change_after_listing(change=deleted.unlink, cache_types=True), + patch( + "pyrit.registry.file_document_storage.os.lstat", + side_effect=PermissionError(13, "Permission denied"), + ), + ): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing(unreadable=(UnreadableDocument(name="deleted", message="Permission denied"),)) + def test_reports_a_link_loop_and_keeps_the_others(self, tmp_path: Path) -> None: (tmp_path / "good.json").write_bytes(b"ok") try: From ad21f0c3d9026cead00e5089f397b114cbab78a7 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 22:56:17 +0000 Subject: [PATCH 35/38] Give only targets the path parameters their class lets callers set --- pyrit/registry/registry_recipe_restore.py | 8 +++--- .../registry/test_registry_recipe_restore.py | 25 ++++++++++++++++++- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index 2d6578d882..b5e9484dac 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -394,10 +394,10 @@ def _build(self, *, stored: StoredRecipe) -> RestoreFailure | None: self._build_technique(name=stored.ref.name, recipe=recipe) else: registry = self._get_registry(stored.ref.kind) - # As in the API, a target may let callers set some path parameters, such as a working folder. - allowed_path_params = getattr( - registry.get_class(recipe.class_name), "caller_path_parameters", frozenset() - ) + allowed_path_params: frozenset[str] = frozenset() + if stored.ref.kind is RegistryRecipeKind.TARGETS: + # As in the target API, a target may let callers set some path parameters, such as a working folder. + allowed_path_params = registry.get_class(recipe.class_name).caller_path_parameters registry.create_named_instance( name=stored.ref.name, type_name=recipe.class_name, diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index ff58ac6d92..49ff9993d0 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -16,7 +16,7 @@ from pydantic import ValidationError from pyrit.backend.models.techniques import CreateTechniqueRequest -from pyrit.converter import Base64Converter +from pyrit.converter import Base64Converter, ColloquialWordswapConverter from pyrit.models.catalog.registry_recipe import ( RegistryRecipeKind, RestoreFailure, @@ -702,6 +702,29 @@ def test_lets_a_target_set_the_path_parameters_its_class_allows_callers(self, tm assert target is not None assert target._working_directory == str(workspace.resolve()) + def test_gives_only_targets_the_path_parameters_their_class_allows_callers(self, tmp_path: Path) -> None: + class CallerPathWordswap(ColloquialWordswapConverter): + caller_path_parameters = frozenset({"wordswap_path"}) + + ConverterRegistry.get_registry_singleton().register_class(CallerPathWordswap) + _save( + tmp_path, + RegistryRecipeKind.CONVERTERS, + "wordswap", + { + "class_module": CallerPathWordswap.__module__, + "class_name": "CallerPathWordswap", + "params": {"wordswap_path": "/etc/passwd"}, + }, + ) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (failure,) = report.failures + assert failure.reason is RestoreFailureReason.BUILD_ERROR + assert "wordswap_path" in failure.message + assert "wordswap" not in ConverterRegistry.get_registry_singleton().instances + def test_reports_a_name_taken_at_startup_and_fails_its_dependents(self, tmp_path: Path) -> None: startup_target = OpenAIChatTarget( endpoint="https://startup.openai.azure.com/openai/v1", model_name="gpt-4o", auth_mode="identity" From db245dcfbef51a74008deb3d01c5f92181a942f2 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Wed, 4 Nov 2026 00:08:59 +0000 Subject: [PATCH 36/38] Test a file replaced just before it's opened and stop a blocked read with a watchdog --- .../registry/test_file_document_storage.py | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index 2c52fdea3e..03926a16dc 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -141,7 +141,7 @@ def test_skips_a_link_deleted_after_listing(self, tmp_path: Path) -> None: assert listing == DocumentListing(documents={"kept": b"kept"}) @pytest.mark.parametrize("replacement", ["pipe", "folder", "dangling link"]) - @pytest.mark.timeout(60) + @pytest.mark.timeout(60, method="thread") def test_reports_a_non_file_put_in_place_of_a_listed_document_without_reading_it( self, tmp_path: Path, replacement: str ) -> None: @@ -175,6 +175,7 @@ def replace() -> None: assert _count_open_files() == open_files @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + @pytest.mark.timeout(60, method="thread") def test_checks_the_opened_file_rather_than_the_path(self, tmp_path: Path) -> None: document = tmp_path / "document.json" document.write_bytes(b"original") @@ -190,6 +191,25 @@ def open_then_replace(path: str, flags: int, *args: int) -> int: listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing == DocumentListing(documents={"document": b"original"}) + assert document.is_fifo() + + @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + @pytest.mark.timeout(60, method="thread") + def test_checks_a_file_replaced_just_before_it_is_opened(self, tmp_path: Path) -> None: + (tmp_path / "document.json").write_bytes(b"original") + real_open = os.open + + def replace_then_open(path: str, flags: int, *args: int) -> int: + os.unlink(path) + os.mkfifo(path) + return real_open(path, flags, *args) + + with patch("pyrit.registry.file_document_storage.os.open", side_effect=replace_then_open): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing( + unreadable=(UnreadableDocument(name="document", message="Not a regular file"),) + ) def test_reports_entries_that_are_not_regular_files_without_reading_them(self, tmp_path: Path) -> None: (tmp_path / "folder.json").mkdir() From 80754a28e6849d086f428f169cd588aa07b57b6c Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Wed, 4 Nov 2026 06:44:58 +0000 Subject: [PATCH 37/38] Read saved recipes separately from building them --- pyrit/registry/registry_recipe_restore.py | 67 ++++++++++++------- .../registry/test_registry_recipe_restore.py | 13 ++++ 2 files changed, 54 insertions(+), 26 deletions(-) diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index b5e9484dac..f919e75798 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -29,6 +29,14 @@ from pyrit.registry.registry import InstanceHoldingRegistry +@dataclass(frozen=True) +class SavedRecipes: + """What reading the saved recipes found: each folder that was read, and a failure for each one that wasn't.""" + + listings: tuple[RecipeListing, ...] + failures: tuple[RestoreFailure, ...] + + @dataclass(frozen=True) class _PlannedEntry: """A saved entry whose class and references checked out, waiting to be built.""" @@ -71,35 +79,15 @@ def __init__(self, *, source: str | None = None) -> None: """ self._storages = {kind: RegistryRecipeStorage(kind=kind, source=source) for kind in RegistryRecipeKind} - def restore(self) -> RegistryRestoreReport: + def read(self) -> SavedRecipes: """ - Rebuild every saved entry that can be rebuilt and register it. + Read every registry's folder without building anything. - Call it when the registries hold only startup entries, as they do at startup and after an - apply resets them; an entry restored earlier would otherwise clash with itself by name. If - any registry's folder can't be read, nothing is built, since entries in the other folders - may name entries in the unread one. - - Returns: - RegistryRestoreReport: What was restored, what failed and why, and which files were ignored. - """ - listings, failures = self._read_all() - ignored = [file for listing in listings for file in listing.ignored] - failures.extend(failure for listing in listings for failure in listing.failures) - if len(listings) < len(self._storages): - return RegistryRestoreReport(failures=self._sorted(failures), ignored=ignored) - stored = {recipe.ref: recipe for listing in listings for recipe in listing.recipes.values()} - planned = self._plan_all(stored=stored, failures=failures) - restored = self._build_in_order(planned=planned, failures=failures) - return RegistryRestoreReport(restored=restored, failures=self._sorted(failures), ignored=ignored) - - def _read_all(self) -> tuple[list[RecipeListing], list[RestoreFailure]]: - """ - Read every registry's folder, turning a folder that can't be listed into a failure. + Reading can wait on storage, so a caller on an event loop can run this in a worker thread and + pass the result to ``restore`` on the loop. A folder that can't be listed becomes a failure. Returns: - tuple[list[RecipeListing], list[RestoreFailure]]: The folders that were read, and one - failure per folder that wasn't. + SavedRecipes: The folders that were read, and a failure for each one that wasn't. """ listings: list[RecipeListing] = [] failures: list[RestoreFailure] = [] @@ -111,7 +99,34 @@ def _read_all(self) -> tuple[list[RecipeListing], list[RestoreFailure]]: failures.append( RestoreFailure(kind=kind, reason=RestoreFailureReason.UNREADABLE_STORAGE, message=message) ) - return listings, failures + return SavedRecipes(listings=tuple(listings), failures=tuple(failures)) + + def restore(self, *, saved: SavedRecipes | None = None) -> RegistryRestoreReport: + """ + Rebuild every saved entry that can be rebuilt and register it. + + Entries are built on the calling thread, as the target API builds them, because a constructor + can need the running event loop. Call it when the registries hold only startup entries, as they + do at startup and after an apply resets them; an entry restored earlier would otherwise clash + with itself by name. If any registry's folder can't be read, nothing is built, since entries + in the other folders may name entries in the unread one. + + Args: + saved (SavedRecipes | None): What ``read`` returned. Defaults to reading now. + + Returns: + RegistryRestoreReport: What was restored, what failed and why, and which files were ignored. + """ + saved = saved if saved is not None else self.read() + listings, failures = list(saved.listings), list(saved.failures) + ignored = [file for listing in listings for file in listing.ignored] + failures.extend(failure for listing in listings for failure in listing.failures) + if len(listings) < len(self._storages): + return RegistryRestoreReport(failures=self._sorted(failures), ignored=ignored) + stored = {recipe.ref: recipe for listing in listings for recipe in listing.recipes.values()} + planned = self._plan_all(stored=stored, failures=failures) + restored = self._build_in_order(planned=planned, failures=failures) + return RegistryRestoreReport(restored=restored, failures=self._sorted(failures), ignored=ignored) def _plan_all( self, *, stored: Mapping[SavedEntryRef, StoredRecipe], failures: list[RestoreFailure] diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index 49ff9993d0..4dda5eb604 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -144,6 +144,19 @@ def test_finds_a_startup_entry_by_its_exact_name(self, tmp_path: Path) -> None: assert report.failures == [] assert isinstance(registry.instances.get("a_pool"), RoundRobinTarget) + def test_reads_without_building_and_restores_what_was_read(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) + restorer = RegistryRecipeRestorer(source=str(tmp_path)) + + with patch.object(TargetRegistry, "create_instance_from_external_input") as build: + saved = restorer.read() + build.assert_not_called() + with patch.object(RegistryRecipeStorage, "list_recipes", side_effect=AssertionError("read again")): + report = restorer.restore(saved=saved) + + assert report.failures == [] + assert [ref.name for ref in report.restored] == ["contoso_gpt4o"] + def test_orders_entries_of_the_same_kind(self, tmp_path: Path) -> None: _save(tmp_path, RegistryRecipeKind.TARGETS, "a_pool", _round_robin("b_member", "c_member")) _save(tmp_path, RegistryRecipeKind.TARGETS, "b_member", _IDENTITY_TARGET) From bd84ef0eb5b546afa0814774b8bbdb2e94ec83cf Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Wed, 4 Nov 2026 06:58:08 +0000 Subject: [PATCH 38/38] Build restored entries on the runtime loop and read them in a worker thread --- pyrit/backend/services/runtime_lifecycle.py | 6 ++- tests/unit/backend/test_runtime_lifecycle.py | 56 ++++++++++++++++++-- 2 files changed, 57 insertions(+), 5 deletions(-) diff --git a/pyrit/backend/services/runtime_lifecycle.py b/pyrit/backend/services/runtime_lifecycle.py index a7e9f778ae..c9ec97a5db 100644 --- a/pyrit/backend/services/runtime_lifecycle.py +++ b/pyrit/backend/services/runtime_lifecycle.py @@ -141,7 +141,11 @@ async def _restore_saved_entries_async(self, config: ConfigurationLoader) -> boo bool: Whether every saved entry was restored. """ restorer = RegistryRecipeRestorer(source=config.registry_recipes_source) - report = await asyncio.to_thread(restorer.restore) + # Reading can wait on storage, so it runs in a worker thread. Entries are built on this loop, as + # the target API builds them, because a constructor such as HuggingFaceChatTarget's starts a task + # on it. A slow constructor holds the loop while it runs. + saved = await asyncio.to_thread(restorer.read) + report = restorer.restore(saved=saved) self.restore_failures, self.restore_ignored = report.failures, report.ignored if not report.failures: return True diff --git a/tests/unit/backend/test_runtime_lifecycle.py b/tests/unit/backend/test_runtime_lifecycle.py index 7ac560f20d..8524bd3b29 100644 --- a/tests/unit/backend/test_runtime_lifecycle.py +++ b/tests/unit/backend/test_runtime_lifecycle.py @@ -10,7 +10,7 @@ from contextlib import ExitStack from datetime import UTC, datetime from pathlib import Path -from threading import Event +from threading import Event, Thread, current_thread from unittest.mock import AsyncMock, MagicMock, call, patch import httpx @@ -33,10 +33,10 @@ from pyrit.prompt_target import OpenAIChatTarget from pyrit.registry.components.converter_registry import ConverterRegistry from pyrit.registry.components.target_registry import TargetRegistry -from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer +from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer, SavedRecipes from pyrit.scenario import Scenario from pyrit.setup.configuration_loader import ConfigurationLoader -from unit.mocks import make_scenario_result +from unit.mocks import MockPromptTarget, make_scenario_result @pytest.fixture @@ -986,6 +986,54 @@ async def test_startup_restores_an_entry_that_uses_a_startup_target_async( assert converter._converter_target is startup_chat +class _LoopBoundTarget(MockPromptTarget): + """A target whose constructor needs the running event loop, as HuggingFaceChatTarget's does.""" + + def __init__(self) -> None: + super().__init__() + self.loop = asyncio.get_running_loop() + + +@pytest.mark.usefixtures("patch_central_database", "startup_chat") +async def test_startup_builds_restored_entries_on_the_runtime_loop_async( + runtime: RuntimeLifecycle, tmp_path: Path +) -> None: + TargetRegistry.get_registry_singleton().register_class(_LoopBoundTarget) + recipes = tmp_path / "recipes" + _save_recipe( + recipes, + kind="targets", + name="loop_bound", + document={"class_module": _LoopBoundTarget.__module__, "class_name": "_LoopBoundTarget", "params": {}}, + ) + + await _start_with_recipes_async(runtime, recipes=recipes) + + assert (runtime.state, runtime.outcome) == ("ready", "success") + target = TargetRegistry.get_registry_singleton().instances.get("loop_bound") + assert isinstance(target, _LoopBoundTarget) + assert target.loop is asyncio.get_running_loop() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_startup_reads_saved_entries_off_the_runtime_loop_async( + runtime: RuntimeLifecycle, tmp_path: Path +) -> None: + read = RegistryRecipeRestorer.read + read_threads: list[Thread] = [] + + def read_noting_thread(restorer: RegistryRecipeRestorer) -> SavedRecipes: + read_threads.append(current_thread()) + return read(restorer) + + with patch.object(RegistryRecipeRestorer, "read", read_noting_thread): + await _start_with_recipes_async(runtime, recipes=tmp_path / "recipes") + + assert (runtime.state, runtime.outcome) == ("ready", "success") + assert len(read_threads) == 1 + assert read_threads[0] is not current_thread() + + @pytest.mark.usefixtures("patch_central_database") async def test_apply_restores_an_entry_that_uses_a_startup_target_async( runtime: RuntimeLifecycle, tmp_path: Path, startup_chat: OpenAIChatTarget @@ -1116,7 +1164,7 @@ async def test_apply_with_an_unrestorable_entry_keeps_the_new_version_async( async def test_restore_cancellation_propagates_from_startup_async(runtime: RuntimeLifecycle, tmp_path: Path) -> None: runtime.state = "initializing" - with patch.object(RegistryRecipeRestorer, "restore", side_effect=asyncio.CancelledError): + with patch.object(RegistryRecipeRestorer, "read", side_effect=asyncio.CancelledError): with pytest.raises(asyncio.CancelledError): await _start_with_recipes_async(runtime, recipes=tmp_path / "recipes")