diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 51488ee0d8..4fbad18e24 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -161,6 +161,21 @@ allow_custom_initializers: false # as scenario_presets_source. Defaults to ~/.pyrit/scenario_suites. # scenario_suites_source: https://account.blob.core.windows.net/container/scenario_suites +# Optional storage for saved targets, converters, scorers, and techniques, so they +# come back when the backend restarts. Each kind has its own folder under the source. +# This supports the same local directory and Azure Blob sources as +# scenario_presets_source. If any saved entry can't be restored, the backend stays +# not ready and lists the failures in its runtime status. Fix or delete the listed +# files, then apply the configuration or restart the backend. Defaults to +# ~/.pyrit/registry_recipes. +# +# ⚠️ WARNING: The backend rebuilds every saved entry in this source as an +# administrator, at startup and live reinitialization. Anyone who can write to this +# source can make the backend run a model's repository code or use its identity. +# Restrict write access to trusted administrators, and use a dedicated folder or +# blob prefix. +# registry_recipes_source: https://account.blob.core.windows.net/container/registry_recipes + # Upload Folder for API-Created Targets # ------------------------------------- # Folder that targets created through the REST API may upload local files from diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index bba64b45d4..c00c7baf78 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -194,6 +194,20 @@ Suites store ordered references to scenario presets and resolve each preset's cu A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/scenario_suites`. Local writes use the same hidden lock-file and atomic replacement behavior as scenario presets. Startup scans the configured source so malformed suite documents are reported immediately without preventing the backend from starting. +### `registry_recipes_source` + +Stores saved targets, converters, scorers, and techniques in a local directory or Azure Blob container, so they come back when the backend restarts. An Azure URI may include a blob-name prefix: + +```yaml +registry_recipes_source: https://account.blob.core.windows.net/pyrit-storage/registry-recipes +``` + +Each kind has its own folder under the source (`targets/`, `converters/`, `scorers/`, `techniques/`), with one JSON file per saved entry named after the entry. A file holds the create request the entry was made with, and restore builds the entry through the API's create for its kind; restore refuses a file that stores a value for a parameter its class marks as a credential, such as `api_key`. A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/registry_recipes`. + +A saved file doesn't record who saved it, so the backend rebuilds every entry in this source as an administrator, at startup and after a live apply. A file here can, for example, make the backend run a model's repository code, or use the server's identity on a storage container that the file names. Write access to this source is administrator authority on the backend: grant it only to principals you already trust as backend administrators, and use a dedicated directory or blob prefix that less-trusted users and services cannot write to. + +The backend restores saved entries at startup and after a live apply, once the configured initializers have registered theirs. If any saved entry can't be restored, the backend doesn't become ready (state `failed`, outcome `restore-failed`) and lists each failure with its reason in `GET /api/config/runtime`, so a scan can't quietly run without an entry it expects. Until then, only the configuration and status routes answer; the rest return 503. Fix or delete the listed files, or, when a folder couldn't be listed, check that the backend can read this source. Then apply the configuration, if `enable_live_reinitialization` is on, or restart the backend. Files whose names aren't valid entry names (up to 64 lowercase letters, digits, and underscores, starting with a letter) are skipped without blocking startup; administrators can see them under `restore_ignored` in `GET /api/config/runtime`. + ### `initialization_scripts` Local paths to custom Python scripts containing `PyRITInitializer` subclasses. Paths can be absolute or relative to the current working directory. diff --git a/frontend/e2e/fixtures/manual_send_backend.py b/frontend/e2e/fixtures/manual_send_backend.py index 0c46ed2398..8be9b44941 100644 --- a/frontend/e2e/fixtures/manual_send_backend.py +++ b/frontend/e2e/fixtures/manual_send_backend.py @@ -4,6 +4,7 @@ """Offline backend for manual-send browser tests, with no real configuration or providers.""" import os +import tempfile from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from unittest.mock import AsyncMock, patch @@ -27,19 +28,22 @@ async def _initialize_offline_async(self: ConfigurationLoader, *, raise_on_initi @asynccontextmanager async def _offline_lifespan_async(application: FastAPI) -> AsyncGenerator[None, None]: - config = ConfigurationLoader(memory_db_type="in_memory", env_files=[], env_akv_ref=[]) - with ( - patch.dict(os.environ, {"PYRIT_AUTH_MODE": "local", "PYRIT_DEV_MODE": "true"}, clear=True), - patch.object( - ConfigurationFileService, - "_read_source_async", - new=AsyncMock(return_value="memory_db_type: in_memory\nenv_files: []\ninitializers: []\n"), - ), - patch.object(ConfigurationLoader, "load_with_overrides", return_value=config), - patch.object(ConfigurationLoader, "initialize_pyrit_async", new=_initialize_offline_async), - ): - async with lifespan(application): - yield + with tempfile.TemporaryDirectory() as recipes: + config = ConfigurationLoader( + memory_db_type="in_memory", env_files=[], env_akv_ref=[], registry_recipes_source=recipes + ) + with ( + patch.dict(os.environ, {"PYRIT_AUTH_MODE": "local", "PYRIT_DEV_MODE": "true"}, clear=True), + patch.object( + ConfigurationFileService, + "_read_source_async", + new=AsyncMock(return_value="memory_db_type: in_memory\nenv_files: []\ninitializers: []\n"), + ), + patch.object(ConfigurationLoader, "load_with_overrides", return_value=config), + patch.object(ConfigurationLoader, "initialize_pyrit_async", new=_initialize_offline_async), + ): + async with lifespan(application): + yield app.router.lifespan_context = _offline_lifespan_async diff --git a/frontend/src/components/Configuration/Reinitialize.test.tsx b/frontend/src/components/Configuration/Reinitialize.test.tsx index 78614257fd..4d74590ca0 100644 --- a/frontend/src/components/Configuration/Reinitialize.test.tsx +++ b/frontend/src/components/Configuration/Reinitialize.test.tsx @@ -124,6 +124,26 @@ describe('Reinitialize', () => { expect(screen.getByRole('button', { name: 'Reinitialize PyRIT' })).toBeDisabled() }) + it('should let a failed restore be retried by reinitializing', async () => { + const user = userEvent.setup() + api.getRuntimeStatus.mockResolvedValue({ + ...ready, + state: 'failed', + outcome: 'restore-failed', + message: "1 saved registry entry couldn't be restored.", + }) + render( + + + , + ) + expect(await screen.findByText(/Runtime: failed\. 1 saved registry entry/)).toBeInTheDocument() + await waitFor(() => expect(screen.getByRole('button', { name: 'Reinitialize PyRIT' })).toBeEnabled()) + await user.click(screen.getByRole('button', { name: 'Reinitialize PyRIT' })) + await user.click(within(await screen.findByRole('dialog')).getByRole('button', { name: 'Reinitialize PyRIT' })) + await waitFor(() => expect(api.reinitialize).toHaveBeenCalledWith('saved-v1')) + }) + describe('runtime status polling', () => { beforeEach(() => { jest.useFakeTimers() diff --git a/pyrit/backend/README.md b/pyrit/backend/README.md index 444cd2be87..ab3aae9ee6 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -691,8 +691,9 @@ The chat preserves failed drafts across replacement, but requires converter choi to be reviewed again rather than restoring outputs tied to the previous registry. The offline browser fixture `frontend.e2e.fixtures.manual_send_backend:app` uses the real -backend lifecycle with isolated in-memory SQLite, no environment files, and no default -providers. The recovery tests register deterministic loopback targets. For example, start +backend lifecycle with isolated in-memory SQLite, no environment files, no saved registry +entries, and no default providers. The recovery tests register deterministic loopback +targets. For example, start it with `uv run python -m uvicorn frontend.e2e.fixtures.manual_send_backend:app --host 127.0.0.1 --port 18213` (set `PYTHONUTF8=1` on Windows). In another shell, set `PYRIT_BACKEND_URL=http://127.0.0.1:18213` and `E2E_FRONTEND_PORT=31213`, then run `npx playwright test chat-recovery --project seeded --workers 1` diff --git a/pyrit/backend/middleware/error_handlers.py b/pyrit/backend/middleware/error_handlers.py index 0beaff62cf..6258e6fd05 100644 --- a/pyrit/backend/middleware/error_handlers.py +++ b/pyrit/backend/middleware/error_handlers.py @@ -15,6 +15,7 @@ from pyrit.backend.models import FieldError, ProblemDetail from pyrit.exceptions.analytics_exception import AnalyticsException from pyrit.exceptions.dataset_exception import DatasetAPIException +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError, DocumentUnreadableError logger = logging.getLogger(__name__) @@ -192,6 +193,38 @@ 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: # pyrit-keyword-only-exempt Starlette exception handler (request, exc) + """ + 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. It leaves out the traceback, + since the error behind it can carry a SAS URL. 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})") + if isinstance(exc, DocumentUnreadableError): + detail = f"{exc.label} '{exc.name}' can't be read: {exc.reason}" + else: + detail = f"{exc.label} storage is unavailable" + problem = ProblemDetail( + type="/errors/storage-unavailable", + title="Service Unavailable", + status=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=detail, + instance=str(request.url.path), + ) + + return JSONResponse( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + content=problem.model_dump(exclude_none=True), + ) + @app.exception_handler(NotImplementedError) async def not_implemented_handler( # pyrit-async-suffix-exempt request: Request, diff --git a/pyrit/backend/routes/initializers.py b/pyrit/backend/routes/initializers.py index 50fa79db85..529f92d641 100644 --- a/pyrit/backend/routes/initializers.py +++ b/pyrit/backend/routes/initializers.py @@ -14,7 +14,6 @@ DELETE /api/initializers/{name} — unregister an initializer """ -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from pyrit.backend.middleware.auth import require_admin @@ -35,19 +34,6 @@ router = APIRouter(prefix="/initializers", tags=["initializers"]) -def _custom_storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable custom initializer storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Custom initializer storage is temporarily unavailable", - ) - - def _configured_initializers(request: Request) -> list[ConfiguredInitializerSetting]: """ Read the initializer list captured from ``.pyrit_conf`` at backend startup. @@ -138,10 +124,7 @@ async def list_custom_initializers(request: Request) -> CustomInitializerListRes CustomInitializerListResponse: The configured source and stored scripts. """ _check_custom_initializers_allowed(request) - try: - return await get_initializer_service().list_custom_initializers_async() - except AzureError as exc: - raise _custom_storage_unavailable() from exc + return await get_initializer_service().list_custom_initializers_async() @router.get( @@ -205,8 +188,6 @@ async def register_initializer( # pyrit-async-suffix-exempt try: return await service.register_initializer_async(name=body.name, script_content=body.script_content) - except AzureError as exc: - raise _custom_storage_unavailable() from exc except ValueError as e: detail = str(e) if "already registered" in detail: @@ -243,8 +224,6 @@ async def unregister_initializer( # pyrit-async-suffix-exempt try: await service.unregister_initializer_async(initializer_name=initializer_name) - except AzureError as exc: - raise _custom_storage_unavailable() from exc except ValueError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, diff --git a/pyrit/backend/routes/scenario_presets.py b/pyrit/backend/routes/scenario_presets.py index ad2880dd31..6d727491eb 100644 --- a/pyrit/backend/routes/scenario_presets.py +++ b/pyrit/backend/routes/scenario_presets.py @@ -19,7 +19,6 @@ from typing import Annotated -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, status from pyrit.backend.middleware.auth import current_user_name, require_admin @@ -38,19 +37,6 @@ router = APIRouter(prefix="/scenario-presets", tags=["scenario-presets"]) -def _storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable preset storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Scenario preset storage is temporarily unavailable", - ) - - def _not_found(name: str) -> HTTPException: """ Create a not-found response for a missing preset. @@ -79,8 +65,6 @@ async def _load_preset_or_404_async(name: str) -> ScenarioPresetResponse: """ try: preset = await get_scenario_preset_service().get_preset_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -106,10 +90,7 @@ async def list_scenario_presets( # pyrit-async-suffix-exempt Returns: ScenarioPresetListResponse: The configured source and the stored presets. """ - try: - return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) - except AzureError as exc: - raise _storage_unavailable() from exc + return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) @router.get( @@ -177,8 +158,6 @@ async def create_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{preset.name}' already exists", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.put( @@ -239,8 +218,6 @@ async def update_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{name}' changed since it was read; re-read it and retry", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.delete( @@ -278,8 +255,6 @@ async def delete_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, TimeoutError) as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -329,8 +304,6 @@ async def resolve_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{name}' changed since it was read; re-read it and retry", ) from None - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None diff --git a/pyrit/backend/routes/scenario_suites.py b/pyrit/backend/routes/scenario_suites.py index 42501471c4..8cd0112916 100644 --- a/pyrit/backend/routes/scenario_suites.py +++ b/pyrit/backend/routes/scenario_suites.py @@ -6,7 +6,6 @@ from typing import Annotated from uuid import UUID -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from pyrit.backend.middleware.auth import current_user_name, require_admin @@ -18,6 +17,7 @@ ) from pyrit.backend.models.common import IdentifierStr from pyrit.backend.services import get_scenario_suite_execution_service, get_scenario_suite_service +from pyrit.backend.services.scenario_run_service import SavedEntryLaunchConflictError from pyrit.backend.services.scenario_suite_execution_service import ScenarioSuiteExecutionCapacityError from pyrit.models import ( LaunchScenarioSuiteRequest, @@ -37,19 +37,6 @@ router = APIRouter(prefix="/scenario-suites", tags=["scenario-suites"]) -def _storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable suite storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Scenario suite storage is temporarily unavailable", - ) - - def _not_found(name: str) -> HTTPException: """ Create a not-found response for a missing suite. @@ -72,8 +59,6 @@ async def _load_suite_or_404_async(name: str) -> ScenarioSuiteResponse: """ try: suite = await get_scenario_suite_service().get_suite_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -90,10 +75,7 @@ async def list_scenario_suites() -> ScenarioSuiteListResponse: # pyrit-async-su Returns: ScenarioSuiteListResponse: Suites from both origins with current entry statuses. """ - try: - return await get_scenario_suite_service().list_suites_async() - except AzureError as exc: - raise _storage_unavailable() from exc + return await get_scenario_suite_service().list_suites_async() @router.post("/{name}/readiness", response_model=ScenarioReadiness) @@ -116,8 +98,6 @@ async def get_scenario_suite_readiness_async( raise HTTPException(status_code=409, detail=str(exc)) from None except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from None - except AzureError as exc: - raise _storage_unavailable() from exc @router.get( @@ -173,7 +153,10 @@ async def cancel_scenario_suite_execution( # pyrit-async-suffix-exempt responses={ 400: {"model": ProblemDetail, "description": "The suite has unresolved or invalid entries"}, 404: {"model": ProblemDetail, "description": "Scenario suite not found"}, - 409: {"model": ProblemDetail, "description": "The suite or a referenced preset changed"}, + 409: { + "model": ProblemDetail, + "description": "The suite or a referenced preset changed, or a launch script would replace a saved entry", + }, 429: {"model": ProblemDetail, "description": "Active suite execution capacity is full"}, 503: {"model": ProblemDetail, "description": "Runtime or definition storage unavailable"}, }, @@ -200,7 +183,7 @@ async def launch_scenario_suite( # pyrit-async-suffix-exempt return await get_scenario_suite_execution_service().launch_async(name=name, request=body) except ScenarioSuiteExecutionNotFoundError as exc: raise _not_found(name) from exc - except ScenarioSuiteLaunchConflictError as exc: + except (ScenarioSuiteLaunchConflictError, SavedEntryLaunchConflictError) as exc: raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(exc)) from None except ScenarioSuiteConflictError as exc: raise HTTPException( @@ -211,8 +194,6 @@ async def launch_scenario_suite( # pyrit-async-suffix-exempt raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None except ScenarioSuiteExecutionCapacityError as exc: raise HTTPException(status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=str(exc)) from None - except AzureError as exc: - raise _storage_unavailable() from exc except RuntimeError as exc: raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(exc)) from None @@ -269,8 +250,6 @@ async def create_scenario_suite( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario suite '{suite.name}' already exists", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.put( @@ -319,8 +298,6 @@ async def update_scenario_suite( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario suite '{name}' changed since it was read; re-read it and retry", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.delete( @@ -350,8 +327,6 @@ async def delete_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 None - except (AzureError, TimeoutError) as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index 8ac2cf2d7e..c7339e0146 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -248,7 +248,10 @@ async def load_scenario_datasets_async( responses={ 400: {"model": ProblemDetail, "description": "Invalid request (bad scenario/target/technique)"}, 404: {"model": ProblemDetail, "description": "Saved run not found"}, - 409: {"model": ProblemDetail, "description": "Saved run cannot be resumed"}, + 409: { + "model": ProblemDetail, + "description": "Saved run cannot be resumed, or a launch script would replace a saved entry", + }, }, ) async def start_scenario_run( # pyrit-async-suffix-exempt diff --git a/pyrit/backend/services/runtime_lifecycle.py b/pyrit/backend/services/runtime_lifecycle.py index daef21f0ed..ebeecf30c8 100644 --- a/pyrit/backend/services/runtime_lifecycle.py +++ b/pyrit/backend/services/runtime_lifecycle.py @@ -7,7 +7,8 @@ import logging import os import uuid -from typing import Any +from functools import partial +from typing import TYPE_CHECKING, Any from fastapi import FastAPI @@ -15,6 +16,7 @@ from pyrit.backend.services.analytics_service import get_analytics_service from pyrit.backend.services.configuration_file_service import ConfigurationFileService from pyrit.backend.services.environment_file_service import EnvironmentFileService +from pyrit.backend.services.saved_entry_restore import build_saved_entry_async from pyrit.backend.services.scenario_run_service import get_scenario_run_service, peek_scenario_run_service from pyrit.backend.services.scenario_suite_execution_service import peek_scenario_suite_execution_service from pyrit.backend.services.service_lifecycle import ( @@ -26,10 +28,14 @@ from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH from pyrit.memory import CentralMemory from pyrit.registry import InitializerRegistry, ScenarioPresetRegistry, ScenarioSuiteRegistry +from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer from pyrit.setup import ConfigurationLoader from pyrit.setup.environment_loading import resolve_environment_async from pyrit.setup.initialization import validate_reinitialization_memory +if TYPE_CHECKING: + from pyrit.models.catalog import IgnoredRecipeFile, RestoreFailure + logger = logging.getLogger(__name__) @@ -46,6 +52,8 @@ def __init__(self, *, app: FastAPI, source: ConfigurationFileService) -> None: self.version: str | None = None self.outcome = "starting" self.message = "" + self.restore_failures: list[RestoreFailure] = [] + self.restore_ignored: list[IgnoredRecipeFile] = [] self.operations: set[asyncio.Task[None]] = set() self.management_operations: set[asyncio.Task[None]] = set() self.apply_task: asyncio.Task[None] | None = None @@ -75,6 +83,8 @@ def status(self) -> dict[str, Any]: "message": self.message, "enabled": self.topology_supported, "applying": self.apply_task is not None and not self.apply_task.done(), + "restore_failures": [failure.model_dump(mode="json") for failure in self.restore_failures], + "restore_ignored": [file.model_dump(mode="json") for file in self.restore_ignored], } async def _load_async(self) -> ConfigurationLoader: @@ -119,6 +129,52 @@ def _publish(self, config: ConfigurationLoader) -> None: self.generation = str(uuid.uuid4()) self.state, self.outcome, self.message = "ready", "success", "PyRIT is ready." + async def _restore_saved_entries_async(self, config: ConfigurationLoader) -> bool: + """ + Rebuild saved registry entries after the startup scripts have registered theirs. + + A failure keeps the runtime not ready, since a scan that would have used a missing entry + could quietly fall back to a different one. Fixing or deleting the stored file, then applying + or restarting, recovers. + + Args: + config (ConfigurationLoader): The loaded configuration naming the storage source. + + Returns: + bool: Whether every saved entry was restored. + """ + restorer = RegistryRecipeRestorer(source=config.registry_recipes_source) + # Reading can wait on storage, so it runs in a worker thread. Entries are built on this loop, through + # the API's create for each kind, because a constructor such as HuggingFaceChatTarget's starts a task + # on it. A slow constructor holds the loop while it runs. + saved = await asyncio.to_thread(restorer.read) + report = await restorer.restore_async( + saved=saved, build_async=partial(build_saved_entry_async, upload_directory=config.target_upload_directory) + ) + self.restore_failures, self.restore_ignored = report.failures, report.ignored + if not report.failures: + return True + self.state, self.outcome = "failed", "restore-failed" + next_step = ( + "apply the configuration or restart the backend" + if config.enable_live_reinitialization and self.topology_supported + else "restart the backend" + ) + unread = sorted({failure.kind.value for failure in report.failures if failure.name is None}) + if unread: + self.message = ( + f"The saved {' and '.join(unread)} couldn't be read, so no saved entries were restored. " + f"Check that the backend can read registry_recipes_source, then {next_step}." + ) + else: + count = len(report.failures) + self.message = ( + f"{count} saved registry {'entry' if count == 1 else 'entries'} couldn't be restored. " + f"Fix or delete the files named in the server log or GET /api/config/runtime, then {next_step}." + ) + logger.error("PyRIT isn't ready: %s", self.message) + return False + async def startup_async(self) -> None: """Keep management alive even when parsing or initialization fails.""" self.app.state.allow_custom_initializers = False @@ -137,7 +193,8 @@ async def startup_async(self) -> None: await config.initialize_pyrit_async(raise_on_initializer_error=True) await get_scenario_run_service().reconcile_interrupted_runs_async() _, self.version = await self.source.read_with_version_async() - self._publish(config) + if await self._restore_saved_entries_async(config): + self._publish(config) except Exception: self.state, self.outcome = "restart-required", "initialization-failed" self.message = "PyRIT setup failed. Repair saved configuration, then restart the backend." @@ -158,7 +215,7 @@ def begin_apply(self, *, version: str) -> dict[str, Any]: "outcome": "unsupported", "message": ("Reinitialization requires one backend worker and one replica."), } - if self.state in ("failed", "restart-required"): + if self.state == "restart-required": return {**self.status(), "outcome": "restart-required"} if (self.apply_task and not self.apply_task.done()) or self.edit_lock.locked() or self.management_operations: return {**self.status(), "outcome": "busy"} @@ -179,7 +236,9 @@ async def _apply_async(self, *, version: str) -> None: if not config.enable_live_reinitialization: self.outcome = "unsupported" self.message = ( - "Set enable_live_reinitialization: true in the saved configuration to use live apply." + "Live apply is off. Restart the backend to restore the saved entries again." + if previous_state == "failed" + else "Set enable_live_reinitialization: true in the saved configuration to use live apply." ) return # Type compatibility is checked before contacting environment sources. @@ -219,7 +278,8 @@ async def _apply_async(self, *, version: str) -> None: await config.apply_prepared_reinitialization_async(prepared=prepared) await self._management_async(config) self.version = current_version - self._publish(config) + if await self._restore_saved_entries_async(config): + self._publish(config) except (Exception, asyncio.CancelledError, BaseExceptionGroup) as error: logger.exception( "PyRIT live apply failed phase=%s generation=%s", diff --git a/pyrit/backend/services/saved_entry_restore.py b/pyrit/backend/services/saved_entry_restore.py new file mode 100644 index 0000000000..a4ac353988 --- /dev/null +++ b/pyrit/backend/services/saved_entry_restore.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Build saved registry entries through the API's own create for each kind.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from pydantic import ValidationError + +from pyrit.backend.models import CreateConverterRequest, CreateScorerRequest, CreateTargetRequest +from pyrit.backend.models.techniques import CreateTechniqueRequest +from pyrit.backend.services.converter_service import get_converter_service +from pyrit.backend.services.scorer_service import get_scorer_service +from pyrit.backend.services.target_service import get_target_service +from pyrit.backend.services.technique_service import get_technique_service +from pyrit.models.catalog import RegistryRecipeKind, RestoreFailureReason, TechniqueRecipe +from pyrit.registry.registry_recipe_restore import RecipeRejectedError + +if TYPE_CHECKING: + from pyrit.models.catalog import StoredRecipe + + +async def build_saved_entry_async(stored: StoredRecipe, *, upload_directory: str | None) -> None: + """ + Build one saved entry with the create the API uses for its kind, and register it under its name. + + The entry passes the same checks a create request does. Restore runs as the server, and a saved file + doesn't record who saved it, so a target is built as an administrator would build it. A target that + uploads files gets this server's upload folder, as it does through the API. + + Args: + stored (StoredRecipe): The checked recipe to build. + upload_directory (str | None): The folder this server lets targets upload files from, or None. + + Raises: + RecipeRejectedError: If the recipe doesn't make a valid create request. + """ + request = _create_request(stored) + if isinstance(request, CreateTechniqueRequest): + await get_technique_service().create_async(request) + elif isinstance(request, CreateTargetRequest): + await get_target_service().create_target_async( + request=request, upload_directory=upload_directory, is_admin=True + ) + elif isinstance(request, CreateConverterRequest): + await get_converter_service().create_converter_async(request=request) + else: + await get_scorer_service().create_scorer_async(request=request) + + +def _create_request( + stored: StoredRecipe, +) -> CreateTargetRequest | CreateConverterRequest | CreateScorerRequest | CreateTechniqueRequest: + """ + Turn a saved recipe back into the create request it was saved from. + + Returns: + CreateTargetRequest | CreateConverterRequest | CreateScorerRequest | CreateTechniqueRequest: The request. + + Raises: + RecipeRejectedError: If the recipe doesn't make a valid request, such as a technique with a reserved name. + """ + recipe = stored.recipe + fields: dict[str, Any] = {"name": stored.ref.name} + try: + if isinstance(recipe, TechniqueRecipe): + return CreateTechniqueRequest.model_validate( + fields | recipe.model_dump(exclude_unset=True, exclude={"schema_version"}) + ) + fields |= {"type": recipe.class_name, "params": recipe.params} + if stored.ref.kind is RegistryRecipeKind.TARGETS: + if recipe.auth_mode is not None: + fields["auth_mode"] = recipe.auth_mode + return CreateTargetRequest.model_validate(fields) + if stored.ref.kind is RegistryRecipeKind.CONVERTERS: + return CreateConverterRequest.model_validate(fields) + return CreateScorerRequest.model_validate(fields) + except ValidationError as error: + problems = "; ".join( + f"{'.'.join(map(str, issue['loc'])) or 'request'}: {issue['msg']}" for issue in error.errors() + ) + raise RecipeRejectedError( + reason=RestoreFailureReason.INVALID_RECIPE, message=f"Not a valid create request: {problems}" + ) from error diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 1b50f68e5f..e5d5a3e6fc 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -103,6 +103,7 @@ from pyrit.models.scenario_readiness import ScenarioConfigurationError from pyrit.prompt_target import PromptTarget from pyrit.registry import InitializerRegistry, ScenarioRegistry +from pyrit.registry.instance_registry import SavedEntryReplacementError, record_saved_entry_refusals from pyrit.registry.resolution import reject_non_external_params, resolve_declared_params from pyrit.scenario import IncompatibleTechniquePolicy, Scenario from pyrit.scenario.core import get_default_adversarial_target, override_default_adversarial_target @@ -192,6 +193,10 @@ class ScenarioRunConflictError(ValueError): """A saved run cannot be admitted to the scheduler in its current state.""" +class SavedEntryLaunchConflictError(ScenarioRunConflictError): + """A launch script would replace an entry rebuilt from a saved recipe.""" + + class ScenarioRunNotFoundError(ValueError): """The requested saved run is not in the active memory database.""" @@ -1657,6 +1662,7 @@ async def _run_initializers_async(self, *, request: RunScenarioRequest) -> None: Raises: ValueError: If an initializer name is not found in the registry. + SavedEntryLaunchConflictError: If an initializer would replace a saved entry. """ if not request.initializers: return @@ -1670,7 +1676,42 @@ async def _run_initializers_async(self, *, request: RunScenarioRequest) -> None: ) except KeyError as e: raise ValueError(f"Initializer not found: {e}") from None - await instance.initialize_async() + # The registries refuse to replace a saved entry; the record also catches a refusal + # that the script itself caught, which would leave the saved entry in its place. + with record_saved_entry_refusals() as refusals: + try: + await instance.initialize_async() + except SavedEntryReplacementError as error: + refusals.append((error.kind, error.name)) + except Exception as error: + # A script can wrap the refusal in an error of its own; the clash is still why it stopped. + if refusals: + raise self._saved_entry_conflict( + initializer_name=initializer_name, refusals=refusals + ) from error + raise + if refusals: + raise self._saved_entry_conflict(initializer_name=initializer_name, refusals=refusals) + + @staticmethod + def _saved_entry_conflict( + *, initializer_name: str, refusals: list[tuple[str, str]] + ) -> SavedEntryLaunchConflictError: + """ + Describe the saved entries a launch script would have replaced. + + Args: + initializer_name (str): The launch script's name. + refusals (list[tuple[str, str]]): The kind and name of each saved entry that was refused. + + Returns: + SavedEntryLaunchConflictError: The error to raise. + """ + entries = ", ".join(f"{kind.removesuffix('s')} '{name}'" for kind, name in dict.fromkeys(refusals)) + return SavedEntryLaunchConflictError( + f"The launch script '{initializer_name}' would replace the saved {entries}. " + "Leave that script out of the launch." + ) async def _initialize_scenario_async( self, diff --git a/pyrit/backend/services/scenario_suite_service.py b/pyrit/backend/services/scenario_suite_service.py index 65cf984e94..ce1e436c2b 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__) @@ -217,12 +216,9 @@ def _resolve_preset_projection( ) stored_version = registry.get_preset_version(name=preset_name) - except (AzureError, OSError): - logger.warning( - "Could not project scenario preset '%s' because preset storage is unavailable.", - preset_name, - exc_info=True, - ) + except DocumentStorageUnavailableError as error: + # The message leaves out credentials; the error behind it, which a traceback would show, may not. + logger.warning(f"Could not project scenario preset '{preset_name}': {error}") return _PresetProjection( origin=None, version=None, diff --git a/pyrit/common/azure_storage.py b/pyrit/common/azure_storage.py index d3ed6eeb4f..dfce6b2e08 100644 --- a/pyrit/common/azure_storage.py +++ b/pyrit/common/azure_storage.py @@ -3,8 +3,16 @@ """Lightweight Azure Storage URI helpers.""" +import re from urllib.parse import parse_qs, urlparse +# A scheme followed by "://" and everything up to whitespace, so a quote can't end a URL before its query. +# A scheme starts only where one can't continue from the character before it, such as after "_" or a +# space, so each run of scheme characters is scanned once and long text stays fast. +_URL_IN_TEXT = re.compile(r"(? bool: """ @@ -48,9 +56,30 @@ def has_sas_signature(value: str) -> bool: def redact_url_credentials(value: str) -> str: """ - Remove query parameters and fragments from a URL for display or logging. + Remove the user name and password, query parameters, and fragment from a URL for display or logging. + + Returns: + str: The URL without them. + """ + parsed = urlparse(value) + return parsed._replace(netloc=parsed.netloc.rpartition("@")[2], query="", fragment="").geturl() + + +def redact_urls(text: str) -> str: + """ + Remove the user name and password, query parameters, and fragment from every URL in a text. + + A query or user name can hold a credential, such as a SAS token. A URL that can't be parsed is + hidden entirely, and a SAS signature is hidden wherever it appears. Returns: - str: The URL without query parameters or a fragment. + str: The text with each URL redacted. """ - return urlparse(value)._replace(query="", fragment="").geturl() + + def redact(match: re.Match[str]) -> str: + try: + return redact_url_credentials(match.group(0)) + except ValueError: + return "