diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index c00c7baf78..af4d551a4f 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -196,15 +196,15 @@ A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredentia ### `registry_recipes_source` -Stores saved targets, converters, scorers, and techniques in a local directory or Azure Blob container, so they come back when the backend restarts. An Azure URI may include a blob-name prefix: +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; restore refuses a file that stores a value for a parameter its class marks as a credential, such as `api_key`. A SAS query string may be included; otherwise, PyRIT uses `DefaultAzureCredential`. The default is `~/.pyrit/registry_recipes`. +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`. -A saved file doesn't record who saved it, so the backend rebuilds every entry in this source as an administrator, at startup and after a live apply. A file here can, for example, make the backend run a model's repository code, or use the server's identity on a storage container that the file names. Write access to this source is administrator authority on the backend: grant it only to principals you already trust as backend administrators, and use a dedicated directory or blob prefix that less-trusted users and services cannot write to. +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`. 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/pyrit/backend/README.md b/pyrit/backend/README.md index ab3aae9ee6..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. 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/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/scenario_suites.py b/pyrit/backend/routes/scenario_suites.py index 79ea691149..8cd0112916 100644 --- a/pyrit/backend/routes/scenario_suites.py +++ b/pyrit/backend/routes/scenario_suites.py @@ -17,6 +17,7 @@ ) from pyrit.backend.models.common import IdentifierStr from pyrit.backend.services import get_scenario_suite_execution_service, get_scenario_suite_service +from pyrit.backend.services.scenario_run_service import SavedEntryLaunchConflictError from pyrit.backend.services.scenario_suite_execution_service import ScenarioSuiteExecutionCapacityError from pyrit.models import ( LaunchScenarioSuiteRequest, @@ -152,7 +153,10 @@ async def cancel_scenario_suite_execution( # pyrit-async-suffix-exempt responses={ 400: {"model": ProblemDetail, "description": "The suite has unresolved or invalid entries"}, 404: {"model": ProblemDetail, "description": "Scenario suite not found"}, - 409: {"model": ProblemDetail, "description": "The suite or a referenced preset changed"}, + 409: { + "model": ProblemDetail, + "description": "The suite or a referenced preset changed, or a launch script would replace a saved entry", + }, 429: {"model": ProblemDetail, "description": "Active suite execution capacity is full"}, 503: {"model": ProblemDetail, "description": "Runtime or definition storage unavailable"}, }, @@ -179,7 +183,7 @@ async def launch_scenario_suite( # pyrit-async-suffix-exempt return await get_scenario_suite_execution_service().launch_async(name=name, request=body) except ScenarioSuiteExecutionNotFoundError as exc: raise _not_found(name) from exc - except ScenarioSuiteLaunchConflictError as exc: + except (ScenarioSuiteLaunchConflictError, SavedEntryLaunchConflictError) as exc: raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(exc)) from None except ScenarioSuiteConflictError as exc: raise HTTPException( 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 ebeecf30c8..efe118884a 100644 --- a/pyrit/backend/services/runtime_lifecycle.py +++ b/pyrit/backend/services/runtime_lifecycle.py @@ -27,7 +27,9 @@ ) 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 @@ -106,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() @@ -179,6 +182,7 @@ 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 = {} @@ -190,7 +194,10 @@ 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() if await self._restore_saved_entries_async(config): @@ -275,7 +282,8 @@ 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 if await self._restore_saved_entries_async(config): diff --git a/pyrit/backend/services/saved_entry_restore.py b/pyrit/backend/services/saved_entry_restore.py index a4ac353988..ac89699d36 100644 --- a/pyrit/backend/services/saved_entry_restore.py +++ b/pyrit/backend/services/saved_entry_restore.py @@ -26,8 +26,8 @@ async def build_saved_entry_async(stored: StoredRecipe, *, upload_directory: str """ Build one saved entry with the create the API uses for its kind, and register it under its name. - The entry passes the same checks a create request does. Restore runs as the server, and a saved file - doesn't record who saved it, so a target is built as an administrator would build it. A target that + 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: 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/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/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index a113a8bef0..a1f0c4575a 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -23,10 +23,12 @@ from pyrit.models.catalog.registry_recipe import ( ComponentRecipe, IgnoredRecipeFile, + RegistryEntryOrigin, RegistryRecipeKind, RegistryRestoreReport, RestoreFailure, RestoreFailureReason, + SavedEntryOwner, SavedEntryRef, StoredRecipe, TechniqueRecipe, @@ -77,9 +79,11 @@ "PresetIssue": "pyrit.models.catalog.scenario_preset", "RegisteredInitializer": "pyrit.models.catalog.initializer", "RegistryRecipeKind": "pyrit.models.catalog.registry_recipe", + "RegistryEntryOrigin": "pyrit.models.catalog.registry_recipe", "RegistryRestoreReport": "pyrit.models.catalog.registry_recipe", "RestoreFailure": "pyrit.models.catalog.registry_recipe", "RestoreFailureReason": "pyrit.models.catalog.registry_recipe", + "SavedEntryOwner": "pyrit.models.catalog.registry_recipe", "SavedEntryRef": "pyrit.models.catalog.registry_recipe", "ScorerInstance": "pyrit.models.catalog.scorer", "RegisteredScenario": "pyrit.models.catalog.scenario", diff --git a/pyrit/models/catalog/registry_recipe.py b/pyrit/models/catalog/registry_recipe.py index c381f1c763..e019318cf6 100644 --- a/pyrit/models/catalog/registry_recipe.py +++ b/pyrit/models/catalog/registry_recipe.py @@ -19,6 +19,14 @@ _TechniqueTag = Annotated[str, Field(pattern=r"^[A-Za-z][A-Za-z0-9_]{0,63}$")] +class RegistryEntryOrigin(str, Enum): + """Where a registered entry comes from, and so whether it is still there after a restart.""" + + SAVED = "saved" + CONFIGURATION = "configuration" + TEMPORARY = "temporary" + + class RegistryRecipeKind(str, Enum): """The registries whose entries can be saved, named after their storage folders.""" @@ -38,6 +46,15 @@ class _RegistryRecipe(BaseModel): schema_version: int = Field(1, ge=1, le=RECIPE_SCHEMA_VERSION, description="Shape of this document") +class SavedEntryOwner(BaseModel): + """The signed-in user who saved an entry.""" + + model_config = ConfigDict(extra="forbid", strict=True, frozen=True) + + oid: _Identifier = Field(..., min_length=1, description="The user's Microsoft Entra object ID") + name: _Identifier = Field(..., description="The user's display name, or email when no name is set") + + class ComponentRecipe(_RegistryRecipe): """How to rebuild one saved target, converter, or scorer.""" @@ -53,6 +70,7 @@ class ComponentRecipe(_RegistryRecipe): max_length=MAX_ITEMS, description="Constructor parameters as the API accepted them; other entries are named, never embedded", ) + owner: SavedEntryOwner | None = Field(None, description="Who saved the entry; None when sign-in is off") class TechniqueRecipe(_RegistryRecipe): diff --git a/pyrit/models/catalog/target.py b/pyrit/models/catalog/target.py index 1d1dd3496a..7a6fad1a7b 100644 --- a/pyrit/models/catalog/target.py +++ b/pyrit/models/catalog/target.py @@ -24,6 +24,7 @@ from pydantic import BaseModel, Field +from pyrit.models.catalog.registry_recipe import RegistryEntryOrigin from pyrit.models.identifiers.component_identifier import JSONValue from pyrit.models.identifiers.target_identifier import TargetIdentifier from pyrit.models.target.request_generation_config import GenerationRange @@ -70,3 +71,12 @@ class TargetInstance(BaseModel): None, description="Inner targets for composite targets like RoundRobinTarget (full instances, not just identifiers)", ) + origin: RegistryEntryOrigin | None = Field( + None, + description=( + "Where the target comes from: 'saved' (restored at every start), 'configuration' (registered by " + "the server's startup scripts), or 'temporary' (created through the API without saving, or " + "registered by a scan's launch scripts; gone after a restart or a configuration apply). None for " + "an inner target" + ), + ) diff --git a/pyrit/prompt_target/a2a_target.py b/pyrit/prompt_target/a2a_target.py index a57d863889..e41e947566 100644 --- a/pyrit/prompt_target/a2a_target.py +++ b/pyrit/prompt_target/a2a_target.py @@ -8,10 +8,11 @@ import math import time import uuid -from collections.abc import AsyncGenerator, Awaitable, Callable # noqa: TC003 - registry annotation resolution +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping # noqa: TC003 - registry annotation resolution from dataclasses import dataclass from email.utils import parsedate_to_datetime -from typing import TYPE_CHECKING, Any, Literal +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, ClassVar, Literal from weakref import WeakValueDictionary import httpx @@ -79,6 +80,10 @@ class A2ATarget(PromptTarget): `input-required` tasks are extracted from the task status message. """ + # The endpoint has no server setting to fall back to, and the card path has a fixed default. + destination_parameters: ClassVar[Mapping[str, bool]] = MappingProxyType( + {"endpoint": False, "agent_card_path": False} + ) _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( capabilities=TargetCapabilities( supports_multi_turn=True, diff --git a/pyrit/prompt_target/azure_blob_storage_target.py b/pyrit/prompt_target/azure_blob_storage_target.py index cc10fc667d..33bb081cf9 100644 --- a/pyrit/prompt_target/azure_blob_storage_target.py +++ b/pyrit/prompt_target/azure_blob_storage_target.py @@ -3,6 +3,8 @@ import logging import posixpath +from collections.abc import Mapping +from types import MappingProxyType from typing import ClassVar from urllib.parse import urlparse @@ -53,6 +55,8 @@ class AzureBlobStorageTarget(PromptTarget): # the target authenticates with the server's own credentials, so the create-target API gates # it on administrator access. server_credential_parameter: ClassVar[str | None] = "sas_token" + server_fallback_parameters: ClassVar[Mapping[str, str | None]] = MappingProxyType({"sas_token": None}) + destination_parameters: ClassVar[Mapping[str, bool]] = MappingProxyType({"container_url": True}) _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( capabilities=TargetCapabilities( diff --git a/pyrit/prompt_target/azure_ml_chat_target.py b/pyrit/prompt_target/azure_ml_chat_target.py index cdb4c9524d..38d9a76ee3 100644 --- a/pyrit/prompt_target/azure_ml_chat_target.py +++ b/pyrit/prompt_target/azure_ml_chat_target.py @@ -2,7 +2,8 @@ # Licensed under the MIT license. import logging -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping +from types import MappingProxyType from typing import Any, ClassVar, cast from httpx import HTTPStatusError @@ -51,6 +52,8 @@ class AzureMLChatTarget(PromptTarget): # AML managed online endpoints can authenticate with a Microsoft Entra ID # token scoped to AML, so this target supports both auth modes. supported_auth_modes: ClassVar[tuple[AuthMode, ...]] = ("api_key", "identity") + server_fallback_parameters: ClassVar[Mapping[str, str | None]] = MappingProxyType({"api_key": None}) + destination_parameters: ClassVar[Mapping[str, bool]] = MappingProxyType({"endpoint": True}) # Entra ID token scope for Azure Machine Learning managed online endpoints. _AZURE_ML_SCOPE: ClassVar[str] = "https://ml.azure.com/.default" diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index d2b75a0ee7..de1e2499c4 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -6,6 +6,7 @@ from collections.abc import Generator, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar +from types import MappingProxyType from typing import Any, ClassVar, final from pyrit.common import AuthMode, print_deprecation_message @@ -120,6 +121,28 @@ class PromptTarget(Identifiable): # unaffected. server_credential_parameter: ClassVar[str | None] = None + # Credential parameters that, left empty, let the target fall back to a credential from this + # server (an environment variable, or the server's Azure identity), each mapped to the value + # that turns that off, or None when only an auth mode does: ``"identity"`` doesn't read a + # parameter mapped to None, and ``"none"`` reads none of them. A saved target never falls back + # to this server's credentials, so a saved entry that would is refused. Choosing ``"identity"`` + # isn't a fallback. + server_fallback_parameters: ClassVar[Mapping[str, str | None]] = MappingProxyType({}) + + # Parameters that say where the target sends requests (a URL, or a path the target adds to + # one), each mapped to whether an empty value takes this server's setting. A saved target sets + # each one that would, since a server setting can carry a credential, and none can carry a + # credential itself: a user name or password, a SAS token, a fragment, or a query other than + # ``api-version``. + destination_parameters: ClassVar[Mapping[str, bool]] = MappingProxyType({}) + + # Environment variables the target's client reads extra headers from, whatever the target's own + # settings say. A server can put a credential there, so a saved target isn't built while one is set. + server_header_variables: ClassVar[tuple[str, ...]] = () + + # Why this target type can't be saved even without a credential, or None when it can. + not_saveable_reason: ClassVar[str | None] = None + def __init_subclass__(cls, **kwargs: object) -> None: """ Validate that subclasses follow the keyword-only ``__init__`` contract. diff --git a/pyrit/prompt_target/github_copilot_target.py b/pyrit/prompt_target/github_copilot_target.py index 1557b522db..e70a50e337 100644 --- a/pyrit/prompt_target/github_copilot_target.py +++ b/pyrit/prompt_target/github_copilot_target.py @@ -4,8 +4,10 @@ import asyncio import logging import math +from collections.abc import Mapping from dataclasses import dataclass, field from pathlib import Path +from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar from uuid import uuid4 @@ -46,6 +48,7 @@ class GitHubCopilotTarget(PromptTarget): GITHUB_TOKEN_ENVIRONMENT_VARIABLE: str = "GITHUB_TOKEN" caller_path_parameters: ClassVar[frozenset[str]] = frozenset({"working_directory"}) + server_fallback_parameters: ClassVar[Mapping[str, str | None]] = MappingProxyType({"github_token": None}) _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( capabilities=TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) diff --git a/pyrit/prompt_target/hack_a_prompt_target.py b/pyrit/prompt_target/hack_a_prompt_target.py index d58cfacf1e..29a4397b0c 100644 --- a/pyrit/prompt_target/hack_a_prompt_target.py +++ b/pyrit/prompt_target/hack_a_prompt_target.py @@ -4,7 +4,9 @@ import enum import json import logging -from typing import Any +from collections.abc import Mapping +from types import MappingProxyType +from typing import Any, ClassVar from pyrit.common import default_values, net_utility from pyrit.models import ComponentIdentifier, Message, construct_response_from_request @@ -278,6 +280,9 @@ class HackAPromptTarget(PromptTarget): BASE_URL: str = "https://www.hackaprompt.com" session_id_environment_variable: str = "HACK_A_PROMPT_SESSION_ID" cookie_environment_variable: str = "HACK_A_PROMPT_COOKIE" + server_fallback_parameters: ClassVar[Mapping[str, str | None]] = MappingProxyType( + {"session_id": None, "cookie": None} + ) # Prefixes of the data-stream parts this target reads, see _parse_stream. # ``0:`` carries model text and ``3:`` carries an error the server chose to diff --git a/pyrit/prompt_target/http_target/http_target.py b/pyrit/prompt_target/http_target/http_target.py index 5b7e616f58..37b2f31314 100644 --- a/pyrit/prompt_target/http_target/http_target.py +++ b/pyrit/prompt_target/http_target/http_target.py @@ -67,6 +67,8 @@ class HTTPTarget(PromptTarget): """ + not_saveable_reason: ClassVar[str | None] = "Its raw HTTP request can carry credentials." + _CREDENTIAL_HEADERS: ClassVar[frozenset[str]] = frozenset( { "authorization", diff --git a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py index 3488893403..5592757e29 100644 --- a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py +++ b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py @@ -5,8 +5,10 @@ import json import logging import warnings +from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from pathlib import Path +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, cast from pyrit.common import default_values @@ -41,6 +43,10 @@ class HuggingFaceChatTarget(PromptTarget): ) code_execution_parameter: ClassVar[str | None] = "trust_remote_code" + server_fallback_parameters: ClassVar[Mapping[str, str | None]] = MappingProxyType({"hf_access_token": None}) + not_saveable_reason: ClassVar[str | None] = ( + "It needs a Hugging Face token or model files on this server, and saved targets can hold neither yet." + ) # Class-level cache for model and tokenizer _cached_model: Any = None diff --git a/pyrit/prompt_target/litellm_chat_target.py b/pyrit/prompt_target/litellm_chat_target.py index b6056d31b0..43a68c35b5 100644 --- a/pyrit/prompt_target/litellm_chat_target.py +++ b/pyrit/prompt_target/litellm_chat_target.py @@ -5,7 +5,8 @@ import logging import os from collections.abc import Awaitable, Callable, Mapping, MutableSequence, Sequence -from typing import Any, NoReturn, cast +from types import MappingProxyType +from typing import Any, ClassVar, NoReturn, cast from pyrit.auth import ensure_async_token_provider from pyrit.exceptions import ( @@ -202,6 +203,8 @@ class LiteLLMChatTarget(PromptTarget): sends a fresh ``traceparent`` in ``extra_headers``. """ + server_fallback_parameters: ClassVar[Mapping[str, str | None]] = MappingProxyType({"api_key": None}) + # Fallback only. The real per-instance configuration is normally derived from LiteLLM's # model metadata at construction time (see ``_derive_capabilities_from_litellm``). _SUPPORTS_TOOL_CALL_HISTORY = True diff --git a/pyrit/prompt_target/openai/openai_target.py b/pyrit/prompt_target/openai/openai_target.py index e784618d58..d6fdf1aa89 100644 --- a/pyrit/prompt_target/openai/openai_target.py +++ b/pyrit/prompt_target/openai/openai_target.py @@ -5,8 +5,9 @@ import logging import re from abc import abstractmethod -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping from copy import deepcopy +from types import MappingProxyType from typing import Any, ClassVar, TypedDict from urllib.parse import urlparse @@ -159,6 +160,13 @@ class OpenAITarget(PromptTarget): # endpoint (see ``is_azure_openai_endpoint``), and OpenAI-compatible servers # such as Ollama or vLLM may need no key at all. supported_auth_modes: ClassVar[tuple[AuthMode, ...]] = ("api_key", "identity", "none") + # Empty headers read ``OPENAI_ADDITIONAL_REQUEST_HEADERS`` unless the mode is "none"; "{}" sends none. + server_fallback_parameters: ClassVar[Mapping[str, str | None]] = MappingProxyType( + {"api_key": None, "headers": "{}"} + ) + destination_parameters: ClassVar[Mapping[str, bool]] = MappingProxyType({"endpoint": True}) + # The OpenAI SDK adds these headers to every client it creates. + server_header_variables: ClassVar[tuple[str, ...]] = ("OPENAI_CUSTOM_HEADERS",) model_name_environment_variable: str endpoint_environment_variable: str diff --git a/pyrit/prompt_target/prompt_shield_target.py b/pyrit/prompt_target/prompt_shield_target.py index 6c6b83ed4f..407ed16ce5 100644 --- a/pyrit/prompt_target/prompt_shield_target.py +++ b/pyrit/prompt_target/prompt_shield_target.py @@ -3,7 +3,8 @@ import json import logging -from collections.abc import Callable +from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Literal, assert_never from pyrit.auth import ( @@ -57,6 +58,8 @@ class PromptShieldTarget(PromptTarget): # A subscription key is the "api_key"; auth_mode="identity" mints an Entra ID token # provider for a recognized Azure Content Safety endpoint instead. supported_auth_modes: ClassVar[tuple[AuthMode, ...]] = ("api_key", "identity") + server_fallback_parameters: ClassVar[Mapping[str, str | None]] = MappingProxyType({"api_key": None}) + destination_parameters: ClassVar[Mapping[str, bool]] = MappingProxyType({"endpoint": True}) _endpoint: str _api_key: str | Callable[[], str] | None diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 4bba0b7d05..176af12a04 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -7,7 +7,7 @@ import pathlib import uuid from enum import IntEnum -from typing import Any +from typing import Any, ClassVar import httpx import websockets @@ -94,6 +94,7 @@ class WebSocketCopilotTarget(PromptTarget): RESPONSE_TIMEOUT_SECONDS: int = 60 CONNECTION_TIMEOUT_SECONDS: int = 30 + not_saveable_reason: ClassVar[str | None] = "It signs in with the Copilot account this server is set up with." _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( capabilities=TargetCapabilities( supports_multi_turn=True, diff --git a/pyrit/registry/components/attack_technique_registry.py b/pyrit/registry/components/attack_technique_registry.py index c7f449f174..0675f95d85 100644 --- a/pyrit/registry/components/attack_technique_registry.py +++ b/pyrit/registry/components/attack_technique_registry.py @@ -31,7 +31,12 @@ from pyrit.models import AttackIdentifier, ComponentType from pyrit.models.component_spec import SourceInstanceSpec -from pyrit.registry.instance_registry import DefaultInstanceRegistry, RegistryNameTakenError +from pyrit.registry.instance_registry import ( + DefaultInstanceRegistry, + RegistryNameTakenError, + _record_saved_entry_refusal, + refuse_saved_entry, +) from pyrit.registry.registry import Registry from pyrit.registry.registry_metadata import RegistryMetadata @@ -143,10 +148,14 @@ class TechniqueAdmissionError(ValueError): class TechniqueInstanceRegistry(DefaultInstanceRegistry["AttackTechniqueFactory"]): """Factory storage with an atomic runtime-admission path.""" - def register_runtime(self, factory: AttackTechniqueFactory) -> None: + def register_runtime(self, factory: AttackTechniqueFactory, *, metadata: dict[str, Any] | None = None) -> None: """ Check selector collisions while holding the same lock as registration. + Args: + factory (AttackTechniqueFactory): The technique to register. + metadata (dict[str, Any] | None): Optional per-entry metadata. Defaults to None. + Raises: TechniqueSelectorConflictError: If the name or a tag conflicts with another technique's name or tag. TechniqueAdmissionError: If a tag repeats the technique's own name, a selector can't become a scenario @@ -162,6 +171,8 @@ def register_runtime(self, factory: AttackTechniqueFactory) -> None: tags = {tag.casefold(): tag for entry in entries for tag in entry.instance.technique_tags} folded_name = factory.name.casefold() if folded_name in names: + # Code that catches this error would leave a saved technique standing in, so the clash is recorded. + _record_saved_entry_refusal(self.get_entry(names[folded_name])) raise TechniqueSelectorConflictError( f"Technique '{factory.name}' already exists (names are case-insensitive)" ) @@ -186,7 +197,7 @@ def register_runtime(self, factory: AttackTechniqueFactory) -> None: ) except ValueError as error: raise TechniqueAdmissionError(str(error)) from error - self.register(factory, name=factory.name, tags=factory.technique_tags) + self.register(factory, name=factory.name, tags=factory.technique_tags, metadata=metadata) class AttackTechniqueRegistry(Registry["AttackTechniqueFactory", AttackTechniqueMetadata]): @@ -769,8 +780,13 @@ def register_from_factories( factories (list[AttackTechniqueFactory]): Self-describing factories to register. Each factory's ``name`` and ``technique_tags`` properties are used directly. + + Raises: + SavedEntryReplacementError: If a factory's name belongs to a technique rebuilt + from a saved recipe, which skipping would quietly use in its place. """ for factory in factories: + refuse_saved_entry(self.instances.get_entry(factory.name)) if factory.name not in self.instances: tags: dict[str, str] = dict.fromkeys(factory.technique_tags, "") self.register_technique( diff --git a/pyrit/registry/instance_registry.py b/pyrit/registry/instance_registry.py index d83c13b284..4c1067294b 100644 --- a/pyrit/registry/instance_registry.py +++ b/pyrit/registry/instance_registry.py @@ -23,11 +23,15 @@ from __future__ import annotations import threading +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable +from pyrit.models.catalog import RegistryEntryOrigin + if TYPE_CHECKING: - from collections.abc import Callable, Iterator + from collections.abc import Callable, Generator, Iterator, Mapping from pyrit.models import ComponentIdentifier from pyrit.registry.tag_query import TagQuery @@ -43,6 +47,18 @@ def get_identifier(self) -> ComponentIdentifier: T = TypeVar("T", bound=SupportsIdentifier) # The type of items stored +#: Entry metadata naming the kind of saved recipe an entry was rebuilt from, such as ``"targets"``. +SAVED_ENTRY_METADATA_KEY = "saved_recipe_kind" + +#: Entry metadata marking an entry that the server's startup scripts registered, so its name is registered again +#: after a restart. Entries with neither this nor the saved marker are gone after a restart. +CONFIGURATION_ENTRY_METADATA_KEY = "configuration_entry" + +_saved_entry_refusals: ContextVar[list[tuple[str, str]] | None] = ContextVar("_saved_entry_refusals", default=None) +_registration_origin: ContextVar[RegistryEntryOrigin | None] = ContextVar("_registration_origin", default=None) +# The kind and name of the saved entry being rebuilt in this context, so registering that name marks the entry. +_saved_entry_registration: ContextVar[tuple[str, str] | None] = ContextVar("_saved_entry_registration", default=None) + class RegistryNameTakenError(ValueError): """A name is already registered, so another entry can't take it.""" @@ -70,6 +86,148 @@ class RegistryEntry(Generic[T]): metadata: dict[str, Any] = field(default_factory=dict) +class SavedEntryReplacementError(RuntimeError): + """ + Raised when code registers an entry under the name of an entry rebuilt from a saved recipe. + + It isn't a ``ValueError``, so a startup script that skips the entries it can't build doesn't + swallow it. + """ + + def __init__(self, *, kind: str, name: str) -> None: + """ + Describe the saved entry that would have been replaced. + + Args: + kind (str): The saved recipe kind, such as ``"targets"``. + name (str): The entry's name. + """ + super().__init__(f"Saved {kind.removesuffix('s')} '{name}' can't be replaced.") + self.kind = kind + self.name = name + + +def refuse_saved_entry(entry: RegistryEntry[Any] | None) -> None: + """ + Refuse to replace, or stand in for, an entry rebuilt from a saved recipe. + + Args: + entry (RegistryEntry[Any] | None): The entry registered under the name, if any. + + Raises: + SavedEntryReplacementError: If the entry was rebuilt from a saved recipe. + """ + kind = _record_saved_entry_refusal(entry) + if entry is not None and kind is not None: + raise SavedEntryReplacementError(kind=kind, name=entry.name) + + +def _record_saved_entry_refusal(entry: RegistryEntry[Any] | None) -> str | None: + """ + Record a refused use of an entry rebuilt from a saved recipe, when refusals are being recorded. + + Args: + entry (RegistryEntry[Any] | None): The entry registered under the name, if any. + + Returns: + str | None: The entry's saved recipe kind, or None when it wasn't rebuilt from a saved recipe. + """ + if entry is None or SAVED_ENTRY_METADATA_KEY not in entry.metadata: + return None + kind = str(entry.metadata[SAVED_ENTRY_METADATA_KEY]) + refusals = _saved_entry_refusals.get() + if refusals is not None: + refusals.append((kind, entry.name)) + return kind + + +def registry_entry_origin(metadata: Mapping[str, Any]) -> RegistryEntryOrigin: + """ + Tell where a registered entry comes from, from its metadata. + + Args: + metadata (Mapping[str, Any]): The entry's registry metadata. + + Returns: + RegistryEntryOrigin: Saved when rebuilt from a saved recipe, configuration when the server's startup scripts + registered it, and temporary otherwise, such as an entry created through the API without saving it or + one a scan's launch scripts registered. + """ + if SAVED_ENTRY_METADATA_KEY in metadata: + return RegistryEntryOrigin.SAVED + if metadata.get(CONFIGURATION_ENTRY_METADATA_KEY): + return RegistryEntryOrigin.CONFIGURATION + return RegistryEntryOrigin.TEMPORARY + + +@contextmanager +def registration_origin(origin: RegistryEntryOrigin) -> Generator[None, None, None]: + """ + Mark what code run in this context registers as coming from the server's startup scripts or a scan. + + Entries registered while the startup scripts run are marked as configuration. Classes that first + appear while a scan's launch scripts run are temporary, since those scripts don't run at startup. + + Args: + origin (RegistryEntryOrigin): ``CONFIGURATION`` for startup scripts, ``TEMPORARY`` for launch scripts. + + Yields: + None: Control while the scripts run. + """ + token = _registration_origin.set(origin) + try: + yield + finally: + _registration_origin.reset(token) + + +def current_registration_origin() -> RegistryEntryOrigin | None: + """ + Return the origin set by the innermost ``registration_origin`` around the caller, if any. + + Returns: + RegistryEntryOrigin | None: The origin, or None outside any. + """ + return _registration_origin.get() + + +@contextmanager +def record_saved_entry_refusals() -> Generator[list[tuple[str, str]], None, None]: + """ + Collect the saved entries that code run in this context tried to replace. + + The list holds every refusal, including ones the code caught itself. + + Yields: + list[tuple[str, str]]: The kind and name of each saved entry that was refused. + """ + refusals: list[tuple[str, str]] = [] + token = _saved_entry_refusals.set(refusals) + try: + yield refusals + finally: + _saved_entry_refusals.reset(token) + + +@contextmanager +def saved_entry_registration(*, kind: str, name: str) -> Generator[None, None, None]: + """ + Mark the entry registered under ``name`` in this context as rebuilt from a saved recipe. + + Args: + kind (str): The saved recipe kind, such as ``"targets"``. + name (str): The saved entry's name. + + Yields: + None: Control returns to the caller while the mark applies. + """ + token = _saved_entry_registration.set((kind, name)) + try: + yield + finally: + _saved_entry_registration.reset(token) + + @runtime_checkable class InstanceRegistry(Protocol[T]): """ @@ -278,6 +436,9 @@ def register( """ Register a pre-configured instance. + An entry registered while the server's startup scripts run is marked as configuration + (see ``registration_origin``). + Args: instance (T): The instance to register. name (str | None): The registry name. Defaults to the instance's @@ -292,6 +453,8 @@ def register( ``instance`` is not of that type. ValueError: If the name is reserved. RegistryNameTakenError: If the name is already registered and ``replace`` is False. + SavedEntryReplacementError: If ``replace`` would replace an entry rebuilt + from a saved recipe. """ expected_type = self._resolve_instance_type() if expected_type is not None and not isinstance(instance, expected_type): @@ -302,17 +465,25 @@ def register( if name is None: name = instance.get_identifier().unique_name + entry_metadata = dict(metadata or {}) + if _registration_origin.get() is RegistryEntryOrigin.CONFIGURATION: + entry_metadata[CONFIGURATION_ENTRY_METADATA_KEY] = True + restoring = _saved_entry_registration.get() + if restoring is not None and restoring[1] == name: + entry_metadata[SAVED_ENTRY_METADATA_KEY] = restoring[0] entry = RegistryEntry( name=name, instance=instance, tags=self._normalize_tags(tags), - metadata=metadata or {}, + metadata=entry_metadata, ) with self._lock: if name in self._reserved_names: raise ValueError(f"Instance name '{name}' is reserved") if not replace: self.validate_name_available(name) + else: + refuse_saved_entry(self._registry_items.get(name)) self._registry_items[name] = entry self._metadata_cache = None self._instance_version += 1 @@ -345,6 +516,8 @@ def validate_name_available(self, name: str) -> None: if name in self._reserved_names: raise ValueError(f"Instance name '{name}' is reserved") if name in self._registry_items: + # Code that catches this error would leave a saved entry standing in, so the clash is recorded. + _record_saved_entry_refusal(self._registry_items[name]) raise RegistryNameTakenError(f"Instance '{name}' already exists") def unregister(self, name: str, *, expected_entry: RegistryEntry[T] | None = None) -> T | None: diff --git a/pyrit/registry/registry.py b/pyrit/registry/registry.py index d5a8ca6e69..f9b74f8828 100644 --- a/pyrit/registry/registry.py +++ b/pyrit/registry/registry.py @@ -32,7 +32,13 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable -from pyrit.registry.instance_registry import DefaultInstanceRegistry, InstanceRegistry, SupportsIdentifier +from pyrit.models.catalog import RegistryEntryOrigin +from pyrit.registry.instance_registry import ( + DefaultInstanceRegistry, + InstanceRegistry, + SupportsIdentifier, + current_registration_origin, +) from pyrit.registry.registry_metadata import RegistryMetadata from pyrit.registry.resolution import ( derive_parameters, @@ -212,16 +218,17 @@ def __init__(self, *, lazy_discovery: bool = True) -> None: # so catalog reads (__contains__, get_class, ...) are never parked behind a build. self._metadata_build_lock = threading.RLock() self._classes: dict[str, type[T]] = {} + self._temporary_class_names: set[str] = set() self._metadata_cache: dict[str, MetadataT] | None = None # Bumped on every catalog mutation so an in-flight metadata build can detect a # registration that landed mid-build and retry instead of caching a stale snapshot. self._catalog_version = 0 self._discovered = False + self._discovering = False self._lazy_discovery = lazy_discovery if not lazy_discovery: - self._discover() - self._discovered = True + self._ensure_discovered() @classmethod def get_registry_singleton(cls) -> Self: @@ -281,7 +288,11 @@ def _ensure_discovered(self) -> None: """Ensure discovery has been performed. Runs discovery on first access.""" with self._catalog_lock: if not self._discovered: - self._discover() + self._discovering = True + try: + self._discover() + finally: + self._discovering = False self._discovered = True def _base_type(self) -> type[T]: @@ -598,10 +609,29 @@ def register_class(self, cls: type[T], *, name: str | None = None) -> None: if name is None: name = self._get_registry_name(cls) self._validate_class(cls) + if current_registration_origin() is RegistryEntryOrigin.TEMPORARY and not self._discovering: + # A class a scan's launch scripts add, or put in place of another, is gone after a restart. + # The classes discovery finds come back, so they are found first. + self._ensure_discovered() + if self._classes.get(name) is not cls: + self._temporary_class_names.add(name) self._classes[name] = cls self._metadata_cache = None self._catalog_version += 1 + def is_temporary_class(self, name: str) -> bool: + """ + Return whether a scan's launch scripts added the class under this name, so it's gone after a restart. + + Args: + name (str): The catalog name. + + Returns: + bool: True for a class a scan's launch scripts added, or put in place of another one. + """ + with self._catalog_lock: + return name in self._temporary_class_names + def get_class(self, name: str) -> type[T]: """ Get a registered class by name. diff --git a/pyrit/registry/registry_recipe_restore.py b/pyrit/registry/registry_recipe_restore.py index e892c207f8..03acbc0189 100644 --- a/pyrit/registry/registry_recipe_restore.py +++ b/pyrit/registry/registry_recipe_restore.py @@ -26,7 +26,8 @@ TechniqueRecipe, ) from pyrit.registry.file_document_storage import DocumentStorageUnavailableError -from pyrit.registry.instance_registry import RegistryNameTakenError +from pyrit.registry.instance_registry import RegistryNameTakenError, saved_entry_registration +from pyrit.registry.registry_recipe_rules import SavedEntryRefusedError, check_saved_entry_params from pyrit.registry.registry_recipe_storage import RecipeListing, RegistryRecipeStorage, safe_failure_text if TYPE_CHECKING: @@ -73,13 +74,13 @@ class RegistryRecipeRestorer: is built after the saved entries it names. The caller supplies the builder, so an entry is built the way the caller's create builds one; the backend's builder calls the API's create for each kind. A file never chooses what gets imported: classes and attack types are looked up among the ones already - registered. Server paths are refused as they are for API input, and a recipe never holds a value for - a parameter its class declares as a credential. Entries that can't come back are reported with the - reason instead of being skipped. + registered. Server paths are refused as they are for API input, and ``check_saved_entry_params`` + refuses a recipe that stores a credential, or a target recipe that would let its target use one from + this server or send requests where this server's settings say. Entries that can't come back are + reported with the reason instead of being skipped. - Stored files are otherwise trusted the way configuration is. Rules for credentials a class - reads from server settings, and for fields that can carry a credential without being declared - one, need each class to declare them, so they arrive with saving and apply to restore too. + Stored files are otherwise trusted the way configuration is. Converters and scorers don't + declare the credentials they read from server settings yet, so those rules cover targets only. """ _KIND_BY_COMPONENT_TYPE: ClassVar[dict[ComponentType, RegistryRecipeKind]] = { @@ -241,7 +242,7 @@ def _plan( if isinstance(recipe, TechniqueRecipe) else self._component_references(ref=stored.ref, recipe=recipe) ) - except RecipeRejectedError as rejection: + except (RecipeRejectedError, SavedEntryRefusedError) as rejection: return self._failure(stored=stored, reason=rejection.reason, message=str(rejection)) except Exception as error: # A bug here must not stop the others from being checked, or blame the file. causes[(stored.ref.kind, stored.ref.name)] = error @@ -270,8 +271,9 @@ def _component_references(self, *, ref: SavedEntryRef, recipe: ComponentRecipe) list[SavedEntryRef]: The referenced entries, in parameter order. Raises: - RecipeRejectedError: If the class is unknown or is now one from outside PyRIT, the recipe holds - a credential or sets a server file or folder, or the name is taken or can't be used. + RecipeRejectedError: If the class is unknown or is now one from outside PyRIT, the recipe sets a + server file or folder, or the name is taken or can't be used. + SavedEntryRefusedError: If the class can't be saved, or the recipe holds or would use a credential. """ registry = self._get_registry(ref.kind) try: @@ -288,7 +290,7 @@ def _component_references(self, *, ref: SavedEntryRef, recipe: ComponentRecipe) ) raise RecipeRejectedError(reason=RestoreFailureReason.MODULE_MISMATCH, message=message) parameters = self._class_parameters(registry=registry, name=recipe.class_name, cls=cls) - self._check_stored_credentials(parameters=parameters, params=recipe.params) + check_saved_entry_params(cls=cls, parameters=parameters, params=recipe.params, auth_mode=recipe.auth_mode) # As in the target API, a target may let callers set some path parameters, such as a working folder. allowed_paths = ( getattr(cls, "caller_path_parameters", frozenset()) @@ -326,7 +328,8 @@ def _technique_references(self, *, ref: SavedEntryRef, recipe: TechniqueRecipe) list[SavedEntryRef]: The adversarial target, converters, and attack parameter references. Raises: - RecipeRejectedError: If the attack type isn't registered, or the recipe holds a credential. + RecipeRejectedError: If the attack type isn't registered. + SavedEntryRefusedError: If the recipe holds a credential. """ from pyrit.registry.components.attack_registry import AttackRegistry @@ -341,7 +344,7 @@ def _technique_references(self, *, ref: SavedEntryRef, recipe: TechniqueRecipe) if recipe.adversarial_chat is not None: references.append(SavedEntryRef(kind=RegistryRecipeKind.TARGETS, name=recipe.adversarial_chat)) parameters = self._class_parameters(registry=attack_registry, name=recipe.type, cls=attack_class) - self._check_stored_credentials(parameters=parameters, params=recipe.params) + check_saved_entry_params(cls=attack_class, parameters=parameters, params=recipe.params) return references + self._parameter_references(parameters=parameters, params=recipe.params) @staticmethod @@ -355,25 +358,6 @@ def _class_parameters(*, registry: Registry[Any, Any], name: str, cls: type) -> metadata: RegistryMetadata = registry.get_registered_class_metadata(name) or registry.get_class_metadata(cls) return metadata.parameters - @staticmethod - def _check_stored_credentials(*, parameters: Sequence[Parameter], params: Mapping[str, Any]) -> None: - """ - Refuse a recipe that stores a value for a parameter its class declares as a credential. - - Raises: - RecipeRejectedError: If a declared credential parameter has a value. - """ - stored = sorted( - parameter.name - for parameter in parameters - if parameter.sensitive and params.get(parameter.name) not in (None, "") - ) - if stored: - raise RecipeRejectedError( - reason=RestoreFailureReason.UNSUPPORTED_CREDENTIAL, - message=f"Saved entries can't hold credentials ('{stored[0]}').", - ) - def _parameter_references( self, *, parameters: Sequence[Parameter], params: Mapping[str, Any] ) -> list[SavedEntryRef]: @@ -518,8 +502,9 @@ async def _build_async( from pyrit.registry.components.attack_technique_registry import TechniqueAdmissionError try: - await build_async(stored) - except RecipeRejectedError as rejection: + with saved_entry_registration(kind=stored.ref.kind.value, name=stored.ref.name): + await build_async(stored) + except (RecipeRejectedError, SavedEntryRefusedError) as rejection: return self._failure(stored=stored, reason=rejection.reason, message=str(rejection)) except RegistryNameTakenError as error: return self._failure(stored=stored, reason=RestoreFailureReason.NAME_CLASH, message=str(error)) diff --git a/pyrit/registry/registry_recipe_rules.py b/pyrit/registry/registry_recipe_rules.py new file mode 100644 index 0000000000..169d894c23 --- /dev/null +++ b/pyrit/registry/registry_recipe_rules.py @@ -0,0 +1,282 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""What a saved registry entry may hold, checked before a saved entry is built.""" + +from __future__ import annotations + +import os +import re +import unicodedata +from typing import TYPE_CHECKING, Any +from urllib.parse import unquote, urlsplit + +from pyrit.models.catalog import RestoreFailureReason + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from pyrit.models import Parameter + +_URL_SCHEMES = frozenset({"http", "https", "ws", "wss"}) +_URL_QUERY_KEYS = frozenset({"api-version"}) +_QUERY_SEPARATORS = re.compile(r"[&;]") +_QUERY_VALUE = re.compile(r"[A-Za-z0-9._-]*") +_PATH_PARAMETERS = re.compile(r"[;&=]") +_MAX_URL_READINGS = 8 + + +class SavedEntryRefusedError(ValueError): + """A recipe holds, or would make its entry use, a credential, or builds a class that can't be saved.""" + + def __init__(self, *, reason: RestoreFailureReason, message: str) -> None: + """ + Initialize the error with the reason to report. + + Args: + reason (RestoreFailureReason): Why the recipe is refused. + message (str): Why, in words safe to show users. It never includes a parameter's value. + """ + super().__init__(message) + self.reason = reason + + +def check_saved_entry_params( + *, cls: type, parameters: Sequence[Parameter], params: Mapping[str, Any], auth_mode: str | None = None +) -> None: + """ + Refuse a recipe that a saved entry of ``cls`` can't hold. + + A saved entry is rebuilt from its recipe at every startup, so the recipe can't hold a credential + and can't leave the entry to use one from this server: an empty key the class would read from the + server's settings or identity, extra headers the server adds, a destination the server's settings + choose, or a credential inside a URL. A saved target authenticates with Microsoft Entra ID + (``auth_mode`` ``"identity"``) or needs no key. + + Args: + cls (type): The class the recipe builds. + parameters (Sequence[Parameter]): The class's constructor parameters. + params (Mapping[str, Any]): The recipe's parameter values. + auth_mode (str | None): The target recipe's auth mode. None means the API's default, ``"api_key"``. + + Raises: + SavedEntryRefusedError: If the class can't be saved, or the values hold or would use a credential. + """ + _check_class(cls) + _check_server_headers(cls) + _check_stored_credentials(parameters=parameters, params=params) + _check_server_fallbacks(cls=cls, params=params, auth_mode=auth_mode) + _check_destinations(cls=cls, params=params) + + +def _check_class(cls: type) -> None: + """ + Refuse a class that a saved entry can't build. + + Raises: + SavedEntryRefusedError: If the class needs the server's upload folder or declares why it can't be saved. + """ + if getattr(cls, "upload_directory_parameter", None) is not None: + message = f"'{cls.__name__}' uploads files from a folder the server chooses, which a saved entry can't use yet." + raise SavedEntryRefusedError(reason=RestoreFailureReason.UNSUPPORTED_CLASS, message=message) + reason = getattr(cls, "not_saveable_reason", None) + if reason is not None: + raise SavedEntryRefusedError( + reason=RestoreFailureReason.UNSUPPORTED_CLASS, message=f"'{cls.__name__}' can't be saved. {reason}" + ) + + +def _check_server_headers(cls: type) -> None: + """ + Refuse a class whose client would send headers this server sets in its environment. + + Raises: + SavedEntryRefusedError: If an environment variable the class's client reads headers from is set. + """ + for variable in getattr(cls, "server_header_variables", ()): + if os.environ.get(variable): + message = ( + f"This server sets {variable}, and '{cls.__name__}' would send its headers with every request. " + "A saved target can't use this server's settings, so it isn't built while that is set." + ) + raise SavedEntryRefusedError(reason=RestoreFailureReason.UNSUPPORTED_CREDENTIAL, message=message) + + +def _check_stored_credentials(*, parameters: Sequence[Parameter], params: Mapping[str, Any]) -> None: + """ + Refuse a value for a parameter the class declares as a credential. + + Raises: + SavedEntryRefusedError: If a declared credential parameter has a value. + """ + stored = sorted( + parameter.name + for parameter in parameters + if parameter.sensitive and params.get(parameter.name) not in (None, "") + ) + if stored: + raise SavedEntryRefusedError( + reason=RestoreFailureReason.UNSUPPORTED_CREDENTIAL, + message=f"Saved entries can't hold credentials ('{stored[0]}').", + ) + + +def _check_server_fallbacks(*, cls: type, params: Mapping[str, Any], auth_mode: str | None) -> None: + """ + Refuse values that would let the entry use a credential from this server. + + Raises: + SavedEntryRefusedError: If the class doesn't support the auth mode, or a parameter the class reads + from the server is left empty in an auth mode that reads it, or holds a value other than the + one that turns the server's value off. + """ + if auth_mode is None: + auth_mode = "api_key" + supported_auth_modes = getattr(cls, "supported_auth_modes", None) + if supported_auth_modes is not None and not supported_auth_modes: + # A class that declares no auth modes, such as TextTarget, takes no credentials of its own. + supported_auth_modes = ("api_key",) + if supported_auth_modes is not None and auth_mode not in supported_auth_modes: + modes = ", ".join(f"'{mode}'" for mode in supported_auth_modes) + raise SavedEntryRefusedError( + reason=RestoreFailureReason.UNSUPPORTED_CREDENTIAL, + message=f"'{cls.__name__}' only supports auth_mode {modes}.", + ) + fallbacks: Mapping[str, str | None] = getattr(cls, "server_fallback_parameters", {}) + for name, turned_off in fallbacks.items(): + value = params.get(name) + if value in (None, ""): + if auth_mode == "none" or (auth_mode == "identity" and turned_off is None): + continue + message = f"'{name}' is empty, so '{cls.__name__}' can use this server's credential instead." + if turned_off is not None: + message += f" Saved entries set it to '{turned_off}'." + raise SavedEntryRefusedError(reason=RestoreFailureReason.UNSUPPORTED_CREDENTIAL, message=message) + if value != turned_off: + raise SavedEntryRefusedError( + reason=RestoreFailureReason.UNSUPPORTED_CREDENTIAL, + message=f"'{name}' can carry a credential, which saved entries can't hold.", + ) + + +def _check_destinations(*, cls: type, params: Mapping[str, Any]) -> None: + """ + Refuse a destination that this server's settings would choose, or one that can carry a credential. + + Raises: + SavedEntryRefusedError: If a destination that would take this server's setting is empty or isn't an + http(s) or ws(s) address, or a destination can carry a credential. + """ + destinations: Mapping[str, bool] = getattr(cls, "destination_parameters", {}) + for name, takes_server_setting in destinations.items(): + value = params.get(name) + if value is None or (isinstance(value, str) and not value.strip()): + if takes_server_setting: + message = ( + f"'{name}' is empty, so '{cls.__name__}' would send requests where this server's settings " + "say, with any credential in them. A saved target sets it." + ) + raise SavedEntryRefusedError(reason=RestoreFailureReason.UNSUPPORTED_CREDENTIAL, message=message) + continue + if takes_server_setting and not (isinstance(value, str) and _is_web_address(value)): + raise SavedEntryRefusedError( + reason=RestoreFailureReason.INVALID_RECIPE, + message=f"'{name}' must be a full http(s) or ws(s) address in a saved target.", + ) + if not isinstance(value, str) or _can_carry_credential(value, address=takes_server_setting): + message = ( + f"'{name}' can carry a credential in a URL, which saved entries can't hold: a user name, " + "password, SAS token, fragment, ';', '&' or '=' in the path, or a query other than 'api-version'." + ) + raise SavedEntryRefusedError(reason=RestoreFailureReason.UNSUPPORTED_CREDENTIAL, message=message) + + +def _is_web_address(text: str) -> bool: + """ + Return whether text, as written, is an http(s) or ws(s) address with a host. + + Returns: + bool: True for an address such as ``https://contoso.openai.azure.com/openai/v1``. + """ + try: + parts = urlsplit(text.strip()) + except ValueError: + return False + return parts.scheme in _URL_SCHEMES and bool(parts.netloc) + + +def _can_carry_credential(text: str, *, address: bool) -> bool: + """ + Return whether a URL or path can carry a credential, in any reading of it. + + A value whose percent-encoding or look-alike characters don't settle within a few readings counts as + one that can, since an encoded token is still usable. + + Args: + text (str): The URL or path. + address (bool): Whether it must be a full http(s) or ws(s) address rather than a path. + + Returns: + bool: True when a reading isn't a plain address or path, or holds a SAS token, user info, a + fragment, a path parameter, or a query other than ``api-version``. + """ + readings = _readings(text) + return readings is None or any(_reading_can_carry_credential(reading, address=address) for reading in readings) + + +def _reading_can_carry_credential(text: str, *, address: bool) -> bool: + """ + Return whether one reading of a URL or path can carry a credential. + + An address has an http(s) or ws(s) scheme and a host without user info. A path has neither, and no + query. Neither may have a fragment or ';', '&' or '=' in its path, and an address's query may hold + only ``api-version``. + + Args: + text (str): One reading of the URL or path. + address (bool): Whether it must be a full address rather than a path. + + Returns: + bool: True when the reading isn't a plain address or path, or can carry a credential. + """ + if "sig=" in text.casefold(): + return True + try: + parts = urlsplit(text.strip()) + except ValueError: + return True + if address or parts.scheme or parts.netloc: + if parts.scheme not in _URL_SCHEMES or not parts.netloc or "@" in parts.netloc: + return True + elif parts.query: + return True + if parts.fragment or _PATH_PARAMETERS.search(parts.path): + return True + return not all(_is_version_query(segment) for segment in _QUERY_SEPARATORS.split(parts.query) if segment) + + +def _is_version_query(segment: str) -> bool: + """ + Return whether one query segment is ``api-version`` with a version-shaped value. + + Returns: + bool: True for ``api-version=2024-10-21`` and similar. + """ + key, _, value = segment.partition("=") + return key.strip().casefold() in _URL_QUERY_KEYS and _QUERY_VALUE.fullmatch(value.strip()) is not None + + +def _readings(text: str) -> list[str] | None: + """ + List the text as written and as each round of look-alike replacement and percent-decoding reads it. + + Returns: + list[str] | None: The readings until they stop changing, or None if they still change after the limit. + """ + readings = [text] + while len(readings) <= _MAX_URL_READINGS: + reading = unquote(unicodedata.normalize("NFKC", readings[-1])) + if reading == readings[-1]: + return readings + readings.append(reading) + return None diff --git a/pyrit/registry/registry_recipe_storage.py b/pyrit/registry/registry_recipe_storage.py index 616f58ce58..5e5b81b6d0 100644 --- a/pyrit/registry/registry_recipe_storage.py +++ b/pyrit/registry/registry_recipe_storage.py @@ -188,6 +188,56 @@ def list_recipes(self) -> RecipeListing: causes=MappingProxyType(causes), ) + def create_recipe(self, *, name: str, recipe: ComponentRecipe | TechniqueRecipe) -> str: + """ + Save a new recipe under a name that no saved recipe of this kind has. + + Args: + name (str): The entry's name, which is also its document name. + recipe (ComponentRecipe | TechniqueRecipe): The recipe to save. + + Returns: + str: The version of the saved document. + + Raises: + DocumentConflictError: If a recipe with this name is already saved. + ValueError: If the name isn't a legal registry name. + DocumentStorageUnavailableError: If the folder or container can't be written, or a local write + could not acquire the document lock. + """ + return self._save_document_conditional(name=name, content=self._serialize(recipe), expected_version=None) + + def delete_recipe(self, name: str) -> bool: + """ + Delete a saved recipe, such as one whose entry couldn't be registered after it was saved. + + Args: + name (str): The entry's name. + + Returns: + bool: True if the recipe existed and was deleted, False if it was already gone. + + Raises: + ValueError: If the name isn't a legal registry name. + DocumentStorageUnavailableError: If the folder or container can't be written, or a local delete + could not acquire the document lock. + """ + return self._delete_document(name) + + @staticmethod + def _serialize(recipe: ComponentRecipe | TechniqueRecipe) -> bytes: + """ + Write a recipe as JSON, labeled with the format version this PyRIT writes. + + A recipe built in memory reads as version 1 unless it says otherwise, so the label is set here; + otherwise, once the version moves past 1, new files would be labeled with the old version. + + Returns: + bytes: The document. + """ + current = recipe.model_copy(update={"schema_version": RECIPE_SCHEMA_VERSION}) + return current.model_dump_json(indent=2, exclude_none=True).encode("utf-8") + def _parse_recipe( self, *, name: str, content: bytes, causes: dict[str, BaseException] | None = None ) -> StoredRecipe | RestoreFailure: diff --git a/pyrit/setup/initializers/scorers.py b/pyrit/setup/initializers/scorers.py index 2dcdda7c35..aabb077b7e 100644 --- a/pyrit/setup/initializers/scorers.py +++ b/pyrit/setup/initializers/scorers.py @@ -26,6 +26,7 @@ from pyrit.common.async_compatibility import run_legacy_sync_async from pyrit.models import Parameter, SeedPrompt from pyrit.registry import ScorerRegistry, TargetRegistry +from pyrit.registry.instance_registry import refuse_saved_entry from pyrit.score import ( AzureContentFilterScorer, DecisionScorer, @@ -513,6 +514,8 @@ def _tag_best_per_category(self) -> None: for best_tag, (preferred_name, category_tag) in self._PREFERRED_BEST.items(): entry = scorer_registry.instances.get_entry(preferred_name) if entry is not None: + # A saved scorer under this name would stand in for the one this script couldn't build. + refuse_saved_entry(entry) scorer_registry.instances.add_tags(name=preferred_name, tags=[best_tag]) logger.info(f"Tagged {preferred_name} as {best_tag}") continue @@ -635,7 +638,10 @@ def _tag_best_objective(self) -> None: # Fall back: prefer scale_and_refusal, then first composite best_tags: list[str] = [ScorerInitializerTags.BEST_OBJECTIVE, ScorerInitializerTags.DEFAULT_OBJECTIVE_SCORER] - if scorer_registry.instances.get_entry(self.SCALE_AND_REFUSAL): + scale_and_refusal = scorer_registry.instances.get_entry(self.SCALE_AND_REFUSAL) + if scale_and_refusal is not None: + # A saved scorer under this name would stand in for the one this script couldn't build. + refuse_saved_entry(scale_and_refusal) fallback_name = self.SCALE_AND_REFUSAL else: composites = scorer_registry.instances.get_by_tag(tag=ScorerInitializerTags.OBJECTIVE_COMPOSITE) diff --git a/pyrit/setup/initializers/targets.py b/pyrit/setup/initializers/targets.py index 35d05043fe..92c00655ed 100644 --- a/pyrit/setup/initializers/targets.py +++ b/pyrit/setup/initializers/targets.py @@ -37,6 +37,7 @@ RoundRobinTarget, ) from pyrit.registry import TargetRegistry +from pyrit.registry.instance_registry import refuse_saved_entry from pyrit.setup.pyrit_initializer import PyRITInitializer logger = logging.getLogger(__name__) @@ -827,6 +828,8 @@ def _auto_group_targets(self) -> None: rr_name = generate_rr_name(key) if rr_name in registry.instances: + # Skipping would leave a saved target standing in for this group. + refuse_saved_entry(registry.instances.get_entry(rr_name)) logger.debug(f"Skipping auto-group {rr_name}: name already exists in registry") continue diff --git a/pyrit/setup/initializers/tools.py b/pyrit/setup/initializers/tools.py index 0c3dbe0554..d6f4a23f78 100644 --- a/pyrit/setup/initializers/tools.py +++ b/pyrit/setup/initializers/tools.py @@ -8,6 +8,7 @@ from pyrit.executor.attack import AttackConverterConfig, PromptSendingAttack from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.registry import AttackTechniqueRegistry, ConverterRegistry +from pyrit.registry.instance_registry import refuse_saved_entry from pyrit.scenario import AttackTechniqueFactory from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -18,8 +19,11 @@ class ToolsInitializer(PyRITInitializer): async def initialize_async(self) -> None: """Register tool-use techniques using the configured adversarial chat.""" registry = AttackTechniqueRegistry.get_registry_singleton() + # Returning early would leave a saved technique standing in for this one. + refuse_saved_entry(registry.instances.get_entry("role_play_tool_calls")) if "role_play_tool_calls" in registry.instances: return + refuse_saved_entry(ConverterRegistry.get_registry_singleton().instances.get_entry("role_play_tool_calls")) converter = await asyncio.to_thread( ConverterRegistry.get_registry_singleton().create_instance, "RolePlayToolCallConverter", diff --git a/tests/end_to_end/conftest.py b/tests/end_to_end/conftest.py index 9200c85eae..ab1aa9ee78 100644 --- a/tests/end_to_end/conftest.py +++ b/tests/end_to_end/conftest.py @@ -13,6 +13,7 @@ from __future__ import annotations import asyncio +import json from pathlib import Path import pytest @@ -23,18 +24,25 @@ @pytest.fixture(scope="session", autouse=True) -def _pyrit_backend_server(): +def _pyrit_backend_server(tmp_path_factory: pytest.TempPathFactory): """Launch ``pyrit_backend`` for the duration of the e2e test session. Uses the same ``test_config.yaml`` the scenario tests pass to the CLI - client so server-side and client-side memory/config stay in sync. + client so server-side and client-side memory/config stay in sync. The + backend's copy also points saved targets, converters, scorers, and + techniques at a fresh folder, so a backend this fixture starts never reads + or writes the ones saved on the machine running it. Tear down via ``launcher.stop()`` so a lingering subprocess does not occupy port 8000 between local pytest runs. The launcher is a no-op if a backend is already healthy on the port (e.g. a developer has one running), so this fixture also remains friendly to local iteration. """ + folder = tmp_path_factory.mktemp("pyrit_backend") + config_file = folder / _E2E_CONFIG_FILE.name + recipes_source = json.dumps(str(folder / "registry_recipes")) + config_file.write_text(f"{_E2E_CONFIG_FILE.read_text()}\nregistry_recipes_source: {recipes_source}\n") launcher = ServerLauncher() - asyncio.run(launcher.start_async(config_file=_E2E_CONFIG_FILE)) + asyncio.run(launcher.start_async(config_file=config_file)) yield launcher.stop() diff --git a/tests/unit/backend/test_api_routes.py b/tests/unit/backend/test_api_routes.py index f0487227de..cdbf59942f 100644 --- a/tests/unit/backend/test_api_routes.py +++ b/tests/unit/backend/test_api_routes.py @@ -52,14 +52,16 @@ from pyrit.backend.routes import version as version_routes from pyrit.backend.routes.common import parse_label_query_params from pyrit.backend.routes.scores import _get_user_identifier +from pyrit.backend.routes.targets import _saved_entry_owner from pyrit.backend.services import AttackService from pyrit.backend.services.attack_service import AttackObjectiveConflictError from pyrit.backend.services.chat_scoring_service import ChatScoringError from pyrit.backend.services.manual_send_scheduler import ManualSendConflictError, ManualSendQueueFullError +from pyrit.backend.services.target_service import SavedTargetConflictError, SavedTargetStorageError from pyrit.memory import MemoryInterface from pyrit.memory.memory_interface import AttackStateConflictError from pyrit.models import AttackOutcome, ConverterIdentifier, MessagePiece, Score, TargetCapabilities, TargetIdentifier -from pyrit.models.catalog import TargetInstance +from pyrit.models.catalog import SavedEntryOwner, TargetInstance def _make_message_view(*, role: str = "user", value: str = "hello", sequence: int = 1) -> MessageView: @@ -1394,6 +1396,90 @@ def test_create_target_forbids_code_execution_for_non_admin(self, client: TestCl assert response.status_code == status.HTTP_403_FORBIDDEN assert "administrator access" in response.json()["detail"] + @pytest.mark.parametrize( + ("error", "status_code"), + [ + (SavedTargetConflictError("A saved target named 'contoso' already exists."), status.HTTP_409_CONFLICT), + (SavedTargetStorageError("Saved targets can't be written right now."), status.HTTP_503_SERVICE_UNAVAILABLE), + ], + ) + def test_create_target_maps_save_errors(self, client: TestClient, error: Exception, status_code: int) -> None: + with patch("pyrit.backend.routes.targets.get_target_service") as mock_get_service: + mock_get_service.return_value.create_target_async = AsyncMock(side_effect=error) + + response = client.post( + "/api/targets", json={"name": "contoso", "type": "TextTarget", "params": {}, "persist": True} + ) + + assert response.status_code == status_code + assert response.json()["detail"] == str(error) + + def test_create_target_passes_the_saved_entries_source_and_no_owner_when_sign_in_is_off( + self, client: TestClient + ) -> None: + with ( + patch("pyrit.backend.routes.targets.get_target_service") as mock_get_service, + patch.object(client.app.state, "registry_recipes_source", "/srv/registry-recipes", create=True), + ): + mock_get_service.return_value.create_target_async = AsyncMock(side_effect=ValueError("stop")) + + client.post("/api/targets", json={"name": "contoso", "type": "TextTarget", "params": {}, "persist": True}) + + kwargs = mock_get_service.return_value.create_target_async.call_args.kwargs + assert (kwargs["recipes_source"], kwargs["owner"], kwargs["request"].persist) == ( + "/srv/registry-recipes", + None, + True, + ) + + def test_create_target_records_the_signed_in_user_as_the_owner(self, client: TestClient) -> None: + owner = SavedEntryOwner(oid="user-oid", name="Ada") + with ( + patch("pyrit.backend.routes.targets.get_target_service") as mock_get_service, + patch("pyrit.backend.routes.targets._saved_entry_owner", return_value=owner), + ): + mock_get_service.return_value.create_target_async = AsyncMock(side_effect=ValueError("stop")) + + client.post("/api/targets", json={"name": "contoso", "type": "TextTarget", "params": {}, "persist": True}) + + assert mock_get_service.return_value.create_target_async.call_args.kwargs["owner"] == owner + + def test_create_target_records_no_owner_when_not_saving(self, client: TestClient) -> None: + with ( + patch("pyrit.backend.routes.targets.get_target_service") as mock_get_service, + patch("pyrit.backend.routes.targets._saved_entry_owner") as saved_entry_owner, + ): + mock_get_service.return_value.create_target_async = AsyncMock(side_effect=ValueError("stop")) + + client.post("/api/targets", json={"name": "contoso", "type": "TextTarget", "params": {}}) + + assert mock_get_service.return_value.create_target_async.call_args.kwargs["owner"] is None + saved_entry_owner.assert_not_called() + + @pytest.mark.parametrize( + ("fields", "message"), + [ + ({}, "persist requires a name"), + ({"name": "Contoso-GPT"}, "Invalid registry name"), + ({"name": "contoso", "register": False}, "persist requires register=true"), + ], + ) + def test_create_target_checks_what_persist_needs( + self, client: TestClient, fields: dict[str, object], message: str + ) -> None: + response = client.post("/api/targets", json={"type": "TextTarget", "params": {}, "persist": True, **fields}) + + assert response.status_code == status.HTTP_422_UNPROCESSABLE_CONTENT + assert message in json.dumps(response.json()) + + @pytest.mark.parametrize(("name", "expected"), [("Ada", "Ada"), ("", "ada@example.invalid")]) + def test_saved_entry_owner_is_the_signed_in_user(self, name: str, expected: str) -> None: + request = Request({"type": "http"}) + request.state.user = AuthenticatedUser(oid="user-oid", name=name, email="ada@example.invalid", groups=[]) + + assert _saved_entry_owner(request) == SavedEntryOwner(oid="user-oid", name=expected) + assert _saved_entry_owner(Request({"type": "http"})) is None + def test_create_target_passes_admin_flag_from_request(self, client: TestClient) -> None: """The route resolves administrator access and forwards it to the service.""" with ( diff --git a/tests/unit/backend/test_mappers.py b/tests/unit/backend/test_mappers.py index 337bfad634..f5c07b4a2c 100644 --- a/tests/unit/backend/test_mappers.py +++ b/tests/unit/backend/test_mappers.py @@ -44,6 +44,7 @@ PromptDataType, Score, ) +from pyrit.models.catalog import RegistryEntryOrigin from pyrit.prompt_target import ( CapabilityName, PromptTarget, @@ -1590,6 +1591,15 @@ def test_target_instance_reports_effective_history_replay(mode: str) -> None: assert not result.capabilities.supports_editable_history +class TestTargetInstanceOrigin: + @pytest.mark.usefixtures("patch_central_database") + def test_carries_the_origin_onto_the_instance(self) -> None: + result = target_object_to_instance("t-1", MockPromptTarget(), origin=RegistryEntryOrigin.SAVED) + + assert result.origin is RegistryEntryOrigin.SAVED + assert target_object_to_instance("t-1", MockPromptTarget()).origin is None + + class TestTargetObjectToInstance: """Tests for target_object_to_instance function.""" diff --git a/tests/unit/backend/test_runtime_lifecycle.py b/tests/unit/backend/test_runtime_lifecycle.py index 94c47971ac..62ed0cfd4a 100644 --- a/tests/unit/backend/test_runtime_lifecycle.py +++ b/tests/unit/backend/test_runtime_lifecycle.py @@ -23,7 +23,7 @@ from pyrit.backend.authentication_policy import AuthenticationPolicy from pyrit.backend.middleware.auth import require_admin from pyrit.backend.middleware.runtime import RuntimeAdmissionMiddleware -from pyrit.backend.routes import configuration, health +from pyrit.backend.routes import configuration, health, targets from pyrit.backend.services import ScenarioRunService, get_converter_service, get_target_service from pyrit.backend.services.configuration_file_service import ConfigurationFileService from pyrit.backend.services.manual_send_scheduler import get_manual_send_scheduler @@ -33,10 +33,12 @@ from pyrit.backend.services.technique_service import get_technique_service from pyrit.memory import CentralMemory, MemoryInterface from pyrit.models import ScenarioRunState -from pyrit.models.catalog import RunScenarioRequest -from pyrit.prompt_target import OpenAIChatTarget +from pyrit.models.catalog import ComponentRecipe, RunScenarioRequest +from pyrit.prompt_target import OpenAIChatTarget, TextTarget from pyrit.registry import AttackTechniqueRegistry, ConverterRegistry, ScorerRegistry, TargetRegistry +from pyrit.registry.instance_registry import CONFIGURATION_ENTRY_METADATA_KEY from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer, SavedRecipes +from pyrit.registry.registry_recipe_storage import RegistryRecipeStorage from pyrit.scenario import Scenario from pyrit.setup import ConfigurationLoader from pyrit.setup.initialization import reset_setup_registries @@ -702,11 +704,17 @@ async def test_management_publishes_target_upload_directory(tmp_path: Path) -> N app = FastAPI() with patch.dict(os.environ, {}, clear=True): service = RuntimeLifecycle(app=app, source=ConfigurationFileService(config_file_value=None)) - config = ConfigurationLoader(memory_db_type="in_memory", env_files=[], target_upload_directory=str(tmp_path)) + config = ConfigurationLoader( + memory_db_type="in_memory", + env_files=[], + target_upload_directory=str(tmp_path), + registry_recipes_source=str(tmp_path / "recipes"), + ) await service._management_async(config) assert app.state.target_upload_directory == str(tmp_path) + assert app.state.registry_recipes_source == str(tmp_path / "recipes") async def test_invalid_cold_configuration_requires_restart_but_retains_raw_repair(tmp_path: Path) -> None: @@ -719,6 +727,7 @@ async def test_invalid_cold_configuration_requires_restart_but_retains_raw_repai assert service.state == "restart-required" assert not app.state.allow_custom_initializers assert app.state.target_upload_directory is None + assert app.state.registry_recipes_source is None assert app.state.environment_file_service is None content, version = await source.read_with_version_async() assert content == "broken: [" @@ -1062,6 +1071,26 @@ async def test_startup_restores_an_entry_that_uses_a_startup_target_async( assert converter._converter_target is startup_chat +@pytest.mark.usefixtures("patch_central_database") +async def test_startup_marks_the_targets_its_scripts_register_but_not_saved_ones_async( + runtime: RuntimeLifecycle, tmp_path: Path, startup_chat: OpenAIChatTarget +) -> None: + recipes = tmp_path / "recipes" + text_target = {"class_module": TextTarget.__module__, "class_name": "TextTarget", "params": {}} + _save_recipe(recipes, kind="targets", name="saved_text", document=text_target) + targets = TargetRegistry.get_registry_singleton() + + await _start_with_recipes_async( + runtime, recipes=recipes, initialize=lambda **_: targets.instances.register(startup_chat, name="startup_chat") + ) + + assert (runtime.state, runtime.outcome) == ("ready", "success") + startup_entry, saved_entry = targets.instances.get_entry("startup_chat"), targets.instances.get_entry("saved_text") + assert startup_entry is not None and saved_entry is not None + assert CONFIGURATION_ENTRY_METADATA_KEY in startup_entry.metadata + assert CONFIGURATION_ENTRY_METADATA_KEY not in saved_entry.metadata + + class _LoopBoundTarget(MockPromptTarget): """A target whose constructor needs the running event loop, as HuggingFaceChatTarget's does.""" @@ -1154,6 +1183,9 @@ def reinitialize(**_: object) -> None: assert converter is not None assert converter._converter_target is startup_chat assert ScorerRegistry.get_registry_singleton().instances.get("refusal") is not None + startup_entry = TargetRegistry.get_registry_singleton().instances.get_entry("startup_chat") + assert startup_entry is not None + assert CONFIGURATION_ENTRY_METADATA_KEY in startup_entry.metadata class _UnreachableTarget(MockPromptTarget): @@ -1376,3 +1408,36 @@ async def test_restore_cancellation_propagates_from_startup_async(runtime: Runti await _start_with_recipes_async(runtime, recipes=tmp_path / "recipes") assert runtime.state == "initializing" + + +@pytest.mark.usefixtures("patch_central_database") +async def test_a_target_save_in_flight_makes_apply_busy_async(runtime: RuntimeLifecycle, tmp_path: Path) -> None: + TargetRegistry.reset_registry_singleton() + writing, finish_writing = Event(), Event() + create_recipe = RegistryRecipeStorage.create_recipe + + def create_recipe_slowly(storage: RegistryRecipeStorage, *, name: str, recipe: ComponentRecipe) -> str: + writing.set() + finish_writing.wait(TEST_WATCHDOG_SECONDS) + return create_recipe(storage, name=name, recipe=recipe) + + runtime.app.include_router(targets.router, prefix="/api") + runtime.app.state.registry_recipes_source = str(tmp_path / "recipes") + body = {"name": "saved_text", "type": "TextTarget", "params": {}, "persist": True} + try: + with patch.object(RegistryRecipeStorage, "create_recipe", create_recipe_slowly): + transport = httpx.ASGITransport(app=runtime.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + save = asyncio.create_task(client.post("/api/targets", json=body)) + assert await asyncio.to_thread(writing.wait, TEST_WATCHDOG_SECONDS) + await apply_async(runtime) + assert (runtime.outcome, runtime.state, runtime.generation) == ("busy", "ready", "original") + finish_writing.set() + response = await save + + assert response.status_code == 201 + assert (tmp_path / "recipes" / "targets" / "saved_text.json").exists() + lifecycle_module.close_services_async.assert_not_awaited() + finally: + finish_writing.set() + TargetRegistry.reset_registry_singleton() diff --git a/tests/unit/backend/test_saved_entry_restore.py b/tests/unit/backend/test_saved_entry_restore.py index dde2721798..bca2c955ca 100644 --- a/tests/unit/backend/test_saved_entry_restore.py +++ b/tests/unit/backend/test_saved_entry_restore.py @@ -7,7 +7,7 @@ from collections.abc import Iterator from functools import partial from pathlib import Path -from types import ModuleType +from types import MappingProxyType, ModuleType from typing import Any from unittest.mock import patch @@ -15,7 +15,7 @@ from pyrit.backend.services.saved_entry_restore import build_saved_entry_async from pyrit.models.catalog import RegistryRecipeKind, RegistryRestoreReport, RestoreFailureReason -from pyrit.prompt_target import HTTPXAPITarget, OpenAIChatTarget +from pyrit.prompt_target import GitHubCopilotTarget, OpenAIChatTarget from pyrit.registry import AttackTechniqueRegistry, ConverterRegistry, ScorerRegistry, TargetRegistry from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer @@ -23,7 +23,7 @@ "class_module": "pyrit.prompt_target.openai.openai_chat_target", "class_name": "OpenAIChatTarget", "auth_mode": "identity", - "params": {"endpoint": "https://contoso.openai.azure.com/openai/v1", "model_name": "gpt-4o"}, + "params": {"endpoint": "https://contoso.openai.azure.com/openai/v1", "headers": "{}", "model_name": "gpt-4o"}, } _REGISTRIES = (TargetRegistry, ConverterRegistry, ScorerRegistry, AttackTechniqueRegistry) @@ -114,7 +114,12 @@ async def test_a_target_its_create_refuses_fails_only_that_entry(tmp_path: Path) ) _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) - with patch.dict(sys.modules, {"copilot": ModuleType("copilot")}), patch.dict(os.environ, {"GITHUB_TOKEN": ""}): + with ( + patch.dict(sys.modules, {"copilot": ModuleType("copilot")}), + patch.dict(os.environ, {"GITHUB_TOKEN": ""}), + # A token-less Copilot target would use the server's token, which saved entries refuse before building. + patch.object(GitHubCopilotTarget, "server_fallback_parameters", MappingProxyType({})), + ): report = await _restore(tmp_path) (failure,) = report.failures @@ -123,49 +128,6 @@ async def test_a_target_its_create_refuses_fails_only_that_entry(tmp_path: Path) assert [ref.name for ref in report.restored] == ["contoso_gpt4o"] -@pytest.mark.usefixtures("patch_central_database") -async def test_a_target_that_uploads_files_gets_this_servers_upload_folder(tmp_path: Path) -> None: - uploads = tmp_path / "uploads" - uploads.mkdir() - recipes = tmp_path / "recipes" - _save( - recipes, - RegistryRecipeKind.TARGETS, - "uploader", - { - "class_module": "pyrit.prompt_target.http_target.httpx_api_target", - "class_name": "HTTPXAPITarget", - "params": {"http_url": "https://example.test/upload"}, - }, - ) - - report = await _restore(recipes, upload_directory=str(uploads)) - - assert report.failures == [] - target = TargetRegistry.get_registry_singleton().instances.get("uploader") - assert isinstance(target, HTTPXAPITarget) - assert target.allowed_upload_directory == uploads.resolve() - - -@pytest.mark.usefixtures("patch_central_database") -async def test_a_target_that_uploads_files_fails_where_the_server_sets_no_upload_folder(tmp_path: Path) -> None: - _save( - tmp_path, - RegistryRecipeKind.TARGETS, - "uploader", - { - "class_module": "pyrit.prompt_target.http_target.httpx_api_target", - "class_name": "HTTPXAPITarget", - "params": {"http_url": "https://example.test/upload"}, - }, - ) - - (failure,) = (await _restore(tmp_path)).failures - - assert (failure.name, failure.reason) == ("uploader", RestoreFailureReason.BUILD_ERROR) - assert "target_upload_directory" in failure.message - - async def test_a_technique_its_create_request_refuses_is_invalid(tmp_path: Path) -> None: _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "types", {"type": "RedTeamingAttack"}) diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index bd95653556..d8097cae99 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -9,6 +9,7 @@ import logging import threading import uuid +from collections.abc import Iterator from dataclasses import replace from datetime import UTC, datetime, timedelta from types import SimpleNamespace @@ -22,9 +23,10 @@ import pyrit.backend.services.scenario_run_service as _svc_mod from pyrit.backend.services import ScenarioRunService from pyrit.backend.services.scenario_progress_read_model import ScenarioPlanLookup, ScenarioProgressReadModel -from pyrit.backend.services.scenario_run_service import ScenarioRunConflictError +from pyrit.backend.services.scenario_run_service import SavedEntryLaunchConflictError, ScenarioRunConflictError from pyrit.common.utils import to_sha256 -from pyrit.converter import Converter +from pyrit.converter import Base64Converter, Converter, ROT13Converter +from pyrit.executor.attack import PromptSendingAttack from pyrit.memory import ( AttackResultKeysetCursor, AzureSQLMemory, @@ -62,10 +64,17 @@ SeedObjective, config_hash, ) -from pyrit.models.catalog import RunScenarioRequest -from pyrit.prompt_target import TargetCapabilities -from pyrit.registry import ConverterRegistry, ScenarioMetadata, ScenarioRegistry -from pyrit.scenario import IncompatibleTechniquePolicy, Scenario, ScenarioTechnique +from pyrit.models.catalog import RegistryEntryOrigin, RunScenarioRequest +from pyrit.prompt_target import TargetCapabilities, TextTarget +from pyrit.registry import ( + AttackTechniqueRegistry, + ConverterRegistry, + ScenarioMetadata, + ScenarioRegistry, + TargetRegistry, +) +from pyrit.registry.instance_registry import SAVED_ENTRY_METADATA_KEY, registry_entry_origin +from pyrit.scenario import AttackTechniqueFactory, IncompatibleTechniquePolicy, Scenario, ScenarioTechnique from pyrit.scenario.core import ( DatasetAttackConfiguration, DatasetConfiguration, @@ -73,6 +82,7 @@ ) from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides from pyrit.scenario.scenarios.airt import Jailbreak, Scam +from pyrit.setup.initializers import ConverterInitializer, ToolsInitializer from unit.async_utils import wait_for_completion_async from unit.mocks import ( MockPromptTarget, @@ -263,6 +273,155 @@ async def prepare_async() -> _svc_mod._PreparedRun: await service.close_async() +@pytest.fixture +def converter_registry() -> Iterator[ConverterRegistry]: + ConverterRegistry.reset_registry_singleton() + yield ConverterRegistry.get_registry_singleton() + ConverterRegistry.reset_registry_singleton() + + +class _LaunchScript: + """A launch script that registers one converter, optionally swallowing or wrapping errors.""" + + def __init__(self, *, name: str, swallow: bool = False, wrap: bool = False) -> None: + self._name = name + self._swallow = swallow + self._wrap = wrap + + async def initialize_async(self) -> None: + registry = ConverterRegistry.get_registry_singleton() + try: + registry.instances.register(Base64Converter(), name=self._name, replace=True) + except Exception as error: + if self._wrap: + raise ValueError("The launch script failed.") from error + if not self._swallow: + raise + + +async def _run_launch_scripts_async(scripts: dict[str, Any]) -> None: + initializers = MagicMock() + initializers.create_and_configure.side_effect = lambda name, initializer_params=None: scripts[name] + with patch(f"{_REGISTRY_PATCH_BASE}.InitializerRegistry.get_registry_singleton", return_value=initializers): + await ScenarioRunService()._run_initializers_async(request=_make_request(initializers=list(scripts))) + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + "script", + [ + ConverterInitializer(), + _LaunchScript(name="rot13"), + _LaunchScript(name="rot13", swallow=True), + _LaunchScript(name="rot13", wrap=True), + ], +) +async def test_launch_scripts_cannot_replace_a_saved_entry_async( + converter_registry: ConverterRegistry, script: Any +) -> None: + saved = ROT13Converter() + converter_registry.instances.register(saved, name="rot13", metadata={SAVED_ENTRY_METADATA_KEY: "converters"}) + + with pytest.raises(SavedEntryLaunchConflictError, match="'converters' would replace the saved converter 'rot13'"): + await _run_launch_scripts_async({"converters": script}) + + assert converter_registry.instances.get("rot13") is saved + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tools_launch_script_cannot_use_a_saved_technique_in_its_place_async() -> None: + AttackTechniqueRegistry.reset_registry_singleton() + registry = AttackTechniqueRegistry.get_registry_singleton() + saved = AttackTechniqueFactory(name="role_play_tool_calls", attack_class=PromptSendingAttack) + registry.instances.register_runtime(saved, metadata={SAVED_ENTRY_METADATA_KEY: "techniques"}) + try: + with pytest.raises( + SavedEntryLaunchConflictError, match="'tools' would replace the saved technique 'role_play_tool_calls'" + ): + await _run_launch_scripts_async({"tools": ToolsInitializer()}) + + assert registry.instances.get("role_play_tool_calls") is saved + finally: + AttackTechniqueRegistry.reset_registry_singleton() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tools_launch_script_cannot_replace_a_saved_converter_async( + converter_registry: ConverterRegistry, +) -> None: + AttackTechniqueRegistry.reset_registry_singleton() + saved = ROT13Converter() + converter_registry.instances.register( + saved, name="role_play_tool_calls", metadata={SAVED_ENTRY_METADATA_KEY: "converters"} + ) + try: + with pytest.raises( + SavedEntryLaunchConflictError, match="'tools' would replace the saved converter 'role_play_tool_calls'" + ): + await _run_launch_scripts_async({"tools": ToolsInitializer()}) + + assert converter_registry.instances.get("role_play_tool_calls") is saved + assert "role_play_tool_calls" not in AttackTechniqueRegistry.get_registry_singleton().instances + finally: + AttackTechniqueRegistry.reset_registry_singleton() + + +class _FailingLaunchScript: + """A launch script that fails for a reason unrelated to saved entries.""" + + async def initialize_async(self) -> None: + raise ValueError("The launch script failed.") + + +@pytest.mark.usefixtures("patch_central_database") +async def test_launch_script_errors_without_a_saved_entry_clash_pass_through_async( + converter_registry: ConverterRegistry, +) -> None: + with pytest.raises(ValueError, match="The launch script failed") as raised: + await _run_launch_scripts_async({"failing": _FailingLaunchScript()}) + + assert not isinstance(raised.value, SavedEntryLaunchConflictError) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_launch_scripts_still_replace_startup_entries_async(converter_registry: ConverterRegistry) -> None: + startup = ROT13Converter() + converter_registry.instances.register(startup, name="rot13") + + await _run_launch_scripts_async({"converters": _LaunchScript(name="rot13")}) + + assert isinstance(converter_registry.instances.get("rot13"), Base64Converter) + + +class _LaunchScriptTextTarget(TextTarget): + """A target class that only a launch script registers.""" + + +class _TargetLaunchScript: + """A launch script that adds a target class and a target.""" + + async def initialize_async(self) -> None: + registry = TargetRegistry.get_registry_singleton() + registry.register_class(_LaunchScriptTextTarget, name="launch_script_text") + registry.instances.register(TextTarget(), name="launched_text") + + +@pytest.mark.usefixtures("patch_central_database") +async def test_what_launch_scripts_add_is_temporary_async() -> None: + TargetRegistry.reset_registry_singleton() + try: + await _run_launch_scripts_async({"targets": _TargetLaunchScript()}) + + registry = TargetRegistry.get_registry_singleton() + entry = registry.instances.get_entry("launched_text") + assert entry is not None + assert registry_entry_origin(entry.metadata) is RegistryEntryOrigin.TEMPORARY + assert registry.is_temporary_class("launch_script_text") + assert not registry.is_temporary_class("TextTarget") + finally: + TargetRegistry.reset_registry_singleton() + + def _make_request( *, scenario_name: str = "foundry.red_team_agent", diff --git a/tests/unit/backend/test_scenario_suite_routes.py b/tests/unit/backend/test_scenario_suite_routes.py index 177ad3677a..ce659e0ed0 100644 --- a/tests/unit/backend/test_scenario_suite_routes.py +++ b/tests/unit/backend/test_scenario_suite_routes.py @@ -22,14 +22,18 @@ ScenarioSuiteResponse, ) from pyrit.backend.services import ( + ScenarioRunService, ScenarioSuiteExecutionService, ScenarioSuiteService, ) from pyrit.backend.services.scenario_suite_execution_service import ( ScenarioSuiteExecutionCapacityError, ) +from pyrit.converter import Base64Converter from pyrit.models import ( + PresetEntry, PresetIssue, + ScenarioPreset, ScenarioSuite, ScenarioSuiteEntry, ScenarioSuiteExecutionEntrySummary, @@ -39,6 +43,8 @@ ) from pyrit.models.catalog import ScenarioSuiteRegistryEntry from pyrit.registry import ( + ConverterRegistry, + InitializerRegistry, ScenarioPresetRegistry, ScenarioSuiteConflictError, ScenarioSuiteExecutionNotFoundError, @@ -49,6 +55,8 @@ ScenarioSuiteStorage, ) from pyrit.registry.file_document_storage import DocumentStorageUnavailableError +from pyrit.registry.instance_registry import SAVED_ENTRY_METADATA_KEY +from pyrit.setup.initializers import ConverterInitializer SUITE_NAME = "weekly" ENTRY_ID = UUID("11111111-1111-4111-8111-111111111111") @@ -207,6 +215,14 @@ def signed_in_user() -> Iterator[None]: app.dependency_overrides.pop(current_user_name, None) +@pytest.fixture +def converter_registry() -> Iterator[ConverterRegistry]: + """Give the test a fresh converter registry and drop what launch scripts registered.""" + ConverterRegistry.reset_registry_singleton() + yield ConverterRegistry.get_registry_singleton() + ConverterRegistry.reset_registry_singleton() + + def test_list_is_open_and_returns_entry_issues(anonymous_client: TestClient, service: MagicMock) -> None: """Test open reads with entry-scoped advisory issues.""" issue = PresetIssue(field="preset_name", message="Preset is missing.") @@ -285,6 +301,49 @@ def test_launch_maps_preflight_failures( assert response.json()["detail"] == "wait for an active suite" +@pytest.mark.usefixtures("patch_central_database") +def test_launch_refuses_a_launch_script_that_would_replace_a_saved_entry( + *, anonymous_client: TestClient, converter_registry: ConverterRegistry, tmp_path: Path +) -> None: + """Test a launch script that would replace a saved entry is a conflict before any suite execution or scan.""" + saved = Base64Converter() + converter_registry.instances.register(saved, name="rot13", metadata={SAVED_ENTRY_METADATA_KEY: "converters"}) + suite_registry = ScenarioSuiteRegistry(storage=ScenarioSuiteStorage(source=str(tmp_path))) + stored = suite_registry.save_suite(suite=_suite(), expected_version=None) + preset_registry = MagicMock(spec=ScenarioPresetRegistry) + preset_registry.get_preset.return_value = PresetEntry( + preset=ScenarioPreset(name="nightly", scenario_name="scenario.nightly"), origin="stored", version="preset-v1" + ) + preset_registry.check_preset.return_value = [] + run_service = ScenarioRunService() + execution_service = ScenarioSuiteExecutionService( + suite_registry=suite_registry, preset_registry=preset_registry, run_service=run_service + ) + initializers = MagicMock(spec=InitializerRegistry) + initializers.create_and_configure.return_value = ConverterInitializer() + + with ( + patch( + "pyrit.backend.routes.scenario_suites.get_scenario_suite_execution_service", + return_value=execution_service, + ), + patch.object(InitializerRegistry, "get_registry_singleton", return_value=initializers), + patch.object(run_service, "start_run_async", new_callable=AsyncMock) as start_run, + ): + response = anonymous_client.post( + f"/api/scenario-suites/{SUITE_NAME}/launch", + json={**_launch_body(), "expected_version": stored.version, "initializers": ["converters"]}, + ) + + assert response.status_code == 409 + assert response.json()["detail"] == ( + "The launch script 'converters' would replace the saved converter 'rot13'. Leave that script out of the launch." + ) + assert execution_service._executions == {} + start_run.assert_not_awaited() + assert converter_registry.instances.get("rot13") is saved + + def test_launch_limits_the_expected_version(anonymous_client: TestClient, execution_service: MagicMock) -> None: """Test the reviewed suite version has the identifier limit.""" execution_service.launch_async = AsyncMock(return_value=_execution_summary()) diff --git a/tests/unit/backend/test_target_service.py b/tests/unit/backend/test_target_service.py index 4283a91226..99bb14f8ad 100644 --- a/tests/unit/backend/test_target_service.py +++ b/tests/unit/backend/test_target_service.py @@ -5,19 +5,37 @@ Tests for backend target service. """ +import asyncio +import json import os +from functools import partial from pathlib import Path from types import ModuleType -from unittest.mock import MagicMock, patch +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from azure.core.exceptions import AzureError from pyrit.backend.models import CreateTargetRequest from pyrit.backend.services import TargetService, get_target_service +from pyrit.backend.services.saved_entry_restore import build_saved_entry_async +from pyrit.backend.services.target_service import SavedTargetConflictError, SavedTargetStorageError from pyrit.models import ComponentIdentifier +from pyrit.models.catalog import ( + ComponentRecipe, + RegistryEntryOrigin, + RegistryRecipeKind, + RegistryRestoreReport, + SavedEntryOwner, + TargetInstance, +) from pyrit.models.component_spec import SourceInstanceSpec -from pyrit.prompt_target import GitHubCopilotTarget, PromptTarget, TargetCapabilities +from pyrit.prompt_target import GitHubCopilotTarget, PromptTarget, RoundRobinTarget, TargetCapabilities, TextTarget from pyrit.registry import TargetRegistry +from pyrit.registry.instance_registry import RegistryNameTakenError, registration_origin +from pyrit.registry.registry_recipe_restore import RegistryRecipeRestorer +from pyrit.registry.registry_recipe_storage import RegistryRecipeStorage from unit.mocks import MockPromptTarget @@ -1733,3 +1751,306 @@ async def test_create_target_with_no_key_rejects_unsupported_type(self, sqlite_i with pytest.raises(ValueError, match="does not support no-key authentication"): await service.create_target_async(request=request) + + +_AZURE_ENDPOINT = "https://contoso.openai.azure.com/openai/v1" + + +def _save_request( + *, + name: str = "contoso_gpt4o", + params: dict[str, Any] | None = None, + auth_mode: str = "identity", + persist: bool = True, +) -> CreateTargetRequest: + return CreateTargetRequest( + name=name, + type="OpenAIChatTarget", + params=params if params is not None else {"endpoint": _AZURE_ENDPOINT, "model_name": "gpt-4o"}, + auth_mode=auth_mode, + persist=persist, + ) + + +def _saved_document(source: Path, name: str) -> dict[str, Any]: + return json.loads((source / "targets" / f"{name}.json").read_text()) + + +async def _restore(source: Path) -> RegistryRestoreReport: + """Restore the saved targets through the API's create, as startup does.""" + get_target_service.cache_clear() + build_async = partial(build_saved_entry_async, upload_directory=None) + return await RegistryRecipeRestorer(source=str(source)).restore_async(build_async=build_async) + + +@pytest.mark.usefixtures("sqlite_instance") +class TestSaveTarget: + async def test_saves_an_identity_target_that_restore_rebuilds(self, tmp_path: Path) -> None: + service = TargetService() + owner = SavedEntryOwner(oid="00000000-0000-0000-0000-000000000001", name="Ada") + + target = await service.create_target_async(request=_save_request(), owner=owner, recipes_source=str(tmp_path)) + + assert target.origin is RegistryEntryOrigin.SAVED + assert _saved_document(tmp_path, "contoso_gpt4o") == { + "schema_version": 1, + "class_module": "pyrit.prompt_target.openai.openai_chat_target", + "class_name": "OpenAIChatTarget", + "auth_mode": "identity", + "params": {"endpoint": _AZURE_ENDPOINT, "model_name": "gpt-4o", "headers": "{}"}, + "owner": {"oid": owner.oid, "name": "Ada"}, + } + saved_hash = service.get_target_object(target_registry_name="contoso_gpt4o").get_identifier().hash + TargetRegistry.reset_registry_singleton() + report = await _restore(tmp_path) + assert report.failures == [] + restored = TargetRegistry.get_registry_singleton().instances.get("contoso_gpt4o") + assert restored is not None + assert restored.get_identifier().hash == saved_hash + + async def test_saves_a_no_key_target_without_an_owner_when_sign_in_is_off(self, tmp_path: Path) -> None: + request = _save_request( + name="local_llm", params={"endpoint": "http://localhost:11434/v1", "model_name": "llama3"}, auth_mode="none" + ) + + target = await TargetService().create_target_async(request=request, recipes_source=str(tmp_path)) + + assert target.origin is RegistryEntryOrigin.SAVED + document = _saved_document(tmp_path, "local_llm") + assert document["auth_mode"] == "none" + assert "auth_mode" not in document["params"] + assert "owner" not in document + + @pytest.mark.parametrize( + ("auth_mode", "params", "message"), + [ + ("api_key", {"api_key": "sk-typed-secret"}, "can't hold credentials"), + ("api_key", {}, "'api_key' is empty"), + ("identity", {"headers": '{"api-key": "sk-header-secret"}'}, "'headers' can carry a credential"), + ("identity", {"endpoint": f"{_AZURE_ENDPOINT}?code=sk-url-secret"}, "can carry a credential in a URL"), + ], + ) + async def test_refuses_to_save_a_target_that_would_hold_or_use_a_key( + self, tmp_path: Path, auth_mode: str, params: dict[str, Any], message: str + ) -> None: + service = TargetService() + request = _save_request( + params={"endpoint": _AZURE_ENDPOINT, "model_name": "gpt-4o", **params}, auth_mode=auth_mode + ) + + with ( + patch.dict(os.environ, {"OPENAI_CHAT_KEY": "sk-server-key"}), + pytest.raises(ValueError, match=message) as refusal, + ): + await service.create_target_async(request=request, recipes_source=str(tmp_path)) + + assert str(refusal.value).endswith("It can still be created without saving it.") + assert "secret" not in str(refusal.value) + assert "sk-server-key" not in str(refusal.value) + assert service.get_target_object(target_registry_name="contoso_gpt4o") is None + assert not (tmp_path / "targets").exists() + + async def test_a_taken_name_is_a_conflict_only_when_saving(self, tmp_path: Path) -> None: + service = TargetService() + await service.create_target_async(request=_save_request(persist=False)) + + with pytest.raises(SavedTargetConflictError, match="already exists"): + await service.create_target_async(request=_save_request(), recipes_source=str(tmp_path)) + with pytest.raises(ValueError, match="already exists") as temporary_clash: + await service.create_target_async(request=_save_request(persist=False)) + + assert not isinstance(temporary_clash.value, SavedTargetConflictError) + assert not (tmp_path / "targets").exists() + + async def test_a_saved_file_with_the_name_is_a_conflict_that_leaves_the_file(self, tmp_path: Path) -> None: + service = TargetService() + storage = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)) + storage.create_recipe(name="contoso_gpt4o", recipe=ComponentRecipe(class_module="m", class_name="TextTarget")) + before = (tmp_path / "targets" / "contoso_gpt4o.json").read_bytes() + + with ( + patch("pyrit.backend.services.target_service.release_component_async", new_callable=AsyncMock) as release, + pytest.raises(SavedTargetConflictError, match="A saved target named 'contoso_gpt4o' already exists."), + ): + await service.create_target_async(request=_save_request(), recipes_source=str(tmp_path)) + + release.assert_awaited_once() + assert service.get_target_object(target_registry_name="contoso_gpt4o") is None + assert (tmp_path / "targets" / "contoso_gpt4o.json").read_bytes() == before + + @pytest.mark.parametrize( + ("error", "raised"), + [ + (RegistryNameTakenError("Instance 'contoso_gpt4o' already exists"), SavedTargetConflictError), + (ValueError("Instance name 'contoso_gpt4o' is reserved"), ValueError), + (RuntimeError("x"), RuntimeError), + ], + ) + async def test_deletes_the_saved_file_when_the_target_cant_be_registered( + self, tmp_path: Path, error: Exception, raised: type[Exception] + ) -> None: + service = TargetService() + + with ( + patch.object(service._registry.instances, "register", side_effect=error), + pytest.raises(raised) as refusal, + ): + await service.create_target_async(request=_save_request(), recipes_source=str(tmp_path)) + + assert type(refusal.value) is raised + assert not (tmp_path / "targets" / "contoso_gpt4o.json").exists() + + @pytest.mark.parametrize( + ("error", "reason"), + [ + (PermissionError(13, "Permission denied", "/srv/recipes/targets/x.json"), "Permission denied"), + (OSError("disk full"), "OSError"), + (AzureError("unreachable"), "AzureError"), + ], + ) + async def test_reports_storage_failures_as_unavailable(self, tmp_path: Path, error: Exception, reason: str) -> None: + service = TargetService() + + with ( + patch.object(RegistryRecipeStorage, "_save_local_conditional", side_effect=error), + patch("pyrit.backend.services.target_service.logger") as log, + pytest.raises(SavedTargetStorageError, match="Saved targets can't be written right now."), + ): + await service.create_target_async(request=_save_request(), recipes_source=str(tmp_path)) + + assert service.get_target_object(target_registry_name="contoso_gpt4o") is None + message = log.error.call_args.args[0] % log.error.call_args.args[1:] + assert f"'contoso_gpt4o' couldn't be written to {tmp_path / 'targets'} ({reason})" in message + assert not any(detail in message for detail in ("/srv/recipes", "disk full", "unreachable")) + assert log.error.call_args.kwargs.get("exc_info") is None + + async def test_two_saves_of_one_name_save_one_target(self, tmp_path: Path) -> None: + service = TargetService() + mini = {"endpoint": _AZURE_ENDPOINT, "model_name": "gpt-4o-mini"} + + results = await asyncio.gather( + service.create_target_async(request=_save_request(), recipes_source=str(tmp_path)), + service.create_target_async(request=_save_request(params=mini), recipes_source=str(tmp_path)), + return_exceptions=True, + ) + + assert sorted(type(result).__name__ for result in results) == ["SavedTargetConflictError", "TargetInstance"] + registered = service.get_target_object(target_registry_name="contoso_gpt4o") + assert _saved_document(tmp_path, "contoso_gpt4o")["params"]["model_name"] == registered._model_name + + async def test_saves_a_round_robin_only_over_targets_that_come_back_after_a_restart(self, tmp_path: Path) -> None: + service = TargetService() + with registration_origin(RegistryEntryOrigin.CONFIGURATION): + for name in ("startup_one", "startup_two"): + await service.create_target_async(request=_save_request(name=name, persist=False)) + await service.create_target_async(request=_save_request(name="launch_member", persist=False)) + await service.create_target_async(request=_save_request(name="saved_member"), recipes_source=str(tmp_path)) + + def pool(name: str, *members: str) -> CreateTargetRequest: + return CreateTargetRequest( + name=name, type="RoundRobinTarget", params={"targets": list(members)}, persist=True + ) + + with pytest.raises(ValueError, match="Uses temporary target 'launch_member'"): + await service.create_target_async( + request=pool("pool", "startup_one", "launch_member"), recipes_source=str(tmp_path) + ) + startup_pool = await service.create_target_async( + request=pool("startup_pool", "startup_one", "startup_two"), recipes_source=str(tmp_path) + ) + mixed_pool = await service.create_target_async( + request=pool("mixed_pool", "startup_one", "saved_member"), recipes_source=str(tmp_path) + ) + + assert (startup_pool.origin, mixed_pool.origin) == (RegistryEntryOrigin.SAVED, RegistryEntryOrigin.SAVED) + assert _saved_document(tmp_path, "mixed_pool")["params"] == {"targets": ["startup_one", "saved_member"]} + assert not (tmp_path / "targets" / "pool.json").exists() + + async def test_saves_a_class_registered_under_another_name_so_restore_finds_it(self, tmp_path: Path) -> None: + class AliasedTextTarget(TextTarget): + pass + + TargetRegistry.get_registry_singleton().register_class(AliasedTextTarget, name="contoso_text") + request = CreateTargetRequest(name="aliased", type="contoso_text", params={}, persist=True) + + await TargetService().create_target_async(request=request, recipes_source=str(tmp_path)) + + assert _saved_document(tmp_path, "aliased")["class_name"] == "contoso_text" + TargetRegistry.reset_registry_singleton() + TargetRegistry.get_registry_singleton().register_class(AliasedTextTarget, name="contoso_text") + assert (await _restore(tmp_path)).failures == [] + + async def test_refuses_to_save_a_target_whose_class_a_launch_script_added(self, tmp_path: Path) -> None: + class LaunchScriptTextTarget(TextTarget): + pass + + with registration_origin(RegistryEntryOrigin.TEMPORARY): + TargetRegistry.get_registry_singleton().register_class(LaunchScriptTextTarget, name="launch_text") + service = TargetService() + request = CreateTargetRequest(name="launched", type="launch_text", params={}, persist=True) + + with pytest.raises(ValueError, match="'launch_text' was added by a scan's launch script"): + await service.create_target_async(request=request, recipes_source=str(tmp_path)) + await service.create_target_async(request=request.model_copy(update={"persist": False})) + + assert not (tmp_path / "targets").exists() + assert service.get_target_object(target_registry_name="launched") is not None + + async def test_checks_a_class_without_cached_metadata_the_way_restore_does(self) -> None: + service = TargetService() + await service.create_target_async(request=_save_request(name="temporary_member", persist=False)) + + with ( + patch.object(TargetRegistry, "get_registered_class_metadata", return_value=None), + pytest.raises(ValueError, match="Uses temporary target 'temporary_member'"), + ): + await service._prepare_saved_params_async( + target_type="RoundRobinTarget", + target_cls=RoundRobinTarget, + params={"targets": ["temporary_member"]}, + auth_mode="api_key", + ) + + async def test_reports_storage_trouble_when_a_failed_save_cant_be_undone(self, tmp_path: Path) -> None: + service = TargetService() + + with ( + patch.object(service._registry.instances, "register", side_effect=ValueError("taken")), + patch.object(RegistryRecipeStorage, "delete_recipe", side_effect=OSError("disk gone")), + patch("pyrit.backend.services.target_service.logger") as log, + pytest.raises(SavedTargetStorageError), + ): + await service.create_target_async(request=_save_request(), recipes_source=str(tmp_path)) + + assert (tmp_path / "targets" / "contoso_gpt4o.json").exists() + message = log.error.call_args.args[0] % log.error.call_args.args[1:] + assert "'contoso_gpt4o'" in message + assert "couldn't be deleted (OSError)" in message + assert "disk gone" not in message + + async def test_a_reserved_name_is_refused_as_invalid_not_as_taken(self, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="reserved") as refusal: + await TargetService().create_target_async(request=_save_request(name="types"), recipes_source=str(tmp_path)) + + assert not isinstance(refusal.value, SavedTargetConflictError) + + async def test_lists_where_each_target_comes_from(self, tmp_path: Path) -> None: + service = TargetService() + with registration_origin(RegistryEntryOrigin.CONFIGURATION): + service._registry.instances.register(MockPromptTarget(), name="from_startup_script") + with registration_origin(RegistryEntryOrigin.TEMPORARY): + service._registry.instances.register(MockPromptTarget(), name="from_launch_script") + await service.create_target_async(request=_save_request(name="temporary_one", persist=False)) + await service.create_target_async(request=_save_request(name="saved_one"), recipes_source=str(tmp_path)) + + listed = await service.list_targets_async() + + assert {item.target_registry_name: item.origin for item in listed.items} == { + "from_startup_script": RegistryEntryOrigin.CONFIGURATION, + "from_launch_script": RegistryEntryOrigin.TEMPORARY, + "temporary_one": RegistryEntryOrigin.TEMPORARY, + "saved_one": RegistryEntryOrigin.SAVED, + } + fetched = await service.get_target_async(target_registry_name="temporary_one") + assert isinstance(fetched, TargetInstance) + assert fetched.origin is RegistryEntryOrigin.TEMPORARY diff --git a/tests/unit/registry/test_attack_technique_registry.py b/tests/unit/registry/test_attack_technique_registry.py index 14d8915a90..6f4809390c 100644 --- a/tests/unit/registry/test_attack_technique_registry.py +++ b/tests/unit/registry/test_attack_technique_registry.py @@ -14,6 +14,11 @@ from pyrit.models import ComponentIdentifier from pyrit.prompt_target import PromptTarget, TargetRequirements from pyrit.registry import AttackTechniqueRegistry, Registry, TargetRegistry +from pyrit.registry.instance_registry import ( + SAVED_ENTRY_METADATA_KEY, + SavedEntryReplacementError, + record_saved_entry_refusals, +) from pyrit.scenario import AttackTechniqueFactory, ScenarioTechnique from pyrit.scenario.core import ScorerOverridePolicy from pyrit.setup.initializers.techniques import build_technique_factories @@ -362,6 +367,35 @@ def test_policy_passed_to_factories_via_register_from_factories(self): stored = self.registry.instances.get_entry("stub_policy").instance assert stored._scorer_override_policy == ScorerOverridePolicy.WARN + def test_register_from_factories_refuses_to_skip_a_saved_technique(self): + saved = AttackTechniqueFactory(name="formal", attack_class=_StubAttack, technique_tags=["custom"]) + self.registry.instances.register_runtime(saved, metadata={SAVED_ENTRY_METADATA_KEY: "techniques"}) + + with pytest.raises(SavedEntryReplacementError, match="Saved technique 'formal'"): + self.registry.register_from_factories([AttackTechniqueFactory(name="formal", attack_class=_StubAttack)]) + + assert self.registry.instances.get("formal") is saved + + def test_register_runtime_records_a_caught_clash_with_a_saved_technique(self): + saved = AttackTechniqueFactory(name="formal", attack_class=_StubAttack, technique_tags=["custom"]) + self.registry.instances.register_runtime(saved, metadata={SAVED_ENTRY_METADATA_KEY: "techniques"}) + + with record_saved_entry_refusals() as refusals: + with pytest.raises(ValueError, match="names are case-insensitive"): + clash = AttackTechniqueFactory(name="FORMAL", attack_class=_StubAttack) + self.registry.instances.register_runtime(clash) + + assert refusals == [("techniques", "formal")] + assert self.registry.instances.get("formal") is saved + + def test_register_from_factories_still_skips_an_existing_startup_technique(self): + startup = AttackTechniqueFactory(name="startup", attack_class=_StubAttack) + self.registry.register_from_factories([startup]) + + self.registry.register_from_factories([AttackTechniqueFactory(name="startup", attack_class=_StubAttack)]) + + assert self.registry.instances.get("startup") is startup + SCENARIO_FACTORIES_FIXTURE: list[AttackTechniqueFactory] = [] diff --git a/tests/unit/registry/test_converter_registry.py b/tests/unit/registry/test_converter_registry.py index db7a5cccf9..0df5345925 100644 --- a/tests/unit/registry/test_converter_registry.py +++ b/tests/unit/registry/test_converter_registry.py @@ -36,6 +36,7 @@ ConverterRegistry, TargetRegistry, ) +from pyrit.registry.instance_registry import SAVED_ENTRY_METADATA_KEY, record_saved_entry_refusals from pyrit.registry.resolution import derive_parameters @@ -197,6 +198,17 @@ def test_create_named_instance_builds_and_stores_converter(self, registry: Conve assert isinstance(converter, Base64Converter) assert registry.instances.get("base64") is converter + def test_create_named_instance_records_a_caught_clash_with_a_saved_entry(self, registry: ConverterRegistry): + saved = Base64Converter() + registry.instances.register(saved, name="base64", metadata={SAVED_ENTRY_METADATA_KEY: "converters"}) + + with record_saved_entry_refusals() as refusals: + with pytest.raises(ValueError, match="already exists"): + registry.create_named_instance(name="base64", type_name="Base64Converter") + + assert refusals == [("converters", "base64")] + assert registry.instances.get("base64") is saved + def test_create_named_instance_stores_registry_metadata(self, registry: ConverterRegistry): converter = registry.create_named_instance( name="base64", diff --git a/tests/unit/registry/test_instance_registry.py b/tests/unit/registry/test_instance_registry.py index c6a2959f96..94a60cede2 100644 --- a/tests/unit/registry/test_instance_registry.py +++ b/tests/unit/registry/test_instance_registry.py @@ -7,6 +7,7 @@ ``SupportsInstances`` protocols. """ +import contextlib from concurrent.futures import ThreadPoolExecutor from threading import Barrier, Event from typing import assert_type @@ -15,6 +16,7 @@ import pytest from pyrit.models import ComponentIdentifier, Identifiable +from pyrit.models.catalog import RegistryEntryOrigin from pyrit.registry import ( DefaultInstanceRegistry, InstanceRegistry, @@ -22,6 +24,17 @@ SupportsInstances, TagQuery, ) +from pyrit.registry.instance_registry import ( + CONFIGURATION_ENTRY_METADATA_KEY, + SAVED_ENTRY_METADATA_KEY, + RegistryNameTakenError, + SavedEntryReplacementError, + current_registration_origin, + record_saved_entry_refusals, + registration_origin, + registry_entry_origin, + saved_entry_registration, +) class _TestItem(Identifiable): @@ -89,17 +102,124 @@ def test_register_multiple_instances(self, registry: DefaultInstanceRegistry[_Te def test_register_rejects_existing_name(self, registry: DefaultInstanceRegistry[_TestItem]): registry.register(_item("original"), name="name") - with pytest.raises(ValueError, match="already exists"): + with pytest.raises(RegistryNameTakenError, match="already exists"): registry.register(_item("updated"), name="name") assert registry.get("name") == "original" + def test_a_reserved_name_is_refused_but_not_as_taken(self): + registry: DefaultInstanceRegistry[_TestItem] = DefaultInstanceRegistry(reserved_names={"types"}) + + for check in ( + lambda: registry.register(_item("x"), name="types"), + lambda: registry.validate_name_available("types"), + ): + with pytest.raises(ValueError, match="reserved") as refusal: + check() + assert not isinstance(refusal.value, RegistryNameTakenError) + def test_register_can_explicitly_replace_existing(self, registry: DefaultInstanceRegistry[_TestItem]): registry.register(_item("original"), name="name") registry.register(_item("updated"), name="name", replace=True) assert registry.get("name") == "updated" + def test_register_refuses_to_replace_a_saved_entry(self, registry: DefaultInstanceRegistry[_TestItem]): + registry.register(_item("saved"), name="name", metadata={SAVED_ENTRY_METADATA_KEY: "targets"}) + revision = registry.revision + + with pytest.raises(SavedEntryReplacementError, match="Saved target 'name'"): + registry.register(_item("updated"), name="name", replace=True) + + assert registry.get("name") == "saved" + assert registry.revision == revision + + def test_refused_replacements_are_recorded_even_when_caught(self, registry: DefaultInstanceRegistry[_TestItem]): + registry.register(_item("saved"), name="name", metadata={SAVED_ENTRY_METADATA_KEY: "converters"}) + + with record_saved_entry_refusals() as refusals: + with contextlib.suppress(SavedEntryReplacementError): + registry.register(_item("updated"), name="name", replace=True) + registry.register(_item("other"), name="other", replace=True) + + assert refusals == [("converters", "name")] + assert registry.get("other") == "other" + + def test_marks_entries_registered_while_the_startup_scripts_run(self, registry: DefaultInstanceRegistry[_TestItem]): + startup_metadata = {"supports_tool_calls": True} + with registration_origin(RegistryEntryOrigin.CONFIGURATION): + registry.register(_item("startup"), name="startup", metadata=startup_metadata) + registry.register(_item("replaced"), name="replaced") + with registration_origin(RegistryEntryOrigin.TEMPORARY): + registry.register(_item("launched"), name="launched") + registry.register(_item("swapped in"), name="replaced", replace=True) + registry.register(_item("created"), name="created") + + origins = {entry.name: registry_entry_origin(entry.metadata) for entry in registry.get_all_instances()} + assert origins == { + "created": RegistryEntryOrigin.TEMPORARY, + "launched": RegistryEntryOrigin.TEMPORARY, + "replaced": RegistryEntryOrigin.TEMPORARY, + "startup": RegistryEntryOrigin.CONFIGURATION, + } + assert startup_metadata == {"supports_tool_calls": True} + + def test_the_registration_origin_ends_with_its_block_even_when_a_script_fails(self): + with pytest.raises(RuntimeError, match="script failed"), registration_origin(RegistryEntryOrigin.CONFIGURATION): + with registration_origin(RegistryEntryOrigin.TEMPORARY): + assert current_registration_origin() is RegistryEntryOrigin.TEMPORARY + assert current_registration_origin() is RegistryEntryOrigin.CONFIGURATION + raise RuntimeError("script failed") + + assert current_registration_origin() is None + + @pytest.mark.parametrize( + ("metadata", "origin"), + [ + ({SAVED_ENTRY_METADATA_KEY: "targets"}, RegistryEntryOrigin.SAVED), + ({SAVED_ENTRY_METADATA_KEY: "targets", CONFIGURATION_ENTRY_METADATA_KEY: True}, RegistryEntryOrigin.SAVED), + ({CONFIGURATION_ENTRY_METADATA_KEY: True}, RegistryEntryOrigin.CONFIGURATION), + ({}, RegistryEntryOrigin.TEMPORARY), + ({"supports_tool_calls": True}, RegistryEntryOrigin.TEMPORARY), + ], + ) + def test_tells_where_an_entry_comes_from(self, metadata: dict[str, object], origin: RegistryEntryOrigin): + assert registry_entry_origin(metadata) is origin + + def test_a_restore_scope_marks_only_the_entry_it_names(self, registry: DefaultInstanceRegistry[_TestItem]): + with saved_entry_registration(kind="targets", name="saved"): + registry.register(_item("saved"), name="saved", metadata={"owner": "admin"}) + registry.register(_item("helper"), name="helper") + registry.register(_item("later"), name="later") + + assert registry.get_entry("saved").metadata == {"owner": "admin", SAVED_ENTRY_METADATA_KEY: "targets"} + assert SAVED_ENTRY_METADATA_KEY not in registry.get_entry("helper").metadata + assert SAVED_ENTRY_METADATA_KEY not in registry.get_entry("later").metadata + + def test_a_caught_duplicate_name_for_a_saved_entry_is_recorded(self, registry: DefaultInstanceRegistry[_TestItem]): + registry.register(_item("saved"), name="name", metadata={SAVED_ENTRY_METADATA_KEY: "targets"}) + registry.register(_item("startup"), name="startup") + + with record_saved_entry_refusals() as refusals: + for name in ("name", "startup"): + with pytest.raises(ValueError, match="already exists"): + registry.register(_item("updated"), name=name) + + assert refusals == [("targets", "name")] + assert registry.get("name") == "saved" + + def test_a_caught_name_check_against_a_saved_entry_is_recorded(self, registry: DefaultInstanceRegistry[_TestItem]): + registry.register(_item("saved"), name="name", metadata={SAVED_ENTRY_METADATA_KEY: "converters"}) + registry.register(_item("startup"), name="startup") + + with record_saved_entry_refusals() as refusals: + for name in ("name", "startup"): + with pytest.raises(ValueError, match="already exists"): + registry.validate_name_available(name) + registry.validate_name_available("free") + + assert refusals == [("converters", "name")] + def test_register_rejects_reserved_name(self): registry: DefaultInstanceRegistry[_TestItem] = DefaultInstanceRegistry(reserved_names={"types"}) diff --git a/tests/unit/registry/test_registry_recipe_restore.py b/tests/unit/registry/test_registry_recipe_restore.py index 4eef924a28..3a6e886bc2 100644 --- a/tests/unit/registry/test_registry_recipe_restore.py +++ b/tests/unit/registry/test_registry_recipe_restore.py @@ -3,6 +3,7 @@ import dataclasses import json +import os import sys import threading from collections.abc import Iterator @@ -27,6 +28,7 @@ ) from pyrit.prompt_target import OpenAIChatTarget, RoundRobinTarget from pyrit.registry import AttackRegistry, AttackTechniqueRegistry, ConverterRegistry, ScorerRegistry, TargetRegistry +from pyrit.registry.instance_registry import SAVED_ENTRY_METADATA_KEY from pyrit.registry.registry_recipe_restore import RecipeBuilder, RecipeRejectedError, RegistryRecipeRestorer from pyrit.registry.registry_recipe_storage import MAX_FAILURE_MESSAGE_LENGTH, RecipeListing, RegistryRecipeStorage from unit.mocks import MockPromptTarget @@ -36,7 +38,11 @@ "class_module": _OPENAI_MODULE, "class_name": "OpenAIChatTarget", "auth_mode": "identity", - "params": {"endpoint": "https://contoso.openai.azure.com/openai/v1", "model_name": "gpt-4o"}, + "params": { + "endpoint": "https://contoso.openai.azure.com/openai/v1", + "headers": "{}", + "model_name": "gpt-4o", + }, } _ROUND_ROBIN_MODULE = "pyrit.prompt_target.round_robin_target" _TECHNIQUE = { @@ -189,6 +195,49 @@ async def test_restores_every_kind_after_the_entries_it_uses(self, tmp_path: Pat assert factory is not None assert factory.technique_tags == ["custom"] + async def test_marks_each_restored_entry_with_its_kind(self, tmp_path: Path) -> None: + _save(tmp_path, RegistryRecipeKind.TECHNIQUES, "formal_red_team", _TECHNIQUE) + _save( + tmp_path, + RegistryRecipeKind.SCORERS, + "refusal", + { + "class_module": "pyrit.score.true_false.self_ask_refusal_scorer", + "class_name": "SelfAskRefusalScorer", + "params": {"chat_target": "contoso_gpt4o"}, + }, + ) + _save( + tmp_path, + RegistryRecipeKind.CONVERTERS, + "tone_formal", + { + "class_module": "pyrit.converter.tone_converter", + "class_name": "ToneConverter", + "params": {"converter_target": "contoso_gpt4o", "tone": "formal"}, + }, + ) + _save(tmp_path, RegistryRecipeKind.TARGETS, "contoso_gpt4o", _IDENTITY_TARGET) + ConverterRegistry.get_registry_singleton().instances.register(Base64Converter(), name="startup_encoder") + + report = await _restore(tmp_path) + + assert report.failures == [] + entries = { + "contoso_gpt4o": TargetRegistry.get_registry_singleton().instances.get_entry("contoso_gpt4o"), + "tone_formal": ConverterRegistry.get_registry_singleton().instances.get_entry("tone_formal"), + "refusal": ScorerRegistry.get_registry_singleton().instances.get_entry("refusal"), + "formal_red_team": AttackTechniqueRegistry.get_registry_singleton().instances.get_entry("formal_red_team"), + "startup_encoder": ConverterRegistry.get_registry_singleton().instances.get_entry("startup_encoder"), + } + assert {name: entry.metadata.get(SAVED_ENTRY_METADATA_KEY) for name, entry in entries.items() if entry} == { + "contoso_gpt4o": "targets", + "tone_formal": "converters", + "refusal": "scorers", + "formal_red_team": "techniques", + "startup_encoder": None, + } + async def test_finds_a_startup_entry_by_its_exact_name(self, tmp_path: Path) -> None: registry = TargetRegistry.get_registry_singleton() # Each name has a different normalized form, so a lookup that normalized names would miss them. @@ -427,6 +476,45 @@ async def test_refuses_a_stored_credential(self, tmp_path: Path) -> None: assert "sk-secret-value" not in failure.message assert [ref.name for ref in report.restored] == ["good"] + @pytest.mark.parametrize( + ("name", "left_out"), + [("server_key", "auth_mode"), ("server_headers", "headers"), ("server_endpoint", "endpoint")], + ) + async def test_refuses_a_target_that_would_use_this_servers_credentials( + self, tmp_path: Path, name: str, left_out: str + ) -> None: + recipe = {key: value for key, value in _IDENTITY_TARGET.items() if key != left_out} + recipe["params"] = {key: value for key, value in _IDENTITY_TARGET["params"].items() if key != left_out} + _save(tmp_path, RegistryRecipeKind.TARGETS, name, recipe) + server_settings = { + "OPENAI_CHAT_KEY": "sk-server-key", + "OPENAI_ADDITIONAL_REQUEST_HEADERS": '{"api-key": "sk-server-key"}', + "OPENAI_CHAT_ENDPOINT": "https://operator:sk-server-key@contoso.openai.azure.com/openai/v1", + } + + with patch.dict(os.environ, server_settings): + report = await _restore(tmp_path) + + (failure,) = report.failures + assert (failure.name, failure.reason) == (name, RestoreFailureReason.UNSUPPORTED_CREDENTIAL) + assert f"'{left_out if left_out != 'auth_mode' else 'api_key'}' is empty" in failure.message + assert "sk-server-key" not in failure.message + assert name not in TargetRegistry.get_registry_singleton().instances.get_names() + + async def test_refuses_a_target_that_signs_in_with_this_servers_account(self, tmp_path: Path) -> None: + _save( + tmp_path, + RegistryRecipeKind.TARGETS, + "copilot", + {"class_module": "pyrit.prompt_target.websocket_copilot_target", "class_name": "WebSocketCopilotTarget"}, + ) + + report = await _restore(tmp_path) + + (failure,) = report.failures + assert (failure.name, failure.reason) == ("copilot", RestoreFailureReason.UNSUPPORTED_CLASS) + assert "Copilot account" in failure.message + async def test_restores_a_class_whose_optional_credentials_are_left_empty(self, tmp_path: Path) -> None: TargetRegistry.get_registry_singleton().register_class(_OptionalKeyTarget) _save( @@ -447,6 +535,39 @@ async def test_restores_a_class_whose_optional_credentials_are_left_empty(self, assert isinstance(target, _OptionalKeyTarget) assert (target.endpoint, target.api_key) == ("https://agent.example.test", None) + async def test_refuses_a_class_that_needs_the_server_upload_folder(self, tmp_path: Path) -> None: + _save( + tmp_path, + RegistryRecipeKind.TARGETS, + "uploader", + { + "class_module": "pyrit.prompt_target.http_target.httpx_api_target", + "class_name": "HTTPXAPITarget", + "params": {"http_url": "https://example.test/upload", "file_path": "/etc/hosts"}, + }, + ) + + report = await _restore(tmp_path) + + (failure,) = report.failures + assert failure.reason is RestoreFailureReason.UNSUPPORTED_CLASS + assert "uploader" not in TargetRegistry.get_registry_singleton().instances.get_names() + + @pytest.mark.parametrize("auth_mode", [[], {}, ["identity"]]) + async def test_reports_malformed_values_without_stopping_the_restore( + self, tmp_path: Path, auth_mode: object + ) -> None: + _save(tmp_path, RegistryRecipeKind.TARGETS, "malformed", {**_IDENTITY_TARGET, "auth_mode": auth_mode}) + _save(tmp_path, RegistryRecipeKind.TARGETS, "good", _IDENTITY_TARGET) + + report = await _restore(tmp_path) + + assert [(failure.name, failure.reason) for failure in report.failures] == [ + ("malformed", RestoreFailureReason.INVALID_RECIPE) + ] + assert "auth_mode" in report.failures[0].message + assert [ref.name for ref in report.restored] == ["good"] + async def test_a_planning_bug_fails_only_that_entry_as_a_logged_build_error( self, tmp_path: Path, caplog: pytest.LogCaptureFixture ) -> None: diff --git a/tests/unit/registry/test_registry_recipe_rules.py b/tests/unit/registry/test_registry_recipe_rules.py new file mode 100644 index 0000000000..955613614e --- /dev/null +++ b/tests/unit/registry/test_registry_recipe_rules.py @@ -0,0 +1,335 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from collections.abc import Iterator +from typing import Any +from urllib.parse import quote + +import pytest + +from pyrit.models.catalog import RestoreFailureReason +from pyrit.registry import ConverterRegistry, TargetRegistry +from pyrit.registry.registry_recipe_rules import SavedEntryRefusedError, check_saved_entry_params + +_ENDPOINT = "https://contoso.openai.azure.com/openai/v1" +_OPENAI_CLASSES = ( + "OpenAIChatTarget", + "OpenAICompletionTarget", + "OpenAIImageTarget", + "OpenAIResponseTarget", + "OpenAITTSTarget", + "OpenAIVideoTarget", + "RealtimeTarget", +) + + +@pytest.fixture(autouse=True) +def isolated_target_registry() -> Iterator[None]: + for registry in (TargetRegistry, ConverterRegistry): + registry.reset_registry_singleton() + yield + for registry in (TargetRegistry, ConverterRegistry): + registry.reset_registry_singleton() + + +def _check(class_name: str, params: dict[str, Any]) -> None: + """Check a target recipe; an ``auth_mode`` key stands for the recipe's top-level auth mode.""" + registry = TargetRegistry.get_registry_singleton() + cls = registry.get_class(class_name) + values = dict(params) + auth_mode = values.pop("auth_mode", None) + check_saved_entry_params( + cls=cls, parameters=registry.get_class_metadata(cls).parameters, params=values, auth_mode=auth_mode + ) + + +def test_check_saved_entry_params_accepts_a_target_that_takes_no_credentials() -> None: + assert TargetRegistry.get_registry_singleton().get_class("TextTarget").supported_auth_modes == () + + _check("TextTarget", {}) + + +def _refusal(class_name: str, params: dict[str, Any]) -> SavedEntryRefusedError: + with pytest.raises(SavedEntryRefusedError) as refusal: + _check(class_name, params) + return refusal.value + + +@pytest.mark.parametrize( + ("class_name", "params"), + [ + *((name, {"auth_mode": "identity", "endpoint": _ENDPOINT, "headers": "{}"}) for name in _OPENAI_CLASSES), + ("OpenAIChatTarget", {"auth_mode": "none", "endpoint": "http://localhost:11434/v1"}), + ("OpenAIChatTarget", {"auth_mode": "none", "endpoint": "http://localhost:11434/v1", "headers": "{}"}), + ("AzureMLChatTarget", {"auth_mode": "identity", "endpoint": "https://contoso.inference.ml.azure.com/score"}), + ("PromptShieldTarget", {"auth_mode": "identity", "endpoint": "https://contoso.cognitiveservices.azure.com/"}), + ( + "AzureBlobStorageTarget", + {"auth_mode": "identity", "container_url": "https://contoso.blob.core.windows.net/prompts"}, + ), + ("A2ATarget", {"endpoint": "https://agent.example.test"}), + ("TextTarget", {}), + ("RoundRobinTarget", {"targets": ["contoso_gpt4o", "contoso_gpt4o_mini"]}), + ], +) +def test_check_saved_entry_params_accepts_entra_id_and_no_key(class_name: str, params: dict[str, Any]) -> None: + _check(class_name, params) + + +@pytest.mark.parametrize("class_name", _OPENAI_CLASSES) +@pytest.mark.parametrize("auth_mode", [None, "api_key"]) +def test_check_saved_entry_params_refuses_an_openai_target_without_a_key_or_identity( + class_name: str, auth_mode: str | None +) -> None: + params = {"endpoint": _ENDPOINT, "headers": "{}"} | ({"auth_mode": auth_mode} if auth_mode else {}) + + refusal = _refusal(class_name, params) + + assert refusal.reason is RestoreFailureReason.UNSUPPORTED_CREDENTIAL + assert str(refusal) == f"'api_key' is empty, so '{class_name}' can use this server's credential instead." + + +@pytest.mark.parametrize( + ("class_name", "params", "name"), + [ + ("AzureMLChatTarget", {"endpoint": "https://contoso.inference.ml.azure.com/score"}, "api_key"), + ("AzureMLChatTarget", {"endpoint": "https://contoso.inference.ml.azure.com/score", "api_key": ""}, "api_key"), + ("PromptShieldTarget", {"endpoint": "https://contoso.cognitiveservices.azure.com/"}, "api_key"), + ("AzureBlobStorageTarget", {"container_url": "https://contoso.blob.core.windows.net/prompts"}, "sas_token"), + ("GitHubCopilotTarget", {"model_name": "gpt-4o"}, "github_token"), + ("HackAPromptTarget", {"challenge_slug": "pathogen", "competition_slug": "cbrne"}, "session_id"), + ("LiteLLMChatTarget", {"model_name": "gpt-4o"}, "api_key"), + ("LiteLLMDecisionTarget", {"model_name": "gpt-4o"}, "api_key"), + ], +) +def test_check_saved_entry_params_refuses_an_empty_key_the_class_reads_from_this_server( + class_name: str, params: dict[str, Any], name: str +) -> None: + refusal = _refusal(class_name, params) + + assert refusal.reason is RestoreFailureReason.UNSUPPORTED_CREDENTIAL + assert str(refusal).startswith(f"'{name}' is empty, so '{class_name}' can use this server's credential") + + +def test_check_saved_entry_params_refuses_identity_without_headers() -> None: + refusal = _refusal("OpenAIChatTarget", {"auth_mode": "identity", "endpoint": _ENDPOINT}) + + assert refusal.reason is RestoreFailureReason.UNSUPPORTED_CREDENTIAL + assert str(refusal).endswith("Saved entries set it to '{}'.") + + +@pytest.mark.parametrize("auth_mode", ["identity", "none"]) +def test_check_saved_entry_params_refuses_headers_that_can_carry_a_credential(auth_mode: str) -> None: + headers = '{"api-key": "sk-header-secret"}' + + refusal = _refusal("OpenAIChatTarget", {"auth_mode": auth_mode, "endpoint": _ENDPOINT, "headers": headers}) + + assert refusal.reason is RestoreFailureReason.UNSUPPORTED_CREDENTIAL + assert str(refusal) == "'headers' can carry a credential, which saved entries can't hold." + + +@pytest.mark.parametrize( + ("class_name", "phrase"), + [ + ("WebSocketCopilotTarget", "Copilot account"), + ("HTTPTarget", "raw HTTP request"), + ("HTTPXAPITarget", "uploads files"), + ("HuggingFaceChatTarget", "Hugging Face token or model files"), + ], +) +def test_check_saved_entry_params_refuses_a_class_that_cant_be_saved(class_name: str, phrase: str) -> None: + refusal = _refusal(class_name, {}) + + assert refusal.reason is RestoreFailureReason.UNSUPPORTED_CLASS + assert phrase in str(refusal) + + +@pytest.mark.parametrize( + "url", + [ + _ENDPOINT, + "https://contoso.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21", + "https://contoso.openai.azure.com/openai/v1?API-Version=preview", + "http://localhost:11434/v1", + ], +) +def test_check_saved_entry_params_accepts_a_url_without_a_credential(url: str) -> None: + _check("OpenAIChatTarget", {"auth_mode": "identity", "endpoint": url, "headers": "{}"}) + + +@pytest.mark.parametrize( + "url", + [ + "https://contoso.blob.core.windows.net/prompts?sv=2024-11-04&sig=SECRET", + "https://contoso.blob.core.windows.net/prompts%3Fsv%3D2024-11-04%26sig%3DSECRET", + "https://contoso.blob.core.windows.net/prompts%253Fsig%253DSECRET", + "https://contoso.blob.core.windows.net/prompts&sig=SECRET", + "https://user:SECRET@contoso.example/v1", + "https://SECRET@contoso.example/v1", + "https://contoso.example/v1#access_token=SECRET", + "https://contoso.example/v1?code=SECRET", + "https://contoso.example/v1\uff1fcode=SECRET", + "https://contoso.example/v1?api-version=2024-10-21;code=SECRET", + "https://contoso.example/v1%2525253Fsig%2525253DSECRET", + "https://contoso.example/v1\uff053Fsig\uff053DSECRET", + "https://contoso.example/v1?=SECRET", + "https://contoso.example/v1?api-version=2024-10-21&=SECRET", + "https://contoso.example/v1?api-version=2024-10-21?code=SECRET", + "https://contoso.example/v1;code=SECRET", + "https://contoso.example/v1;SECRET", + "https://contoso.example/v1&SECRET", + "https://contoso.blob.core.windows.net;sig=SECRET/prompts", + "https://contoso.blob.core.windows.net;SIG=SECRET/prompts", + "https://contoso.example/v1?sv=SECRET", + "https://contoso.openai.azure.com/openai/v1?api-key=SECRET", + "https://contoso.cognitiveservices.azure.com/?subscription-key=SECRET", + "https://contoso.example/v1?api-version=2024-10-21=SECRET", + "https://contoso.example/v1/key=SECRET", + ], +) +def test_check_saved_entry_params_refuses_a_url_that_can_carry_a_credential(url: str) -> None: + refusal = _refusal("OpenAIChatTarget", {"auth_mode": "identity", "endpoint": url, "headers": "{}"}) + + assert refusal.reason is RestoreFailureReason.UNSUPPORTED_CREDENTIAL + assert str(refusal).startswith("'endpoint' can carry a credential in a URL") + assert "SECRET" not in str(refusal) + + +def test_check_saved_entry_params_refuses_a_path_it_cant_finish_reading() -> None: + card_path = "%41" + for _ in range(10): + card_path = quote(card_path, safe="") + + refusal = _refusal("A2ATarget", {"endpoint": "https://agent.example.test", "agent_card_path": card_path}) + + assert str(refusal).startswith("'agent_card_path' can carry a credential in a URL") + + +@pytest.mark.parametrize( + ("class_name", "params", "name"), + [ + ("OpenAIChatTarget", {"auth_mode": "identity", "headers": "{}"}, "endpoint"), + ("OpenAIChatTarget", {"auth_mode": "identity", "endpoint": " ", "headers": "{}"}, "endpoint"), + ("AzureMLChatTarget", {"auth_mode": "identity"}, "endpoint"), + ("PromptShieldTarget", {"auth_mode": "identity"}, "endpoint"), + ("AzureBlobStorageTarget", {"auth_mode": "identity"}, "container_url"), + ], +) +def test_check_saved_entry_params_refuses_a_destination_left_to_this_server( + class_name: str, params: dict[str, Any], name: str +) -> None: + refusal = _refusal(class_name, params) + + assert refusal.reason is RestoreFailureReason.UNSUPPORTED_CREDENTIAL + assert str(refusal).startswith(f"'{name}' is empty, so '{class_name}' would send requests where this server's") + + +@pytest.mark.parametrize("card_path", [None, "/.well-known/agent-card.json", "agents/redteam/card.json"]) +def test_check_saved_entry_params_accepts_an_agent_card_path_without_a_credential(card_path: str | None) -> None: + params = {"endpoint": "https://agent.example.test"} | ({"agent_card_path": card_path} if card_path else {}) + + _check("A2ATarget", params) + + +@pytest.mark.parametrize( + "card_path", + [ + "card?access_token=SECRET", + "/.well-known/agent.json?api_key=SECRET", + "card#SECRET", + "//SECRET@h", + "card.json&access_token=SECRET", + "card.json;session=SECRET", + "card.json;SECRET", + "card.json&SECRET", + "card.json?api-version=1", + ], +) +def test_check_saved_entry_params_refuses_an_agent_card_path_that_can_carry_a_credential(card_path: str) -> None: + refusal = _refusal("A2ATarget", {"endpoint": "https://agent.example.test", "agent_card_path": card_path}) + + assert str(refusal).startswith("'agent_card_path' can carry a credential in a URL") + assert "SECRET" not in str(refusal) + + +@pytest.mark.parametrize( + "endpoint", + [ + "ftp://agent.example.test", + "file://localhost/tmp/card.json", + "https:agent.example.test/card", + "https://[agent.example.test?access_token=SECRET", + 123, + ], +) +def test_check_saved_entry_params_refuses_an_agent_endpoint_that_isnt_a_plain_address(endpoint: object) -> None: + refusal = _refusal("A2ATarget", {"endpoint": endpoint}) + + assert str(refusal).startswith("'endpoint' can carry a credential in a URL") + + +@pytest.mark.parametrize( + ("class_name", "params"), + [ + ("AzureMLChatTarget", {"auth_mode": "none", "endpoint": "https://contoso.inference.ml.azure.com/score"}), + ("GitHubCopilotTarget", {"auth_mode": "identity", "model_name": "gpt-4o"}), + ("OpenAIChatTarget", {"auth_mode": ["identity"], "endpoint": _ENDPOINT, "headers": "{}"}), + ], +) +def test_check_saved_entry_params_refuses_an_auth_mode_the_class_doesnt_support( + class_name: str, params: dict[str, Any] +) -> None: + refusal = _refusal(class_name, params) + + assert refusal.reason is RestoreFailureReason.UNSUPPORTED_CREDENTIAL + assert f"'{class_name}' only supports auth_mode" in str(refusal) + + +def test_check_saved_entry_params_leaves_text_that_isnt_a_destination_alone() -> None: + registry = ConverterRegistry.get_registry_singleton() + cls = registry.get_class("PromptTemplateConverter") + params = {"template": "https://example.com/search?q={{ prompt }}&sig=1#top"} + + check_saved_entry_params(cls=cls, parameters=registry.get_class_metadata(cls).parameters, params=params) + + +@pytest.mark.parametrize( + "endpoint", + [ + "contoso.openai.azure.com/openai/v1", + "/openai/v1", + "//contoso.openai.azure.com/openai/v1", + "//user:SECRET@contoso.example/v1", + "https://[contoso.example/v1?key=SECRET", + "https:/contoso.example/v1?code=SECRET", + "https:/user:SECRET@contoso.example/v1", + "ftp://contoso.example/v1", + "contoso.example/v1?code=SECRET", + "https://user:SECRET\uff20contoso.example/v1", + ], +) +def test_check_saved_entry_params_refuses_an_endpoint_that_isnt_a_full_address(endpoint: str) -> None: + refusal = _refusal("OpenAIChatTarget", {"auth_mode": "identity", "endpoint": endpoint, "headers": "{}"}) + + assert refusal.reason is RestoreFailureReason.INVALID_RECIPE + assert str(refusal) == "'endpoint' must be a full http(s) or ws(s) address in a saved target." + + +@pytest.mark.parametrize("class_name", _OPENAI_CLASSES) +def test_check_saved_entry_params_refuses_an_openai_target_while_the_server_sets_sdk_headers( + class_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Ocp-Apim-Subscription-Key: server-secret") + + refusal = _refusal(class_name, {"auth_mode": "identity", "endpoint": _ENDPOINT, "headers": "{}"}) + + assert refusal.reason is RestoreFailureReason.UNSUPPORTED_CREDENTIAL + assert "OPENAI_CUSTOM_HEADERS" in str(refusal) + assert "server-secret" not in str(refusal) + + +def test_check_saved_entry_params_ignores_sdk_headers_for_other_targets(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Ocp-Apim-Subscription-Key: server-secret") + + _check("A2ATarget", {"endpoint": "https://agent.example.test"}) + _check("TextTarget", {}) diff --git a/tests/unit/registry/test_registry_recipe_storage.py b/tests/unit/registry/test_registry_recipe_storage.py index 28ce0142a9..7eadfc4c8c 100644 --- a/tests/unit/registry/test_registry_recipe_storage.py +++ b/tests/unit/registry/test_registry_recipe_storage.py @@ -22,7 +22,11 @@ TechniqueRecipe, ) from pyrit.models.catalog.registry_recipe import RECIPE_SCHEMA_VERSION -from pyrit.registry.file_document_storage import DocumentStorageUnavailableError, _read_regular_file +from pyrit.registry.file_document_storage import ( + DocumentConflictError, + DocumentStorageUnavailableError, + _read_regular_file, +) from pyrit.registry.registry_recipe_storage import MAX_FAILURE_MESSAGE_LENGTH, RegistryRecipeStorage _TARGET = { @@ -355,3 +359,62 @@ def test_raises_when_the_folder_cannot_be_listed(self, tmp_path: Path) -> None: with pytest.raises(DocumentStorageUnavailableError): RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)).list_recipes() + + +class TestRegistryRecipeStorageWrites: + def test_creates_a_recipe_that_lists_back_with_its_owner_and_version(self, tmp_path: Path) -> None: + storage = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)) + recipe = ComponentRecipe.model_validate(_TARGET | {"owner": {"oid": "user-oid", "name": "Ada"}}) + + version = storage.create_recipe(name="contoso_gpt4o", recipe=recipe) + + stored = storage.list_recipes().recipes["contoso_gpt4o"] + assert (stored.recipe, stored.version) == (recipe, version) + assert json.loads((tmp_path / "targets" / "contoso_gpt4o.json").read_text())["owner"] == { + "oid": "user-oid", + "name": "Ada", + } + + def test_leaves_out_an_owner_that_was_not_recorded(self, tmp_path: Path) -> None: + storage = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)) + + storage.create_recipe(name="contoso_gpt4o", recipe=ComponentRecipe.model_validate(_TARGET)) + + assert "owner" not in json.loads((tmp_path / "targets" / "contoso_gpt4o.json").read_text()) + + @pytest.mark.parametrize( + ("kind", "recipe"), + [ + (RegistryRecipeKind.TARGETS, ComponentRecipe(class_module="m", class_name="TextTarget")), + (RegistryRecipeKind.TECHNIQUES, TechniqueRecipe.model_validate({"type": "RedTeamingAttack"})), + ], + ) + def test_labels_a_new_recipe_with_the_format_version_this_pyrit_writes( + self, tmp_path: Path, kind: RegistryRecipeKind, recipe: ComponentRecipe | TechniqueRecipe + ) -> None: + storage = RegistryRecipeStorage(kind=kind, source=str(tmp_path)) + + with patch("pyrit.registry.registry_recipe_storage.RECIPE_SCHEMA_VERSION", 2): + storage.create_recipe(name="contoso", recipe=recipe) + + assert json.loads((tmp_path / kind.value / "contoso.json").read_text())["schema_version"] == 2 + assert recipe.schema_version == 1 + + def test_refuses_to_create_a_recipe_that_is_already_saved(self, tmp_path: Path) -> None: + storage = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)) + storage.create_recipe(name="contoso_gpt4o", recipe=ComponentRecipe.model_validate(_TARGET)) + before = (tmp_path / "targets" / "contoso_gpt4o.json").read_bytes() + other = ComponentRecipe.model_validate(_TARGET | {"params": {}}) + + with pytest.raises(DocumentConflictError): + storage.create_recipe(name="contoso_gpt4o", recipe=other) + + assert (tmp_path / "targets" / "contoso_gpt4o.json").read_bytes() == before + + def test_deletes_a_recipe_once(self, tmp_path: Path) -> None: + storage = RegistryRecipeStorage(kind=RegistryRecipeKind.TARGETS, source=str(tmp_path)) + storage.create_recipe(name="contoso_gpt4o", recipe=ComponentRecipe.model_validate(_TARGET)) + + assert storage.delete_recipe("contoso_gpt4o") is True + assert storage.delete_recipe("contoso_gpt4o") is False + assert storage.list_recipes().recipes == {} diff --git a/tests/unit/registry/test_target_registry.py b/tests/unit/registry/test_target_registry.py index 2a4b8fddad..66b19ba022 100644 --- a/tests/unit/registry/test_target_registry.py +++ b/tests/unit/registry/test_target_registry.py @@ -11,6 +11,7 @@ import pytest from pyrit.models import ComponentIdentifier, ComponentType, Message, MessagePiece +from pyrit.models.catalog import RegistryEntryOrigin from pyrit.prompt_target import ( CopilotType, PlaywrightCopilotTarget, @@ -18,8 +19,10 @@ RoundRobinTarget, TargetCapabilities, TargetConfiguration, + TextTarget, ) from pyrit.registry import TargetMetadata, TargetRegistry +from pyrit.registry.instance_registry import registration_origin from pyrit.registry.resolution import derive_parameters @@ -166,6 +169,38 @@ class NotATarget: assert len(registry.instances) == 0 +class TestTargetRegistryTemporaryClasses: + """Tests for telling which classes a scan's launch scripts added.""" + + def test_a_class_that_first_appears_in_a_launch_script_is_temporary(self, registry: TargetRegistry): + registry.register_class(MockPromptChatTarget) + with registration_origin(RegistryEntryOrigin.TEMPORARY): + registry.register_class(TextTarget) + registry.register_class(MockPromptTarget) + registry.register_class(MockPromptChatTarget) + with registration_origin(RegistryEntryOrigin.CONFIGURATION): + registry.register_class(MockPromptTarget, name="startup_mock") + + assert registry.is_temporary_class("MockPromptTarget") + assert not any( + registry.is_temporary_class(name) + for name in ("TextTarget", "OpenAIChatTarget", "MockPromptChatTarget", "startup_mock") + ) + + def test_classes_found_on_first_use_in_a_launch_script_are_not_temporary(self, registry: TargetRegistry): + with registration_origin(RegistryEntryOrigin.TEMPORARY): + assert registry.get_class("TextTarget") is TextTarget + + assert not registry.is_temporary_class("TextTarget") + + def test_a_class_a_launch_script_puts_in_place_of_another_is_temporary(self, registry: TargetRegistry): + with registration_origin(RegistryEntryOrigin.TEMPORARY): + registry.register_class(MockPromptTarget, name="TextTarget") + + assert registry.get_class("TextTarget") is MockPromptTarget + assert registry.is_temporary_class("TextTarget") + + @pytest.mark.usefixtures("patch_central_database") class TestTargetRegistryGetInstanceByName: """Tests for instance lookup via ``instances.get``.""" diff --git a/tests/unit/setup/test_converter_initializer.py b/tests/unit/setup/test_converter_initializer.py index 667624cfb1..9fb3138156 100644 --- a/tests/unit/setup/test_converter_initializer.py +++ b/tests/unit/setup/test_converter_initializer.py @@ -12,6 +12,7 @@ from pyrit.converter import Base64Converter, LeetspeakConverter, ROT13Converter, VariationConverter from pyrit.registry import ConverterRegistry, InitializerRegistry, TargetRegistry +from pyrit.registry.instance_registry import SAVED_ENTRY_METADATA_KEY, SavedEntryReplacementError from pyrit.setup.initializers import ConverterInitializer from pyrit.setup.initializers.converters import ConverterConfig @@ -80,6 +81,20 @@ async def test_initialize_replaces_converter_with_current_target() -> None: assert converter._converter_target is not first_target +async def test_initialize_does_not_skip_a_saved_converter_it_would_replace() -> None: + registry = ConverterRegistry.get_registry_singleton() + saved = Base64Converter() + registry.instances.register(saved, name="rot13", metadata={SAVED_ENTRY_METADATA_KEY: "converters"}) + + with ( + patch.object(ConverterInitializer, "CONFIGS", _get_configs("base64", "rot13")), + pytest.raises(SavedEntryReplacementError, match="Saved converter 'rot13'"), + ): + await ConverterInitializer().initialize_async() + + assert registry.instances.get("rot13") is saved + + async def test_initialize_skips_converter_without_declared_target() -> None: registry = ConverterRegistry.get_registry_singleton() diff --git a/tests/unit/setup/test_scorer_initializer.py b/tests/unit/setup/test_scorer_initializer.py index af33ae7801..4c4984f0d7 100644 --- a/tests/unit/setup/test_scorer_initializer.py +++ b/tests/unit/setup/test_scorer_initializer.py @@ -7,9 +7,11 @@ from unittest.mock import MagicMock, patch import pytest +from unit.mocks import MockPromptTarget from pyrit.prompt_target import LiteLLMDecisionTarget, OpenAIChatTarget from pyrit.registry import ScorerRegistry, TargetRegistry +from pyrit.registry.instance_registry import SAVED_ENTRY_METADATA_KEY, SavedEntryReplacementError from pyrit.score import ( DecisionScorer, FloatScaleThresholdScorer, @@ -27,6 +29,7 @@ GPT4O_UNSAFE_TEMP9_TARGET, GPT5_1_TARGET, GPT5_4_TARGET, + REFUSAL_GPT5_4, ) @@ -364,6 +367,37 @@ def _register_mock_target(self, *, name: str, underlying_model: str = "gpt-4o") registry.instances.register(target, name=name) return target + @patch("pyrit.setup.initializers.scorers.find_objective_metrics_by_eval_hash", return_value=None) + async def test_refuses_to_tag_a_saved_scorer_it_could_not_build(self, mock_find_metrics) -> None: + self._register_mock_target(name=GPT4O_TARGET) + saved = SubStringScorer(substring="saved") + registry = ScorerRegistry.get_registry_singleton() + registry.instances.register(saved, name=REFUSAL_GPT5_4, metadata={SAVED_ENTRY_METADATA_KEY: "scorers"}) + + with pytest.raises(SavedEntryReplacementError, match=f"Saved scorer '{REFUSAL_GPT5_4}'"): + await ScorerInitializer().initialize_async() + + entry = registry.instances.get_entry(REFUSAL_GPT5_4) + assert entry is not None + assert entry.instance is saved + assert ScorerInitializerTags.BEST_REFUSAL not in entry.tags + + @patch("pyrit.setup.initializers.scorers.find_objective_metrics_by_eval_hash", return_value=None) + async def test_refuses_to_make_a_saved_scorer_the_default_objective(self, mock_find_metrics) -> None: + TargetRegistry.get_registry_singleton().instances.register(MockPromptTarget(), name="unrelated_target") + saved = SubStringScorer(substring="saved") + registry = ScorerRegistry.get_registry_singleton() + name = ScorerInitializer.SCALE_AND_REFUSAL + registry.instances.register(saved, name=name, metadata={SAVED_ENTRY_METADATA_KEY: "scorers"}) + + with pytest.raises(SavedEntryReplacementError, match=f"Saved scorer '{name}'"): + await ScorerInitializer().initialize_async() + + entry = registry.instances.get_entry(name) + assert entry is not None + assert entry.instance is saved + assert ScorerInitializerTags.BEST_OBJECTIVE not in entry.tags + @patch("pyrit.setup.initializers.scorers.find_objective_metrics_by_eval_hash") async def test_best_objective_tags_best_scorer(self, mock_find_metrics: MagicMock) -> None: """Test that _tag_best_objective tags the scorer with highest F1.""" diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index a88e43e259..7342ada631 100644 --- a/tests/unit/setup/test_targets_initializer.py +++ b/tests/unit/setup/test_targets_initializer.py @@ -9,6 +9,7 @@ from pyrit.prompt_target import LiteLLMDecisionTarget, OpenAIChatTarget from pyrit.registry import TargetRegistry +from pyrit.registry.instance_registry import SAVED_ENTRY_METADATA_KEY, SavedEntryReplacementError from pyrit.setup.initializers import TargetInitializer from pyrit.setup.initializers.targets import TARGET_CONFIGS, generate_rr_name, get_behavioral_key @@ -932,6 +933,20 @@ async def test_individual_targets_still_accessible_after_auto_group(self) -> Non assert registry.instances.get("azure_openai_gpt4o") is not None assert registry.instances.get("azure_openai_gpt4o2") is not None + async def test_refuses_to_leave_a_saved_target_in_place_of_a_group(self) -> None: + os.environ.update(self.GPT4O_1_ENV) + os.environ.update(self.GPT4O_2_ENV) + registry = TargetRegistry.get_registry_singleton() + saved = OpenAIChatTarget(endpoint="https://saved.test/v1", model_name="gpt-4o", api_key="saved-key") + registry.instances.register( + saved, name="OpenAIChatTarget_gpt-4o_rr", metadata={SAVED_ENTRY_METADATA_KEY: "targets"} + ) + + with pytest.raises(SavedEntryReplacementError, match="Saved target 'OpenAIChatTarget_gpt-4o_rr'"): + await TargetInitializer().initialize_async() + + assert registry.instances.get("OpenAIChatTarget_gpt-4o_rr") is saved + async def test_no_round_robin_when_single_target(self) -> None: """Test that no round-robin is created when only one target has a given model.""" os.environ.update(self.GPT4O_1_ENV)