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 e8196da582..b2fb7ed98f 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; 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 (up to 64 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/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 a42514bb7c..668e6599e7 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -271,8 +271,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` diff --git a/pyrit/backend/services/runtime_lifecycle.py b/pyrit/backend/services/runtime_lifecycle.py index 9767c56749..c9ec97a5db 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,40 @@ 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) + # 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 + 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 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", + [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 +179,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 +262,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/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index 102788e2a8..a113a8bef0 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -20,6 +20,17 @@ if TYPE_CHECKING: from pyrit.models.catalog.initializer import RegisteredInitializer + from pyrit.models.catalog.registry_recipe import ( + ComponentRecipe, + IgnoredRecipeFile, + RegistryRecipeKind, + RegistryRestoreReport, + RestoreFailure, + RestoreFailureReason, + SavedEntryRef, + StoredRecipe, + TechniqueRecipe, + ) from pyrit.models.catalog.scenario import ( AttackErrorSummary, AttackRetrySummary, @@ -60,9 +71,16 @@ _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", + "RegistryRestoreReport": "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 +105,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 new file mode 100644 index 0000000000..9b0c2e66d3 --- /dev/null +++ b/pyrit/models/catalog/registry_recipe.py @@ -0,0 +1,167 @@ +# 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, PlainSerializer, field_validator + +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}$")] + + +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") + + +# 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): + """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 = 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) + 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: _DisplayText + + +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: _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") + 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: _DisplayText + message: _DisplayText + + +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/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/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index 56432d80b7..c75d398b59 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -12,13 +12,16 @@ from __future__ import annotations +import fnmatch import hashlib import logging import os +import stat import sys 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 +36,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. @@ -81,6 +110,31 @@ 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. + """ + descriptor = os.open(path, _READ_FLAGS) + try: + if not stat.S_ISREG(os.fstat(descriptor).st_mode): + return None + with open(descriptor, "rb", closefd=False) as file: + return file.read() + finally: + os.close(descriptor) + + class DocumentConflictError(ValueError): """A stored document changed after the caller read it.""" @@ -199,6 +253,160 @@ 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`` applies the platform's case rule, as ``glob`` does. + with os.scandir(directory) as entries: + candidates = sorted( + (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 entry in candidates: + name = Path(entry.name).stem + problem = self._name_problem(name) + if problem is not None: + 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. + + 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: + content = _read_regular_file(entry.path) if entry.is_file() else None + except FileNotFoundError: + content = None + except OSError as error: + return UnreadableDocument(name=name, message=self._describe_read_error(error)) + if content is not None: + return content + 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") + + 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 + from azure.storage.blob import StorageErrorCode + + 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 AzureError as error: + 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: + """ + 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/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py new file mode 100644 index 0000000000..f919e75798 --- /dev/null +++ b/pyrit/registry/registry_recipe_restore.py @@ -0,0 +1,578 @@ +# 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 unicodedata +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar + +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 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.""" + + 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"}) + _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 read(self) -> SavedRecipes: + """ + Read every registry's folder without building anything. + + 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: + SavedRecipes: The folders that were read, and a failure for each one that wasn't. + """ + listings: list[RecipeListing] = [] + failures: list[RestoreFailure] = [] + for kind, storage in self._storages.items(): + try: + listings.append(storage.list_recipes()) + 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) + ) + 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] + ) -> 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, the attack type isn't registered, or the + recipe holds a credential. + """ + 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 + self._check_stored_credentials(parameters=parameters, params=recipe.params) + 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: + registry = self._get_registry(stored.ref.kind) + 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, + 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)) + 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") + + @staticmethod + def _safe_message(error: Exception) -> str: + """ + 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. 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: + 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: + # A plain str, so no string method a subclass overrides runs on the text below. + text = str.__str__(str(error)) + except Exception: + return name + 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=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: + """ + 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's compatibility form contains one of the marks. + """ + folded = unicodedata.normalize("NFKC", char) + return any(mark in folded for mark in marks) + + +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/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/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/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/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_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"]) diff --git a/tests/unit/backend/test_runtime_lifecycle.py b/tests/unit/backend/test_runtime_lifecycle.py index 5bc4187327..8524bd3b29 100644 --- a/tests/unit/backend/test_runtime_lifecycle.py +++ b/tests/unit/backend/test_runtime_lifecycle.py @@ -4,12 +4,13 @@ """Admission and fail-fast runtime replacement invariants.""" 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 -from threading import Event +from threading import Event, Thread, current_thread from unittest.mock import AsyncMock, MagicMock, call, patch import httpx @@ -29,9 +30,13 @@ 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, 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 @@ -917,3 +922,250 @@ 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, 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(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 + + +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 +) -> 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"}) + + 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 + 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"] + + +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_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: + 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, "read", 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/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py new file mode 100644 index 0000000000..03926a16dc --- /dev/null +++ b/tests/unit/registry/test_file_document_storage.py @@ -0,0 +1,379 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import os +import sys +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 + +import pytest +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, + _read_regular_file, +) + +_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 + + +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`, + # 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("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() + + 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] = [] + + def read(path: str) -> bytes | None: + read_names.append(Path(path).name) + return _read_regular_file(path) + + 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"} + 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"{}\r\n") + (tmp_path / "empty.json").write_bytes(b"") + deleted = tmp_path / "deleted.json" + deleted.write_bytes(b"gone") + + 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"}) + + 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.parametrize("replacement", ["pipe", "folder", "dangling link"]) + @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: + 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() -> None: + swapped.unlink() + 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") + @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") + 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"}) + 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() + 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("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.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: + (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") + + 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() + + 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_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 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 + + 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"}, + 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.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 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..4dda5eb604 --- /dev/null +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -0,0 +1,820 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import dataclasses +import json +import os +import sys +import unicodedata +from collections.abc import Iterator +from pathlib import Path +from types import ModuleType +from typing import Any +from unittest.mock import patch + +import pytest +from pydantic import ValidationError + +from pyrit.backend.models.techniques import CreateTechniqueRequest +from pyrit.converter import Base64Converter, ColloquialWordswapConverter +from pyrit.models.catalog.registry_recipe import ( + RegistryRecipeKind, + RestoreFailure, + RestoreFailureReason, + 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 +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_finds_a_startup_entry_by_its_exact_name(self, tmp_path: Path) -> None: + registry = TargetRegistry.get_registry_singleton() + for name in ("ghost\ud800", "other\udcff"): + startup_target = OpenAIChatTarget( + endpoint="https://startup.openai.azure.com/openai/v1", model_name="gpt-4o", auth_mode="identity" + ) + registry.instances.register(startup_target, name=name) + _save(tmp_path, RegistryRecipeKind.TARGETS, "a_pool", _round_robin("ghost\ud800", "other\udcff")) + + report = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + 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) + _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"}) + ) + + 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", + [ + {"tags": ["all"]}, + {"tags": ["Custom", "custom"]}, + {"tags": ["bad-tag"]}, + {"type": ""}, + {"params": {"x" * 300: 1}}, + {"request_converters": ["c"] * 101}, + {"params": {"temperature": float("nan")}}, + {"description": "x" * 100_001}, + {"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") + + @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"} + _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_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, + allowed_path_params=frozenset(), + ) + + 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"] + + @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", "'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: + 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 + if shown is None: + assert failure.message == "ValueError (message hidden because it may include a credential)" + else: + assert failure.message.startswith("ValueError: Identity-based authentication requires") + 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 + + 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, + "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"] + + @pytest.mark.parametrize( + ("text", "expected"), + [ + ( + "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)"), + (" \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']", + "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)"), + ("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)"), + ( + "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: + assert RegistryRecipeRestorer._safe_message(ValueError(text)) == expected + + @pytest.mark.parametrize( + "mark", ["?", "#", "\u2047", "\u2048", "\u2049", "\ufe16", "\ufe56", "\ufe5f", "\uff03", "\uff1f"] + ) + 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 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(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): + def __str__(self) -> str: + raise RuntimeError("cannot print") + + assert RegistryRecipeRestorer._safe_message(UnprintableError()) == "UnprintableError" + + def test_safe_message_ignores_string_methods_an_error_overrides(self) -> None: + class KeepEverything(str): + __slots__ = () + + def __getitem__(self, key: Any) -> str: + return str.__str__(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_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") + sent = json.loads(json.dumps(report.model_dump(mode="json"), ensure_ascii=False).encode("utf-8")) + assert sent["failures"][0]["dependency"]["name"] == "ghost\\ud800" + + 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, + 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_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_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" + ) + 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_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.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 = RegistryRecipeRestorer(source=str(tmp_path)).restore() + + (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) + 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_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) + _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"] 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..67b641abed --- /dev/null +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -0,0 +1,219 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +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, + SavedEntryRef, + TechniqueRecipe, +) +from pyrit.registry.file_document_storage import _read_regular_file +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 TestRestoreReportText: + 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" + ) + failure = RestoreFailure( + kind=RegistryRecipeKind.TARGETS, + name="x\udcff", + reason=RestoreFailureReason.MISSING_DEPENDENCY, + message="got \udcff", + dependency=SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800"), + ) + + 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.model_dump(mode="json")["name"] == "caf\u00e9 \uff10\uff0e\uff15 \U0001f600" + + +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_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) == ("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: + _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_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 == 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" + _write(folder, "locked", json.dumps(_TARGET)) + _write(folder, "Bad-Name", "{}") + + def read(path: str) -> bytes | None: + if Path(path).name == "locked.json": + raise PermissionError(13, "Permission denied", path) + return _read_regular_file(path) + + with patch("pyrit.registry.file_document_storage._read_regular_file", side_effect=read): + 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 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")