diff --git a/pyrit/backend/middleware/error_handlers.py b/pyrit/backend/middleware/error_handlers.py index 5973d6289e..dbd03c36b6 100644 --- a/pyrit/backend/middleware/error_handlers.py +++ b/pyrit/backend/middleware/error_handlers.py @@ -14,6 +14,7 @@ from pyrit.backend.middleware.request_size import RequestTooLargeError, request_too_large_response from pyrit.backend.models import FieldError, ProblemDetail from pyrit.exceptions.analytics_exception import AnalyticsException +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError, DocumentUnreadableError logger = logging.getLogger(__name__) @@ -169,6 +170,39 @@ async def permission_error_handler( # pyrit-async-suffix-exempt content=problem.model_dump(exclude_none=True), ) + @app.exception_handler(DocumentStorageUnavailableError) + async def storage_unavailable_handler( # pyrit-async-suffix-exempt + request: Request, + exc: DocumentStorageUnavailableError, + ) -> JSONResponse: + """ + Handle stored documents that can't be read or written as 503 Service Unavailable. + + The log names the storage location, without credentials, and the reason, with the + traceback at debug level. The response never carries a path or URL. + + Returns: + JSONResponse: RFC 7807 problem detail response with 503 status. + """ + logger.warning(f"{exc} ({request.method} {request.url.path})") + logger.debug("Storage error details", exc_info=exc) + if isinstance(exc, DocumentUnreadableError): + detail = f"{exc.label} '{exc.name}' can't be read: {exc.reason}" + else: + detail = f"{exc.label} storage is unavailable" + problem = ProblemDetail( + type="/errors/storage-unavailable", + title="Service Unavailable", + status=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=detail, + instance=str(request.url.path), + ) + + return JSONResponse( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + content=problem.model_dump(exclude_none=True), + ) + @app.exception_handler(NotImplementedError) async def not_implemented_handler( # pyrit-async-suffix-exempt request: Request, diff --git a/pyrit/backend/routes/initializers.py b/pyrit/backend/routes/initializers.py index 50fa79db85..529f92d641 100644 --- a/pyrit/backend/routes/initializers.py +++ b/pyrit/backend/routes/initializers.py @@ -14,7 +14,6 @@ DELETE /api/initializers/{name} — unregister an initializer """ -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from pyrit.backend.middleware.auth import require_admin @@ -35,19 +34,6 @@ router = APIRouter(prefix="/initializers", tags=["initializers"]) -def _custom_storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable custom initializer storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Custom initializer storage is temporarily unavailable", - ) - - def _configured_initializers(request: Request) -> list[ConfiguredInitializerSetting]: """ Read the initializer list captured from ``.pyrit_conf`` at backend startup. @@ -138,10 +124,7 @@ async def list_custom_initializers(request: Request) -> CustomInitializerListRes CustomInitializerListResponse: The configured source and stored scripts. """ _check_custom_initializers_allowed(request) - try: - return await get_initializer_service().list_custom_initializers_async() - except AzureError as exc: - raise _custom_storage_unavailable() from exc + return await get_initializer_service().list_custom_initializers_async() @router.get( @@ -205,8 +188,6 @@ async def register_initializer( # pyrit-async-suffix-exempt try: return await service.register_initializer_async(name=body.name, script_content=body.script_content) - except AzureError as exc: - raise _custom_storage_unavailable() from exc except ValueError as e: detail = str(e) if "already registered" in detail: @@ -243,8 +224,6 @@ async def unregister_initializer( # pyrit-async-suffix-exempt try: await service.unregister_initializer_async(initializer_name=initializer_name) - except AzureError as exc: - raise _custom_storage_unavailable() from exc except ValueError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, diff --git a/pyrit/backend/routes/scenario_presets.py b/pyrit/backend/routes/scenario_presets.py index 4b3fc1db00..973112df7b 100644 --- a/pyrit/backend/routes/scenario_presets.py +++ b/pyrit/backend/routes/scenario_presets.py @@ -19,7 +19,6 @@ from typing import Annotated -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, status from pyrit.backend.middleware.auth import current_user_name, require_admin @@ -38,19 +37,6 @@ router = APIRouter(prefix="/scenario-presets", tags=["scenario-presets"]) -def _storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable preset storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Scenario preset storage is temporarily unavailable", - ) - - def _not_found(name: str) -> HTTPException: """ Create a not-found response for a missing preset. @@ -79,8 +65,6 @@ async def _load_preset_or_404_async(name: str) -> ScenarioPresetResponse: """ try: preset = await get_scenario_preset_service().get_preset_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -106,10 +90,7 @@ async def list_scenario_presets( # pyrit-async-suffix-exempt Returns: ScenarioPresetListResponse: The configured source and the stored presets. """ - try: - return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) - except AzureError as exc: - raise _storage_unavailable() from exc + return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) @router.get( @@ -177,8 +158,6 @@ async def create_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{preset.name}' already exists", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.put( @@ -239,8 +218,6 @@ async def update_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{name}' changed since it was read; re-read it and retry", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.delete( @@ -261,8 +238,6 @@ async def delete_scenario_preset(name: IdentifierStr) -> None: # pyrit-async-su """ try: deleted = await get_scenario_preset_service().delete_preset_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -312,8 +287,6 @@ async def resolve_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{name}' changed since it was read; re-read it and retry", ) from None - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None diff --git a/pyrit/backend/routes/scenario_suites.py b/pyrit/backend/routes/scenario_suites.py index 0b4fa7e09a..f1ec16872e 100644 --- a/pyrit/backend/routes/scenario_suites.py +++ b/pyrit/backend/routes/scenario_suites.py @@ -6,7 +6,6 @@ from typing import Annotated from uuid import UUID -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Request, status from pyrit.backend.middleware.auth import current_user_name, require_admin @@ -31,19 +30,6 @@ router = APIRouter(prefix="/scenario-suites", tags=["scenario-suites"]) -def _storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable suite storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Scenario suite storage is temporarily unavailable", - ) - - def _not_found(name: str) -> HTTPException: """ Create a not-found response for a missing suite. @@ -66,8 +52,6 @@ async def _load_suite_or_404_async(name: str) -> ScenarioSuiteResponse: """ try: suite = await get_scenario_suite_service().get_suite_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -84,10 +68,7 @@ async def list_scenario_suites() -> ScenarioSuiteListResponse: # pyrit-async-su Returns: ScenarioSuiteListResponse: Suites from both origins with current entry statuses. """ - try: - return await get_scenario_suite_service().list_suites_async() - except AzureError as exc: - raise _storage_unavailable() from exc + return await get_scenario_suite_service().list_suites_async() @router.get( @@ -181,8 +162,6 @@ async def launch_scenario_suite( # pyrit-async-suffix-exempt raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None except ScenarioSuiteExecutionCapacityError as exc: raise HTTPException(status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=str(exc)) from None - except AzureError as exc: - raise _storage_unavailable() from exc except RuntimeError as exc: raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(exc)) from None @@ -239,8 +218,6 @@ async def create_scenario_suite( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario suite '{suite.name}' already exists", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.put( @@ -289,8 +266,6 @@ async def update_scenario_suite( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario suite '{name}' changed since it was read; re-read it and retry", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.delete( @@ -303,8 +278,6 @@ async def delete_scenario_suite(name: IdentifierStr) -> None: # pyrit-async-suf """Delete one stored scenario suite.""" try: deleted = await get_scenario_suite_service().delete_suite_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None diff --git a/pyrit/backend/services/scenario_suite_service.py b/pyrit/backend/services/scenario_suite_service.py index 4e668bb4eb..2636d842d6 100644 --- a/pyrit/backend/services/scenario_suite_service.py +++ b/pyrit/backend/services/scenario_suite_service.py @@ -9,8 +9,6 @@ from functools import lru_cache from typing import Literal -from azure.core.exceptions import AzureError - from pyrit.backend.models import ( ScenarioSuiteEntryStatus, ScenarioSuiteListResponse, @@ -25,6 +23,7 @@ ScenarioSuiteRegistryEntry, ) from pyrit.registry import ScenarioPresetRegistry, ScenarioSuiteRegistry +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError logger = logging.getLogger(__name__) @@ -209,7 +208,7 @@ def _resolve_preset_projection( ) stored_version = registry.get_preset_version(name=preset_name) - except (AzureError, OSError): + except DocumentStorageUnavailableError: logger.warning( "Could not project scenario preset '%s' because preset storage is unavailable.", preset_name, diff --git a/pyrit/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index 102788e2a8..94c1abb948 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -20,6 +20,16 @@ if TYPE_CHECKING: from pyrit.models.catalog.initializer import RegisteredInitializer + from pyrit.models.catalog.registry_recipe import ( + ComponentRecipe, + IgnoredRecipeFile, + RegistryRecipeKind, + RestoreFailure, + RestoreFailureReason, + SavedEntryRef, + StoredRecipe, + TechniqueRecipe, + ) from pyrit.models.catalog.scenario import ( AttackErrorSummary, AttackRetrySummary, @@ -60,9 +70,15 @@ _LAZY_EXPORTS: dict[str, str] = { "AttackErrorSummary": "pyrit.models.catalog.scenario", "AttackRetrySummary": "pyrit.models.catalog.scenario", + "ComponentRecipe": "pyrit.models.catalog.registry_recipe", + "IgnoredRecipeFile": "pyrit.models.catalog.registry_recipe", "PresetEntry": "pyrit.models.catalog.scenario_preset", "PresetIssue": "pyrit.models.catalog.scenario_preset", "RegisteredInitializer": "pyrit.models.catalog.initializer", + "RegistryRecipeKind": "pyrit.models.catalog.registry_recipe", + "RestoreFailure": "pyrit.models.catalog.registry_recipe", + "RestoreFailureReason": "pyrit.models.catalog.registry_recipe", + "SavedEntryRef": "pyrit.models.catalog.registry_recipe", "ScorerInstance": "pyrit.models.catalog.scorer", "RegisteredScenario": "pyrit.models.catalog.scenario", "RunScenarioRequest": "pyrit.models.catalog.scenario", @@ -87,8 +103,10 @@ "ScenarioRunSummary": "pyrit.models.catalog.scenario", "ScenarioTechniqueSummary": "pyrit.models.catalog.scenario", "StoredPreset": "pyrit.models.catalog.scenario_preset", + "StoredRecipe": "pyrit.models.catalog.registry_recipe", "StoredScenarioSuite": "pyrit.models.catalog.scenario_suite", "TargetInstance": "pyrit.models.catalog.target", + "TechniqueRecipe": "pyrit.models.catalog.registry_recipe", } __all__ = list(_LAZY_EXPORTS) diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py new file mode 100644 index 0000000000..d8846147d2 --- /dev/null +++ b/pyrit/models/catalog/registry_recipe.py @@ -0,0 +1,176 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Saved registry entries: the recipe kept for each one and what restoring them found.""" + +from __future__ import annotations + +from enum import Enum +from typing import Annotated, ClassVar + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_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}$")] + + +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) + + # A literal 1 rather than CURRENT_SCHEMA_VERSION: a document saved without this field is a version 1 + # document and must keep reading as one after the constant is bumped, or a migration would skip it. + schema_version: int = Field(1, ge=1, description="Shape of this document") + + @field_validator("schema_version") + @classmethod + def _validate_schema_version(cls, value: int) -> int: + """ + Reject a document saved by a newer PyRIT than this one. + + Returns: + int: The schema version. + + Raises: + ValueError: If the document's format is newer than this PyRIT reads. + """ + if value > cls.CURRENT_SCHEMA_VERSION: + raise ValueError( + f"Format version {value} is newer than this PyRIT reads (up to {cls.CURRENT_SCHEMA_VERSION})" + ) + return value + + +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: str + + +class StoredRecipe(BaseModel): + """A saved recipe and the version of the document it was read from.""" + + ref: SavedEntryRef + recipe: ComponentRecipe | TechniqueRecipe + version: str = Field(..., description="Opaque version of the document this recipe was read from") + + @model_validator(mode="after") + def _validate_recipe_kind(self) -> StoredRecipe: + """ + Reject a technique recipe filed under another registry, or another registry's recipe filed as a technique. + + Returns: + StoredRecipe: The stored recipe. + + Raises: + ValueError: If the recipe's type doesn't match the registry it's filed under. + """ + expects_technique = self.ref.kind is RegistryRecipeKind.TECHNIQUES + if expects_technique != isinstance(self.recipe, TechniqueRecipe): + expected = TechniqueRecipe if expects_technique else ComponentRecipe + raise ValueError(f"A saved {self.ref.kind.value} entry needs a {expected.__name__}") + return self + + +class RestoreFailureReason(str, Enum): + """Why a saved entry couldn't be restored.""" + + UNREADABLE_STORAGE = "unreadable_storage" + UNREADABLE_RECIPE = "unreadable_recipe" + INVALID_RECIPE = "invalid_recipe" + NEWER_SCHEMA_VERSION = "newer_schema_version" + UNKNOWN_CLASS = "unknown_class" + MODULE_MISMATCH = "module_mismatch" + UNSUPPORTED_CLASS = "unsupported_class" + UNSUPPORTED_CREDENTIAL = "unsupported_credential" + NAME_CLASH = "name_clash" + MISSING_DEPENDENCY = "missing_dependency" + FAILED_DEPENDENCY = "failed_dependency" + DEPENDENCY_LOOP = "dependency_loop" + BUILD_ERROR = "build_error" + + +class RestoreFailure(BaseModel): + """One saved entry, or one whole storage folder, that couldn't be restored.""" + + model_config = ConfigDict(frozen=True) + + kind: RegistryRecipeKind + name: str | None = Field(None, description="Entry name; None when the whole folder couldn't be read") + reason: RestoreFailureReason + message: str = Field(..., description="Why, in words safe to show users") + version: str | None = Field(None, description="Version of the stored document, when it could be read") + dependency: SavedEntryRef | None = Field(None, description="The entry this one needed that is missing or failed") + + +class IgnoredRecipeFile(BaseModel): + """A stored file whose name can't be an entry name, so it was never read.""" + + model_config = ConfigDict(frozen=True) + + kind: RegistryRecipeKind + file_name: str + message: str diff --git a/pyrit/registry/components/initializer_registry.py b/pyrit/registry/components/initializer_registry.py index 6e1e283d2d..53d06fb8a5 100644 --- a/pyrit/registry/components/initializer_registry.py +++ b/pyrit/registry/components/initializer_registry.py @@ -420,6 +420,7 @@ def register_from_content(self, *, name: str, script_content: str) -> str: ValueError: If the source cannot be compiled, does not define a concrete ``PyRITInitializer`` subclass, or *name* collides with an existing entry. + DocumentStorageUnavailableError: If the script storage can't be written. """ self._ensure_discovered() @@ -429,11 +430,7 @@ def register_from_content(self, *, name: str, script_content: str) -> str: raise ValueError(f"Initializer '{name}' is already registered. Unregister it first to replace it.") discovered = self._load_custom_initializer_class(name=name, script_content=script_content) - storage = self._get_custom_storage() - try: - storage.save_script(name=name, content=script_content) - except OSError as error: - raise ValueError(f"Failed to write initializer script: {error}") from error + self._get_custom_storage().save_script(name=name, content=script_content) self.register_class(discovered, name=name) logger.info(f"Registered custom initializer: {name} ({discovered.__name__})") @@ -484,10 +481,12 @@ def unregister_and_cleanup(self, name: str) -> None: Raises: KeyError: If no stored custom initializer has the requested name. ValueError: If the name refers to a built-in initializer. + DocumentStorageUnavailableError: If the script storage can't be used. """ self._ensure_discovered() storage = self._get_custom_storage() - stored_names = storage.list_scripts().keys() + # A script that can't be read or decoded is still deleted, so the API can clean it up. + stored_names = sorted(storage.list_script_names()) if name not in stored_names: if name in self._builtin_names: raise ValueError(f"Cannot remove built-in initializer '{name}'.") diff --git a/pyrit/registry/custom_initializer_storage.py b/pyrit/registry/custom_initializer_storage.py index 572c838c00..90ec10695f 100644 --- a/pyrit/registry/custom_initializer_storage.py +++ b/pyrit/registry/custom_initializer_storage.py @@ -67,6 +67,16 @@ def list_scripts(self) -> dict[str, str]: logger.warning(f"Skipping stored initializer '{name}': it is not valid UTF-8 text.") return scripts + def list_script_names(self) -> set[str]: + """ + List the names of every stored script, including ones that can't be read or decoded. + + Returns: + set[str]: The registry names of the stored scripts. + """ + listing = self._list_documents_with_issues() + return set(listing.documents) | {unreadable.name for unreadable in listing.unreadable} + def save_script(self, *, name: str, content: str) -> None: """Persist one custom initializer script.""" self._save_document(name=name, content=content.encode("utf-8")) diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index f71309a25f..c18a34a2cb 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -12,14 +12,19 @@ from __future__ import annotations +import errno +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 types import MappingProxyType from typing import TYPE_CHECKING from urllib.parse import unquote, urlparse @@ -27,12 +32,100 @@ from pyrit.models import validate_registry_name if TYPE_CHECKING: - from collections.abc import Generator + from collections.abc import Generator, Mapping from azure.storage.blob import ContainerClient logger = logging.getLogger(__name__) +# Called through these names so tests can change the folder between steps without patching ``os`` for the +# whole process. +_open = os.open +_stat = os.stat +_lstat = os.lstat +_scandir = os.scandir + + +@dataclass(frozen=True) +class UnreadableDocument: + """ + A stored document with a valid name whose content could not be read. + + ``message`` is a short reason that is safe to show. ``error`` is what stopped the read, kept for logging and + for telling a passing failure from a broken file; it is None when nothing raised, such as for a folder. + """ + + name: str + file_name: str + message: str + error: BaseException | None = field(default=None, compare=False, repr=False) + + +@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: Mapping[str, bytes] = field(default_factory=lambda: MappingProxyType({})) + unreadable: tuple[UnreadableDocument, ...] = () + ignored: tuple[IgnoredDocument, ...] = () + + +class DocumentStorageUnavailableError(Exception): + """ + The folder or container that holds the documents couldn't be used, so nothing was read or written. + + The message names the location for logs; ``display_source`` already hides a SAS token. ``reason`` is short and + carries no path or URL, so an API response can show it. + """ + + def __init__(self, *, label: str, source: str, reason: str, message: str | None = None) -> None: + """ + Initialize the error. + + Args: + label (str): The kind of document, such as "Scenario preset". + source (str): The storage location, without credentials. + reason (str): A short reason without a path or URL, such as "Permission denied". + message (str | None): The full message, or None for the default one. Defaults to None. + """ + super().__init__(message or f"{label} storage at {source} is unavailable: {reason}") + self.label = label + self.source = source + self.reason = reason + + +class DocumentUnreadableError(DocumentStorageUnavailableError): + """One stored document exists but couldn't be read, such as a file without read permission or a folder.""" + + def __init__(self, *, label: str, source: str, name: str, file_name: str, reason: str) -> None: + """ + Initialize the error. + + Args: + label (str): The kind of document, such as "Scenario preset". + source (str): The storage location, without credentials. + name (str): The document's name. + file_name (str): The file or blob name in storage. + reason (str): A short reason without a path or URL, such as "Not a regular file". + """ + super().__init__( + label=label, + source=source, + reason=reason, + message=f"{label} '{file_name}' in {source} can't be read: {reason}", + ) + self.name = name + self.file_name = file_name + + # 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 +174,84 @@ 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. On Windows, opening a folder raises ``PermissionError``, + so a folder swapped in after listing is reported as "Permission denied" there. + + Returns: + bytes | None: The content, or None if the path names something other than a regular file. + """ + descriptor = _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) + + +def _affects_whole_store(error: Exception) -> bool: + """ + Return whether a Blob error means the container can't be used, rather than one blob. + + Returns: + bool: True for sign-in and permission failures, a missing container, and network errors. + """ + from azure.core.exceptions import ClientAuthenticationError, ServiceRequestError, ServiceResponseError + from azure.storage.blob import StorageErrorCode + + if isinstance(error, (ClientAuthenticationError, ServiceRequestError, ServiceResponseError)): + return True + if getattr(error, "status_code", None) in (401, 403): + return True + error_code = getattr(error, "error_code", None) + return isinstance(error_code, str) and error_code == StorageErrorCode.CONTAINER_NOT_FOUND + + +def _is_missing_blob(error: Exception) -> bool: + """ + Return whether a Blob error says the blob itself doesn't exist, as opposed to its container. + + Returns: + bool: True only for ``BlobNotFound``. + """ + from azure.storage.blob import StorageErrorCode + + error_code = getattr(error, "error_code", None) + return isinstance(error_code, str) and error_code == StorageErrorCode.BLOB_NOT_FOUND + + +def describe_storage_error(error: Exception) -> str: + """ + Describe a storage error in a few words, leaving out any path or URL it carries. + + Args: + error (Exception): The error a local folder or Azure Blob Storage raised. + + 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__ + + class DocumentConflictError(ValueError): """A stored document changed after the caller read it.""" @@ -133,6 +304,13 @@ class FileDocumentStorage: an OS advisory lock held across the whole read-compare-replace locally. That lock binds only writers using this class; a direct edit is caught by the version instead. + A failure of the storage itself (a folder that can't be read, a missing container, a + sign-in or network error, a lock that stays taken) raises ``DocumentStorageUnavailableError``. + Its message names the location for logs, using ``display_source`` so a SAS token never + appears; an API response shows only its short reason. One document that exists but + can't be read raises ``DocumentUnreadableError`` from a single read, and is reported in + a listing's ``unreadable`` instead. + Subclasses supply the extension and an error-message label. """ @@ -191,72 +369,262 @@ def _list_documents(self) -> dict[str, bytes]: """ Read every stored document that the single-document operations can address. - A name is checked before its content is read, so a file this class cannot address - is never opened and cannot fail the listing on its way out. + The documents ``_list_documents_with_issues`` can't use are logged and left out. A storage + location that can't be listed at all raises instead of looking empty. Returns: dict[str, bytes]: Document content keyed by name. + + Raises: + DocumentStorageUnavailableError: If the folder or container can't be listed. + """ + listing = self._list_documents_with_issues() + for ignored in listing.ignored: + logger.warning( + f"Ignoring stored document '{ignored.file_name}' in {self.display_source}: {ignored.message}" + ) + for unreadable in listing.unreadable: + logger.warning( + f"Skipping unreadable document '{unreadable.file_name}' in {self.display_source}: {unreadable.message}" + ) + return dict(listing.documents) + + def _list_documents_with_issues(self) -> DocumentListing: + """ + Read every stored document and report the ones that can't be used. + + 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, and so does a Blob error that affects the whole container. + + 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. + + Raises: + DocumentStorageUnavailableError: If the folder or container can't be listed or read. """ - return self._list_blob_documents() if self._is_blob else self._list_local_documents() + with self._unavailable_on_storage_error(): + if self._is_blob: + return self._list_blob_documents_with_issues() + return self._list_local_documents_with_issues() - def _list_local_documents(self) -> dict[str, bytes]: + def _list_local_documents_with_issues(self) -> DocumentListing: """ - Read addressable documents from the configured local directory. + Read the local directory, reporting unreadable and badly named documents. Returns: - dict[str, bytes]: Document content keyed by file stem. + 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 _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] = {} - for path in sorted(directory.glob(f"*{self._extension}")): - if not self._is_addressable_name(path.stem): + 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 try: - documents[path.stem] = path.read_bytes() - except OSError as error: - logger.warning(f"Skipping unreadable document '{path.name}' in {self.display_source}: {error}") - return documents + content = self._read_local_document(path=entry.path, name=name, file_name=entry.name) + except DocumentUnreadableError as error: + unreadable.append( + UnreadableDocument(name=name, file_name=entry.name, message=error.reason, error=error.__cause__) + ) + continue + if content is not None: + documents[name] = content + return DocumentListing( + documents=MappingProxyType(documents), unreadable=tuple(unreadable), ignored=tuple(ignored) + ) + + def _read_local_document(self, *, path: str, name: str, file_name: str) -> bytes | None: + """ + Read one local document, the same way for a listing and for a single read. + + The folder can change after it is listed, so what is opened is checked again: an entry + deleted in between counts as missing, whether or not it was a link, and anything that isn't + a regular file is reported without being read. + + Returns: + bytes | None: The content, or None if nothing exists at the path. + + Raises: + DocumentUnreadableError: If something exists at the path but can't be read as a regular file. + """ + try: + if stat.S_ISREG(_stat(path).st_mode): + content = _read_regular_file(path) + if content is not None: + return content + except FileNotFoundError: + pass + except OSError as error: + raise self._unreadable(name=name, file_name=file_name, reason=describe_storage_error(error)) from error + try: + _lstat(path) + except FileNotFoundError: + return None + except OSError as error: + raise self._unreadable(name=name, file_name=file_name, reason=describe_storage_error(error)) from error + # A broken link, a folder, or a pipe, which would block if read. + raise self._unreadable(name=name, file_name=file_name, reason="Not a regular file") + + def _list_blob_documents_with_issues(self) -> DocumentListing: + """ + Read the blob container, reporting unreadable and badly named documents. - def _is_addressable_name(self, name: str) -> bool: + Returns: + DocumentListing: The documents and issues found under the configured prefix. + + Raises: + DocumentStorageUnavailableError: If a blob error affects the whole container. """ - Return whether a discovered document name is one this storage can address. + from azure.core.exceptions import AzureError - Ordinary files such as ``__init__.py`` or ``My-Script.py`` can sit alongside - valid documents, and a caller that fed such a name back into a read, write, or - delete would get a ``ValueError`` it has no way to anticipate. + 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 _is_missing_blob(error): + # Deleted after the container was listed, so it is no longer part of the listing. + continue + if _affects_whole_store(error): + raise self._unavailable(error) from error + unreadable.append( + UnreadableDocument( + name=name, file_name=file_name, message=describe_storage_error(error), error=error + ) + ) + return DocumentListing( + documents=MappingProxyType(documents), unreadable=tuple(unreadable), ignored=tuple(ignored) + ) + + @staticmethod + def _name_problem(name: str) -> str | None: + """ + Return why a discovered document name can't be addressed, or None if it can. Returns: - bool: Whether the name is a legal registry name. + str | None: The naming-rule error, or ``None`` for a legal name. """ try: validate_registry_name(name) except ValueError as error: - logger.warning(f"Ignoring stored document '{name}{self._extension}' in {self.display_source}: {error}") - return False - return True + return str(error) + return None def _read_document_bytes(self, name: str) -> bytes | None: """ - Read the raw bytes of one document. + Read the raw bytes of one document, agreeing with what a listing reports. Returns: bytes | None: Document content, or ``None`` if it does not exist. Raises: ValueError: If *name* is not a legal registry name. + DocumentUnreadableError: If the document exists but can't be read, such as a folder or a file + without read permission. + DocumentStorageUnavailableError: If the folder or container can't be used. """ validate_registry_name(name) - if self._is_blob: - from azure.core.exceptions import ResourceNotFoundError + file_name = f"{name}{self._extension}" + with self._unavailable_on_storage_error(): + if self._is_blob: + return self._read_blob_document(name=name, file_name=file_name) + path = self._local_directory() / file_name + return self._read_local_document(path=str(path), name=name, file_name=file_name) - with self._open_container_client() as client: - try: - return client.download_blob(self._get_blob_name(name)).readall() - except ResourceNotFoundError: + def _read_blob_document(self, *, name: str, file_name: str) -> bytes | None: + """ + Read one blob, treating only a missing blob as missing. + + Returns: + bytes | None: The content, or None if the blob doesn't exist. + + Raises: + DocumentUnreadableError: If the blob exists but can't be read. + DocumentStorageUnavailableError: If the error affects the whole container. + """ + from azure.core.exceptions import AzureError + + with self._open_container_client() as client: + try: + return client.download_blob(self._get_blob_name(name)).readall() + except AzureError as error: + if _is_missing_blob(error): return None + if _affects_whole_store(error): + raise self._unavailable(error) from error + raise self._unreadable(name=name, file_name=file_name, reason=describe_storage_error(error)) from error + + def _unreadable(self, *, name: str, file_name: str, reason: str) -> DocumentUnreadableError: + """ + Build the error for one document that exists but can't be read. + + Returns: + DocumentUnreadableError: The error naming the document and the storage location. + """ + return DocumentUnreadableError( + label=self._source_label, source=self.display_source, name=name, file_name=file_name, reason=reason + ) + + @contextmanager + def _unavailable_on_storage_error(self) -> Generator[None, None, None]: + """ + Report a failure of the folder or container itself as ``DocumentStorageUnavailableError``. + + Yields: + None: Control while storage is used. + + Raises: + DocumentStorageUnavailableError: If an ``OSError`` or Azure error escapes the block. + """ + from azure.core.exceptions import AzureError + + try: + yield + except DocumentStorageUnavailableError: + raise + except (OSError, AzureError) as error: + raise self._unavailable(error) from error - path = self._local_directory() / f"{name}{self._extension}" - return path.read_bytes() if path.is_file() else None + def _unavailable(self, error: Exception) -> DocumentStorageUnavailableError: + """ + Build the error for a folder or container that can't be used. + + Returns: + DocumentStorageUnavailableError: The error naming the storage location and a short reason. + """ + return DocumentStorageUnavailableError( + label=self._source_label, source=self.display_source, reason=describe_storage_error(error) + ) def _save_document(self, *, name: str, content: bytes) -> None: """ @@ -264,14 +632,16 @@ def _save_document(self, *, name: str, content: bytes) -> None: Raises: ValueError: If *name* is not a legal registry name. + DocumentStorageUnavailableError: If the folder or container can't be written. """ validate_registry_name(name) - if self._is_blob: - with self._open_container_client() as client: - client.upload_blob(name=self._get_blob_name(name), data=content, overwrite=True) - else: - directory = self._local_directory(create=True) - self._replace_file(path=directory / f"{name}{self._extension}", content=content) + with self._unavailable_on_storage_error(): + if self._is_blob: + with self._open_container_client() as client: + client.upload_blob(name=self._get_blob_name(name), data=content, overwrite=True) + else: + directory = self._local_directory(create=True) + self._replace_file(path=directory / f"{name}{self._extension}", content=content) @staticmethod def _compute_version(content: bytes) -> str: @@ -321,12 +691,14 @@ def _save_document_conditional(self, *, name: str, content: bytes, expected_vers Raises: DocumentConflictError: If the stored version does not match *expected_version*. ValueError: If *name* is not a legal registry name. - TimeoutError: If a local write could not acquire the document lock. + DocumentStorageUnavailableError: If the folder or container can't be written, or a local write + could not acquire the document lock. """ validate_registry_name(name) - if self._is_blob: - return self._save_blob_conditional(name=name, content=content, expected_version=expected_version) - return self._save_local_conditional(name=name, content=content, expected_version=expected_version) + with self._unavailable_on_storage_error(): + if self._is_blob: + return self._save_blob_conditional(name=name, content=content, expected_version=expected_version) + return self._save_local_conditional(name=name, content=content, expected_version=expected_version) def _save_local_conditional(self, *, name: str, content: bytes, expected_version: str | None) -> str: """ @@ -337,11 +709,12 @@ def _save_local_conditional(self, *, name: str, content: bytes, expected_version Raises: DocumentConflictError: If the stored version does not match *expected_version*. + DocumentUnreadableError: If something that can't be read is stored under the name. TimeoutError: If the document lock could not be acquired. """ path = self._local_directory(create=True) / f"{name}{self._extension}" with self._local_document_lock(path): - stored = path.read_bytes() if path.is_file() else None + stored = self._read_local_document(path=str(path), name=name, file_name=path.name) actual_version = None if stored is None else self._compute_version(stored) if actual_version != expected_version: raise self._conflict_error(name=name, expected_version=expected_version, actual_version=actual_version) @@ -360,7 +733,7 @@ def _save_blob_conditional(self, *, name: str, content: bytes, expected_version: Raises: DocumentConflictError: If the stored version does not match *expected_version*. - AzureError: If the container is unreachable or does not exist. + DocumentStorageUnavailableError: If the container is missing rather than the blob. """ from azure.core import MatchConditions from azure.core.exceptions import ResourceExistsError, ResourceModifiedError, ResourceNotFoundError @@ -388,7 +761,9 @@ def _save_blob_conditional(self, *, name: str, content: bytes, expected_version: etag=etag, match_condition=MatchConditions.IfNotModified, ) - except (ResourceModifiedError, ResourceNotFoundError): + except (ResourceModifiedError, ResourceNotFoundError) as error: + if isinstance(error, ResourceNotFoundError) and not _is_missing_blob(error): + raise self._unavailable(error) from error raise self._conflict_error( name=name, expected_version=expected_version, @@ -410,12 +785,15 @@ def _read_unmodified_blob_etag( Raises: DocumentConflictError: If the blob is absent or no longer holds that version. + DocumentStorageUnavailableError: If the container is missing rather than the blob. """ from azure.core.exceptions import ResourceNotFoundError try: downloader = client.download_blob(blob_name) - except ResourceNotFoundError: + except ResourceNotFoundError as error: + if not _is_missing_blob(error): + raise self._unavailable(error) from error raise self._conflict_error(name=name, expected_version=expected_version, actual_version=None) from None actual_version = self._compute_version(downloader.readall()) if actual_version != expected_version: @@ -428,12 +806,17 @@ def _read_blob_version(self, *, client: ContainerClient, blob_name: str) -> str Returns: str | None: The version token, or ``None`` if the blob does not exist. + + Raises: + DocumentStorageUnavailableError: If the container is missing rather than the blob. """ from azure.core.exceptions import ResourceNotFoundError try: return self._compute_version(client.download_blob(blob_name).readall()) - except ResourceNotFoundError: + except ResourceNotFoundError as error: + if not _is_missing_blob(error): + raise self._unavailable(error) from error return None @contextmanager @@ -480,7 +863,9 @@ def _acquire_document_lock(self, *, descriptor: int, lock_path: Path) -> None: deadline = time.monotonic() + self.LOCK_TIMEOUT_SECONDS while not _try_acquire_exclusive_lock(descriptor): if time.monotonic() >= deadline: - raise TimeoutError(f"Timed out waiting to write '{lock_path.name}'; another writer still holds it.") + raise TimeoutError( + errno.ETIMEDOUT, f"Timed out waiting for another writer to finish with '{lock_path.name}'" + ) time.sleep(self.LOCK_POLL_SECONDS) @staticmethod @@ -529,27 +914,31 @@ def _delete_document(self, name: str) -> bool: Raises: ValueError: If *name* is not a legal registry name. - TimeoutError: If a local delete could not acquire the document lock. + DocumentStorageUnavailableError: If the folder or container can't be written, or a local delete + could not acquire the document lock. """ validate_registry_name(name) - if self._is_blob: - from azure.core.exceptions import ResourceNotFoundError - - with self._open_container_client() as client: + with self._unavailable_on_storage_error(): + if self._is_blob: + from azure.core.exceptions import ResourceNotFoundError + + with self._open_container_client() as client: + try: + client.delete_blob(self._get_blob_name(name)) + except ResourceNotFoundError as error: + if not _is_missing_blob(error): + raise self._unavailable(error) from error + return False + return True + + path = self._local_directory(create=True) / f"{name}{self._extension}" + with self._local_document_lock(path): try: - client.delete_blob(self._get_blob_name(name)) - except ResourceNotFoundError: + path.unlink() + except FileNotFoundError: return False return True - path = self._local_directory(create=True) / f"{name}{self._extension}" - with self._local_document_lock(path): - try: - path.unlink() - except FileNotFoundError: - return False - return True - def _local_directory(self, *, create: bool = False) -> Path: """ Resolve the configured local directory. @@ -562,32 +951,6 @@ def _local_directory(self, *, create: bool = False) -> Path: directory.mkdir(parents=True, exist_ok=True) return directory - def _list_blob_documents(self) -> dict[str, bytes]: - """ - Read addressable documents from the configured Azure Blob container. - - Returns: - dict[str, bytes]: Document content keyed by blob stem. - """ - from azure.core.exceptions import AzureError - - documents: dict[str, bytes] = {} - 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: - name = PurePosixPath(blob_name.removeprefix(prefix or "")).stem - if not self._is_addressable_name(name): - continue - try: - documents[name] = client.download_blob(blob_name).readall() - except AzureError as error: - logger.warning(f"Skipping unreadable document '{blob_name}' in {self.display_source}: {error}") - return documents - def _parse_blob_source(self) -> tuple[str, str]: """ Split the configured source into a container URL and blob prefix. diff --git a/pyrit/registry/registry_recipe_storage.py b/pyrit/registry/registry_recipe_storage.py new file mode 100644 index 0000000000..f0a222e3c5 --- /dev/null +++ b/pyrit/registry/registry_recipe_storage.py @@ -0,0 +1,240 @@ +# 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 types import MappingProxyType +from typing import TYPE_CHECKING +from urllib.parse import urlparse + +from pydantic import ValidationError + +from pyrit.common.azure_storage import is_azure_blob_uri +from pyrit.models.catalog import ( + ComponentRecipe, + IgnoredRecipeFile, + RegistryRecipeKind, + RestoreFailure, + RestoreFailureReason, + SavedEntryRef, + StoredRecipe, + TechniqueRecipe, +) +from pyrit.registry.file_document_storage import FileDocumentStorage + +if TYPE_CHECKING: + from collections.abc import Mapping + +MAX_FAILURE_MESSAGE_LENGTH = 1_024 +_MAX_LOCATION_PART_LENGTH = 64 + + +def shorten_text(text: str, *, limit: int) -> str: + """ + Cut text that's longer than *limit*, marking the cut. + + Returns: + str: The text, at most *limit* characters long. + """ + return text if len(text) <= limit else f"{text[: limit - 3]}..." + + +def _escape_invalid_text(value: str) -> str: + """ + Escape what can't be written as UTF-8, such as the bytes of a file name that isn't UTF-8. + + Returns: + str: The text, with each such character written as a backslash escape. + """ + return value.encode("utf-8", "backslashreplace").decode("utf-8") + + +@dataclass(frozen=True) +class RecipeListing: + """Every saved recipe of one kind, with the ones that couldn't be read or parsed.""" + + recipes: Mapping[str, StoredRecipe] = field(default_factory=lambda: MappingProxyType({})) + failures: tuple[RestoreFailure, ...] = () + ignored: tuple[IgnoredRecipeFile, ...] = () + # The error behind each recipe that couldn't be read, by name, for the log only. + causes: Mapping[str, BaseException] = field(default_factory=lambda: MappingProxyType({}), compare=False, repr=False) + + +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. + + Text from the OS, such as a file name that isn't UTF-8, is escaped here so a report can be sent as JSON. + + Returns: + RecipeListing: Parsed recipes keyed by name, plus the failures and ignored files. + + Raises: + DocumentStorageUnavailableError: If the folder or 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=_escape_invalid_text(issue.message), + ) + for issue in listing.unreadable + ] + causes = {issue.name: issue.error for issue in listing.unreadable if issue.error is not None} + 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=_escape_invalid_text(issue.file_name), + message=_escape_invalid_text(issue.message), + ) + for issue in listing.ignored + ) + return RecipeListing( + recipes=MappingProxyType(recipes), + failures=tuple(failures), + ignored=ignored, + causes=MappingProxyType(causes), + ) + + 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 + too_deep = "Nested too deeply to read." + try: + # Recipes can be written by hand, and Windows tools often start a file with a byte-order mark. + payload = json.loads(content.decode("utf-8-sig")) + 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=too_deep) + if not isinstance(payload, dict): + return self._failure(name=name, version=version, reason=invalid, message="Not a JSON object.") + try: + json.dumps(payload, ensure_ascii=False).encode("utf-8") + except UnicodeEncodeError: + message = "Has a \\u escape that isn't a character (an unpaired surrogate)." + return self._failure(name=name, version=version, reason=invalid, message=message) + except RecursionError: + return self._failure(name=name, version=version, reason=invalid, message=too_deep) + model = TechniqueRecipe if self._kind is RegistryRecipeKind.TECHNIQUES else ComponentRecipe + schema_version = payload.get("schema_version") + current = model.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 + ) + 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=too_deep) + 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, naming each field and what's wrong with it, without the stored values. + + A field's location can repeat a stored key, such as an unknown one, so each part of a location + is shortened, and so is the whole summary. + + Returns: + str: The summary, at most ``MAX_FAILURE_MESSAGE_LENGTH`` characters long. + """ + problems = [] + for issue in error.errors(): + location = ".".join(shorten_text(str(part), limit=_MAX_LOCATION_PART_LENGTH) for part in issue["loc"]) + problems.append(f"{location or 'document'}: {issue['msg']}") + return shorten_text("Not a valid saved recipe: " + "; ".join(problems), limit=MAX_FAILURE_MESSAGE_LENGTH) diff --git a/tests/unit/backend/test_error_handlers.py b/tests/unit/backend/test_error_handlers.py index 2992a232c6..b45c781abd 100644 --- a/tests/unit/backend/test_error_handlers.py +++ b/tests/unit/backend/test_error_handlers.py @@ -5,12 +5,15 @@ Tests for backend error handler middleware. """ +import logging + import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from pydantic import BaseModel from pyrit.backend.middleware import register_error_handlers +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError, DocumentUnreadableError class TestErrorHandlers: @@ -125,6 +128,53 @@ async def test_endpoint() -> dict: assert data["title"] == "Forbidden" assert data["status"] == 403 + def test_storage_failure_returns_503_and_names_the_location_only_in_the_log( + self, app: FastAPI, client: TestClient, caplog: pytest.LogCaptureFixture + ) -> None: + """Test that the log names the storage location and the response doesn't.""" + error = DocumentStorageUnavailableError( + label="Scenario preset", source="/srv/presets", reason="Permission denied" + ) + + @app.get("/test") + async def test_endpoint() -> dict: + raise error + + with caplog.at_level(logging.DEBUG, logger="pyrit.backend.middleware.error_handlers"): + response = client.get("/test") + + assert response.status_code == 503 + assert response.json() == { + "type": "/errors/storage-unavailable", + "title": "Service Unavailable", + "status": 503, + "detail": "Scenario preset storage is unavailable", + "instance": "/test", + } + assert "Scenario preset storage at /srv/presets is unavailable: Permission denied (GET /test)" in ( + caplog.messages + ) + assert any(record.exc_info and record.exc_info[1] is error for record in caplog.records) + + def test_unreadable_document_returns_503_with_the_reason(self, app: FastAPI, client: TestClient) -> None: + """Test that one document that can't be read is named with its reason.""" + + @app.get("/test") + async def test_endpoint() -> dict: + raise DocumentUnreadableError( + label="Scenario preset", + source="/srv/presets", + name="nightly", + file_name="nightly.json", + reason="Not a regular file", + ) + + response = client.get("/test") + + assert response.status_code == 503 + assert response.json()["detail"] == "Scenario preset 'nightly' can't be read: Not a regular file" + assert "/srv/presets" not in response.text + def test_not_implemented_error_returns_501(self, app: FastAPI, client: TestClient) -> None: """Test that NotImplementedError returns 501 with RFC 7807 format.""" diff --git a/tests/unit/backend/test_initializer_service.py b/tests/unit/backend/test_initializer_service.py index e1d06ed5bf..18c6ea3f2e 100644 --- a/tests/unit/backend/test_initializer_service.py +++ b/tests/unit/backend/test_initializer_service.py @@ -5,11 +5,11 @@ Tests for backend initializer service and routes. """ +import logging from collections.abc import Iterator from unittest.mock import AsyncMock, MagicMock, patch import pytest -from azure.core.exceptions import AzureError from fastapi import HTTPException, status from fastapi.testclient import TestClient @@ -26,6 +26,7 @@ from pyrit.backend.services import InitializerService, get_initializer_service from pyrit.models import Parameter from pyrit.registry import InitializerMetadata +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError, DocumentUnreadableError @pytest.fixture @@ -470,13 +471,20 @@ def test_custom_initializer_routes_return_503_for_storage_failure( self, client_with_custom_initializers_enabled: TestClient, operation: str, + caplog: pytest.LogCaptureFixture, ) -> None: - """Test Blob failures are returned without exposing Azure SDK details.""" - with patch("pyrit.backend.routes.initializers.get_initializer_service") as mock_get_service: + """Test that the log names the storage location and the response doesn't.""" + error = DocumentStorageUnavailableError( + label="Custom initializer", source="/srv/initializers", reason="Permission denied" + ) + with ( + patch("pyrit.backend.routes.initializers.get_initializer_service") as mock_get_service, + caplog.at_level(logging.WARNING, logger="pyrit.backend.middleware.error_handlers"), + ): mock_service = MagicMock() - mock_service.list_custom_initializers_async = AsyncMock(side_effect=AzureError("credential details")) - mock_service.register_initializer_async = AsyncMock(side_effect=AzureError("credential details")) - mock_service.unregister_initializer_async = AsyncMock(side_effect=AzureError("credential details")) + mock_service.list_custom_initializers_async = AsyncMock(side_effect=error) + mock_service.register_initializer_async = AsyncMock(side_effect=error) + mock_service.unregister_initializer_async = AsyncMock(side_effect=error) mock_get_service.return_value = mock_service if operation == "list": response = client_with_custom_initializers_enabled.get("/api/initializers/custom") @@ -489,7 +497,29 @@ def test_custom_initializer_routes_return_503_for_storage_failure( response = client_with_custom_initializers_enabled.delete("/api/initializers/custom") assert response.status_code == status.HTTP_503_SERVICE_UNAVAILABLE - assert response.json()["detail"] == "Custom initializer storage is temporarily unavailable" + assert response.json()["detail"] == "Custom initializer storage is unavailable" + assert "/srv/initializers" not in response.text + assert "Custom initializer storage at /srv/initializers is unavailable: Permission denied" in caplog.text + + def test_deleting_a_script_that_cant_be_read_returns_503_with_the_reason( + self, client_with_custom_initializers_enabled: TestClient + ) -> None: + error = DocumentUnreadableError( + label="Custom initializer", + source="/srv/initializers", + name="custom", + file_name="custom.py", + reason="Not a regular file", + ) + with patch("pyrit.backend.routes.initializers.get_initializer_service") as mock_get_service: + mock_service = MagicMock() + mock_service.unregister_initializer_async = AsyncMock(side_effect=error) + mock_get_service.return_value = mock_service + + response = client_with_custom_initializers_enabled.delete("/api/initializers/custom") + + assert response.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + assert response.json()["detail"] == "Custom initializer 'custom' can't be read: Not a regular file" def test_post_root_still_registers_runtime_initializer( self, client_with_custom_initializers_enabled: TestClient diff --git a/tests/unit/backend/test_scenario_preset_routes.py b/tests/unit/backend/test_scenario_preset_routes.py index 92722eaacb..31750c98e9 100644 --- a/tests/unit/backend/test_scenario_preset_routes.py +++ b/tests/unit/backend/test_scenario_preset_routes.py @@ -3,11 +3,11 @@ """Tests for the scenario preset routes.""" +import logging from collections.abc import Iterator from unittest.mock import AsyncMock, MagicMock, patch import pytest -from azure.core.exceptions import AzureError from fastapi.testclient import TestClient from pyrit.backend.main import app @@ -24,6 +24,7 @@ ScenarioPresetNameRegisteredError, resolve_scenario_preset_run_request, ) +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError PRESET_NAME = "quick_scan" SCENARIO_NAME = "foundry.red_team_agent" @@ -146,13 +147,22 @@ def test_list_is_readable_without_admin(self, anonymous_client: TestClient, serv assert anonymous_client.get("/api/scenario-presets").status_code == 200 - def test_storage_failure_is_reported_without_sdk_detail(self, client: TestClient, service: MagicMock) -> None: - service.list_presets_async = AsyncMock(side_effect=AzureError("container 'x' key=secret")) + def test_storage_failure_names_the_source_only_in_the_log( + self, client: TestClient, service: MagicMock, caplog: pytest.LogCaptureFixture + ) -> None: + service.list_presets_async = AsyncMock( + side_effect=DocumentStorageUnavailableError( + label="Scenario preset", source="/srv/presets", reason="Permission denied" + ) + ) - response = client.get("/api/scenario-presets") + with caplog.at_level(logging.WARNING, logger="pyrit.backend.middleware.error_handlers"): + response = client.get("/api/scenario-presets") assert response.status_code == 503 - assert "secret" not in response.text + assert response.json()["detail"] == "Scenario preset storage is unavailable" + assert "/srv/presets" not in response.text + assert "Scenario preset storage at /srv/presets is unavailable: Permission denied" in caplog.text def test_each_preset_carries_its_own_run_size(self, client: TestClient, service: MagicMock) -> None: estimate = _estimate(42) @@ -416,7 +426,9 @@ def test_a_missing_preset_still_conflicts_instead_of_being_recreated( def test_a_failed_author_read_is_reported_without_storage_details( self, *, client: TestClient, service: MagicMock ) -> None: - service.get_preset_async.side_effect = AzureError("key=secret") + service.get_preset_async.side_effect = DocumentStorageUnavailableError( + label="Scenario preset", source="/srv/presets", reason="Permission denied" + ) response = client.put( f"/api/scenario-presets/{PRESET_NAME}", @@ -424,7 +436,7 @@ def test_a_failed_author_read_is_reported_without_storage_details( ) assert response.status_code == 503 - assert "secret" not in response.text + assert "/srv/presets" not in response.text service.save_preset_async.assert_not_called() def test_a_body_naming_a_different_preset_is_rejected(self, client: TestClient, service: MagicMock) -> None: diff --git a/tests/unit/backend/test_scenario_suite_routes.py b/tests/unit/backend/test_scenario_suite_routes.py index e53fa141b8..fd6f1345fc 100644 --- a/tests/unit/backend/test_scenario_suite_routes.py +++ b/tests/unit/backend/test_scenario_suite_routes.py @@ -3,6 +3,7 @@ """Tests for scenario suite definition routes.""" +import logging from collections.abc import Iterator from datetime import UTC, datetime from pathlib import Path @@ -11,7 +12,6 @@ from uuid import UUID import pytest -from azure.core.exceptions import AzureError from fastapi.testclient import TestClient from pyrit.backend.main import app @@ -48,6 +48,7 @@ ScenarioSuiteRegistry, ScenarioSuiteStorage, ) +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError SUITE_NAME = "weekly" ENTRY_ID = UUID("11111111-1111-4111-8111-111111111111") @@ -181,7 +182,9 @@ def unavailable_preset_backend( ) -> Iterator[tuple[ScenarioSuiteService, ScenarioSuiteStorage]]: """Run routes against real suite storage and an unavailable preset registry.""" registry = MagicMock(spec=ScenarioPresetRegistry) - registry.get_preset.side_effect = AzureError("expired preset container signature") + registry.get_preset.side_effect = DocumentStorageUnavailableError( + label="Scenario preset", source="/srv/presets", reason="Permission denied" + ) storage = ScenarioSuiteStorage(source=str(tmp_path)) scenario_service = ScenarioSuiteService(registry=ScenarioSuiteRegistry(storage=storage)) with ( @@ -257,7 +260,10 @@ def test_launch_is_rejected_during_runtime_reinitialization( (ScenarioSuiteLaunchConflictError("changed"), 409), (ScenarioSuiteNotRunnableError("not runnable"), 400), (ScenarioSuiteExecutionCapacityError("wait for an active suite"), 429), - (AzureError("storage unavailable"), 503), + ( + DocumentStorageUnavailableError(label="Scenario suite", source="/srv/suites", reason="Permission denied"), + 503, + ), (RuntimeError("scheduling is stopping"), 503), ], ) @@ -273,8 +279,8 @@ def test_launch_maps_preflight_failures( response = anonymous_client.post(f"/api/scenario-suites/{SUITE_NAME}/launch", json=_launch_body()) assert response.status_code == expected_status - if isinstance(error, AzureError): - assert "storage unavailable" not in response.text + if isinstance(error, DocumentStorageUnavailableError): + assert response.json()["detail"] == "Scenario suite storage is unavailable" if isinstance(error, ScenarioSuiteExecutionCapacityError): assert response.json()["detail"] == "wait for an active suite" @@ -369,14 +375,25 @@ def test_get_registered_suite_returns_origin_and_null_version(client: TestClient assert response.json()["version"] is None -def test_storage_failure_is_sanitized(client: TestClient, service: MagicMock) -> None: - """Test that Azure SDK details do not reach API callers.""" - service.list_suites_async = AsyncMock(side_effect=AzureError("container key=secret")) +def test_storage_failure_names_the_source_only_in_the_log( + client: TestClient, service: MagicMock, caplog: pytest.LogCaptureFixture +) -> None: + """Test that the log names the storage location and the response doesn't.""" + service.list_suites_async = AsyncMock( + side_effect=DocumentStorageUnavailableError( + label="Scenario suite", + source="https://account.blob.core.windows.net/suites", + reason="ClientAuthenticationError", + ) + ) - response = client.get("/api/scenario-suites") + with caplog.at_level(logging.WARNING, logger="pyrit.backend.middleware.error_handlers"): + response = client.get("/api/scenario-suites") assert response.status_code == 503 - assert "secret" not in response.text + assert response.json()["detail"] == "Scenario suite storage is unavailable" + assert "account.blob.core.windows.net" not in response.text + assert "https://account.blob.core.windows.net/suites is unavailable: ClientAuthenticationError" in caplog.text def test_create_succeeds_when_preset_storage_is_unavailable( diff --git a/tests/unit/backend/test_scenario_suite_service.py b/tests/unit/backend/test_scenario_suite_service.py index 1b80a3325e..f1e60b5990 100644 --- a/tests/unit/backend/test_scenario_suite_service.py +++ b/tests/unit/backend/test_scenario_suite_service.py @@ -9,7 +9,6 @@ from uuid import UUID import pytest -from azure.core.exceptions import AzureError from pyrit.backend.services import ScenarioSuiteService, get_scenario_suite_service from pyrit.models import PresetEntry, PresetIssue, ScenarioPreset, ScenarioSuite, ScenarioSuiteEntry @@ -19,6 +18,7 @@ ScenarioSuiteRegistry, ScenarioSuiteStorage, ) +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError, DocumentUnreadableError ENTRY_ONE = UUID("11111111-1111-4111-8111-111111111111") ENTRY_TWO = UUID("22222222-2222-4222-8222-222222222222") @@ -119,8 +119,17 @@ async def test_registered_suite_is_listed_with_nullable_version( @pytest.mark.parametrize( "error", - [AzureError("expired preset credentials"), OSError("preset directory unavailable")], - ids=["azure", "local"], + [ + DocumentStorageUnavailableError(label="Scenario preset", source="/srv/presets", reason="Permission denied"), + DocumentUnreadableError( + label="Scenario preset", + source="/srv/presets", + name="nightly", + file_name="nightly.json", + reason="Not a regular file", + ), + ], + ids=["storage", "one_document"], ) async def test_preset_storage_failure_is_advisory_after_committed_save( service: ScenarioSuiteService, 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..48c6fa68ca --- /dev/null +++ b/tests/unit/registry/test_file_document_storage.py @@ -0,0 +1,610 @@ +# 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 ( + ClientAuthenticationError, + HttpResponseError, + ResourceNotFoundError, + ServiceRequestError, + ServiceResponseError, +) +from azure.storage.blob import ContainerClient, StorageErrorCode + +import pyrit.registry.file_document_storage as storage_module +from pyrit.registry.file_document_storage import ( + DocumentConflictError, + DocumentListing, + DocumentStorageUnavailableError, + DocumentUnreadableError, + FileDocumentStorage, + UnreadableDocument, + _read_regular_file, + describe_storage_error, +) + +_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 _blob_error( + *, status_code: int, error_code: str, message: str = "Storage error.", missing: bool = False +) -> HttpResponseError: + error_type = ResourceNotFoundError if missing else HttpResponseError + error = error_type(message=message) + error.status_code = status_code + error.error_code = error_code + return error + + +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.object(storage_module, "_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") + + def read(path: str) -> bytes | None: + if Path(path).name == "second.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 = _local_storage(tmp_path)._list_documents_with_issues() + documents = _local_storage(tmp_path)._list_documents() + + assert listing == DocumentListing( + documents={"first": b"one"}, + unreadable=(UnreadableDocument(name="second", file_name="second.json", message="Permission denied"),), + ) + assert isinstance(listing.unreadable[0].error, PermissionError) + assert documents == {"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", file_name="swapped.json", message="Not a regular file"),), + ) + assert listing.unreadable[0].error is None + assert _count_open_files() == open_files + + @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + @pytest.mark.timeout(60, method="thread") + def test_checks_the_opened_file_rather_than_the_path(self, tmp_path: Path) -> None: + document = tmp_path / "document.json" + document.write_bytes(b"original") + 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.object(storage_module, "_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.object(storage_module, "_open", side_effect=replace_then_open): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing( + unreadable=(UnreadableDocument(name="document", file_name="document.json", 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", file_name="broken_link.json", message="Not a regular file"), + UnreadableDocument(name="folder", file_name="folder.json", 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.object(storage_module, "_lstat", side_effect=PermissionError(13, "Permission denied")), + ): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing( + unreadable=(UnreadableDocument(name="deleted", file_name="deleted.json", 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_with_the_platform_case_rule(self, tmp_path: Path) -> None: + (tmp_path / "lower.json").write_bytes(b"one") + (tmp_path / "upper.JSON").write_bytes(b"two") + + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert sorted(listing.documents) == sorted(path.stem for path in tmp_path.glob("*.json")) + + 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.object(storage_module, "_scandir", side_effect=PermissionError(13, "denied")): + with pytest.raises(DocumentStorageUnavailableError) as raised: + _local_storage(tmp_path)._list_documents_with_issues() + + assert isinstance(raised.value.__cause__, PermissionError) + assert (raised.value.label, raised.value.source, raised.value.reason) == ( + "Test documents", + str(tmp_path), + "denied", + ) + assert str(raised.value) == f"Test documents storage at {tmp_path} is unavailable: denied" + + +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": + raise _blob_error(status_code=409, error_code="BlobArchived", message=f"Archived: {_SAS_SOURCE}") + 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", file_name="broken.json", message="HttpResponseError (409, BlobArchived)"), + ) + assert isinstance(listing.unreadable[0].error, HttpResponseError) + 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"]) + + def download_blob(blob_name: str) -> SimpleNamespace: + if blob_name == "saved/deleted.json": + raise _blob_error(status_code=404, error_code=StorageErrorCode.BLOB_NOT_FOUND, missing=True) + 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"}) + + @pytest.mark.parametrize( + "error", + [ + _blob_error(status_code=404, error_code=StorageErrorCode.CONTAINER_NOT_FOUND, missing=True), + _blob_error(status_code=403, error_code=StorageErrorCode.AUTHORIZATION_FAILURE), + _blob_error(status_code=401, error_code="InvalidAuthenticationInfo"), + ClientAuthenticationError(message="no token"), + ServiceRequestError(message="connection refused"), + ServiceResponseError(message="connection reset"), + ], + ) + def test_raises_when_an_error_affects_the_whole_container(self, error: Exception) -> None: + client = _blob_client(names=["saved/first.json", "saved/second.json"]) + client.download_blob.side_effect = error + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + with pytest.raises(DocumentStorageUnavailableError) as raised: + _blob_storage()._list_documents_with_issues() + + assert raised.value.__cause__ is error + assert "secret-signature" not in str(raised.value) + + 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(DocumentStorageUnavailableError) as raised: + _blob_storage()._list_documents_with_issues() + + assert isinstance(raised.value.__cause__, HttpResponseError) + + +class TestReadDocument: + def test_reads_a_document_and_treats_a_missing_one_as_missing(self, tmp_path: Path) -> None: + (tmp_path / "present.json").write_bytes(b"content") + storage = _local_storage(tmp_path) + + assert storage._read_document_bytes("present") == b"content" + assert storage._read_document_bytes("absent") is None + assert _local_storage(tmp_path / "missing_folder")._read_document_bytes("absent") is None + + @pytest.mark.parametrize("entry", ["folder", "broken link", "pipe"]) + @pytest.mark.timeout(60, method="thread") + def test_reports_what_the_listing_reports_as_unreadable(self, tmp_path: Path, entry: str) -> None: + path = tmp_path / "odd.json" + if entry == "folder": + path.mkdir() + elif entry == "broken link": + try: + path.symlink_to(tmp_path / "missing_target.json") + except OSError: + pytest.skip("Creating symbolic links isn't permitted here") + elif not hasattr(os, "mkfifo"): + pytest.skip("Named pipes need os.mkfifo") + else: + os.mkfifo(path) + storage = _local_storage(tmp_path) + + with pytest.raises(DocumentUnreadableError) as raised: + storage._read_document_bytes("odd") + + assert (raised.value.name, raised.value.file_name, raised.value.reason) == ( + "odd", + "odd.json", + "Not a regular file", + ) + assert [issue.name for issue in storage._list_documents_with_issues().unreadable] == ["odd"] + + def test_reports_a_file_it_cannot_open(self, tmp_path: Path) -> None: + (tmp_path / "locked.json").write_bytes(b"secret") + + with patch.object(storage_module, "_open", side_effect=PermissionError(13, "Permission denied")): + with pytest.raises(DocumentUnreadableError) as raised: + _local_storage(tmp_path)._read_document_bytes("locked") + + assert raised.value.reason == "Permission denied" + assert isinstance(raised.value.__cause__, PermissionError) + assert str(raised.value) == f"Test documents 'locked.json' in {tmp_path} can't be read: Permission denied" + + def test_blob_treats_only_a_missing_blob_as_missing(self) -> None: + client = _blob_client(names=[]) + client.download_blob.side_effect = _blob_error( + status_code=404, error_code=StorageErrorCode.BLOB_NOT_FOUND, missing=True + ) + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + assert _blob_storage()._read_document_bytes("absent") is None + + def test_blob_raises_when_the_container_is_gone(self) -> None: + client = _blob_client(names=[]) + client.download_blob.side_effect = _blob_error( + status_code=404, error_code=StorageErrorCode.CONTAINER_NOT_FOUND, missing=True + ) + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + with pytest.raises(DocumentStorageUnavailableError) as raised: + _blob_storage()._read_document_bytes("present") + + assert not isinstance(raised.value, DocumentUnreadableError) + + def test_blob_reports_a_blob_it_cannot_read(self) -> None: + client = _blob_client(names=[]) + client.download_blob.side_effect = _blob_error(status_code=409, error_code="BlobArchived") + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + with pytest.raises(DocumentUnreadableError) as raised: + _blob_storage()._read_document_bytes("archived") + + assert raised.value.reason == "HttpResponseError (409, BlobArchived)" + + +class TestDeleteDocument: + def test_blob_reports_a_missing_blob_as_already_gone(self) -> None: + client = _blob_client(names=[]) + client.delete_blob.side_effect = _blob_error( + status_code=404, error_code=StorageErrorCode.BLOB_NOT_FOUND, missing=True + ) + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + assert _blob_storage()._delete_document("absent") is False + + def test_blob_raises_when_the_container_is_gone(self) -> None: + client = _blob_client(names=[]) + client.delete_blob.side_effect = _blob_error( + status_code=404, error_code=StorageErrorCode.CONTAINER_NOT_FOUND, missing=True + ) + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + with pytest.raises(DocumentStorageUnavailableError): + _blob_storage()._delete_document("present") + + +class TestConditionalSave: + def test_local_refuses_to_replace_something_it_cannot_read(self, tmp_path: Path) -> None: + (tmp_path / "odd.json").mkdir() + + with pytest.raises(DocumentUnreadableError): + _local_storage(tmp_path)._save_document_conditional(name="odd", content=b"new", expected_version=None) + + assert (tmp_path / "odd.json").is_dir() + + def test_blob_update_of_a_deleted_blob_is_a_conflict(self) -> None: + client = _blob_client(names=[]) + client.download_blob.side_effect = _blob_error( + status_code=404, error_code=StorageErrorCode.BLOB_NOT_FOUND, missing=True + ) + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + with pytest.raises(DocumentConflictError) as raised: + _blob_storage()._save_document_conditional(name="gone", content=b"new", expected_version="v1") + + assert raised.value.actual_version is None + + def test_blob_update_raises_when_the_container_is_gone(self) -> None: + client = _blob_client(names=[]) + client.download_blob.side_effect = _blob_error( + status_code=404, error_code=StorageErrorCode.CONTAINER_NOT_FOUND, missing=True + ) + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + with pytest.raises(DocumentStorageUnavailableError) as raised: + _blob_storage()._save_document_conditional(name="present", content=b"new", expected_version="v1") + + assert not isinstance(raised.value, DocumentUnreadableError) + client.upload_blob.assert_not_called() + + +def test_a_listing_cannot_be_changed_by_its_reader(tmp_path: Path) -> None: + (tmp_path / "first.json").write_bytes(b"one") + + listing = _local_storage(tmp_path)._list_documents_with_issues() + + with pytest.raises(TypeError): + listing.documents["second"] = b"two" # type: ignore[index] + + +class TestListDocuments: + def test_logs_what_it_leaves_out_and_returns_the_rest(self, tmp_path: Path) -> None: + (tmp_path / "good.json").write_bytes(b"ok") + (tmp_path / "Bad-Name.json").write_bytes(b"never read") + (tmp_path / "folder.json").mkdir() + + with patch("pyrit.registry.file_document_storage.logger") as log: + documents = _local_storage(tmp_path)._list_documents() + + assert documents == {"good": b"ok"} + messages = [call.args[0] for call in log.warning.call_args_list] + assert len(messages) == 2 + assert messages[0].startswith(f"Ignoring stored document 'Bad-Name.json' in {tmp_path}: ") + assert messages[1] == f"Skipping unreadable document 'folder.json' in {tmp_path}: Not a regular file" + + @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + @pytest.mark.timeout(60, method="thread") + def test_leaves_out_a_pipe_named_like_a_document_without_blocking(self, tmp_path: Path) -> None: + (tmp_path / "good.json").write_bytes(b"ok") + os.mkfifo(tmp_path / "pipe.json") + + assert _local_storage(tmp_path)._list_documents() == {"good": b"ok"} + + def test_raises_when_the_directory_cannot_be_listed(self, tmp_path: Path) -> None: + with patch.object(storage_module, "_scandir", side_effect=PermissionError(13, "denied")): + with pytest.raises(DocumentStorageUnavailableError): + _local_storage(tmp_path)._list_documents() + + def test_logs_an_unreadable_blob_without_exposing_the_url(self) -> None: + client = _blob_client(names=["saved/good.json", "saved/broken.json"]) + + def download_blob(blob_name: str) -> SimpleNamespace: + if blob_name == "saved/broken.json": + raise HttpResponseError(message=f"Forbidden: {_SAS_SOURCE}") + 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), + patch("pyrit.registry.file_document_storage.logger") as log, + ): + documents = _blob_storage()._list_documents() + + assert documents == {"good": b"content"} + message = log.warning.call_args.args[0] + assert message.startswith( + "Skipping unreadable document 'broken.json' in https://account.blob.core.windows.net/" + ) + assert message.endswith(": HttpResponseError") + assert "secret-signature" not in message + + +class TestDescribeStorageError: + @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 describe_storage_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 describe_storage_error(error) == expected diff --git a/tests/unit/registry/test_initializer_registry.py b/tests/unit/registry/test_initializer_registry.py index e2762ce122..a57321845e 100644 --- a/tests/unit/registry/test_initializer_registry.py +++ b/tests/unit/registry/test_initializer_registry.py @@ -180,7 +180,7 @@ def test_unregister_and_cleanup_removes_entry_and_file(lazy_registry): def test_unregister_and_cleanup_keeps_entry_when_storage_delete_fails(lazy_registry: InitializerRegistry) -> None: """Test that failed storage cleanup leaves the runtime registration retryable.""" storage = MagicMock(spec=CustomInitializerStorage) - storage.list_scripts.return_value = {"custom": _VALID_SCRIPT} + storage.list_script_names.return_value = {"custom"} storage.delete_script.side_effect = OSError("storage unavailable") lazy_registry._custom_storage = storage lazy_registry._classes["custom"] = _ParamInitializer @@ -194,7 +194,7 @@ def test_unregister_and_cleanup_keeps_entry_when_storage_delete_fails(lazy_regis def test_unregister_and_cleanup_removes_invalid_stored_source(lazy_registry: InitializerRegistry) -> None: """Test invalid stored scripts can be removed without a runtime registration.""" storage = MagicMock(spec=CustomInitializerStorage) - storage.list_scripts.return_value = {"invalid": "not Python"} + storage.list_script_names.return_value = {"invalid"} lazy_registry._custom_storage = storage lazy_registry.unregister_and_cleanup("invalid") @@ -205,7 +205,7 @@ def test_unregister_and_cleanup_removes_invalid_stored_source(lazy_registry: Ini def test_unregister_and_cleanup_removes_source_shadowed_by_builtin(lazy_registry: InitializerRegistry) -> None: """Test deleting stored source preserves a built-in registration with the same name.""" storage = MagicMock(spec=CustomInitializerStorage) - storage.list_scripts.return_value = {"builtin": _VALID_SCRIPT} + storage.list_script_names.return_value = {"builtin"} lazy_registry._custom_storage = storage lazy_registry._classes["builtin"] = _ParamInitializer lazy_registry._builtin_names.add("builtin") @@ -216,6 +216,24 @@ def test_unregister_and_cleanup_removes_source_shadowed_by_builtin(lazy_registry assert lazy_registry.get_class("builtin") is _ParamInitializer +def test_unregister_and_cleanup_removes_a_script_that_cannot_be_read( + lazy_registry: InitializerRegistry, tmp_path: Path +) -> None: + """Test that a stored script the listing reports as unreadable, or can't decode, can still be deleted.""" + lazy_registry.configure_custom_scripts_source(str(tmp_path)) + (tmp_path / "not_utf8.py").write_bytes(b"\xff\xfe") + try: + (tmp_path / "broken_link.py").symlink_to(tmp_path / "missing_target.py") + except OSError: + pytest.skip("Creating symbolic links isn't permitted here") + + lazy_registry.unregister_and_cleanup("not_utf8") + lazy_registry.unregister_and_cleanup("broken_link") + + assert not (tmp_path / "not_utf8.py").exists() + assert not (tmp_path / "broken_link.py").is_symlink() + + def test_register_from_content_uses_configured_storage(lazy_registry: InitializerRegistry, tmp_path: Path) -> None: """Test runtime registration persists source in the configured directory.""" lazy_registry.configure_custom_scripts_source(str(tmp_path)) 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..077169efe3 --- /dev/null +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -0,0 +1,279 @@ +# 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 pydantic import ValidationError + +from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH +from pyrit.models.catalog import ( + ComponentRecipe, + RegistryRecipeKind, + RestoreFailureReason, + SavedEntryRef, + StoredRecipe, + TechniqueRecipe, +) +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError, _read_regular_file +from pyrit.registry.registry_recipe_storage import MAX_FAILURE_MESSAGE_LENGTH, 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 TestRecipeModels: + def test_a_document_without_a_schema_version_is_version_one(self) -> None: + document = {key: value for key, value in _TARGET.items() if key != "schema_version"} + + with patch.object(ComponentRecipe, "CURRENT_SCHEMA_VERSION", 2): + assert ComponentRecipe.model_validate(document).schema_version == 1 + + def test_rejects_a_schema_version_newer_than_the_model_reads(self) -> None: + with pytest.raises(ValidationError, match="newer than this PyRIT reads"): + ComponentRecipe.model_validate({**_TARGET, "schema_version": 2}) + + @pytest.mark.parametrize( + ("kind", "recipe"), + [ + (RegistryRecipeKind.TECHNIQUES, ComponentRecipe.model_validate(_TARGET)), + (RegistryRecipeKind.TARGETS, TechniqueRecipe.model_validate({"type": "RedTeamingAttack"})), + ], + ) + def test_a_stored_recipe_must_match_its_registry( + self, kind: RegistryRecipeKind, recipe: ComponentRecipe | TechniqueRecipe + ) -> None: + with pytest.raises(ValidationError, match="entry needs a"): + StoredRecipe(ref=SavedEntryRef(kind=kind, name="mismatched"), recipe=recipe, version="v1") + + +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 == "bad\\udcffname.json" + assert "bad\\udcffname" in ignored.message + assert json.loads(ignored.model_dump_json())["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 + assert isinstance(listing.causes["locked"], PermissionError) + (ignored,) = listing.ignored + assert (ignored.kind, ignored.file_name) == (RegistryRecipeKind.SCORERS, "Bad-Name.json") + assert "Bad-Name" in ignored.message + + def test_accepts_a_byte_order_mark(self, tmp_path: Path) -> None: + _write(tmp_path / "targets", "from_notepad", b"\xef\xbb\xbf" + json.dumps(_TARGET).encode("utf-8")) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + assert list(listing.recipes) == ["from_notepad"] + + @pytest.mark.parametrize( + "document", + [ + '{"class_module": "m", "class_name": "C", "params": {"text": "bad \\udcff"}}', + '{"class_module": "m", "class_name": "C", "params": {"bad \\ud800": 1}}', + ], + ids=["value", "key"], + ) + def test_reports_an_unpaired_surrogate_as_invalid(self, tmp_path: Path, document: str) -> None: + _write(tmp_path / "targets", "odd_text", document) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + (failure,) = listing.failures + assert (failure.name, failure.reason) == ("odd_text", RestoreFailureReason.INVALID_RECIPE) + assert "unpaired surrogate" in failure.message + + def test_keeps_text_that_is_valid_unicode(self, tmp_path: Path) -> None: + document = {**_TARGET, "params": {"greeting": "caf\u00e9 \U0001f600"}} + _write(tmp_path / "targets", "friendly", json.dumps(document)) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + assert listing.recipes["friendly"].recipe.params["greeting"] == "caf\u00e9 \U0001f600" + + def test_checks_a_technique_against_the_technique_format_version(self, tmp_path: Path) -> None: + _write(tmp_path / "techniques", "next_format", json.dumps({"schema_version": 2, "type": "RedTeamingAttack"})) + _write(tmp_path / "targets", "next_format", json.dumps({**_TARGET, "schema_version": 2})) + + with patch.object(TechniqueRecipe, "CURRENT_SCHEMA_VERSION", 2): + techniques = RegistryRecipeStorage(kind=RegistryRecipeKind.TECHNIQUES, source=str(tmp_path)).list_recipes() + targets = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + assert list(techniques.recipes) == ["next_format"] + assert targets.failures[0].reason is RestoreFailureReason.NEWER_SCHEMA_VERSION + + def test_shortens_what_a_validation_error_repeats(self, tmp_path: Path) -> None: + long_key = "k" * 10_000 + many_keys = {f"unknown_{index:03d}_{'x' * 40}": True for index in range(100)} + _write(tmp_path / "targets", "long_key", json.dumps({**_TARGET, long_key: True})) + _write(tmp_path / "targets", "many_keys", json.dumps({**_TARGET, **many_keys})) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + messages = {failure.name: failure.message for failure in listing.failures} + assert "k" * 61 + "..." in messages["long_key"] + assert "k" * 65 not in messages["long_key"] + assert len(messages["many_keys"]) == MAX_FAILURE_MESSAGE_LENGTH + assert messages["many_keys"].endswith("...") + + def test_a_listing_cannot_be_changed_by_its_reader(self, tmp_path: Path) -> None: + _write(tmp_path / "targets", "contoso_gpt4o", json.dumps(_TARGET)) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + with pytest.raises(TypeError): + listing.recipes["other"] = listing.recipes["contoso_gpt4o"] # type: ignore[index] + + def test_raises_when_the_folder_cannot_be_listed(self, tmp_path: Path) -> None: + (tmp_path / "targets").write_text("a file where the folder should be") + + with pytest.raises(DocumentStorageUnavailableError): + RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() diff --git a/tests/unit/registry/test_scenario_preset_storage.py b/tests/unit/registry/test_scenario_preset_storage.py index 137ef3bfd3..2dcb6af285 100644 --- a/tests/unit/registry/test_scenario_preset_storage.py +++ b/tests/unit/registry/test_scenario_preset_storage.py @@ -15,7 +15,11 @@ from pyrit.models import DatasetSourceInput, ScenarioDatasetRequest, ScenarioPreset from pyrit.models.catalog.scenario_preset import PRESET_SCHEMA_VERSION from pyrit.registry import ScenarioPresetConflictError, ScenarioPresetStorage -from pyrit.registry.file_document_storage import _release_exclusive_lock, _try_acquire_exclusive_lock +from pyrit.registry.file_document_storage import ( + DocumentStorageUnavailableError, + _release_exclusive_lock, + _try_acquire_exclusive_lock, +) def _make_preset(**overrides: object) -> ScenarioPreset: @@ -415,6 +419,7 @@ def test_blob_storage_reads_missing_document_as_none() -> None: client = MagicMock() client.__enter__.return_value = client client.download_blob.side_effect = ResourceNotFoundError("missing") + client.download_blob.side_effect.error_code = "BlobNotFound" storage = ScenarioPresetStorage(source="https://account.blob.core.windows.net/presets?sig=secret") with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): @@ -440,6 +445,7 @@ def test_blob_delete_of_a_missing_document_reports_absence_instead_of_raising() client = MagicMock() client.__enter__.return_value = client client.delete_blob.side_effect = ResourceNotFoundError("missing") + client.delete_blob.side_effect.error_code = "BlobNotFound" storage = ScenarioPresetStorage(source="https://account.blob.core.windows.net/presets?sig=secret") with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): @@ -552,7 +558,7 @@ def fail(source: object, target: object) -> None: raise OSError(28, "No space left on device") with patch("os.replace", fail): - with pytest.raises(OSError): + with pytest.raises(DocumentStorageUnavailableError): storage.save_preset(preset=_make_preset(description="second"), expected_version=first.version) loaded = storage.load_preset("nightly") @@ -627,7 +633,7 @@ def test_a_second_writer_cannot_enter_while_a_save_is_in_flight(tmp_path: Path) real_replace_file = ScenarioPresetStorage._replace_file def replace_once_a_competitor_has_tried(*, path: Path, content: bytes) -> None: - with pytest.raises(TimeoutError) as error: + with pytest.raises(DocumentStorageUnavailableError) as error: competitor.save_preset(preset=_make_preset(description="racer"), expected_version=first.version) refusals.append(str(error.value)) real_replace_file(path=path, content=content) @@ -651,7 +657,7 @@ def test_a_delete_cannot_enter_while_a_save_is_in_flight(tmp_path: Path) -> None real_replace_file = ScenarioPresetStorage._replace_file def replace_once_a_deleter_has_tried(*, path: Path, content: bytes) -> None: - with pytest.raises(TimeoutError) as error: + with pytest.raises(DocumentStorageUnavailableError) as error: competitor.delete_preset("nightly") refusals.append(str(error.value)) real_replace_file(path=path, content=content) @@ -699,7 +705,7 @@ def test_a_lock_held_by_a_live_writer_is_respected(tmp_path: Path) -> None: try: with patch.object(ScenarioPresetStorage, "LOCK_TIMEOUT_SECONDS", 0.1): - with pytest.raises(TimeoutError, match="nightly.json.lock"): + with pytest.raises(DocumentStorageUnavailableError, match="nightly.json.lock"): storage.save_preset(preset=_make_preset(), expected_version=None) finally: _release_exclusive_lock(descriptor) @@ -772,7 +778,7 @@ def test_blob_create_refuses_to_overwrite_a_document_another_writer_just_created def test_blob_create_against_a_missing_container_is_not_reported_as_a_conflict() -> None: """Test that an absent container surfaces as a storage error instead of a phantom concurrent edit.""" - from azure.core.exceptions import AzureError, ResourceNotFoundError + from azure.core.exceptions import ResourceNotFoundError client = MagicMock() client.__enter__.return_value = client @@ -780,7 +786,7 @@ def test_blob_create_against_a_missing_container_is_not_reported_as_a_conflict() storage = ScenarioPresetStorage(source="https://account.blob.core.windows.net/absent?sig=secret") with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): - with pytest.raises(AzureError): + with pytest.raises(DocumentStorageUnavailableError): storage.save_preset(preset=_make_preset(), expected_version=None) @@ -822,6 +828,7 @@ def test_blob_update_of_a_deleted_document_is_a_conflict() -> None: client = MagicMock() client.__enter__.return_value = client client.download_blob.side_effect = ResourceNotFoundError("missing") + client.download_blob.side_effect.error_code = "BlobNotFound" storage = ScenarioPresetStorage(source="https://account.blob.core.windows.net/presets?sig=secret") with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): diff --git a/tests/unit/registry/test_scenario_suite_storage.py b/tests/unit/registry/test_scenario_suite_storage.py index 8de826ac48..52f2585d76 100644 --- a/tests/unit/registry/test_scenario_suite_storage.py +++ b/tests/unit/registry/test_scenario_suite_storage.py @@ -15,6 +15,7 @@ from pyrit.models import ScenarioSuite, ScenarioSuiteEntry from pyrit.models.catalog.scenario_suite import SUITE_SCHEMA_VERSION from pyrit.registry import ScenarioSuiteConflictError, ScenarioSuiteStorage +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError ENTRY_ONE = UUID("11111111-1111-4111-8111-111111111111") ENTRY_TWO = UUID("22222222-2222-4222-8222-222222222222") @@ -95,7 +96,7 @@ def test_failed_atomic_replace_preserves_previous_suite(tmp_path: Path) -> None: created = storage.save_suite(suite=_suite(description="original"), expected_version=None) with patch("os.replace", side_effect=OSError(28, "No space left on device")): - with pytest.raises(OSError): + with pytest.raises(DocumentStorageUnavailableError): storage.save_suite(suite=_suite(description="replacement"), expected_version=created.version) loaded = storage.load_suite("weekly") @@ -112,7 +113,7 @@ def test_local_lock_excludes_a_second_writer(tmp_path: Path) -> None: real_replace = ScenarioSuiteStorage._replace_file def replace_after_competitor_times_out(*, path: Path, content: bytes) -> None: - with pytest.raises(TimeoutError): + with pytest.raises(DocumentStorageUnavailableError): competitor.save_suite(suite=_suite(description="racer"), expected_version=created.version) real_replace(path=path, content=content)