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/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