diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 51488ee0d8..f31a5feaec 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -161,6 +161,14 @@ 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 a saved entry can't be restored, the backend still +# starts, lists the failure in its runtime status, and refuses to use that entry's +# name until the file is fixed or deleted. Defaults to ~/.pyrit/registry_recipes. +# 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..1f08ef06e8 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -194,6 +194,18 @@ 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 what's needed to rebuild the entry, and never 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": "{}"` 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`. + +The backend restores saved entries at startup and after a live apply, once the configured initializers have registered theirs. If a saved entry can't be restored, the backend still becomes ready and lists each failure with its reason in `GET /api/config/runtime`. The failed entry keeps its name: looking it up, or starting a run that uses it, fails with the reason (409) instead of finding nothing or something else, and nothing else can take the name. While a saved technique is broken, a launch has to name its techniques, because choosing them by `all`, `default`, or a tag could quietly leave it out. If a whole folder can't be listed, nothing is restored and every name of that kind that isn't registered is held the same way. Fix or delete the listed files, then apply or restart. Files whose names aren't valid entry names (up to 64 lowercase letters, digits, and underscores, starting with a letter) are listed as ignored. + ### `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 ee94076b9e..983f794fce 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", "**/scenario-hierarchy.spec.ts", "**/scenario-datasets.spec.ts", "**/scenario-techniques.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/pyrit/backend/README.md b/pyrit/backend/README.md index d0b8ea33ab..2b20158b9a 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -54,6 +54,39 @@ its database collation. ### Targets - `GET /api/targets` - List available prompt targets - `GET /api/targets/{id}` - Get target details +- `POST /api/targets` - Create a target + +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 only a scan's launch scripts added | **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** | +| takes the name of a saved entry that couldn't be restored | **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 its class and constructor parameters as the API accepted them, so it 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 its name stays held, so nothing +else stands in for it. ### Manual attack creation @@ -337,8 +370,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 0fcfb98932..73248443b6 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 TargetIdentifier -from pyrit.models.catalog import TargetInstance +from pyrit.models.catalog import RegistryEntryOrigin, TargetInstance from pyrit.prompt_target import CapabilityName, PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.registry import Reconstructable @@ -55,7 +55,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. @@ -66,6 +68,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. @@ -100,6 +103,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 5973d6289e..9186f41531 100644 --- a/pyrit/backend/middleware/error_handlers.py +++ b/pyrit/backend/middleware/error_handlers.py @@ -11,9 +11,12 @@ from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse +from pyrit.backend.middleware.auth import has_admin_access from pyrit.backend.middleware.request_size import RequestTooLargeError, request_too_large_response from pyrit.backend.models import FieldError, ProblemDetail from pyrit.exceptions.analytics_exception import AnalyticsException +from pyrit.registry import SavedEntryUnavailableError, TechniqueSelectionUnavailableError +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError, DocumentUnreadableError logger = logging.getLogger(__name__) @@ -169,6 +172,68 @@ async def permission_error_handler( # pyrit-async-suffix-exempt content=problem.model_dump(exclude_none=True), ) + @app.exception_handler(DocumentStorageUnavailableError) + async def storage_unavailable_handler( # pyrit-async-suffix-exempt + request: Request, + exc: DocumentStorageUnavailableError, + ) -> JSONResponse: + """ + Handle stored documents that can't be read or written as 503 Service Unavailable. + + The log names the storage location, without credentials, and the reason, with the + traceback at debug level. The response never carries a path or URL. + + Returns: + JSONResponse: RFC 7807 problem detail response with 503 status. + """ + logger.warning(f"{exc} ({request.method} {request.url.path})") + logger.debug("Storage error details", exc_info=exc) + if isinstance(exc, DocumentUnreadableError): + detail = f"{exc.label} '{exc.name}' can't be read: {exc.reason}" + else: + detail = f"{exc.label} storage is unavailable" + problem = ProblemDetail( + type="/errors/storage-unavailable", + title="Service Unavailable", + status=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=detail, + instance=str(request.url.path), + ) + + return JSONResponse( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + content=problem.model_dump(exclude_none=True), + ) + + @app.exception_handler(SavedEntryUnavailableError) + async def saved_entry_unavailable_handler( # pyrit-async-suffix-exempt + request: Request, + exc: SavedEntryUnavailableError, + ) -> JSONResponse: + """ + Handle a name held for a saved entry that couldn't be restored as 409 Conflict. + + Administrators get the full message; anyone else gets a fixed explanation of the reason, + since the message can quote a server setting. + + Returns: + JSONResponse: RFC 7807 problem detail response with 409 status. + """ + return _conflict(request=request, detail=exc.describe(with_details=_shows_details(request))) + + @app.exception_handler(TechniqueSelectionUnavailableError) + async def technique_selection_unavailable_handler( # pyrit-async-suffix-exempt + request: Request, + exc: TechniqueSelectionUnavailableError, + ) -> JSONResponse: + """ + Handle a broad technique selection made while a saved technique couldn't be restored as 409 Conflict. + + Returns: + JSONResponse: RFC 7807 problem detail response with 409 status. + """ + return _conflict(request=request, detail=str(exc)) + @app.exception_handler(NotImplementedError) async def not_implemented_handler( # pyrit-async-suffix-exempt request: Request, @@ -223,3 +288,33 @@ async def generic_exception_handler( # pyrit-async-suffix-exempt status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=problem.model_dump(exclude_none=True), ) + + +def _conflict(*, request: Request, detail: str) -> JSONResponse: + """ + Build a 409 Conflict problem response. + + Returns: + JSONResponse: RFC 7807 problem detail response with 409 status. + """ + problem = ProblemDetail( + type="/errors/conflict", + title="Conflict", + status=status.HTTP_409_CONFLICT, + detail=detail, + instance=str(request.url.path), + ) + return JSONResponse(status_code=status.HTTP_409_CONFLICT, content=problem.model_dump(exclude_none=True)) + + +def _shows_details(request: Request) -> bool: + """ + Return whether the caller may see a restore failure's full message, which can quote a server setting. + + Returns: + bool: True for an administrator; False otherwise, including when no authentication policy is set. + """ + try: + return has_admin_access(request) + except RuntimeError: + return False diff --git a/pyrit/backend/models/targets.py b/pyrit/backend/models/targets.py index 61e79de672..6be3d197d2 100644 --- a/pyrit/backend/models/targets.py +++ b/pyrit/backend/models/targets.py @@ -12,7 +12,7 @@ from pyrit.backend.models.common import MAX_ITEMS, REGISTRY_INSTANCE_NAME_PATTERN, IdentifierStr, PaginationInfo from pyrit.common import AuthMode -from pyrit.models import JSONValue, Parameter, TargetCapabilities, TargetIdentifier +from pyrit.models import JSONValue, Parameter, TargetCapabilities, TargetIdentifier, validate_registry_name from pyrit.models.catalog import TargetInstance from pyrit.models.component_spec import SourceInstanceSpec @@ -68,6 +68,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( @@ -91,6 +99,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/attacks.py b/pyrit/backend/routes/attacks.py index 51cc1d4dc0..76bac6db9f 100644 --- a/pyrit/backend/routes/attacks.py +++ b/pyrit/backend/routes/attacks.py @@ -45,6 +45,7 @@ from pyrit.backend.services.media_persistence import MediaAccessDeniedError from pyrit.common import print_deprecation_message from pyrit.memory.memory_interface import AttackStateConflictError +from pyrit.registry import SavedEntryUnavailableError logger = logging.getLogger(__name__) @@ -666,6 +667,8 @@ async def add_message( # pyrit-async-suffix-exempt status_code=status.HTTP_400_BAD_REQUEST, detail=error_msg, ) from e + except SavedEntryUnavailableError: + raise except Exception as e: logger.exception("Failed to add message to attack '%s'", attack_result_id) raise HTTPException( diff --git a/pyrit/backend/routes/converters.py b/pyrit/backend/routes/converters.py index 10d129394e..d97841170e 100644 --- a/pyrit/backend/routes/converters.py +++ b/pyrit/backend/routes/converters.py @@ -24,6 +24,7 @@ UnregisteredConverter, ) from pyrit.backend.services import get_converter_service +from pyrit.registry import SavedEntryUnavailableError router = APIRouter(prefix="/converters", tags=["converters"]) @@ -89,6 +90,8 @@ async def create_converter( status_code=status.HTTP_400_BAD_REQUEST, detail=str(e), ) from e + except SavedEntryUnavailableError: + raise except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -165,6 +168,8 @@ async def preview_conversion(request: ConverterPreviewRequest) -> ConverterPrevi status_code=status.HTTP_400_BAD_REQUEST, detail=str(e), ) from e + except SavedEntryUnavailableError: + raise except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 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/message_sends.py b/pyrit/backend/routes/message_sends.py index 297783c99e..7441f28af5 100644 --- a/pyrit/backend/routes/message_sends.py +++ b/pyrit/backend/routes/message_sends.py @@ -15,6 +15,7 @@ from pyrit.backend.services.manual_send_scheduler import ManualSendConflictError, ManualSendQueueFullError from pyrit.backend.services.message_send_service import MessageSendNotFoundError, get_message_send_service from pyrit.memory.memory_interface import AttackStateConflictError +from pyrit.registry import SavedEntryUnavailableError logger = logging.getLogger(__name__) router = APIRouter(prefix="/attacks", tags=["attacks"]) @@ -49,6 +50,8 @@ async def submit_message_send_async( except ValueError as exc: is_not_found = "not found" in str(exc).lower() or "does not exist" in str(exc).lower() raise HTTPException(status_code=404 if is_not_found else 400, detail=str(exc)) from exc + except SavedEntryUnavailableError: + raise except Exception as exc: logger.exception("Failed to accept message send for attack '%s'", attack_result_id) raise HTTPException(status_code=500, detail="Could not accept the message send. Check server logs.") from exc diff --git a/pyrit/backend/routes/scenario_presets.py b/pyrit/backend/routes/scenario_presets.py index 4b3fc1db00..973112df7b 100644 --- a/pyrit/backend/routes/scenario_presets.py +++ b/pyrit/backend/routes/scenario_presets.py @@ -19,7 +19,6 @@ from typing import Annotated -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Query, status from pyrit.backend.middleware.auth import current_user_name, require_admin @@ -38,19 +37,6 @@ router = APIRouter(prefix="/scenario-presets", tags=["scenario-presets"]) -def _storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable preset storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Scenario preset storage is temporarily unavailable", - ) - - def _not_found(name: str) -> HTTPException: """ Create a not-found response for a missing preset. @@ -79,8 +65,6 @@ async def _load_preset_or_404_async(name: str) -> ScenarioPresetResponse: """ try: preset = await get_scenario_preset_service().get_preset_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -106,10 +90,7 @@ async def list_scenario_presets( # pyrit-async-suffix-exempt Returns: ScenarioPresetListResponse: The configured source and the stored presets. """ - try: - return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) - except AzureError as exc: - raise _storage_unavailable() from exc + return await get_scenario_preset_service().list_presets_async(include_estimates=include_estimates) @router.get( @@ -177,8 +158,6 @@ async def create_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{preset.name}' already exists", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.put( @@ -239,8 +218,6 @@ async def update_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{name}' changed since it was read; re-read it and retry", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.delete( @@ -261,8 +238,6 @@ async def delete_scenario_preset(name: IdentifierStr) -> None: # pyrit-async-su """ try: deleted = await get_scenario_preset_service().delete_preset_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -312,8 +287,6 @@ async def resolve_scenario_preset( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario preset '{name}' changed since it was read; re-read it and retry", ) from None - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None diff --git a/pyrit/backend/routes/scenario_suites.py b/pyrit/backend/routes/scenario_suites.py index 0b4fa7e09a..f1ec16872e 100644 --- a/pyrit/backend/routes/scenario_suites.py +++ b/pyrit/backend/routes/scenario_suites.py @@ -6,7 +6,6 @@ from typing import Annotated from uuid import UUID -from azure.core.exceptions import AzureError from fastapi import APIRouter, Depends, HTTPException, Request, status from pyrit.backend.middleware.auth import current_user_name, require_admin @@ -31,19 +30,6 @@ router = APIRouter(prefix="/scenario-suites", tags=["scenario-suites"]) -def _storage_unavailable() -> HTTPException: - """ - Create a sanitized response for unavailable suite storage. - - Returns: - HTTPException: A service-unavailable response without SDK details. - """ - return HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Scenario suite storage is temporarily unavailable", - ) - - def _not_found(name: str) -> HTTPException: """ Create a not-found response for a missing suite. @@ -66,8 +52,6 @@ async def _load_suite_or_404_async(name: str) -> ScenarioSuiteResponse: """ try: suite = await get_scenario_suite_service().get_suite_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None @@ -84,10 +68,7 @@ async def list_scenario_suites() -> ScenarioSuiteListResponse: # pyrit-async-su Returns: ScenarioSuiteListResponse: Suites from both origins with current entry statuses. """ - try: - return await get_scenario_suite_service().list_suites_async() - except AzureError as exc: - raise _storage_unavailable() from exc + return await get_scenario_suite_service().list_suites_async() @router.get( @@ -181,8 +162,6 @@ async def launch_scenario_suite( # pyrit-async-suffix-exempt raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None except ScenarioSuiteExecutionCapacityError as exc: raise HTTPException(status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=str(exc)) from None - except AzureError as exc: - raise _storage_unavailable() from exc except RuntimeError as exc: raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(exc)) from None @@ -239,8 +218,6 @@ async def create_scenario_suite( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario suite '{suite.name}' already exists", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.put( @@ -289,8 +266,6 @@ async def update_scenario_suite( # pyrit-async-suffix-exempt status_code=status.HTTP_409_CONFLICT, detail=f"Scenario suite '{name}' changed since it was read; re-read it and retry", ) from exc - except AzureError as exc: - raise _storage_unavailable() from exc @router.delete( @@ -303,8 +278,6 @@ async def delete_scenario_suite(name: IdentifierStr) -> None: # pyrit-async-suf """Delete one stored scenario suite.""" try: deleted = await get_scenario_suite_service().delete_suite_async(name=name) - except AzureError as exc: - raise _storage_unavailable() from exc except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index f7cc97bba4..945f5a3834 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -220,7 +220,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/scorers.py b/pyrit/backend/routes/scorers.py index 445f2830b3..65287aea6e 100644 --- a/pyrit/backend/routes/scorers.py +++ b/pyrit/backend/routes/scorers.py @@ -17,6 +17,7 @@ ) from pyrit.backend.services.scorer_service import get_scorer_service from pyrit.models.catalog import ScorerInstance +from pyrit.registry import SavedEntryUnavailableError router = APIRouter(prefix="/scorers", tags=["scorers"]) @@ -65,6 +66,8 @@ async def create_scorer( return await get_scorer_service().create_scorer_async(request=request) except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc + except SavedEntryUnavailableError: + raise except Exception as exc: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, diff --git a/pyrit/backend/routes/targets.py b/pyrit/backend/routes/targets.py index c520a42e05..faa51c4d65 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, @@ -22,7 +22,9 @@ 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.registry import SavedEntryUnavailableError router = APIRouter(prefix="/targets", tags=["targets"]) @@ -41,6 +43,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, @@ -96,6 +114,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( @@ -121,7 +144,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, @@ -132,6 +161,8 @@ async def create_target( status_code=status.HTTP_400_BAD_REQUEST, detail=str(e), ) from e + except SavedEntryUnavailableError: + raise except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, diff --git a/pyrit/backend/services/runtime_lifecycle.py b/pyrit/backend/services/runtime_lifecycle.py index daef21f0ed..fdbc7b0b3c 100644 --- a/pyrit/backend/services/runtime_lifecycle.py +++ b/pyrit/backend/services/runtime_lifecycle.py @@ -7,7 +7,7 @@ import logging import os import uuid -from typing import Any +from typing import TYPE_CHECKING, Any from fastapi import FastAPI @@ -25,11 +25,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 +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: @@ -96,6 +106,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 +130,37 @@ 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) -> None: + """ + Rebuild saved registry entries after the startup scripts have registered theirs. + + The runtime becomes ready even when some can't be restored. Each one that failed keeps its name, + so looking it up or starting a run that uses it fails with the reason instead of finding nothing + or something else, and the failures are listed in the runtime status. + + Args: + config (ConfigurationLoader): The loaded configuration naming the storage source. + """ + 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, as + # the target API builds them, 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 = restorer.restore(saved=saved) + self.restore_failures, self.restore_ignored = report.failures, report.ignored + if report.failures: + count = len(report.failures) + logger.warning( + "%d saved registry %s couldn't be restored. Fix or delete them; see GET /api/config/runtime.", + count, + "entry" if count == 1 else "entries", + ) + 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,9 +172,13 @@ 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() + await self._restore_saved_entries_async(config) self._publish(config) except Exception: self.state, self.outcome = "restart-required", "initialization-failed" @@ -216,9 +258,11 @@ 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 + await self._restore_saved_entries_async(config) self._publish(config) except (Exception, asyncio.CancelledError, BaseExceptionGroup) as error: logger.exception( diff --git a/pyrit/backend/services/scenario_configuration_resolver.py b/pyrit/backend/services/scenario_configuration_resolver.py index 42a78d86f6..fa5d941657 100644 --- a/pyrit/backend/services/scenario_configuration_resolver.py +++ b/pyrit/backend/services/scenario_configuration_resolver.py @@ -7,7 +7,15 @@ from typing import TYPE_CHECKING, Any -from pyrit.registry import ScenarioRegistry, ScorerRegistry, TargetRegistry +from pyrit.registry import ( + AttackTechniqueRegistry, + SavedEntryUnavailableError, + ScenarioRegistry, + ScorerRegistry, + TargetRegistry, + TechniqueSelectionUnavailableError, + UnavailableEntry, +) from pyrit.scenario.core import resolve_technique_tokens from pyrit.scenario.core.scenario_target_defaults import validate_default_adversarial_target from pyrit.score import TrueFalseScorer @@ -105,6 +113,41 @@ def resolve_target(*, target_name: str) -> PromptTarget: f"Target '{target_name}' not found in registry. Available targets: {', '.join(available_names)}" ) + @staticmethod + def _refuse_broad_selection_while_techniques_are_held( + *, + scenario_class: type[Scenario], + techniques: list[str] | None, + technique_override: TechniqueOverride | None, + ) -> None: + """ + Refuse choosing techniques by all, default, or a tag while a saved technique couldn't be restored. + + A broken technique's tags may be unknown, so such a selection could quietly leave it out. Naming + techniques still works, and scenarios whose techniques don't come from the registry aren't affected. + + Raises: + SavedEntryUnavailableError: If a named technique is held for a saved technique that couldn't be restored. + TechniqueSelectionUnavailableError: If the selection is broad while a saved technique is held. + """ + held = raise_if_a_named_technique_is_held(techniques) + if not held or "scenario_techniques" not in {p.name for p in scenario_class.supported_parameters()}: + return + selectors = [token.partition(":")[0] for token in techniques or []] + if not selectors and technique_override is not None: + return + try: + technique_class = scenario_class()._technique_class # type: ignore[ty:missing-argument] + except Exception: + technique_class = None + if technique_class is not None and getattr(technique_class, "_selection_catalog", None) is None: + return + aggregates = {"all", "default"} + if technique_class is not None: + aggregates |= {tag.casefold() for tag in technique_class.get_aggregate_tags()} + if not selectors or any(selector.casefold() in aggregates for selector in selectors): + raise TechniqueSelectionUnavailableError(names=[entry.name for entry in held if entry.name is not None]) + @classmethod def resolve_configuration( cls, @@ -131,6 +174,9 @@ def resolve_configuration( Raises: ValueError: If techniques or dataset overrides are invalid. """ + cls._refuse_broad_selection_while_techniques_are_held( + scenario_class=scenario_class, techniques=techniques, technique_override=technique_override + ) resolved: dict[str, Any] = {} if objective_target is not None: resolved["objective_target"] = objective_target @@ -213,3 +259,25 @@ def resolve_techniques_and_converters( if resolution.issues: raise ValueError(resolution.issues[0].message) return resolution.techniques, resolution.technique_converters + + +def raise_if_a_named_technique_is_held(techniques: list[str] | None) -> list[UnavailableEntry]: + """ + Refuse technique tokens that name a saved technique that couldn't be restored, in any letter case. + + Args: + techniques (list[str] | None): Technique tokens, each optionally followed by ``:`` and converters. + + Returns: + list[UnavailableEntry]: Every held saved technique, and a held folder of them. + + Raises: + SavedEntryUnavailableError: If a token names a held technique. + """ + held = AttackTechniqueRegistry.get_registry_singleton().instances.get_unavailable_entries() + held_by_name = {entry.name.casefold(): entry for entry in held if entry.name is not None} + for token in techniques or []: + name = token.partition(":")[0] + if (entry := held_by_name.get(name.casefold())) is not None: + raise SavedEntryUnavailableError(entry=entry, name=name) + return held diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index ed3b7d5aba..220c1dce11 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -46,7 +46,10 @@ normalize_label_filters, ) from pyrit.backend.services.saved_conversation_read_service import SavedConversationReadService -from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver +from pyrit.backend.services.scenario_configuration_resolver import ( + ScenarioConfigurationResolver, + raise_if_a_named_technique_is_held, +) from pyrit.backend.services.scenario_progress_read_model import ( ScenarioPlanLookup, ScenarioProgressReadModel, @@ -96,10 +99,21 @@ 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.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 @@ -179,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.""" @@ -443,6 +461,7 @@ def _selected_launch_factories( technique_class = metadata.technique_class if not isinstance(technique_class, type): return [local_factory] if local_factory is not None else [] + raise_if_a_named_technique_is_held(request.techniques) if configured_techniques is None: selectors = ( [technique_class(token.partition(":")[0]) for token in request.techniques] @@ -1614,6 +1633,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 @@ -1627,7 +1647,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 1b85e89580..4691281ee6 100644 --- a/pyrit/backend/services/scenario_suite_execution_service.py +++ b/pyrit/backend/services/scenario_suite_execution_service.py @@ -13,7 +13,11 @@ from typing import TYPE_CHECKING from uuid import UUID, uuid4 -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, @@ -25,9 +29,11 @@ ) from pyrit.registry import ( ResolvedScenarioSuiteEntry, + SavedEntryUnavailableError, ScenarioPresetRegistry, ScenarioSuiteLaunchResolver, ScenarioSuiteRegistry, + TechniqueSelectionUnavailableError, ) if TYPE_CHECKING: @@ -261,7 +267,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 @@ -290,7 +296,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: @@ -547,6 +553,23 @@ 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. + """ + # These name only the launch script, entries and techniques, so they're safe to show. + if isinstance(error, (SavedEntryLaunchConflictError, TechniqueSelectionUnavailableError)): + return str(error) + if isinstance(error, SavedEntryUnavailableError): + return error.describe(with_details=False) + 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 4e668bb4eb..2636d842d6 100644 --- a/pyrit/backend/services/scenario_suite_service.py +++ b/pyrit/backend/services/scenario_suite_service.py @@ -9,8 +9,6 @@ from functools import lru_cache from typing import Literal -from azure.core.exceptions import AzureError - from pyrit.backend.models import ( ScenarioSuiteEntryStatus, ScenarioSuiteListResponse, @@ -25,6 +23,7 @@ ScenarioSuiteRegistryEntry, ) from pyrit.registry import ScenarioPresetRegistry, ScenarioSuiteRegistry +from pyrit.registry.file_document_storage import DocumentStorageUnavailableError logger = logging.getLogger(__name__) @@ -209,7 +208,7 @@ def _resolve_preset_projection( ) stored_version = registry.get_preset_version(name=preset_name) - except (AzureError, OSError): + except DocumentStorageUnavailableError: logger.warning( "Could not project scenario preset '%s' because preset storage is unavailable.", preset_name, diff --git a/pyrit/backend/services/target_service.py b/pyrit/backend/services/target_service.py index aae49d1886..7651edee1d 100644 --- a/pyrit/backend/services/target_service.py +++ b/pyrit/backend/services/target_service.py @@ -15,6 +15,7 @@ import asyncio import logging import uuid +from collections.abc import Mapping, Sequence from functools import lru_cache from typing import Any @@ -31,11 +32,21 @@ ) from pyrit.backend.services.component_lifecycle import release_component_async from pyrit.common import AUTH_MODES, REQUIRED_VALUE, AuthMode, default_values -from pyrit.models import Parameter, TargetIdentifier -from pyrit.models.catalog import TargetInstance +from pyrit.models import ComponentType, Parameter, 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, PromptTarget 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__) @@ -52,6 +63,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. @@ -65,14 +84,16 @@ def __init__(self) -> None: """Initialize the target service.""" self._registry = TargetRegistry.get_registry_singleton() - 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, @@ -91,7 +112,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) @@ -132,10 +155,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: """ @@ -352,7 +379,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. @@ -375,7 +408,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. @@ -383,6 +417,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. @@ -392,11 +428,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( @@ -467,12 +506,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 + ) # 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( @@ -499,13 +547,153 @@ 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, + 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] + ) -> 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) + 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, + 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. + + 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 + ) + recipe = ComponentRecipe(class_module=target_cls.__module__, class_name=type_name, params=params, 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..0195a24956 100644 --- a/pyrit/common/azure_storage.py +++ b/pyrit/common/azure_storage.py @@ -3,8 +3,14 @@ """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. +_URL_IN_TEXT = re.compile(r"\b[A-Za-z][A-Za-z0-9+.-]*://\S+") +# The signature of a SAS token, which can also appear in text without a scheme. +_SAS_SIGNATURE = re.compile(r"(?i)\bsig=[^&\s]+") + def is_azure_blob_uri(value: str, *, min_path_segments: int = 1) -> bool: """ @@ -48,9 +54,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 "