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..af4d551a4f 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. Targets created through `POST /api/targets` with `"persist": true` are saved here. 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. A file never holds a credential: restore refuses a file that stores a value for a parameter its class marks as a credential, such as `api_key`. A saved target also can't use a key from this server: it signs in with Microsoft Entra ID (`"auth_mode": "identity"`), which uses this server's Azure identity for everyone using the app, or needs no key. Targets from the server's configuration that it uses, such as a round-robin's members, keep their own sign-in. So restore refuses a target file that would read an empty key from an environment variable or leave an address to this server's settings. The addresses each target type declares (such as `endpoint`, `container_url` for Blob storage, or an A2A agent card path) can't hold a user name, password, SAS token, fragment, path parameters (`;`, `&`, `=`), or a query other than `api-version`, and a path can't have a query at all. An address the server would otherwise fill in must be a full http(s) or ws(s) address. In identity mode, an OpenAI target file sets `"headers": "{}"` in its `params` so that the server's `OPENAI_ADDITIONAL_REQUEST_HEADERS` aren't added. The OpenAI SDK adds the headers in `OPENAI_CUSTOM_HEADERS` to every request, so while the server sets it, OpenAI targets can't be saved and saved ones aren't restored. The SDK's `OPENAI_ORG_ID` and `OPENAI_PROJECT_ID` still apply to every OpenAI target on this server. GitHub Copilot, Hugging Face, HackAPrompt, LiteLLM, HTTP, and Copilot WebSocket targets can't be saved yet. A SAS query string may be included in the source; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/registry_recipes`. + +Restore can't tell who wrote a saved file, even one that names the user 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/playwright.config.ts b/frontend/playwright.config.ts index 23140b23a0..547d6b29e3 100644 --- a/frontend/playwright.config.ts +++ b/frontend/playwright.config.ts @@ -1,3 +1,7 @@ +import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import path from "node:path"; + import { defineConfig, devices } from "@playwright/test"; import type { UserPreferences } from "./src/types"; @@ -24,6 +28,28 @@ const VITE_SERVER_OPTIONS = const E2E_BACKEND_PORT = process.env.PYRIT_E2E_BACKEND_PORT ?? "18000"; const E2E_BACKEND_URL = `http://127.0.0.1:${E2E_BACKEND_PORT}`; +/** + * Copies the seeded backend's configuration into a fresh folder that also holds + * its saved targets, converters, scorers, and techniques, so a run never reads + * or writes the ones saved on the machine running it. + */ +function createSeededBackendConfigFile(): string { + const folder = mkdtempSync(path.join(tmpdir(), "pyrit-seeded-backend-")); + // The runner stops the backend before it exits, so nothing still uses the folder by then. + process.once("exit", () => rmSync(folder, { recursive: true, force: true })); + const configFile = path.join(folder, "browser_config.yaml"); + const config = readFileSync(new URL("../tests/end_to_end/browser_config.yaml", import.meta.url), "utf-8"); + const recipesSource = JSON.stringify(path.join(folder, "registry_recipes")); + writeFileSync(configFile, `${config}\nregistry_recipes_source: ${recipesSource}\n`); + return configFile; +} + +// Workers load this file too, but only the main process starts the backend. +const SEEDED_BACKEND_CONFIG_FILE = + SEEDED_MODE && process.env.TEST_WORKER_INDEX === undefined + ? createSeededBackendConfigFile() + : ""; + export default defineConfig({ testDir: "./e2e", testIgnore: ["**/analytics.spec.ts", "**/paired-outcomes.spec.ts", "**/scenario-hierarchy.spec.ts", "**/scenario-datasets.spec.ts", "**/scenario-techniques.spec.ts", "**/scan-launch.seeded.spec.ts"], @@ -113,7 +139,7 @@ export default defineConfig({ command: `cd .. && uv run --no-sync python -m pyrit.backend.pyrit_backend ` + `--auth-mode local --host 127.0.0.1 --port ${E2E_BACKEND_PORT} --log-level warning ` + - "--config-file tests/end_to_end/browser_config.yaml", + `--config-file "${SEEDED_BACKEND_CONFIG_FILE}"`, env: { PYRIT_DEV_MODE: "true" }, url: `${E2E_BACKEND_URL}/api/health`, reuseExistingServer: false, 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..b9dd3d3d9e 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -54,6 +54,7 @@ its database collation. ### Targets - `GET /api/targets` - List available prompt targets - `GET /api/targets/{id}` - Get target details +- `POST /api/targets` - Create a target - `POST /api/targets/{id}/detect-capabilities` - Run live capability checks and apply the result to the registered target. Send `{"expected_hash": ""}`. Detection uses a private target, has a four-minute operation deadline, and @@ -82,6 +83,38 @@ overrides; reconstruction preserves declarations from the source target. Per-send values use [`request_generation_config` metadata](../../doc/code/targets/0_prompt_targets.md#runtime-generation-configuration). +A created target is Temporary: a restart or a configuration apply removes it. With `"persist": true` (and `register` +left on), the target is saved under `registry_recipes_source` and restored at every start, and the signed-in user is +recorded as the one who saved it. A saved target can't hold or use a key: it signs in with Microsoft Entra ID +(`"auth_mode": "identity"`, through this server's Azure identity) or needs no key. An OpenAI target's empty `headers` +is set to `"{}"` first, so a saved one doesn't send this server's `OPENAI_ADDITIONAL_REQUEST_HEADERS`. + +| A request with `persist` that... | Gets | +|---|---| +| names the target with anything but up to 64 lowercase letters, digits, and underscores, starting with a letter | **422** | +| would save a key, or use one of this server's keys or address settings | **400** | +| gives an address that can carry a credential, or, where this server's setting would apply, one that isn't a full http(s) or ws(s) address | **400** | +| saves a target type that can't be saved, or one that a scan's launch scripts added or put in place of another | **400** | +| saves an OpenAI target while this server sets `OPENAI_CUSTOM_HEADERS`, whose headers the OpenAI SDK adds to every request | **400** | +| uses a Temporary target | **400** | +| takes a name that a target or a saved file already has | **409** | +| can't be written to storage | **503** | + +A request refused with **400** can be sent again without `persist`. When storage can't be written, the server log names +the target and its folder, since the target comes back at the next start if the write reached storage anyway. + +Each listed target's `origin` is `saved`, `configuration` (registered by the server's startup scripts), or `temporary` +(gone after a restart or a configuration apply, such as a target created without `persist` or registered by a scan's +launch scripts). A saved target can only use targets that come back after a restart; targets from configuration keep +their own sign-in. + +A saved target records the create request without its name: the target type, `auth_mode`, and the parameters as the +API accepted them. Restore sends it through the same create, so it passes the same checks and follows the class's +constructor. A class that moves within PyRIT still restores. Renaming or removing a class or a constructor parameter +changes the saved format: such a change raises `schema_version` and adds a migration for older files, with a test that +restores one. A file that still doesn't match fails to restore with the reason, and the backend stays not ready until +the file is fixed or deleted, so nothing runs without the saved target. + ### Dataset creation, batch mutation, and deletion `POST /api/datasets` creates a named dataset synchronously from multipart form data. @@ -691,8 +724,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/mappers/target_mappers.py b/pyrit/backend/mappers/target_mappers.py index ef8f1aaa2d..3c424ca3dc 100644 --- a/pyrit/backend/mappers/target_mappers.py +++ b/pyrit/backend/mappers/target_mappers.py @@ -15,7 +15,7 @@ from typing import Any from pyrit.models import CapabilityName, TargetIdentifier -from pyrit.models.catalog import TargetInstance +from pyrit.models.catalog import RegistryEntryOrigin, TargetInstance from pyrit.prompt_target import PromptTarget # Capability flag names that should never be surfaced as identifier-level params: @@ -52,7 +52,9 @@ def _target_specific_params(target_identifier: TargetIdentifier) -> dict[str, An return {**extra, **explicit_specific} or None -def target_object_to_instance(target_registry_name: str, target_obj: PromptTarget) -> TargetInstance: +def target_object_to_instance( + target_registry_name: str, target_obj: PromptTarget, *, origin: RegistryEntryOrigin | None = None +) -> TargetInstance: """ Build a TargetInstance DTO from a registry target object. @@ -63,6 +65,7 @@ def target_object_to_instance(target_registry_name: str, target_obj: PromptTarge Args: target_registry_name: The human-friendly target registry name. target_obj: The domain PromptTarget object from the registry. + origin: Where the registered target comes from, or None for an inner target. Returns: TargetInstance DTO with metadata derived from the object. @@ -83,6 +86,7 @@ def target_object_to_instance(target_registry_name: str, target_obj: PromptTarge supports_history_replay=supports_history_replay, target_specific_params=_target_specific_params(target_identifier), inner_targets=_build_inner_targets(target_obj), + origin=origin, ) 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/models/targets.py b/pyrit/backend/models/targets.py index d4ecb93c4a..efc6f24976 100644 --- a/pyrit/backend/models/targets.py +++ b/pyrit/backend/models/targets.py @@ -12,7 +12,14 @@ from pyrit.backend.models.common import MAX_ITEMS, REGISTRY_INSTANCE_NAME_PATTERN, IdentifierStr, PaginationInfo from pyrit.common import AuthMode -from pyrit.models import GenerationRange, JSONValue, Parameter, TargetCapabilities, TargetIdentifier +from pyrit.models import ( + GenerationRange, + JSONValue, + Parameter, + TargetCapabilities, + TargetIdentifier, + validate_registry_name, +) from pyrit.models.catalog import TargetInstance from pyrit.models.component_spec import SourceInstanceSpec @@ -75,6 +82,14 @@ class CreateTargetRequest(BaseModel): description="Unique registry name; omitted only for legacy UI compatibility", ) register: bool = Field(True, description="Register the object; false returns only its descriptor") + persist: bool = Field( + False, + description=( + "Save the target so it comes back after the backend restarts or is reinitialized (true), or keep it " + "only until then (false). A saved target needs a name of lowercase letters, digits, and underscores, " + "and can't hold or use a key: it authenticates with Microsoft Entra ID ('identity') or needs no key." + ), + ) source: SourceInstanceSpec | None = None type: IdentifierStr = Field(..., description="Target type (e.g., 'OpenAIChatTarget')") params: dict[IdentifierStr, JSONValue] = Field( @@ -98,6 +113,17 @@ def _validate_source(self) -> "CreateTargetRequest": raise ValueError("source requires register=false and overrides, including auth_mode, in source.params") return self + @model_validator(mode="after") + def _validate_persist(self) -> "CreateTargetRequest": + if not self.persist: + return self + if not self.register: + raise ValueError("persist requires register=true") + if self.name is None: + raise ValueError("persist requires a name") + validate_registry_name(self.name) + return self + class UnregisteredTarget(BaseModel): """A constructed target descriptor with no registry key or object handle.""" 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..79ea691149 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 @@ -37,19 +36,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 +58,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 +74,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 +97,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( @@ -211,8 +190,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 +246,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 +294,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 +323,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/routes/targets.py b/pyrit/backend/routes/targets.py index 1acd074d6f..d56f6fd30e 100644 --- a/pyrit/backend/routes/targets.py +++ b/pyrit/backend/routes/targets.py @@ -10,7 +10,7 @@ from fastapi import APIRouter, HTTPException, Query, Request, status -from pyrit.backend.middleware.auth import has_admin_access +from pyrit.backend.middleware.auth import AuthenticatedUser, has_admin_access from pyrit.backend.models import ( CreateTargetRequest, ProblemDetail, @@ -23,7 +23,8 @@ UnregisteredTarget, ) from pyrit.backend.services import get_target_service -from pyrit.models.catalog import TargetInstance +from pyrit.backend.services.target_service import SavedTargetConflictError, SavedTargetStorageError +from pyrit.models.catalog import SavedEntryOwner, TargetInstance from pyrit.prompt_target.common.discover_target_capabilities import CapabilityDiscoveryError router = APIRouter(prefix="/targets", tags=["targets"]) @@ -81,6 +82,22 @@ def _target_upload_directory(request: Request) -> str | None: return directory +def _saved_entry_owner(request: Request) -> SavedEntryOwner | None: + """ + Return the signed-in user to record on a saved target. + + Args: + request: The incoming FastAPI request. + + Returns: + SavedEntryOwner | None: The user's object ID and name, or None when sign-in is off. + """ + user = getattr(request.state, "user", None) + if not isinstance(user, AuthenticatedUser): + return None + return SavedEntryOwner(oid=user.oid, name=user.name or user.email) + + @router.get( "", response_model=TargetListResponse, @@ -136,6 +153,11 @@ async def list_target_types(request: Request) -> TargetTypeResponse: # pyrit-as "description": "Administrator access is required to run model repository code or to use the " "server's stored credentials", }, + 409: { + "model": ProblemDetail, + "description": "A target to save has the name of a registered or saved target", + }, + 503: {"model": ProblemDetail, "description": "Saved targets can't be written right now"}, }, ) async def create_target( @@ -161,7 +183,13 @@ async def create_target( request=body, upload_directory=_target_upload_directory(request), is_admin=has_admin_access(request), + owner=_saved_entry_owner(request) if body.persist else None, + recipes_source=getattr(request.app.state, "registry_recipes_source", None), ) + except SavedTargetConflictError as e: + raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(e)) from e + except SavedTargetStorageError as e: + raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(e)) from e except PermissionError as e: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, diff --git a/pyrit/backend/services/runtime_lifecycle.py b/pyrit/backend/services/runtime_lifecycle.py index daef21f0ed..efe118884a 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 ( @@ -25,11 +27,17 @@ ) from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH from pyrit.memory import CentralMemory +from pyrit.models.catalog import RegistryEntryOrigin from pyrit.registry import InitializerRegistry, ScenarioPresetRegistry, ScenarioSuiteRegistry +from pyrit.registry.instance_registry import registration_origin +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 +54,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 +85,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: @@ -96,6 +108,7 @@ async def _management_async(self, config: ConfigurationLoader) -> None: ) self.app.state.allow_custom_initializers = config.allow_custom_initializers self.app.state.target_upload_directory = config.target_upload_directory + self.app.state.registry_recipes_source = config.registry_recipes_source registry = await asyncio.to_thread(InitializerRegistry.get_registry_singleton) registry.configure_custom_scripts_source(config.custom_initializers_source) preset_registry = ScenarioPresetRegistry.get_registry_singleton() @@ -119,10 +132,57 @@ 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 self.app.state.target_upload_directory = None + self.app.state.registry_recipes_source = None self.app.state.environment_file_service = None self.app.state.configured_initializers = [] self.app.state.default_labels = {} @@ -134,10 +194,14 @@ async def startup_async(self) -> None: if config.allow_custom_initializers: logger.warning("Custom initializer registration is ENABLED (allow_custom_initializers: true).") await asyncio.to_thread(registry.register_stored_initializers, strict=True) - await config.initialize_pyrit_async(raise_on_initializer_error=True) + # Entries the startup scripts register are marked as configuration, so they can be told apart + # from entries that are gone after a restart, such as ones a scan's launch scripts register. + with registration_origin(RegistryEntryOrigin.CONFIGURATION): + 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 +222,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 +243,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. @@ -216,10 +282,12 @@ async def _apply_async(self, *, version: str) -> None: mutated = True await close_services_async() - await config.apply_prepared_reinitialization_async(prepared=prepared) + with registration_origin(RegistryEntryOrigin.CONFIGURATION): + 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..ac89699d36 --- /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 can't tell who + wrote a saved file, 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..c46c44e5be 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -98,11 +98,22 @@ TargetIdentifier, project_behavioral_identity, ) -from pyrit.models.catalog import AttackErrorSummary, AttackRetrySummary, RunScenarioRequest, ScenarioRunSummary +from pyrit.models.catalog import ( + AttackErrorSummary, + AttackRetrySummary, + RegistryEntryOrigin, + RunScenarioRequest, + ScenarioRunSummary, +) from pyrit.models.catalog.scenario import ScenarioOverloadSummary, ScenarioTargetSummary 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, + registration_origin, +) 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 +203,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 +1672,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 +1686,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, registration_origin(RegistryEntryOrigin.TEMPORARY): + 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_execution_service.py b/pyrit/backend/services/scenario_suite_execution_service.py index 66f3bbe968..87a63fe126 100644 --- a/pyrit/backend/services/scenario_suite_execution_service.py +++ b/pyrit/backend/services/scenario_suite_execution_service.py @@ -14,7 +14,11 @@ from uuid import UUID, uuid4 from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver -from pyrit.backend.services.scenario_run_service import ScenarioRunService, get_scenario_run_service +from pyrit.backend.services.scenario_run_service import ( + SavedEntryLaunchConflictError, + ScenarioRunService, + get_scenario_run_service, +) from pyrit.models import ( LaunchScenarioSuiteRequest, ScenarioRunState, @@ -319,7 +323,7 @@ async def _submit_entries_async(self, *, state: _ExecutionState) -> None: ) async with state.lock: entry.state = ScenarioSuiteSubmissionState.SUBMISSION_FAILED - entry.error = self.SUBMISSION_FAILURE_MESSAGE + entry.error = self._describe_submission_failure(exc) entry.error_type = type(exc).__name__ state.updated_at = datetime.now(UTC) continue @@ -348,7 +352,7 @@ async def _submit_entries_async(self, *, state: _ExecutionState) -> None: except Exception as exc: async with state.lock: entry.state = ScenarioSuiteSubmissionState.SUBMISSION_FAILED - entry.error = self.SUBMISSION_FAILURE_MESSAGE + entry.error = self._describe_submission_failure(exc) entry.error_type = type(exc).__name__ state.updated_at = datetime.now(UTC) else: @@ -605,6 +609,21 @@ def _get_run_service(self) -> ScenarioRunService: """Return the injected run service or process singleton.""" return self._run_service if self._run_service is not None else get_scenario_run_service() + def _describe_submission_failure(self, error: Exception) -> str: + """ + Describe why an entry's run couldn't be created, without exposing internal error text. + + Args: + error (Exception): The error the run service raised. + + Returns: + str: The message to show for the entry. + """ + # A clash with a saved entry names only the launch script and the entry, so it's safe to show. + if isinstance(error, SavedEntryLaunchConflictError): + return str(error) + return self.SUBMISSION_FAILURE_MESSAGE + _service_instance: ScenarioSuiteExecutionService | None = None 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/backend/services/target_service.py b/pyrit/backend/services/target_service.py index a3d413e411..812b0e56a1 100644 --- a/pyrit/backend/services/target_service.py +++ b/pyrit/backend/services/target_service.py @@ -15,7 +15,7 @@ import asyncio import logging import uuid -from collections.abc import Generator +from collections.abc import Generator, Mapping, Sequence from contextlib import contextmanager from functools import lru_cache from typing import Any @@ -34,8 +34,14 @@ ) from pyrit.backend.services.component_lifecycle import construct_component_async, release_component_async from pyrit.common import AUTH_MODES, REQUIRED_VALUE, AuthMode, default_values -from pyrit.models import Parameter, PromptDataType, TargetCapabilities, TargetIdentifier -from pyrit.models.catalog import TargetInstance +from pyrit.models import ComponentType, Parameter, PromptDataType, TargetCapabilities, TargetIdentifier +from pyrit.models.catalog import ( + ComponentRecipe, + RegistryEntryOrigin, + RegistryRecipeKind, + SavedEntryOwner, + TargetInstance, +) from pyrit.models.component_spec import SourceInstanceSpec, TargetBinding from pyrit.prompt_target import ( GitHubCopilotTarget, @@ -49,6 +55,10 @@ ) from pyrit.prompt_target.common.tool_call_history import TOOL_CALL_INPUT_MODALITIES from pyrit.registry import TargetRegistry +from pyrit.registry.file_document_storage import DocumentConflictError, DocumentStorageUnavailableError +from pyrit.registry.instance_registry import SAVED_ENTRY_METADATA_KEY, RegistryNameTakenError, registry_entry_origin +from pyrit.registry.registry_recipe_rules import SavedEntryRefusedError, check_saved_entry_params +from pyrit.registry.registry_recipe_storage import RegistryRecipeStorage logger = logging.getLogger(__name__) @@ -65,6 +75,14 @@ _AUTH_MODE_PARAM = "auth_mode" +class SavedTargetConflictError(ValueError): + """A target can't be saved under a name that a registered or saved target already has.""" + + +class SavedTargetStorageError(RuntimeError): + """Saved targets can't be written right now.""" + + class TargetService: """ Service for managing target instances. @@ -199,14 +217,16 @@ async def _probe_capabilities_async(self, target: PromptTarget) -> TargetCapabil finally: await release_component_async(probe) - def _build_instance_from_object(self, *, target_registry_name: str, target_obj: Any) -> TargetInstance: + def _build_instance_from_object( + self, *, target_registry_name: str, target_obj: Any, origin: RegistryEntryOrigin | None = None + ) -> TargetInstance: """ Build a TargetInstance from a registry object. Returns: TargetInstance with metadata derived from the object. """ - return target_object_to_instance(target_registry_name, target_obj) + return target_object_to_instance(target_registry_name, target_obj, origin=origin) async def list_targets_async( self, @@ -225,7 +245,9 @@ async def list_targets_async( TargetListResponse containing paginated targets. """ items = [ - self._build_instance_from_object(target_registry_name=entry.name, target_obj=entry.instance) + self._build_instance_from_object( + target_registry_name=entry.name, target_obj=entry.instance, origin=registry_entry_origin(entry.metadata) + ) for entry in self._registry.instances.get_all_instances() ] page, has_more = self._paginate(items=items, cursor=cursor, limit=limit) @@ -266,10 +288,14 @@ async def get_target_async(self, *, target_registry_name: str) -> TargetInstance Returns: TargetInstance if found, None otherwise. """ - obj = self._registry.instances.get(target_registry_name) - if obj is None: + entry = self._registry.instances.get_entry(target_registry_name) + if entry is None: return None - return self._build_instance_from_object(target_registry_name=target_registry_name, target_obj=obj) + return self._build_instance_from_object( + target_registry_name=target_registry_name, + target_obj=entry.instance, + origin=registry_entry_origin(entry.metadata), + ) def get_target_object(self, *, target_registry_name: str) -> Any | None: """ @@ -508,7 +534,13 @@ def _apply_upload_directory( params[parameter_name] = upload_directory async def create_target_async( - self, *, request: CreateTargetRequest, upload_directory: str | None = None, is_admin: bool = False + self, + *, + request: CreateTargetRequest, + upload_directory: str | None = None, + is_admin: bool = False, + owner: SavedEntryOwner | None = None, + recipes_source: str | None = None, ) -> TargetInstance | UnregisteredTarget: """ Create a new target instance from API request. @@ -531,7 +563,8 @@ async def create_target_async( repository code, or that would authenticate with this server's stored credentials, is built only for an administrator. The response is built before the target is registered, so a failed request leaves no registered - target. + target. With ``persist``, the target is saved so it is restored at every + start; otherwise it is Temporary. Args: request: The create target request with type, params, and auth_mode. @@ -539,6 +572,8 @@ async def create_target_async( or None when the server configures none. is_admin: Whether the caller may create a target that runs model repository code or authenticates with this server's stored credentials. + owner: The signed-in user, recorded on a saved target, or None when sign-in is off. + recipes_source: Where saved entries are stored, or None for the default folder. Returns: TargetInstance with the new target's details. @@ -548,11 +583,14 @@ async def create_target_async( ``auth_mode`` that conflicts with the request-level choice, identity or no-key auth is requested but unsupported by the target type, or a target that uploads local files is requested with its upload folder or while - the server configures none. Construction errors (unknown params, + the server configures none, or a target to save would hold or use a credential + or names a Temporary target. Construction errors (unknown params, incompatible inner targets, unrecognized identity endpoints) are raised by the registry / target classes. PermissionError: If a non-administrator enables a target's model code execution, or creates a target that would authenticate with this server's stored credentials. + SavedTargetConflictError: If a target to save has the name of a registered or saved target. + SavedTargetStorageError: If a target to save can't be written to storage. """ if request.type not in self._registry: raise ValueError( @@ -635,12 +673,21 @@ async def create_target_async( self._enforce_server_credential_permission( target_type=request.type, params=permission_params, is_admin=is_admin ) + if request.persist: + params = await self._prepare_saved_params_async( + target_type=request.type, target_cls=target_cls, params=params, auth_mode=request.auth_mode + ) # LEGACY COMPATIBILITY: The current configuration UI omits the name. # Remove this generated fallback after that UI sends an explicit name. target_registry_name = request.name or f"compat_{uuid.uuid4().hex}" if request.register: - self._registry.instances.validate_name_available(target_registry_name) + try: + self._registry.instances.validate_name_available(target_registry_name) + except RegistryNameTakenError as error: + if request.persist: + raise SavedTargetConflictError(str(error)) from error + raise if ( request.source and type( @@ -668,13 +715,163 @@ async def create_target_async( finally: await release_component_async(target_obj) try: - target = self._build_instance_from_object(target_registry_name=target_registry_name, target_obj=target_obj) + if request.persist: + return await self._save_and_register_async( + name=target_registry_name, + type_name=request.type, + target_cls=target_cls, + target_obj=target_obj, + auth_mode=request.auth_mode, + params=params, + owner=owner, + recipes_source=recipes_source, + ) + target = self._build_instance_from_object( + target_registry_name=target_registry_name, target_obj=target_obj, origin=RegistryEntryOrigin.TEMPORARY + ) self._registry.instances.register(target_obj, name=target_registry_name) return target except Exception: await release_component_async(target_obj) raise + async def _prepare_saved_params_async( + self, *, target_type: str, target_cls: type[PromptTarget], params: dict[str, Any], auth_mode: AuthMode + ) -> dict[str, Any]: + """ + Turn off the server values a target to save would read, then refuse what it still can't hold. + + Returns: + dict[str, Any]: The parameters to build and save the target with. + + Raises: + ValueError: If the target would hold or use a credential, or names a Temporary target. + """ + saved_params = dict(params) + for name, turned_off in target_cls.server_fallback_parameters.items(): + if turned_off is not None and saved_params.get(name) in (None, ""): + saved_params[name] = turned_off + if self._registry.is_temporary_class(target_type): + raise ValueError( + f"'{target_type}' was added by a scan's launch script, so a saved target of it couldn't be " + "restored after a restart. It can still be created without saving it." + ) + metadata = await asyncio.to_thread(self._registry.get_registered_class_metadata, target_type) + if metadata is None: + # Restore falls back to the class itself the same way, so saving never checks less than restoring. + metadata = await asyncio.to_thread(self._registry.get_class_metadata, target_cls) + parameters = metadata.parameters + try: + check_saved_entry_params(cls=target_cls, parameters=parameters, params=saved_params, auth_mode=auth_mode) + except SavedEntryRefusedError as error: + raise ValueError(f"{error} It can still be created without saving it.") from error + self._refuse_temporary_references(parameters=parameters, params=saved_params) + return saved_params + + def _refuse_temporary_references(self, *, parameters: Sequence[Parameter], params: Mapping[str, Any]) -> None: + """ + Refuse a target to save that uses a Temporary target, which won't exist after a restart. + + Raises: + ValueError: If a target parameter names a Temporary target. + """ + for parameter in parameters: + if parameter.reference is None or parameter.reference.component_type is not ComponentType.TARGET: + continue + value = params.get(parameter.name) + for name in value if isinstance(value, list) else [value]: + entry = self._registry.instances.get_entry(name) if isinstance(name, str) else None + if entry is not None and registry_entry_origin(entry.metadata) is RegistryEntryOrigin.TEMPORARY: + raise ValueError( + f"Uses temporary target '{name}', which won't exist after a restart. " + "It can still be created without saving it." + ) + + async def _save_and_register_async( + self, + *, + name: str, + type_name: str, + target_cls: type[PromptTarget], + target_obj: PromptTarget, + auth_mode: AuthMode, + params: dict[str, Any], + owner: SavedEntryOwner | None, + recipes_source: str | None, + ) -> TargetInstance: + """ + Save the target's recipe, then register the target as a saved entry. + + The recipe is written only if no saved target has the name, so of two saves under one name + only one succeeds. If registering then fails, the recipe is deleted again. The recipe names + the class by the registry name the request used, which restore looks up, and keeps the auth + mode apart from the parameters, as the create request does. + + Returns: + TargetInstance: The saved target. + + Raises: + SavedTargetConflictError: If a saved or registered target already has the name. + SavedTargetStorageError: If the recipe can't be written, or can't be deleted after registering failed. + """ + target = self._build_instance_from_object( + target_registry_name=name, target_obj=target_obj, origin=RegistryEntryOrigin.SAVED + ) + auth_parameters = target_cls.get_auth_mode_parameters(auth_mode=auth_mode) + recipe = ComponentRecipe( + class_module=target_cls.__module__, + class_name=type_name, + auth_mode=auth_mode, + params={key: value for key, value in params.items() if key not in auth_parameters}, + owner=owner, + ) + storage = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=recipes_source) + try: + await asyncio.to_thread(storage.create_recipe, name=name, recipe=recipe) + except DocumentConflictError as error: + raise SavedTargetConflictError(f"A saved target named '{name}' already exists.") from error + except DocumentStorageUnavailableError as error: + logger.error( + "Saved target '%s' couldn't be written to %s (%s). If the write reached storage anyway, the " + "target is restored at the next start.", + name, + storage.display_source, + error.reason, + ) + raise SavedTargetStorageError("Saved targets can't be written right now.") from error + try: + self._registry.instances.register( + target_obj, name=name, metadata={SAVED_ENTRY_METADATA_KEY: RegistryRecipeKind.TARGETS.value} + ) + except Exception as error: + if not await self._delete_recipe_async(storage=storage, name=name): + raise SavedTargetStorageError("Saved targets can't be written right now.") from error + if isinstance(error, RegistryNameTakenError): + raise SavedTargetConflictError(str(error)) from error + raise + return target + + @staticmethod + async def _delete_recipe_async(*, storage: RegistryRecipeStorage, name: str) -> bool: + """ + Delete a recipe whose target couldn't be registered. + + Returns: + bool: False if the recipe couldn't be deleted, which is logged for an operator to fix. + """ + try: + await asyncio.to_thread(storage.delete_recipe, name) + except Exception as error: + logger.error( + "Saved target '%s' couldn't be registered, and its saved recipe in %s couldn't be deleted (%s). " + "Delete it there, or the target is restored at the next start.", + name, + storage.display_source, + error.reason if isinstance(error, DocumentStorageUnavailableError) else type(error).__name__, + ) + return False + return True + async def _enforce_code_execution_permission_async( self, *, target_type: str, params: dict[str, Any], is_admin: bool ) -> 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 "