From 66fd155e705676037e31f0cfcfeaac8a80557978 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 19:29:26 +0000 Subject: [PATCH 01/18] FEAT: Report the documents a storage listing cannot use --- pyrit/registry/file_document_storage.py | 125 ++++++++++++++ .../registry/test_file_document_storage.py | 158 ++++++++++++++++++ 2 files changed, 283 insertions(+) create mode 100644 tests/unit/registry/test_file_document_storage.py diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index 56432d80b7..bd16601de0 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -12,6 +12,7 @@ from __future__ import annotations +import fnmatch import hashlib import logging import os @@ -19,6 +20,7 @@ import tempfile import time from contextlib import contextmanager, suppress +from dataclasses import dataclass, field from pathlib import Path, PurePosixPath from typing import TYPE_CHECKING from urllib.parse import unquote, urlparse @@ -33,6 +35,32 @@ logger = logging.getLogger(__name__) + +@dataclass(frozen=True) +class UnreadableDocument: + """A stored document with a valid name whose content could not be read.""" + + name: str + message: str + + +@dataclass(frozen=True) +class IgnoredDocument: + """A stored file whose name breaks the naming rule, so no single-document call can address it.""" + + file_name: str + message: str + + +@dataclass(frozen=True) +class DocumentListing: + """Every document a listing found, including the ones it could not use.""" + + documents: dict[str, bytes] = field(default_factory=dict) + unreadable: tuple[UnreadableDocument, ...] = () + ignored: tuple[IgnoredDocument, ...] = () + + # Both platforms lock a single byte at offset zero, which they allow past the end of an # empty file. The lock belongs to the open file handle, so it excludes other threads in # this process as well as other processes, and the kernel drops it if the holder exits. @@ -199,6 +227,103 @@ def _list_documents(self) -> dict[str, bytes]: """ return self._list_blob_documents() if self._is_blob else self._list_local_documents() + def _list_documents_with_issues(self) -> DocumentListing: + """ + Read every stored document and report the ones that can't be used. + + Unlike ``_list_documents``, nothing is skipped quietly: a document that can't be read + is reported as unreadable, and a file whose name breaks the naming rule is reported as + ignored without being read, because no single-document call could address it. A + storage location that can't be listed at all raises instead of looking empty. + + Messages carry only a short reason, never a path or URL, so they can be shown to users + without exposing server paths or a SAS token. + + Returns: + DocumentListing: Readable documents keyed by name, plus the unreadable and ignored ones. + """ + return self._list_blob_documents_with_issues() if self._is_blob else self._list_local_documents_with_issues() + + def _list_local_documents_with_issues(self) -> DocumentListing: + """ + Read the local directory, reporting unreadable and badly named documents. + + Returns: + DocumentListing: The documents and issues found in the directory. + """ + directory = self._local_directory(create=True) + # ``Path.glob`` returns nothing for a directory it can't read; listing it directly raises instead. + # ``fnmatch`` keeps glob's per-platform case rule, so both listings see the same files. + with os.scandir(directory) as entries: + file_names = sorted(entry.name for entry in entries if fnmatch.fnmatch(entry.name, f"*{self._extension}")) + documents: dict[str, bytes] = {} + unreadable: list[UnreadableDocument] = [] + ignored: list[IgnoredDocument] = [] + for file_name in file_names: + name = Path(file_name).stem + problem = self._name_problem(name) + if problem is not None: + ignored.append(IgnoredDocument(file_name=file_name, message=problem)) + continue + try: + documents[name] = (directory / file_name).read_bytes() + except FileNotFoundError: + # Deleted after the directory was listed, so it is no longer part of the listing. + continue + except OSError as error: + unreadable.append(UnreadableDocument(name=name, message=error.strerror or type(error).__name__)) + return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) + + def _list_blob_documents_with_issues(self) -> DocumentListing: + """ + Read the blob container, reporting unreadable and badly named documents. + + Returns: + DocumentListing: The documents and issues found under the configured prefix. + """ + from azure.core.exceptions import AzureError, ResourceNotFoundError + + documents: dict[str, bytes] = {} + unreadable: list[UnreadableDocument] = [] + ignored: list[IgnoredDocument] = [] + with self._open_container_client() as client: + prefix = f"{self._blob_prefix}/" if self._blob_prefix else None + blobs = client.list_blobs(name_starts_with=prefix) if prefix else client.list_blobs() + blob_names = sorted( + blob.name for blob in blobs if self._is_direct_document_blob(blob_name=blob.name, prefix=prefix) + ) + for blob_name in blob_names: + file_name = blob_name.removeprefix(prefix or "") + name = PurePosixPath(file_name).stem + problem = self._name_problem(name) + if problem is not None: + ignored.append(IgnoredDocument(file_name=file_name, message=problem)) + continue + try: + documents[name] = client.download_blob(blob_name).readall() + except ResourceNotFoundError: + # Deleted after the container was listed, so it is no longer part of the listing. + continue + except AzureError as error: + status = getattr(error, "status_code", None) + reason = type(error).__name__ if status is None else f"{type(error).__name__} ({status})" + unreadable.append(UnreadableDocument(name=name, message=reason)) + return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) + + @staticmethod + def _name_problem(name: str) -> str | None: + """ + Return why a discovered document name can't be addressed, or None if it can. + + Returns: + str | None: The naming-rule error, or ``None`` for a legal name. + """ + try: + validate_registry_name(name) + except ValueError as error: + return str(error) + return None + def _list_local_documents(self) -> dict[str, bytes]: """ Read addressable documents from the configured local directory. diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py new file mode 100644 index 0000000000..e6d67bec77 --- /dev/null +++ b/tests/unit/registry/test_file_document_storage.py @@ -0,0 +1,158 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from azure.core.exceptions import HttpResponseError, ResourceNotFoundError +from azure.storage.blob import ContainerClient + +from pyrit.registry.file_document_storage import DocumentListing, FileDocumentStorage, UnreadableDocument + +_SAS_SOURCE = "https://account.blob.core.windows.net/documents/saved?sp=rl&sig=secret-signature" + + +def _local_storage(directory: Path) -> FileDocumentStorage: + return FileDocumentStorage(source=str(directory), extension=".json", source_label="Test documents") + + +def _blob_storage() -> FileDocumentStorage: + return FileDocumentStorage(source=_SAS_SOURCE, extension=".json", source_label="Test documents") + + +def _blob_client(*, names: list[str]) -> MagicMock: + client = MagicMock(spec=ContainerClient) + client.__enter__.return_value = client + client.list_blobs.return_value = [SimpleNamespace(name=name) for name in names] + return client + + +class TestLocalListingWithIssues: + def test_reports_unreadable_document_and_keeps_the_others(self, tmp_path: Path) -> None: + (tmp_path / "first.json").write_bytes(b"one") + (tmp_path / "second.json").write_bytes(b"two") + original_read_bytes = Path.read_bytes + + def read_bytes(path: Path) -> bytes: + if path.name == "second.json": + raise PermissionError(13, "Permission denied", str(path)) + return original_read_bytes(path) + + with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + listing = _local_storage(tmp_path)._list_documents_with_issues() + legacy = _local_storage(tmp_path)._list_documents() + + assert listing == DocumentListing( + documents={"first": b"one"}, + unreadable=(UnreadableDocument(name="second", message="Permission denied"),), + ) + assert legacy == {"first": b"one"} + + def test_reports_badly_named_files_as_ignored_without_reading_them(self, tmp_path: Path) -> None: + (tmp_path / "valid_name.json").write_bytes(b"ok") + (tmp_path / "Bad-Name.json").write_bytes(b"never read") + (tmp_path / "notes.txt").write_bytes(b"other extension") + (tmp_path / "nested").mkdir() + (tmp_path / "nested" / "inner.json").write_bytes(b"not a direct child") + read_names: list[str] = [] + original_read_bytes = Path.read_bytes + + def read_bytes(path: Path) -> bytes: + read_names.append(path.name) + return original_read_bytes(path) + + with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing.documents == {"valid_name": b"ok"} + assert listing.unreadable == () + assert [issue.file_name for issue in listing.ignored] == ["Bad-Name.json"] + assert "Bad-Name" in listing.ignored[0].message + assert read_names == ["valid_name.json"] + assert _local_storage(tmp_path)._list_documents() == listing.documents + + def test_skips_a_document_deleted_after_listing(self, tmp_path: Path) -> None: + (tmp_path / "kept.json").write_bytes(b"kept") + (tmp_path / "deleted.json").write_bytes(b"gone") + original_read_bytes = Path.read_bytes + + def read_bytes(path: Path) -> bytes: + if path.name == "deleted.json": + raise FileNotFoundError(2, "No such file or directory", str(path)) + return original_read_bytes(path) + + with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing(documents={"kept": b"kept"}) + + def test_creates_missing_directory_and_lists_nothing(self, tmp_path: Path) -> None: + listing = _local_storage(tmp_path / "missing")._list_documents_with_issues() + + assert listing == DocumentListing() + assert (tmp_path / "missing").is_dir() + + def test_raises_when_the_directory_cannot_be_listed(self, tmp_path: Path) -> None: + (tmp_path / "first.json").write_bytes(b"one") + + with patch("pyrit.registry.file_document_storage.os.scandir", side_effect=PermissionError(13, "denied")): + with pytest.raises(PermissionError): + _local_storage(tmp_path)._list_documents_with_issues() + + +class TestBlobListingWithIssues: + def test_reports_unreadable_and_ignored_blobs_without_exposing_the_url(self) -> None: + client = _blob_client( + names=[ + "saved/good.json", + "saved/broken.json", + "saved/Bad-Name.json", + "saved/notes.txt", + "saved/deeper/inner.json", + ] + ) + + def download_blob(blob_name: str) -> SimpleNamespace: + if blob_name == "saved/broken.json": + error = HttpResponseError(message=f"Forbidden: {_SAS_SOURCE}") + error.status_code = 403 + raise error + return SimpleNamespace(readall=lambda: b"content") + + client.download_blob.side_effect = download_blob + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + listing = _blob_storage()._list_documents_with_issues() + + assert listing.documents == {"good": b"content"} + assert listing.unreadable == (UnreadableDocument(name="broken", message="HttpResponseError (403)"),) + assert [issue.file_name for issue in listing.ignored] == ["Bad-Name.json"] + assert all("secret-signature" not in issue.message for issue in (*listing.unreadable, *listing.ignored)) + client.list_blobs.assert_called_once_with(name_starts_with="saved/") + downloaded = [call.args[0] for call in client.download_blob.call_args_list] + assert downloaded == ["saved/broken.json", "saved/good.json"] + + def test_skips_a_blob_deleted_after_listing(self) -> None: + client = _blob_client(names=["saved/kept.json", "saved/deleted.json"]) + + def download_blob(blob_name: str) -> SimpleNamespace: + if blob_name == "saved/deleted.json": + raise ResourceNotFoundError(message="The specified blob does not exist.") + return SimpleNamespace(readall=lambda: b"kept") + + client.download_blob.side_effect = download_blob + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + listing = _blob_storage()._list_documents_with_issues() + + assert listing == DocumentListing(documents={"kept": b"kept"}) + + def test_raises_when_the_container_cannot_be_listed(self) -> None: + client = _blob_client(names=[]) + client.list_blobs.side_effect = HttpResponseError(message="listing failed") + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + with pytest.raises(HttpResponseError): + _blob_storage()._list_documents_with_issues() From 2d8c465c9dfefb8efa0b1e47cf0d1ca1d8199927 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 20:12:11 +0000 Subject: [PATCH 02/18] Report non-files and missing containers instead of skipping them --- pyrit/registry/file_document_storage.py | 58 +++++++++---- .../registry/test_file_document_storage.py | 84 +++++++++++++++++-- 2 files changed, 116 insertions(+), 26 deletions(-) diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index bd16601de0..17f2219ebb 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -253,25 +253,30 @@ def _list_local_documents_with_issues(self) -> DocumentListing: """ directory = self._local_directory(create=True) # ``Path.glob`` returns nothing for a directory it can't read; listing it directly raises instead. - # ``fnmatch`` keeps glob's per-platform case rule, so both listings see the same files. + # ``fnmatch`` applies the platform's case rule, as ``glob`` does. with os.scandir(directory) as entries: - file_names = sorted(entry.name for entry in entries if fnmatch.fnmatch(entry.name, f"*{self._extension}")) + candidates = sorted( + (entry.name, entry.is_file()) for entry in entries if fnmatch.fnmatch(entry.name, f"*{self._extension}") + ) documents: dict[str, bytes] = {} unreadable: list[UnreadableDocument] = [] ignored: list[IgnoredDocument] = [] - for file_name in file_names: + for file_name, is_file in candidates: name = Path(file_name).stem problem = self._name_problem(name) if problem is not None: ignored.append(IgnoredDocument(file_name=file_name, message=problem)) - continue - try: - documents[name] = (directory / file_name).read_bytes() - except FileNotFoundError: - # Deleted after the directory was listed, so it is no longer part of the listing. - continue - except OSError as error: - unreadable.append(UnreadableDocument(name=name, message=error.strerror or type(error).__name__)) + elif not is_file: + # A broken link, a folder, or a pipe, which would block if read. + unreadable.append(UnreadableDocument(name=name, message="Not a regular file")) + else: + try: + documents[name] = (directory / file_name).read_bytes() + except FileNotFoundError: + # Deleted after the directory was listed, so it is no longer part of the listing. + continue + except OSError as error: + unreadable.append(UnreadableDocument(name=name, message=self._describe_read_error(error))) return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) def _list_blob_documents_with_issues(self) -> DocumentListing: @@ -281,7 +286,8 @@ def _list_blob_documents_with_issues(self) -> DocumentListing: Returns: DocumentListing: The documents and issues found under the configured prefix. """ - from azure.core.exceptions import AzureError, ResourceNotFoundError + from azure.core.exceptions import AzureError + from azure.storage.blob import StorageErrorCode documents: dict[str, bytes] = {} unreadable: list[UnreadableDocument] = [] @@ -301,15 +307,31 @@ def _list_blob_documents_with_issues(self) -> DocumentListing: continue try: documents[name] = client.download_blob(blob_name).readall() - except ResourceNotFoundError: - # Deleted after the container was listed, so it is no longer part of the listing. - continue except AzureError as error: - status = getattr(error, "status_code", None) - reason = type(error).__name__ if status is None else f"{type(error).__name__} ({status})" - unreadable.append(UnreadableDocument(name=name, message=reason)) + if getattr(error, "error_code", None) == StorageErrorCode.BLOB_NOT_FOUND: + # Deleted after the container was listed, so it is no longer part of the listing. + continue + unreadable.append(UnreadableDocument(name=name, message=self._describe_read_error(error))) return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) + @staticmethod + def _describe_read_error(error: Exception) -> str: + """ + Describe a storage error in a few words, leaving out any path or URL it carries. + + Returns: + str: For example ``Permission denied`` or ``HttpResponseError (403, AuthorizationPermissionMismatch)``. + """ + if isinstance(error, OSError) and error.strerror: + return error.strerror + error_code = getattr(error, "error_code", None) + details = [ + str(detail) + for detail in (getattr(error, "status_code", None), getattr(error_code, "value", error_code)) + if detail is not None + ] + return f"{type(error).__name__} ({', '.join(details)})" if details else type(error).__name__ + @staticmethod def _name_problem(name: str) -> str | None: """ diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index e6d67bec77..face54fe80 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -1,13 +1,14 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from collections.abc import Iterator from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest -from azure.core.exceptions import HttpResponseError, ResourceNotFoundError -from azure.storage.blob import ContainerClient +from azure.core.exceptions import HttpResponseError, ResourceNotFoundError, ServiceRequestError +from azure.storage.blob import ContainerClient, StorageErrorCode from pyrit.registry.file_document_storage import DocumentListing, FileDocumentStorage, UnreadableDocument @@ -88,6 +89,30 @@ def read_bytes(path: Path) -> bytes: assert listing == DocumentListing(documents={"kept": b"kept"}) + def test_reports_entries_that_are_not_regular_files_without_reading_them(self, tmp_path: Path) -> None: + (tmp_path / "folder.json").mkdir() + try: + (tmp_path / "broken_link.json").symlink_to(tmp_path / "missing_target.json") + except OSError: + pytest.skip("Creating symbolic links isn't permitted here") + + with patch.object(Path, "read_bytes", autospec=True) as read_bytes: + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing.unreadable == ( + UnreadableDocument(name="broken_link", message="Not a regular file"), + UnreadableDocument(name="folder", message="Not a regular file"), + ) + read_bytes.assert_not_called() + + def test_matches_extensions_like_the_legacy_listing(self, tmp_path: Path) -> None: + (tmp_path / "lower.json").write_bytes(b"one") + (tmp_path / "upper.JSON").write_bytes(b"two") + + storage = _local_storage(tmp_path) + + assert storage._list_documents_with_issues().documents == storage._list_documents() + def test_creates_missing_directory_and_lists_nothing(self, tmp_path: Path) -> None: listing = _local_storage(tmp_path / "missing")._list_documents_with_issues() @@ -134,12 +159,19 @@ def download_blob(blob_name: str) -> SimpleNamespace: downloaded = [call.args[0] for call in client.download_blob.call_args_list] assert downloaded == ["saved/broken.json", "saved/good.json"] - def test_skips_a_blob_deleted_after_listing(self) -> None: - client = _blob_client(names=["saved/kept.json", "saved/deleted.json"]) + def test_skips_only_a_blob_deleted_after_listing(self) -> None: + client = _blob_client(names=["saved/kept.json", "saved/deleted.json", "saved/container_gone.json"]) + error_codes = { + "saved/deleted.json": StorageErrorCode.BLOB_NOT_FOUND, + "saved/container_gone.json": StorageErrorCode.CONTAINER_NOT_FOUND, + } def download_blob(blob_name: str) -> SimpleNamespace: - if blob_name == "saved/deleted.json": - raise ResourceNotFoundError(message="The specified blob does not exist.") + if blob_name in error_codes: + error = ResourceNotFoundError(message="Not found.") + error.status_code = 404 + error.error_code = error_codes[blob_name] + raise error return SimpleNamespace(readall=lambda: b"kept") client.download_blob.side_effect = download_blob @@ -147,12 +179,48 @@ def download_blob(blob_name: str) -> SimpleNamespace: with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): listing = _blob_storage()._list_documents_with_issues() - assert listing == DocumentListing(documents={"kept": b"kept"}) + assert listing == DocumentListing( + documents={"kept": b"kept"}, + unreadable=( + UnreadableDocument(name="container_gone", message="ResourceNotFoundError (404, ContainerNotFound)"), + ), + ) def test_raises_when_the_container_cannot_be_listed(self) -> None: + def failing_pages() -> Iterator[SimpleNamespace]: + raise HttpResponseError(message="listing failed") + yield SimpleNamespace(name="saved/never.json") + client = _blob_client(names=[]) - client.list_blobs.side_effect = HttpResponseError(message="listing failed") + client.list_blobs.return_value = failing_pages() with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): with pytest.raises(HttpResponseError): _blob_storage()._list_documents_with_issues() + + +class TestDescribeReadError: + @pytest.mark.parametrize( + ("error", "expected"), + [ + (PermissionError(13, "Permission denied", "/server/path/secret.json"), "Permission denied"), + (OSError(), "OSError"), + (ServiceRequestError(message="https://account.blob.core.windows.net/x?sig=secret"), "ServiceRequestError"), + ], + ) + def test_describes_errors_without_their_details(self, error: Exception, expected: str) -> None: + assert FileDocumentStorage._describe_read_error(error) == expected + + @pytest.mark.parametrize( + ("error_code", "expected"), + [ + (StorageErrorCode.AUTHORIZATION_FAILURE, "HttpResponseError (403, AuthorizationFailure)"), + ("AuthorizationPermissionMismatch", "HttpResponseError (403, AuthorizationPermissionMismatch)"), + ], + ) + def test_includes_the_azure_status_and_error_code(self, error_code: str, expected: str) -> None: + error = HttpResponseError(message="https://account.blob.core.windows.net/x?sig=secret") + error.status_code = 403 + error.error_code = error_code + + assert FileDocumentStorage._describe_read_error(error) == expected From 5dec828431f959e3e1f68852f6b68680130954ea Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 21:10:09 +0000 Subject: [PATCH 03/18] Report an entry that can't be checked instead of failing the whole listing --- pyrit/registry/file_document_storage.py | 44 ++++++++++++------- .../registry/test_file_document_storage.py | 12 +++++ 2 files changed, 41 insertions(+), 15 deletions(-) diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index 17f2219ebb..4ba3f0949d 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -256,29 +256,43 @@ def _list_local_documents_with_issues(self) -> DocumentListing: # ``fnmatch`` applies the platform's case rule, as ``glob`` does. with os.scandir(directory) as entries: candidates = sorted( - (entry.name, entry.is_file()) for entry in entries if fnmatch.fnmatch(entry.name, f"*{self._extension}") + (entry for entry in entries if fnmatch.fnmatch(entry.name, f"*{self._extension}")), + key=lambda entry: entry.name, ) documents: dict[str, bytes] = {} unreadable: list[UnreadableDocument] = [] ignored: list[IgnoredDocument] = [] - for file_name, is_file in candidates: - name = Path(file_name).stem + for entry in candidates: + name = Path(entry.name).stem problem = self._name_problem(name) if problem is not None: - ignored.append(IgnoredDocument(file_name=file_name, message=problem)) - elif not is_file: - # A broken link, a folder, or a pipe, which would block if read. - unreadable.append(UnreadableDocument(name=name, message="Not a regular file")) - else: - try: - documents[name] = (directory / file_name).read_bytes() - except FileNotFoundError: - # Deleted after the directory was listed, so it is no longer part of the listing. - continue - except OSError as error: - unreadable.append(UnreadableDocument(name=name, message=self._describe_read_error(error))) + ignored.append(IgnoredDocument(file_name=entry.name, message=problem)) + continue + content = self._read_local_entry(entry=entry, name=name) + if isinstance(content, bytes): + documents[name] = content + elif content is not None: + unreadable.append(content) return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) + def _read_local_entry(self, *, entry: os.DirEntry[str], name: str) -> bytes | UnreadableDocument | None: + """ + Read one listed local document. + + Returns: + bytes | UnreadableDocument | None: The content, why it can't be read, or None if it was + deleted after the directory was listed. + """ + try: + if not entry.is_file(): + # A broken link, a folder, or a pipe, which would block if read. + return UnreadableDocument(name=name, message="Not a regular file") + return Path(entry.path).read_bytes() + except FileNotFoundError: + return None + except OSError as error: + return UnreadableDocument(name=name, message=self._describe_read_error(error)) + def _list_blob_documents_with_issues(self) -> DocumentListing: """ Read the blob container, reporting unreadable and badly named documents. diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index face54fe80..63bc7601dc 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -105,6 +105,18 @@ def test_reports_entries_that_are_not_regular_files_without_reading_them(self, t ) read_bytes.assert_not_called() + def test_reports_a_link_loop_and_keeps_the_others(self, tmp_path: Path) -> None: + (tmp_path / "good.json").write_bytes(b"ok") + try: + (tmp_path / "loop.json").symlink_to(tmp_path / "loop.json") + except OSError: + pytest.skip("Creating symbolic links isn't permitted here") + + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing.documents == {"good": b"ok"} + assert [issue.name for issue in listing.unreadable] == ["loop"] + def test_matches_extensions_like_the_legacy_listing(self, tmp_path: Path) -> None: (tmp_path / "lower.json").write_bytes(b"one") (tmp_path / "upper.JSON").write_bytes(b"two") From bf662b71e028aabb159cb7fa47051d3b0d1283db Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 21:50:58 +0000 Subject: [PATCH 04/18] FEAT: Read saved registry recipes --- pyrit/models/catalog/registry_recipe.py | 143 ++++++++++++++ pyrit/registry/registry_recipe_storage.py | 182 ++++++++++++++++++ .../registry/test_registry_recipe_storage.py | 150 +++++++++++++++ 3 files changed, 475 insertions(+) create mode 100644 pyrit/models/catalog/registry_recipe.py create mode 100644 pyrit/registry/registry_recipe_storage.py create mode 100644 tests/unit/registry/test_registry_recipe_storage.py diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py new file mode 100644 index 0000000000..ce2e544588 --- /dev/null +++ b/pyrit/models/catalog/registry_recipe.py @@ -0,0 +1,143 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Saved registry entries: the recipe kept for each one and what restoring them found.""" + +from __future__ import annotations + +from enum import Enum +from typing import Annotated, ClassVar + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator + +from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS + +_Identifier = Annotated[str, Field(max_length=MAX_IDENTIFIER_LENGTH)] +_TechniqueTag = Annotated[str, Field(pattern=r"^[A-Za-z][A-Za-z0-9_]{0,63}$")] + + +class RegistryRecipeKind(str, Enum): + """The registries whose entries can be saved, named after their storage folders.""" + + TARGETS = "targets" + CONVERTERS = "converters" + SCORERS = "scorers" + TECHNIQUES = "techniques" + + +class _RegistryRecipe(BaseModel): + """Fields every saved recipe shares. The entry's name is its document name, so it isn't repeated.""" + + CURRENT_SCHEMA_VERSION: ClassVar[int] = 1 + + model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False, frozen=True) + + schema_version: int = Field( + CURRENT_SCHEMA_VERSION, ge=1, le=CURRENT_SCHEMA_VERSION, description="Shape of this document" + ) + + +class ComponentRecipe(_RegistryRecipe): + """How to rebuild one saved target, converter, or scorer.""" + + class_module: _Identifier = Field( + ..., min_length=1, description="Module that defined the class when the entry was saved" + ) + class_name: _Identifier = Field(..., min_length=1, description="Registered class name") + params: dict[_Identifier, JsonValue] = Field( + default_factory=dict, + max_length=MAX_ITEMS, + description="Constructor parameters as the API accepted them; other entries are named, never embedded", + ) + + +class TechniqueRecipe(_RegistryRecipe): + """ + How to rebuild one saved technique: the fields of its create request, without the name. + + The fields and their limits mirror the technique API's create request, which this library + can't import. + """ + + type: str = Field(..., min_length=1, max_length=256, description="Registered attack type") + params: dict[_Identifier, JsonValue] = Field( + default_factory=dict, max_length=MAX_ITEMS, description="Attack constructor parameters" + ) + description: str | None = None + tags: list[_TechniqueTag] = Field(default_factory=list, max_length=MAX_ITEMS) + request_converters: list[_Identifier] | None = Field(default=None, max_length=MAX_ITEMS) + response_converters: list[_Identifier] | None = Field(default=None, max_length=MAX_ITEMS) + adversarial_chat: _Identifier | None = None + adversarial_system_prompt: str | None = None + adversarial_seed_prompt: str | None = None + adversarial_prompt_template: str | None = None + + @field_validator("tags") + @classmethod + def _validate_tags(cls, value: list[str]) -> list[str]: + folded = {tag.casefold() for tag in value} + if len(folded) != len(value): + raise ValueError("Tags must be unique, including letter case") + if {"all", "default"} & folded: + raise ValueError("all and default are reserved selectors") + return value + + +class SavedEntryRef(BaseModel): + """Points at one saved entry.""" + + model_config = ConfigDict(frozen=True) + + kind: RegistryRecipeKind + name: str + + +class StoredRecipe(BaseModel): + """A saved recipe and the version of the document it was read from.""" + + model_config = ConfigDict(frozen=True) + + ref: SavedEntryRef + recipe: ComponentRecipe | TechniqueRecipe + version: str = Field(..., description="Opaque version of the document this recipe was read from") + + +class RestoreFailureReason(str, Enum): + """Why a saved entry couldn't be restored.""" + + UNREADABLE_STORAGE = "unreadable_storage" + UNREADABLE_RECIPE = "unreadable_recipe" + INVALID_RECIPE = "invalid_recipe" + NEWER_SCHEMA_VERSION = "newer_schema_version" + UNKNOWN_CLASS = "unknown_class" + MODULE_MISMATCH = "module_mismatch" + UNSUPPORTED_CLASS = "unsupported_class" + UNSUPPORTED_CREDENTIAL = "unsupported_credential" + NAME_CLASH = "name_clash" + MISSING_DEPENDENCY = "missing_dependency" + FAILED_DEPENDENCY = "failed_dependency" + DEPENDENCY_LOOP = "dependency_loop" + BUILD_ERROR = "build_error" + + +class RestoreFailure(BaseModel): + """One saved entry, or one whole storage folder, that couldn't be restored.""" + + model_config = ConfigDict(frozen=True) + + kind: RegistryRecipeKind + name: str | None = Field(None, description="Entry name; None when the whole folder couldn't be read") + reason: RestoreFailureReason + message: str = Field(..., description="Why, in words safe to show users") + version: str | None = Field(None, description="Version of the stored document, when it could be read") + dependency: SavedEntryRef | None = Field(None, description="The entry this one needed that is missing or failed") + + +class IgnoredRecipeFile(BaseModel): + """A stored file whose name can't be an entry name, so it was never read.""" + + model_config = ConfigDict(frozen=True) + + kind: RegistryRecipeKind + file_name: str + message: str diff --git a/pyrit/registry/registry_recipe_storage.py b/pyrit/registry/registry_recipe_storage.py new file mode 100644 index 0000000000..0a956a19b8 --- /dev/null +++ b/pyrit/registry/registry_recipe_storage.py @@ -0,0 +1,182 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Storage for saved registry recipes, one folder per registry.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path +from urllib.parse import urlparse + +from pydantic import ValidationError + +from pyrit.common.azure_storage import is_azure_blob_uri +from pyrit.models.catalog.registry_recipe import ( + ComponentRecipe, + IgnoredRecipeFile, + RegistryRecipeKind, + RestoreFailure, + RestoreFailureReason, + SavedEntryRef, + StoredRecipe, + TechniqueRecipe, +) +from pyrit.registry.file_document_storage import FileDocumentStorage + + +@dataclass(frozen=True) +class RecipeListing: + """Every saved recipe of one kind, with the ones that couldn't be read or parsed.""" + + recipes: dict[str, StoredRecipe] = field(default_factory=dict) + failures: tuple[RestoreFailure, ...] = () + ignored: tuple[IgnoredRecipeFile, ...] = () + + +class RegistryRecipeStorage(FileDocumentStorage): + """ + Read the saved recipes of one registry as JSON documents. + + Each registry has its own folder under the configured source (``targets``, ``converters``, + ``scorers``, ``techniques``), so a target and a converter can share a name. Storage is read + through on every call, like scenario presets, because the folder can be shared between + processes. + """ + + def __init__(self, *, kind: RegistryRecipeKind, source: str | None = None) -> None: + """ + Initialize storage for one registry's folder. + + Args: + kind (RegistryRecipeKind): The registry whose recipes this storage reads. + source (str | None): Local directory or Azure Blob URI that holds the per-registry + folders. Defaults to ``registry_recipes`` under the PyRIT configuration directory. + + Raises: + ValueError: If the source has an unsupported URI scheme. + """ + root = source or str(self._get_default_storage_dir()) + super().__init__( + source=self._get_kind_source(root=root, kind=kind), + extension=".json", + source_label=f"Saved {kind.value}", + ) + self._kind = kind + + @staticmethod + def _get_default_storage_dir() -> Path: + """ + Get the default directory that holds the per-registry folders. + + Returns: + Path: Path to ``~/.pyrit/registry_recipes/``. + """ + # Deferred: importing pyrit.common.path triggers pyrit __init__.py + from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH + + return CONFIGURATION_DIRECTORY_PATH / "registry_recipes" + + @staticmethod + def _get_kind_source(*, root: str, kind: RegistryRecipeKind) -> str: + """ + Append the registry's folder to the configured source. + + For Azure Blob the folder goes on the URL path, so a SAS token in the query string is kept. + + Returns: + str: The local directory or Blob URI for the registry's folder. + """ + if is_azure_blob_uri(root): + parsed = urlparse(root) + return parsed._replace(path=f"{parsed.path.rstrip('/')}/{kind.value}").geturl() + return str(Path(root) / kind.value) + + def list_recipes(self) -> RecipeListing: + """ + Read every saved recipe of this kind and report the ones that can't be used. + + Returns: + RecipeListing: Parsed recipes keyed by name, plus the failures and ignored files. + + Raises: + OSError: If the local folder can't be listed. + AzureError: If the Blob container can't be listed. + """ + listing = self._list_documents_with_issues() + recipes: dict[str, StoredRecipe] = {} + failures = [ + self._failure(name=issue.name, reason=RestoreFailureReason.UNREADABLE_RECIPE, message=issue.message) + for issue in listing.unreadable + ] + for name, content in listing.documents.items(): + parsed = self._parse_recipe(name=name, content=content) + if isinstance(parsed, RestoreFailure): + failures.append(parsed) + else: + recipes[name] = parsed + ignored = tuple( + IgnoredRecipeFile(kind=self._kind, file_name=issue.file_name, message=issue.message) + for issue in listing.ignored + ) + return RecipeListing(recipes=recipes, failures=tuple(failures), ignored=ignored) + + def _parse_recipe(self, *, name: str, content: bytes) -> StoredRecipe | RestoreFailure: + """ + Parse one stored document, checking its schema version before anything else. + + Returns: + StoredRecipe | RestoreFailure: The recipe, or why it can't be used. + """ + version = self._compute_version(content) + invalid = RestoreFailureReason.INVALID_RECIPE + try: + payload = json.loads(content.decode("utf-8")) + except ValueError: + return self._failure(name=name, version=version, reason=invalid, message="Not valid UTF-8 JSON.") + except RecursionError: + return self._failure(name=name, version=version, reason=invalid, message="Nested too deeply to read.") + if not isinstance(payload, dict): + return self._failure(name=name, version=version, reason=invalid, message="Not a JSON object.") + schema_version = payload.get("schema_version") + current = ComponentRecipe.CURRENT_SCHEMA_VERSION + if isinstance(schema_version, int) and schema_version > current: + message = f"Saved by a newer PyRIT (format version {schema_version}; this one reads up to {current})." + return self._failure( + name=name, version=version, reason=RestoreFailureReason.NEWER_SCHEMA_VERSION, message=message + ) + model = TechniqueRecipe if self._kind is RegistryRecipeKind.TECHNIQUES else ComponentRecipe + try: + recipe = model.model_validate(payload) + except ValidationError as error: + return self._failure( + name=name, version=version, reason=invalid, message=self._describe_validation_error(error) + ) + except RecursionError: + return self._failure(name=name, version=version, reason=invalid, message="Nested too deeply to read.") + return StoredRecipe(ref=SavedEntryRef(kind=self._kind, name=name), recipe=recipe, version=version) + + def _failure( + self, *, name: str, reason: RestoreFailureReason, message: str, version: str | None = None + ) -> RestoreFailure: + """ + Build a failure for one entry in this registry's folder. + + Returns: + RestoreFailure: The failure. + """ + return RestoreFailure(kind=self._kind, name=name, reason=reason, message=message, version=version) + + @staticmethod + def _describe_validation_error(error: ValidationError) -> str: + """ + Summarize a validation error without echoing the stored values. + + Returns: + str: One line per problem, naming the field and what's wrong with it. + """ + problems = [ + f"{'.'.join(str(part) for part in issue['loc']) or 'document'}: {issue['msg']}" for issue in error.errors() + ] + return "Not a valid saved recipe: " + "; ".join(problems) diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py new file mode 100644 index 0000000000..58851e2b45 --- /dev/null +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -0,0 +1,150 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from azure.storage.blob import ContainerClient + +from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH +from pyrit.models.catalog.registry_recipe import ( + ComponentRecipe, + RegistryRecipeKind, + RestoreFailureReason, + TechniqueRecipe, +) +from pyrit.registry.registry_recipe_storage import RegistryRecipeStorage + +_TARGET = { + "schema_version": 1, + "class_module": "pyrit.prompt_target.openai.openai_chat_target", + "class_name": "OpenAIChatTarget", + "params": {"auth_mode": "identity", "endpoint": "https://contoso.openai.azure.com/openai/v1"}, +} + + +def _write(directory: Path, name: str, content: str | bytes) -> bytes: + directory.mkdir(parents=True, exist_ok=True) + data = content.encode("utf-8") if isinstance(content, str) else content + (directory / f"{name}.json").write_bytes(data) + return data + + +class TestRegistryRecipeStorageSource: + def test_local_source_uses_a_folder_per_registry(self, tmp_path: Path) -> None: + storage = RegistryRecipeStorage(kind=RegistryRecipeKind.CONVERTERS, source=str(tmp_path)) + + assert storage.display_source == str(tmp_path / "converters") + + def test_default_source_is_under_the_pyrit_configuration_directory(self) -> None: + storage = RegistryRecipeStorage(kind=RegistryRecipeKind.SCORERS) + + assert storage.display_source == str(CONFIGURATION_DIRECTORY_PATH / "registry_recipes" / "scorers") + + def test_blob_source_adds_the_folder_and_keeps_the_sas_token(self) -> None: + source = "https://account.blob.core.windows.net/recipes/team?sp=rl&sig=secret" + + kind_source = RegistryRecipeStorage._get_kind_source(root=source, kind=RegistryRecipeKind.TARGETS) + + assert kind_source == "https://account.blob.core.windows.net/recipes/team/targets?sp=rl&sig=secret" + + def test_blob_container_source_uses_the_folder_as_prefix(self) -> None: + client = MagicMock(spec=ContainerClient) + client.__enter__.return_value = client + client.list_blobs.return_value = [SimpleNamespace(name="techniques/formal.json")] + client.download_blob.return_value.readall.return_value = json.dumps({"type": "RedTeamingAttack"}).encode() + storage = RegistryRecipeStorage( + kind=RegistryRecipeKind.TECHNIQUES, source="https://account.blob.core.windows.net/recipes?sig=secret" + ) + + with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): + listing = storage.list_recipes() + + client.list_blobs.assert_called_once_with(name_starts_with="techniques/") + assert list(listing.recipes) == ["formal"] + + +class TestRegistryRecipeStorageListing: + def test_reads_component_recipes_with_their_versions(self, tmp_path: Path) -> None: + content = _write(tmp_path / "targets", "contoso_gpt4o", json.dumps(_TARGET)) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + stored = listing.recipes["contoso_gpt4o"] + assert stored.ref.kind is RegistryRecipeKind.TARGETS + assert stored.ref.name == "contoso_gpt4o" + assert stored.recipe == ComponentRecipe.model_validate(_TARGET) + assert stored.version == hashlib.sha256(content).hexdigest() + assert listing.failures == () + + def test_reads_technique_recipes_as_technique_recipes(self, tmp_path: Path) -> None: + document = {"type": "RedTeamingAttack", "tags": ["custom"], "adversarial_chat": "contoso_gpt4o"} + _write(tmp_path / "techniques", "formal_red_team", json.dumps(document)) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TECHNIQUES, source=str(tmp_path)).list_recipes() + + assert listing.recipes["formal_red_team"].recipe == TechniqueRecipe.model_validate(document) + + def test_reports_documents_that_are_not_recipes(self, tmp_path: Path) -> None: + folder = tmp_path / "targets" + _write(folder, "not_json", "{oops") + _write(folder, "not_utf8", b"\xff\xfe") + _write(folder, "not_object", "[1, 2]") + _write(folder, "extra_field", json.dumps({**_TARGET, "credentials": {"api_key": "sk-secret-value"}})) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + assert listing.recipes == {} + assert {failure.name for failure in listing.failures} == {"not_json", "not_utf8", "not_object", "extra_field"} + assert all(failure.reason is RestoreFailureReason.INVALID_RECIPE for failure in listing.failures) + assert all(failure.version for failure in listing.failures) + extra = next(failure for failure in listing.failures if failure.name == "extra_field") + assert "credentials" in extra.message + assert "sk-secret-value" not in extra.message + + def test_reports_a_document_nested_too_deeply(self, tmp_path: Path) -> None: + _write(tmp_path / "targets", "deep", "[" * 100_000 + "]" * 100_000) + _write(tmp_path / "targets", "contoso_gpt4o", json.dumps(_TARGET)) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + (failure,) = listing.failures + assert (failure.name, failure.reason) == ("deep", RestoreFailureReason.INVALID_RECIPE) + assert list(listing.recipes) == ["contoso_gpt4o"] + + def test_reports_a_newer_schema_version_before_checking_fields(self, tmp_path: Path) -> None: + _write(tmp_path / "targets", "future", json.dumps({"schema_version": 2, "unknown_field": True})) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + (failure,) = listing.failures + assert failure.reason is RestoreFailureReason.NEWER_SCHEMA_VERSION + assert "newer PyRIT" in failure.message + + def test_reports_unreadable_and_ignored_files(self, tmp_path: Path) -> None: + folder = tmp_path / "scorers" + _write(folder, "locked", json.dumps(_TARGET)) + _write(folder, "Bad-Name", "{}") + original_read_bytes = Path.read_bytes + + def read_bytes(path: Path) -> bytes: + if path.name == "locked.json": + raise PermissionError(13, "Permission denied", str(path)) + return original_read_bytes(path) + + with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.SCORERS, source=str(tmp_path)).list_recipes() + + (failure,) = listing.failures + assert (failure.name, failure.reason, failure.message) == ( + "locked", + RestoreFailureReason.UNREADABLE_RECIPE, + "Permission denied", + ) + assert failure.version is None + (ignored,) = listing.ignored + assert (ignored.kind, ignored.file_name) == (RegistryRecipeKind.SCORERS, "Bad-Name.json") + assert "Bad-Name" in ignored.message From cfa4fc03ab6acfd6ef5c2d90c304ce4ebb05345d Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 2 Nov 2026 23:41:45 +0000 Subject: [PATCH 05/18] Limit technique descriptions like the create request and export the recipe models --- pyrit/models/catalog/__init__.py | 18 ++++++++++++++++++ pyrit/models/catalog/registry_recipe.py | 4 ++-- pyrit/models/request_limits.py | 1 + .../registry/test_registry_recipe_storage.py | 10 ++++++++++ 4 files changed, 31 insertions(+), 2 deletions(-) diff --git a/pyrit/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index 102788e2a8..94c1abb948 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -20,6 +20,16 @@ if TYPE_CHECKING: from pyrit.models.catalog.initializer import RegisteredInitializer + from pyrit.models.catalog.registry_recipe import ( + ComponentRecipe, + IgnoredRecipeFile, + RegistryRecipeKind, + RestoreFailure, + RestoreFailureReason, + SavedEntryRef, + StoredRecipe, + TechniqueRecipe, + ) from pyrit.models.catalog.scenario import ( AttackErrorSummary, AttackRetrySummary, @@ -60,9 +70,15 @@ _LAZY_EXPORTS: dict[str, str] = { "AttackErrorSummary": "pyrit.models.catalog.scenario", "AttackRetrySummary": "pyrit.models.catalog.scenario", + "ComponentRecipe": "pyrit.models.catalog.registry_recipe", + "IgnoredRecipeFile": "pyrit.models.catalog.registry_recipe", "PresetEntry": "pyrit.models.catalog.scenario_preset", "PresetIssue": "pyrit.models.catalog.scenario_preset", "RegisteredInitializer": "pyrit.models.catalog.initializer", + "RegistryRecipeKind": "pyrit.models.catalog.registry_recipe", + "RestoreFailure": "pyrit.models.catalog.registry_recipe", + "RestoreFailureReason": "pyrit.models.catalog.registry_recipe", + "SavedEntryRef": "pyrit.models.catalog.registry_recipe", "ScorerInstance": "pyrit.models.catalog.scorer", "RegisteredScenario": "pyrit.models.catalog.scenario", "RunScenarioRequest": "pyrit.models.catalog.scenario", @@ -87,8 +103,10 @@ "ScenarioRunSummary": "pyrit.models.catalog.scenario", "ScenarioTechniqueSummary": "pyrit.models.catalog.scenario", "StoredPreset": "pyrit.models.catalog.scenario_preset", + "StoredRecipe": "pyrit.models.catalog.registry_recipe", "StoredScenarioSuite": "pyrit.models.catalog.scenario_suite", "TargetInstance": "pyrit.models.catalog.target", + "TechniqueRecipe": "pyrit.models.catalog.registry_recipe", } __all__ = list(_LAZY_EXPORTS) diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index ce2e544588..5ac95be552 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -10,7 +10,7 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator -from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS +from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS, MAX_TEXT_LENGTH _Identifier = Annotated[str, Field(max_length=MAX_IDENTIFIER_LENGTH)] _TechniqueTag = Annotated[str, Field(pattern=r"^[A-Za-z][A-Za-z0-9_]{0,63}$")] @@ -63,7 +63,7 @@ class TechniqueRecipe(_RegistryRecipe): params: dict[_Identifier, JsonValue] = Field( default_factory=dict, max_length=MAX_ITEMS, description="Attack constructor parameters" ) - description: str | None = None + description: str | None = Field(default=None, max_length=MAX_TEXT_LENGTH) tags: list[_TechniqueTag] = Field(default_factory=list, max_length=MAX_ITEMS) request_converters: list[_Identifier] | None = Field(default=None, max_length=MAX_ITEMS) response_converters: list[_Identifier] | None = Field(default=None, max_length=MAX_ITEMS) diff --git a/pyrit/models/request_limits.py b/pyrit/models/request_limits.py index c9a20196ce..033c8ef9bd 100644 --- a/pyrit/models/request_limits.py +++ b/pyrit/models/request_limits.py @@ -7,3 +7,4 @@ MAX_IDENTIFIER_LENGTH = 256 MAX_LABEL_KEY_LENGTH = 128 MAX_LABEL_VALUE_LENGTH = 1_024 +MAX_TEXT_LENGTH = 100_000 diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index 58851e2b45..dfcecc970e 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -115,6 +115,16 @@ def test_reports_a_document_nested_too_deeply(self, tmp_path: Path) -> None: assert (failure.name, failure.reason) == ("deep", RestoreFailureReason.INVALID_RECIPE) assert list(listing.recipes) == ["contoso_gpt4o"] + def test_reports_a_technique_description_over_the_limit(self, tmp_path: Path) -> None: + document = {"type": "RedTeamingAttack", "description": "x" * 100_001} + _write(tmp_path / "techniques", "long_description", json.dumps(document)) + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TECHNIQUES, source=str(tmp_path)).list_recipes() + + (failure,) = listing.failures + assert (failure.name, failure.reason) == ("long_description", RestoreFailureReason.INVALID_RECIPE) + assert "description" in failure.message + def test_reports_a_newer_schema_version_before_checking_fields(self, tmp_path: Path) -> None: _write(tmp_path / "targets", "future", json.dumps({"schema_version": 2, "unknown_field": True})) From e45865b3fbfddcc541ffdf17936753f9e39aa271 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 01:05:23 +0000 Subject: [PATCH 06/18] Test that a description at the limit is accepted --- tests/unit/registry/test_registry_recipe_storage.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index dfcecc970e..af55c36e64 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -115,14 +115,16 @@ def test_reports_a_document_nested_too_deeply(self, tmp_path: Path) -> None: assert (failure.name, failure.reason) == ("deep", RestoreFailureReason.INVALID_RECIPE) assert list(listing.recipes) == ["contoso_gpt4o"] - def test_reports_a_technique_description_over_the_limit(self, tmp_path: Path) -> None: - document = {"type": "RedTeamingAttack", "description": "x" * 100_001} - _write(tmp_path / "techniques", "long_description", json.dumps(document)) + def test_limits_a_technique_description_like_the_create_request(self, tmp_path: Path) -> None: + folder = tmp_path / "techniques" + _write(folder, "at_limit", json.dumps({"type": "RedTeamingAttack", "description": "x" * 100_000})) + _write(folder, "over_limit", json.dumps({"type": "RedTeamingAttack", "description": "x" * 100_001})) listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TECHNIQUES, source=str(tmp_path)).list_recipes() + assert list(listing.recipes) == ["at_limit"] (failure,) = listing.failures - assert (failure.name, failure.reason) == ("long_description", RestoreFailureReason.INVALID_RECIPE) + assert (failure.name, failure.reason) == ("over_limit", RestoreFailureReason.INVALID_RECIPE) assert "description" in failure.message def test_reports_a_newer_schema_version_before_checking_fields(self, tmp_path: Path) -> None: From 9a419b4ab11205dbd8483c0a8d30ddea782b7fdc Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 03:45:48 +0000 Subject: [PATCH 07/18] Escape text that isn't UTF-8 in restore reports so they can be sent as JSON --- pyrit/models/catalog/registry_recipe.py | 21 +++++++++--- .../registry/test_registry_recipe_storage.py | 34 +++++++++++++++++++ 2 files changed, 51 insertions(+), 4 deletions(-) diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index 5ac95be552..c069ccede5 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -8,7 +8,7 @@ from enum import Enum from typing import Annotated, ClassVar -from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, JsonValue, field_validator from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS, MAX_TEXT_LENGTH @@ -16,6 +16,19 @@ _TechniqueTag = Annotated[str, Field(pattern=r"^[A-Za-z][A-Za-z0-9_]{0,63}$")] +def _escape_invalid_text(value: str) -> str: + """ + Escape characters that can't be written as UTF-8, such as bytes from a file name that isn't UTF-8. + + Returns: + str: The text, safe to send in a JSON response. + """ + return value.encode("utf-8", "backslashreplace").decode("utf-8") + + +_DisplayText = Annotated[str, AfterValidator(_escape_invalid_text)] + + class RegistryRecipeKind(str, Enum): """The registries whose entries can be saved, named after their storage folders.""" @@ -128,7 +141,7 @@ class RestoreFailure(BaseModel): kind: RegistryRecipeKind name: str | None = Field(None, description="Entry name; None when the whole folder couldn't be read") reason: RestoreFailureReason - message: str = Field(..., description="Why, in words safe to show users") + message: _DisplayText = Field(..., description="Why, in words safe to show users") version: str | None = Field(None, description="Version of the stored document, when it could be read") dependency: SavedEntryRef | None = Field(None, description="The entry this one needed that is missing or failed") @@ -139,5 +152,5 @@ class IgnoredRecipeFile(BaseModel): model_config = ConfigDict(frozen=True) kind: RegistryRecipeKind - file_name: str - message: str + file_name: _DisplayText + message: _DisplayText diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index af55c36e64..e7ea4f0d91 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -3,16 +3,20 @@ import hashlib import json +import os from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, patch +import pytest from azure.storage.blob import ContainerClient from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH from pyrit.models.catalog.registry_recipe import ( ComponentRecipe, + IgnoredRecipeFile, RegistryRecipeKind, + RestoreFailure, RestoreFailureReason, TechniqueRecipe, ) @@ -67,6 +71,22 @@ def test_blob_container_source_uses_the_folder_as_prefix(self) -> None: assert list(listing.recipes) == ["formal"] +class TestRestoreReportText: + def test_text_that_is_not_utf8_is_escaped_so_it_can_be_sent_as_json(self) -> None: + ignored = IgnoredRecipeFile( + kind=RegistryRecipeKind.TARGETS, file_name="bad\udcffname.json", message="bad\udcffname" + ) + failure = RestoreFailure( + kind=RegistryRecipeKind.TARGETS, name="x", reason=RestoreFailureReason.BUILD_ERROR, message="got \udcff" + ) + + assert ignored.file_name == "bad\\udcffname.json" + assert failure.message == "got \\udcff" + json.dumps([ignored.model_dump(mode="json"), failure.model_dump(mode="json")], ensure_ascii=False).encode( + "utf-8" + ) + + class TestRegistryRecipeStorageListing: def test_reads_component_recipes_with_their_versions(self, tmp_path: Path) -> None: content = _write(tmp_path / "targets", "contoso_gpt4o", json.dumps(_TARGET)) @@ -136,6 +156,20 @@ def test_reports_a_newer_schema_version_before_checking_fields(self, tmp_path: P assert failure.reason is RestoreFailureReason.NEWER_SCHEMA_VERSION assert "newer PyRIT" in failure.message + def test_reports_a_file_name_that_is_not_utf8(self, tmp_path: Path) -> None: + folder = tmp_path / "targets" + folder.mkdir() + try: + (folder / os.fsdecode(b"bad\xffname.json")).write_text("{}") + except (OSError, UnicodeError): + pytest.skip("This file system only allows UTF-8 file names") + + listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + (ignored,) = listing.ignored + assert ignored.file_name == "bad\\udcffname.json" + json.dumps(ignored.model_dump(mode="json"), ensure_ascii=False).encode("utf-8") + def test_reports_unreadable_and_ignored_files(self, tmp_path: Path) -> None: folder = tmp_path / "scorers" _write(folder, "locked", json.dumps(_TARGET)) From 274b4ce57129e6a7d271ad3b274628df8f5a6da1 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 05:10:06 +0000 Subject: [PATCH 08/18] Escape entry names in restore reports too --- pyrit/models/catalog/registry_recipe.py | 4 ++-- .../unit/registry/test_registry_recipe_storage.py | 15 +++++++++++++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index c069ccede5..82be6a14e6 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -102,7 +102,7 @@ class SavedEntryRef(BaseModel): model_config = ConfigDict(frozen=True) kind: RegistryRecipeKind - name: str + name: _DisplayText class StoredRecipe(BaseModel): @@ -139,7 +139,7 @@ class RestoreFailure(BaseModel): model_config = ConfigDict(frozen=True) kind: RegistryRecipeKind - name: str | None = Field(None, description="Entry name; None when the whole folder couldn't be read") + name: _DisplayText | None = Field(None, description="Entry name; None when the whole folder couldn't be read") reason: RestoreFailureReason message: _DisplayText = Field(..., description="Why, in words safe to show users") version: str | None = Field(None, description="Version of the stored document, when it could be read") diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index e7ea4f0d91..1f49ac1487 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -18,6 +18,7 @@ RegistryRecipeKind, RestoreFailure, RestoreFailureReason, + SavedEntryRef, TechniqueRecipe, ) from pyrit.registry.registry_recipe_storage import RegistryRecipeStorage @@ -77,15 +78,25 @@ def test_text_that_is_not_utf8_is_escaped_so_it_can_be_sent_as_json(self) -> Non kind=RegistryRecipeKind.TARGETS, file_name="bad\udcffname.json", message="bad\udcffname" ) failure = RestoreFailure( - kind=RegistryRecipeKind.TARGETS, name="x", reason=RestoreFailureReason.BUILD_ERROR, message="got \udcff" + kind=RegistryRecipeKind.TARGETS, + name="x\udcff", + reason=RestoreFailureReason.MISSING_DEPENDENCY, + message="got \udcff", + dependency=SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800"), ) assert ignored.file_name == "bad\\udcffname.json" - assert failure.message == "got \\udcff" + assert (failure.name, failure.message) == ("x\\udcff", "got \\udcff") + assert failure.dependency == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\\ud800") json.dumps([ignored.model_dump(mode="json"), failure.model_dump(mode="json")], ensure_ascii=False).encode( "utf-8" ) + def test_valid_text_is_kept_as_written(self) -> None: + ref = SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="caf\u00e9 \uff10\uff0e\uff15 \U0001f600") + + assert ref.name == "caf\u00e9 \uff10\uff0e\uff15 \U0001f600" + class TestRegistryRecipeStorageListing: def test_reads_component_recipes_with_their_versions(self, tmp_path: Path) -> None: From c391b0d81124d64a9c44aa7968348742a1794e69 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 07:06:15 +0000 Subject: [PATCH 09/18] Escape report text only when it is sent as JSON --- pyrit/models/catalog/registry_recipe.py | 5 ++-- .../registry/test_registry_recipe_storage.py | 30 +++++++++++++------ 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index 82be6a14e6..43dd832ef1 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -8,7 +8,7 @@ from enum import Enum from typing import Annotated, ClassVar -from pydantic import AfterValidator, BaseModel, ConfigDict, Field, JsonValue, field_validator +from pydantic import BaseModel, ConfigDict, Field, JsonValue, PlainSerializer, field_validator from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS, MAX_TEXT_LENGTH @@ -26,7 +26,8 @@ def _escape_invalid_text(value: str) -> str: return value.encode("utf-8", "backslashreplace").decode("utf-8") -_DisplayText = Annotated[str, AfterValidator(_escape_invalid_text)] +# The model keeps the text it was given, so a name still matches the entry it names; only JSON output is escaped. +_DisplayText = Annotated[str, PlainSerializer(_escape_invalid_text, return_type=str, when_used="json")] class RegistryRecipeKind(str, Enum): diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index 1f49ac1487..49ef62c6e1 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -73,7 +73,7 @@ def test_blob_container_source_uses_the_folder_as_prefix(self) -> None: class TestRestoreReportText: - def test_text_that_is_not_utf8_is_escaped_so_it_can_be_sent_as_json(self) -> None: + def test_text_that_is_not_utf8_is_escaped_when_sent_as_json(self) -> None: ignored = IgnoredRecipeFile( kind=RegistryRecipeKind.TARGETS, file_name="bad\udcffname.json", message="bad\udcffname" ) @@ -85,17 +85,28 @@ def test_text_that_is_not_utf8_is_escaped_so_it_can_be_sent_as_json(self) -> Non dependency=SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800"), ) - assert ignored.file_name == "bad\\udcffname.json" - assert (failure.name, failure.message) == ("x\\udcff", "got \\udcff") - assert failure.dependency == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\\ud800") - json.dumps([ignored.model_dump(mode="json"), failure.model_dump(mode="json")], ensure_ascii=False).encode( - "utf-8" + sent_ignored, sent_failure = json.loads( + json.dumps([ignored.model_dump(mode="json"), failure.model_dump(mode="json")], ensure_ascii=False).encode( + "utf-8" + ) ) + assert (sent_ignored["file_name"], sent_ignored["message"]) == ("bad\\udcffname.json", "bad\\udcffname") + assert (sent_failure["name"], sent_failure["message"]) == ("x\\udcff", "got \\udcff") + assert sent_failure["dependency"]["name"] == "ghost\\ud800" + assert json.loads(failure.model_dump_json())["dependency"]["name"] == "ghost\\ud800" + + def test_names_keep_the_text_they_were_given(self) -> None: + ref = SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800") + + assert ref.name == "ghost\ud800" + assert ref == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800") + assert ref != SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\\ud800") + def test_valid_text_is_kept_as_written(self) -> None: ref = SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="caf\u00e9 \uff10\uff0e\uff15 \U0001f600") - assert ref.name == "caf\u00e9 \uff10\uff0e\uff15 \U0001f600" + assert ref.model_dump(mode="json")["name"] == "caf\u00e9 \uff10\uff0e\uff15 \U0001f600" class TestRegistryRecipeStorageListing: @@ -178,8 +189,9 @@ def test_reports_a_file_name_that_is_not_utf8(self, tmp_path: Path) -> None: listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() (ignored,) = listing.ignored - assert ignored.file_name == "bad\\udcffname.json" - json.dumps(ignored.model_dump(mode="json"), ensure_ascii=False).encode("utf-8") + assert ignored.file_name == os.fsdecode(b"bad\xffname.json") + sent = json.loads(json.dumps(ignored.model_dump(mode="json"), ensure_ascii=False).encode("utf-8")) + assert sent["file_name"] == "bad\\udcffname.json" def test_reports_unreadable_and_ignored_files(self, tmp_path: Path) -> None: folder = tmp_path / "scorers" From 507567b7be04f37b4b88e7c7d3f18261aea1c10a Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 21:03:42 +0000 Subject: [PATCH 10/18] Check the opened file before reading it and skip links deleted since the listing --- pyrit/registry/file_document_storage.py | 39 +++++++- .../registry/test_file_document_storage.py | 99 +++++++++++++++---- 2 files changed, 115 insertions(+), 23 deletions(-) diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index 4ba3f0949d..c025c77eff 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -16,6 +16,7 @@ import hashlib import logging import os +import stat import sys import tempfile import time @@ -109,6 +110,27 @@ def _release_exclusive_lock(descriptor: int) -> None: fcntl.flock(descriptor, fcntl.LOCK_UN) +# Opening without blocking keeps a pipe from hanging the read; it changes nothing for a regular +# file. Windows has no O_NONBLOCK and needs O_BINARY so the bytes aren't translated. +_READ_FLAGS = os.O_RDONLY | getattr(os, "O_NONBLOCK", 0) | getattr(os, "O_BINARY", 0) + + +def _read_regular_file(path: str) -> bytes | None: + """ + Read a file only if what was opened is a regular file. + + The opened file is checked rather than the path, so the check still holds if the path + was replaced after it was listed. + + Returns: + bytes | None: The content, or None if the path names something other than a regular file. + """ + with open(os.open(path, _READ_FLAGS), "rb") as file: + if not stat.S_ISREG(os.fstat(file.fileno()).st_mode): + return None + return file.read() + + class DocumentConflictError(ValueError): """A stored document changed after the caller read it.""" @@ -279,19 +301,26 @@ def _read_local_entry(self, *, entry: os.DirEntry[str], name: str) -> bytes | Un """ Read one listed local document. + The directory can change after it is listed, so what is opened is checked again: an entry + deleted in between is skipped, whether or not it was a link, and anything that isn't a regular + file is reported without being read. + Returns: bytes | UnreadableDocument | None: The content, why it can't be read, or None if it was deleted after the directory was listed. """ try: - if not entry.is_file(): - # A broken link, a folder, or a pipe, which would block if read. - return UnreadableDocument(name=name, message="Not a regular file") - return Path(entry.path).read_bytes() + content = _read_regular_file(entry.path) if entry.is_file() else None except FileNotFoundError: - return None + content = None except OSError as error: return UnreadableDocument(name=name, message=self._describe_read_error(error)) + if content is not None: + return content + if not os.path.lexists(entry.path): + return None + # A broken link, a folder, or a pipe, which would block if read. + return UnreadableDocument(name=name, message="Not a regular file") def _list_blob_documents_with_issues(self) -> DocumentListing: """ diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index 63bc7601dc..a6e264a45f 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -1,7 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from collections.abc import Iterator +import os +from collections.abc import Callable, Generator, Iterator +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -10,7 +12,12 @@ from azure.core.exceptions import HttpResponseError, ResourceNotFoundError, ServiceRequestError from azure.storage.blob import ContainerClient, StorageErrorCode -from pyrit.registry.file_document_storage import DocumentListing, FileDocumentStorage, UnreadableDocument +from pyrit.registry.file_document_storage import ( + DocumentListing, + FileDocumentStorage, + UnreadableDocument, + _read_regular_file, +) _SAS_SOURCE = "https://account.blob.core.windows.net/documents/saved?sp=rl&sig=secret-signature" @@ -30,18 +37,46 @@ def _blob_client(*, names: list[str]) -> MagicMock: return client +@contextmanager +def _change_after_listing(*, change: Callable[[], None], cache_types: bool = False) -> Generator[None, None, None]: + # Runs `change` after the folder is listed and before its entries are read. With `cache_types`, + # each entry's type is looked up first, so it is out of date once `change` runs. + real_scandir = os.scandir + + @contextmanager + def scandir(path: Path) -> Generator[Iterator[os.DirEntry[str]], None, None]: + with real_scandir(path) as entries: + listed = list(entries) + if cache_types: + for entry in listed: + entry.is_file() + yield iter(listed) + change() + + with patch("pyrit.registry.file_document_storage.os.scandir", side_effect=scandir): + yield + + class TestLocalListingWithIssues: def test_reports_unreadable_document_and_keeps_the_others(self, tmp_path: Path) -> None: (tmp_path / "first.json").write_bytes(b"one") (tmp_path / "second.json").write_bytes(b"two") original_read_bytes = Path.read_bytes + def read(path: str) -> bytes | None: + if Path(path).name == "second.json": + raise PermissionError(13, "Permission denied", path) + return _read_regular_file(path) + def read_bytes(path: Path) -> bytes: if path.name == "second.json": raise PermissionError(13, "Permission denied", str(path)) return original_read_bytes(path) - with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + with ( + patch("pyrit.registry.file_document_storage._read_regular_file", side_effect=read), + patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes), + ): listing = _local_storage(tmp_path)._list_documents_with_issues() legacy = _local_storage(tmp_path)._list_documents() @@ -58,13 +93,12 @@ def test_reports_badly_named_files_as_ignored_without_reading_them(self, tmp_pat (tmp_path / "nested").mkdir() (tmp_path / "nested" / "inner.json").write_bytes(b"not a direct child") read_names: list[str] = [] - original_read_bytes = Path.read_bytes - def read_bytes(path: Path) -> bytes: - read_names.append(path.name) - return original_read_bytes(path) + def read(path: str) -> bytes | None: + read_names.append(Path(path).name) + return _read_regular_file(path) - with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + with patch("pyrit.registry.file_document_storage._read_regular_file", side_effect=read): listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing.documents == {"valid_name": b"ok"} @@ -75,20 +109,49 @@ def read_bytes(path: Path) -> bytes: assert _local_storage(tmp_path)._list_documents() == listing.documents def test_skips_a_document_deleted_after_listing(self, tmp_path: Path) -> None: - (tmp_path / "kept.json").write_bytes(b"kept") - (tmp_path / "deleted.json").write_bytes(b"gone") - original_read_bytes = Path.read_bytes + (tmp_path / "kept.json").write_bytes(b"{}\r\n") + (tmp_path / "empty.json").write_bytes(b"") + deleted = tmp_path / "deleted.json" + deleted.write_bytes(b"gone") - def read_bytes(path: Path) -> bytes: - if path.name == "deleted.json": - raise FileNotFoundError(2, "No such file or directory", str(path)) - return original_read_bytes(path) + with _change_after_listing(change=deleted.unlink, cache_types=True): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing(documents={"empty": b"", "kept": b"{}\r\n"}) - with patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes): + def test_skips_a_link_deleted_after_listing(self, tmp_path: Path) -> None: + (tmp_path / "kept.json").write_bytes(b"kept") + (tmp_path / "target.txt").write_bytes(b"linked") + link = tmp_path / "linked.json" + try: + link.symlink_to(tmp_path / "target.txt") + except OSError: + pytest.skip("Creating symbolic links isn't permitted here") + + with _change_after_listing(change=link.unlink): listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing == DocumentListing(documents={"kept": b"kept"}) + @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + @pytest.mark.timeout(60) + def test_reports_a_pipe_put_in_place_of_a_listed_document_without_blocking(self, tmp_path: Path) -> None: + (tmp_path / "kept.json").write_bytes(b"kept") + swapped = tmp_path / "swapped.json" + swapped.write_bytes(b"replaced") + + def replace_with_pipe() -> None: + swapped.unlink() + os.mkfifo(swapped) + + with _change_after_listing(change=replace_with_pipe, cache_types=True): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing( + documents={"kept": b"kept"}, + unreadable=(UnreadableDocument(name="swapped", message="Not a regular file"),), + ) + def test_reports_entries_that_are_not_regular_files_without_reading_them(self, tmp_path: Path) -> None: (tmp_path / "folder.json").mkdir() try: @@ -96,14 +159,14 @@ def test_reports_entries_that_are_not_regular_files_without_reading_them(self, t except OSError: pytest.skip("Creating symbolic links isn't permitted here") - with patch.object(Path, "read_bytes", autospec=True) as read_bytes: + with patch("pyrit.registry.file_document_storage._read_regular_file") as read: listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing.unreadable == ( UnreadableDocument(name="broken_link", message="Not a regular file"), UnreadableDocument(name="folder", message="Not a regular file"), ) - read_bytes.assert_not_called() + read.assert_not_called() def test_reports_a_link_loop_and_keeps_the_others(self, tmp_path: Path) -> None: (tmp_path / "good.json").write_bytes(b"ok") From a45ac239ba57f9a70b6805908b27c2a2cadf2996 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 3 Nov 2026 22:46:18 +0000 Subject: [PATCH 11/18] Close a file that isn't regular before reporting it and report entries that can't be checked --- pyrit/registry/file_document_storage.py | 26 +++++-- .../registry/test_file_document_storage.py | 70 +++++++++++++++++-- 2 files changed, 86 insertions(+), 10 deletions(-) diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index c025c77eff..c75d398b59 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -125,10 +125,14 @@ def _read_regular_file(path: str) -> bytes | None: Returns: bytes | None: The content, or None if the path names something other than a regular file. """ - with open(os.open(path, _READ_FLAGS), "rb") as file: - if not stat.S_ISREG(os.fstat(file.fileno()).st_mode): + descriptor = os.open(path, _READ_FLAGS) + try: + if not stat.S_ISREG(os.fstat(descriptor).st_mode): return None - return file.read() + with open(descriptor, "rb", closefd=False) as file: + return file.read() + finally: + os.close(descriptor) class DocumentConflictError(ValueError): @@ -317,8 +321,22 @@ def _read_local_entry(self, *, entry: os.DirEntry[str], name: str) -> bytes | Un return UnreadableDocument(name=name, message=self._describe_read_error(error)) if content is not None: return content - if not os.path.lexists(entry.path): + return self._describe_non_file(path=entry.path, name=name) + + @classmethod + def _describe_non_file(cls, *, path: str, name: str) -> UnreadableDocument | None: + """ + Report a listed entry that isn't a regular file, unless it was deleted after the listing. + + Returns: + UnreadableDocument | None: Why it can't be read, or None if it no longer exists. + """ + try: + os.lstat(path) + except FileNotFoundError: return None + except OSError as error: + return UnreadableDocument(name=name, message=cls._describe_read_error(error)) # A broken link, a folder, or a pipe, which would block if read. return UnreadableDocument(name=name, message="Not a regular file") diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index a6e264a45f..2c52fdea3e 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. import os +import sys from collections.abc import Callable, Generator, Iterator from contextlib import contextmanager from pathlib import Path @@ -37,6 +38,12 @@ def _blob_client(*, names: list[str]) -> MagicMock: return client +def _count_open_files() -> int | None: + # Linux lists this process's open files here; elsewhere the count is skipped. + descriptors = Path("/proc/self/fd") + return len(list(descriptors.iterdir())) if descriptors.is_dir() else None + + @contextmanager def _change_after_listing(*, change: Callable[[], None], cache_types: bool = False) -> Generator[None, None, None]: # Runs `change` after the folder is listed and before its entries are read. With `cache_types`, @@ -133,24 +140,56 @@ def test_skips_a_link_deleted_after_listing(self, tmp_path: Path) -> None: assert listing == DocumentListing(documents={"kept": b"kept"}) - @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + @pytest.mark.parametrize("replacement", ["pipe", "folder", "dangling link"]) @pytest.mark.timeout(60) - def test_reports_a_pipe_put_in_place_of_a_listed_document_without_blocking(self, tmp_path: Path) -> None: + def test_reports_a_non_file_put_in_place_of_a_listed_document_without_reading_it( + self, tmp_path: Path, replacement: str + ) -> None: + if replacement == "pipe" and not hasattr(os, "mkfifo"): + pytest.skip("Named pipes need os.mkfifo") + if replacement == "folder" and sys.platform == "win32": + pytest.skip("Windows can't open a folder as a file") + if replacement == "dangling link" and sys.platform == "win32": + pytest.skip("Creating symbolic links may not be permitted on Windows") (tmp_path / "kept.json").write_bytes(b"kept") swapped = tmp_path / "swapped.json" swapped.write_bytes(b"replaced") - def replace_with_pipe() -> None: + def replace() -> None: swapped.unlink() - os.mkfifo(swapped) - - with _change_after_listing(change=replace_with_pipe, cache_types=True): + if replacement == "pipe": + os.mkfifo(swapped) + elif replacement == "folder": + swapped.mkdir() + else: + swapped.symlink_to(tmp_path / "missing_target.json") + + open_files = _count_open_files() + with _change_after_listing(change=replace, cache_types=True): listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing == DocumentListing( documents={"kept": b"kept"}, unreadable=(UnreadableDocument(name="swapped", message="Not a regular file"),), ) + assert _count_open_files() == open_files + + @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + def test_checks_the_opened_file_rather_than_the_path(self, tmp_path: Path) -> None: + document = tmp_path / "document.json" + document.write_bytes(b"original") + real_open = os.open + + def open_then_replace(path: str, flags: int, *args: int) -> int: + descriptor = real_open(path, flags, *args) + os.unlink(path) + os.mkfifo(path) + return descriptor + + with patch("pyrit.registry.file_document_storage.os.open", side_effect=open_then_replace): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing(documents={"document": b"original"}) def test_reports_entries_that_are_not_regular_files_without_reading_them(self, tmp_path: Path) -> None: (tmp_path / "folder.json").mkdir() @@ -168,6 +207,25 @@ def test_reports_entries_that_are_not_regular_files_without_reading_them(self, t ) read.assert_not_called() + @pytest.mark.skipif(sys.platform == "win32", reason="The null device is a POSIX character device here") + def test_reads_only_regular_files(self) -> None: + assert _read_regular_file(os.devnull) is None + + def test_reports_an_entry_that_cannot_be_checked_after_listing(self, tmp_path: Path) -> None: + deleted = tmp_path / "deleted.json" + deleted.write_bytes(b"gone") + + with ( + _change_after_listing(change=deleted.unlink, cache_types=True), + patch( + "pyrit.registry.file_document_storage.os.lstat", + side_effect=PermissionError(13, "Permission denied"), + ), + ): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing(unreadable=(UnreadableDocument(name="deleted", message="Permission denied"),)) + def test_reports_a_link_loop_and_keeps_the_others(self, tmp_path: Path) -> None: (tmp_path / "good.json").write_bytes(b"ok") try: From db245dcfbef51a74008deb3d01c5f92181a942f2 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Wed, 4 Nov 2026 00:08:59 +0000 Subject: [PATCH 12/18] Test a file replaced just before it's opened and stop a blocked read with a watchdog --- .../registry/test_file_document_storage.py | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index 2c52fdea3e..03926a16dc 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -141,7 +141,7 @@ def test_skips_a_link_deleted_after_listing(self, tmp_path: Path) -> None: assert listing == DocumentListing(documents={"kept": b"kept"}) @pytest.mark.parametrize("replacement", ["pipe", "folder", "dangling link"]) - @pytest.mark.timeout(60) + @pytest.mark.timeout(60, method="thread") def test_reports_a_non_file_put_in_place_of_a_listed_document_without_reading_it( self, tmp_path: Path, replacement: str ) -> None: @@ -175,6 +175,7 @@ def replace() -> None: assert _count_open_files() == open_files @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + @pytest.mark.timeout(60, method="thread") def test_checks_the_opened_file_rather_than_the_path(self, tmp_path: Path) -> None: document = tmp_path / "document.json" document.write_bytes(b"original") @@ -190,6 +191,25 @@ def open_then_replace(path: str, flags: int, *args: int) -> int: listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing == DocumentListing(documents={"document": b"original"}) + assert document.is_fifo() + + @pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="Named pipes need os.mkfifo") + @pytest.mark.timeout(60, method="thread") + def test_checks_a_file_replaced_just_before_it_is_opened(self, tmp_path: Path) -> None: + (tmp_path / "document.json").write_bytes(b"original") + real_open = os.open + + def replace_then_open(path: str, flags: int, *args: int) -> int: + os.unlink(path) + os.mkfifo(path) + return real_open(path, flags, *args) + + with patch("pyrit.registry.file_document_storage.os.open", side_effect=replace_then_open): + listing = _local_storage(tmp_path)._list_documents_with_issues() + + assert listing == DocumentListing( + unreadable=(UnreadableDocument(name="document", message="Not a regular file"),) + ) def test_reports_entries_that_are_not_regular_files_without_reading_them(self, tmp_path: Path) -> None: (tmp_path / "folder.json").mkdir() From d1bdd82830bae5a89df4249744b306bfa16b3917 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Thu, 5 Nov 2026 21:25:14 +0000 Subject: [PATCH 13/18] List presets, suites, and initializers through the checked listing --- pyrit/registry/file_document_storage.py | 86 ++++--------------- .../registry/test_file_document_storage.py | 74 ++++++++++++---- 2 files changed, 76 insertions(+), 84 deletions(-) diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index c75d398b59..45150c8d72 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -245,22 +245,32 @@ 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. """ - return self._list_blob_documents() if self._is_blob else self._list_local_documents() + 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.name}{self._extension}' in {self.display_source}: " + f"{unreadable.message}" + ) + return listing.documents def _list_documents_with_issues(self) -> DocumentListing: """ Read every stored document and report the ones that can't be used. - Unlike ``_list_documents``, nothing is skipped quietly: a document that can't be read - is reported as unreadable, and a file whose name breaks the naming rule is reported as - ignored without being read, because no single-document call could address it. A - storage location that can't be listed at all raises instead of looking empty. + A document that can't be read is reported as unreadable, and a file whose name breaks + the naming rule is reported as ignored without being read, because no single-document + call could address it. A storage location that can't be listed at all raises instead of + looking empty. Messages carry only a short reason, never a path or URL, so they can be shown to users without exposing server paths or a SAS token. @@ -407,42 +417,6 @@ def _name_problem(name: str) -> str | None: return str(error) return None - def _list_local_documents(self) -> dict[str, bytes]: - """ - Read addressable documents from the configured local directory. - - Returns: - dict[str, bytes]: Document content keyed by file stem. - """ - directory = self._local_directory(create=True) - documents: dict[str, bytes] = {} - for path in sorted(directory.glob(f"*{self._extension}")): - if not self._is_addressable_name(path.stem): - 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 - - def _is_addressable_name(self, name: str) -> bool: - """ - Return whether a discovered document name is one this storage can address. - - 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. - - Returns: - bool: Whether the name is a legal registry 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 - def _read_document_bytes(self, name: str) -> bytes | None: """ Read the raw bytes of one document. @@ -770,32 +744,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/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index 03926a16dc..966be0729d 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -68,30 +68,21 @@ class TestLocalListingWithIssues: def test_reports_unreadable_document_and_keeps_the_others(self, tmp_path: Path) -> None: (tmp_path / "first.json").write_bytes(b"one") (tmp_path / "second.json").write_bytes(b"two") - original_read_bytes = Path.read_bytes def read(path: str) -> bytes | None: if Path(path).name == "second.json": raise PermissionError(13, "Permission denied", path) return _read_regular_file(path) - def read_bytes(path: Path) -> bytes: - if path.name == "second.json": - raise PermissionError(13, "Permission denied", str(path)) - return original_read_bytes(path) - - with ( - patch("pyrit.registry.file_document_storage._read_regular_file", side_effect=read), - patch.object(Path, "read_bytes", autospec=True, side_effect=read_bytes), - ): + with patch("pyrit.registry.file_document_storage._read_regular_file", side_effect=read): listing = _local_storage(tmp_path)._list_documents_with_issues() - legacy = _local_storage(tmp_path)._list_documents() + documents = _local_storage(tmp_path)._list_documents() assert listing == DocumentListing( documents={"first": b"one"}, unreadable=(UnreadableDocument(name="second", message="Permission denied"),), ) - assert legacy == {"first": b"one"} + 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") @@ -258,13 +249,13 @@ def test_reports_a_link_loop_and_keeps_the_others(self, tmp_path: Path) -> None: assert listing.documents == {"good": b"ok"} assert [issue.name for issue in listing.unreadable] == ["loop"] - def test_matches_extensions_like_the_legacy_listing(self, tmp_path: Path) -> None: + 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") - storage = _local_storage(tmp_path) + listing = _local_storage(tmp_path)._list_documents_with_issues() - assert storage._list_documents_with_issues().documents == storage._list_documents() + 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() @@ -352,6 +343,59 @@ def failing_pages() -> Iterator[SimpleNamespace]: _blob_storage()._list_documents_with_issues() +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("pyrit.registry.file_document_storage.os.scandir", side_effect=PermissionError(13, "denied")): + with pytest.raises(PermissionError): + _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 TestDescribeReadError: @pytest.mark.parametrize( ("error", "expected"), From 76dbf9c098e14dce17cf3092814155074dc09813 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Thu, 5 Nov 2026 22:38:18 +0000 Subject: [PATCH 14/18] Report a preset, suite, or initializer folder that can't be read as unavailable storage --- pyrit/backend/routes/initializers.py | 4 ++-- pyrit/backend/routes/scenario_presets.py | 2 +- pyrit/backend/routes/scenario_suites.py | 2 +- .../unit/backend/test_initializer_service.py | 21 +++++++++++++++++++ .../backend/test_scenario_preset_routes.py | 10 +++++++-- .../backend/test_scenario_suite_routes.py | 9 +++++--- 6 files changed, 39 insertions(+), 9 deletions(-) diff --git a/pyrit/backend/routes/initializers.py b/pyrit/backend/routes/initializers.py index 6c99ca8067..65734b037a 100644 --- a/pyrit/backend/routes/initializers.py +++ b/pyrit/backend/routes/initializers.py @@ -137,7 +137,7 @@ async def list_custom_initializers(request: Request) -> CustomInitializerListRes _check_custom_initializers_allowed(request) try: return await get_initializer_service().list_custom_initializers_async() - except AzureError as exc: + except (AzureError, OSError) as exc: raise _custom_storage_unavailable() from exc @@ -240,7 +240,7 @@ async def unregister_initializer( # pyrit-async-suffix-exempt try: await service.unregister_initializer_async(initializer_name=initializer_name) - except AzureError as exc: + except (AzureError, OSError) as exc: raise _custom_storage_unavailable() from exc except ValueError as e: raise HTTPException( diff --git a/pyrit/backend/routes/scenario_presets.py b/pyrit/backend/routes/scenario_presets.py index 223e679154..aa6721df8e 100644 --- a/pyrit/backend/routes/scenario_presets.py +++ b/pyrit/backend/routes/scenario_presets.py @@ -106,7 +106,7 @@ async def list_scenario_presets( # pyrit-async-suffix-exempt """ try: return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) - except AzureError as exc: + except (AzureError, OSError) as exc: raise _storage_unavailable() from exc diff --git a/pyrit/backend/routes/scenario_suites.py b/pyrit/backend/routes/scenario_suites.py index e9da0ce2c9..ff11d2ac3a 100644 --- a/pyrit/backend/routes/scenario_suites.py +++ b/pyrit/backend/routes/scenario_suites.py @@ -89,7 +89,7 @@ async def list_scenario_suites() -> ScenarioSuiteListResponse: # pyrit-async-su """ try: return await get_scenario_suite_service().list_suites_async() - except AzureError as exc: + except (AzureError, OSError) as exc: raise _storage_unavailable() from exc diff --git a/tests/unit/backend/test_initializer_service.py b/tests/unit/backend/test_initializer_service.py index c6f7125a4d..03fd7c5de3 100644 --- a/tests/unit/backend/test_initializer_service.py +++ b/tests/unit/backend/test_initializer_service.py @@ -489,6 +489,27 @@ def test_custom_initializer_routes_return_503_for_storage_failure( assert response.status_code == status.HTTP_503_SERVICE_UNAVAILABLE assert response.json()["detail"] == "Custom initializer storage is temporarily unavailable" + @pytest.mark.parametrize("operation", ["list", "delete"]) + def test_custom_initializer_routes_return_503_for_a_folder_that_cant_be_read( + self, + client_with_custom_initializers_enabled: TestClient, + operation: str, + ) -> None: + error = PermissionError(13, "Permission denied", "/srv/secret/initializers") + with patch("pyrit.backend.routes.initializers.get_initializer_service") as mock_get_service: + mock_service = MagicMock() + mock_service.list_custom_initializers_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") + else: + 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 "/srv/secret" not in response.text + def test_post_root_still_registers_runtime_initializer( self, client_with_custom_initializers_enabled: TestClient ) -> None: diff --git a/tests/unit/backend/test_scenario_preset_routes.py b/tests/unit/backend/test_scenario_preset_routes.py index 3e418cc92d..f67f952a03 100644 --- a/tests/unit/backend/test_scenario_preset_routes.py +++ b/tests/unit/backend/test_scenario_preset_routes.py @@ -147,8 +147,14 @@ 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")) + @pytest.mark.parametrize( + "error", + [AzureError("container 'x' key=secret"), PermissionError(13, "Permission denied", "/srv/secret/presets")], + ) + def test_storage_failure_is_reported_without_sdk_detail( + self, client: TestClient, service: MagicMock, error: Exception + ) -> None: + service.list_presets_async = AsyncMock(side_effect=error) response = client.get("/api/scenario-presets") diff --git a/tests/unit/backend/test_scenario_suite_routes.py b/tests/unit/backend/test_scenario_suite_routes.py index a0d879a21b..2ef323aab7 100644 --- a/tests/unit/backend/test_scenario_suite_routes.py +++ b/tests/unit/backend/test_scenario_suite_routes.py @@ -333,9 +333,12 @@ 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")) +@pytest.mark.parametrize( + "error", [AzureError("container key=secret"), PermissionError(13, "Permission denied", "/srv/secret/suites")] +) +def test_storage_failure_is_sanitized(client: TestClient, service: MagicMock, error: Exception) -> None: + """Test that storage details do not reach API callers.""" + service.list_suites_async = AsyncMock(side_effect=error) response = client.get("/api/scenario-suites") From 7d9ac2fdaf834aa63e551f868eb81c852b4e95b9 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Sat, 7 Nov 2026 01:23:02 +0000 Subject: [PATCH 15/18] Log why storage is unavailable, without its path or URL --- pyrit/backend/routes/initializers.py | 19 ++++++-- pyrit/backend/routes/scenario_presets.py | 24 +++++++--- pyrit/backend/routes/scenario_suites.py | 24 +++++++--- pyrit/registry/file_document_storage.py | 45 ++++++++++--------- .../unit/backend/test_initializer_service.py | 9 +++- .../backend/test_scenario_preset_routes.py | 9 +++- .../backend/test_scenario_suite_routes.py | 13 ++++-- .../registry/test_file_document_storage.py | 7 +-- 8 files changed, 102 insertions(+), 48 deletions(-) diff --git a/pyrit/backend/routes/initializers.py b/pyrit/backend/routes/initializers.py index 65734b037a..8e27952150 100644 --- a/pyrit/backend/routes/initializers.py +++ b/pyrit/backend/routes/initializers.py @@ -14,6 +14,8 @@ DELETE /api/initializers/{name} — unregister an initializer """ +import logging + from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, Request, status @@ -28,17 +30,26 @@ ) from pyrit.backend.services.initializer_service import get_initializer_service from pyrit.models.catalog.initializer import RegisteredInitializer +from pyrit.registry.file_document_storage import describe_storage_error + +logger = logging.getLogger(__name__) router = APIRouter(prefix="/initializers", tags=["initializers"]) -def _custom_storage_unavailable() -> HTTPException: +def _custom_storage_unavailable(error: Exception) -> HTTPException: """ Create a sanitized response for unavailable custom initializer storage. + The cause is logged without any path or URL it carries. + + Args: + error (Exception): The storage error. + Returns: HTTPException: A service-unavailable response without SDK details. """ + logger.warning(f"Custom initializer storage is unavailable: {describe_storage_error(error)}") return HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Custom initializer storage is temporarily unavailable", @@ -138,7 +149,7 @@ async def list_custom_initializers(request: Request) -> CustomInitializerListRes try: return await get_initializer_service().list_custom_initializers_async() except (AzureError, OSError) as exc: - raise _custom_storage_unavailable() from exc + raise _custom_storage_unavailable(exc) from exc @router.get( @@ -203,7 +214,7 @@ 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 + raise _custom_storage_unavailable(exc) from exc except ValueError as e: detail = str(e) if "already registered" in detail: @@ -241,7 +252,7 @@ async def unregister_initializer( # pyrit-async-suffix-exempt try: await service.unregister_initializer_async(initializer_name=initializer_name) except (AzureError, OSError) as exc: - raise _custom_storage_unavailable() from exc + raise _custom_storage_unavailable(exc) 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 aa6721df8e..986adb49ff 100644 --- a/pyrit/backend/routes/scenario_presets.py +++ b/pyrit/backend/routes/scenario_presets.py @@ -17,6 +17,7 @@ POST /api/scenario-presets/{name}/resolve — combine a preset with launch fields """ +import logging from typing import Annotated from azure.core.exceptions import AzureError @@ -32,17 +33,26 @@ from pyrit.backend.services.scenario_preset_service import get_scenario_preset_service from pyrit.models.catalog import ResolveScenarioPresetRequest, RunScenarioRequest, ScenarioPreset from pyrit.registry import ScenarioPresetConflictError, ScenarioPresetNameRegisteredError +from pyrit.registry.file_document_storage import describe_storage_error + +logger = logging.getLogger(__name__) router = APIRouter(prefix="/scenario-presets", tags=["scenario-presets"]) -def _storage_unavailable() -> HTTPException: +def _storage_unavailable(error: Exception) -> HTTPException: """ Create a sanitized response for unavailable preset storage. + The cause is logged without any path or URL it carries. + + Args: + error (Exception): The storage error. + Returns: HTTPException: A service-unavailable response without SDK details. """ + logger.warning(f"Scenario preset storage is unavailable: {describe_storage_error(error)}") return HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Scenario preset storage is temporarily unavailable", @@ -78,7 +88,7 @@ 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 + raise _storage_unavailable(exc) from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -107,7 +117,7 @@ async def list_scenario_presets( # pyrit-async-suffix-exempt try: return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) except (AzureError, OSError) as exc: - raise _storage_unavailable() from exc + raise _storage_unavailable(exc) from exc @router.get( @@ -176,7 +186,7 @@ async def create_scenario_preset( # pyrit-async-suffix-exempt detail=f"Scenario preset '{preset.name}' already exists", ) from exc except AzureError as exc: - raise _storage_unavailable() from exc + raise _storage_unavailable(exc) from exc @router.put( @@ -238,7 +248,7 @@ async def update_scenario_preset( # pyrit-async-suffix-exempt 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 + raise _storage_unavailable(exc) from exc @router.delete( @@ -260,7 +270,7 @@ async def delete_scenario_preset(name: str) -> None: # pyrit-async-suffix-exemp try: deleted = await get_scenario_preset_service().delete_preset_async(name=name) except AzureError as exc: - raise _storage_unavailable() from exc + raise _storage_unavailable(exc) from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -311,7 +321,7 @@ async def resolve_scenario_preset( # pyrit-async-suffix-exempt 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 + raise _storage_unavailable(exc) 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 ff11d2ac3a..3d7150914a 100644 --- a/pyrit/backend/routes/scenario_suites.py +++ b/pyrit/backend/routes/scenario_suites.py @@ -3,6 +3,7 @@ """Scenario suite definition API routes.""" +import logging from typing import Annotated from uuid import UUID @@ -30,17 +31,26 @@ ScenarioSuiteNameRegisteredError, ScenarioSuiteNotRunnableError, ) +from pyrit.registry.file_document_storage import describe_storage_error + +logger = logging.getLogger(__name__) router = APIRouter(prefix="/scenario-suites", tags=["scenario-suites"]) -def _storage_unavailable() -> HTTPException: +def _storage_unavailable(error: Exception) -> HTTPException: """ Create a sanitized response for unavailable suite storage. + The cause is logged without any path or URL it carries. + + Args: + error (Exception): The storage error. + Returns: HTTPException: A service-unavailable response without SDK details. """ + logger.warning(f"Scenario suite storage is unavailable: {describe_storage_error(error)}") return HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Scenario suite storage is temporarily unavailable", @@ -70,7 +80,7 @@ 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 + raise _storage_unavailable(exc) from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -90,7 +100,7 @@ async def list_scenario_suites() -> ScenarioSuiteListResponse: # pyrit-async-su try: return await get_scenario_suite_service().list_suites_async() except (AzureError, OSError) as exc: - raise _storage_unavailable() from exc + raise _storage_unavailable(exc) from exc @router.get( @@ -185,7 +195,7 @@ async def launch_scenario_suite( # pyrit-async-suffix-exempt 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 + raise _storage_unavailable(exc) from exc except RuntimeError as exc: raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(exc)) from None @@ -243,7 +253,7 @@ async def create_scenario_suite( # pyrit-async-suffix-exempt detail=f"Scenario suite '{suite.name}' already exists", ) from exc except AzureError as exc: - raise _storage_unavailable() from exc + raise _storage_unavailable(exc) from exc @router.put( @@ -293,7 +303,7 @@ async def update_scenario_suite( # pyrit-async-suffix-exempt 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 + raise _storage_unavailable(exc) from exc @router.delete( @@ -307,7 +317,7 @@ async def delete_scenario_suite(name: str) -> None: # pyrit-async-suffix-exempt try: deleted = await get_scenario_suite_service().delete_suite_async(name=name) except AzureError as exc: - raise _storage_unavailable() from exc + raise _storage_unavailable(exc) from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None diff --git a/pyrit/registry/file_document_storage.py b/pyrit/registry/file_document_storage.py index 45150c8d72..36a8e42da5 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -135,6 +135,27 @@ def _read_regular_file(path: str) -> bytes | None: os.close(descriptor) +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.""" @@ -328,7 +349,7 @@ def _read_local_entry(self, *, entry: os.DirEntry[str], name: str) -> bytes | Un except FileNotFoundError: content = None except OSError as error: - return UnreadableDocument(name=name, message=self._describe_read_error(error)) + return UnreadableDocument(name=name, message=describe_storage_error(error)) if content is not None: return content return self._describe_non_file(path=entry.path, name=name) @@ -346,7 +367,7 @@ def _describe_non_file(cls, *, path: str, name: str) -> UnreadableDocument | Non except FileNotFoundError: return None except OSError as error: - return UnreadableDocument(name=name, message=cls._describe_read_error(error)) + return UnreadableDocument(name=name, message=describe_storage_error(error)) # A broken link, a folder, or a pipe, which would block if read. return UnreadableDocument(name=name, message="Not a regular file") @@ -382,27 +403,9 @@ def _list_blob_documents_with_issues(self) -> DocumentListing: if getattr(error, "error_code", None) == StorageErrorCode.BLOB_NOT_FOUND: # Deleted after the container was listed, so it is no longer part of the listing. continue - unreadable.append(UnreadableDocument(name=name, message=self._describe_read_error(error))) + unreadable.append(UnreadableDocument(name=name, message=describe_storage_error(error))) return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) - @staticmethod - def _describe_read_error(error: Exception) -> str: - """ - Describe a storage error in a few words, leaving out any path or URL it carries. - - Returns: - str: For example ``Permission denied`` or ``HttpResponseError (403, AuthorizationPermissionMismatch)``. - """ - if isinstance(error, OSError) and error.strerror: - return error.strerror - error_code = getattr(error, "error_code", None) - details = [ - str(detail) - for detail in (getattr(error, "status_code", None), getattr(error_code, "value", error_code)) - if detail is not None - ] - return f"{type(error).__name__} ({', '.join(details)})" if details else type(error).__name__ - @staticmethod def _name_problem(name: str) -> str | None: """ diff --git a/tests/unit/backend/test_initializer_service.py b/tests/unit/backend/test_initializer_service.py index 03fd7c5de3..fdd79cc9d4 100644 --- a/tests/unit/backend/test_initializer_service.py +++ b/tests/unit/backend/test_initializer_service.py @@ -5,6 +5,7 @@ Tests for backend initializer service and routes. """ +import logging from collections.abc import Iterator from unittest.mock import AsyncMock, MagicMock, patch @@ -494,9 +495,13 @@ def test_custom_initializer_routes_return_503_for_a_folder_that_cant_be_read( self, client_with_custom_initializers_enabled: TestClient, operation: str, + caplog: pytest.LogCaptureFixture, ) -> None: error = PermissionError(13, "Permission denied", "/srv/secret/initializers") - with patch("pyrit.backend.routes.initializers.get_initializer_service") as mock_get_service: + with ( + patch("pyrit.backend.routes.initializers.get_initializer_service") as mock_get_service, + caplog.at_level(logging.WARNING, logger="pyrit.backend.routes.initializers"), + ): mock_service = MagicMock() mock_service.list_custom_initializers_async = AsyncMock(side_effect=error) mock_service.unregister_initializer_async = AsyncMock(side_effect=error) @@ -509,6 +514,8 @@ def test_custom_initializer_routes_return_503_for_a_folder_that_cant_be_read( assert response.status_code == status.HTTP_503_SERVICE_UNAVAILABLE assert response.json()["detail"] == "Custom initializer storage is temporarily unavailable" assert "/srv/secret" not in response.text + assert "Custom initializer storage is unavailable: Permission denied" in caplog.messages + assert "/srv/secret" not in caplog.text 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 f67f952a03..5151fc9e9b 100644 --- a/tests/unit/backend/test_scenario_preset_routes.py +++ b/tests/unit/backend/test_scenario_preset_routes.py @@ -3,6 +3,7 @@ """Tests for the scenario preset routes.""" +import logging from collections.abc import Iterator from unittest.mock import AsyncMock, MagicMock, patch @@ -25,6 +26,7 @@ ScenarioPresetNameRegisteredError, resolve_scenario_preset_run_request, ) +from pyrit.registry.file_document_storage import describe_storage_error PRESET_NAME = "quick_scan" SCENARIO_NAME = "foundry.red_team_agent" @@ -152,14 +154,17 @@ def test_list_is_readable_without_admin(self, anonymous_client: TestClient, serv [AzureError("container 'x' key=secret"), PermissionError(13, "Permission denied", "/srv/secret/presets")], ) def test_storage_failure_is_reported_without_sdk_detail( - self, client: TestClient, service: MagicMock, error: Exception + self, client: TestClient, service: MagicMock, error: Exception, caplog: pytest.LogCaptureFixture ) -> None: service.list_presets_async = AsyncMock(side_effect=error) - response = client.get("/api/scenario-presets") + with caplog.at_level(logging.WARNING, logger="pyrit.backend.routes.scenario_presets"): + response = client.get("/api/scenario-presets") assert response.status_code == 503 assert "secret" not in response.text + assert f"Scenario preset storage is unavailable: {describe_storage_error(error)}" in caplog.messages + assert "secret" not in caplog.text def test_each_preset_carries_its_own_run_size(self, client: TestClient, service: MagicMock) -> None: estimate = _estimate(42) diff --git a/tests/unit/backend/test_scenario_suite_routes.py b/tests/unit/backend/test_scenario_suite_routes.py index 2ef323aab7..d3680f4b36 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 @@ -43,6 +44,7 @@ ScenarioSuiteRegistry, ScenarioSuiteStorage, ) +from pyrit.registry.file_document_storage import describe_storage_error SUITE_NAME = "weekly" ENTRY_ID = UUID("11111111-1111-4111-8111-111111111111") @@ -336,14 +338,19 @@ def test_get_registered_suite_returns_origin_and_null_version(client: TestClient @pytest.mark.parametrize( "error", [AzureError("container key=secret"), PermissionError(13, "Permission denied", "/srv/secret/suites")] ) -def test_storage_failure_is_sanitized(client: TestClient, service: MagicMock, error: Exception) -> None: - """Test that storage details do not reach API callers.""" +def test_storage_failure_is_sanitized( + client: TestClient, service: MagicMock, error: Exception, caplog: pytest.LogCaptureFixture +) -> None: + """Test that storage details reach neither API callers nor the log.""" service.list_suites_async = AsyncMock(side_effect=error) - response = client.get("/api/scenario-suites") + with caplog.at_level(logging.WARNING, logger="pyrit.backend.routes.scenario_suites"): + response = client.get("/api/scenario-suites") assert response.status_code == 503 assert "secret" not in response.text + assert f"Scenario suite storage is unavailable: {describe_storage_error(error)}" in caplog.messages + assert "secret" not in caplog.text def test_create_succeeds_when_preset_storage_is_unavailable( diff --git a/tests/unit/registry/test_file_document_storage.py b/tests/unit/registry/test_file_document_storage.py index 966be0729d..9a8aed9843 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -18,6 +18,7 @@ FileDocumentStorage, UnreadableDocument, _read_regular_file, + describe_storage_error, ) _SAS_SOURCE = "https://account.blob.core.windows.net/documents/saved?sp=rl&sig=secret-signature" @@ -396,7 +397,7 @@ def download_blob(blob_name: str) -> SimpleNamespace: assert "secret-signature" not in message -class TestDescribeReadError: +class TestDescribeStorageError: @pytest.mark.parametrize( ("error", "expected"), [ @@ -406,7 +407,7 @@ class TestDescribeReadError: ], ) def test_describes_errors_without_their_details(self, error: Exception, expected: str) -> None: - assert FileDocumentStorage._describe_read_error(error) == expected + assert describe_storage_error(error) == expected @pytest.mark.parametrize( ("error_code", "expected"), @@ -420,4 +421,4 @@ def test_includes_the_azure_status_and_error_code(self, error_code: str, expecte error.status_code = 403 error.error_code = error_code - assert FileDocumentStorage._describe_read_error(error) == expected + assert describe_storage_error(error) == expected From 19205b12116ebc6b57fc1d2ec640bfe6e7d50616 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Sun, 8 Nov 2026 00:26:17 +0000 Subject: [PATCH 16/18] Import the saved recipe models from their package, as main now requires --- pyrit/registry/registry_recipe_storage.py | 2 +- tests/unit/registry/test_registry_recipe_storage.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyrit/registry/registry_recipe_storage.py b/pyrit/registry/registry_recipe_storage.py index 0a956a19b8..4a4cf7da15 100644 --- a/pyrit/registry/registry_recipe_storage.py +++ b/pyrit/registry/registry_recipe_storage.py @@ -13,7 +13,7 @@ from pydantic import ValidationError from pyrit.common.azure_storage import is_azure_blob_uri -from pyrit.models.catalog.registry_recipe import ( +from pyrit.models.catalog import ( ComponentRecipe, IgnoredRecipeFile, RegistryRecipeKind, diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index 67b641abed..7df79c2851 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -12,7 +12,7 @@ from azure.storage.blob import ContainerClient from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH -from pyrit.models.catalog.registry_recipe import ( +from pyrit.models.catalog import ( ComponentRecipe, IgnoredRecipeFile, RegistryRecipeKind, From 622bbb88b65414c49977fcd9c55a7d89bc226c28 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 9 Nov 2026 23:02:47 +0000 Subject: [PATCH 17/18] Report storage failures with one error and answer them with 503 FileDocumentStorage now raises DocumentStorageUnavailableError when the folder or container can't be used, and DocumentUnreadableError when one document exists but can't be read. The message names the location, with any SAS token hidden, and a short reason. One backend handler maps both to 503 and logs the location and reason, with the traceback at debug level, so the per-route storage helpers are gone. A single read, a conditional save and a delete now follow the listing's rules: a folder, pipe or broken link named like a document is unreadable rather than missing, and on Blob only a deleted blob counts as missing. Blob sign-in, permission, network and missing-container errors stop the listing instead of marking every file unreadable. Each unreadable report keeps the real file name and the original error, and the listing's documents are read-only. Initializer delete removes a script that can't be read or decoded, and registering one no longer turns a storage write failure into a 400. Tests patch module-level aliases instead of os. --- pyrit/backend/middleware/error_handlers.py | 34 ++ pyrit/backend/routes/initializers.py | 34 +- pyrit/backend/routes/scenario_presets.py | 39 +- pyrit/backend/routes/scenario_suites.py | 39 +- .../services/scenario_suite_service.py | 5 +- .../components/initializer_registry.py | 11 +- pyrit/registry/custom_initializer_storage.py | 10 + pyrit/registry/file_document_storage.py | 374 ++++++++++++++---- tests/unit/backend/test_error_handlers.py | 50 +++ .../unit/backend/test_initializer_service.py | 56 +-- .../backend/test_scenario_preset_routes.py | 31 +- .../backend/test_scenario_suite_routes.py | 41 +- .../backend/test_scenario_suite_service.py | 15 +- .../registry/test_file_document_storage.py | 264 +++++++++++-- .../registry/test_initializer_registry.py | 24 +- .../registry/test_scenario_preset_storage.py | 21 +- .../registry/test_scenario_suite_storage.py | 5 +- 17 files changed, 737 insertions(+), 316 deletions(-) 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 2747c4a9b3..529f92d641 100644 --- a/pyrit/backend/routes/initializers.py +++ b/pyrit/backend/routes/initializers.py @@ -14,9 +14,6 @@ DELETE /api/initializers/{name} — unregister an initializer """ -import logging - -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from pyrit.backend.middleware.auth import require_admin @@ -33,32 +30,10 @@ ) from pyrit.backend.services import get_initializer_service from pyrit.models.catalog import RegisteredInitializer -from pyrit.registry.file_document_storage import describe_storage_error - -logger = logging.getLogger(__name__) router = APIRouter(prefix="/initializers", tags=["initializers"]) -def _custom_storage_unavailable(error: Exception) -> HTTPException: - """ - Create a sanitized response for unavailable custom initializer storage. - - The cause is logged without any path or URL it carries. - - Args: - error (Exception): The storage error. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - logger.warning(f"Custom initializer storage is unavailable: {describe_storage_error(error)}") - 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. @@ -149,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, OSError) as exc: - raise _custom_storage_unavailable(exc) from exc + return await get_initializer_service().list_custom_initializers_async() @router.get( @@ -216,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(exc) from exc except ValueError as e: detail = str(e) if "already registered" in detail: @@ -254,8 +224,6 @@ async def unregister_initializer( # pyrit-async-suffix-exempt try: await service.unregister_initializer_async(initializer_name=initializer_name) - except (AzureError, OSError) as exc: - raise _custom_storage_unavailable(exc) 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 a4e45598b1..973112df7b 100644 --- a/pyrit/backend/routes/scenario_presets.py +++ b/pyrit/backend/routes/scenario_presets.py @@ -17,10 +17,8 @@ POST /api/scenario-presets/{name}/resolve — combine a preset with launch fields """ -import logging 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 @@ -35,32 +33,10 @@ from pyrit.models import ResolveScenarioPresetRequest, ScenarioPreset from pyrit.models.catalog import RunScenarioRequest from pyrit.registry import ScenarioPresetConflictError, ScenarioPresetNameRegisteredError -from pyrit.registry.file_document_storage import describe_storage_error - -logger = logging.getLogger(__name__) router = APIRouter(prefix="/scenario-presets", tags=["scenario-presets"]) -def _storage_unavailable(error: Exception) -> HTTPException: - """ - Create a sanitized response for unavailable preset storage. - - The cause is logged without any path or URL it carries. - - Args: - error (Exception): The storage error. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - logger.warning(f"Scenario preset storage is unavailable: {describe_storage_error(error)}") - 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. @@ -89,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(exc) from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -116,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, OSError) as exc: - raise _storage_unavailable(exc) from exc + return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) @router.get( @@ -187,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(exc) from exc @router.put( @@ -249,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(exc) from exc @router.delete( @@ -271,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(exc) from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -322,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(exc) 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 aecd04c608..f1ec16872e 100644 --- a/pyrit/backend/routes/scenario_suites.py +++ b/pyrit/backend/routes/scenario_suites.py @@ -3,11 +3,9 @@ """Scenario suite definition API routes.""" -import logging 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 @@ -28,32 +26,10 @@ ScenarioSuiteNameRegisteredError, ScenarioSuiteNotRunnableError, ) -from pyrit.registry.file_document_storage import describe_storage_error - -logger = logging.getLogger(__name__) router = APIRouter(prefix="/scenario-suites", tags=["scenario-suites"]) -def _storage_unavailable(error: Exception) -> HTTPException: - """ - Create a sanitized response for unavailable suite storage. - - The cause is logged without any path or URL it carries. - - Args: - error (Exception): The storage error. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - logger.warning(f"Scenario suite storage is unavailable: {describe_storage_error(error)}") - 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. @@ -76,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(exc) from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -94,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, OSError) as exc: - raise _storage_unavailable(exc) from exc + return await get_scenario_suite_service().list_suites_async() @router.get( @@ -191,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(exc) from exc except RuntimeError as exc: raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(exc)) from None @@ -249,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(exc) from exc @router.put( @@ -299,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(exc) from exc @router.delete( @@ -313,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(exc) 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/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 5341898cf5..c18a34a2cb 100644 --- a/pyrit/registry/file_document_storage.py +++ b/pyrit/registry/file_document_storage.py @@ -12,6 +12,7 @@ from __future__ import annotations +import errno import fnmatch import hashlib import logging @@ -23,6 +24,7 @@ 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 @@ -30,19 +32,33 @@ 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.""" + """ + 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) @@ -57,11 +73,59 @@ class IgnoredDocument: class DocumentListing: """Every document a listing found, including the ones it could not use.""" - documents: dict[str, bytes] = field(default_factory=dict) + 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. @@ -120,12 +184,13 @@ 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. + 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 = os.open(path, _READ_FLAGS) + descriptor = _open(path, _READ_FLAGS) try: if not stat.S_ISREG(os.fstat(descriptor).st_mode): return None @@ -135,6 +200,37 @@ def _read_regular_file(path: str) -> bytes | None: 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. @@ -208,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. """ @@ -271,6 +374,9 @@ def _list_documents(self) -> dict[str, bytes]: 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: @@ -279,10 +385,9 @@ def _list_documents(self) -> dict[str, bytes]: ) for unreadable in listing.unreadable: logger.warning( - f"Skipping unreadable document '{unreadable.name}{self._extension}' in {self.display_source}: " - f"{unreadable.message}" + f"Skipping unreadable document '{unreadable.file_name}' in {self.display_source}: {unreadable.message}" ) - return listing.documents + return dict(listing.documents) def _list_documents_with_issues(self) -> DocumentListing: """ @@ -291,15 +396,21 @@ def _list_documents_with_issues(self) -> DocumentListing: 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. + 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_with_issues() if self._is_blob else self._list_local_documents_with_issues() + 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_with_issues(self) -> DocumentListing: """ @@ -311,7 +422,7 @@ def _list_local_documents_with_issues(self) -> DocumentListing: directory = self._local_directory(create=True) # ``Path.glob`` returns nothing for a directory it can't read; listing it directly raises instead. # ``fnmatch`` applies the platform's case rule, as ``glob`` does. - with os.scandir(directory) as entries: + 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, @@ -325,51 +436,50 @@ def _list_local_documents_with_issues(self) -> DocumentListing: if problem is not None: ignored.append(IgnoredDocument(file_name=entry.name, message=problem)) continue - content = self._read_local_entry(entry=entry, name=name) - if isinstance(content, bytes): + try: + 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 - elif content is not None: - unreadable.append(content) - return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) + return DocumentListing( + documents=MappingProxyType(documents), unreadable=tuple(unreadable), ignored=tuple(ignored) + ) - def _read_local_entry(self, *, entry: os.DirEntry[str], name: str) -> bytes | UnreadableDocument | None: + def _read_local_document(self, *, path: str, name: str, file_name: str) -> bytes | None: """ - Read one listed local document. + Read one local document, the same way for a listing and for a single read. - The directory can change after it is listed, so what is opened is checked again: an entry - deleted in between is skipped, whether or not it was a link, and anything that isn't a regular - file is reported without being read. + 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 | UnreadableDocument | None: The content, why it can't be read, or None if it was - deleted after the directory was listed. + 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: - content = _read_regular_file(entry.path) if entry.is_file() else None + if stat.S_ISREG(_stat(path).st_mode): + content = _read_regular_file(path) + if content is not None: + return content except FileNotFoundError: - content = None + pass except OSError as error: - return UnreadableDocument(name=name, message=describe_storage_error(error)) - if content is not None: - return content - return self._describe_non_file(path=entry.path, name=name) - - @classmethod - def _describe_non_file(cls, *, path: str, name: str) -> UnreadableDocument | None: - """ - Report a listed entry that isn't a regular file, unless it was deleted after the listing. - - Returns: - UnreadableDocument | None: Why it can't be read, or None if it no longer exists. - """ + raise self._unreadable(name=name, file_name=file_name, reason=describe_storage_error(error)) from error try: - os.lstat(path) + _lstat(path) except FileNotFoundError: return None except OSError as error: - return UnreadableDocument(name=name, message=describe_storage_error(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. - return UnreadableDocument(name=name, message="Not a regular file") + raise self._unreadable(name=name, file_name=file_name, reason="Not a regular file") def _list_blob_documents_with_issues(self) -> DocumentListing: """ @@ -377,9 +487,11 @@ def _list_blob_documents_with_issues(self) -> DocumentListing: Returns: DocumentListing: The documents and issues found under the configured prefix. + + Raises: + DocumentStorageUnavailableError: If a blob error affects the whole container. """ from azure.core.exceptions import AzureError - from azure.storage.blob import StorageErrorCode documents: dict[str, bytes] = {} unreadable: list[UnreadableDocument] = [] @@ -400,11 +512,19 @@ def _list_blob_documents_with_issues(self) -> DocumentListing: try: documents[name] = client.download_blob(blob_name).readall() except AzureError as error: - if getattr(error, "error_code", None) == StorageErrorCode.BLOB_NOT_FOUND: + if _is_missing_blob(error): # Deleted after the container was listed, so it is no longer part of the listing. continue - unreadable.append(UnreadableDocument(name=name, message=describe_storage_error(error))) - return DocumentListing(documents=documents, unreadable=tuple(unreadable), ignored=tuple(ignored)) + 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: @@ -422,26 +542,89 @@ def _name_problem(name: str) -> str | 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 + ) - path = self._local_directory() / f"{name}{self._extension}" - return path.read_bytes() if path.is_file() else None + @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 + + 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: """ @@ -449,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: @@ -506,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: """ @@ -522,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) @@ -545,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 @@ -573,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, @@ -595,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: @@ -613,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 @@ -665,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 @@ -714,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. 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 ed9ede3cf1..18c6ea3f2e 100644 --- a/tests/unit/backend/test_initializer_service.py +++ b/tests/unit/backend/test_initializer_service.py @@ -10,7 +10,6 @@ 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 @@ -27,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 @@ -471,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") @@ -490,34 +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 - @pytest.mark.parametrize("operation", ["list", "delete"]) - def test_custom_initializer_routes_return_503_for_a_folder_that_cant_be_read( - self, - client_with_custom_initializers_enabled: TestClient, - operation: str, - caplog: pytest.LogCaptureFixture, + def test_deleting_a_script_that_cant_be_read_returns_503_with_the_reason( + self, client_with_custom_initializers_enabled: TestClient ) -> None: - error = PermissionError(13, "Permission denied", "/srv/secret/initializers") - with ( - patch("pyrit.backend.routes.initializers.get_initializer_service") as mock_get_service, - caplog.at_level(logging.WARNING, logger="pyrit.backend.routes.initializers"), - ): + 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.list_custom_initializers_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") - else: - response = client_with_custom_initializers_enabled.delete("/api/initializers/custom") + + 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 "/srv/secret" not in response.text - assert "Custom initializer storage is unavailable: Permission denied" in caplog.messages - assert "/srv/secret" not in caplog.text + 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 c3b244e464..31750c98e9 100644 --- a/tests/unit/backend/test_scenario_preset_routes.py +++ b/tests/unit/backend/test_scenario_preset_routes.py @@ -8,7 +8,6 @@ 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 @@ -25,7 +24,7 @@ ScenarioPresetNameRegisteredError, resolve_scenario_preset_run_request, ) -from pyrit.registry.file_document_storage import describe_storage_error +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError PRESET_NAME = "quick_scan" SCENARIO_NAME = "foundry.red_team_agent" @@ -148,22 +147,22 @@ def test_list_is_readable_without_admin(self, anonymous_client: TestClient, serv assert anonymous_client.get("/api/scenario-presets").status_code == 200 - @pytest.mark.parametrize( - "error", - [AzureError("container 'x' key=secret"), PermissionError(13, "Permission denied", "/srv/secret/presets")], - ) - def test_storage_failure_is_reported_without_sdk_detail( - self, client: TestClient, service: MagicMock, error: Exception, caplog: pytest.LogCaptureFixture + 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=error) + service.list_presets_async = AsyncMock( + side_effect=DocumentStorageUnavailableError( + label="Scenario preset", source="/srv/presets", reason="Permission denied" + ) + ) - with caplog.at_level(logging.WARNING, logger="pyrit.backend.routes.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 f"Scenario preset storage is unavailable: {describe_storage_error(error)}" in caplog.messages - assert "secret" not in caplog.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) @@ -427,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}", @@ -435,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 7d7acdaef6..fd6f1345fc 100644 --- a/tests/unit/backend/test_scenario_suite_routes.py +++ b/tests/unit/backend/test_scenario_suite_routes.py @@ -12,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 @@ -49,7 +48,7 @@ ScenarioSuiteRegistry, ScenarioSuiteStorage, ) -from pyrit.registry.file_document_storage import describe_storage_error +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError SUITE_NAME = "weekly" ENTRY_ID = UUID("11111111-1111-4111-8111-111111111111") @@ -183,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 ( @@ -259,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), ], ) @@ -275,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" @@ -371,22 +375,25 @@ def test_get_registered_suite_returns_origin_and_null_version(client: TestClient assert response.json()["version"] is None -@pytest.mark.parametrize( - "error", [AzureError("container key=secret"), PermissionError(13, "Permission denied", "/srv/secret/suites")] -) -def test_storage_failure_is_sanitized( - client: TestClient, service: MagicMock, error: Exception, caplog: pytest.LogCaptureFixture +def test_storage_failure_names_the_source_only_in_the_log( + client: TestClient, service: MagicMock, caplog: pytest.LogCaptureFixture ) -> None: - """Test that storage details reach neither API callers nor the log.""" - service.list_suites_async = AsyncMock(side_effect=error) + """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", + ) + ) - with caplog.at_level(logging.WARNING, logger="pyrit.backend.routes.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 f"Scenario suite storage is unavailable: {describe_storage_error(error)}" in caplog.messages - assert "secret" not in caplog.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 index 9a8aed9843..48c6fa68ca 100644 --- a/tests/unit/registry/test_file_document_storage.py +++ b/tests/unit/registry/test_file_document_storage.py @@ -10,11 +10,21 @@ from unittest.mock import MagicMock, patch import pytest -from azure.core.exceptions import HttpResponseError, ResourceNotFoundError, ServiceRequestError +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, @@ -39,6 +49,16 @@ def _blob_client(*, names: list[str]) -> MagicMock: 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") @@ -61,7 +81,7 @@ def scandir(path: Path) -> Generator[Iterator[os.DirEntry[str]], None, None]: yield iter(listed) change() - with patch("pyrit.registry.file_document_storage.os.scandir", side_effect=scandir): + with patch.object(storage_module, "_scandir", side_effect=scandir): yield @@ -81,8 +101,9 @@ def read(path: str) -> bytes | None: assert listing == DocumentListing( documents={"first": b"one"}, - unreadable=(UnreadableDocument(name="second", message="Permission denied"),), + 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: @@ -162,8 +183,9 @@ def replace() -> None: assert listing == DocumentListing( documents={"kept": b"kept"}, - unreadable=(UnreadableDocument(name="swapped", message="Not a regular file"),), + 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") @@ -179,7 +201,7 @@ def open_then_replace(path: str, flags: int, *args: int) -> int: os.mkfifo(path) return descriptor - with patch("pyrit.registry.file_document_storage.os.open", side_effect=open_then_replace): + 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"}) @@ -196,11 +218,11 @@ def replace_then_open(path: str, flags: int, *args: int) -> int: os.mkfifo(path) return real_open(path, flags, *args) - with patch("pyrit.registry.file_document_storage.os.open", side_effect=replace_then_open): + 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", message="Not a regular file"),) + 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: @@ -214,8 +236,8 @@ def test_reports_entries_that_are_not_regular_files_without_reading_them(self, t listing = _local_storage(tmp_path)._list_documents_with_issues() assert listing.unreadable == ( - UnreadableDocument(name="broken_link", message="Not a regular file"), - UnreadableDocument(name="folder", message="Not a regular file"), + 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() @@ -229,14 +251,13 @@ def test_reports_an_entry_that_cannot_be_checked_after_listing(self, tmp_path: P with ( _change_after_listing(change=deleted.unlink, cache_types=True), - patch( - "pyrit.registry.file_document_storage.os.lstat", - side_effect=PermissionError(13, "Permission denied"), - ), + 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", message="Permission denied"),)) + 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") @@ -267,10 +288,18 @@ def test_creates_missing_directory_and_lists_nothing(self, tmp_path: Path) -> No def test_raises_when_the_directory_cannot_be_listed(self, tmp_path: Path) -> None: (tmp_path / "first.json").write_bytes(b"one") - with patch("pyrit.registry.file_document_storage.os.scandir", side_effect=PermissionError(13, "denied")): - with pytest.raises(PermissionError): + 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: @@ -286,9 +315,7 @@ def test_reports_unreadable_and_ignored_blobs_without_exposing_the_url(self) -> def download_blob(blob_name: str) -> SimpleNamespace: if blob_name == "saved/broken.json": - error = HttpResponseError(message=f"Forbidden: {_SAS_SOURCE}") - error.status_code = 403 - raise error + 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 @@ -297,7 +324,10 @@ def download_blob(blob_name: str) -> SimpleNamespace: listing = _blob_storage()._list_documents_with_issues() assert listing.documents == {"good": b"content"} - assert listing.unreadable == (UnreadableDocument(name="broken", message="HttpResponseError (403)"),) + 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/") @@ -305,18 +335,11 @@ def download_blob(blob_name: str) -> SimpleNamespace: assert downloaded == ["saved/broken.json", "saved/good.json"] def test_skips_only_a_blob_deleted_after_listing(self) -> None: - client = _blob_client(names=["saved/kept.json", "saved/deleted.json", "saved/container_gone.json"]) - error_codes = { - "saved/deleted.json": StorageErrorCode.BLOB_NOT_FOUND, - "saved/container_gone.json": StorageErrorCode.CONTAINER_NOT_FOUND, - } + client = _blob_client(names=["saved/kept.json", "saved/deleted.json"]) def download_blob(blob_name: str) -> SimpleNamespace: - if blob_name in error_codes: - error = ResourceNotFoundError(message="Not found.") - error.status_code = 404 - error.error_code = error_codes[blob_name] - raise error + 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 @@ -324,12 +347,29 @@ def download_blob(blob_name: str) -> SimpleNamespace: with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): listing = _blob_storage()._list_documents_with_issues() - assert listing == DocumentListing( - documents={"kept": b"kept"}, - unreadable=( - UnreadableDocument(name="container_gone", message="ResourceNotFoundError (404, ContainerNotFound)"), - ), - ) + 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]: @@ -340,9 +380,155 @@ def failing_pages() -> Iterator[SimpleNamespace]: client.list_blobs.return_value = failing_pages() with patch("azure.storage.blob.ContainerClient.from_container_url", return_value=client): - with pytest.raises(HttpResponseError): + 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: @@ -368,8 +554,8 @@ def test_leaves_out_a_pipe_named_like_a_document_without_blocking(self, tmp_path 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("pyrit.registry.file_document_storage.os.scandir", side_effect=PermissionError(13, "denied")): - with pytest.raises(PermissionError): + 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: 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_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) From eba7bf2dbe8cead44a245b0bd6744ddfda281707 Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Mon, 9 Nov 2026 23:14:43 +0000 Subject: [PATCH 18/18] Escape file names once and simplify the saved recipe models List recipes escapes the OS text it reports, such as a file name that isn't UTF-8, and a recipe whose JSON holds an unpaired surrogate escape is reported as invalid, so the models no longer need a JSON-only serializer. Recipes may start with a byte-order mark. A document without schema_version reads as version 1, and each file is checked against the format version of the model that reads it. Validation summaries shorten each location part and the whole message, since a location can repeat a stored key. Recipes and StoredRecipe are no longer frozen, since their params and tags never were, and StoredRecipe checks that its recipe type matches its registry. The listing's recipes are read-only, and it keeps the error behind each recipe that couldn't be read, for the log. --- pyrit/models/catalog/registry_recipe.py | 71 +++++---- pyrit/registry/registry_recipe_storage.py | 92 +++++++++--- .../registry/test_registry_recipe_storage.py | 138 +++++++++++++----- 3 files changed, 219 insertions(+), 82 deletions(-) diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index 43dd832ef1..d8846147d2 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -8,7 +8,7 @@ from enum import Enum from typing import Annotated, ClassVar -from pydantic import BaseModel, ConfigDict, Field, JsonValue, PlainSerializer, field_validator +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 @@ -16,20 +16,6 @@ _TechniqueTag = Annotated[str, Field(pattern=r"^[A-Za-z][A-Za-z0-9_]{0,63}$")] -def _escape_invalid_text(value: str) -> str: - """ - Escape characters that can't be written as UTF-8, such as bytes from a file name that isn't UTF-8. - - Returns: - str: The text, safe to send in a JSON response. - """ - return value.encode("utf-8", "backslashreplace").decode("utf-8") - - -# The model keeps the text it was given, so a name still matches the entry it names; only JSON output is escaped. -_DisplayText = Annotated[str, PlainSerializer(_escape_invalid_text, return_type=str, when_used="json")] - - class RegistryRecipeKind(str, Enum): """The registries whose entries can be saved, named after their storage folders.""" @@ -44,11 +30,29 @@ class _RegistryRecipe(BaseModel): CURRENT_SCHEMA_VERSION: ClassVar[int] = 1 - model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False, frozen=True) + model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False) - schema_version: int = Field( - CURRENT_SCHEMA_VERSION, ge=1, le=CURRENT_SCHEMA_VERSION, description="Shape of this document" - ) + # 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): @@ -103,18 +107,33 @@ class SavedEntryRef(BaseModel): model_config = ConfigDict(frozen=True) kind: RegistryRecipeKind - name: _DisplayText + name: str class StoredRecipe(BaseModel): """A saved recipe and the version of the document it was read from.""" - model_config = ConfigDict(frozen=True) - ref: SavedEntryRef recipe: ComponentRecipe | TechniqueRecipe version: str = Field(..., description="Opaque version of the document this recipe was read from") + @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.""" @@ -140,9 +159,9 @@ class RestoreFailure(BaseModel): model_config = ConfigDict(frozen=True) kind: RegistryRecipeKind - name: _DisplayText | None = Field(None, description="Entry name; None when the whole folder couldn't be read") + name: str | None = Field(None, description="Entry name; None when the whole folder couldn't be read") reason: RestoreFailureReason - message: _DisplayText = Field(..., description="Why, in words safe to show users") + 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") @@ -153,5 +172,5 @@ class IgnoredRecipeFile(BaseModel): model_config = ConfigDict(frozen=True) kind: RegistryRecipeKind - file_name: _DisplayText - message: _DisplayText + file_name: str + message: str diff --git a/pyrit/registry/registry_recipe_storage.py b/pyrit/registry/registry_recipe_storage.py index 4a4cf7da15..f0a222e3c5 100644 --- a/pyrit/registry/registry_recipe_storage.py +++ b/pyrit/registry/registry_recipe_storage.py @@ -8,6 +8,8 @@ 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 @@ -25,14 +27,42 @@ ) 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: dict[str, StoredRecipe] = field(default_factory=dict) + 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): @@ -97,19 +127,25 @@ 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: - OSError: If the local folder can't be listed. - AzureError: If the Blob container can't be listed. + 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=issue.message) + 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): @@ -117,10 +153,19 @@ def list_recipes(self) -> RecipeListing: else: recipes[name] = parsed ignored = tuple( - IgnoredRecipeFile(kind=self._kind, file_name=issue.file_name, message=issue.message) + 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=recipes, failures=tuple(failures), ignored=ignored) + return RecipeListing( + recipes=MappingProxyType(recipes), + failures=tuple(failures), + ignored=ignored, + causes=MappingProxyType(causes), + ) def _parse_recipe(self, *, name: str, content: bytes) -> StoredRecipe | RestoreFailure: """ @@ -131,22 +176,31 @@ def _parse_recipe(self, *, name: str, content: bytes) -> StoredRecipe | RestoreF """ version = self._compute_version(content) invalid = RestoreFailureReason.INVALID_RECIPE + too_deep = "Nested too deeply to read." try: - payload = json.loads(content.decode("utf-8")) + # 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="Nested too deeply to read.") + 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 = ComponentRecipe.CURRENT_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 ) - model = TechniqueRecipe if self._kind is RegistryRecipeKind.TECHNIQUES else ComponentRecipe try: recipe = model.model_validate(payload) except ValidationError as error: @@ -154,7 +208,7 @@ def _parse_recipe(self, *, name: str, content: bytes) -> StoredRecipe | RestoreF name=name, version=version, reason=invalid, message=self._describe_validation_error(error) ) except RecursionError: - return self._failure(name=name, version=version, reason=invalid, message="Nested too deeply to read.") + return 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( @@ -171,12 +225,16 @@ def _failure( @staticmethod def _describe_validation_error(error: ValidationError) -> str: """ - Summarize a validation error without echoing the stored values. + 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: One line per problem, naming the field and what's wrong with it. + str: The summary, at most ``MAX_FAILURE_MESSAGE_LENGTH`` characters long. """ - problems = [ - f"{'.'.join(str(part) for part in issue['loc']) or 'document'}: {issue['msg']}" for issue in error.errors() - ] - return "Not a valid saved recipe: " + "; ".join(problems) + 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/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index 7df79c2851..077169efe3 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -10,19 +10,19 @@ 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, - IgnoredRecipeFile, RegistryRecipeKind, - RestoreFailure, RestoreFailureReason, SavedEntryRef, + StoredRecipe, TechniqueRecipe, ) -from pyrit.registry.file_document_storage import _read_regular_file -from pyrit.registry.registry_recipe_storage import RegistryRecipeStorage +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, @@ -73,41 +73,29 @@ def test_blob_container_source_uses_the_folder_as_prefix(self) -> None: assert list(listing.recipes) == ["formal"] -class TestRestoreReportText: - def test_text_that_is_not_utf8_is_escaped_when_sent_as_json(self) -> None: - ignored = IgnoredRecipeFile( - kind=RegistryRecipeKind.TARGETS, file_name="bad\udcffname.json", message="bad\udcffname" - ) - failure = RestoreFailure( - kind=RegistryRecipeKind.TARGETS, - name="x\udcff", - reason=RestoreFailureReason.MISSING_DEPENDENCY, - message="got \udcff", - dependency=SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800"), - ) - - sent_ignored, sent_failure = json.loads( - json.dumps([ignored.model_dump(mode="json"), failure.model_dump(mode="json")], ensure_ascii=False).encode( - "utf-8" - ) - ) - - assert (sent_ignored["file_name"], sent_ignored["message"]) == ("bad\\udcffname.json", "bad\\udcffname") - assert (sent_failure["name"], sent_failure["message"]) == ("x\\udcff", "got \\udcff") - assert sent_failure["dependency"]["name"] == "ghost\\ud800" - assert json.loads(failure.model_dump_json())["dependency"]["name"] == "ghost\\ud800" +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"} - def test_names_keep_the_text_they_were_given(self) -> None: - ref = SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800") + with patch.object(ComponentRecipe, "CURRENT_SCHEMA_VERSION", 2): + assert ComponentRecipe.model_validate(document).schema_version == 1 - assert ref.name == "ghost\ud800" - assert ref == SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\ud800") - assert ref != SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="ghost\\ud800") + def test_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}) - def test_valid_text_is_kept_as_written(self) -> None: - ref = SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name="caf\u00e9 \uff10\uff0e\uff15 \U0001f600") - - assert ref.model_dump(mode="json")["name"] == "caf\u00e9 \uff10\uff0e\uff15 \U0001f600" + @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: @@ -190,9 +178,9 @@ def test_reports_a_file_name_that_is_not_utf8(self, tmp_path: Path) -> None: listing = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() (ignored,) = listing.ignored - assert ignored.file_name == os.fsdecode(b"bad\xffname.json") - sent = json.loads(json.dumps(ignored.model_dump(mode="json"), ensure_ascii=False).encode("utf-8")) - assert sent["file_name"] == "bad\\udcffname.json" + 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" @@ -214,6 +202,78 @@ def read(path: str) -> bytes | None: "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()