From 9f409abce269546ede61594857f50bb13246d3ea Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 7 Jul 2026 16:06:08 -0700 Subject: [PATCH 01/34] feat(setup): add plug-in loading to initialize_pyrit_async Adds a plug-in mechanism so an operator can run non-disclosable scenarios (and self-contained datasets) from a stock public PyRIT install, from a pre-built wheel referenced by PLUGIN_WHEEL, without those components living in the public repo. Plug-in loading runs as a guaranteed-first phase inside initialize_pyrit_async: after central memory is set and before the configured initializers, so plug-in datasets/scenarios register before LoadDefaultDatasets and PreloadScenarioMetadata read the registry. Ordering is true by construction, independent of .pyrit_conf list position. The wheel is extracted (stdlib zipfile - never pip/.venv) to .plugin//, prepended to sys.path, imported (dataset providers self-register), and its bootstrap (a top-level register() callable or a shipped PyRITInitializer subclass) is run; the loader then asserts the plug-in registered something. No-op when PLUGIN_WHEEL is unset. Fails closed by default; fail-open via PLUGIN_FAIL_OPEN or initialize_pyrit_async(plugin_fail_open=True). Guards against silent-failure modes: extraction (not zipimport) so __file__-relative datasets resolve, atomic extraction, submodule discovery so components register even if the package __init__ does not import them, a package-shadowing guard when an installed package of the same name is importable, a loud warning when a plug-in dataset name collides with an existing name (the resolver is memory-authoritative, so the plug-in copy would otherwise be silently ignored), and rollback of sys.path / sys.modules / registries on failure. Wiring: .gitignore keeps .plugin/ (ignores its contents); .env_example documents the PLUGIN_* variables. Tests build a mock wheel (no dependency on any real plug-in) covering extraction, registration, ordering, collisions, and failure modes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .env_example | 30 + .gitignore | 5 + .plugin/.gitkeep | 2 + .pyrit_conf_example | 4 + pyrit/setup/initialization.py | 13 + pyrit/setup/plugin_loader.py | 624 ++++++++++++++++++++ tests/unit/setup/test_plugin_loader.py | 772 +++++++++++++++++++++++++ 7 files changed, 1450 insertions(+) create mode 100644 .plugin/.gitkeep create mode 100644 pyrit/setup/plugin_loader.py create mode 100644 tests/unit/setup/test_plugin_loader.py diff --git a/.env_example b/.env_example index e6d8f403e9..01ae5e8108 100644 --- a/.env_example +++ b/.env_example @@ -250,6 +250,36 @@ AZURE_OPENAI_VIDEO_MODEL="sora-2" AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" OPENAI_VIDEO_ENDPOINT = ${AZURE_OPENAI_VIDEO_ENDPOINT} + +################################## +# PLUG-INS +# +# PyRIT can load a non-disclosable plug-in (extra datasets + scenarios) from a +# pre-built wheel at initialization time. Plug-in loading runs automatically as a +# guaranteed-first phase inside initialize_pyrit_async (before the configured +# initializers) — there is nothing to add to .pyrit_conf. The wheel is extracted to +# .plugin//, prepended to sys.path, imported, and its bootstrap is run so its +# datasets/scenarios register like built-ins. Extraction only — never pip/.venv. +# +# WARNING: loading a plug-in executes third-party code in the PyRIT process. +# Whoever can set PLUGIN_* / write this .env can run code on the host. Treat this +# file as sensitive. +################################### + +# Path to a pre-built plug-in wheel on disk. Setting it enables plug-in loading; +# leaving it unset is a no-op. +# PLUGIN_WHEEL="/abs/path/to/my_plugin-0.0.0-py3-none-any.whl" + +# Optional: the wheel's top-level import package. Auto-detected from the wheel +# when omitted; set this to disambiguate multi-package wheels. +# PLUGIN_PACKAGE="my_plugin" + +# Optional: continue (with a warning) instead of failing when the plug-in cannot +# be loaded. Equivalent to initialize_pyrit_async(plugin_fail_open=True). +# PLUGIN_FAIL_OPEN="false" + +# Optional: override the base extraction directory (defaults to /.plugin). +# PLUGIN_DIR="/abs/path/to/.plugin" OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" diff --git a/.gitignore b/.gitignore index f92df639f0..6908547442 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,11 @@ dbdata/ eval/ default_memory.json.memory +# Plug-in extraction working directory (see DefaultPluginInitializer / PLUGIN_WHEEL). +# The directory is kept; extracted plug-in artifacts are ignored. +.plugin/* +!.plugin/.gitkeep + # Frontend build artifacts copied to backend for packaging pyrit/backend/frontend/ diff --git a/.plugin/.gitkeep b/.plugin/.gitkeep new file mode 100644 index 0000000000..0ed2780a0d --- /dev/null +++ b/.plugin/.gitkeep @@ -0,0 +1,2 @@ +# Keeps the .plugin/ directory in version control while its contents stay ignored. +# PyRIT extracts plug-in wheels here at runtime (see PLUGIN_WHEEL). diff --git a/.pyrit_conf_example b/.pyrit_conf_example index bbf1e98903..95c1b928f5 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -28,6 +28,10 @@ memory_db_type: sqlite # - load_default_datasets: Loads datasets into memory so scenarios can run # - preload_scenario_metadata: Preloads scenario metadata into the registry # +# Note: plug-ins (PLUGIN_WHEEL) are NOT configured here. They load automatically as a +# guaranteed-first phase inside initialize_pyrit_async — before these initializers — so +# there is no ordering to get right. See the PLUGIN_* variables in .env_example. +# # Each initializer can be specified as: # - A simple string (name only) # - A dictionary with 'name' and optional 'args' for parameters diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 43e7bfd4e0..a549338cc7 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -203,6 +203,7 @@ async def initialize_pyrit_async( env_files: Sequence[pathlib.Path] | None = None, env_akv_ref: Sequence[str] | None = None, silent: bool = False, + plugin_fail_open: bool | None = None, **memory_instance_kwargs: Any, ) -> None: """ @@ -224,6 +225,9 @@ async def initialize_pyrit_async( so local files take precedence over AKV. Requires ``azure-keyvault-secrets``. silent (bool): If True, suppresses print statements about environment file loading and schema migration. Defaults to False. + plugin_fail_open (bool | None): Overrides ``PLUGIN_FAIL_OPEN`` for plug-in loading. When True, + a plug-in (``PLUGIN_WHEEL``) that fails to load is skipped with a warning instead of raising. + Defaults to None (use the ``PLUGIN_FAIL_OPEN`` environment variable, else fail-closed). **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. Raises: @@ -259,6 +263,15 @@ async def initialize_pyrit_async( CentralMemory.set_memory_instance(memory) + # Load a configured plug-in (PLUGIN_WHEEL) as a guaranteed-first phase: after memory + # is set (a plug-in bootstrap may use it) and BEFORE any configured initializers run, + # so plug-in datasets/scenarios are registered before LoadDefaultDatasets and + # PreloadScenarioMetadata read the registry. Ordering is true by construction here, so + # it does not depend on .pyrit_conf list position. No-op unless PLUGIN_WHEEL is set. + from pyrit.setup.plugin_loader import load_plugin_if_configured_async + + await load_plugin_if_configured_async(fail_open=plugin_fail_open) + # Combine directly provided initializers with those loaded from scripts all_initializers = list(initializers) if initializers else [] diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py new file mode 100644 index 0000000000..c632d1a242 --- /dev/null +++ b/pyrit/setup/plugin_loader.py @@ -0,0 +1,624 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Load a non-disclosable PyRIT plug-in from a pre-built wheel at initialization time. + +A plug-in is a pure-Python wheel that ships dataset providers and/or scenarios that +must not live in the public PyRIT repo. The loader **extracts** the wheel (stdlib +``zipfile`` — never ``pip``/``.venv``) into ``.plugin//``, prepends that directory +to ``sys.path``, imports the package (so ``SeedDatasetProvider`` subclasses self-register), +and runs the plug-in's bootstrap (a top-level ``register()`` callable or a shipped +``PyRITInitializer`` subclass) which registers the plug-in's scenarios. + +``load_plugin_if_configured_async`` is invoked as a guaranteed-first phase inside +``initialize_pyrit_async`` — after central memory is set and **before** any configured +initializers run — so plug-in datasets and scenarios are registered before +``LoadDefaultDatasets`` / ``PreloadScenarioMetadata`` read the registry. Ordering is +therefore true by construction, without relying on ``.pyrit_conf`` list position. It is a +no-op when ``PLUGIN_WHEEL`` is unset. +""" + +from __future__ import annotations + +import importlib +import inspect +import logging +import os +import pkgutil +import shutil +import sys +import zipfile +from pathlib import Path +from typing import TYPE_CHECKING + +from pyrit.setup.pyrit_initializer import PyRITInitializer + +if TYPE_CHECKING: + from collections.abc import Mapping + from types import ModuleType + + from pyrit.registry import ScenarioRegistry + +logger = logging.getLogger(__name__) + +_TRUE_TOKENS = frozenset({"1", "true", "yes", "on"}) + + +async def load_plugin_if_configured_async(*, fail_open: bool | None = None) -> None: + """ + Load the plug-in referenced by ``PLUGIN_WHEEL`` if one is configured. + + Convenience entry point invoked by ``initialize_pyrit_async``. A no-op when + ``PLUGIN_WHEEL`` is unset. + + Args: + fail_open: If provided, overrides the ``PLUGIN_FAIL_OPEN`` environment variable. + When True, a plug-in that fails to load is skipped with a warning. + + Raises: + PluginLoadError: If the plug-in is configured but fails to load and fail-open is + not enabled. + """ + await PluginLoader(fail_open=fail_open).load_async() + + +def _name_owned_by(module_name: str, package_name: str) -> bool: + """ + Return whether a module name belongs to the given plug-in package. + + Args: + module_name: A dotted module name. + package_name: The plug-in's top-level package name. + + Returns: + bool: True if ``module_name`` is the package or one of its submodules. + """ + return module_name == package_name or module_name.startswith(f"{package_name}.") + + +def _module_owned_by(cls: type, package_name: str) -> bool: + """ + Return whether ``cls`` is defined within the given plug-in package. + + Args: + cls: The class to check. + package_name: The plug-in's top-level package name. + + Returns: + bool: True if the class's module is the package or one of its submodules. + """ + return _name_owned_by(cls.__module__ or "", package_name) + + +class PluginLoadError(RuntimeError): + """Raised when a configured plug-in fails to load and fail-open is not enabled.""" + + +class PluginLoader: + """ + Extract and register a PyRIT plug-in wheel referenced by ``PLUGIN_WHEEL``. + + No-op unless ``PLUGIN_WHEEL`` is set. When set, the wheel is extracted to + ``.plugin//`` (never installed), imported, and its bootstrap is run so its + datasets and scenarios register like built-ins. Fails closed by default; set + ``fail_open`` (constructor / ``initialize_pyrit_async`` param) or ``PLUGIN_FAIL_OPEN`` + to continue without the plug-in when it cannot be loaded. + """ + + def __init__(self, *, fail_open: bool | None = None) -> None: + """ + Initialize the loader. + + Args: + fail_open: If provided, overrides ``PLUGIN_FAIL_OPEN``. When True, a plug-in + that fails to load is skipped with a warning instead of raising. + """ + self._explicit_fail_open = fail_open + + async def load_async(self) -> None: + """ + Load the plug-in referenced by ``PLUGIN_WHEEL`` (no-op when unset). + + Raises: + PluginLoadError: If the plug-in is configured but fails to load and fail-open + is not enabled. + """ + wheel_env = os.getenv("PLUGIN_WHEEL") + if not wheel_env: + logger.debug("PLUGIN_WHEEL is not set; plug-in loading is a no-op.") + return + + fail_open = self._resolve_fail_open() + + try: + await self._load_plugin_async(wheel_path=Path(wheel_env).expanduser()) + except Exception as exc: + if fail_open: + logger.warning( + "Plug-in from PLUGIN_WHEEL='%s' failed to load; fail_open is set so continuing without it: %s", + wheel_env, + exc, + ) + return + raise PluginLoadError( + f"Failed to load plug-in from PLUGIN_WHEEL='{wheel_env}': {exc}. " + "Remove the plug-in configuration, or enable fail-open (PLUGIN_FAIL_OPEN=true, or the " + "initialize_pyrit_async(plugin_fail_open=True) parameter) to continue without it." + ) from exc + + async def _load_plugin_async(self, *, wheel_path: Path) -> None: + """ + Extract, import, bootstrap, and verify a single plug-in wheel. + + Global state (``sys.path``, imported plug-in modules, and the provider/scenario + registries) is rolled back if the load fails, so a failed or fail-open load + leaves no partial trace. + + Args: + wheel_path: Path to the pre-built plug-in wheel on disk. + + Raises: + FileNotFoundError: If ``wheel_path`` does not point to an existing file. + ValueError: If the file is not a ``.whl`` or the plug-in registered nothing. + """ + if not wheel_path.is_file(): + raise FileNotFoundError(f"PLUGIN_WHEEL does not point to an existing file: {wheel_path}") + if wheel_path.suffix != ".whl": + raise ValueError(f"PLUGIN_WHEEL must point to a .whl file, got: {wheel_path}") + + extract_dir = self._extract_wheel(wheel_path=wheel_path) + package_name = self._resolve_package_name(extract_dir=extract_dir) + + from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider + from pyrit.registry import ScenarioRegistry + + scenario_registry = ScenarioRegistry.get_registry_singleton() + provider_snapshot = dict(SeedDatasetProvider._registry) + scenario_snapshot = dict(scenario_registry._classes) + modules_snapshot = {name for name in sys.modules if _name_owned_by(name, package_name)} + syspath_entry = str(extract_dir) + added_to_syspath = syspath_entry not in sys.path + if added_to_syspath: + sys.path.insert(0, syspath_entry) + + try: + logger.info("Importing plug-in package '%s'", package_name) + module = importlib.import_module(package_name) + self._verify_module_location(module=module, extract_dir=extract_dir, package_name=package_name) + self._import_submodules(module=module, package_name=package_name) + + await self._run_bootstrap_async(package_name=package_name, module=module) + + provider_count, scenario_count = self._count_registered( + package_name=package_name, scenario_registry=scenario_registry + ) + if not provider_count and not scenario_count: + raise ValueError( + f"Plug-in package '{package_name}' imported successfully but registered no datasets or " + "scenarios. The wheel is likely mis-packaged (imports cleanly yet loads nothing)." + ) + + dataset_collisions = self._warn_on_dataset_name_collisions(package_name=package_name) + except Exception: + self._rollback( + package_name=package_name, + syspath_entry=syspath_entry if added_to_syspath else None, + modules_snapshot=modules_snapshot, + provider_snapshot=provider_snapshot, + scenario_registry=scenario_registry, + scenario_snapshot=scenario_snapshot, + ) + raise + + logger.info( + "Loaded plug-in '%s': %d dataset provider(s), %d scenario(s) registered; %d dataset name collision(s)%s.", + package_name, + provider_count, + scenario_count, + len(dataset_collisions), + " (see PLUGIN DATASET SHADOWED warnings above)" if dataset_collisions else "", + ) + + def _extract_wheel(self, *, wheel_path: Path) -> Path: + """ + Extract the wheel into ``.plugin//``, reusing a cached extraction. + + Extraction is atomic: the wheel is unpacked into a temporary sibling directory and + moved into place only on success, so a crash mid-extraction never leaves a partial + tree that would later be treated as a valid cache. + + Args: + wheel_path: Path to the plug-in wheel. + + Returns: + Path: The directory the wheel was extracted to. + """ + base_dir = self._plugin_base_dir() + base_dir.mkdir(parents=True, exist_ok=True) + + extract_dir = base_dir / wheel_path.stem + if extract_dir.is_dir() and any(extract_dir.iterdir()): + logger.info("Reusing cached plug-in extraction at %s", extract_dir) + return extract_dir + + tmp_dir = base_dir / f".{wheel_path.stem}.tmp-{os.getpid()}" + if tmp_dir.exists(): + shutil.rmtree(tmp_dir) + tmp_dir.mkdir(parents=True) + try: + with zipfile.ZipFile(wheel_path) as wheel_zip: + wheel_zip.extractall(tmp_dir) + if extract_dir.exists(): + shutil.rmtree(extract_dir) + os.replace(tmp_dir, extract_dir) + finally: + if tmp_dir.exists(): + shutil.rmtree(tmp_dir, ignore_errors=True) + + logger.info("Extracted plug-in wheel '%s' to %s", wheel_path.name, extract_dir) + return extract_dir + + @staticmethod + def _plugin_base_dir() -> Path: + """ + Resolve the base directory for plug-in extractions. + + Uses ``PLUGIN_DIR`` when set, otherwise ``/.plugin``. + + Returns: + Path: The resolved plug-in base directory. + """ + override = os.getenv("PLUGIN_DIR") + if override: + return Path(override).expanduser().resolve() + from pyrit.common import path + + return Path(path.HOME_PATH, ".plugin").resolve() + + @staticmethod + def _resolve_package_name(*, extract_dir: Path) -> str: + """ + Determine the plug-in's top-level import package. + + Resolution order: ``PLUGIN_PACKAGE`` env var, then ``*.dist-info/top_level.txt``, + then the single importable top-level directory in the extraction. + + Args: + extract_dir: The directory the wheel was extracted to. + + Returns: + str: The top-level package name to import. + + Raises: + ValueError: If the package cannot be unambiguously determined. + """ + explicit = os.getenv("PLUGIN_PACKAGE") + if explicit: + return explicit + + for dist_info in sorted(extract_dir.glob("*.dist-info")): + top_level = dist_info / "top_level.txt" + if top_level.is_file(): + for line in top_level.read_text(encoding="utf-8").splitlines(): + name = line.strip() + if name: + return name + + candidates = sorted( + child.name + for child in extract_dir.iterdir() + if child.is_dir() + and not child.name.endswith(".dist-info") + and not child.name.endswith(".data") + and (child / "__init__.py").is_file() + ) + if len(candidates) == 1: + return candidates[0] + if not candidates: + raise ValueError( + f"Could not find an importable top-level package in {extract_dir}. " + "Set PLUGIN_PACKAGE to the plug-in's package name." + ) + raise ValueError( + f"Found multiple top-level packages in {extract_dir}: {candidates}. Set PLUGIN_PACKAGE to disambiguate." + ) + + @staticmethod + def _verify_module_location(*, module: ModuleType, extract_dir: Path, package_name: str) -> None: + """ + Verify the imported package resolves inside the extraction directory. + + Guards against an installed package of the same name shadowing the extracted + plug-in — a silent failure where import succeeds but the wheel's code/data is + ignored. + + Args: + module: The imported plug-in package module. + extract_dir: The directory the wheel was extracted to. + package_name: The plug-in's top-level package name. + + Raises: + ValueError: If the imported package resolves outside ``extract_dir``. + """ + extract_resolved = extract_dir.resolve() + raw_locations = list(getattr(module, "__path__", []) or []) + module_file = getattr(module, "__file__", None) + if module_file: + raw_locations.append(module_file) + + locations = [Path(location).resolve() for location in raw_locations if location] + if not locations: + return + if not any(location.is_relative_to(extract_resolved) for location in locations): + raise ValueError( + f"Imported package '{package_name}' resolved to {locations[0]} which is outside the " + f"plug-in extraction directory {extract_resolved}. An installed package with the same " + "name is likely shadowing the plug-in; set PLUGIN_PACKAGE or resolve the name conflict." + ) + + @staticmethod + def _import_submodules(*, module: ModuleType, package_name: str) -> None: + """ + Import every submodule of the plug-in package. + + Ensures dataset providers self-register and bootstrap initializers become + discoverable even when the package ``__init__`` does not import them. Import + errors surface (plug-in dependencies must be pre-satisfied — fail loud). + + Args: + module: The imported plug-in package module. + package_name: The plug-in's top-level package name. + """ + module_path = getattr(module, "__path__", None) + if not module_path: + return # Single-module plug-in (not a package); nothing to walk. + + def _raise_on_error(name: str) -> None: + raise ImportError(f"Failed to import plug-in submodule '{name}'") + + for submodule in pkgutil.walk_packages(module_path, prefix=f"{package_name}.", onerror=_raise_on_error): + importlib.import_module(submodule.name) + + async def _run_bootstrap_async(self, *, package_name: str, module: ModuleType) -> None: + """ + Run the plug-in's bootstrap so its scenarios register. + + Prefers a top-level ``register()`` callable on the package, then any + ``PyRITInitializer`` subclass defined within the package. If neither exists the + plug-in is assumed to register everything on import (datasets-only plug-ins). + + Args: + package_name: The plug-in's top-level package name. + module: The imported plug-in package module. + """ + register = getattr(module, "register", None) + if callable(register): + logger.info("Running plug-in bootstrap register() from '%s'", package_name) + result = register() + if inspect.isawaitable(result): + await result + return + + initializer_classes = self._find_plugin_initializers(package_name=package_name) + if initializer_classes: + for initializer_class in initializer_classes: + logger.info("Running plug-in bootstrap initializer %s", initializer_class.__name__) + await initializer_class().initialize_async() + return + + logger.info( + "Plug-in '%s' exposes no register() or PyRITInitializer bootstrap; relying on " + "import-time registration only.", + package_name, + ) + + @staticmethod + def _find_plugin_initializers(*, package_name: str) -> list[type[PyRITInitializer]]: + """ + Find concrete ``PyRITInitializer`` subclasses defined within the plug-in package. + + Args: + package_name: The plug-in's top-level package name. + + Returns: + list[type[PyRITInitializer]]: Bootstrap initializer classes owned by the plug-in. + """ + prefix = f"{package_name}." + found: list[type[PyRITInitializer]] = [] + seen: set[type[PyRITInitializer]] = set() + + stack: list[type[PyRITInitializer]] = list(PyRITInitializer.__subclasses__()) + while stack: + cls = stack.pop() + if cls in seen: + continue + seen.add(cls) + stack.extend(cls.__subclasses__()) + + module_name = cls.__module__ or "" + if inspect.isabstract(cls): + continue + if module_name == package_name or module_name.startswith(prefix): + found.append(cls) + return found + + @staticmethod + def _count_registered(*, package_name: str, scenario_registry: ScenarioRegistry) -> tuple[int, int]: + """ + Count providers and scenarios registered by the plug-in package. + + Both are counted by matching each registered class's module against the plug-in + package, so the check is precise to this plug-in and safe to re-run. + + Args: + package_name: The plug-in's top-level package name. + scenario_registry: The scenario registry singleton the bootstrap registered into. + + Returns: + tuple[int, int]: (dataset provider count, scenario count) owned by the plug-in. + """ + from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider + + provider_count = sum( + 1 for cls in SeedDatasetProvider.get_all_providers().values() if _module_owned_by(cls, package_name) + ) + + # Read the raw class catalog directly: this snapshot must not trigger built-in + # discovery, and the plug-in's register_class writes straight into it. + scenario_count = sum(1 for cls in scenario_registry._classes.values() if _module_owned_by(cls, package_name)) + + return provider_count, scenario_count + + @staticmethod + def _warn_on_dataset_name_collisions(*, package_name: str) -> list[str]: + """ + Warn loudly when a plug-in dataset name collides with an existing dataset name. + + The dataset resolver treats central memory as authoritative and only consults a + provider when memory has no seeds for that ``dataset_name``. Once a same-named + dataset is in memory, a scan uses it and never consults the plug-in's provider, so + the plug-in's copy is silently bypassed. Any collision with **another** registered + provider's ``dataset_name`` is surfaced prominently at load time so the mismatch is + never silent. + + This compares the **provider registry**, not live memory, on purpose. At this phase + memory is not populated yet, and a live-memory check would false-positive on the + plug-in's own datasets persisted from a prior run (the seed rows carry no trustworthy + source, so "already in memory" cannot be told apart from "this plug-in loaded it last + run" — it is fundamentally undecidable at load time). The registry check is a + **conservative proxy** for the shadowing that ``LoadDefaultDatasets`` will cause by + loading provider datasets into memory: if the operator's config does not run + ``load_default_datasets`` (or loads only a tag subset), a built-in name may not actually + land in memory and this warning can fire without real shadowing. Over-warning is the + safe direction — do NOT "fix" this into a memory check (it reintroduces the false + positives). Governing principle: a guard's value is its precision — a check that cries + wolf on legitimate plug-in data every run desensitizes operators and defeats itself for + the real collision, so false-positive-free with a documented gap beats high-recall-but- + noisy. Hard enforcement that can tell a real mismatch from a harmless name coincidence + belongs to the scenario's declared required-dataset-names / expected-source mechanism + (which knows the operator's intent), not this loader, and is intentionally not gated + behind ``PLUGIN_FAIL_OPEN``. + + Args: + package_name: The plug-in's top-level package name. + + Returns: + list[str]: The sorted colliding dataset names (empty when there are none). + """ + from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider + + def _safe_name(provider_class: type[SeedDatasetProvider]) -> str | None: + try: + return provider_class().dataset_name + except Exception: + return None + + providers = SeedDatasetProvider.get_all_providers() + + # Map dataset_name -> owning provider class name(s), split into the plug-in's own + # providers vs. everything else. "Other" deliberately EXCLUDES the plug-in's own + # providers, so a plug-in shipping multiple datasets (or a re-run) never self-flags. + plugin_owned: dict[str, str] = {} + other_owned: dict[str, list[str]] = {} + for class_name, provider_class in providers.items(): + name = _safe_name(provider_class) + if name is None: + continue + if _module_owned_by(provider_class, package_name): + plugin_owned.setdefault(name, class_name) + else: + other_owned.setdefault(name, []).append(class_name) + + collisions = sorted(plugin_owned.keys() & other_owned.keys()) + for name in collisions: + logger.warning( + "PLUGIN DATASET SHADOWED: plug-in '%s' provider %s registers dataset_name '%s', which is " + "already provided by %s. Central memory is authoritative, so a scan will use the existing " + "dataset and the plug-in's copy will NOT take effect. Rename the plug-in dataset to a unique name.", + package_name, + plugin_owned[name], + name, + ", ".join(sorted(other_owned[name])), + ) + return collisions + + @staticmethod + def _rollback( + *, + package_name: str, + syspath_entry: str | None, + modules_snapshot: set[str], + provider_snapshot: Mapping[str, type], + scenario_registry: ScenarioRegistry, + scenario_snapshot: Mapping[str, type], + ) -> None: + """ + Undo the partial global-state changes made while loading a plug-in. + + Removes the plug-in's ``sys.path`` entry, the modules it newly imported, and any + provider/scenario registrations it added, so a failed (or fail-open) load leaves + PyRIT as if the plug-in had never been loaded. State present before the load — + including modules that already existed and built-ins discovered meanwhile — is + preserved. + + Args: + package_name: The plug-in's top-level package name. + syspath_entry: The ``sys.path`` entry to remove, or None if it was already present. + modules_snapshot: Package-owned module names present before the load. + provider_snapshot: Provider registry contents captured before the load. + scenario_registry: The scenario registry singleton to clean up. + scenario_snapshot: Scenario catalog contents captured before the load. + """ + from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider + + if syspath_entry and syspath_entry in sys.path: + sys.path.remove(syspath_entry) + + for name in [m for m in sys.modules if _name_owned_by(m, package_name) and m not in modules_snapshot]: + del sys.modules[name] + + for key, cls in list(SeedDatasetProvider._registry.items()): + if key not in provider_snapshot and _module_owned_by(cls, package_name): + del SeedDatasetProvider._registry[key] + + removed_scenario = False + for name, cls in list(scenario_registry._classes.items()): + if name not in scenario_snapshot and _module_owned_by(cls, package_name): + del scenario_registry._classes[name] + removed_scenario = True + if removed_scenario: + scenario_registry._metadata_cache = None + + def _resolve_fail_open(self) -> bool: + """ + Resolve the fail-open setting from the explicit value or the environment. + + Precedence: the explicit ``fail_open`` passed to the constructor (e.g. from + ``initialize_pyrit_async``), then the ``PLUGIN_FAIL_OPEN`` environment variable, + otherwise fail-closed. + + Returns: + bool: True if a failed plug-in load should be skipped with a warning. + """ + if self._explicit_fail_open is not None: + return self._explicit_fail_open + + env_value = os.getenv("PLUGIN_FAIL_OPEN") + if env_value is not None: + return self._coerce_bool(env_value) + + return False + + @staticmethod + def _coerce_bool(value: str) -> bool: + """ + Interpret a string as a boolean flag. + + Args: + value: The raw string value. + + Returns: + bool: True for common truthy tokens (1/true/yes/on, case-insensitive). + """ + return str(value).strip().lower() in _TRUE_TOKENS diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py new file mode 100644 index 0000000000..7bf2951c60 --- /dev/null +++ b/tests/unit/setup/test_plugin_loader.py @@ -0,0 +1,772 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Unit tests for the PyRIT plug-in loader. + +These tests build a **mock plug-in wheel** at test time (no dependency on any real +plug-in) and exercise the full consumer mechanism: extract -> sys.path -> +import -> bootstrap -> assert-loaded, plus the fail-open/closed policy and the +silent-failure guards called out in the design brief. +""" + +import logging +import os +import sys +import textwrap +import uuid +import zipfile +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider +from pyrit.memory import CentralMemory +from pyrit.models import SeedDataset +from pyrit.registry import ScenarioRegistry +from pyrit.setup.initialization import IN_MEMORY, initialize_pyrit_async +from pyrit.setup.plugin_loader import PluginLoader, PluginLoadError, load_plugin_if_configured_async + +# --------------------------------------------------------------------------- +# Mock-wheel builder +# --------------------------------------------------------------------------- + + +class MockWheel: + """Handle describing a built mock plug-in wheel.""" + + def __init__(self, *, path: Path, package: str, scenario_name: str, dataset_name: str) -> None: + self.path = path + self.package = package + self.scenario_name = scenario_name + self.dataset_name = dataset_name + + +def _unique_package_name() -> str: + """Return a unique, import-safe mock package name.""" + return f"mock_plugin_{uuid.uuid4().hex[:8]}" + + +def build_mock_wheel( + dest_dir: Path, + *, + bootstrap: str = "initializer", + include_provider: bool = True, + include_scenario: bool = True, + wire_init: bool = True, + package_name: str | None = None, +) -> MockWheel: + """ + Build a mock plug-in wheel in ``dest_dir`` and return a handle to it. + + Args: + dest_dir: Directory to write the wheel source tree and .whl into. + bootstrap: Bootstrap style: "initializer" (a PyRITInitializer subclass), + "register" (a top-level register() callable), or "none". + include_provider: Whether to ship a self-registering SeedDatasetProvider. + include_scenario: Whether to ship a Scenario subclass. + wire_init: Whether __init__.py imports the submodules. When False, submodules are + shipped but not imported by __init__, exercising the loader's submodule walk. + package_name: Optional explicit package name; a unique one is generated otherwise. + + Returns: + MockWheel: The built wheel handle (path + package/scenario/dataset names). + """ + package_name = package_name or _unique_package_name() + scenario_name = f"airt.{package_name}" + dataset_name = f"{package_name}_dataset" + + src = dest_dir / f"{package_name}_src" + pkg = src / package_name + pkg.mkdir(parents=True, exist_ok=True) + + imports = [] + if include_provider: + imports.append("provider") + if include_scenario: + imports.append("scenario") + if bootstrap in ("initializer", "initializer_raises", "register"): + imports.append("bootstrap") + + init_lines = [] + if wire_init and imports: + init_lines.append(f"from . import {', '.join(imports)} # noqa: F401") + if wire_init and bootstrap == "register": + init_lines.append("from .bootstrap import register # noqa: F401") + (pkg / "__init__.py").write_text(("\n".join(init_lines) + "\n") if init_lines else "", encoding="utf-8") + + # __file__-relative dataset path (as a real plug-in ships). Only resolves on real disk. + (pkg / "paths.py").write_text( + textwrap.dedent( + """\ + from pathlib import Path + + MOCK_ROOT = Path(__file__, "..").resolve() + MOCK_DATASETS_PATH = Path(MOCK_ROOT, "datasets").resolve() + """ + ), + encoding="utf-8", + ) + + if include_provider: + datasets = pkg / "datasets" + datasets.mkdir(parents=True, exist_ok=True) + (pkg / "provider.py").write_text( + textwrap.dedent( + f"""\ + from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider + from pyrit.models.seeds.seed_dataset import SeedDataset + + from .paths import MOCK_DATASETS_PATH + + + class MockProvider(SeedDatasetProvider): + @property + def dataset_name(self) -> str: + return "{dataset_name}" + + async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: + return SeedDataset.from_yaml_file(MOCK_DATASETS_PATH / "seed.yaml") + """ + ), + encoding="utf-8", + ) + (datasets / "seed.yaml").write_text( + textwrap.dedent( + f"""\ + dataset_name: {dataset_name} + harm_categories: + - mock + data_type: text + description: mock dataset for plugin test + authors: + - tester + groups: + - test + seeds: + - value: mock prompt one + - value: mock prompt two + - value: mock prompt three + """ + ), + encoding="utf-8", + ) + + if include_scenario: + (pkg / "scenario.py").write_text( + textwrap.dedent( + """\ + from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + + + class MockScenario(RapidResponse): + \"\"\"Mock plugin scenario for registration test.\"\"\" + """ + ), + encoding="utf-8", + ) + + if bootstrap == "initializer": + (pkg / "bootstrap.py").write_text( + textwrap.dedent( + f"""\ + from pyrit.registry import ScenarioRegistry + from pyrit.setup.pyrit_initializer import PyRITInitializer + + from .scenario import MockScenario + + + class MockBootstrapInitializer(PyRITInitializer): + \"\"\"Register the mock plugin scenario.\"\"\" + + async def initialize_async(self) -> None: + ScenarioRegistry.get_registry_singleton().register_class( + MockScenario, name="{scenario_name}" + ) + """ + ), + encoding="utf-8", + ) + elif bootstrap == "initializer_raises": + (pkg / "bootstrap.py").write_text( + textwrap.dedent( + f"""\ + from pyrit.registry import ScenarioRegistry + from pyrit.setup.pyrit_initializer import PyRITInitializer + + from .scenario import MockScenario + + + class MockBootstrapInitializer(PyRITInitializer): + \"\"\"Register the scenario, then fail to test rollback.\"\"\" + + async def initialize_async(self) -> None: + ScenarioRegistry.get_registry_singleton().register_class( + MockScenario, name="{scenario_name}" + ) + raise RuntimeError("bootstrap failed after registering") + """ + ), + encoding="utf-8", + ) + elif bootstrap == "register": + (pkg / "bootstrap.py").write_text( + textwrap.dedent( + f"""\ + from pyrit.registry import ScenarioRegistry + + from .scenario import MockScenario + + + def register() -> None: + ScenarioRegistry.get_registry_singleton().register_class( + MockScenario, name="{scenario_name}" + ) + """ + ), + encoding="utf-8", + ) + + # Minimal dist-info without top_level.txt so package name inference is exercised. + distinfo = src / f"{package_name}-0.0.1.dist-info" + distinfo.mkdir(parents=True, exist_ok=True) + (distinfo / "METADATA").write_text( + f"Metadata-Version: 2.1\nName: {package_name}\nVersion: 0.0.1\n", encoding="utf-8" + ) + (distinfo / "WHEEL").write_text( + "Wheel-Version: 1.0\nGenerator: test\nRoot-Is-Purelib: true\nTag: py3-none-any\n", encoding="utf-8" + ) + (distinfo / "RECORD").write_text("", encoding="utf-8") + + wheel = dest_dir / f"{package_name}-0.0.1-py3-none-any.whl" + with zipfile.ZipFile(wheel, "w", zipfile.ZIP_DEFLATED) as archive: + for root, _, files in os.walk(src): + for file_name in files: + file_path = Path(root) / file_name + archive.write(file_path, str(file_path.relative_to(src))) + + return MockWheel(path=wheel, package=package_name, scenario_name=scenario_name, dataset_name=dataset_name) + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def plugin_sandbox() -> Iterator[None]: + """Snapshot and restore global import + registry state around each test.""" + sys_path_snapshot = list(sys.path) + provider_snapshot = dict(SeedDatasetProvider._registry) + + yield + + sys.path[:] = sys_path_snapshot + # Only drop the mock plug-in modules this suite imports; leave real pyrit modules + # in place so re-imports don't create duplicate class objects for other tests. + for name in [m for m in sys.modules if m == "mock_plugin" or m.startswith("mock_plugin_")]: + del sys.modules[name] + SeedDatasetProvider._registry.clear() + SeedDatasetProvider._registry.update(provider_snapshot) + ScenarioRegistry.reset_registry_singleton() + + +@contextmanager +def plugin_env(**overrides: str) -> Iterator[None]: + """Patch os.environ so only the given PLUGIN_* overrides are present.""" + with patch.dict(os.environ, overrides, clear=False): + for key in ("PLUGIN_WHEEL", "PLUGIN_DIR", "PLUGIN_PACKAGE", "PLUGIN_FAIL_OPEN"): + if key not in overrides: + os.environ.pop(key, None) + yield + + +async def load_plugin( + wheel: MockWheel, + plugin_dir: Path, + *, + fail_open: bool | None = None, + extra_env: dict[str, str] | None = None, +) -> None: + """Run the plug-in loader against a mock wheel with an isolated env.""" + env = {"PLUGIN_WHEEL": str(wheel.path), "PLUGIN_DIR": str(plugin_dir)} + if extra_env: + env.update(extra_env) + + with plugin_env(**env): + await load_plugin_if_configured_async(fail_open=fail_open) + + +# --------------------------------------------------------------------------- +# Loader phase inside initialize_pyrit_async +# --------------------------------------------------------------------------- + + +async def test_plugin_phase_runs_after_memory_before_initializers() -> None: + """initialize_pyrit_async loads the plug-in after memory is set, before initializers.""" + manager = MagicMock() + manager.attach_mock(MagicMock(), "set_memory") + manager.attach_mock(AsyncMock(), "load_plugin") + manager.attach_mock(AsyncMock(), "execute") + + with ( + patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), + patch.object(CentralMemory, "set_memory_instance", manager.set_memory), + patch("pyrit.setup.plugin_loader.load_plugin_if_configured_async", manager.load_plugin), + patch("pyrit.setup.initialization._execute_initializers_async", manager.execute), + ): + await initialize_pyrit_async(IN_MEMORY, initializers=[MagicMock()], env_files=[], silent=True) + + order = [call[0] for call in manager.mock_calls if call[0] in {"set_memory", "load_plugin", "execute"}] + assert order.index("set_memory") < order.index("load_plugin") < order.index("execute") + + +async def test_plugin_phase_forwards_fail_open_param() -> None: + """initialize_pyrit_async forwards plugin_fail_open to the loader.""" + load_plugin_mock = AsyncMock() + with ( + patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), + patch.object(CentralMemory, "set_memory_instance"), + patch("pyrit.setup.plugin_loader.load_plugin_if_configured_async", load_plugin_mock), + ): + await initialize_pyrit_async(IN_MEMORY, env_files=[], silent=True, plugin_fail_open=True) + + load_plugin_mock.assert_awaited_once_with(fail_open=True) + + +# --------------------------------------------------------------------------- +# No-op behavior +# --------------------------------------------------------------------------- + + +async def test_no_op_when_plugin_wheel_unset() -> None: + """With no PLUGIN_WHEEL the loader does nothing and registers nothing.""" + providers_before = dict(SeedDatasetProvider.get_all_providers()) + path_before = list(sys.path) + + with plugin_env(): + await load_plugin_if_configured_async() + + assert SeedDatasetProvider.get_all_providers() == providers_before + assert sys.path == path_before + + +# --------------------------------------------------------------------------- +# Silent-failure trap: extraction, not zipimport +# --------------------------------------------------------------------------- + + +def test_raw_wheel_on_syspath_loses_datasets(tmp_path: Path) -> None: + """A raw .whl on sys.path imports but __file__-relative datasets vanish (regression guard).""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_scenario=False) + + sys.path.insert(0, str(wheel.path)) + module = __import__(wheel.package) + paths_module = __import__(f"{wheel.package}.paths", fromlist=["MOCK_DATASETS_PATH"]) + + assert ".whl" in (module.__file__ or "") + assert not paths_module.MOCK_DATASETS_PATH.exists() + assert list(paths_module.MOCK_DATASETS_PATH.glob("**/*.yaml")) == [] + + +def test_extracted_wheel_loads_datasets(tmp_path: Path) -> None: + """Extracting the wheel to disk makes __file__-relative datasets resolve and load.""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_scenario=False) + extract_dir = tmp_path / "extracted" + extract_dir.mkdir() + with zipfile.ZipFile(wheel.path) as archive: + archive.extractall(extract_dir) + + sys.path.insert(0, str(extract_dir)) + paths_module = __import__(f"{wheel.package}.paths", fromlist=["MOCK_DATASETS_PATH"]) + + yamls = list(paths_module.MOCK_DATASETS_PATH.glob("**/*.yaml")) + assert len(yamls) == 1 + + dataset = SeedDataset.from_yaml_file(yamls[0]) + assert len(dataset.seeds) == 3 + + +# --------------------------------------------------------------------------- +# Loading via the initializer +# --------------------------------------------------------------------------- + + +async def test_load_registers_provider_on_import(tmp_path: Path) -> None: + """Importing the plug-in package self-registers its SeedDatasetProvider.""" + wheel = build_mock_wheel(tmp_path) + + await load_plugin(wheel, tmp_path / ".plugin") + + assert "MockProvider" in SeedDatasetProvider.get_all_providers() + + +async def test_load_extracts_to_plugin_dir(tmp_path: Path) -> None: + """The wheel is extracted (not installed) under the configured plug-in dir.""" + wheel = build_mock_wheel(tmp_path) + plugin_dir = tmp_path / ".plugin" + + await load_plugin(wheel, plugin_dir) + + extract_dir = plugin_dir / wheel.path.stem + assert (extract_dir / wheel.package / "__init__.py").is_file() + assert (extract_dir / wheel.package / "datasets" / "seed.yaml").is_file() + + +async def test_scenario_registration_survives_discovery(tmp_path: Path) -> None: + """A plug-in scenario registered before discovery coexists with built-ins afterwards.""" + wheel = build_mock_wheel(tmp_path, bootstrap="initializer") + + await load_plugin(wheel, tmp_path / ".plugin") + + registry = ScenarioRegistry.get_registry_singleton() + assert registry._discovered is False # register_class must not trigger discovery + + names = registry.get_class_names() # triggers built-in discovery + assert wheel.scenario_name in names + assert "airt.rapid_response" in names + + mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario + assert registry.get_class(wheel.scenario_name) is mock_scenario + + +async def test_register_callable_bootstrap(tmp_path: Path) -> None: + """A plug-in exposing a top-level register() callable is bootstrapped too.""" + wheel = build_mock_wheel(tmp_path, bootstrap="register") + + await load_plugin(wheel, tmp_path / ".plugin") + + names = ScenarioRegistry.get_registry_singleton().get_class_names() + assert wheel.scenario_name in names + + +async def test_ordering_scenario_visible_to_preload(tmp_path: Path) -> None: + """The plug-in scenario is registered before a later PreloadScenarioMetadata read.""" + wheel = build_mock_wheel(tmp_path, bootstrap="initializer") + + await load_plugin(wheel, tmp_path / ".plugin") + + # get_class_names() is exactly what PreloadScenarioMetadata iterates; the plug-in + # scenario being present proves it registered before that read would happen. + names = ScenarioRegistry.get_registry_singleton().get_class_names() + assert wheel.scenario_name in names + + +async def test_datasets_only_plugin_loads_without_bootstrap(tmp_path: Path) -> None: + """A datasets-only plug-in (no bootstrap, no scenario) loads via import-time registration.""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_scenario=False) + + await load_plugin(wheel, tmp_path / ".plugin") + + assert "MockProvider" in SeedDatasetProvider.get_all_providers() + + +async def test_submodule_walk_discovers_unwired_components(tmp_path: Path) -> None: + """Provider + bootstrap register even when __init__.py does not import them.""" + wheel = build_mock_wheel(tmp_path, bootstrap="initializer", wire_init=False) + + await load_plugin(wheel, tmp_path / ".plugin") + + assert "MockProvider" in SeedDatasetProvider.get_all_providers() + assert wheel.scenario_name in ScenarioRegistry.get_registry_singleton().get_class_names() + + +async def test_shadowing_installed_package_is_rejected(tmp_path: Path) -> None: + """An installed package of the same name shadowing the plug-in fails loudly.""" + wheel = build_mock_wheel(tmp_path) + + # PLUGIN_PACKAGE points at a stdlib package that imports from outside the extraction dir. + with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(tmp_path / ".plugin"), PLUGIN_PACKAGE="json"): + with pytest.raises(PluginLoadError, match="shadowing"): + await load_plugin_if_configured_async() + + +# --------------------------------------------------------------------------- +# Dataset name collision (memory-authoritative resolver guard) +# --------------------------------------------------------------------------- + + +async def test_colliding_dataset_name_warns_loudly(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + """A plug-in dataset_name that collides with an existing provider's name warns at load.""" + wheel = build_mock_wheel(tmp_path) + colliding_name = wheel.dataset_name + + class CollidingProvider(SeedDatasetProvider): + """Non-plug-in provider that already claims the plug-in's dataset name.""" + + @property + def dataset_name(self) -> str: + return colliding_name + + async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: + raise NotImplementedError + + with caplog.at_level(logging.WARNING, logger="pyrit.setup.plugin_loader"): + await load_plugin(wheel, tmp_path / ".plugin") + + messages = [record.getMessage() for record in caplog.records] + # Un-missable: greppable prefix, names the colliding dataset and BOTH providers. + assert any( + "PLUGIN DATASET SHADOWED:" in message + and colliding_name in message + and "MockProvider" in message + and "CollidingProvider" in message + for message in messages + ) + + +async def test_unique_dataset_name_does_not_warn(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + """A plug-in whose dataset name is unique produces no collision warning.""" + wheel = build_mock_wheel(tmp_path) + + with caplog.at_level(logging.WARNING, logger="pyrit.setup.plugin_loader"): + await load_plugin(wheel, tmp_path / ".plugin") + + assert not any("PLUGIN DATASET SHADOWED:" in record.getMessage() for record in caplog.records) + + +async def test_reload_does_not_self_flag(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + """Loading the same plug-in twice must not flag its own provider as a collision.""" + wheel = build_mock_wheel(tmp_path) + plugin_dir = tmp_path / ".plugin" + + await load_plugin(wheel, plugin_dir) + with caplog.at_level(logging.WARNING, logger="pyrit.setup.plugin_loader"): + await load_plugin(wheel, plugin_dir) + + assert not any("PLUGIN DATASET SHADOWED:" in record.getMessage() for record in caplog.records) + + +# --------------------------------------------------------------------------- +# Rollback on failure +# --------------------------------------------------------------------------- + + +async def test_failed_load_rolls_back_syspath(tmp_path: Path) -> None: + """A failed load removes its own sys.path entry (fail-closed leaves no trace).""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) + plugin_dir = tmp_path / ".plugin" + + with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(plugin_dir)): + with pytest.raises(PluginLoadError): + await load_plugin_if_configured_async() + + extract_dir = str(plugin_dir / wheel.path.stem) + assert extract_dir not in sys.path + + +async def test_failing_bootstrap_rolls_back_partial_registration(tmp_path: Path) -> None: + """A bootstrap that registers then raises has its registration rolled back.""" + wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") + plugin_dir = tmp_path / ".plugin" + + with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(plugin_dir)): + with pytest.raises(PluginLoadError): + await load_plugin_if_configured_async() + + registry = ScenarioRegistry.get_registry_singleton() + assert wheel.scenario_name not in registry._classes + assert "MockProvider" not in SeedDatasetProvider.get_all_providers() + assert str(plugin_dir / wheel.path.stem) not in sys.path + + +async def test_fail_open_rolls_back_partial_registration(tmp_path: Path) -> None: + """Under fail_open, a partially-registered failed plug-in is still fully rolled back.""" + wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") + plugin_dir = tmp_path / ".plugin" + + await load_plugin(wheel, plugin_dir, fail_open=True) # must not raise + + registry = ScenarioRegistry.get_registry_singleton() + assert wheel.scenario_name not in registry._classes + + +# --------------------------------------------------------------------------- +# Extraction cache +# --------------------------------------------------------------------------- + + +def test_extract_wheel_reuses_cached_extraction(tmp_path: Path) -> None: + """A second extraction of an unchanged wheel reuses the cached directory.""" + wheel = build_mock_wheel(tmp_path) + plugin_dir = tmp_path / ".plugin" + + with plugin_env(PLUGIN_DIR=str(plugin_dir)): + initializer = PluginLoader() + first = initializer._extract_wheel(wheel_path=wheel.path) + marker = first / "cache_marker.txt" + marker.write_text("kept", encoding="utf-8") + + second = initializer._extract_wheel(wheel_path=wheel.path) + + assert first == second + assert marker.is_file() # not wiped -> cached, not re-extracted + + +# --------------------------------------------------------------------------- +# Package name resolution +# --------------------------------------------------------------------------- + + +def test_resolve_package_name_prefers_env(tmp_path: Path) -> None: + """PLUGIN_PACKAGE takes precedence over inference.""" + (tmp_path / "some_pkg").mkdir() + (tmp_path / "some_pkg" / "__init__.py").write_text("", encoding="utf-8") + + with plugin_env(PLUGIN_PACKAGE="explicit_pkg"): + assert PluginLoader._resolve_package_name(extract_dir=tmp_path) == "explicit_pkg" + + +def test_resolve_package_name_infers_single_package(tmp_path: Path) -> None: + """The single importable top-level directory is inferred when no env/top_level.txt exists.""" + (tmp_path / "the_pkg").mkdir() + (tmp_path / "the_pkg" / "__init__.py").write_text("", encoding="utf-8") + (tmp_path / "the_pkg-0.0.1.dist-info").mkdir() + + with plugin_env(): + assert PluginLoader._resolve_package_name(extract_dir=tmp_path) == "the_pkg" + + +def test_resolve_package_name_uses_top_level_txt(tmp_path: Path) -> None: + """top_level.txt is consulted before directory inference.""" + (tmp_path / "pkg_a").mkdir() + (tmp_path / "pkg_a" / "__init__.py").write_text("", encoding="utf-8") + (tmp_path / "pkg_b").mkdir() + (tmp_path / "pkg_b" / "__init__.py").write_text("", encoding="utf-8") + distinfo = tmp_path / "thing-0.0.1.dist-info" + distinfo.mkdir() + (distinfo / "top_level.txt").write_text("pkg_b\n", encoding="utf-8") + + with plugin_env(): + assert PluginLoader._resolve_package_name(extract_dir=tmp_path) == "pkg_b" + + +def test_resolve_package_name_none_raises(tmp_path: Path) -> None: + """No importable package raises a clear error pointing at PLUGIN_PACKAGE.""" + with plugin_env(), pytest.raises(ValueError, match="PLUGIN_PACKAGE"): + PluginLoader._resolve_package_name(extract_dir=tmp_path) + + +def test_resolve_package_name_multiple_raises(tmp_path: Path) -> None: + """Multiple top-level packages require PLUGIN_PACKAGE to disambiguate.""" + for name in ("pkg_a", "pkg_b"): + (tmp_path / name).mkdir() + (tmp_path / name / "__init__.py").write_text("", encoding="utf-8") + + with plugin_env(), pytest.raises(ValueError, match="disambiguate"): + PluginLoader._resolve_package_name(extract_dir=tmp_path) + + +# --------------------------------------------------------------------------- +# Failure modes +# --------------------------------------------------------------------------- + + +async def test_missing_wheel_fails_closed() -> None: + """A configured-but-missing wheel raises by default (fail-closed).""" + with plugin_env(PLUGIN_WHEEL=str(Path("does_not_exist.whl"))): + with pytest.raises(PluginLoadError, match="Failed to load plug-in"): + await load_plugin_if_configured_async() + + +async def test_missing_wheel_fail_open_param_proceeds() -> None: + """fail_open via the explicit param skips a broken plug-in with a warning.""" + with plugin_env(PLUGIN_WHEEL=str(Path("does_not_exist.whl"))): + await load_plugin_if_configured_async(fail_open=True) # must not raise + + +async def test_missing_wheel_fail_open_env_proceeds() -> None: + """fail_open via PLUGIN_FAIL_OPEN env skips a broken plug-in with a warning.""" + with plugin_env(PLUGIN_WHEEL=str(Path("does_not_exist.whl")), PLUGIN_FAIL_OPEN="true"): + await load_plugin_if_configured_async() # must not raise + + +async def test_empty_wheel_is_loud(tmp_path: Path) -> None: + """A wheel that imports cleanly but registers nothing fails loudly.""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) + + with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(tmp_path / ".plugin")): + with pytest.raises(PluginLoadError, match="registered no datasets or scenarios"): + await load_plugin_if_configured_async() + + +async def test_empty_wheel_fail_open_proceeds(tmp_path: Path) -> None: + """An empty wheel under fail_open proceeds instead of raising.""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) + + await load_plugin(wheel, tmp_path / ".plugin", fail_open=True) # must not raise + + +async def test_non_whl_path_fails_closed(tmp_path: Path) -> None: + """PLUGIN_WHEEL that is not a .whl file fails closed.""" + not_a_wheel = tmp_path / "plugin.zip" + not_a_wheel.write_text("not a wheel", encoding="utf-8") + + with plugin_env(PLUGIN_WHEEL=str(not_a_wheel)): + with pytest.raises(PluginLoadError, match="Failed to load plug-in"): + await load_plugin_if_configured_async() + + +# --------------------------------------------------------------------------- +# No-arg-instantiable contract +# --------------------------------------------------------------------------- + + +def test_non_no_arg_scenario_fails_metadata_cleanly() -> None: + """A registered scenario that is not no-arg instantiable fails metadata build clearly.""" + from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + + class BadScenario(RapidResponse): + """Scenario that violates the no-arg-instantiable contract.""" + + def __init__(self, *, required_value: str) -> None: + super().__init__() + self._required_value = required_value + + registry = ScenarioRegistry() + registry.register_class(BadScenario, name="airt.bad") # signature-only validation passes + + with pytest.raises(TypeError, match="no arguments"): + registry._build_metadata("airt.bad", BadScenario) + + +# --------------------------------------------------------------------------- +# fail_open resolution +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "value,expected", + [("true", True), ("True", True), ("1", True), ("yes", True), ("on", True), ("false", False), ("0", False)], +) +def test_resolve_fail_open_from_env_tokens(value: str, expected: bool) -> None: + """fail_open resolves from PLUGIN_FAIL_OPEN across truthy/falsey tokens.""" + with plugin_env(PLUGIN_FAIL_OPEN=value): + assert PluginLoader()._resolve_fail_open() is expected + + +def test_resolve_fail_open_explicit_true() -> None: + """An explicit fail_open=True resolves to True.""" + with plugin_env(PLUGIN_FAIL_OPEN="false"): + assert PluginLoader(fail_open=True)._resolve_fail_open() is True + + +def test_resolve_fail_open_explicit_overrides_env() -> None: + """An explicit fail_open value takes precedence over the env var.""" + with plugin_env(PLUGIN_FAIL_OPEN="true"): + assert PluginLoader(fail_open=False)._resolve_fail_open() is False + + +def test_resolve_fail_open_from_env_when_no_explicit() -> None: + """fail_open falls back to PLUGIN_FAIL_OPEN when no explicit value is set.""" + with plugin_env(PLUGIN_FAIL_OPEN="true"): + assert PluginLoader()._resolve_fail_open() is True + + +def test_resolve_fail_open_defaults_false() -> None: + """fail_open defaults to False (fail-closed).""" + with plugin_env(): + assert PluginLoader()._resolve_fail_open() is False From eca59287ea07fbd46670762ea5bf2a07fe913b47 Mon Sep 17 00:00:00 2001 From: Victor Valbuena <50061128+ValbuenaVC@users.noreply.github.com> Date: Wed, 8 Jul 2026 10:13:07 -0700 Subject: [PATCH 02/34] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- pyrit/setup/plugin_loader.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index c632d1a242..3e5d9b3c07 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -248,6 +248,11 @@ def _extract_wheel(self, *, wheel_path: Path) -> Path: tmp_dir.mkdir(parents=True) try: with zipfile.ZipFile(wheel_path) as wheel_zip: + tmp_dir_resolved = tmp_dir.resolve() + for member in wheel_zip.infolist(): + member_path = (tmp_dir / member.filename).resolve() + if not member_path.is_relative_to(tmp_dir_resolved): + raise ValueError(f"Wheel contains unsafe path: {member.filename}") wheel_zip.extractall(tmp_dir) if extract_dir.exists(): shutil.rmtree(extract_dir) From e6159e518b5e47d7963ebe0f2beb312fc8affce4 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Jul 2026 11:13:14 -0700 Subject: [PATCH 03/34] test(setup): verify plug-in scenarios are discovered end-to-end Adds an integration test that loads a plug-in wheel through initialize_pyrit_async and asserts every scenario the wheel ships is registered in ScenarioRegistry. Two cases: a self-contained case builds a small wheel at test time and checks its scenarios are discovered (guards the mechanism in public CI); an injected case loads a wheel supplied out-of-band via PLUGIN_TEST_WHEEL (+ PLUGIN_TEST_SCENARIO_DIRS or PLUGIN_TEST_PACKAGE) and verifies all of its scenarios are picked up. The injected case is skipped unless those variables are set, so the committed test depends on no specific external package; a downstream CI job can point it at a real scenario wheel to guarantee full discovery. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- tests/integration/setup/__init__.py | 2 + .../setup/test_plugin_loader_integration.py | 282 ++++++++++++++++++ 2 files changed, 284 insertions(+) create mode 100644 tests/integration/setup/__init__.py create mode 100644 tests/integration/setup/test_plugin_loader_integration.py diff --git a/tests/integration/setup/__init__.py b/tests/integration/setup/__init__.py new file mode 100644 index 0000000000..9a0454564d --- /dev/null +++ b/tests/integration/setup/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. diff --git a/tests/integration/setup/test_plugin_loader_integration.py b/tests/integration/setup/test_plugin_loader_integration.py new file mode 100644 index 0000000000..d481bc1e88 --- /dev/null +++ b/tests/integration/setup/test_plugin_loader_integration.py @@ -0,0 +1,282 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Integration test: plug-in scenarios are discovered by ``ScenarioRegistry`` after init. + +This exercises the full public loading path end-to-end: point ``PLUGIN_WHEEL`` at a +built wheel, run ``initialize_pyrit_async``, and assert every scenario the wheel ships +is registered in ``ScenarioRegistry``. + +Two cases: + +* A self-contained case builds a small wheel at test time and verifies its scenarios + are discovered. This always runs and guards the mechanism in public CI. +* An injected case loads a wheel supplied out-of-band via environment variables and + verifies all of its scenarios are discovered. It is skipped unless those variables + are set, so the committed test depends on no specific external package. Point it at a + real scenario wheel (for example in a downstream/private CI job) to guarantee that + every scenario that wheel ships is picked up after initialization:: + + PLUGIN_TEST_WHEEL=/path/to/plugin.whl + PLUGIN_TEST_PACKAGE=the_plugin_package # optional; enables package enumeration + PLUGIN_TEST_SCENARIO_DIRS=/path/to/scenarios # optional; os.pathsep-separated +""" + +import inspect +import os +import sys +import textwrap +import uuid +import zipfile +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +import pytest + +from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider +from pyrit.registry import ScenarioRegistry +from pyrit.registry.discovery import discover_in_directory +from pyrit.scenario.core import Scenario +from pyrit.setup import IN_MEMORY, initialize_pyrit_async + +# (module stem, class name, registry name) for the scenarios the self-contained wheel ships. +_MOCK_SCENARIOS = [ + ("alpha", "MockAlphaScenario", "airt.mock_alpha"), + ("beta", "MockBetaScenario", "airt.mock_beta"), + ("gamma", "MockGammaScenario", "airt.mock_gamma"), +] + + +# REV Should this be a Pytest fixture? +def _build_scenario_plugin_wheel(dest_dir: Path, *, package: str) -> Path: + """ + Build a wheel whose package ships several scenarios and a ``register()`` bootstrap. + + Args: + dest_dir: Directory to write the source tree and .whl into. + package: Top-level package name for the wheel. + + Returns: + Path: The built wheel. + """ + src = dest_dir / f"{package}_src" + pkg = src / package + scenarios_pkg = pkg / "scenarios" + scenarios_pkg.mkdir(parents=True, exist_ok=True) + + (pkg / "__init__.py").write_text("from .bootstrap import register # noqa: F401\n", encoding="utf-8") + (scenarios_pkg / "__init__.py").write_text("", encoding="utf-8") + + for stem, class_name, _registry_name in _MOCK_SCENARIOS: + (scenarios_pkg / f"{stem}.py").write_text( + textwrap.dedent( + f"""\ + from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + + + class {class_name}(RapidResponse): + \"\"\"Mock plug-in scenario {stem}.\"\"\" + """ + ), + encoding="utf-8", + ) + + imports = "\n".join(f"from .scenarios.{stem} import {class_name}" for stem, class_name, _ in _MOCK_SCENARIOS) + registrations = "\n".join( + f' registry.register_class({class_name}, name="{registry_name}")' + for _stem, class_name, registry_name in _MOCK_SCENARIOS + ) + (pkg / "bootstrap.py").write_text( + textwrap.dedent( + """\ + from pyrit.registry import ScenarioRegistry + + {imports} + + + def register() -> None: + registry = ScenarioRegistry.get_registry_singleton() + {registrations} + """ + ).format(imports=imports, registrations=registrations), + encoding="utf-8", + ) + + distinfo = src / f"{package}-0.0.1.dist-info" + distinfo.mkdir(parents=True, exist_ok=True) + (distinfo / "METADATA").write_text(f"Metadata-Version: 2.1\nName: {package}\nVersion: 0.0.1\n", encoding="utf-8") + (distinfo / "WHEEL").write_text( + "Wheel-Version: 1.0\nGenerator: test\nRoot-Is-Purelib: true\nTag: py3-none-any\n", encoding="utf-8" + ) + (distinfo / "RECORD").write_text("", encoding="utf-8") + + wheel = dest_dir / f"{package}-0.0.1-py3-none-any.whl" + with zipfile.ZipFile(wheel, "w", zipfile.ZIP_DEFLATED) as archive: + for root, _, files in os.walk(src): + for file_name in files: + file_path = Path(root) / file_name + archive.write(file_path, str(file_path.relative_to(src))) + return wheel + + +def _registered_scenario_class_names() -> set[str]: + """Return the class names of every scenario currently registered in ScenarioRegistry.""" + registry = ScenarioRegistry.get_registry_singleton() + return {registry.get_class(name).__name__ for name in registry.get_class_names()} + + +def all_scenario_class_names(from_dirs: list[Path]) -> set[str]: + """ + Walk source directories and collect the class name of every ``Scenario`` subclass. + + This is the expected-set builder: given the scenario source directories to cover, it + enumerates the scenario classes defined there so the test can assert each is + discovered after initialization. + + Args: + from_dirs: Directories to walk recursively for ``Scenario`` subclasses. + + Returns: + set[str]: The scenario class names found across the directories. + """ + names: set[str] = set() + for directory in from_dirs: + for _stem, _path, cls in discover_in_directory(directory=directory, base_class=Scenario, recursive=True): + names.add(cls.__name__) + return names + + +def _scenario_class_names_under_package(package_prefix: str) -> set[str]: + """ + Return class names of concrete ``Scenario`` subclasses whose module is under a package. + + Uses in-memory subclass enumeration (reliable after the plug-in has been imported), + which sidesteps the standalone-import problems a filesystem walk can hit for a + package that uses relative imports. + + Args: + package_prefix: The plug-in's top-level package name. + + Returns: + set[str]: Scenario class names owned by that package. + """ + prefix = f"{package_prefix}." + found: set[str] = set() + seen: set[type] = set() + stack: list[type] = list(Scenario.__subclasses__()) + while stack: + cls = stack.pop() + if cls in seen: + continue + seen.add(cls) + stack.extend(cls.__subclasses__()) + module = cls.__module__ or "" + if not inspect.isabstract(cls) and (module == package_prefix or module.startswith(prefix)): + found.add(cls.__name__) + return found + + +@contextmanager +def _plugin_env(*, wheel: Path, plugin_dir: Path | None, extra: dict[str, str] | None = None) -> Iterator[None]: + """Set PLUGIN_* env for the duration of a load, then restore the prior values.""" + values: dict[str, str] = {"PLUGIN_WHEEL": str(wheel)} + if plugin_dir is not None: + values["PLUGIN_DIR"] = str(plugin_dir) + if extra: + values.update(extra) + + saved: dict[str, str | None] = {key: os.environ.get(key) for key in values} + os.environ.update(values) + try: + yield + finally: + for key, previous in saved.items(): + if previous is None: + os.environ.pop(key, None) + else: + os.environ[key] = previous + + +@pytest.fixture +def plugin_sandbox() -> Iterator[None]: + """Snapshot and restore global import + registry state so a load does not leak.""" + sys_path_snapshot = list(sys.path) + modules_snapshot = set(sys.modules) + provider_snapshot = dict(SeedDatasetProvider._registry) + + yield + + sys.path[:] = sys_path_snapshot + for name in set(sys.modules) - modules_snapshot: + del sys.modules[name] + SeedDatasetProvider._registry.clear() + SeedDatasetProvider._registry.update(provider_snapshot) + ScenarioRegistry.reset_registry_singleton() + + +@pytest.mark.run_only_if_all_tests +async def test_built_wheel_scenarios_are_discovered(tmp_path: Path, plugin_sandbox: None) -> None: + """A built wheel's scenarios are all registered in ScenarioRegistry after init.""" + package = f"mock_scenario_plugin_{uuid.uuid4().hex[:8]}" + wheel = _build_scenario_plugin_wheel(tmp_path, package=package) + + registry = ScenarioRegistry.get_registry_singleton() + before = set(registry.get_class_names()) + + with _plugin_env(wheel=wheel, plugin_dir=tmp_path / ".plugin"): + await initialize_pyrit_async(IN_MEMORY) + + registry = ScenarioRegistry.get_registry_singleton() + after = set(registry.get_class_names()) + + expected_registry_names = {registry_name for _stem, _cls, registry_name in _MOCK_SCENARIOS} + # The plug-in adds exactly its scenarios and removes none of the built-ins. + assert after - before == expected_registry_names + assert before <= after + + # And the scenario classes themselves resolve to the plug-in's classes. + expected_class_names = {class_name for _stem, class_name, _ in _MOCK_SCENARIOS} + assert expected_class_names <= _registered_scenario_class_names() + + +@pytest.mark.run_only_if_all_tests +async def test_injected_wheel_scenarios_are_discovered(plugin_sandbox: None) -> None: + """Every scenario in an out-of-band wheel is discovered after initialization. + + Skipped unless ``PLUGIN_TEST_WHEEL`` is set, so the committed test names no external + package. Provide ``PLUGIN_TEST_SCENARIO_DIRS`` (preferred) or ``PLUGIN_TEST_PACKAGE`` + to tell the test which scenarios to expect. + """ + wheel_env = os.getenv("PLUGIN_TEST_WHEEL") + if not wheel_env: + pytest.skip("PLUGIN_TEST_WHEEL is not set; skipping injected-wheel scenario discovery test.") + + wheel = Path(wheel_env).expanduser() + assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" + + scenario_dirs_env = os.getenv("PLUGIN_TEST_SCENARIO_DIRS") + package = os.getenv("PLUGIN_TEST_PACKAGE") + if not scenario_dirs_env and not package: + pytest.skip("Set PLUGIN_TEST_SCENARIO_DIRS or PLUGIN_TEST_PACKAGE to define the expected scenarios.") + + extra_env: dict[str, str] = {} + if package: + extra_env["PLUGIN_PACKAGE"] = package + + with _plugin_env(wheel=wheel, plugin_dir=None, extra=extra_env): + await initialize_pyrit_async(IN_MEMORY) + + found = _registered_scenario_class_names() + + if scenario_dirs_env: + dirs = [Path(p) for p in scenario_dirs_env.split(os.pathsep) if p] + expected = all_scenario_class_names(dirs) + else: + assert package is not None + expected = _scenario_class_names_under_package(package) + + assert expected, "No plug-in scenarios were found to verify; check the injected wheel/scenario source." + missing = expected - found + assert not missing, f"Plug-in scenarios not discovered by ScenarioRegistry: {sorted(missing)}" From f8e761ddeeddbbddc801ec48af673007382b3f6c Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Jul 2026 11:18:56 -0700 Subject: [PATCH 04/34] fix(setup): restore overwritten registry entries on plug-in rollback Both the dataset provider registry (keyed by class name) and the scenario catalog (keyed by registry name) assign unconditionally, so a plug-in whose provider or scenario name collides with an existing one silently replaces the original on registration. Rollback previously only deleted the keys the plug-in added, leaving a collided entry permanently replaced after a failed (or fail-open) load. Rollback now restores the pre-load value for any entry the plug-in overwrote, and only deletes keys the plug-in newly added. Entries present before the load (including built-ins discovered meanwhile) are still preserved. Adds unit tests covering the provider and scenario overwrite-then-restore paths. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/setup/plugin_loader.py | 41 ++++++++++++++--------- tests/unit/setup/test_plugin_loader.py | 46 ++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 15 deletions(-) diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 3e5d9b3c07..2cecbef4d0 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -561,11 +561,14 @@ def _rollback( """ Undo the partial global-state changes made while loading a plug-in. - Removes the plug-in's ``sys.path`` entry, the modules it newly imported, and any - provider/scenario registrations it added, so a failed (or fail-open) load leaves - PyRIT as if the plug-in had never been loaded. State present before the load — - including modules that already existed and built-ins discovered meanwhile — is - preserved. + Removes the plug-in's ``sys.path`` entry, the modules it newly imported, and the + provider/scenario registrations it added, and **restores any entries the plug-in + overwrote**. Both registries key by a name the plug-in does not control + (``SeedDatasetProvider`` by ``cls.__name__``; the scenario catalog by registry + name) and assign unconditionally, so a plug-in whose provider/scenario name + collides with an existing one silently replaces it; rollback must put the original + back, not just drop the new key. State present before the load — including modules + that already existed and built-ins discovered meanwhile — is preserved. Args: package_name: The plug-in's top-level package name. @@ -583,16 +586,24 @@ def _rollback( for name in [m for m in sys.modules if _name_owned_by(m, package_name) and m not in modules_snapshot]: del sys.modules[name] - for key, cls in list(SeedDatasetProvider._registry.items()): - if key not in provider_snapshot and _module_owned_by(cls, package_name): - del SeedDatasetProvider._registry[key] - - removed_scenario = False - for name, cls in list(scenario_registry._classes.items()): - if name not in scenario_snapshot and _module_owned_by(cls, package_name): - del scenario_registry._classes[name] - removed_scenario = True - if removed_scenario: + # For every entry the plug-in now owns: restore the pre-load value if the key + # existed before (the plug-in overwrote it), otherwise drop the key it added. + for key in list(SeedDatasetProvider._registry): + if _module_owned_by(SeedDatasetProvider._registry[key], package_name): + if key in provider_snapshot: + SeedDatasetProvider._registry[key] = provider_snapshot[key] # type: ignore[ty:invalid-assignment] + else: + del SeedDatasetProvider._registry[key] + + changed_scenarios = False + for name in list(scenario_registry._classes): + if _module_owned_by(scenario_registry._classes[name], package_name): + if name in scenario_snapshot: + scenario_registry._classes[name] = scenario_snapshot[name] # type: ignore[ty:invalid-assignment] + else: + del scenario_registry._classes[name] + changed_scenarios = True + if changed_scenarios: scenario_registry._metadata_cache = None def _resolve_fail_open(self) -> bool: diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index 7bf2951c60..2e530f55cd 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -584,6 +584,52 @@ async def test_fail_open_rolls_back_partial_registration(tmp_path: Path) -> None assert wheel.scenario_name not in registry._classes +async def test_rollback_restores_overwritten_provider(tmp_path: Path) -> None: + """A failed load restores a provider entry the plug-in overwrote (name collision).""" + wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") + + # SeedDatasetProvider keys by class name and the mock provider is "MockProvider"; + # occupy that key so the plug-in's import overwrites it. + class _PreexistingProvider(SeedDatasetProvider): + should_register = False + + @property + def dataset_name(self) -> str: + return "preexisting" + + async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: + raise NotImplementedError + + SeedDatasetProvider._registry["MockProvider"] = _PreexistingProvider + + with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(tmp_path / ".plugin")): + with pytest.raises(PluginLoadError): + await load_plugin_if_configured_async() + + # The original provider is restored, not deleted or left replaced by the plug-in's. + assert SeedDatasetProvider._registry["MockProvider"] is _PreexistingProvider + + +async def test_rollback_restores_overwritten_scenario(tmp_path: Path) -> None: + """A failed load restores a scenario entry the plug-in overwrote (name collision).""" + from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + + wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") + + class _PreexistingScenario(RapidResponse): + """Sentinel scenario occupying the plug-in's registry name.""" + + registry = ScenarioRegistry.get_registry_singleton() + registry.register_class(_PreexistingScenario, name=wheel.scenario_name) + + with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(tmp_path / ".plugin")): + with pytest.raises(PluginLoadError): + await load_plugin_if_configured_async() + + # The original scenario is restored, not deleted or left replaced by the plug-in's. + assert registry._classes[wheel.scenario_name] is _PreexistingScenario + + # --------------------------------------------------------------------------- # Extraction cache # --------------------------------------------------------------------------- From ed3fc0d6d30f5d840a164c6328d03404cf6cf8de Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Jul 2026 11:32:52 -0700 Subject: [PATCH 05/34] refactor(setup): extract plug-in wheels with safe_extract_zip Replace the raw zipfile.extractall (and the hand-rolled path-traversal check) in the wheel extractor with pyrit.common.safe_extract.safe_extract_zip, the existing helper used for untrusted archive extraction. It validates every member before writing anything: path traversal, absolute/drive paths, symlink and other non-regular entry types, per-file and total size caps, entry count, and compression ratio (zip bomb), then sanitizes extracted permissions. Extraction stays atomic (temp dir then os.replace) and the cache-reuse path is unchanged. Adds a unit test asserting a wheel with a path-traversal member is rejected and nothing is written outside the extraction directory. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/setup/plugin_loader.py | 16 ++++++---------- tests/unit/setup/test_plugin_loader.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 2cecbef4d0..3082d14c79 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -28,7 +28,6 @@ import pkgutil import shutil import sys -import zipfile from pathlib import Path from typing import TYPE_CHECKING @@ -226,7 +225,9 @@ def _extract_wheel(self, *, wheel_path: Path) -> Path: Extraction is atomic: the wheel is unpacked into a temporary sibling directory and moved into place only on success, so a crash mid-extraction never leaves a partial - tree that would later be treated as a valid cache. + tree that would later be treated as a valid cache. ``safe_extract_zip`` validates + every member first (path traversal, symlinks, and size / entry-count / compression + caps) so a tampered wheel cannot escape the extraction directory or exhaust disk. Args: wheel_path: Path to the plug-in wheel. @@ -234,6 +235,8 @@ def _extract_wheel(self, *, wheel_path: Path) -> Path: Returns: Path: The directory the wheel was extracted to. """ + from pyrit.common.safe_extract import safe_extract_zip + base_dir = self._plugin_base_dir() base_dir.mkdir(parents=True, exist_ok=True) @@ -245,15 +248,8 @@ def _extract_wheel(self, *, wheel_path: Path) -> Path: tmp_dir = base_dir / f".{wheel_path.stem}.tmp-{os.getpid()}" if tmp_dir.exists(): shutil.rmtree(tmp_dir) - tmp_dir.mkdir(parents=True) try: - with zipfile.ZipFile(wheel_path) as wheel_zip: - tmp_dir_resolved = tmp_dir.resolve() - for member in wheel_zip.infolist(): - member_path = (tmp_dir / member.filename).resolve() - if not member_path.is_relative_to(tmp_dir_resolved): - raise ValueError(f"Wheel contains unsafe path: {member.filename}") - wheel_zip.extractall(tmp_dir) + safe_extract_zip(source=wheel_path, dest_dir=tmp_dir) if extract_dir.exists(): shutil.rmtree(extract_dir) os.replace(tmp_dir, extract_dir) diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index 2e530f55cd..2491a4d1bb 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -756,6 +756,21 @@ async def test_non_whl_path_fails_closed(tmp_path: Path) -> None: await load_plugin_if_configured_async() +async def test_wheel_with_path_traversal_member_fails_closed(tmp_path: Path) -> None: + """A wheel containing a path-traversal member is rejected during safe extraction.""" + malicious = tmp_path / "evil-0.0.1-py3-none-any.whl" + with zipfile.ZipFile(malicious, "w") as archive: + archive.writestr("evil_pkg/__init__.py", "") + archive.writestr("../escape.py", "compromised = True") + + with plugin_env(PLUGIN_WHEEL=str(malicious), PLUGIN_DIR=str(tmp_path / ".plugin")): + with pytest.raises(PluginLoadError, match="Failed to load plug-in"): + await load_plugin_if_configured_async() + + # The traversal target was not written outside the extraction directory. + assert not (tmp_path / "escape.py").exists() + + # --------------------------------------------------------------------------- # No-arg-instantiable contract # --------------------------------------------------------------------------- From fb822e05f4860f051b34b35151acafd243e2060d Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Jul 2026 11:35:14 -0700 Subject: [PATCH 06/34] perf(setup): run blocking plug-in extraction off the event loop _load_plugin_async is reached from the async initialization path but performed wheel extraction and directory scanning inline, which blocks the event loop and serializes unrelated async work during startup. Wrap the two filesystem-bound steps (_extract_wheel and _resolve_package_name) in asyncio.to_thread. Package import and bootstrap stay on the loop thread since they mutate global registries and sys state and represent the registration step rather than pure I/O. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/setup/plugin_loader.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 3082d14c79..4d7b7e72f4 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -21,6 +21,7 @@ from __future__ import annotations +import asyncio import importlib import inspect import logging @@ -166,8 +167,10 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: if wheel_path.suffix != ".whl": raise ValueError(f"PLUGIN_WHEEL must point to a .whl file, got: {wheel_path}") - extract_dir = self._extract_wheel(wheel_path=wheel_path) - package_name = self._resolve_package_name(extract_dir=extract_dir) + # Wheel extraction and directory scanning are blocking filesystem work; run them + # off the event loop so init does not stall unrelated async tasks. + extract_dir = await asyncio.to_thread(self._extract_wheel, wheel_path=wheel_path) + package_name = await asyncio.to_thread(self._resolve_package_name, extract_dir=extract_dir) from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider from pyrit.registry import ScenarioRegistry From 95de68c6dccf97c7715c0fb44c4bfe0aa1b09c6e Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Jul 2026 11:47:40 -0700 Subject: [PATCH 07/34] feat(setup): add granular PluginLoadError subclasses A single PluginLoadError collapsed every failure mode into one type. Add subclasses so callers can distinguish causes: PluginWheelNotFoundError (PLUGIN_WHEEL missing or not a .whl), PluginImportError (package import failed, incl. an installed package shadowing the wheel), and PluginRegisteredNothingError (imported but registered nothing). All subclass PluginLoadError, so existing `except PluginLoadError` handling and the fail-open path are unchanged. The load boundary preserves the specific type while still prefixing the standard "Failed to load plug-in ... " guidance; unknown errors (a raising bootstrap, or an unsafe archive) remain a plain PluginLoadError. Tests assert the specific type per failure mode and the shared hierarchy. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/setup/plugin_loader.py | 52 +++++++++++++++++++------- tests/unit/setup/test_plugin_loader.py | 27 +++++++++---- 2 files changed, 58 insertions(+), 21 deletions(-) diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 4d7b7e72f4..b92d0d2f99 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -91,8 +91,26 @@ def _module_owned_by(cls: type, package_name: str) -> bool: return _name_owned_by(cls.__module__ or "", package_name) +_REMEDIATION = ( + "Remove the plug-in configuration, or enable fail-open (PLUGIN_FAIL_OPEN=true, or the " + "initialize_pyrit_async(plugin_fail_open=True) parameter) to continue without it." +) + + class PluginLoadError(RuntimeError): - """Raised when a configured plug-in fails to load and fail-open is not enabled.""" + """Base error raised when a configured plug-in fails to load and fail-open is not enabled.""" + + +class PluginWheelNotFoundError(PluginLoadError): + """``PLUGIN_WHEEL`` does not point to a readable ``.whl`` file.""" + + +class PluginImportError(PluginLoadError): + """The plug-in package could not be imported (missing dependency, or an installed package shadows it).""" + + +class PluginRegisteredNothingError(PluginLoadError): + """The plug-in imported cleanly but registered no datasets or scenarios.""" class PluginLoader: @@ -141,11 +159,13 @@ async def load_async(self) -> None: exc, ) return - raise PluginLoadError( - f"Failed to load plug-in from PLUGIN_WHEEL='{wheel_env}': {exc}. " - "Remove the plug-in configuration, or enable fail-open (PLUGIN_FAIL_OPEN=true, or the " - "initialize_pyrit_async(plugin_fail_open=True) parameter) to continue without it." - ) from exc + message = f"Failed to load plug-in from PLUGIN_WHEEL='{wheel_env}': {exc} {_REMEDIATION}" + # Preserve the specific failure type so callers can distinguish modes, while + # always surfacing the remediation guidance. Unknown errors (e.g. a raising + # bootstrap or an unsafe archive) become a plain PluginLoadError. + if isinstance(exc, PluginLoadError): + raise type(exc)(message) from exc + raise PluginLoadError(message) from exc async def _load_plugin_async(self, *, wheel_path: Path) -> None: """ @@ -159,13 +179,14 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: wheel_path: Path to the pre-built plug-in wheel on disk. Raises: - FileNotFoundError: If ``wheel_path`` does not point to an existing file. - ValueError: If the file is not a ``.whl`` or the plug-in registered nothing. + PluginWheelNotFoundError: If ``wheel_path`` is not an existing ``.whl`` file. + PluginImportError: If the plug-in package cannot be imported. + PluginRegisteredNothingError: If the plug-in registered no datasets or scenarios. """ if not wheel_path.is_file(): - raise FileNotFoundError(f"PLUGIN_WHEEL does not point to an existing file: {wheel_path}") + raise PluginWheelNotFoundError(f"PLUGIN_WHEEL does not point to an existing file: {wheel_path}") if wheel_path.suffix != ".whl": - raise ValueError(f"PLUGIN_WHEEL must point to a .whl file, got: {wheel_path}") + raise PluginWheelNotFoundError(f"PLUGIN_WHEEL must point to a .whl file, got: {wheel_path}") # Wheel extraction and directory scanning are blocking filesystem work; run them # off the event loop so init does not stall unrelated async tasks. @@ -186,9 +207,12 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: try: logger.info("Importing plug-in package '%s'", package_name) - module = importlib.import_module(package_name) - self._verify_module_location(module=module, extract_dir=extract_dir, package_name=package_name) - self._import_submodules(module=module, package_name=package_name) + try: + module = importlib.import_module(package_name) + self._verify_module_location(module=module, extract_dir=extract_dir, package_name=package_name) + self._import_submodules(module=module, package_name=package_name) + except Exception as exc: + raise PluginImportError(f"Could not import plug-in package '{package_name}': {exc}") from exc await self._run_bootstrap_async(package_name=package_name, module=module) @@ -196,7 +220,7 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: package_name=package_name, scenario_registry=scenario_registry ) if not provider_count and not scenario_count: - raise ValueError( + raise PluginRegisteredNothingError( f"Plug-in package '{package_name}' imported successfully but registered no datasets or " "scenarios. The wheel is likely mis-packaged (imports cleanly yet loads nothing)." ) diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index 2491a4d1bb..f144ee2068 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -28,7 +28,14 @@ from pyrit.models import SeedDataset from pyrit.registry import ScenarioRegistry from pyrit.setup.initialization import IN_MEMORY, initialize_pyrit_async -from pyrit.setup.plugin_loader import PluginLoader, PluginLoadError, load_plugin_if_configured_async +from pyrit.setup.plugin_loader import ( + PluginImportError, + PluginLoader, + PluginLoadError, + PluginRegisteredNothingError, + PluginWheelNotFoundError, + load_plugin_if_configured_async, +) # --------------------------------------------------------------------------- # Mock-wheel builder @@ -480,7 +487,7 @@ async def test_shadowing_installed_package_is_rejected(tmp_path: Path) -> None: # PLUGIN_PACKAGE points at a stdlib package that imports from outside the extraction dir. with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(tmp_path / ".plugin"), PLUGIN_PACKAGE="json"): - with pytest.raises(PluginLoadError, match="shadowing"): + with pytest.raises(PluginImportError, match="shadowing"): await load_plugin_if_configured_async() @@ -712,9 +719,9 @@ def test_resolve_package_name_multiple_raises(tmp_path: Path) -> None: async def test_missing_wheel_fails_closed() -> None: - """A configured-but-missing wheel raises by default (fail-closed).""" + """A configured-but-missing wheel raises PluginWheelNotFoundError by default (fail-closed).""" with plugin_env(PLUGIN_WHEEL=str(Path("does_not_exist.whl"))): - with pytest.raises(PluginLoadError, match="Failed to load plug-in"): + with pytest.raises(PluginWheelNotFoundError, match="Failed to load plug-in"): await load_plugin_if_configured_async() @@ -735,7 +742,7 @@ async def test_empty_wheel_is_loud(tmp_path: Path) -> None: wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(tmp_path / ".plugin")): - with pytest.raises(PluginLoadError, match="registered no datasets or scenarios"): + with pytest.raises(PluginRegisteredNothingError, match="registered no datasets or scenarios"): await load_plugin_if_configured_async() @@ -747,15 +754,21 @@ async def test_empty_wheel_fail_open_proceeds(tmp_path: Path) -> None: async def test_non_whl_path_fails_closed(tmp_path: Path) -> None: - """PLUGIN_WHEEL that is not a .whl file fails closed.""" + """PLUGIN_WHEEL that is not a .whl file fails closed with PluginWheelNotFoundError.""" not_a_wheel = tmp_path / "plugin.zip" not_a_wheel.write_text("not a wheel", encoding="utf-8") with plugin_env(PLUGIN_WHEEL=str(not_a_wheel)): - with pytest.raises(PluginLoadError, match="Failed to load plug-in"): + with pytest.raises(PluginWheelNotFoundError, match="Failed to load plug-in"): await load_plugin_if_configured_async() +def test_error_subclasses_are_plugin_load_errors() -> None: + """All specific plug-in errors subclass PluginLoadError so one except still catches them.""" + for error_cls in (PluginWheelNotFoundError, PluginImportError, PluginRegisteredNothingError): + assert issubclass(error_cls, PluginLoadError) + + async def test_wheel_with_path_traversal_member_fails_closed(tmp_path: Path) -> None: """A wheel containing a path-traversal member is rejected during safe extraction.""" malicious = tmp_path / "evil-0.0.1-py3-none-any.whl" From c10857828b51a93200612fbba13c4bcdb4deb5cf Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Jul 2026 12:15:36 -0700 Subject: [PATCH 08/34] fix(setup): use a unique temp dir for plug-in wheel extraction _extract_wheel used a temp dir keyed only by wheel stem + PID. Once extraction was moved off the event loop with asyncio.to_thread, two concurrent loads of the same wheel in one process could collide on that shared path, where one thread's rmtree clobbers the other's in-progress extraction (crash or corrupted cache). Allocate a unique temp dir per extraction with tempfile.mkdtemp (atomic, 0700) so concurrent extractions never share a path. Surfaced by security review of the to_thread change. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/setup/plugin_loader.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index b92d0d2f99..a05be9ccdd 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -29,6 +29,7 @@ import pkgutil import shutil import sys +import tempfile from pathlib import Path from typing import TYPE_CHECKING @@ -272,9 +273,10 @@ def _extract_wheel(self, *, wheel_path: Path) -> Path: logger.info("Reusing cached plug-in extraction at %s", extract_dir) return extract_dir - tmp_dir = base_dir / f".{wheel_path.stem}.tmp-{os.getpid()}" - if tmp_dir.exists(): - shutil.rmtree(tmp_dir) + # Unique per-extraction temp dir so concurrent loads of the same wheel in one + # process cannot collide on a shared path (mkdtemp is atomic and 0700). safe_extract + # into it, then atomically move into place. + tmp_dir = Path(tempfile.mkdtemp(prefix=f".{wheel_path.stem}.tmp-", dir=base_dir)) try: safe_extract_zip(source=wheel_path, dest_dir=tmp_dir) if extract_dir.exists(): From c028f8bf00a47412248c8240cc4117e36dc720ce Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 9 Jul 2026 09:53:24 -0700 Subject: [PATCH 09/34] refactor(setup): rename plugin_fail_open to plugin_accept_load_failures The prior name did not convey what the flag does. Rename the parameter, env var (PLUGIN_FAIL_OPEN -> PLUGIN_ACCEPT_LOAD_FAILURES), and internals so the intent - accept a plug-in load failure and continue - is clear at the call site. Behavior is unchanged: fail-closed by default; explicit param overrides env; env overrides default. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .env_example | 4 +- pyrit/setup/initialization.py | 11 ++-- pyrit/setup/plugin_loader.py | 63 ++++++++++---------- tests/unit/setup/test_plugin_loader.py | 80 +++++++++++++------------- 4 files changed, 81 insertions(+), 77 deletions(-) diff --git a/.env_example b/.env_example index 01ae5e8108..cede548027 100644 --- a/.env_example +++ b/.env_example @@ -275,8 +275,8 @@ OPENAI_VIDEO_ENDPOINT = ${AZURE_OPENAI_VIDEO_ENDPOINT} # PLUGIN_PACKAGE="my_plugin" # Optional: continue (with a warning) instead of failing when the plug-in cannot -# be loaded. Equivalent to initialize_pyrit_async(plugin_fail_open=True). -# PLUGIN_FAIL_OPEN="false" +# be loaded. Equivalent to initialize_pyrit_async(plugin_accept_load_failures=True). +# PLUGIN_ACCEPT_LOAD_FAILURES="false" # Optional: override the base extraction directory (defaults to /.plugin). # PLUGIN_DIR="/abs/path/to/.plugin" diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index a549338cc7..cec450ffdf 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -203,7 +203,7 @@ async def initialize_pyrit_async( env_files: Sequence[pathlib.Path] | None = None, env_akv_ref: Sequence[str] | None = None, silent: bool = False, - plugin_fail_open: bool | None = None, + plugin_accept_load_failures: bool | None = None, **memory_instance_kwargs: Any, ) -> None: """ @@ -225,9 +225,10 @@ async def initialize_pyrit_async( so local files take precedence over AKV. Requires ``azure-keyvault-secrets``. silent (bool): If True, suppresses print statements about environment file loading and schema migration. Defaults to False. - plugin_fail_open (bool | None): Overrides ``PLUGIN_FAIL_OPEN`` for plug-in loading. When True, - a plug-in (``PLUGIN_WHEEL``) that fails to load is skipped with a warning instead of raising. - Defaults to None (use the ``PLUGIN_FAIL_OPEN`` environment variable, else fail-closed). + plugin_accept_load_failures (bool | None): Overrides ``PLUGIN_ACCEPT_LOAD_FAILURES`` for plug-in + loading. When True, a plug-in (``PLUGIN_WHEEL``) that fails to load is skipped with a warning + instead of raising. Defaults to None (use the ``PLUGIN_ACCEPT_LOAD_FAILURES`` environment + variable, else fail-closed). **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. Raises: @@ -270,7 +271,7 @@ async def initialize_pyrit_async( # it does not depend on .pyrit_conf list position. No-op unless PLUGIN_WHEEL is set. from pyrit.setup.plugin_loader import load_plugin_if_configured_async - await load_plugin_if_configured_async(fail_open=plugin_fail_open) + await load_plugin_if_configured_async(accept_load_failures=plugin_accept_load_failures) # Combine directly provided initializers with those loaded from scripts all_initializers = list(initializers) if initializers else [] diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index a05be9ccdd..d241ff2204 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -46,7 +46,7 @@ _TRUE_TOKENS = frozenset({"1", "true", "yes", "on"}) -async def load_plugin_if_configured_async(*, fail_open: bool | None = None) -> None: +async def load_plugin_if_configured_async(*, accept_load_failures: bool | None = None) -> None: """ Load the plug-in referenced by ``PLUGIN_WHEEL`` if one is configured. @@ -54,14 +54,15 @@ async def load_plugin_if_configured_async(*, fail_open: bool | None = None) -> N ``PLUGIN_WHEEL`` is unset. Args: - fail_open: If provided, overrides the ``PLUGIN_FAIL_OPEN`` environment variable. - When True, a plug-in that fails to load is skipped with a warning. + accept_load_failures: If provided, overrides the ``PLUGIN_ACCEPT_LOAD_FAILURES`` + environment variable. When True, a plug-in that fails to load is skipped with + a warning. Raises: - PluginLoadError: If the plug-in is configured but fails to load and fail-open is - not enabled. + PluginLoadError: If the plug-in is configured but fails to load and load failures + are not accepted. """ - await PluginLoader(fail_open=fail_open).load_async() + await PluginLoader(accept_load_failures=accept_load_failures).load_async() def _name_owned_by(module_name: str, package_name: str) -> bool: @@ -93,13 +94,13 @@ def _module_owned_by(cls: type, package_name: str) -> bool: _REMEDIATION = ( - "Remove the plug-in configuration, or enable fail-open (PLUGIN_FAIL_OPEN=true, or the " - "initialize_pyrit_async(plugin_fail_open=True) parameter) to continue without it." + "Remove the plug-in configuration, or accept load failures (PLUGIN_ACCEPT_LOAD_FAILURES=true, or the " + "initialize_pyrit_async(plugin_accept_load_failures=True) parameter) to continue without it." ) class PluginLoadError(RuntimeError): - """Base error raised when a configured plug-in fails to load and fail-open is not enabled.""" + """Base error raised when a configured plug-in fails to load and load failures are not accepted.""" class PluginWheelNotFoundError(PluginLoadError): @@ -121,41 +122,43 @@ class PluginLoader: No-op unless ``PLUGIN_WHEEL`` is set. When set, the wheel is extracted to ``.plugin//`` (never installed), imported, and its bootstrap is run so its datasets and scenarios register like built-ins. Fails closed by default; set - ``fail_open`` (constructor / ``initialize_pyrit_async`` param) or ``PLUGIN_FAIL_OPEN`` - to continue without the plug-in when it cannot be loaded. + ``accept_load_failures`` (constructor / ``initialize_pyrit_async`` param) or + ``PLUGIN_ACCEPT_LOAD_FAILURES`` to continue without the plug-in when it cannot be loaded. """ - def __init__(self, *, fail_open: bool | None = None) -> None: + def __init__(self, *, accept_load_failures: bool | None = None) -> None: """ Initialize the loader. Args: - fail_open: If provided, overrides ``PLUGIN_FAIL_OPEN``. When True, a plug-in - that fails to load is skipped with a warning instead of raising. + accept_load_failures: If provided, overrides ``PLUGIN_ACCEPT_LOAD_FAILURES``. + When True, a plug-in that fails to load is skipped with a warning instead + of raising. """ - self._explicit_fail_open = fail_open + self._explicit_accept_load_failures = accept_load_failures async def load_async(self) -> None: """ Load the plug-in referenced by ``PLUGIN_WHEEL`` (no-op when unset). Raises: - PluginLoadError: If the plug-in is configured but fails to load and fail-open - is not enabled. + PluginLoadError: If the plug-in is configured but fails to load and load + failures are not accepted. """ wheel_env = os.getenv("PLUGIN_WHEEL") if not wheel_env: logger.debug("PLUGIN_WHEEL is not set; plug-in loading is a no-op.") return - fail_open = self._resolve_fail_open() + accept_load_failures = self._resolve_accept_load_failures() try: await self._load_plugin_async(wheel_path=Path(wheel_env).expanduser()) except Exception as exc: - if fail_open: + if accept_load_failures: logger.warning( - "Plug-in from PLUGIN_WHEEL='%s' failed to load; fail_open is set so continuing without it: %s", + "Plug-in from PLUGIN_WHEEL='%s' failed to load; load failures are accepted so " + "continuing without it: %s", wheel_env, exc, ) @@ -173,7 +176,7 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: Extract, import, bootstrap, and verify a single plug-in wheel. Global state (``sys.path``, imported plug-in modules, and the provider/scenario - registries) is rolled back if the load fails, so a failed or fail-open load + registries) is rolled back if the load fails, so a failed or accepted-failure load leaves no partial trace. Args: @@ -528,7 +531,7 @@ def _warn_on_dataset_name_collisions(*, package_name: str) -> list[str]: noisy. Hard enforcement that can tell a real mismatch from a harmless name coincidence belongs to the scenario's declared required-dataset-names / expected-source mechanism (which knows the operator's intent), not this loader, and is intentionally not gated - behind ``PLUGIN_FAIL_OPEN``. + behind ``PLUGIN_ACCEPT_LOAD_FAILURES``. Args: package_name: The plug-in's top-level package name. @@ -631,21 +634,21 @@ def _rollback( if changed_scenarios: scenario_registry._metadata_cache = None - def _resolve_fail_open(self) -> bool: + def _resolve_accept_load_failures(self) -> bool: """ - Resolve the fail-open setting from the explicit value or the environment. + Resolve the accept-load-failures setting from the explicit value or the environment. - Precedence: the explicit ``fail_open`` passed to the constructor (e.g. from - ``initialize_pyrit_async``), then the ``PLUGIN_FAIL_OPEN`` environment variable, - otherwise fail-closed. + Precedence: the explicit ``accept_load_failures`` passed to the constructor (e.g. from + ``initialize_pyrit_async``), then the ``PLUGIN_ACCEPT_LOAD_FAILURES`` environment + variable, otherwise fail-closed. Returns: bool: True if a failed plug-in load should be skipped with a warning. """ - if self._explicit_fail_open is not None: - return self._explicit_fail_open + if self._explicit_accept_load_failures is not None: + return self._explicit_accept_load_failures - env_value = os.getenv("PLUGIN_FAIL_OPEN") + env_value = os.getenv("PLUGIN_ACCEPT_LOAD_FAILURES") if env_value is not None: return self._coerce_bool(env_value) diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index f144ee2068..49565cc17b 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -6,7 +6,7 @@ These tests build a **mock plug-in wheel** at test time (no dependency on any real plug-in) and exercise the full consumer mechanism: extract -> sys.path -> -import -> bootstrap -> assert-loaded, plus the fail-open/closed policy and the +import -> bootstrap -> assert-loaded, plus the accept-load-failures/fail-closed policy and the silent-failure guards called out in the design brief. """ @@ -285,7 +285,7 @@ def plugin_sandbox() -> Iterator[None]: def plugin_env(**overrides: str) -> Iterator[None]: """Patch os.environ so only the given PLUGIN_* overrides are present.""" with patch.dict(os.environ, overrides, clear=False): - for key in ("PLUGIN_WHEEL", "PLUGIN_DIR", "PLUGIN_PACKAGE", "PLUGIN_FAIL_OPEN"): + for key in ("PLUGIN_WHEEL", "PLUGIN_DIR", "PLUGIN_PACKAGE", "PLUGIN_ACCEPT_LOAD_FAILURES"): if key not in overrides: os.environ.pop(key, None) yield @@ -295,7 +295,7 @@ async def load_plugin( wheel: MockWheel, plugin_dir: Path, *, - fail_open: bool | None = None, + accept_load_failures: bool | None = None, extra_env: dict[str, str] | None = None, ) -> None: """Run the plug-in loader against a mock wheel with an isolated env.""" @@ -304,7 +304,7 @@ async def load_plugin( env.update(extra_env) with plugin_env(**env): - await load_plugin_if_configured_async(fail_open=fail_open) + await load_plugin_if_configured_async(accept_load_failures=accept_load_failures) # --------------------------------------------------------------------------- @@ -331,17 +331,17 @@ async def test_plugin_phase_runs_after_memory_before_initializers() -> None: assert order.index("set_memory") < order.index("load_plugin") < order.index("execute") -async def test_plugin_phase_forwards_fail_open_param() -> None: - """initialize_pyrit_async forwards plugin_fail_open to the loader.""" +async def test_plugin_phase_forwards_accept_load_failures_param() -> None: + """initialize_pyrit_async forwards plugin_accept_load_failures to the loader.""" load_plugin_mock = AsyncMock() with ( patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), patch.object(CentralMemory, "set_memory_instance"), patch("pyrit.setup.plugin_loader.load_plugin_if_configured_async", load_plugin_mock), ): - await initialize_pyrit_async(IN_MEMORY, env_files=[], silent=True, plugin_fail_open=True) + await initialize_pyrit_async(IN_MEMORY, env_files=[], silent=True, plugin_accept_load_failures=True) - load_plugin_mock.assert_awaited_once_with(fail_open=True) + load_plugin_mock.assert_awaited_once_with(accept_load_failures=True) # --------------------------------------------------------------------------- @@ -580,12 +580,12 @@ async def test_failing_bootstrap_rolls_back_partial_registration(tmp_path: Path) assert str(plugin_dir / wheel.path.stem) not in sys.path -async def test_fail_open_rolls_back_partial_registration(tmp_path: Path) -> None: - """Under fail_open, a partially-registered failed plug-in is still fully rolled back.""" +async def test_accept_load_failures_rolls_back_partial_registration(tmp_path: Path) -> None: + """When load failures are accepted, a partially-registered failed plug-in is still fully rolled back.""" wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") plugin_dir = tmp_path / ".plugin" - await load_plugin(wheel, plugin_dir, fail_open=True) # must not raise + await load_plugin(wheel, plugin_dir, accept_load_failures=True) # must not raise registry = ScenarioRegistry.get_registry_singleton() assert wheel.scenario_name not in registry._classes @@ -725,15 +725,15 @@ async def test_missing_wheel_fails_closed() -> None: await load_plugin_if_configured_async() -async def test_missing_wheel_fail_open_param_proceeds() -> None: - """fail_open via the explicit param skips a broken plug-in with a warning.""" +async def test_missing_wheel_accept_load_failures_param_proceeds() -> None: + """Accepting load failures via the explicit param skips a broken plug-in with a warning.""" with plugin_env(PLUGIN_WHEEL=str(Path("does_not_exist.whl"))): - await load_plugin_if_configured_async(fail_open=True) # must not raise + await load_plugin_if_configured_async(accept_load_failures=True) # must not raise -async def test_missing_wheel_fail_open_env_proceeds() -> None: - """fail_open via PLUGIN_FAIL_OPEN env skips a broken plug-in with a warning.""" - with plugin_env(PLUGIN_WHEEL=str(Path("does_not_exist.whl")), PLUGIN_FAIL_OPEN="true"): +async def test_missing_wheel_accept_load_failures_env_proceeds() -> None: + """Accepting load failures via PLUGIN_ACCEPT_LOAD_FAILURES env skips a broken plug-in with a warning.""" + with plugin_env(PLUGIN_WHEEL=str(Path("does_not_exist.whl")), PLUGIN_ACCEPT_LOAD_FAILURES="true"): await load_plugin_if_configured_async() # must not raise @@ -746,11 +746,11 @@ async def test_empty_wheel_is_loud(tmp_path: Path) -> None: await load_plugin_if_configured_async() -async def test_empty_wheel_fail_open_proceeds(tmp_path: Path) -> None: - """An empty wheel under fail_open proceeds instead of raising.""" +async def test_empty_wheel_accept_load_failures_proceeds(tmp_path: Path) -> None: + """An empty wheel with load failures accepted proceeds instead of raising.""" wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) - await load_plugin(wheel, tmp_path / ".plugin", fail_open=True) # must not raise + await load_plugin(wheel, tmp_path / ".plugin", accept_load_failures=True) # must not raise async def test_non_whl_path_fails_closed(tmp_path: Path) -> None: @@ -808,7 +808,7 @@ def __init__(self, *, required_value: str) -> None: # --------------------------------------------------------------------------- -# fail_open resolution +# accept_load_failures resolution # --------------------------------------------------------------------------- @@ -816,31 +816,31 @@ def __init__(self, *, required_value: str) -> None: "value,expected", [("true", True), ("True", True), ("1", True), ("yes", True), ("on", True), ("false", False), ("0", False)], ) -def test_resolve_fail_open_from_env_tokens(value: str, expected: bool) -> None: - """fail_open resolves from PLUGIN_FAIL_OPEN across truthy/falsey tokens.""" - with plugin_env(PLUGIN_FAIL_OPEN=value): - assert PluginLoader()._resolve_fail_open() is expected +def test_resolve_accept_load_failures_from_env_tokens(value: str, expected: bool) -> None: + """accept_load_failures resolves from PLUGIN_ACCEPT_LOAD_FAILURES across truthy/falsey tokens.""" + with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES=value): + assert PluginLoader()._resolve_accept_load_failures() is expected -def test_resolve_fail_open_explicit_true() -> None: - """An explicit fail_open=True resolves to True.""" - with plugin_env(PLUGIN_FAIL_OPEN="false"): - assert PluginLoader(fail_open=True)._resolve_fail_open() is True +def test_resolve_accept_load_failures_explicit_true() -> None: + """An explicit accept_load_failures=True resolves to True.""" + with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="false"): + assert PluginLoader(accept_load_failures=True)._resolve_accept_load_failures() is True -def test_resolve_fail_open_explicit_overrides_env() -> None: - """An explicit fail_open value takes precedence over the env var.""" - with plugin_env(PLUGIN_FAIL_OPEN="true"): - assert PluginLoader(fail_open=False)._resolve_fail_open() is False +def test_resolve_accept_load_failures_explicit_overrides_env() -> None: + """An explicit accept_load_failures value takes precedence over the env var.""" + with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): + assert PluginLoader(accept_load_failures=False)._resolve_accept_load_failures() is False -def test_resolve_fail_open_from_env_when_no_explicit() -> None: - """fail_open falls back to PLUGIN_FAIL_OPEN when no explicit value is set.""" - with plugin_env(PLUGIN_FAIL_OPEN="true"): - assert PluginLoader()._resolve_fail_open() is True +def test_resolve_accept_load_failures_from_env_when_no_explicit() -> None: + """accept_load_failures falls back to PLUGIN_ACCEPT_LOAD_FAILURES when no explicit value is set.""" + with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): + assert PluginLoader()._resolve_accept_load_failures() is True -def test_resolve_fail_open_defaults_false() -> None: - """fail_open defaults to False (fail-closed).""" +def test_resolve_accept_load_failures_defaults_false() -> None: + """accept_load_failures defaults to False (fail-closed).""" with plugin_env(): - assert PluginLoader()._resolve_fail_open() is False + assert PluginLoader()._resolve_accept_load_failures() is False From 2ae6c5593d1803179a260c45e4309aa11814eccd Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 9 Jul 2026 10:25:16 -0700 Subject: [PATCH 10/34] feat(setup): configure plug-ins in .pyrit_conf with multi-plug-in support Move plug-in configuration out of .env and into a dedicated .pyrit_conf 'plugins' list. Each entry is a wheel path, a {package: wheel} pair, or an explicit {wheel, package} mapping, parsed into a PluginSpec. initialize_pyrit_async now takes a plugins sequence and loads them in list order as the guaranteed-first phase, so one plug-in behaves identically to several. A configless setup simply has no plug-ins. Plug-in loading stays a fixed phase (after memory is set, before configured initializers) rather than a user-ordered initializer, so 'plug-ins register before anything reads the registry' remains true by construction. A new _verify_plugin_prerequisites health check gates loading on the prior init phases (env, memory) having completed. The PLUGIN_WHEEL/PLUGIN_PACKAGE env vars are removed as config sources; PLUGIN_DIR and PLUGIN_ACCEPT_LOAD_FAILURES remain env-overridable. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .env_example | 29 -- .pyrit_conf_example | 36 ++- pyrit/setup/__init__.py | 2 + pyrit/setup/configuration_loader.py | 45 +++ pyrit/setup/initialization.py | 56 +++- pyrit/setup/plugin_loader.py | 182 ++++++++---- .../setup/test_plugin_loader_integration.py | 22 +- tests/unit/setup/test_configuration_loader.py | 79 ++++++ tests/unit/setup/test_plugin_loader.py | 264 +++++++++++++----- 9 files changed, 535 insertions(+), 180 deletions(-) diff --git a/.env_example b/.env_example index 138a60adfa..0bec2cbb6e 100644 --- a/.env_example +++ b/.env_example @@ -251,35 +251,6 @@ AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" OPENAI_VIDEO_ENDPOINT = ${AZURE_OPENAI_VIDEO_ENDPOINT} -################################## -# PLUG-INS -# -# PyRIT can load a non-disclosable plug-in (extra datasets + scenarios) from a -# pre-built wheel at initialization time. Plug-in loading runs automatically as a -# guaranteed-first phase inside initialize_pyrit_async (before the configured -# initializers) — there is nothing to add to .pyrit_conf. The wheel is extracted to -# .plugin//, prepended to sys.path, imported, and its bootstrap is run so its -# datasets/scenarios register like built-ins. Extraction only — never pip/.venv. -# -# WARNING: loading a plug-in executes third-party code in the PyRIT process. -# Whoever can set PLUGIN_* / write this .env can run code on the host. Treat this -# file as sensitive. -################################### - -# Path to a pre-built plug-in wheel on disk. Setting it enables plug-in loading; -# leaving it unset is a no-op. -# PLUGIN_WHEEL="/abs/path/to/my_plugin-0.0.0-py3-none-any.whl" - -# Optional: the wheel's top-level import package. Auto-detected from the wheel -# when omitted; set this to disambiguate multi-package wheels. -# PLUGIN_PACKAGE="my_plugin" - -# Optional: continue (with a warning) instead of failing when the plug-in cannot -# be loaded. Equivalent to initialize_pyrit_async(plugin_accept_load_failures=True). -# PLUGIN_ACCEPT_LOAD_FAILURES="false" - -# Optional: override the base extraction directory (defaults to /.plugin). -# PLUGIN_DIR="/abs/path/to/.plugin" OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 95c1b928f5..e14a0c599d 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -28,9 +28,9 @@ memory_db_type: sqlite # - load_default_datasets: Loads datasets into memory so scenarios can run # - preload_scenario_metadata: Preloads scenario metadata into the registry # -# Note: plug-ins (PLUGIN_WHEEL) are NOT configured here. They load automatically as a -# guaranteed-first phase inside initialize_pyrit_async — before these initializers — so -# there is no ordering to get right. See the PLUGIN_* variables in .env_example. +# Note: plug-ins are configured below via the dedicated `plugins:` key, NOT as an entry +# in this `initializers:` list. Plug-ins load as a guaranteed-first phase (before these +# initializers), so there is no ordering to get right here. # # Each initializer can be specified as: # - A simple string (name only) @@ -57,6 +57,36 @@ initializers: - name: technique - name: load_default_datasets +# Plug-ins +# -------- +# Non-disclosable plug-ins (extra scenarios + datasets) shipped as pre-built wheels and +# loaded at initialization. Plug-ins load as a guaranteed-first phase — after memory is set +# and BEFORE the initializers above — so their scenarios/datasets register before anything +# reads the registry. They are a dedicated phase, not an entry in `initializers:`, so there +# is no ordering to get right. Loaded in list order; one plug-in behaves identically to +# several. A configless setup simply has no plug-ins. +# +# Each entry can be: +# - a wheel path (package auto-detected): - /abs/path/to/my_plugin.whl +# - a {package: wheel} pair (names the package): - my_plugin: /abs/path/to/my_plugin.whl +# - an explicit mapping: - wheel: /abs/path/to/my_plugin.whl +# package: my_plugin +# +# WARNING: loading a plug-in executes third-party code in the PyRIT process. Whoever can +# write this file can run code on the host. Treat it as sensitive. +# +# Example: +# plugins: +# - my_first_plugin: /abs/path/to/first-0.0.0-py3-none-any.whl +# - my_second_plugin: /abs/path/to/second-0.0.0-py3-none-any.whl + +# Accept plug-in load failures +# ---------------------------- +# When true, a plug-in that fails to load is skipped with a warning instead of aborting +# initialization. Defaults to false (fail-closed). Precedence: this value wins when set; +# the PLUGIN_ACCEPT_LOAD_FAILURES env var is used only when this is omitted. +# plugin_accept_load_failures: false + # Default Scenario # ---------------- # Optional default scenario to run when invoking `pyrit_scan` without a diff --git a/pyrit/setup/__init__.py b/pyrit/setup/__init__.py index 4cac6e1470..1f16003a42 100644 --- a/pyrit/setup/__init__.py +++ b/pyrit/setup/__init__.py @@ -11,6 +11,7 @@ MemoryDatabaseType, initialize_pyrit_async, ) +from pyrit.setup.plugin_loader import PluginSpec __all__ = [ "AZURE_SQL", @@ -20,4 +21,5 @@ "initialize_from_config_async", "MemoryDatabaseType", "ConfigurationLoader", + "PluginSpec", ] diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 416421cb91..6870def2fa 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -24,6 +24,7 @@ ) if TYPE_CHECKING: + from pyrit.setup.plugin_loader import PluginSpec from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -106,6 +107,11 @@ class ConfigurationLoader(YamlLoadable): silent: Whether to suppress initialization messages. operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. + plugins: List of plug-ins to load as the guaranteed-first initialization phase. Each + entry is a wheel path string, a ``{package: wheel}`` mapping, or a + ``{wheel: ..., package: ...}`` mapping. Empty means no plug-ins. + plugin_accept_load_failures: When True, a plug-in that fails to load is skipped with a + warning instead of raising. None defers to ``PLUGIN_ACCEPT_LOAD_FAILURES`` (else fail-closed). Example YAML configuration: memory_db_type: sqlite @@ -149,10 +155,15 @@ class ConfigurationLoader(YamlLoadable): max_concurrent_scenario_runs: int = 3 allow_custom_initializers: bool = False server: dict[str, Any] | None = None + plugins: list[str | dict[str, Any]] = field(default_factory=list) + plugin_accept_load_failures: bool | None = None extensions: dict[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: """Validate and normalize the configuration after loading.""" + # Normalize plug-ins first: they load as the guaranteed-first phase of + # initialization, so a malformed plug-in entry should fail fast before anything else. + self._normalize_plugins() self._normalize_memory_db_type() self._normalize_initializers() self._normalize_scenario() @@ -216,6 +227,22 @@ def _normalize_initializers(self) -> None: raise ValueError(f"Initializer entry must be a string or dict, got: {type(entry).__name__}") self._initializer_configs = normalized + def _normalize_plugins(self) -> None: + """ + Normalize ``plugins`` entries to ``PluginSpec`` instances. + + Each entry is a wheel path string, a single-key ``{package: wheel}`` mapping, or an + explicit ``{wheel: ..., package: ...}`` mapping. Validating here (before any other + normalization) fails fast on a malformed plug-in entry, since plug-ins load as the + guaranteed-first phase of initialization. + + Raises: + ValueError: If a plug-in entry has an unsupported shape. + """ + from pyrit.setup.plugin_loader import PluginSpec + + self._plugin_specs = [PluginSpec.from_config(entry) for entry in self.plugins] + def _normalize_scenario(self) -> None: """ Normalize the optional ``scenario`` block to a ``ScenarioConfig``. @@ -362,6 +389,8 @@ def load_with_overrides( "env_files": None, # None = use defaults "env_akv_ref": None, "silent": False, + "plugins": [], + "plugin_accept_load_failures": None, } # 1. Try loading default config file if it exists @@ -381,6 +410,8 @@ def load_with_overrides( config_data["env_files"] = default_config.env_files config_data["env_akv_ref"] = default_config.env_akv_ref config_data["silent"] = default_config.silent + config_data["plugins"] = list(default_config.plugins) + config_data["plugin_accept_load_failures"] = default_config.plugin_accept_load_failures if default_config.operator: config_data["operator"] = default_config.operator if default_config.operation: @@ -407,6 +438,8 @@ def load_with_overrides( config_data["env_files"] = explicit_config.env_files config_data["env_akv_ref"] = explicit_config.env_akv_ref config_data["silent"] = explicit_config.silent + config_data["plugins"] = list(explicit_config.plugins) + config_data["plugin_accept_load_failures"] = explicit_config.plugin_accept_load_failures if explicit_config.operator: config_data["operator"] = explicit_config.operator if explicit_config.operation: @@ -491,6 +524,15 @@ def resolve_initializers(self) -> Sequence["PyRITInitializer"]: return resolved + def resolve_plugins(self) -> list["PluginSpec"]: + """ + Resolve the configured ``plugins`` entries to ``PluginSpec`` instances. + + Returns: + list[PluginSpec]: The plug-ins to load, in configured order (empty if none). + """ + return list(self._plugin_specs) + def resolve_initialization_scripts(self) -> Sequence[pathlib.Path] | None: """ Resolve initialization script paths. @@ -563,6 +605,7 @@ async def initialize_pyrit_async(self) -> None: resolved_initializers = self.resolve_initializers() resolved_scripts = self.resolve_initialization_scripts() resolved_env_files = self.resolve_env_files() + resolved_plugins = self.resolve_plugins() # Map snake_case memory_db_type to internal constant internal_memory_db_type = self._MEMORY_DB_TYPE_MAP[self.memory_db_type] @@ -574,6 +617,8 @@ async def initialize_pyrit_async(self) -> None: env_files=resolved_env_files, env_akv_ref=self.env_akv_ref, silent=self.silent, + plugins=resolved_plugins if resolved_plugins else None, + plugin_accept_load_failures=self.plugin_accept_load_failures, ) diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index cec450ffdf..efb711316f 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -13,6 +13,7 @@ from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory if TYPE_CHECKING: + from pyrit.setup.plugin_loader import PluginSpec from pyrit.setup.pyrit_initializer import PyRITInitializer logger = logging.getLogger(__name__) @@ -195,6 +196,30 @@ async def _execute_initializers_async(*, initializers: Sequence["PyRITInitialize raise +def _verify_plugin_prerequisites() -> None: + """ + Verify the initialization phases plug-ins depend on completed successfully. + + Plug-in loading runs a third party's bootstrap, which typically constructs targets and + scorers (needing loaded environment variables) and reads or writes central memory. This + is the health gate for that phase: plug-ins are loaded only if the prior initialization + phases were green. Central memory being set is the concrete, verifiable signal that the + environment and memory phases completed — the linear init flow would have raised earlier + otherwise, so a violation here means a broken initialization state, not a plug-in fault, + and is surfaced loudly regardless of ``plugin_accept_load_failures``. + + Raises: + RuntimeError: If central memory is not set, indicating a prior phase did not complete. + """ + try: + CentralMemory.get_memory_instance() + except Exception as exc: + raise RuntimeError( + "Cannot load plug-ins: central memory is not initialized. Plug-in loading requires " + "the environment and memory phases of initialize_pyrit_async to have completed first." + ) from exc + + async def initialize_pyrit_async( memory_db_type: MemoryDatabaseType | str, *, @@ -203,6 +228,7 @@ async def initialize_pyrit_async( env_files: Sequence[pathlib.Path] | None = None, env_akv_ref: Sequence[str] | None = None, silent: bool = False, + plugins: Sequence["PluginSpec"] | None = None, plugin_accept_load_failures: bool | None = None, **memory_instance_kwargs: Any, ) -> None: @@ -225,10 +251,15 @@ async def initialize_pyrit_async( so local files take precedence over AKV. Requires ``azure-keyvault-secrets``. silent (bool): If True, suppresses print statements about environment file loading and schema migration. Defaults to False. + plugins (Sequence[PluginSpec] | None): Optional plug-ins to load as a guaranteed-first phase, + after central memory is set and before the configured initializers run. Each plug-in is a + pre-built wheel shipping datasets/scenarios that register like built-ins. Plug-ins are + loaded in list order, so one plug-in behaves identically to several. Typically populated + from the ``plugins`` key of ``.pyrit_conf``; direct callers pass ``PluginSpec`` instances. plugin_accept_load_failures (bool | None): Overrides ``PLUGIN_ACCEPT_LOAD_FAILURES`` for plug-in - loading. When True, a plug-in (``PLUGIN_WHEEL``) that fails to load is skipped with a warning - instead of raising. Defaults to None (use the ``PLUGIN_ACCEPT_LOAD_FAILURES`` environment - variable, else fail-closed). + loading. When True, a plug-in that fails to load is skipped with a warning instead of + raising. Defaults to None (use the ``PLUGIN_ACCEPT_LOAD_FAILURES`` environment variable, + else fail-closed). **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. Raises: @@ -264,14 +295,19 @@ async def initialize_pyrit_async( CentralMemory.set_memory_instance(memory) - # Load a configured plug-in (PLUGIN_WHEEL) as a guaranteed-first phase: after memory - # is set (a plug-in bootstrap may use it) and BEFORE any configured initializers run, - # so plug-in datasets/scenarios are registered before LoadDefaultDatasets and - # PreloadScenarioMetadata read the registry. Ordering is true by construction here, so - # it does not depend on .pyrit_conf list position. No-op unless PLUGIN_WHEEL is set. - from pyrit.setup.plugin_loader import load_plugin_if_configured_async + # Load configured plug-ins as a guaranteed-first phase: after central memory is set (a + # plug-in bootstrap may use it) and BEFORE any configured initializer runs, so plug-in + # datasets/scenarios are registered before LoadDefaultDatasets and PreloadScenarioMetadata + # read the registry. Because `plugins` is its own always-first phase (not one of the + # ordered `initializers`), this ordering is true by construction. Plug-ins are loaded only + # after a health check confirms the initialization phases they depend on completed. No-op + # when no plug-ins are configured. + if plugins: + _verify_plugin_prerequisites() + + from pyrit.setup.plugin_loader import load_plugins_if_configured_async - await load_plugin_if_configured_async(accept_load_failures=plugin_accept_load_failures) + await load_plugins_if_configured_async(plugins=plugins, accept_load_failures=plugin_accept_load_failures) # Combine directly provided initializers with those loaded from scripts all_initializers = list(initializers) if initializers else [] diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index d241ff2204..44a3d30ff9 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -2,7 +2,7 @@ # Licensed under the MIT license. """ -Load a non-disclosable PyRIT plug-in from a pre-built wheel at initialization time. +Load non-disclosable PyRIT plug-ins from pre-built wheels at initialization time. A plug-in is a pure-Python wheel that ships dataset providers and/or scenarios that must not live in the public PyRIT repo. The loader **extracts** the wheel (stdlib @@ -11,12 +11,13 @@ and runs the plug-in's bootstrap (a top-level ``register()`` callable or a shipped ``PyRITInitializer`` subclass) which registers the plug-in's scenarios. -``load_plugin_if_configured_async`` is invoked as a guaranteed-first phase inside -``initialize_pyrit_async`` — after central memory is set and **before** any configured -initializers run — so plug-in datasets and scenarios are registered before -``LoadDefaultDatasets`` / ``PreloadScenarioMetadata`` read the registry. Ordering is -therefore true by construction, without relying on ``.pyrit_conf`` list position. It is a -no-op when ``PLUGIN_WHEEL`` is unset. +Plug-ins are declared in ``.pyrit_conf`` under the dedicated ``plugins`` key (a list, so +several plug-ins load identically to one). ``load_plugins_if_configured_async`` is invoked +as a guaranteed-first phase inside ``initialize_pyrit_async`` — after central memory is set +and **before** any configured initializer runs — so plug-in datasets and scenarios are +registered before ``LoadDefaultDatasets`` / ``PreloadScenarioMetadata`` read the registry. +Because ``plugins`` is its own always-first phase (not one of the ordered ``initializers``), +this ordering is true by construction. It is a no-op when no plug-ins are configured. """ from __future__ import annotations @@ -30,13 +31,15 @@ import shutil import sys import tempfile +from collections.abc import Mapping +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING from pyrit.setup.pyrit_initializer import PyRITInitializer if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Sequence from types import ModuleType from pyrit.registry import ScenarioRegistry @@ -46,23 +49,104 @@ _TRUE_TOKENS = frozenset({"1", "true", "yes", "on"}) -async def load_plugin_if_configured_async(*, accept_load_failures: bool | None = None) -> None: +@dataclass(frozen=True) +class PluginSpec: """ - Load the plug-in referenced by ``PLUGIN_WHEEL`` if one is configured. + A single plug-in to load: a pre-built wheel plus its optional top-level package name. - Convenience entry point invoked by ``initialize_pyrit_async``. A no-op when - ``PLUGIN_WHEEL`` is unset. + Attributes: + wheel (Path): Filesystem path to the pre-built plug-in wheel (``.whl``). + package (str | None): The wheel's top-level import package. When ``None`` it is + auto-detected from the wheel; set it to disambiguate a multi-package wheel. + """ + + wheel: Path + package: str | None = None + + @classmethod + def from_config(cls, entry: str | Mapping[str, object]) -> PluginSpec: + """ + Build a ``PluginSpec`` from a ``.pyrit_conf`` ``plugins`` list entry. + + Accepted shapes: + + * A bare wheel path string (package auto-detected):: + + - /abs/path/to/plugin.whl + + * A single-key ``{package: wheel}`` mapping (concise, names the package):: + + - my_plugin: /abs/path/to/plugin.whl + + * An explicit ``{wheel: ..., package: ...}`` mapping (``package`` optional):: + + - wheel: /abs/path/to/plugin.whl + package: my_plugin + + Args: + entry: One ``plugins`` list item from the parsed configuration. + + Returns: + PluginSpec: The normalized spec. + + Raises: + ValueError: If the entry shape is not one of the accepted forms. + """ + if isinstance(entry, str): + return cls(wheel=Path(entry).expanduser()) + + if isinstance(entry, Mapping): + mapping = dict(entry) + if "wheel" in mapping: + extra_keys = set(mapping) - {"wheel", "package"} + if extra_keys: + raise ValueError( + f"Plug-in mapping has unexpected key(s) {sorted(extra_keys)}; only 'wheel' and " + f"'package' are allowed. Got: {entry}" + ) + wheel = mapping["wheel"] + package = mapping.get("package") + if not isinstance(wheel, str) or (package is not None and not isinstance(package, str)): + raise ValueError(f"Plug-in entry 'wheel'/'package' must be strings. Got: {entry}") + return cls(wheel=Path(wheel).expanduser(), package=package) + if len(mapping) == 1 and "package" not in mapping: + ((package, wheel),) = mapping.items() + if not isinstance(package, str) or not isinstance(wheel, str): + raise ValueError(f"Plug-in entry must map a package name to a wheel path. Got: {entry}") + return cls(wheel=Path(wheel).expanduser(), package=package) + raise ValueError( + f"Plug-in mapping must be a single {{package_name: wheel}} pair or carry a 'wheel' key. Got: {entry}" + ) + + raise ValueError(f"Plug-in entry must be a wheel path string or mapping, got: {type(entry).__name__}") + + +async def load_plugins_if_configured_async( + *, plugins: Sequence[PluginSpec], accept_load_failures: bool | None = None +) -> None: + """ + Load every configured plug-in, in order. A no-op when ``plugins`` is empty. + + Convenience entry point invoked by ``initialize_pyrit_async`` after memory is set and + before the configured initializers run. Plug-ins are loaded in list order, so one + plug-in behaves identically to several — the single-plug-in case is just a + one-element list. Args: + plugins: The plug-in specs to load, in order. accept_load_failures: If provided, overrides the ``PLUGIN_ACCEPT_LOAD_FAILURES`` environment variable. When True, a plug-in that fails to load is skipped with - a warning. + a warning and the next plug-in still loads. Raises: - PluginLoadError: If the plug-in is configured but fails to load and load failures - are not accepted. + PluginLoadError: If a plug-in fails to load and load failures are not accepted. """ - await PluginLoader(accept_load_failures=accept_load_failures).load_async() + if not plugins: + logger.debug("No plug-ins configured; plug-in loading is a no-op.") + return + + for spec in plugins: + await PluginLoader(spec=spec, accept_load_failures=accept_load_failures).load_async() def _name_owned_by(module_name: str, package_name: str) -> bool: @@ -104,7 +188,7 @@ class PluginLoadError(RuntimeError): class PluginWheelNotFoundError(PluginLoadError): - """``PLUGIN_WHEEL`` does not point to a readable ``.whl`` file.""" + """The configured plug-in wheel path does not point to a readable ``.whl`` file.""" class PluginImportError(PluginLoadError): @@ -117,53 +201,49 @@ class PluginRegisteredNothingError(PluginLoadError): class PluginLoader: """ - Extract and register a PyRIT plug-in wheel referenced by ``PLUGIN_WHEEL``. + Extract and register a single PyRIT plug-in wheel described by a ``PluginSpec``. - No-op unless ``PLUGIN_WHEEL`` is set. When set, the wheel is extracted to - ``.plugin//`` (never installed), imported, and its bootstrap is run so its - datasets and scenarios register like built-ins. Fails closed by default; set - ``accept_load_failures`` (constructor / ``initialize_pyrit_async`` param) or - ``PLUGIN_ACCEPT_LOAD_FAILURES`` to continue without the plug-in when it cannot be loaded. + The wheel is extracted to ``.plugin//`` (never installed), imported, and its + bootstrap is run so its datasets and scenarios register like built-ins. Fails closed + by default; set ``accept_load_failures`` (constructor / ``initialize_pyrit_async`` + param) or ``PLUGIN_ACCEPT_LOAD_FAILURES`` to continue without the plug-in when it + cannot be loaded. """ - def __init__(self, *, accept_load_failures: bool | None = None) -> None: + def __init__(self, *, spec: PluginSpec, accept_load_failures: bool | None = None) -> None: """ - Initialize the loader. + Initialize the loader for one plug-in. Args: + spec: The plug-in to load (wheel path plus optional package name). accept_load_failures: If provided, overrides ``PLUGIN_ACCEPT_LOAD_FAILURES``. When True, a plug-in that fails to load is skipped with a warning instead of raising. """ + self._spec = spec self._explicit_accept_load_failures = accept_load_failures async def load_async(self) -> None: """ - Load the plug-in referenced by ``PLUGIN_WHEEL`` (no-op when unset). + Load the plug-in described by this loader's ``PluginSpec``. Raises: - PluginLoadError: If the plug-in is configured but fails to load and load - failures are not accepted. + PluginLoadError: If the plug-in fails to load and load failures are not accepted. """ - wheel_env = os.getenv("PLUGIN_WHEEL") - if not wheel_env: - logger.debug("PLUGIN_WHEEL is not set; plug-in loading is a no-op.") - return - + wheel = self._spec.wheel accept_load_failures = self._resolve_accept_load_failures() try: - await self._load_plugin_async(wheel_path=Path(wheel_env).expanduser()) + await self._load_plugin_async(wheel_path=wheel.expanduser()) except Exception as exc: if accept_load_failures: logger.warning( - "Plug-in from PLUGIN_WHEEL='%s' failed to load; load failures are accepted so " - "continuing without it: %s", - wheel_env, + "Plug-in wheel '%s' failed to load; load failures are accepted so continuing without it: %s", + wheel, exc, ) return - message = f"Failed to load plug-in from PLUGIN_WHEEL='{wheel_env}': {exc} {_REMEDIATION}" + message = f"Failed to load plug-in from wheel '{wheel}': {exc} {_REMEDIATION}" # Preserve the specific failure type so callers can distinguish modes, while # always surfacing the remediation guidance. Unknown errors (e.g. a raising # bootstrap or an unsafe archive) become a plain PluginLoadError. @@ -188,14 +268,16 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: PluginRegisteredNothingError: If the plug-in registered no datasets or scenarios. """ if not wheel_path.is_file(): - raise PluginWheelNotFoundError(f"PLUGIN_WHEEL does not point to an existing file: {wheel_path}") + raise PluginWheelNotFoundError(f"Plug-in wheel does not point to an existing file: {wheel_path}") if wheel_path.suffix != ".whl": - raise PluginWheelNotFoundError(f"PLUGIN_WHEEL must point to a .whl file, got: {wheel_path}") + raise PluginWheelNotFoundError(f"Plug-in wheel must be a .whl file, got: {wheel_path}") # Wheel extraction and directory scanning are blocking filesystem work; run them # off the event loop so init does not stall unrelated async tasks. extract_dir = await asyncio.to_thread(self._extract_wheel, wheel_path=wheel_path) - package_name = await asyncio.to_thread(self._resolve_package_name, extract_dir=extract_dir) + package_name = await asyncio.to_thread( + self._resolve_package_name, extract_dir=extract_dir, explicit_package=self._spec.package + ) from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider from pyrit.registry import ScenarioRegistry @@ -310,15 +392,16 @@ def _plugin_base_dir() -> Path: return Path(path.HOME_PATH, ".plugin").resolve() @staticmethod - def _resolve_package_name(*, extract_dir: Path) -> str: + def _resolve_package_name(*, extract_dir: Path, explicit_package: str | None) -> str: """ Determine the plug-in's top-level import package. - Resolution order: ``PLUGIN_PACKAGE`` env var, then ``*.dist-info/top_level.txt``, + Resolution order: the spec's ``package`` (when set), then ``*.dist-info/top_level.txt``, then the single importable top-level directory in the extraction. Args: extract_dir: The directory the wheel was extracted to. + explicit_package: The package name from the plug-in's config, if any. Returns: str: The top-level package name to import. @@ -326,9 +409,8 @@ def _resolve_package_name(*, extract_dir: Path) -> str: Raises: ValueError: If the package cannot be unambiguously determined. """ - explicit = os.getenv("PLUGIN_PACKAGE") - if explicit: - return explicit + if explicit_package: + return explicit_package for dist_info in sorted(extract_dir.glob("*.dist-info")): top_level = dist_info / "top_level.txt" @@ -351,10 +433,11 @@ def _resolve_package_name(*, extract_dir: Path) -> str: if not candidates: raise ValueError( f"Could not find an importable top-level package in {extract_dir}. " - "Set PLUGIN_PACKAGE to the plug-in's package name." + "Set the plug-in's 'package' in its .pyrit_conf entry." ) raise ValueError( - f"Found multiple top-level packages in {extract_dir}: {candidates}. Set PLUGIN_PACKAGE to disambiguate." + f"Found multiple top-level packages in {extract_dir}: {candidates}. " + "Set the plug-in's 'package' in its .pyrit_conf entry to disambiguate." ) @staticmethod @@ -387,7 +470,8 @@ def _verify_module_location(*, module: ModuleType, extract_dir: Path, package_na raise ValueError( f"Imported package '{package_name}' resolved to {locations[0]} which is outside the " f"plug-in extraction directory {extract_resolved}. An installed package with the same " - "name is likely shadowing the plug-in; set PLUGIN_PACKAGE or resolve the name conflict." + "name is likely shadowing the plug-in; set the plug-in's 'package' in .pyrit_conf or " + "resolve the name conflict." ) @staticmethod diff --git a/tests/integration/setup/test_plugin_loader_integration.py b/tests/integration/setup/test_plugin_loader_integration.py index d481bc1e88..e65d6d4ccb 100644 --- a/tests/integration/setup/test_plugin_loader_integration.py +++ b/tests/integration/setup/test_plugin_loader_integration.py @@ -39,7 +39,7 @@ from pyrit.registry import ScenarioRegistry from pyrit.registry.discovery import discover_in_directory from pyrit.scenario.core import Scenario -from pyrit.setup import IN_MEMORY, initialize_pyrit_async +from pyrit.setup import IN_MEMORY, PluginSpec, initialize_pyrit_async # (module stem, class name, registry name) for the scenarios the self-contained wheel ships. _MOCK_SCENARIOS = [ @@ -179,13 +179,11 @@ def _scenario_class_names_under_package(package_prefix: str) -> set[str]: @contextmanager -def _plugin_env(*, wheel: Path, plugin_dir: Path | None, extra: dict[str, str] | None = None) -> Iterator[None]: - """Set PLUGIN_* env for the duration of a load, then restore the prior values.""" - values: dict[str, str] = {"PLUGIN_WHEEL": str(wheel)} +def _plugin_dir_env(*, plugin_dir: Path | None) -> Iterator[None]: + """Set PLUGIN_DIR (extraction dir) for the duration of a load, then restore it.""" + values: dict[str, str] = {} if plugin_dir is not None: values["PLUGIN_DIR"] = str(plugin_dir) - if extra: - values.update(extra) saved: dict[str, str | None] = {key: os.environ.get(key) for key in values} os.environ.update(values) @@ -225,8 +223,8 @@ async def test_built_wheel_scenarios_are_discovered(tmp_path: Path, plugin_sandb registry = ScenarioRegistry.get_registry_singleton() before = set(registry.get_class_names()) - with _plugin_env(wheel=wheel, plugin_dir=tmp_path / ".plugin"): - await initialize_pyrit_async(IN_MEMORY) + with _plugin_dir_env(plugin_dir=tmp_path / ".plugin"): + await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel)]) registry = ScenarioRegistry.get_registry_singleton() after = set(registry.get_class_names()) @@ -261,12 +259,8 @@ async def test_injected_wheel_scenarios_are_discovered(plugin_sandbox: None) -> if not scenario_dirs_env and not package: pytest.skip("Set PLUGIN_TEST_SCENARIO_DIRS or PLUGIN_TEST_PACKAGE to define the expected scenarios.") - extra_env: dict[str, str] = {} - if package: - extra_env["PLUGIN_PACKAGE"] = package - - with _plugin_env(wheel=wheel, plugin_dir=None, extra=extra_env): - await initialize_pyrit_async(IN_MEMORY) + with _plugin_dir_env(plugin_dir=None): + await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel, package=package)]) found = _registered_scenario_class_names() diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 3c33fac5dc..804a67be6b 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -710,3 +710,82 @@ def test_server_url_non_string_raises(self): def test_server_non_dict_raises(self): with pytest.raises(ValueError, match="Server entry must be a dict"): ConfigurationLoader(server="http://oops:8000") # type: ignore[arg-type] + + +class TestConfigurationLoaderPlugins: + """Tests for the .pyrit_conf `plugins` list and `plugin_accept_load_failures`.""" + + def test_no_plugins_by_default(self): + """A config with no plugins resolves to an empty plug-in list.""" + config = ConfigurationLoader() + assert config.plugins == [] + assert config.resolve_plugins() == [] + assert config.plugin_accept_load_failures is None + + def test_bare_string_entry(self): + """A bare wheel-path string resolves to a spec with no explicit package.""" + from pyrit.setup.plugin_loader import PluginSpec + + config = ConfigurationLoader(plugins=["/abs/path/plugin.whl"]) + assert config.resolve_plugins() == [PluginSpec(wheel=pathlib.Path("/abs/path/plugin.whl"))] + + def test_concise_package_wheel_pair(self): + """A {package: wheel} mapping names the package.""" + from pyrit.setup.plugin_loader import PluginSpec + + config = ConfigurationLoader(plugins=[{"my_pkg": "/abs/path/plugin.whl"}]) + assert config.resolve_plugins() == [PluginSpec(wheel=pathlib.Path("/abs/path/plugin.whl"), package="my_pkg")] + + def test_explicit_mapping_entry(self): + """An explicit {wheel, package} mapping resolves both fields.""" + from pyrit.setup.plugin_loader import PluginSpec + + config = ConfigurationLoader(plugins=[{"wheel": "/abs/path/plugin.whl", "package": "my_pkg"}]) + assert config.resolve_plugins() == [PluginSpec(wheel=pathlib.Path("/abs/path/plugin.whl"), package="my_pkg")] + + def test_multiple_plugins_preserve_order(self): + """Several plug-ins resolve in configured order.""" + config = ConfigurationLoader(plugins=["/a/first.whl", {"second_pkg": "/b/second.whl"}]) + specs = config.resolve_plugins() + assert [s.wheel.name for s in specs] == ["first.whl", "second.whl"] + assert specs[1].package == "second_pkg" + + def test_malformed_plugin_entry_fails_fast(self): + """A malformed plug-in entry raises during construction (before other setup).""" + with pytest.raises(ValueError): + ConfigurationLoader(plugins=[123]) # type: ignore[list-item] + + def test_from_dict_with_plugins(self): + """from_dict wires the plugins list and accept-load-failures flag onto the loader.""" + config = ConfigurationLoader.from_dict( + {"plugins": [{"pkg": "/x/plugin.whl"}], "plugin_accept_load_failures": True} + ) + assert config.plugin_accept_load_failures is True + assert config.resolve_plugins()[0].package == "pkg" + + async def test_initialize_forwards_plugins_and_flag(self): + """ConfigurationLoader.initialize_pyrit_async forwards resolved plugins and the flag to core.""" + from pyrit.setup.plugin_loader import PluginSpec + + config = ConfigurationLoader( + memory_db_type="in_memory", + plugins=[{"pkg": "/x/plugin.whl"}], + plugin_accept_load_failures=True, + ) + + with mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") as core_init: + await config.initialize_pyrit_async() + + _, kwargs = core_init.call_args + assert kwargs["plugins"] == [PluginSpec(wheel=pathlib.Path("/x/plugin.whl"), package="pkg")] + assert kwargs["plugin_accept_load_failures"] is True + + async def test_initialize_passes_none_plugins_when_empty(self): + """With no plug-ins configured, core is called with plugins=None (no-op path).""" + config = ConfigurationLoader(memory_db_type="in_memory") + + with mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") as core_init: + await config.initialize_pyrit_async() + + _, kwargs = core_init.call_args + assert kwargs["plugins"] is None diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index 49565cc17b..4ed3f9b535 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -27,14 +27,15 @@ from pyrit.memory import CentralMemory from pyrit.models import SeedDataset from pyrit.registry import ScenarioRegistry -from pyrit.setup.initialization import IN_MEMORY, initialize_pyrit_async +from pyrit.setup.initialization import IN_MEMORY, _verify_plugin_prerequisites, initialize_pyrit_async from pyrit.setup.plugin_loader import ( PluginImportError, PluginLoader, PluginLoadError, PluginRegisteredNothingError, + PluginSpec, PluginWheelNotFoundError, - load_plugin_if_configured_async, + load_plugins_if_configured_async, ) # --------------------------------------------------------------------------- @@ -285,26 +286,35 @@ def plugin_sandbox() -> Iterator[None]: def plugin_env(**overrides: str) -> Iterator[None]: """Patch os.environ so only the given PLUGIN_* overrides are present.""" with patch.dict(os.environ, overrides, clear=False): - for key in ("PLUGIN_WHEEL", "PLUGIN_DIR", "PLUGIN_PACKAGE", "PLUGIN_ACCEPT_LOAD_FAILURES"): + for key in ("PLUGIN_DIR", "PLUGIN_ACCEPT_LOAD_FAILURES"): if key not in overrides: os.environ.pop(key, None) yield +def _spec(wheel: MockWheel, *, package: str | None = None) -> PluginSpec: + """Build a PluginSpec from a mock wheel.""" + return PluginSpec(wheel=wheel.path, package=package) + + +def _dummy_loader(*, accept_load_failures: bool | None = None) -> PluginLoader: + """Build a PluginLoader with a placeholder spec for testing policy resolution.""" + return PluginLoader(spec=PluginSpec(wheel=Path("dummy.whl")), accept_load_failures=accept_load_failures) + + async def load_plugin( wheel: MockWheel, plugin_dir: Path, *, accept_load_failures: bool | None = None, + package: str | None = None, extra_env: dict[str, str] | None = None, ) -> None: - """Run the plug-in loader against a mock wheel with an isolated env.""" - env = {"PLUGIN_WHEEL": str(wheel.path), "PLUGIN_DIR": str(plugin_dir)} - if extra_env: - env.update(extra_env) - - with plugin_env(**env): - await load_plugin_if_configured_async(accept_load_failures=accept_load_failures) + """Run the plug-in loader against a single mock wheel with an isolated env.""" + with plugin_env(PLUGIN_DIR=str(plugin_dir), **(extra_env or {})): + await load_plugins_if_configured_async( + plugins=[_spec(wheel, package=package)], accept_load_failures=accept_load_failures + ) # --------------------------------------------------------------------------- @@ -312,36 +322,73 @@ async def load_plugin( # --------------------------------------------------------------------------- +def test_verify_plugin_prerequisites_passes_when_memory_set() -> None: + """The health check passes when central memory is initialized.""" + with patch.object(CentralMemory, "get_memory_instance", return_value=MagicMock()): + _verify_plugin_prerequisites() # must not raise + + +def test_verify_plugin_prerequisites_raises_when_memory_unset() -> None: + """The health check fails loudly when a prerequisite phase (memory) did not complete.""" + with patch.object(CentralMemory, "get_memory_instance", side_effect=ValueError("not set")): + with pytest.raises(RuntimeError, match="central memory is not initialized"): + _verify_plugin_prerequisites() + + async def test_plugin_phase_runs_after_memory_before_initializers() -> None: - """initialize_pyrit_async loads the plug-in after memory is set, before initializers.""" + """initialize_pyrit_async loads plug-ins after memory is set, before initializers.""" manager = MagicMock() manager.attach_mock(MagicMock(), "set_memory") - manager.attach_mock(AsyncMock(), "load_plugin") + manager.attach_mock(AsyncMock(), "load_plugins") manager.attach_mock(AsyncMock(), "execute") with ( patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), patch.object(CentralMemory, "set_memory_instance", manager.set_memory), - patch("pyrit.setup.plugin_loader.load_plugin_if_configured_async", manager.load_plugin), + patch("pyrit.setup.initialization._verify_plugin_prerequisites"), + patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", manager.load_plugins), patch("pyrit.setup.initialization._execute_initializers_async", manager.execute), ): - await initialize_pyrit_async(IN_MEMORY, initializers=[MagicMock()], env_files=[], silent=True) + await initialize_pyrit_async( + IN_MEMORY, + initializers=[MagicMock()], + env_files=[], + silent=True, + plugins=[PluginSpec(wheel=Path("plugin.whl"))], + ) - order = [call[0] for call in manager.mock_calls if call[0] in {"set_memory", "load_plugin", "execute"}] - assert order.index("set_memory") < order.index("load_plugin") < order.index("execute") + order = [call[0] for call in manager.mock_calls if call[0] in {"set_memory", "load_plugins", "execute"}] + assert order.index("set_memory") < order.index("load_plugins") < order.index("execute") async def test_plugin_phase_forwards_accept_load_failures_param() -> None: - """initialize_pyrit_async forwards plugin_accept_load_failures to the loader.""" - load_plugin_mock = AsyncMock() + """initialize_pyrit_async forwards plugins and plugin_accept_load_failures to the loader.""" + load_plugins_mock = AsyncMock() + spec = PluginSpec(wheel=Path("plugin.whl")) + with ( + patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), + patch.object(CentralMemory, "set_memory_instance"), + patch("pyrit.setup.initialization._verify_plugin_prerequisites"), + patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", load_plugins_mock), + ): + await initialize_pyrit_async( + IN_MEMORY, env_files=[], silent=True, plugins=[spec], plugin_accept_load_failures=True + ) + + load_plugins_mock.assert_awaited_once_with(plugins=[spec], accept_load_failures=True) + + +async def test_plugin_phase_skipped_when_no_plugins() -> None: + """initialize_pyrit_async does not touch the plug-in loader when no plug-ins are configured.""" + load_plugins_mock = AsyncMock() with ( patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), patch.object(CentralMemory, "set_memory_instance"), - patch("pyrit.setup.plugin_loader.load_plugin_if_configured_async", load_plugin_mock), + patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", load_plugins_mock), ): - await initialize_pyrit_async(IN_MEMORY, env_files=[], silent=True, plugin_accept_load_failures=True) + await initialize_pyrit_async(IN_MEMORY, env_files=[], silent=True) - load_plugin_mock.assert_awaited_once_with(accept_load_failures=True) + load_plugins_mock.assert_not_awaited() # --------------------------------------------------------------------------- @@ -349,13 +396,13 @@ async def test_plugin_phase_forwards_accept_load_failures_param() -> None: # --------------------------------------------------------------------------- -async def test_no_op_when_plugin_wheel_unset() -> None: - """With no PLUGIN_WHEEL the loader does nothing and registers nothing.""" +async def test_no_op_when_no_plugins() -> None: + """With an empty plug-in list the loader does nothing and registers nothing.""" providers_before = dict(SeedDatasetProvider.get_all_providers()) path_before = list(sys.path) with plugin_env(): - await load_plugin_if_configured_async() + await load_plugins_if_configured_async(plugins=[]) assert SeedDatasetProvider.get_all_providers() == providers_before assert sys.path == path_before @@ -485,10 +532,9 @@ async def test_shadowing_installed_package_is_rejected(tmp_path: Path) -> None: """An installed package of the same name shadowing the plug-in fails loudly.""" wheel = build_mock_wheel(tmp_path) - # PLUGIN_PACKAGE points at a stdlib package that imports from outside the extraction dir. - with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(tmp_path / ".plugin"), PLUGIN_PACKAGE="json"): - with pytest.raises(PluginImportError, match="shadowing"): - await load_plugin_if_configured_async() + # The spec's package points at a stdlib package that imports from outside the extraction dir. + with pytest.raises(PluginImportError, match="shadowing"): + await load_plugin(wheel, tmp_path / ".plugin", package="json") # --------------------------------------------------------------------------- @@ -557,9 +603,8 @@ async def test_failed_load_rolls_back_syspath(tmp_path: Path) -> None: wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) plugin_dir = tmp_path / ".plugin" - with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(plugin_dir)): - with pytest.raises(PluginLoadError): - await load_plugin_if_configured_async() + with pytest.raises(PluginLoadError): + await load_plugin(wheel, plugin_dir) extract_dir = str(plugin_dir / wheel.path.stem) assert extract_dir not in sys.path @@ -570,9 +615,8 @@ async def test_failing_bootstrap_rolls_back_partial_registration(tmp_path: Path) wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") plugin_dir = tmp_path / ".plugin" - with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(plugin_dir)): - with pytest.raises(PluginLoadError): - await load_plugin_if_configured_async() + with pytest.raises(PluginLoadError): + await load_plugin(wheel, plugin_dir) registry = ScenarioRegistry.get_registry_singleton() assert wheel.scenario_name not in registry._classes @@ -609,9 +653,8 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: SeedDatasetProvider._registry["MockProvider"] = _PreexistingProvider - with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(tmp_path / ".plugin")): - with pytest.raises(PluginLoadError): - await load_plugin_if_configured_async() + with pytest.raises(PluginLoadError): + await load_plugin(wheel, tmp_path / ".plugin") # The original provider is restored, not deleted or left replaced by the plug-in's. assert SeedDatasetProvider._registry["MockProvider"] is _PreexistingProvider @@ -629,9 +672,8 @@ class _PreexistingScenario(RapidResponse): registry = ScenarioRegistry.get_registry_singleton() registry.register_class(_PreexistingScenario, name=wheel.scenario_name) - with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(tmp_path / ".plugin")): - with pytest.raises(PluginLoadError): - await load_plugin_if_configured_async() + with pytest.raises(PluginLoadError): + await load_plugin(wheel, tmp_path / ".plugin") # The original scenario is restored, not deleted or left replaced by the plug-in's. assert registry._classes[wheel.scenario_name] is _PreexistingScenario @@ -648,7 +690,7 @@ def test_extract_wheel_reuses_cached_extraction(tmp_path: Path) -> None: plugin_dir = tmp_path / ".plugin" with plugin_env(PLUGIN_DIR=str(plugin_dir)): - initializer = PluginLoader() + initializer = PluginLoader(spec=_spec(wheel)) first = initializer._extract_wheel(wheel_path=wheel.path) marker = first / "cache_marker.txt" marker.write_text("kept", encoding="utf-8") @@ -664,23 +706,21 @@ def test_extract_wheel_reuses_cached_extraction(tmp_path: Path) -> None: # --------------------------------------------------------------------------- -def test_resolve_package_name_prefers_env(tmp_path: Path) -> None: - """PLUGIN_PACKAGE takes precedence over inference.""" +def test_resolve_package_name_prefers_explicit(tmp_path: Path) -> None: + """An explicit package name takes precedence over inference.""" (tmp_path / "some_pkg").mkdir() (tmp_path / "some_pkg" / "__init__.py").write_text("", encoding="utf-8") - with plugin_env(PLUGIN_PACKAGE="explicit_pkg"): - assert PluginLoader._resolve_package_name(extract_dir=tmp_path) == "explicit_pkg" + assert PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package="explicit_pkg") == "explicit_pkg" def test_resolve_package_name_infers_single_package(tmp_path: Path) -> None: - """The single importable top-level directory is inferred when no env/top_level.txt exists.""" + """The single importable top-level directory is inferred when no package/top_level.txt exists.""" (tmp_path / "the_pkg").mkdir() (tmp_path / "the_pkg" / "__init__.py").write_text("", encoding="utf-8") (tmp_path / "the_pkg-0.0.1.dist-info").mkdir() - with plugin_env(): - assert PluginLoader._resolve_package_name(extract_dir=tmp_path) == "the_pkg" + assert PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) == "the_pkg" def test_resolve_package_name_uses_top_level_txt(tmp_path: Path) -> None: @@ -693,24 +733,23 @@ def test_resolve_package_name_uses_top_level_txt(tmp_path: Path) -> None: distinfo.mkdir() (distinfo / "top_level.txt").write_text("pkg_b\n", encoding="utf-8") - with plugin_env(): - assert PluginLoader._resolve_package_name(extract_dir=tmp_path) == "pkg_b" + assert PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) == "pkg_b" def test_resolve_package_name_none_raises(tmp_path: Path) -> None: - """No importable package raises a clear error pointing at PLUGIN_PACKAGE.""" - with plugin_env(), pytest.raises(ValueError, match="PLUGIN_PACKAGE"): - PluginLoader._resolve_package_name(extract_dir=tmp_path) + """No importable package raises a clear error pointing at the plug-in's config.""" + with pytest.raises(ValueError, match="package"): + PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) def test_resolve_package_name_multiple_raises(tmp_path: Path) -> None: - """Multiple top-level packages require PLUGIN_PACKAGE to disambiguate.""" + """Multiple top-level packages require an explicit package to disambiguate.""" for name in ("pkg_a", "pkg_b"): (tmp_path / name).mkdir() (tmp_path / name / "__init__.py").write_text("", encoding="utf-8") - with plugin_env(), pytest.raises(ValueError, match="disambiguate"): - PluginLoader._resolve_package_name(extract_dir=tmp_path) + with pytest.raises(ValueError, match="disambiguate"): + PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) # --------------------------------------------------------------------------- @@ -720,30 +759,31 @@ def test_resolve_package_name_multiple_raises(tmp_path: Path) -> None: async def test_missing_wheel_fails_closed() -> None: """A configured-but-missing wheel raises PluginWheelNotFoundError by default (fail-closed).""" - with plugin_env(PLUGIN_WHEEL=str(Path("does_not_exist.whl"))): + with plugin_env(): with pytest.raises(PluginWheelNotFoundError, match="Failed to load plug-in"): - await load_plugin_if_configured_async() + await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=Path("does_not_exist.whl"))]) async def test_missing_wheel_accept_load_failures_param_proceeds() -> None: """Accepting load failures via the explicit param skips a broken plug-in with a warning.""" - with plugin_env(PLUGIN_WHEEL=str(Path("does_not_exist.whl"))): - await load_plugin_if_configured_async(accept_load_failures=True) # must not raise + with plugin_env(): + await load_plugins_if_configured_async( + plugins=[PluginSpec(wheel=Path("does_not_exist.whl"))], accept_load_failures=True + ) # must not raise async def test_missing_wheel_accept_load_failures_env_proceeds() -> None: """Accepting load failures via PLUGIN_ACCEPT_LOAD_FAILURES env skips a broken plug-in with a warning.""" - with plugin_env(PLUGIN_WHEEL=str(Path("does_not_exist.whl")), PLUGIN_ACCEPT_LOAD_FAILURES="true"): - await load_plugin_if_configured_async() # must not raise + with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): + await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=Path("does_not_exist.whl"))]) # must not raise async def test_empty_wheel_is_loud(tmp_path: Path) -> None: """A wheel that imports cleanly but registers nothing fails loudly.""" wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) - with plugin_env(PLUGIN_WHEEL=str(wheel.path), PLUGIN_DIR=str(tmp_path / ".plugin")): - with pytest.raises(PluginRegisteredNothingError, match="registered no datasets or scenarios"): - await load_plugin_if_configured_async() + with pytest.raises(PluginRegisteredNothingError, match="registered no datasets or scenarios"): + await load_plugin(wheel, tmp_path / ".plugin") async def test_empty_wheel_accept_load_failures_proceeds(tmp_path: Path) -> None: @@ -754,13 +794,13 @@ async def test_empty_wheel_accept_load_failures_proceeds(tmp_path: Path) -> None async def test_non_whl_path_fails_closed(tmp_path: Path) -> None: - """PLUGIN_WHEEL that is not a .whl file fails closed with PluginWheelNotFoundError.""" + """A wheel path that is not a .whl file fails closed with PluginWheelNotFoundError.""" not_a_wheel = tmp_path / "plugin.zip" not_a_wheel.write_text("not a wheel", encoding="utf-8") - with plugin_env(PLUGIN_WHEEL=str(not_a_wheel)): + with plugin_env(): with pytest.raises(PluginWheelNotFoundError, match="Failed to load plug-in"): - await load_plugin_if_configured_async() + await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=not_a_wheel)]) def test_error_subclasses_are_plugin_load_errors() -> None: @@ -776,9 +816,9 @@ async def test_wheel_with_path_traversal_member_fails_closed(tmp_path: Path) -> archive.writestr("evil_pkg/__init__.py", "") archive.writestr("../escape.py", "compromised = True") - with plugin_env(PLUGIN_WHEEL=str(malicious), PLUGIN_DIR=str(tmp_path / ".plugin")): + with plugin_env(PLUGIN_DIR=str(tmp_path / ".plugin")): with pytest.raises(PluginLoadError, match="Failed to load plug-in"): - await load_plugin_if_configured_async() + await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=malicious)]) # The traversal target was not written outside the extraction directory. assert not (tmp_path / "escape.py").exists() @@ -819,28 +859,102 @@ def __init__(self, *, required_value: str) -> None: def test_resolve_accept_load_failures_from_env_tokens(value: str, expected: bool) -> None: """accept_load_failures resolves from PLUGIN_ACCEPT_LOAD_FAILURES across truthy/falsey tokens.""" with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES=value): - assert PluginLoader()._resolve_accept_load_failures() is expected + assert _dummy_loader()._resolve_accept_load_failures() is expected def test_resolve_accept_load_failures_explicit_true() -> None: """An explicit accept_load_failures=True resolves to True.""" with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="false"): - assert PluginLoader(accept_load_failures=True)._resolve_accept_load_failures() is True + assert _dummy_loader(accept_load_failures=True)._resolve_accept_load_failures() is True def test_resolve_accept_load_failures_explicit_overrides_env() -> None: """An explicit accept_load_failures value takes precedence over the env var.""" with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): - assert PluginLoader(accept_load_failures=False)._resolve_accept_load_failures() is False + assert _dummy_loader(accept_load_failures=False)._resolve_accept_load_failures() is False def test_resolve_accept_load_failures_from_env_when_no_explicit() -> None: """accept_load_failures falls back to PLUGIN_ACCEPT_LOAD_FAILURES when no explicit value is set.""" with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): - assert PluginLoader()._resolve_accept_load_failures() is True + assert _dummy_loader()._resolve_accept_load_failures() is True def test_resolve_accept_load_failures_defaults_false() -> None: """accept_load_failures defaults to False (fail-closed).""" with plugin_env(): - assert PluginLoader()._resolve_accept_load_failures() is False + assert _dummy_loader()._resolve_accept_load_failures() is False + + +# --------------------------------------------------------------------------- +# PluginSpec.from_config parsing +# --------------------------------------------------------------------------- + + +def test_plugin_spec_from_config_bare_string() -> None: + """A bare wheel-path string parses to a spec with no explicit package.""" + spec = PluginSpec.from_config("/abs/path/plugin.whl") + assert spec.wheel == Path("/abs/path/plugin.whl") + assert spec.package is None + + +def test_plugin_spec_from_config_single_key_pair() -> None: + """A {package: wheel} mapping names the package.""" + spec = PluginSpec.from_config({"my_pkg": "/abs/path/plugin.whl"}) + assert spec.wheel == Path("/abs/path/plugin.whl") + assert spec.package == "my_pkg" + + +def test_plugin_spec_from_config_explicit_mapping() -> None: + """An explicit {wheel, package} mapping parses both fields; package is optional.""" + spec = PluginSpec.from_config({"wheel": "/abs/path/plugin.whl", "package": "my_pkg"}) + assert spec == PluginSpec(wheel=Path("/abs/path/plugin.whl"), package="my_pkg") + assert PluginSpec.from_config({"wheel": "/abs/path/plugin.whl"}).package is None + + +@pytest.mark.parametrize( + "entry", [123, {"a": "1", "b": "2"}, {"package": "no_wheel"}, {"wheel": "/x.whl", "packge": "typo"}] +) +def test_plugin_spec_from_config_rejects_bad_shapes(entry: object) -> None: + """Unsupported entry shapes (including explicit mappings with unexpected keys) raise ValueError.""" + with pytest.raises(ValueError): + PluginSpec.from_config(entry) # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- +# Multiple plug-ins +# --------------------------------------------------------------------------- + + +async def test_multiple_plugins_all_load(tmp_path: Path) -> None: + """Every configured plug-in loads; a single plug-in is just the one-element case.""" + wheel_a = build_mock_wheel(tmp_path / "a", package_name="mock_plugin_alpha") + wheel_b = build_mock_wheel(tmp_path / "b", package_name="mock_plugin_beta") + + with plugin_env(PLUGIN_DIR=str(tmp_path / ".plugin")): + await load_plugins_if_configured_async(plugins=[_spec(wheel_a), _spec(wheel_b)]) + + names = ScenarioRegistry.get_registry_singleton().get_class_names() + assert wheel_a.scenario_name in names + assert wheel_b.scenario_name in names + + +async def test_plugins_load_in_configured_order(tmp_path: Path) -> None: + """Plug-ins load in list order (later entries load after earlier ones).""" + wheel_a = build_mock_wheel(tmp_path / "a", package_name="mock_plugin_first") + wheel_b = build_mock_wheel(tmp_path / "b", package_name="mock_plugin_second") + + loaded: list[str] = [] + original = PluginLoader.load_async + + async def _record(self: PluginLoader) -> None: + loaded.append(self._spec.package or self._spec.wheel.stem) + await original(self) + + with plugin_env(PLUGIN_DIR=str(tmp_path / ".plugin")): + with patch.object(PluginLoader, "load_async", _record): + await load_plugins_if_configured_async( + plugins=[_spec(wheel_a, package="mock_plugin_first"), _spec(wheel_b, package="mock_plugin_second")] + ) + + assert loaded == ["mock_plugin_first", "mock_plugin_second"] From 88bd76a81a87c03ed282360cbf1256f49b2fb10e Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 9 Jul 2026 12:20:28 -0700 Subject: [PATCH 11/34] feat(registry): auto-register external plug-in Scenario subclasses Add ScenarioRegistry.register_external_subclasses to register a plug-in package's concrete Scenario subclasses by type (without triggering built-in discovery), refactor the shared _register_subclasses_in_package helper out of _discover, and add an overridable _external_registry_name hook. ScenarioRegistry keys plug-in scenarios by their module path relative to the plug-in package, mirroring built-in dotted names so image/text classes with the same leaf name do not collide. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../registry/components/scenario_registry.py | 29 +++++++ pyrit/registry/registry.py | 81 ++++++++++++++++++- 2 files changed, 107 insertions(+), 3 deletions(-) diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py index f22686e755..af4284506d 100644 --- a/pyrit/registry/components/scenario_registry.py +++ b/pyrit/registry/components/scenario_registry.py @@ -109,6 +109,35 @@ def _get_registry_name(self, cls: type[Scenario]) -> str: return relative return class_name_to_snake_case(cls.__name__, suffix="Scenario") + def _external_registry_name(self, cls: type[Scenario], *, package_name: str) -> str: + """ + Key a plug-in scenario by its module path within the plug-in package. + + Mirrors the built-in dotted-name scheme (module path relative to the scenarios + package) so plug-in scenarios read like built-ins and same-named scenarios in + different submodules do not collide on a bare class name. The name is the class's + module relative to the plug-in's top-level ``package_name`` with a leading + ``scenario.scenarios.`` / ``scenarios.`` segment stripped to match built-in style; + it falls back to the suffix-stripped snake_case class name when no module context + is available. + + Args: + cls (type[Scenario]): The plug-in scenario class. + package_name (str): The plug-in's top-level package name. + + Returns: + str: The dotted registry name for the plug-in scenario. + """ + module = cls.__module__ or "" + prefix = f"{package_name}." + relative = module[len(prefix) :] if module.startswith(prefix) else module + for marker in ("scenario.scenarios.", "scenarios."): + index = relative.find(marker) + if index != -1: + relative = relative[index + len(marker) :] + break + return relative or class_name_to_snake_case(cls.__name__, suffix="Scenario") + def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: """ Build metadata for a Scenario class. diff --git a/pyrit/registry/registry.py b/pyrit/registry/registry.py index e159016a0b..578c815334 100644 --- a/pyrit/registry/registry.py +++ b/pyrit/registry/registry.py @@ -272,9 +272,7 @@ def _discover(self) -> None: ``_base_type``/``_discovery_package``. """ package = self._discovery_package() - base = self._base_type() package_name = package.__name__ - package_prefix = f"{package_name}." # Materialize lazily-exported classes so they are loaded before enumeration. # A lazy import backed by an optional dependency may fail; skip it rather # than fail the whole discovery (the class cannot be built without the dep). @@ -283,14 +281,72 @@ def _discover(self) -> None: getattr(package, exported_name) except Exception as exc: logger.debug(f"Skipping lazily-exported '{exported_name}': {exc}") + self._register_subclasses_in_package(package_name=package_name) + + def register_external_subclasses(self, *, package_name: str) -> int: + """ + Register concrete base-type subclasses that live under an external package. + + This is how a loaded plug-in's components (e.g. scenarios shipped in a plug-in + wheel) enter the registry: the plug-in's modules are imported elsewhere, then + this enumerates the base type's in-memory subclasses under ``package_name`` and + registers each, alongside the built-in discovery package. Built-in discovery is + **not** triggered, so lazy discovery is preserved, and classes already registered + (e.g. by the plug-in's own bootstrap) are left untouched. + + Args: + package_name (str): The external (plug-in) top-level package name. + + Returns: + int: The number of classes newly registered from the package. + """ + return self._register_subclasses_in_package( + package_name=package_name, skip_registered_classes=True, external_package=package_name + ) + + def _register_subclasses_in_package( + self, *, package_name: str, skip_registered_classes: bool = False, external_package: str | None = None + ) -> int: + """ + Register every concrete base-type subclass whose module lives under a package. + + Shared by built-in discovery (``_discover``) and external plug-in + registration (``register_external_subclasses``). Enumerates the in-memory + subclasses of ``_base_type()``, keeps those whose module is ``package_name`` or a + submodule of it, and registers each. Built-in classes are named by + ``_get_registry_name``; when ``external_package`` is set, names come from + ``_external_registry_name`` so a plug-in's classes are keyed by their module path + within the plug-in (mirroring built-ins), which keeps same-named classes in + different submodules from colliding. When ``skip_registered_classes`` is set, + classes already in the catalog (by identity) are skipped so an explicit bootstrap + registration is never shadowed by a fallback-named duplicate. + + Args: + package_name (str): The package whose subclasses to register. + skip_registered_classes (bool): Skip classes already registered by identity. + external_package (str | None): When set, the plug-in package the classes belong + to; names are derived relative to it via ``_external_registry_name``. + + Returns: + int: The number of classes newly registered. + """ + base = self._base_type() + package_prefix = f"{package_name}." + already_registered = set(self._classes.values()) if skip_registered_classes else set() + count = 0 for cls in self._iter_concrete_subclasses(base): module = cls.__module__ or "" if module != package_name and not module.startswith(package_prefix): continue + if skip_registered_classes and cls in already_registered: + continue if (cls.__doc__ or "").strip().startswith("Deprecated alias"): logger.debug(f"Skipping deprecated alias: {cls.__name__}") continue - name = self._get_registry_name(cls) + if external_package is not None: + name = self._external_registry_name(cls, package_name=external_package) + else: + name = self._get_registry_name(cls) existing = self._classes.get(name) if existing is not None and existing is not cls: logger.warning( @@ -299,7 +355,26 @@ def _discover(self) -> None: ) continue self.register_class(cls, name=name) + count += 1 logger.debug(f"Registered {base.__name__} class: {name} ({cls.__name__})") + return count + + def _external_registry_name(self, cls: type[T], *, package_name: str) -> str: + """ + Return the catalog name for a class discovered in an external (plug-in) package. + + Defaults to the same name built-in discovery would use. Registries that key + built-ins by module path (e.g. scenarios) override this so plug-in classes get an + equally path-based, collision-resistant name relative to the plug-in package. + + Args: + cls (type[T]): The external class being registered. + package_name (str): The plug-in's top-level package name. + + Returns: + str: The catalog name to register the class under. + """ + return self._get_registry_name(cls) @staticmethod def _iter_concrete_subclasses(base: type[T]) -> list[type[T]]: From b177f49c9160d7bd771fca21018b1d9d578d3fd7 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 9 Jul 2026 12:20:42 -0700 Subject: [PATCH 12/34] feat(setup): tolerate plug-in version drift and re-extract stale wheels Add pyrit/setup/plugin_compat.py with best-effort drift shims: warn_on_version_drift (reads the plug-in's Requires-Dist: pyrit pin and warns loudly on a minor mismatch) and bridge_scenario_extension_points (makes a plug-in Scenario concrete by bridging a renamed async extension point when the class is abstract solely due to that rename). Wire both into PluginLoader._load_plugin_async, and after bootstrap call ScenarioRegistry.register_external_subclasses so scenario-only plug-ins register without an explicit bootstrap. Add a size+mtime extraction-cache fingerprint (.plugin_wheel_fingerprint) so a rebuilt wheel with an unchanged filename re-extracts instead of reusing a stale extraction. Warnings are loud and greppable (PLUGIN VERSION DRIFT / PLUGIN COMPAT SHIM / PLUGIN CACHE STALE); irreconcilable drift still fails closed. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/setup/plugin_compat.py | 255 +++++++++++++++++++++++++++++++++++ pyrit/setup/plugin_loader.py | 97 +++++++++++-- 2 files changed, 343 insertions(+), 9 deletions(-) create mode 100644 pyrit/setup/plugin_compat.py diff --git a/pyrit/setup/plugin_compat.py b/pyrit/setup/plugin_compat.py new file mode 100644 index 0000000000..e418a8b45b --- /dev/null +++ b/pyrit/setup/plugin_compat.py @@ -0,0 +1,255 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Best-effort compatibility shims for plug-ins built against a slightly different PyRIT. + +A plug-in wheel is authored against one PyRIT version, then loaded into whatever PyRIT the +operator is running. When the host API has drifted, small mechanical differences — most +commonly a renamed extension-point method — can leave the plug-in's scenarios abstract and +therefore undiscoverable, even though the plug-in's own logic is unchanged. Rather than +force every plug-in author to re-release on each host bump, the loader applies narrow, loud +heuristics that bridge known mechanical renames so minor drift does not block loading. + +The shims are deliberately conservative. A scenario is bridged only when it is abstract +*solely* because of a recognized rename and a usable predecessor method is present; a class +abstract for any other reason is left alone (and fails loudly downstream). Every bridge and +every detected version mismatch logs a loud, greppable warning so the drift is visible and +the operator knows that rebuilding the plug-in against the running PyRIT removes the shim. +This module owns drift-bridging so the loader stays focused on extract/import/register. +""" + +from __future__ import annotations + +import inspect +import logging +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine + from pathlib import Path + from typing import Any + +logger = logging.getLogger(__name__) + +# Renamed scenario extension points: each maps a current (host) abstract method to the +# older predecessor a plug-in may still define. A predecessor took no ``context`` (it read +# ``self._*`` state the base still populates before the build call), so the bridge can +# delegate to it and ignore ``context``. Add an entry here when a future rename needs the +# same mechanical bridge. +_SCENARIO_METHOD_RENAMES: dict[str, str] = { + "_build_atomic_attacks_async": "_get_atomic_attacks_async", +} + + +def warn_on_version_drift(*, extract_dir: Path) -> str | None: + """ + Log a loud warning when a plug-in was built against a different PyRIT minor version. + + Reads the plug-in's ``Requires-Dist: pyrit`` pin from its wheel ``METADATA`` and + compares it to the running ``pyrit`` version. A mismatch is surfaced but never fatal: + the loader tolerates slight drift and lets the import/registration path (plus the + scenario shims below) decide whether the plug-in actually works. Absence of a pin, or + an unparseable one, is silently ignored — this is a best-effort signal, not a gate. + + Args: + extract_dir: The directory the plug-in wheel was extracted to. + + Returns: + str | None: The plug-in's declared PyRIT pin when a drift warning was emitted, + else ``None``. + """ + declared = _read_required_pyrit_version(extract_dir=extract_dir) + if declared is None: + return None + + import pyrit + + running = getattr(pyrit, "__version__", "") or "" + if _same_minor(declared, running): + return None + + logger.warning( + "PLUGIN VERSION DRIFT: plug-in was built against pyrit %s but pyrit %s is running. " + "Proceeding; compatibility shims will bridge known mechanical differences, but " + "rebuild the plug-in against the running pyrit if scenarios fail to load.", + declared, + running or "(unknown)", + ) + return declared + + +def bridge_scenario_extension_points(*, package_name: str) -> list[str]: + """ + Make a plug-in's scenarios concrete by bridging renamed extension-point methods. + + Enumerates concrete-intent ``Scenario`` subclasses owned by ``package_name`` that are + still abstract, and for each one whose only unimplemented abstract methods are known + renames with a usable predecessor, injects a thin adapter so the class satisfies the + current contract. Classes abstract for any other reason are left untouched. + + Args: + package_name: The plug-in's top-level package name. + + Returns: + list[str]: Human-readable descriptions of the bridges applied, for logging. + """ + from pyrit.scenario.core import Scenario + + prefix = f"{package_name}." + applied: list[str] = [] + + for cls in _iter_owned_subclasses(base=Scenario, package_name=package_name, prefix=prefix): + if not inspect.isabstract(cls): + continue + + missing = set(cls.__abstractmethods__) # type: ignore[ty:unresolved-attribute] + # Only bridge when every missing method is a known rename we can satisfy from a + # predecessor the class actually provides. Otherwise the class is abstract for a + # reason we must not paper over, so leave it (it will fail loudly downstream). + if not missing or not missing.issubset(_SCENARIO_METHOD_RENAMES.keys()): + continue + if not all(_has_concrete_method(cls, _SCENARIO_METHOD_RENAMES[name]) for name in missing): + continue + + for new_name in missing: + predecessor = _SCENARIO_METHOD_RENAMES[new_name] + setattr(cls, new_name, _make_rename_bridge(predecessor_name=predecessor)) + applied.append(f"{cls.__name__}.{predecessor} -> {new_name}") + logger.warning( + "PLUGIN COMPAT SHIM: bridged %s.%s to the renamed %s (plug-in built against " + "an older pyrit). Rebuild the plug-in against the running pyrit to remove this shim.", + cls.__name__, + predecessor, + new_name, + ) + cls.__abstractmethods__ = frozenset(cls.__abstractmethods__ - missing) # type: ignore[ty:unresolved-attribute] + + return applied + + +def _make_rename_bridge(*, predecessor_name: str) -> Callable[..., Coroutine[Any, Any, Any]]: + """ + Build an adapter that satisfies a renamed async extension point via its predecessor. + + Args: + predecessor_name: The older method name the plug-in still defines. + + Returns: + Callable[..., Coroutine[Any, Any, Any]]: An async adapter that ignores the new + ``context`` keyword and delegates to the predecessor (which reads ``self._*`` + state the base populates before the build call). + """ + + async def _bridged_build_async(self, **_kwargs: Any) -> Any: # noqa: ANN001 + return await getattr(self, predecessor_name)() + + _bridged_build_async.__name__ = predecessor_name + _bridged_build_async.__qualname__ = predecessor_name + return _bridged_build_async + + +def _has_concrete_method(cls: type, name: str) -> bool: + """ + Return whether ``cls`` provides a concrete (non-abstract) method of the given name. + + Args: + cls: The class to inspect. + name: The method name to look for. + + Returns: + bool: True when the method exists and is not itself abstract. + """ + method = getattr(cls, name, None) + return callable(method) and name not in getattr(cls, "__abstractmethods__", frozenset()) + + +def _iter_owned_subclasses(*, base: type, package_name: str, prefix: str) -> list[type]: + """ + Return every subclass of ``base`` whose module is owned by the plug-in package. + + Args: + base: The base class to enumerate subclasses of. + package_name: The plug-in's top-level package name. + prefix: ``f"{package_name}."`` (passed in to avoid recomputation). + + Returns: + list[type]: The owned subclasses currently loaded in memory. + """ + seen: set[int] = set() + owned: list[type] = [] + stack = list(base.__subclasses__()) + while stack: + cls = stack.pop() + if id(cls) in seen: + continue + seen.add(id(cls)) + stack.extend(cls.__subclasses__()) + module = cls.__module__ or "" + if module == package_name or module.startswith(prefix): + owned.append(cls) + return owned + + +def _read_required_pyrit_version(*, extract_dir: Path) -> str | None: + """ + Return the pinned ``pyrit`` version from the plug-in's wheel ``METADATA``, if any. + + Args: + extract_dir: The directory the plug-in wheel was extracted to. + + Returns: + str | None: The pinned version string (e.g. ``"0.14.0"``) or ``None`` when no + parseable ``Requires-Dist: pyrit`` pin is present. + """ + import re + + for metadata in extract_dir.glob("*.dist-info/METADATA"): + try: + text = metadata.read_text(encoding="utf-8", errors="replace") + except OSError: + continue + for line in text.splitlines(): + if not line.lower().startswith("requires-dist:"): + continue + if not re.match(r"requires-dist:\s*pyrit\b", line, flags=re.IGNORECASE): + continue + match = re.search(r"==\s*([0-9][0-9A-Za-z.\-]*)", line) + if match: + return match.group(1) + return None + + +def _same_minor(left: str, right: str) -> bool: + """ + Return whether two version strings share the same ``major.minor`` prefix. + + Args: + left: A version string (e.g. ``"0.14.0"``). + right: A version string (e.g. ``"0.15.0.dev0"``). + + Returns: + bool: True when both parse to the same ``(major, minor)``; False otherwise. An + unparseable side compares unequal so drift is surfaced rather than hidden. + """ + return _major_minor(left) == _major_minor(right) and _major_minor(left) is not None + + +def _major_minor(version: str) -> tuple[int, int] | None: + """ + Parse the leading ``major.minor`` from a version string. + + Args: + version: A version string. + + Returns: + tuple[int, int] | None: The ``(major, minor)`` pair, or ``None`` when the string + does not begin with two integer components. + """ + parts = version.split(".") + if len(parts) < 2: + return None + try: + return int(parts[0]), int(parts[1]) + except ValueError: + return None diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 44a3d30ff9..61866b97d0 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -9,7 +9,9 @@ ``zipfile`` — never ``pip``/``.venv``) into ``.plugin//``, prepends that directory to ``sys.path``, imports the package (so ``SeedDatasetProvider`` subclasses self-register), and runs the plug-in's bootstrap (a top-level ``register()`` callable or a shipped -``PyRITInitializer`` subclass) which registers the plug-in's scenarios. +``PyRITInitializer`` subclass). The plug-in's own ``Scenario`` subclasses are then +auto-registered by type (scoped to the plug-in package), so a plug-in's scenarios are +discovered without the bootstrap having to register each one explicitly. Plug-ins are declared in ``.pyrit_conf`` under the dedicated ``plugins`` key (a list, so several plug-ins load identically to one). ``load_plugins_if_configured_async`` is invoked @@ -210,6 +212,8 @@ class PluginLoader: cannot be loaded. """ + _FINGERPRINT_FILE = ".plugin_wheel_fingerprint" + def __init__(self, *, spec: PluginSpec, accept_load_failures: bool | None = None) -> None: """ Initialize the loader for one plug-in. @@ -279,6 +283,12 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: self._resolve_package_name, extract_dir=extract_dir, explicit_package=self._spec.package ) + from pyrit.setup.plugin_compat import bridge_scenario_extension_points, warn_on_version_drift + + # Surface any host/plug-in version drift before importing, so a subsequent import or + # registration failure is read in the light of the mismatch (tolerate slight drift). + warn_on_version_drift(extract_dir=extract_dir) + from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider from pyrit.registry import ScenarioRegistry @@ -300,8 +310,27 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: except Exception as exc: raise PluginImportError(f"Could not import plug-in package '{package_name}': {exc}") from exc + # Bridge known renamed extension points so scenarios built against an older + # PyRIT are made concrete (and thus discoverable) before bootstrap/registration. + bridged = bridge_scenario_extension_points(package_name=package_name) + if bridged: + logger.warning( + "Applied %d plug-in compatibility shim(s) for '%s': %s", + len(bridged), + package_name, + ", ".join(bridged), + ) + await self._run_bootstrap_async(package_name=package_name, module=module) + # Register the plug-in's own Scenario subclasses the same way built-ins are + # discovered (by type, scoped to the plug-in package), so plug-in scenarios are + # picked up without the bootstrap having to register each one explicitly. Classes + # the bootstrap already registered are left untouched. + registered_scenarios = scenario_registry.register_external_subclasses(package_name=package_name) + if registered_scenarios: + logger.info("Auto-registered %d plug-in scenario(s) from '%s'.", registered_scenarios, package_name) + provider_count, scenario_count = self._count_registered( package_name=package_name, scenario_registry=scenario_registry ) @@ -334,13 +363,19 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: def _extract_wheel(self, *, wheel_path: Path) -> Path: """ - Extract the wheel into ``.plugin//``, reusing a cached extraction. + Extract the wheel into ``.plugin//``, reusing a fresh cached extraction. - Extraction is atomic: the wheel is unpacked into a temporary sibling directory and - moved into place only on success, so a crash mid-extraction never leaves a partial - tree that would later be treated as a valid cache. ``safe_extract_zip`` validates - every member first (path traversal, symlinks, and size / entry-count / compression - caps) so a tampered wheel cannot escape the extraction directory or exhaust disk. + A cached extraction is reused only when its recorded fingerprint (wheel size and + modification time) matches the wheel on disk. A wheel rebuilt at the same path + (common during plug-in development, where the version-stamped filename is stable) + has a newer fingerprint, so the stale extraction is discarded and re-extracted + rather than silently reused — the exact silent-staleness trap the plug-in design + guards against. Extraction is atomic: the wheel is unpacked into a temporary + sibling directory and moved into place only on success, so a crash mid-extraction + never leaves a partial tree that would later be treated as a valid cache. + ``safe_extract_zip`` validates every member first (path traversal, symlinks, and + size / entry-count / compression caps) so a tampered wheel cannot escape the + extraction directory or exhaust disk. Args: wheel_path: Path to the plug-in wheel. @@ -354,9 +389,17 @@ def _extract_wheel(self, *, wheel_path: Path) -> Path: base_dir.mkdir(parents=True, exist_ok=True) extract_dir = base_dir / wheel_path.stem + fingerprint = self._wheel_fingerprint(wheel_path=wheel_path) if extract_dir.is_dir() and any(extract_dir.iterdir()): - logger.info("Reusing cached plug-in extraction at %s", extract_dir) - return extract_dir + if self._cached_fingerprint(extract_dir=extract_dir) == fingerprint: + logger.info("Reusing cached plug-in extraction at %s", extract_dir) + return extract_dir + logger.warning( + "PLUGIN CACHE STALE: extraction at %s does not match the current wheel " + "(rebuilt or replaced); re-extracting '%s'.", + extract_dir, + wheel_path.name, + ) # Unique per-extraction temp dir so concurrent loads of the same wheel in one # process cannot collide on a shared path (mkdtemp is atomic and 0700). safe_extract @@ -364,6 +407,7 @@ def _extract_wheel(self, *, wheel_path: Path) -> Path: tmp_dir = Path(tempfile.mkdtemp(prefix=f".{wheel_path.stem}.tmp-", dir=base_dir)) try: safe_extract_zip(source=wheel_path, dest_dir=tmp_dir) + (tmp_dir / self._FINGERPRINT_FILE).write_text(fingerprint, encoding="utf-8") if extract_dir.exists(): shutil.rmtree(extract_dir) os.replace(tmp_dir, extract_dir) @@ -374,6 +418,41 @@ def _extract_wheel(self, *, wheel_path: Path) -> Path: logger.info("Extracted plug-in wheel '%s' to %s", wheel_path.name, extract_dir) return extract_dir + @staticmethod + def _wheel_fingerprint(*, wheel_path: Path) -> str: + """ + Return a cheap change-detecting fingerprint (size and mtime) for a wheel. + + Args: + wheel_path: Path to the plug-in wheel. + + Returns: + str: A ``":"`` fingerprint that changes whenever the wheel is + rebuilt or replaced. + """ + stat = wheel_path.stat() + return f"{stat.st_size}:{stat.st_mtime_ns}" + + @classmethod + def _cached_fingerprint(cls, *, extract_dir: Path) -> str | None: + """ + Return the fingerprint recorded for a cached extraction, or ``None`` if absent. + + A missing marker (e.g. an extraction from before fingerprinting) reads as ``None`` + so the cache is treated as stale and re-extracted rather than trusted blindly. + + Args: + extract_dir: The cached extraction directory. + + Returns: + str | None: The recorded fingerprint, or ``None`` when no valid marker exists. + """ + marker = extract_dir / cls._FINGERPRINT_FILE + try: + return marker.read_text(encoding="utf-8").strip() + except OSError: + return None + @staticmethod def _plugin_base_dir() -> Path: """ From 1538e6eff2d60d66d4d6621b58cc39a887fa956c Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 9 Jul 2026 12:20:59 -0700 Subject: [PATCH 13/34] test(setup): cover plug-in auto-registration, discovery, instantiate, and execute Unit: add two tests that a plug-in's Scenario subclass is auto-registered by type-scoped discovery without a bootstrap, and that a bootstrap-registered scenario is not duplicated by auto-registration. Integration: refactor the injected-wheel test into fixtures (build_mock_wheel, plugin_dir, all_scenarios) and add coverage that every scenario in an out-of-band wheel is discovered and instantiates, plus an env-gated case (PLUGIN_TEST_EXEC_SCENARIO + ADVERSARIAL_CHAT_ENDPOINT) that drives one scenario through the public set_params_from_args/initialize/run pipeline. The execution case tolerates a live-endpoint 'objectives incomplete' outcome (the attack still ran) but re-raises any other exception so a real regression fails loudly. The always-on mock case stays generic with no external package named. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../setup/test_plugin_loader_integration.py | 228 +++++++++++++++--- tests/unit/setup/test_plugin_loader.py | 26 ++ 2 files changed, 221 insertions(+), 33 deletions(-) diff --git a/tests/integration/setup/test_plugin_loader_integration.py b/tests/integration/setup/test_plugin_loader_integration.py index e65d6d4ccb..9045d62616 100644 --- a/tests/integration/setup/test_plugin_loader_integration.py +++ b/tests/integration/setup/test_plugin_loader_integration.py @@ -2,25 +2,29 @@ # Licensed under the MIT license. """ -Integration test: plug-in scenarios are discovered by ``ScenarioRegistry`` after init. +Integration test: plug-in scenarios load, register, instantiate, and execute after init. This exercises the full public loading path end-to-end: point ``PLUGIN_WHEEL`` at a -built wheel, run ``initialize_pyrit_async``, and assert every scenario the wheel ships -is registered in ``ScenarioRegistry``. +built wheel, run ``initialize_pyrit_async``, and assert the wheel's scenarios are +registered in ``ScenarioRegistry``, construct cleanly, and drive through the public +execution pipeline. Two cases: * A self-contained case builds a small wheel at test time and verifies its scenarios are discovered. This always runs and guards the mechanism in public CI. * An injected case loads a wheel supplied out-of-band via environment variables and - verifies all of its scenarios are discovered. It is skipped unless those variables - are set, so the committed test depends on no specific external package. Point it at a - real scenario wheel (for example in a downstream/private CI job) to guarantee that - every scenario that wheel ships is picked up after initialization:: + verifies all of its scenarios are discovered, instantiate, and that at least one + executes. It is skipped unless those variables are set, so the committed test depends + on no specific external package. Point it at a real scenario wheel (for example in a + downstream/private CI job) to guarantee every scenario that wheel ships is picked up:: PLUGIN_TEST_WHEEL=/path/to/plugin.whl - PLUGIN_TEST_PACKAGE=the_plugin_package # optional; enables package enumeration + PLUGIN_TEST_PACKAGE=the_plugin_package # enables package enumeration + instantiation PLUGIN_TEST_SCENARIO_DIRS=/path/to/scenarios # optional; os.pathsep-separated + PLUGIN_TEST_EXEC_SCENARIO=the.registry.name # optional; enables the execution case + ADVERSARIAL_CHAT_ENDPOINT=... # execution target + scorer endpoint + ADVERSARIAL_CHAT_MODEL=... # optional model name """ import inspect @@ -29,16 +33,20 @@ import textwrap import uuid import zipfile -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager from pathlib import Path +from typing import Any import pytest from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider +from pyrit.models import ScenarioResult, ScenarioRunState +from pyrit.prompt_target import OpenAIChatTarget from pyrit.registry import ScenarioRegistry from pyrit.registry.discovery import discover_in_directory -from pyrit.scenario.core import Scenario +from pyrit.scenario.core import DatasetAttackConfiguration, Scenario +from pyrit.score import SelfAskTrueFalseScorer, TrueFalseQuestionPaths from pyrit.setup import IN_MEMORY, PluginSpec, initialize_pyrit_async # (module stem, class name, registry name) for the scenarios the self-contained wheel ships. @@ -48,8 +56,13 @@ ("gamma", "MockGammaScenario", "airt.mock_gamma"), ] +# Substring of the ValueError ``Scenario.run_async`` raises when an atomic attack's +# objective did not complete (for example a target refusal). The attack still ran through +# the public pipeline, so this outcome proves execution while a drift regression -- which +# surfaces a different exception type before or during the run -- still fails loudly. +_OBJECTIVES_INCOMPLETE_MARKER = "objectives incomplete" + -# REV Should this be a Pytest fixture? def _build_scenario_plugin_wheel(dest_dir: Path, *, package: str) -> Path: """ Build a wheel whose package ships several scenarios and a ``register()`` bootstrap. @@ -148,36 +161,83 @@ def all_scenario_class_names(from_dirs: list[Path]) -> set[str]: return names -def _scenario_class_names_under_package(package_prefix: str) -> set[str]: +def _scenario_classes_under_package(package_prefix: str) -> dict[str, type[Scenario]]: """ - Return class names of concrete ``Scenario`` subclasses whose module is under a package. + Return registered concrete ``Scenario`` subclasses owned by a package, keyed by registry name. - Uses in-memory subclass enumeration (reliable after the plug-in has been imported), - which sidesteps the standalone-import problems a filesystem walk can hit for a - package that uses relative imports. + Uses the post-load registry (reliable after the plug-in has been imported and + bootstrapped), which sidesteps the standalone-import problems a filesystem walk can + hit for a package that uses relative imports. Args: package_prefix: The plug-in's top-level package name. Returns: - set[str]: Scenario class names owned by that package. + dict[str, type[Scenario]]: Registry name -> scenario class for that package. """ prefix = f"{package_prefix}." - found: set[str] = set() - seen: set[type] = set() - stack: list[type] = list(Scenario.__subclasses__()) - while stack: - cls = stack.pop() - if cls in seen: - continue - seen.add(cls) - stack.extend(cls.__subclasses__()) + registry = ScenarioRegistry.get_registry_singleton() + found: dict[str, type[Scenario]] = {} + for name in registry.get_class_names(): + cls = registry.get_class(name) module = cls.__module__ or "" if not inspect.isabstract(cls) and (module == package_prefix or module.startswith(prefix)): - found.add(cls.__name__) + found[name] = cls return found +async def _execute_scenario_async( + *, scenario_cls: type[Scenario], endpoint: str, model: str | None +) -> ScenarioResult | None: + """ + Drive one plug-in scenario through the public initialize/run pipeline. + + Constructs an objective target and scorer from the supplied endpoint, initializes the + scenario, and runs it. ``initialize_async`` must succeed and build at least one atomic + attack; that is where dataset-config resolution drift would surface. + + Args: + scenario_cls: The plug-in scenario class to execute. + endpoint: Chat endpoint backing both the objective target and the scorer. + model: Optional model name for the endpoint. + + Returns: + ScenarioResult | None: The completed result, or ``None`` when the run finished with + incomplete objectives (the attack ran end-to-end but the objective was not achieved). + """ + target = OpenAIChatTarget(endpoint=endpoint, model_name=model) + scorer = SelfAskTrueFalseScorer( + chat_target=target, + true_false_question_path=TrueFalseQuestionPaths.TASK_ACHIEVED_REFINED.value, + ) + + try: + scenario = scenario_cls(objective_scorer=scorer, fast_mode=True) # type: ignore[ty:missing-argument, ty:unknown-argument] + except TypeError: + scenario = scenario_cls(objective_scorer=scorer) # type: ignore[ty:missing-argument] + + dataset_config: DatasetAttackConfiguration | None = None + required_datasets = getattr(scenario_cls, "required_datasets", None) + if callable(required_datasets): + names = list(required_datasets()) + if names: + dataset_config = DatasetAttackConfiguration(dataset_names=names, max_dataset_size=1) + + args: dict[str, Any] = {"objective_target": target, "max_concurrency": 1} + if dataset_config is not None: + args["dataset_config"] = dataset_config + scenario.set_params_from_args(args=args) + await scenario.initialize_async() + assert scenario.atomic_attack_count >= 1, "Scenario built no atomic attacks during initialization." + + try: + return await scenario.run_async() + except ValueError as exc: + if _OBJECTIVES_INCOMPLETE_MARKER in str(exc): + return None + raise + + @contextmanager def _plugin_dir_env(*, plugin_dir: Path | None) -> Iterator[None]: """Set PLUGIN_DIR (extraction dir) for the duration of a load, then restore it.""" @@ -197,6 +257,28 @@ def _plugin_dir_env(*, plugin_dir: Path | None) -> Iterator[None]: os.environ[key] = previous +@pytest.fixture +def plugin_dir(tmp_path: Path) -> Path: + """The directory the loader extracts wheels into for a test.""" + return tmp_path / ".plugin" + + +@pytest.fixture +def build_mock_wheel(tmp_path: Path) -> Callable[[str], Path]: + """A builder for the self-contained scenario wheel, parameterized by package name.""" + + def build(package: str) -> Path: + return _build_scenario_plugin_wheel(tmp_path, package=package) + + return build + + +@pytest.fixture +def all_scenarios() -> Callable[[list[Path]], set[str]]: + """The expected-set builder: scenario class names defined under the given source dirs.""" + return all_scenario_class_names + + @pytest.fixture def plugin_sandbox() -> Iterator[None]: """Snapshot and restore global import + registry state so a load does not leak.""" @@ -215,15 +297,19 @@ def plugin_sandbox() -> Iterator[None]: @pytest.mark.run_only_if_all_tests -async def test_built_wheel_scenarios_are_discovered(tmp_path: Path, plugin_sandbox: None) -> None: +async def test_built_wheel_scenarios_are_discovered( + plugin_sandbox: None, + build_mock_wheel: Callable[[str], Path], + plugin_dir: Path, +) -> None: """A built wheel's scenarios are all registered in ScenarioRegistry after init.""" package = f"mock_scenario_plugin_{uuid.uuid4().hex[:8]}" - wheel = _build_scenario_plugin_wheel(tmp_path, package=package) + wheel = build_mock_wheel(package) registry = ScenarioRegistry.get_registry_singleton() before = set(registry.get_class_names()) - with _plugin_dir_env(plugin_dir=tmp_path / ".plugin"): + with _plugin_dir_env(plugin_dir=plugin_dir): await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel)]) registry = ScenarioRegistry.get_registry_singleton() @@ -240,7 +326,10 @@ async def test_built_wheel_scenarios_are_discovered(tmp_path: Path, plugin_sandb @pytest.mark.run_only_if_all_tests -async def test_injected_wheel_scenarios_are_discovered(plugin_sandbox: None) -> None: +async def test_injected_wheel_scenarios_are_discovered( + plugin_sandbox: None, + all_scenarios: Callable[[list[Path]], set[str]], +) -> None: """Every scenario in an out-of-band wheel is discovered after initialization. Skipped unless ``PLUGIN_TEST_WHEEL`` is set, so the committed test names no external @@ -266,11 +355,84 @@ async def test_injected_wheel_scenarios_are_discovered(plugin_sandbox: None) -> if scenario_dirs_env: dirs = [Path(p) for p in scenario_dirs_env.split(os.pathsep) if p] - expected = all_scenario_class_names(dirs) + expected = all_scenarios(dirs) else: assert package is not None - expected = _scenario_class_names_under_package(package) + expected = {cls.__name__ for cls in _scenario_classes_under_package(package).values()} assert expected, "No plug-in scenarios were found to verify; check the injected wheel/scenario source." missing = expected - found assert not missing, f"Plug-in scenarios not discovered by ScenarioRegistry: {sorted(missing)}" + + +@pytest.mark.run_only_if_all_tests +async def test_injected_wheel_scenarios_instantiate(plugin_sandbox: None) -> None: + """Every scenario in an out-of-band wheel constructs cleanly after initialization. + + Skipped unless ``PLUGIN_TEST_WHEEL`` and ``PLUGIN_TEST_PACKAGE`` are set. + """ + wheel_env = os.getenv("PLUGIN_TEST_WHEEL") + package = os.getenv("PLUGIN_TEST_PACKAGE") + if not wheel_env or not package: + pytest.skip("Set PLUGIN_TEST_WHEEL and PLUGIN_TEST_PACKAGE to run the instantiation test.") + + wheel = Path(wheel_env).expanduser() + assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" + + with _plugin_dir_env(plugin_dir=None): + await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel, package=package)]) + + classes = _scenario_classes_under_package(package) + assert classes, f"No plug-in scenarios registered under package {package!r}." + + failures: dict[str, str] = {} + for name, cls in sorted(classes.items()): + try: + instance = cls() # type: ignore[ty:missing-argument] + assert isinstance(instance, Scenario) + except Exception as exc: # noqa: BLE001 + failures[name] = f"{type(exc).__name__}: {exc}" + + assert not failures, f"Plug-in scenarios failed to instantiate: {failures}" + + +@pytest.mark.run_only_if_all_tests +async def test_injected_wheel_scenario_executes(plugin_sandbox: None) -> None: + """One named scenario from an out-of-band wheel runs through the public pipeline. + + Skipped unless ``PLUGIN_TEST_WHEEL``, ``PLUGIN_TEST_PACKAGE``, and + ``PLUGIN_TEST_EXEC_SCENARIO`` are set, and an ``ADVERSARIAL_CHAT_ENDPOINT`` is + configured (it may come from the loaded ``.env``, so it is checked after init). + """ + wheel_env = os.getenv("PLUGIN_TEST_WHEEL") + package = os.getenv("PLUGIN_TEST_PACKAGE") + exec_scenario = os.getenv("PLUGIN_TEST_EXEC_SCENARIO") + if not wheel_env or not package or not exec_scenario: + pytest.skip("Set PLUGIN_TEST_WHEEL, PLUGIN_TEST_PACKAGE, and PLUGIN_TEST_EXEC_SCENARIO to run the exec test.") + + wheel = Path(wheel_env).expanduser() + assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" + + with _plugin_dir_env(plugin_dir=None): + await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel, package=package)]) + + endpoint = os.getenv("ADVERSARIAL_CHAT_ENDPOINT") + if not endpoint: + pytest.skip("ADVERSARIAL_CHAT_ENDPOINT is not configured; skipping plug-in scenario execution.") + + registry = ScenarioRegistry.get_registry_singleton() + assert exec_scenario in registry.get_class_names(), ( + f"PLUGIN_TEST_EXEC_SCENARIO {exec_scenario!r} is not registered; " + f"available: {sorted(registry.get_class_names())}" + ) + scenario_cls = registry.get_class(exec_scenario) + + result = await _execute_scenario_async( + scenario_cls=scenario_cls, + endpoint=endpoint, + model=os.getenv("ADVERSARIAL_CHAT_MODEL"), + ) + + if result is not None: + assert result.scenario_run_state == ScenarioRunState.COMPLETED + assert result.attack_results, "Completed scenario run produced no attack results." diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index 4ed3f9b535..9d93cd5362 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -509,6 +509,32 @@ async def test_ordering_scenario_visible_to_preload(tmp_path: Path) -> None: assert wheel.scenario_name in names +async def test_plugin_scenario_auto_registered_without_bootstrap(tmp_path: Path) -> None: + """A plug-in's Scenario subclass is auto-registered by discovery even with no bootstrap.""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False) + + await load_plugin(wheel, tmp_path / ".plugin") + + # The scenario is picked up purely by type-scoped discovery of the plug-in package, + # so a scenario-only plug-in with no register()/initializer still loads. + registry = ScenarioRegistry.get_registry_singleton() + mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario + assert mock_scenario in registry._classes.values() + assert registry._discovered is False # auto-registration must not trigger built-in discovery + + +async def test_bootstrap_registration_not_duplicated_by_auto_register(tmp_path: Path) -> None: + """A scenario the bootstrap registers is not also re-registered under a fallback name.""" + wheel = build_mock_wheel(tmp_path, bootstrap="register") + + await load_plugin(wheel, tmp_path / ".plugin") + + registry = ScenarioRegistry.get_registry_singleton() + mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario + registered_names = [name for name, cls in registry._classes.items() if cls is mock_scenario] + assert registered_names == [wheel.scenario_name] + + async def test_datasets_only_plugin_loads_without_bootstrap(tmp_path: Path) -> None: """A datasets-only plug-in (no bootstrap, no scenario) loads via import-time registration.""" wheel = build_mock_wheel(tmp_path, bootstrap="none", include_scenario=False) From 9b440cd9df97555234b13c824d7a42f88f0c9f95 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 15:53:05 -0700 Subject: [PATCH 14/34] refactor(exceptions): centralize plug-in load errors Move the plug-in load error hierarchy into pyrit.exceptions and re-export it from the package while preserving the loader's existing exception identities and behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/exceptions/__init__.py | 8 ++++++++ pyrit/exceptions/exception_classes.py | 16 ++++++++++++++++ pyrit/setup/plugin_loader.py | 22 ++++++---------------- 3 files changed, 30 insertions(+), 16 deletions(-) diff --git a/pyrit/exceptions/__init__.py b/pyrit/exceptions/__init__.py index 5b4948a628..cb8d814f89 100644 --- a/pyrit/exceptions/__init__.py +++ b/pyrit/exceptions/__init__.py @@ -10,6 +10,10 @@ ExperimentalWarning, InvalidJsonException, MissingPromptPlaceholderException, + PluginImportError, + PluginLoadError, + PluginRegisteredNothingError, + PluginWheelNotFoundError, PyritException, RateLimitException, get_retry_max_num_attempts, @@ -52,6 +56,10 @@ "handle_bad_request_exception", "InvalidJsonException", "MissingPromptPlaceholderException", + "PluginLoadError", + "PluginWheelNotFoundError", + "PluginImportError", + "PluginRegisteredNothingError", "PyritException", "pyrit_custom_result_retry", "pyrit_json_retry", diff --git a/pyrit/exceptions/exception_classes.py b/pyrit/exceptions/exception_classes.py index 5b4f3de72f..0ab3f0c6a7 100644 --- a/pyrit/exceptions/exception_classes.py +++ b/pyrit/exceptions/exception_classes.py @@ -246,6 +246,22 @@ class ExperimentalWarning(FutureWarning): """ +class PluginLoadError(RuntimeError): + """Base error raised when a configured plug-in fails to load and load failures are not accepted.""" + + +class PluginWheelNotFoundError(PluginLoadError): + """The configured plug-in wheel path does not point to a readable ``.whl`` file.""" + + +class PluginImportError(PluginLoadError): + """The plug-in package could not be imported (missing dependency, or an installed package shadows it).""" + + +class PluginRegisteredNothingError(PluginLoadError): + """The plug-in imported cleanly but registered no datasets or scenarios.""" + + def pyrit_custom_result_retry( retry_function: Callable[..., bool], retry_max_num_attempts: int | None = None ) -> Callable[..., Any]: diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 61866b97d0..15d631d93a 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -38,6 +38,12 @@ from pathlib import Path from typing import TYPE_CHECKING +from pyrit.exceptions.exception_classes import ( + PluginImportError, + PluginLoadError, + PluginRegisteredNothingError, + PluginWheelNotFoundError, +) from pyrit.setup.pyrit_initializer import PyRITInitializer if TYPE_CHECKING: @@ -185,22 +191,6 @@ def _module_owned_by(cls: type, package_name: str) -> bool: ) -class PluginLoadError(RuntimeError): - """Base error raised when a configured plug-in fails to load and load failures are not accepted.""" - - -class PluginWheelNotFoundError(PluginLoadError): - """The configured plug-in wheel path does not point to a readable ``.whl`` file.""" - - -class PluginImportError(PluginLoadError): - """The plug-in package could not be imported (missing dependency, or an installed package shadows it).""" - - -class PluginRegisteredNothingError(PluginLoadError): - """The plug-in imported cleanly but registered no datasets or scenarios.""" - - class PluginLoader: """ Extract and register a single PyRIT plug-in wheel described by a ``PluginSpec``. From 87f1156261f7dcadbd58af60ae843ff3558c6a4d Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 16:08:30 -0700 Subject: [PATCH 15/34] feat(exceptions): define source plug-in failure stages Add source, discovery, validation, and collision errors under PluginLoadError and align the empty-contribution error with the V1 scenario/attack-technique scope. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/exceptions/__init__.py | 12 ++++++++-- pyrit/exceptions/exception_classes.py | 26 ++++++++++++++++---- tests/unit/exceptions/test_exceptions.py | 30 ++++++++++++++++++++++++ 3 files changed, 61 insertions(+), 7 deletions(-) diff --git a/pyrit/exceptions/__init__.py b/pyrit/exceptions/__init__.py index cb8d814f89..93a22f64d3 100644 --- a/pyrit/exceptions/__init__.py +++ b/pyrit/exceptions/__init__.py @@ -10,9 +10,13 @@ ExperimentalWarning, InvalidJsonException, MissingPromptPlaceholderException, + PluginCollisionError, + PluginDiscoveryError, PluginImportError, PluginLoadError, PluginRegisteredNothingError, + PluginSourceNotFoundError, + PluginValidationError, PluginWheelNotFoundError, PyritException, RateLimitException, @@ -56,10 +60,14 @@ "handle_bad_request_exception", "InvalidJsonException", "MissingPromptPlaceholderException", - "PluginLoadError", - "PluginWheelNotFoundError", + "PluginCollisionError", + "PluginDiscoveryError", "PluginImportError", + "PluginLoadError", "PluginRegisteredNothingError", + "PluginSourceNotFoundError", + "PluginValidationError", + "PluginWheelNotFoundError", "PyritException", "pyrit_custom_result_retry", "pyrit_json_retry", diff --git a/pyrit/exceptions/exception_classes.py b/pyrit/exceptions/exception_classes.py index 0ab3f0c6a7..6c14ed24b5 100644 --- a/pyrit/exceptions/exception_classes.py +++ b/pyrit/exceptions/exception_classes.py @@ -247,19 +247,35 @@ class ExperimentalWarning(FutureWarning): class PluginLoadError(RuntimeError): - """Base error raised when a configured plug-in fails to load and load failures are not accepted.""" + """Base error raised when a configured plug-in fails to load.""" -class PluginWheelNotFoundError(PluginLoadError): - """The configured plug-in wheel path does not point to a readable ``.whl`` file.""" +class PluginCollisionError(PluginLoadError): + """A plug-in contribution conflicts with an existing scenario or attack technique.""" + + +class PluginDiscoveryError(PluginLoadError): + """A plug-in imported but its scenario or attack-technique contributions could not be discovered.""" class PluginImportError(PluginLoadError): - """The plug-in package could not be imported (missing dependency, or an installed package shadows it).""" + """A plug-in source module or wheel package could not be imported.""" class PluginRegisteredNothingError(PluginLoadError): - """The plug-in imported cleanly but registered no datasets or scenarios.""" + """The plug-in imported cleanly but contributed no scenarios or attack techniques.""" + + +class PluginSourceNotFoundError(PluginLoadError): + """The configured plug-in source path does not point to a readable Python file or package.""" + + +class PluginValidationError(PluginLoadError): + """A discovered plug-in scenario or attack technique violates the current PyRIT contract.""" + + +class PluginWheelNotFoundError(PluginLoadError): + """The configured plug-in wheel path does not point to a readable ``.whl`` file.""" def pyrit_custom_result_retry( diff --git a/tests/unit/exceptions/test_exceptions.py b/tests/unit/exceptions/test_exceptions.py index e228efed32..e6e35542ec 100644 --- a/tests/unit/exceptions/test_exceptions.py +++ b/tests/unit/exceptions/test_exceptions.py @@ -15,6 +15,14 @@ EmptyResponseException, InvalidJsonException, MissingPromptPlaceholderException, + PluginCollisionError, + PluginDiscoveryError, + PluginImportError, + PluginLoadError, + PluginRegisteredNothingError, + PluginSourceNotFoundError, + PluginValidationError, + PluginWheelNotFoundError, PyritException, RateLimitException, handle_bad_request_exception, @@ -66,6 +74,28 @@ def test_invalid_json_exception_initialization(): assert str(ex) == "Status Code: 500, Message: Invalid JSON Response" +@pytest.mark.parametrize( + "error_type", + [ + PluginSourceNotFoundError, + PluginWheelNotFoundError, + PluginImportError, + PluginDiscoveryError, + PluginValidationError, + PluginCollisionError, + PluginRegisteredNothingError, + ], +) +def test_plugin_specific_errors_inherit_plugin_load_error(error_type: type[PluginLoadError]) -> None: + """Every specific plug-in failure must be catchable through ``PluginLoadError``.""" + assert issubclass(error_type, PluginLoadError) + + +def test_plugin_registered_nothing_error_describes_supported_components() -> None: + """The empty-contribution error must describe the V1 scenario/technique scope.""" + assert "scenarios or attack techniques" in (PluginRegisteredNothingError.__doc__ or "") + + def test_bad_request_exception_process_exception(caplog): ex = BadRequestException() with caplog.at_level(logging.ERROR): From 97e5976028398c149768b2f2c1195285814e011b Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 16:18:19 -0700 Subject: [PATCH 16/34] feat(setup): normalize source and wheel plug-in specs Introduce the explicit V1 PluginSpec schema, resolve artifacts relative to their config file, reject composition, and remove the fail-open config surface while preserving the existing wheel loader behind the normalized spec. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/setup/__init__.py | 3 +- pyrit/setup/configuration_loader.py | 69 ++++++--- pyrit/setup/plugin_loader.py | 77 +--------- pyrit/setup/plugin_spec.py | 138 ++++++++++++++++++ tests/unit/setup/test_configuration_loader.py | 107 +++++++------- tests/unit/setup/test_plugin_loader.py | 37 +---- tests/unit/setup/test_plugin_spec.py | 69 +++++++++ 7 files changed, 311 insertions(+), 189 deletions(-) create mode 100644 pyrit/setup/plugin_spec.py create mode 100644 tests/unit/setup/test_plugin_spec.py diff --git a/pyrit/setup/__init__.py b/pyrit/setup/__init__.py index 1f16003a42..a94f4b3ec4 100644 --- a/pyrit/setup/__init__.py +++ b/pyrit/setup/__init__.py @@ -11,7 +11,7 @@ MemoryDatabaseType, initialize_pyrit_async, ) -from pyrit.setup.plugin_loader import PluginSpec +from pyrit.setup.plugin_spec import PluginFormat, PluginSpec __all__ = [ "AZURE_SQL", @@ -20,6 +20,7 @@ "initialize_pyrit_async", "initialize_from_config_async", "MemoryDatabaseType", + "PluginFormat", "ConfigurationLoader", "PluginSpec", ] diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 6870def2fa..cc0c1d26d4 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -13,7 +13,10 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar +import yaml + from pyrit.common.path import DEFAULT_CONFIG_PATH +from pyrit.common.utils import verify_and_resolve_path from pyrit.common.yaml_loadable import YamlLoadable from pyrit.models import class_name_to_snake_case from pyrit.setup.initialization import ( @@ -24,7 +27,7 @@ ) if TYPE_CHECKING: - from pyrit.setup.plugin_loader import PluginSpec + from pyrit.setup.plugin_spec import PluginSpec from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -110,8 +113,6 @@ class ConfigurationLoader(YamlLoadable): plugins: List of plug-ins to load as the guaranteed-first initialization phase. Each entry is a wheel path string, a ``{package: wheel}`` mapping, or a ``{wheel: ..., package: ...}`` mapping. Empty means no plug-ins. - plugin_accept_load_failures: When True, a plug-in that fails to load is skipped with a - warning instead of raising. None defers to ``PLUGIN_ACCEPT_LOAD_FAILURES`` (else fail-closed). Example YAML configuration: memory_db_type: sqlite @@ -155,9 +156,9 @@ class ConfigurationLoader(YamlLoadable): max_concurrent_scenario_runs: int = 3 allow_custom_initializers: bool = False server: dict[str, Any] | None = None - plugins: list[str | dict[str, Any]] = field(default_factory=list) - plugin_accept_load_failures: bool | None = None + plugins: list[dict[str, Any]] = field(default_factory=list) extensions: dict[str, Any] = field(default_factory=dict) + _plugin_base_dir: pathlib.Path = field(default_factory=pathlib.Path.cwd, repr=False, compare=False) def __post_init__(self) -> None: """Validate and normalize the configuration after loading.""" @@ -231,17 +232,18 @@ def _normalize_plugins(self) -> None: """ Normalize ``plugins`` entries to ``PluginSpec`` instances. - Each entry is a wheel path string, a single-key ``{package: wheel}`` mapping, or an - explicit ``{wheel: ..., package: ...}`` mapping. Validating here (before any other - normalization) fails fast on a malformed plug-in entry, since plug-ins load as the - guaranteed-first phase of initialization. + V1 accepts at most one explicit mapping that selects the ``source`` or ``wheel`` + format. Validating here (before any other normalization) fails fast on a + malformed plug-in entry. Raises: ValueError: If a plug-in entry has an unsupported shape. """ - from pyrit.setup.plugin_loader import PluginSpec + from pyrit.setup.plugin_spec import PluginSpec - self._plugin_specs = [PluginSpec.from_config(entry) for entry in self.plugins] + if len(self.plugins) > 1: + raise ValueError("V1 supports one plug-in at a time; plug-in composition is not supported.") + self._plugin_specs = [PluginSpec.from_config(entry, base_dir=self._plugin_base_dir) for entry in self.plugins] def _normalize_scenario(self) -> None: """ @@ -315,12 +317,13 @@ def scenario_config(self) -> ScenarioConfig | None: return self._scenario_config @classmethod - def from_dict(cls, data: dict[str, Any]) -> "ConfigurationLoader": + def from_dict(cls, data: dict[str, Any], *, plugin_base_dir: pathlib.Path | None = None) -> "ConfigurationLoader": """ Create a ConfigurationLoader from a dictionary. Args: data: Dictionary containing configuration values. + plugin_base_dir: Directory used to resolve relative plug-in paths. Returns: A new ConfigurationLoader instance. @@ -330,7 +333,7 @@ def from_dict(cls, data: dict[str, Any]) -> "ConfigurationLoader": """ # Filter out None values only - empty lists are meaningful ("load nothing") filtered_data = {k: v for k, v in data.items() if v is not None} - known_fields = set(cls.__dataclass_fields__.keys()) + known_fields = {name for name in cls.__dataclass_fields__ if not name.startswith("_")} known_data = {k: v for k, v in filtered_data.items() if k in known_fields and k != "extensions"} extra_data = {k: v for k, v in filtered_data.items() if k not in known_fields} if "extensions" in filtered_data: @@ -338,7 +341,37 @@ def from_dict(cls, data: dict[str, Any]) -> "ConfigurationLoader": if not isinstance(extensions, dict): raise ValueError(f"ConfigurationLoader.extensions must be a dict. Got: {type(extensions).__name__}") extra_data = {**extra_data, **extensions} - return cls(**known_data, extensions=extra_data) + return cls( + **known_data, + extensions=extra_data, + _plugin_base_dir=(plugin_base_dir or pathlib.Path.cwd()).resolve(), + ) + + @classmethod + def from_yaml_file(cls, file: pathlib.Path | str) -> "ConfigurationLoader": + """ + Load configuration while retaining the file directory for plug-in paths. + + Args: + file: The YAML configuration path. + + Returns: + ConfigurationLoader: The parsed and normalized configuration. + + Raises: + FileNotFoundError: If the path does not exist. + ValueError: If the YAML is empty, malformed, or has a non-mapping root. + """ + resolved = verify_and_resolve_path(file) + try: + yaml_data = yaml.safe_load(resolved.read_text(encoding="utf-8")) + except yaml.YAMLError as exc: + raise ValueError(f"Invalid YAML file '{resolved}': {exc}") from exc + if yaml_data is None: + raise ValueError(f"YAML file '{resolved}' is empty.") + if not isinstance(yaml_data, dict): + raise ValueError(f"YAML file '{resolved}' must contain a top-level mapping.") + return cls.from_dict(yaml_data, plugin_base_dir=resolved.parent) @staticmethod def load_with_overrides( @@ -390,7 +423,6 @@ def load_with_overrides( "env_akv_ref": None, "silent": False, "plugins": [], - "plugin_accept_load_failures": None, } # 1. Try loading default config file if it exists @@ -410,8 +442,7 @@ def load_with_overrides( config_data["env_files"] = default_config.env_files config_data["env_akv_ref"] = default_config.env_akv_ref config_data["silent"] = default_config.silent - config_data["plugins"] = list(default_config.plugins) - config_data["plugin_accept_load_failures"] = default_config.plugin_accept_load_failures + config_data["plugins"] = [spec.to_config() for spec in default_config.resolve_plugins()] if default_config.operator: config_data["operator"] = default_config.operator if default_config.operation: @@ -438,8 +469,7 @@ def load_with_overrides( config_data["env_files"] = explicit_config.env_files config_data["env_akv_ref"] = explicit_config.env_akv_ref config_data["silent"] = explicit_config.silent - config_data["plugins"] = list(explicit_config.plugins) - config_data["plugin_accept_load_failures"] = explicit_config.plugin_accept_load_failures + config_data["plugins"] = [spec.to_config() for spec in explicit_config.resolve_plugins()] if explicit_config.operator: config_data["operator"] = explicit_config.operator if explicit_config.operation: @@ -618,7 +648,6 @@ async def initialize_pyrit_async(self) -> None: env_akv_ref=self.env_akv_ref, silent=self.silent, plugins=resolved_plugins if resolved_plugins else None, - plugin_accept_load_failures=self.plugin_accept_load_failures, ) diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 15d631d93a..5ecb3ed326 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -33,8 +33,6 @@ import shutil import sys import tempfile -from collections.abc import Mapping -from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING @@ -47,88 +45,17 @@ from pyrit.setup.pyrit_initializer import PyRITInitializer if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Mapping, Sequence from types import ModuleType from pyrit.registry import ScenarioRegistry + from pyrit.setup.plugin_spec import PluginSpec logger = logging.getLogger(__name__) _TRUE_TOKENS = frozenset({"1", "true", "yes", "on"}) -@dataclass(frozen=True) -class PluginSpec: - """ - A single plug-in to load: a pre-built wheel plus its optional top-level package name. - - Attributes: - wheel (Path): Filesystem path to the pre-built plug-in wheel (``.whl``). - package (str | None): The wheel's top-level import package. When ``None`` it is - auto-detected from the wheel; set it to disambiguate a multi-package wheel. - """ - - wheel: Path - package: str | None = None - - @classmethod - def from_config(cls, entry: str | Mapping[str, object]) -> PluginSpec: - """ - Build a ``PluginSpec`` from a ``.pyrit_conf`` ``plugins`` list entry. - - Accepted shapes: - - * A bare wheel path string (package auto-detected):: - - - /abs/path/to/plugin.whl - - * A single-key ``{package: wheel}`` mapping (concise, names the package):: - - - my_plugin: /abs/path/to/plugin.whl - - * An explicit ``{wheel: ..., package: ...}`` mapping (``package`` optional):: - - - wheel: /abs/path/to/plugin.whl - package: my_plugin - - Args: - entry: One ``plugins`` list item from the parsed configuration. - - Returns: - PluginSpec: The normalized spec. - - Raises: - ValueError: If the entry shape is not one of the accepted forms. - """ - if isinstance(entry, str): - return cls(wheel=Path(entry).expanduser()) - - if isinstance(entry, Mapping): - mapping = dict(entry) - if "wheel" in mapping: - extra_keys = set(mapping) - {"wheel", "package"} - if extra_keys: - raise ValueError( - f"Plug-in mapping has unexpected key(s) {sorted(extra_keys)}; only 'wheel' and " - f"'package' are allowed. Got: {entry}" - ) - wheel = mapping["wheel"] - package = mapping.get("package") - if not isinstance(wheel, str) or (package is not None and not isinstance(package, str)): - raise ValueError(f"Plug-in entry 'wheel'/'package' must be strings. Got: {entry}") - return cls(wheel=Path(wheel).expanduser(), package=package) - if len(mapping) == 1 and "package" not in mapping: - ((package, wheel),) = mapping.items() - if not isinstance(package, str) or not isinstance(wheel, str): - raise ValueError(f"Plug-in entry must map a package name to a wheel path. Got: {entry}") - return cls(wheel=Path(wheel).expanduser(), package=package) - raise ValueError( - f"Plug-in mapping must be a single {{package_name: wheel}} pair or carry a 'wheel' key. Got: {entry}" - ) - - raise ValueError(f"Plug-in entry must be a wheel path string or mapping, got: {type(entry).__name__}") - - async def load_plugins_if_configured_async( *, plugins: Sequence[PluginSpec], accept_load_failures: bool | None = None ) -> None: diff --git a/pyrit/setup/plugin_spec.py b/pyrit/setup/plugin_spec.py new file mode 100644 index 0000000000..2a8c21029a --- /dev/null +++ b/pyrit/setup/plugin_spec.py @@ -0,0 +1,138 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Static plug-in configuration models shared by setup entry points.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Any + +from pyrit.models import validate_registry_name + + +class PluginFormat(str, Enum): + """The supported plug-in artifact formats.""" + + SOURCE = "source" + WHEEL = "wheel" + + +@dataclass(frozen=True) +class PluginSpec: + """A normalized plug-in artifact declaration.""" + + wheel: Path | None = None + package: str | None = None + name: str | None = None + format: PluginFormat | None = None + source: Path | None = None + + def __post_init__(self) -> None: + """ + Infer legacy direct-construction fields while the loader is refactored. + + Raises: + ValueError: If the artifact fields are missing, conflicting, or inconsistent + with the explicit format. + """ + if (self.source is None) == (self.wheel is None): + raise ValueError("PluginSpec requires exactly one of 'source' or 'wheel'.") + + inferred_format = PluginFormat.SOURCE if self.source is not None else PluginFormat.WHEEL + if self.format is not None and self.format is not inferred_format: + raise ValueError(f"Plug-in format '{self.format.value}' does not match its artifact field.") + object.__setattr__(self, "format", inferred_format) + + if self.name is None: + path = self.source or self.wheel + assert path is not None + inferred_name = self.package.split(".", 1)[0] if self.package else path.stem.replace("-", "_") + object.__setattr__(self, "name", inferred_name) + + @property + def artifact_path(self) -> Path: + """The normalized source or wheel artifact path.""" + path = self.source or self.wheel + assert path is not None + return path + + @classmethod + def from_config(cls, entry: dict[str, Any], *, base_dir: Path | None = None) -> PluginSpec: + """ + Normalize one explicit ``.pyrit_conf`` plug-in entry. + + Args: + entry (dict[str, Any]): The YAML plug-in mapping. + base_dir (Path | None): Directory used to resolve relative artifact paths. + + Returns: + PluginSpec: The normalized plug-in declaration. + + Raises: + ValueError: If the mapping is malformed or inconsistent. + """ + if not isinstance(entry, dict): + raise ValueError(f"Plug-in entry must be a mapping, got {type(entry).__name__}.") + + allowed = {"name", "format", "source", "wheel", "package"} + unexpected = set(entry) - allowed + if unexpected: + raise ValueError(f"Plug-in entry has unexpected key(s): {sorted(unexpected)}.") + + name = entry.get("name") + if not isinstance(name, str): + raise ValueError("Plug-in entry requires a string 'name'.") + validate_registry_name(name) + + try: + plugin_format = PluginFormat(entry.get("format")) + except (TypeError, ValueError) as exc: + raise ValueError("Plug-in entry 'format' must be 'source' or 'wheel'.") from exc + + source = entry.get("source") + wheel = entry.get("wheel") + if (source is None) == (wheel is None): + raise ValueError("Plug-in entry requires exactly one of 'source' or 'wheel'.") + artifact_key = plugin_format.value + artifact = entry.get(artifact_key) + other_key = PluginFormat.WHEEL.value if plugin_format is PluginFormat.SOURCE else PluginFormat.SOURCE.value + if not isinstance(artifact, str) or entry.get(other_key) is not None: + raise ValueError(f"Plug-in format '{plugin_format.value}' requires only the '{artifact_key}' field.") + + package = entry.get("package") + if package is not None and ( + not isinstance(package, str) or not all(part.isidentifier() for part in package.split(".")) + ): + raise ValueError("Plug-in 'package' must be a dotted Python identifier.") + + path = Path(artifact).expanduser() + if not path.is_absolute(): + path = (base_dir or Path.cwd()) / path + path = path.resolve() + + return cls( + name=name, + format=plugin_format, + source=path if plugin_format is PluginFormat.SOURCE else None, + wheel=path if plugin_format is PluginFormat.WHEEL else None, + package=package, + ) + + def to_config(self) -> dict[str, str]: + """ + Serialize this spec to the explicit YAML-style mapping. + + Returns: + dict[str, str]: The normalized configuration mapping. + """ + config = { + "name": self.name or "", + "format": self.format.value if self.format else "", + self.format.value if self.format else "source": str(self.artifact_path), + } + if self.package: + config["package"] = self.package + return config diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 804a67be6b..eaabb9a213 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -713,79 +713,72 @@ def test_server_non_dict_raises(self): class TestConfigurationLoaderPlugins: - """Tests for the .pyrit_conf `plugins` list and `plugin_accept_load_failures`.""" + """Tests for the ``.pyrit_conf`` plug-in list.""" def test_no_plugins_by_default(self): """A config with no plugins resolves to an empty plug-in list.""" config = ConfigurationLoader() assert config.plugins == [] assert config.resolve_plugins() == [] - assert config.plugin_accept_load_failures is None - def test_bare_string_entry(self): - """A bare wheel-path string resolves to a spec with no explicit package.""" - from pyrit.setup.plugin_loader import PluginSpec - - config = ConfigurationLoader(plugins=["/abs/path/plugin.whl"]) - assert config.resolve_plugins() == [PluginSpec(wheel=pathlib.Path("/abs/path/plugin.whl"))] - - def test_concise_package_wheel_pair(self): - """A {package: wheel} mapping names the package.""" - from pyrit.setup.plugin_loader import PluginSpec - - config = ConfigurationLoader(plugins=[{"my_pkg": "/abs/path/plugin.whl"}]) - assert config.resolve_plugins() == [PluginSpec(wheel=pathlib.Path("/abs/path/plugin.whl"), package="my_pkg")] - - def test_explicit_mapping_entry(self): - """An explicit {wheel, package} mapping resolves both fields.""" - from pyrit.setup.plugin_loader import PluginSpec - - config = ConfigurationLoader(plugins=[{"wheel": "/abs/path/plugin.whl", "package": "my_pkg"}]) - assert config.resolve_plugins() == [PluginSpec(wheel=pathlib.Path("/abs/path/plugin.whl"), package="my_pkg")] + def test_explicit_source_mapping(self, tmp_path: pathlib.Path): + config = ConfigurationLoader( + plugins=[ + { + "name": "operation_foobar", + "format": "source", + "source": str(tmp_path / "operation_foobar.py"), + } + ] + ) - def test_multiple_plugins_preserve_order(self): - """Several plug-ins resolve in configured order.""" - config = ConfigurationLoader(plugins=["/a/first.whl", {"second_pkg": "/b/second.whl"}]) - specs = config.resolve_plugins() - assert [s.wheel.name for s in specs] == ["first.whl", "second.whl"] - assert specs[1].package == "second_pkg" + spec = config.resolve_plugins()[0] + assert spec.name == "operation_foobar" + assert spec.source == (tmp_path / "operation_foobar.py").resolve() + + def test_multiple_plugins_rejected(self, tmp_path: pathlib.Path): + with pytest.raises(ValueError, match="one plug-in"): + ConfigurationLoader( + plugins=[ + {"name": "first", "format": "source", "source": str(tmp_path / "first.py")}, + {"name": "second", "format": "source", "source": str(tmp_path / "second.py")}, + ] + ) def test_malformed_plugin_entry_fails_fast(self): """A malformed plug-in entry raises during construction (before other setup).""" with pytest.raises(ValueError): ConfigurationLoader(plugins=[123]) # type: ignore[list-item] - def test_from_dict_with_plugins(self): - """from_dict wires the plugins list and accept-load-failures flag onto the loader.""" + def test_from_dict_with_plugins(self, tmp_path: pathlib.Path): + """``from_dict`` wires the explicit plug-in mapping onto the loader.""" config = ConfigurationLoader.from_dict( - {"plugins": [{"pkg": "/x/plugin.whl"}], "plugin_accept_load_failures": True} + { + "plugins": [ + { + "name": "partner", + "format": "wheel", + "wheel": str(tmp_path / "plugin.whl"), + "package": "partner.plugin", + } + ] + } ) - assert config.plugin_accept_load_failures is True - assert config.resolve_plugins()[0].package == "pkg" - - async def test_initialize_forwards_plugins_and_flag(self): - """ConfigurationLoader.initialize_pyrit_async forwards resolved plugins and the flag to core.""" - from pyrit.setup.plugin_loader import PluginSpec - - config = ConfigurationLoader( - memory_db_type="in_memory", - plugins=[{"pkg": "/x/plugin.whl"}], - plugin_accept_load_failures=True, + assert config.resolve_plugins()[0].package == "partner.plugin" + + def test_yaml_plugin_path_resolves_relative_to_config(self, tmp_path: pathlib.Path): + config_dir = tmp_path / "config" + config_dir.mkdir() + config_path = config_dir / ".pyrit_conf" + config_path.write_text( + """plugins: + - name: operation_foobar + format: source + source: ../plugins/operation_foobar.py +""", + encoding="utf-8", ) - with mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") as core_init: - await config.initialize_pyrit_async() - - _, kwargs = core_init.call_args - assert kwargs["plugins"] == [PluginSpec(wheel=pathlib.Path("/x/plugin.whl"), package="pkg")] - assert kwargs["plugin_accept_load_failures"] is True - - async def test_initialize_passes_none_plugins_when_empty(self): - """With no plug-ins configured, core is called with plugins=None (no-op path).""" - config = ConfigurationLoader(memory_db_type="in_memory") - - with mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") as core_init: - await config.initialize_pyrit_async() + config = ConfigurationLoader.from_yaml_file(config_path) - _, kwargs = core_init.call_args - assert kwargs["plugins"] is None + assert config.resolve_plugins()[0].source == (tmp_path / "plugins" / "operation_foobar.py").resolve() diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index 9d93cd5362..7e25bba9ce 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -27,13 +27,13 @@ from pyrit.memory import CentralMemory from pyrit.models import SeedDataset from pyrit.registry import ScenarioRegistry +from pyrit.setup import PluginSpec from pyrit.setup.initialization import IN_MEMORY, _verify_plugin_prerequisites, initialize_pyrit_async from pyrit.setup.plugin_loader import ( PluginImportError, PluginLoader, PluginLoadError, PluginRegisteredNothingError, - PluginSpec, PluginWheelNotFoundError, load_plugins_if_configured_async, ) @@ -912,41 +912,6 @@ def test_resolve_accept_load_failures_defaults_false() -> None: assert _dummy_loader()._resolve_accept_load_failures() is False -# --------------------------------------------------------------------------- -# PluginSpec.from_config parsing -# --------------------------------------------------------------------------- - - -def test_plugin_spec_from_config_bare_string() -> None: - """A bare wheel-path string parses to a spec with no explicit package.""" - spec = PluginSpec.from_config("/abs/path/plugin.whl") - assert spec.wheel == Path("/abs/path/plugin.whl") - assert spec.package is None - - -def test_plugin_spec_from_config_single_key_pair() -> None: - """A {package: wheel} mapping names the package.""" - spec = PluginSpec.from_config({"my_pkg": "/abs/path/plugin.whl"}) - assert spec.wheel == Path("/abs/path/plugin.whl") - assert spec.package == "my_pkg" - - -def test_plugin_spec_from_config_explicit_mapping() -> None: - """An explicit {wheel, package} mapping parses both fields; package is optional.""" - spec = PluginSpec.from_config({"wheel": "/abs/path/plugin.whl", "package": "my_pkg"}) - assert spec == PluginSpec(wheel=Path("/abs/path/plugin.whl"), package="my_pkg") - assert PluginSpec.from_config({"wheel": "/abs/path/plugin.whl"}).package is None - - -@pytest.mark.parametrize( - "entry", [123, {"a": "1", "b": "2"}, {"package": "no_wheel"}, {"wheel": "/x.whl", "packge": "typo"}] -) -def test_plugin_spec_from_config_rejects_bad_shapes(entry: object) -> None: - """Unsupported entry shapes (including explicit mappings with unexpected keys) raise ValueError.""" - with pytest.raises(ValueError): - PluginSpec.from_config(entry) # type: ignore[arg-type] - - # --------------------------------------------------------------------------- # Multiple plug-ins # --------------------------------------------------------------------------- diff --git a/tests/unit/setup/test_plugin_spec.py b/tests/unit/setup/test_plugin_spec.py new file mode 100644 index 0000000000..a4d33ca0f0 --- /dev/null +++ b/tests/unit/setup/test_plugin_spec.py @@ -0,0 +1,69 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from pathlib import Path + +import pytest + +from pyrit.setup import PluginFormat, PluginSpec + + +def test_source_spec_resolves_relative_path(tmp_path: Path) -> None: + spec = PluginSpec.from_config( + { + "name": "operation_foobar", + "format": "source", + "source": "plugins/operation_foobar.py", + }, + base_dir=tmp_path, + ) + + assert spec.name == "operation_foobar" + assert spec.format is PluginFormat.SOURCE + assert spec.source == (tmp_path / "plugins" / "operation_foobar.py").resolve() + assert spec.wheel is None + assert spec.artifact_path == spec.source + + +def test_wheel_spec_resolves_package_and_path(tmp_path: Path) -> None: + spec = PluginSpec.from_config( + { + "name": "partner_scenarios", + "format": "wheel", + "wheel": "plugins/partner.whl", + "package": "partner_scenarios.plugin", + }, + base_dir=tmp_path, + ) + + assert spec.name == "partner_scenarios" + assert spec.format is PluginFormat.WHEEL + assert spec.wheel == (tmp_path / "plugins" / "partner.whl").resolve() + assert spec.source is None + assert spec.package == "partner_scenarios.plugin" + assert spec.artifact_path == spec.wheel + + +@pytest.mark.parametrize( + "entry", + [ + {}, + {"name": "x", "format": "source"}, + {"name": "x", "format": "wheel"}, + {"name": "x", "format": "source", "source": "x.py", "wheel": "x.whl"}, + {"name": "x", "format": "source", "wheel": "x.whl"}, + {"name": "x", "format": "wheel", "source": "x.py"}, + {"name": "Bad-Name", "format": "source", "source": "x.py"}, + {"name": "x", "format": "wheel", "wheel": "x.whl", "package": "bad-package"}, + {"name": "x", "format": "archive", "wheel": "x.whl"}, + {"name": "x", "format": "source", "source": "x.py", "unexpected": True}, + ], +) +def test_plugin_spec_rejects_invalid_config(entry: dict[str, object], tmp_path: Path) -> None: + with pytest.raises(ValueError): + PluginSpec.from_config(entry, base_dir=tmp_path) + + +def test_plugin_spec_rejects_legacy_shorthand() -> None: + with pytest.raises(ValueError, match="mapping"): + PluginSpec.from_config("plugin.whl") # type: ignore[arg-type] From d092a104e6a873c7230a6cef91723375aa4cd132 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 16:59:03 -0700 Subject: [PATCH 17/34] refactor(setup): activate plug-ins through privileged initializer Have ConfigurationLoader inject a framework-owned PluginInitializer ahead of user initializers, remove the special plug-in phase and fail-open API from core initialization, and enforce the single-plugin fail-closed V1 lifecycle. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- .pyrit_conf_example | 7 - pyrit/setup/configuration_loader.py | 7 +- pyrit/setup/initialization.py | 50 ----- pyrit/setup/plugin_loader.py | 109 ++++------ tests/unit/setup/test_configuration_loader.py | 40 ++++ tests/unit/setup/test_plugin_loader.py | 201 ++++-------------- 6 files changed, 124 insertions(+), 290 deletions(-) diff --git a/.pyrit_conf_example b/.pyrit_conf_example index e14a0c599d..0522e20cab 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -80,13 +80,6 @@ initializers: # - my_first_plugin: /abs/path/to/first-0.0.0-py3-none-any.whl # - my_second_plugin: /abs/path/to/second-0.0.0-py3-none-any.whl -# Accept plug-in load failures -# ---------------------------- -# When true, a plug-in that fails to load is skipped with a warning instead of aborting -# initialization. Defaults to false (fail-closed). Precedence: this value wins when set; -# the PLUGIN_ACCEPT_LOAD_FAILURES env var is used only when this is omitted. -# plugin_accept_load_failures: false - # Default Scenario # ---------------- # Optional default scenario to run when invoking `pyrit_scan` without a diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index cc0c1d26d4..cf57abb065 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -632,10 +632,14 @@ async def initialize_pyrit_async(self) -> None: Raises: ValueError: If configuration is invalid or initializers cannot be resolved. """ - resolved_initializers = self.resolve_initializers() + resolved_initializers = list(self.resolve_initializers()) resolved_scripts = self.resolve_initialization_scripts() resolved_env_files = self.resolve_env_files() resolved_plugins = self.resolve_plugins() + if resolved_plugins: + from pyrit.setup.plugin_loader import PluginInitializer + + resolved_initializers.insert(0, PluginInitializer(plugins=resolved_plugins)) # Map snake_case memory_db_type to internal constant internal_memory_db_type = self._MEMORY_DB_TYPE_MAP[self.memory_db_type] @@ -647,7 +651,6 @@ async def initialize_pyrit_async(self) -> None: env_files=resolved_env_files, env_akv_ref=self.env_akv_ref, silent=self.silent, - plugins=resolved_plugins if resolved_plugins else None, ) diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index efb711316f..43e7bfd4e0 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -13,7 +13,6 @@ from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory if TYPE_CHECKING: - from pyrit.setup.plugin_loader import PluginSpec from pyrit.setup.pyrit_initializer import PyRITInitializer logger = logging.getLogger(__name__) @@ -196,30 +195,6 @@ async def _execute_initializers_async(*, initializers: Sequence["PyRITInitialize raise -def _verify_plugin_prerequisites() -> None: - """ - Verify the initialization phases plug-ins depend on completed successfully. - - Plug-in loading runs a third party's bootstrap, which typically constructs targets and - scorers (needing loaded environment variables) and reads or writes central memory. This - is the health gate for that phase: plug-ins are loaded only if the prior initialization - phases were green. Central memory being set is the concrete, verifiable signal that the - environment and memory phases completed — the linear init flow would have raised earlier - otherwise, so a violation here means a broken initialization state, not a plug-in fault, - and is surfaced loudly regardless of ``plugin_accept_load_failures``. - - Raises: - RuntimeError: If central memory is not set, indicating a prior phase did not complete. - """ - try: - CentralMemory.get_memory_instance() - except Exception as exc: - raise RuntimeError( - "Cannot load plug-ins: central memory is not initialized. Plug-in loading requires " - "the environment and memory phases of initialize_pyrit_async to have completed first." - ) from exc - - async def initialize_pyrit_async( memory_db_type: MemoryDatabaseType | str, *, @@ -228,8 +203,6 @@ async def initialize_pyrit_async( env_files: Sequence[pathlib.Path] | None = None, env_akv_ref: Sequence[str] | None = None, silent: bool = False, - plugins: Sequence["PluginSpec"] | None = None, - plugin_accept_load_failures: bool | None = None, **memory_instance_kwargs: Any, ) -> None: """ @@ -251,15 +224,6 @@ async def initialize_pyrit_async( so local files take precedence over AKV. Requires ``azure-keyvault-secrets``. silent (bool): If True, suppresses print statements about environment file loading and schema migration. Defaults to False. - plugins (Sequence[PluginSpec] | None): Optional plug-ins to load as a guaranteed-first phase, - after central memory is set and before the configured initializers run. Each plug-in is a - pre-built wheel shipping datasets/scenarios that register like built-ins. Plug-ins are - loaded in list order, so one plug-in behaves identically to several. Typically populated - from the ``plugins`` key of ``.pyrit_conf``; direct callers pass ``PluginSpec`` instances. - plugin_accept_load_failures (bool | None): Overrides ``PLUGIN_ACCEPT_LOAD_FAILURES`` for plug-in - loading. When True, a plug-in that fails to load is skipped with a warning instead of - raising. Defaults to None (use the ``PLUGIN_ACCEPT_LOAD_FAILURES`` environment variable, - else fail-closed). **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. Raises: @@ -295,20 +259,6 @@ async def initialize_pyrit_async( CentralMemory.set_memory_instance(memory) - # Load configured plug-ins as a guaranteed-first phase: after central memory is set (a - # plug-in bootstrap may use it) and BEFORE any configured initializer runs, so plug-in - # datasets/scenarios are registered before LoadDefaultDatasets and PreloadScenarioMetadata - # read the registry. Because `plugins` is its own always-first phase (not one of the - # ordered `initializers`), this ordering is true by construction. Plug-ins are loaded only - # after a health check confirms the initialization phases they depend on completed. No-op - # when no plug-ins are configured. - if plugins: - _verify_plugin_prerequisites() - - from pyrit.setup.plugin_loader import load_plugins_if_configured_async - - await load_plugins_if_configured_async(plugins=plugins, accept_load_failures=plugin_accept_load_failures) - # Combine directly provided initializers with those loaded from scripts all_initializers = list(initializers) if initializers else [] diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 5ecb3ed326..4e61c7d400 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -53,12 +53,8 @@ logger = logging.getLogger(__name__) -_TRUE_TOKENS = frozenset({"1", "true", "yes", "on"}) - -async def load_plugins_if_configured_async( - *, plugins: Sequence[PluginSpec], accept_load_failures: bool | None = None -) -> None: +async def load_plugins_if_configured_async(*, plugins: Sequence[PluginSpec]) -> None: """ Load every configured plug-in, in order. A no-op when ``plugins`` is empty. @@ -69,19 +65,19 @@ async def load_plugins_if_configured_async( Args: plugins: The plug-in specs to load, in order. - accept_load_failures: If provided, overrides the ``PLUGIN_ACCEPT_LOAD_FAILURES`` - environment variable. When True, a plug-in that fails to load is skipped with - a warning and the next plug-in still loads. Raises: - PluginLoadError: If a plug-in fails to load and load failures are not accepted. + PluginLoadError: If a plug-in fails to load. + ValueError: If more than one plug-in is supplied. """ if not plugins: logger.debug("No plug-ins configured; plug-in loading is a no-op.") return + if len(plugins) > 1: + raise ValueError("V1 supports one plug-in at a time; plug-in composition is not supported.") for spec in plugins: - await PluginLoader(spec=spec, accept_load_failures=accept_load_failures).load_async() + await PluginLoader(spec=spec).load_async() def _name_owned_by(module_name: str, package_name: str) -> bool: @@ -112,10 +108,35 @@ def _module_owned_by(cls: type, package_name: str) -> bool: return _name_owned_by(cls.__module__ or "", package_name) -_REMEDIATION = ( - "Remove the plug-in configuration, or accept load failures (PLUGIN_ACCEPT_LOAD_FAILURES=true, or the " - "initialize_pyrit_async(plugin_accept_load_failures=True) parameter) to continue without it." -) +_REMEDIATION = "Fix or remove the plug-in entry from .pyrit_conf, then restart PyRIT." + + +class PluginInitializer(PyRITInitializer): + """Privileged config-owned initializer that activates one scenario plug-in.""" + + def __init__(self, *, plugins: Sequence[PluginSpec]) -> None: + """ + Initialize the privileged plug-in loader. + + Args: + plugins (Sequence[PluginSpec]): The normalized plug-in declarations. + + Raises: + ValueError: If V1 receives anything other than one plug-in. + """ + super().__init__() + if len(plugins) != 1: + raise ValueError("PluginInitializer requires exactly one plug-in in V1.") + self._plugins = tuple(plugins) + + @property + def plugins(self) -> list[PluginSpec]: + """The normalized plug-in declarations.""" + return list(self._plugins) + + async def initialize_async(self) -> None: + """Activate the configured plug-in before user-configured initializers.""" + await load_plugins_if_configured_async(plugins=self._plugins) class PluginLoader: @@ -124,46 +145,34 @@ class PluginLoader: The wheel is extracted to ``.plugin//`` (never installed), imported, and its bootstrap is run so its datasets and scenarios register like built-ins. Fails closed - by default; set ``accept_load_failures`` (constructor / ``initialize_pyrit_async`` - param) or ``PLUGIN_ACCEPT_LOAD_FAILURES`` to continue without the plug-in when it - cannot be loaded. + by default. """ _FINGERPRINT_FILE = ".plugin_wheel_fingerprint" - def __init__(self, *, spec: PluginSpec, accept_load_failures: bool | None = None) -> None: + def __init__(self, *, spec: PluginSpec) -> None: """ Initialize the loader for one plug-in. Args: spec: The plug-in to load (wheel path plus optional package name). - accept_load_failures: If provided, overrides ``PLUGIN_ACCEPT_LOAD_FAILURES``. - When True, a plug-in that fails to load is skipped with a warning instead - of raising. """ self._spec = spec - self._explicit_accept_load_failures = accept_load_failures async def load_async(self) -> None: """ Load the plug-in described by this loader's ``PluginSpec``. Raises: - PluginLoadError: If the plug-in fails to load and load failures are not accepted. + PluginWheelNotFoundError: If the spec is not a wheel-format plug-in. + PluginLoadError: If the plug-in fails to load. """ wheel = self._spec.wheel - accept_load_failures = self._resolve_accept_load_failures() - + if wheel is None: + raise PluginWheelNotFoundError("The current loader only accepts wheel-format plug-ins.") try: await self._load_plugin_async(wheel_path=wheel.expanduser()) except Exception as exc: - if accept_load_failures: - logger.warning( - "Plug-in wheel '%s' failed to load; load failures are accepted so continuing without it: %s", - wheel, - exc, - ) - return message = f"Failed to load plug-in from wheel '{wheel}': {exc} {_REMEDIATION}" # Preserve the specific failure type so callers can distinguish modes, while # always surfacing the remediation guidance. Unknown errors (e.g. a raising @@ -610,8 +619,7 @@ def _warn_on_dataset_name_collisions(*, package_name: str) -> list[str]: the real collision, so false-positive-free with a documented gap beats high-recall-but- noisy. Hard enforcement that can tell a real mismatch from a harmless name coincidence belongs to the scenario's declared required-dataset-names / expected-source mechanism - (which knows the operator's intent), not this loader, and is intentionally not gated - behind ``PLUGIN_ACCEPT_LOAD_FAILURES``. + (which knows the operator's intent), not this loader. Args: package_name: The plug-in's top-level package name. @@ -713,36 +721,3 @@ def _rollback( changed_scenarios = True if changed_scenarios: scenario_registry._metadata_cache = None - - def _resolve_accept_load_failures(self) -> bool: - """ - Resolve the accept-load-failures setting from the explicit value or the environment. - - Precedence: the explicit ``accept_load_failures`` passed to the constructor (e.g. from - ``initialize_pyrit_async``), then the ``PLUGIN_ACCEPT_LOAD_FAILURES`` environment - variable, otherwise fail-closed. - - Returns: - bool: True if a failed plug-in load should be skipped with a warning. - """ - if self._explicit_accept_load_failures is not None: - return self._explicit_accept_load_failures - - env_value = os.getenv("PLUGIN_ACCEPT_LOAD_FAILURES") - if env_value is not None: - return self._coerce_bool(env_value) - - return False - - @staticmethod - def _coerce_bool(value: str) -> bool: - """ - Interpret a string as a boolean flag. - - Args: - value: The raw string value. - - Returns: - bool: True for common truthy tokens (1/true/yes/on, case-insensitive). - """ - return str(value).strip().lower() in _TRUE_TOKENS diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index eaabb9a213..91723fe174 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -782,3 +782,43 @@ def test_yaml_plugin_path_resolves_relative_to_config(self, tmp_path: pathlib.Pa config = ConfigurationLoader.from_yaml_file(config_path) assert config.resolve_plugins()[0].source == (tmp_path / "plugins" / "operation_foobar.py").resolve() + + async def test_initialize_prepends_privileged_plugin_initializer(self, tmp_path: pathlib.Path): + from pyrit.setup.plugin_loader import PluginInitializer + + user_initializer = mock.MagicMock() + config = ConfigurationLoader( + memory_db_type="in_memory", + plugins=[ + { + "name": "partner", + "format": "wheel", + "wheel": str(tmp_path / "partner.whl"), + "package": "partner.plugin", + } + ], + ) + + with ( + mock.patch.object(config, "resolve_initializers", return_value=[user_initializer]), + mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") as core_init, + ): + await config.initialize_pyrit_async() + + initializers = core_init.call_args.kwargs["initializers"] + assert isinstance(initializers[0], PluginInitializer) + assert initializers[0].plugins == config.resolve_plugins() + assert initializers[1] is user_initializer + assert "plugins" not in core_init.call_args.kwargs + + async def test_initialize_without_plugins_uses_only_configured_initializers(self): + user_initializer = mock.MagicMock() + config = ConfigurationLoader(memory_db_type="in_memory") + + with ( + mock.patch.object(config, "resolve_initializers", return_value=[user_initializer]), + mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") as core_init, + ): + await config.initialize_pyrit_async() + + assert core_init.call_args.kwargs["initializers"] == [user_initializer] diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index 7e25bba9ce..0186f3d989 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -10,6 +10,7 @@ silent-failure guards called out in the design brief. """ +import inspect import logging import os import sys @@ -28,9 +29,10 @@ from pyrit.models import SeedDataset from pyrit.registry import ScenarioRegistry from pyrit.setup import PluginSpec -from pyrit.setup.initialization import IN_MEMORY, _verify_plugin_prerequisites, initialize_pyrit_async +from pyrit.setup.initialization import IN_MEMORY, initialize_pyrit_async from pyrit.setup.plugin_loader import ( PluginImportError, + PluginInitializer, PluginLoader, PluginLoadError, PluginRegisteredNothingError, @@ -286,7 +288,7 @@ def plugin_sandbox() -> Iterator[None]: def plugin_env(**overrides: str) -> Iterator[None]: """Patch os.environ so only the given PLUGIN_* overrides are present.""" with patch.dict(os.environ, overrides, clear=False): - for key in ("PLUGIN_DIR", "PLUGIN_ACCEPT_LOAD_FAILURES"): + for key in ("PLUGIN_DIR",): if key not in overrides: os.environ.pop(key, None) yield @@ -297,98 +299,64 @@ def _spec(wheel: MockWheel, *, package: str | None = None) -> PluginSpec: return PluginSpec(wheel=wheel.path, package=package) -def _dummy_loader(*, accept_load_failures: bool | None = None) -> PluginLoader: - """Build a PluginLoader with a placeholder spec for testing policy resolution.""" - return PluginLoader(spec=PluginSpec(wheel=Path("dummy.whl")), accept_load_failures=accept_load_failures) - - async def load_plugin( wheel: MockWheel, plugin_dir: Path, *, - accept_load_failures: bool | None = None, package: str | None = None, extra_env: dict[str, str] | None = None, ) -> None: """Run the plug-in loader against a single mock wheel with an isolated env.""" with plugin_env(PLUGIN_DIR=str(plugin_dir), **(extra_env or {})): - await load_plugins_if_configured_async( - plugins=[_spec(wheel, package=package)], accept_load_failures=accept_load_failures - ) + await load_plugins_if_configured_async(plugins=[_spec(wheel, package=package)]) # --------------------------------------------------------------------------- -# Loader phase inside initialize_pyrit_async +# Privileged PluginInitializer # --------------------------------------------------------------------------- -def test_verify_plugin_prerequisites_passes_when_memory_set() -> None: - """The health check passes when central memory is initialized.""" - with patch.object(CentralMemory, "get_memory_instance", return_value=MagicMock()): - _verify_plugin_prerequisites() # must not raise +def test_core_initialize_has_no_plugin_parameters() -> None: + parameters = inspect.signature(initialize_pyrit_async).parameters + assert "plugins" not in parameters + assert "plugin_accept_load_failures" not in parameters -def test_verify_plugin_prerequisites_raises_when_memory_unset() -> None: - """The health check fails loudly when a prerequisite phase (memory) did not complete.""" - with patch.object(CentralMemory, "get_memory_instance", side_effect=ValueError("not set")): - with pytest.raises(RuntimeError, match="central memory is not initialized"): - _verify_plugin_prerequisites() +async def test_plugin_initializer_delegates_to_loader() -> None: + load_plugins_mock = AsyncMock() + spec = PluginSpec(wheel=Path("plugin.whl")) + initializer = PluginInitializer(plugins=[spec]) + with patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", load_plugins_mock): + await initializer.initialize_async() -async def test_plugin_phase_runs_after_memory_before_initializers() -> None: - """initialize_pyrit_async loads plug-ins after memory is set, before initializers.""" - manager = MagicMock() - manager.attach_mock(MagicMock(), "set_memory") - manager.attach_mock(AsyncMock(), "load_plugins") - manager.attach_mock(AsyncMock(), "execute") + load_plugins_mock.assert_awaited_once() + assert list(load_plugins_mock.await_args.kwargs["plugins"]) == [spec] - with ( - patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), - patch.object(CentralMemory, "set_memory_instance", manager.set_memory), - patch("pyrit.setup.initialization._verify_plugin_prerequisites"), - patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", manager.load_plugins), - patch("pyrit.setup.initialization._execute_initializers_async", manager.execute), - ): - await initialize_pyrit_async( - IN_MEMORY, - initializers=[MagicMock()], - env_files=[], - silent=True, - plugins=[PluginSpec(wheel=Path("plugin.whl"))], - ) - order = [call[0] for call in manager.mock_calls if call[0] in {"set_memory", "load_plugins", "execute"}] - assert order.index("set_memory") < order.index("load_plugins") < order.index("execute") +async def test_plugin_initializer_runs_after_memory_setup() -> None: + memory = MagicMock() + initializer = PluginInitializer(plugins=[PluginSpec(wheel=Path("plugin.whl"))]) + async def _assert_memory_ready() -> None: + assert CentralMemory.get_memory_instance() is memory -async def test_plugin_phase_forwards_accept_load_failures_param() -> None: - """initialize_pyrit_async forwards plugins and plugin_accept_load_failures to the loader.""" - load_plugins_mock = AsyncMock() - spec = PluginSpec(wheel=Path("plugin.whl")) with ( - patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), - patch.object(CentralMemory, "set_memory_instance"), - patch("pyrit.setup.initialization._verify_plugin_prerequisites"), - patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", load_plugins_mock), + patch("pyrit.setup.initialization.SQLiteMemory", return_value=memory), + patch.object(initializer, "initialize_async", side_effect=_assert_memory_ready) as initialize_mock, ): - await initialize_pyrit_async( - IN_MEMORY, env_files=[], silent=True, plugins=[spec], plugin_accept_load_failures=True - ) + await initialize_pyrit_async(IN_MEMORY, initializers=[initializer], env_files=[], silent=True) - load_plugins_mock.assert_awaited_once_with(plugins=[spec], accept_load_failures=True) + initialize_mock.assert_awaited_once() -async def test_plugin_phase_skipped_when_no_plugins() -> None: - """initialize_pyrit_async does not touch the plug-in loader when no plug-ins are configured.""" - load_plugins_mock = AsyncMock() - with ( - patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), - patch.object(CentralMemory, "set_memory_instance"), - patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", load_plugins_mock), - ): - await initialize_pyrit_async(IN_MEMORY, env_files=[], silent=True) +def test_plugin_initializer_is_not_user_discoverable() -> None: + from pyrit.registry import InitializerRegistry + + InitializerRegistry.reset_registry_singleton() + registry = InitializerRegistry.get_registry_singleton() - load_plugins_mock.assert_not_awaited() + assert "plugin" not in registry.get_class_names() # --------------------------------------------------------------------------- @@ -650,17 +618,6 @@ async def test_failing_bootstrap_rolls_back_partial_registration(tmp_path: Path) assert str(plugin_dir / wheel.path.stem) not in sys.path -async def test_accept_load_failures_rolls_back_partial_registration(tmp_path: Path) -> None: - """When load failures are accepted, a partially-registered failed plug-in is still fully rolled back.""" - wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") - plugin_dir = tmp_path / ".plugin" - - await load_plugin(wheel, plugin_dir, accept_load_failures=True) # must not raise - - registry = ScenarioRegistry.get_registry_singleton() - assert wheel.scenario_name not in registry._classes - - async def test_rollback_restores_overwritten_provider(tmp_path: Path) -> None: """A failed load restores a provider entry the plug-in overwrote (name collision).""" wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") @@ -790,20 +747,6 @@ async def test_missing_wheel_fails_closed() -> None: await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=Path("does_not_exist.whl"))]) -async def test_missing_wheel_accept_load_failures_param_proceeds() -> None: - """Accepting load failures via the explicit param skips a broken plug-in with a warning.""" - with plugin_env(): - await load_plugins_if_configured_async( - plugins=[PluginSpec(wheel=Path("does_not_exist.whl"))], accept_load_failures=True - ) # must not raise - - -async def test_missing_wheel_accept_load_failures_env_proceeds() -> None: - """Accepting load failures via PLUGIN_ACCEPT_LOAD_FAILURES env skips a broken plug-in with a warning.""" - with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): - await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=Path("does_not_exist.whl"))]) # must not raise - - async def test_empty_wheel_is_loud(tmp_path: Path) -> None: """A wheel that imports cleanly but registers nothing fails loudly.""" wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) @@ -812,13 +755,6 @@ async def test_empty_wheel_is_loud(tmp_path: Path) -> None: await load_plugin(wheel, tmp_path / ".plugin") -async def test_empty_wheel_accept_load_failures_proceeds(tmp_path: Path) -> None: - """An empty wheel with load failures accepted proceeds instead of raising.""" - wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) - - await load_plugin(wheel, tmp_path / ".plugin", accept_load_failures=True) # must not raise - - async def test_non_whl_path_fails_closed(tmp_path: Path) -> None: """A wheel path that is not a .whl file fails closed with PluginWheelNotFoundError.""" not_a_wheel = tmp_path / "plugin.zip" @@ -873,79 +809,16 @@ def __init__(self, *, required_value: str) -> None: registry._build_metadata("airt.bad", BadScenario) -# --------------------------------------------------------------------------- -# accept_load_failures resolution -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - "value,expected", - [("true", True), ("True", True), ("1", True), ("yes", True), ("on", True), ("false", False), ("0", False)], -) -def test_resolve_accept_load_failures_from_env_tokens(value: str, expected: bool) -> None: - """accept_load_failures resolves from PLUGIN_ACCEPT_LOAD_FAILURES across truthy/falsey tokens.""" - with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES=value): - assert _dummy_loader()._resolve_accept_load_failures() is expected - - -def test_resolve_accept_load_failures_explicit_true() -> None: - """An explicit accept_load_failures=True resolves to True.""" - with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="false"): - assert _dummy_loader(accept_load_failures=True)._resolve_accept_load_failures() is True - - -def test_resolve_accept_load_failures_explicit_overrides_env() -> None: - """An explicit accept_load_failures value takes precedence over the env var.""" - with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): - assert _dummy_loader(accept_load_failures=False)._resolve_accept_load_failures() is False - - -def test_resolve_accept_load_failures_from_env_when_no_explicit() -> None: - """accept_load_failures falls back to PLUGIN_ACCEPT_LOAD_FAILURES when no explicit value is set.""" - with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): - assert _dummy_loader()._resolve_accept_load_failures() is True - - -def test_resolve_accept_load_failures_defaults_false() -> None: - """accept_load_failures defaults to False (fail-closed).""" - with plugin_env(): - assert _dummy_loader()._resolve_accept_load_failures() is False - - # --------------------------------------------------------------------------- # Multiple plug-ins # --------------------------------------------------------------------------- -async def test_multiple_plugins_all_load(tmp_path: Path) -> None: - """Every configured plug-in loads; a single plug-in is just the one-element case.""" +async def test_multiple_plugins_rejected(tmp_path: Path) -> None: + """V1 rejects plug-in composition.""" wheel_a = build_mock_wheel(tmp_path / "a", package_name="mock_plugin_alpha") wheel_b = build_mock_wheel(tmp_path / "b", package_name="mock_plugin_beta") with plugin_env(PLUGIN_DIR=str(tmp_path / ".plugin")): - await load_plugins_if_configured_async(plugins=[_spec(wheel_a), _spec(wheel_b)]) - - names = ScenarioRegistry.get_registry_singleton().get_class_names() - assert wheel_a.scenario_name in names - assert wheel_b.scenario_name in names - - -async def test_plugins_load_in_configured_order(tmp_path: Path) -> None: - """Plug-ins load in list order (later entries load after earlier ones).""" - wheel_a = build_mock_wheel(tmp_path / "a", package_name="mock_plugin_first") - wheel_b = build_mock_wheel(tmp_path / "b", package_name="mock_plugin_second") - - loaded: list[str] = [] - original = PluginLoader.load_async - - async def _record(self: PluginLoader) -> None: - loaded.append(self._spec.package or self._spec.wheel.stem) - await original(self) - - with plugin_env(PLUGIN_DIR=str(tmp_path / ".plugin")): - with patch.object(PluginLoader, "load_async", _record): - await load_plugins_if_configured_async( - plugins=[_spec(wheel_a, package="mock_plugin_first"), _spec(wheel_b, package="mock_plugin_second")] - ) - - assert loaded == ["mock_plugin_first", "mock_plugin_second"] + with pytest.raises(ValueError, match="one plug-in"): + await load_plugins_if_configured_async(plugins=[_spec(wheel_a), _spec(wheel_b)]) From 38f62fa9df839aa0d6557ce3ea00fede934d638e Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 17:14:41 -0700 Subject: [PATCH 18/34] refactor(setup): isolate wheel plug-in preparation Introduce the common PreparedPlugin artifact model and WheelPluginFormat adapter, route the loader through safe threaded preparation, retain advisory version metadata, and remove duplicated extraction/package-resolution logic. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/setup/plugin_formats.py | 165 +++++++++++++++++++ pyrit/setup/plugin_loader.py | 203 +----------------------- tests/unit/setup/test_plugin_formats.py | 116 ++++++++++++++ tests/unit/setup/test_plugin_loader.py | 74 --------- 4 files changed, 289 insertions(+), 269 deletions(-) create mode 100644 pyrit/setup/plugin_formats.py create mode 100644 tests/unit/setup/test_plugin_formats.py diff --git a/pyrit/setup/plugin_formats.py b/pyrit/setup/plugin_formats.py new file mode 100644 index 0000000000..675ec34aab --- /dev/null +++ b/pyrit/setup/plugin_formats.py @@ -0,0 +1,165 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Artifact-format adapters for scenario plug-in consumption.""" + +from __future__ import annotations + +import asyncio +import hashlib +import os +import shutil +import tempfile +from dataclasses import dataclass +from pathlib import Path + +from pyrit.exceptions import PluginWheelNotFoundError +from pyrit.setup.plugin_spec import PluginFormat, PluginSpec + + +@dataclass(frozen=True) +class PreparedPlugin: + """A prepared import root shared by every plug-in artifact format.""" + + spec: PluginSpec + import_root: Path + entry_modules: tuple[str, ...] + owned_module_prefixes: tuple[str, ...] + artifact_fingerprint: str + declared_pyrit_version: str | None = None + + +class WheelPluginFormat: + """Prepare a pre-built Python wheel for plug-in discovery.""" + + _FINGERPRINT_FILE = ".plugin_wheel_fingerprint" + + def __init__(self, *, base_dir: Path | None = None) -> None: + """ + Initialize the wheel adapter. + + Args: + base_dir (Path | None): Extraction root. Defaults to ``PLUGIN_DIR`` or + PyRIT's ``.plugin`` directory. + """ + self._base_dir = base_dir + + async def prepare_async(self, *, spec: PluginSpec) -> PreparedPlugin: + """ + Validate and extract one wheel-format plug-in. + + Args: + spec (PluginSpec): The normalized wheel declaration. + + Returns: + PreparedPlugin: The prepared wheel import root and ownership metadata. + + Raises: + ValueError: If the spec is not wheel format or package discovery is ambiguous. + PluginWheelNotFoundError: If the wheel path is absent or invalid. + """ + return await asyncio.to_thread(self._prepare, spec=spec) + + def _prepare(self, *, spec: PluginSpec) -> PreparedPlugin: + if spec.format is not PluginFormat.WHEEL or spec.wheel is None: + raise ValueError("WheelPluginFormat requires a wheel-format PluginSpec.") + wheel = spec.wheel.expanduser() + if not wheel.is_file(): + raise PluginWheelNotFoundError(f"Plug-in wheel does not point to an existing file: {wheel}") + if wheel.suffix != ".whl": + raise PluginWheelNotFoundError(f"Plug-in wheel must be a .whl file, got: {wheel}") + + fingerprint = self._fingerprint(wheel=wheel) + extract_dir = self._extract(wheel=wheel, fingerprint=fingerprint) + package = self._resolve_package(extract_dir=extract_dir, explicit_package=spec.package) + return PreparedPlugin( + spec=spec, + import_root=extract_dir, + entry_modules=(package,), + owned_module_prefixes=(package,), + artifact_fingerprint=fingerprint, + declared_pyrit_version=self._read_required_pyrit_version(extract_dir=extract_dir), + ) + + def _extract(self, *, wheel: Path, fingerprint: str) -> Path: + from pyrit.common.safe_extract import safe_extract_zip + + base_dir = self._resolve_base_dir() + base_dir.mkdir(parents=True, exist_ok=True) + extract_dir = base_dir / wheel.stem + marker = extract_dir / self._FINGERPRINT_FILE + if extract_dir.is_dir() and marker.is_file(): + try: + if marker.read_text(encoding="utf-8").strip() == fingerprint: + return extract_dir + except OSError: + pass + + temp_dir = Path(tempfile.mkdtemp(prefix=f".{wheel.stem}.tmp-", dir=base_dir)) + try: + safe_extract_zip(source=wheel, dest_dir=temp_dir) + (temp_dir / self._FINGERPRINT_FILE).write_text(fingerprint, encoding="utf-8") + if extract_dir.exists(): + shutil.rmtree(extract_dir) + os.replace(temp_dir, extract_dir) + finally: + if temp_dir.exists(): + shutil.rmtree(temp_dir, ignore_errors=True) + return extract_dir + + def _resolve_base_dir(self) -> Path: + if self._base_dir is not None: + return self._base_dir.expanduser().resolve() + override = os.getenv("PLUGIN_DIR") + if override: + return Path(override).expanduser().resolve() + from pyrit.common import path + + return (Path(path.HOME_PATH) / ".plugin").resolve() + + @staticmethod + def _fingerprint(*, wheel: Path) -> str: + digest = hashlib.sha256() + with wheel.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + @staticmethod + def _resolve_package(*, extract_dir: Path, explicit_package: str | None) -> str: + if explicit_package: + return explicit_package + for dist_info in sorted(extract_dir.glob("*.dist-info")): + top_level = dist_info / "top_level.txt" + if top_level.is_file(): + names = [line.strip() for line in top_level.read_text(encoding="utf-8").splitlines() if line.strip()] + if len(names) == 1: + return names[0] + + candidates = sorted( + child.name + for child in extract_dir.iterdir() + if child.is_dir() and not child.name.endswith((".dist-info", ".data")) and (child / "__init__.py").is_file() + ) + if len(candidates) == 1: + return candidates[0] + if not candidates: + raise ValueError(f"Could not find an importable top-level package in {extract_dir}.") + raise ValueError(f"Found multiple top-level packages in {extract_dir}: {candidates}.") + + @staticmethod + def _read_required_pyrit_version(*, extract_dir: Path) -> str | None: + import re + + for metadata in extract_dir.glob("*.dist-info/METADATA"): + try: + text = metadata.read_text(encoding="utf-8", errors="replace") + except OSError: + continue + for line in text.splitlines(): + if not re.match(r"Requires-Dist:\s*pyrit\b", line, flags=re.IGNORECASE): + continue + match = re.search(r"==\s*([0-9][0-9A-Za-z.\-]*)", line) + if match: + return match.group(1) + return None diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 4e61c7d400..09915393a6 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -24,15 +24,11 @@ from __future__ import annotations -import asyncio import importlib import inspect import logging -import os import pkgutil -import shutil import sys -import tempfile from pathlib import Path from typing import TYPE_CHECKING @@ -42,6 +38,7 @@ PluginRegisteredNothingError, PluginWheelNotFoundError, ) +from pyrit.setup.plugin_formats import WheelPluginFormat from pyrit.setup.pyrit_initializer import PyRITInitializer if TYPE_CHECKING: @@ -148,8 +145,6 @@ class PluginLoader: by default. """ - _FINGERPRINT_FILE = ".plugin_wheel_fingerprint" - def __init__(self, *, spec: PluginSpec) -> None: """ Initialize the loader for one plug-in. @@ -171,7 +166,7 @@ async def load_async(self) -> None: if wheel is None: raise PluginWheelNotFoundError("The current loader only accepts wheel-format plug-ins.") try: - await self._load_plugin_async(wheel_path=wheel.expanduser()) + await self._load_plugin_async() except Exception as exc: message = f"Failed to load plug-in from wheel '{wheel}': {exc} {_REMEDIATION}" # Preserve the specific failure type so callers can distinguish modes, while @@ -181,7 +176,7 @@ async def load_async(self) -> None: raise type(exc)(message) from exc raise PluginLoadError(message) from exc - async def _load_plugin_async(self, *, wheel_path: Path) -> None: + async def _load_plugin_async(self) -> None: """ Extract, import, bootstrap, and verify a single plug-in wheel. @@ -189,30 +184,17 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: registries) is rolled back if the load fails, so a failed or accepted-failure load leaves no partial trace. - Args: - wheel_path: Path to the pre-built plug-in wheel on disk. - Raises: - PluginWheelNotFoundError: If ``wheel_path`` is not an existing ``.whl`` file. + PluginWheelNotFoundError: If the configured wheel is absent or invalid. PluginImportError: If the plug-in package cannot be imported. PluginRegisteredNothingError: If the plug-in registered no datasets or scenarios. """ - if not wheel_path.is_file(): - raise PluginWheelNotFoundError(f"Plug-in wheel does not point to an existing file: {wheel_path}") - if wheel_path.suffix != ".whl": - raise PluginWheelNotFoundError(f"Plug-in wheel must be a .whl file, got: {wheel_path}") - - # Wheel extraction and directory scanning are blocking filesystem work; run them - # off the event loop so init does not stall unrelated async tasks. - extract_dir = await asyncio.to_thread(self._extract_wheel, wheel_path=wheel_path) - package_name = await asyncio.to_thread( - self._resolve_package_name, extract_dir=extract_dir, explicit_package=self._spec.package - ) + prepared = await WheelPluginFormat().prepare_async(spec=self._spec) + extract_dir = prepared.import_root + package_name = prepared.entry_modules[0] - from pyrit.setup.plugin_compat import bridge_scenario_extension_points, warn_on_version_drift + from pyrit.setup.plugin_compat import warn_on_version_drift - # Surface any host/plug-in version drift before importing, so a subsequent import or - # registration failure is read in the light of the mismatch (tolerate slight drift). warn_on_version_drift(extract_dir=extract_dir) from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider @@ -236,17 +218,6 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: except Exception as exc: raise PluginImportError(f"Could not import plug-in package '{package_name}': {exc}") from exc - # Bridge known renamed extension points so scenarios built against an older - # PyRIT are made concrete (and thus discoverable) before bootstrap/registration. - bridged = bridge_scenario_extension_points(package_name=package_name) - if bridged: - logger.warning( - "Applied %d plug-in compatibility shim(s) for '%s': %s", - len(bridged), - package_name, - ", ".join(bridged), - ) - await self._run_bootstrap_async(package_name=package_name, module=module) # Register the plug-in's own Scenario subclasses the same way built-ins are @@ -287,164 +258,6 @@ async def _load_plugin_async(self, *, wheel_path: Path) -> None: " (see PLUGIN DATASET SHADOWED warnings above)" if dataset_collisions else "", ) - def _extract_wheel(self, *, wheel_path: Path) -> Path: - """ - Extract the wheel into ``.plugin//``, reusing a fresh cached extraction. - - A cached extraction is reused only when its recorded fingerprint (wheel size and - modification time) matches the wheel on disk. A wheel rebuilt at the same path - (common during plug-in development, where the version-stamped filename is stable) - has a newer fingerprint, so the stale extraction is discarded and re-extracted - rather than silently reused — the exact silent-staleness trap the plug-in design - guards against. Extraction is atomic: the wheel is unpacked into a temporary - sibling directory and moved into place only on success, so a crash mid-extraction - never leaves a partial tree that would later be treated as a valid cache. - ``safe_extract_zip`` validates every member first (path traversal, symlinks, and - size / entry-count / compression caps) so a tampered wheel cannot escape the - extraction directory or exhaust disk. - - Args: - wheel_path: Path to the plug-in wheel. - - Returns: - Path: The directory the wheel was extracted to. - """ - from pyrit.common.safe_extract import safe_extract_zip - - base_dir = self._plugin_base_dir() - base_dir.mkdir(parents=True, exist_ok=True) - - extract_dir = base_dir / wheel_path.stem - fingerprint = self._wheel_fingerprint(wheel_path=wheel_path) - if extract_dir.is_dir() and any(extract_dir.iterdir()): - if self._cached_fingerprint(extract_dir=extract_dir) == fingerprint: - logger.info("Reusing cached plug-in extraction at %s", extract_dir) - return extract_dir - logger.warning( - "PLUGIN CACHE STALE: extraction at %s does not match the current wheel " - "(rebuilt or replaced); re-extracting '%s'.", - extract_dir, - wheel_path.name, - ) - - # Unique per-extraction temp dir so concurrent loads of the same wheel in one - # process cannot collide on a shared path (mkdtemp is atomic and 0700). safe_extract - # into it, then atomically move into place. - tmp_dir = Path(tempfile.mkdtemp(prefix=f".{wheel_path.stem}.tmp-", dir=base_dir)) - try: - safe_extract_zip(source=wheel_path, dest_dir=tmp_dir) - (tmp_dir / self._FINGERPRINT_FILE).write_text(fingerprint, encoding="utf-8") - if extract_dir.exists(): - shutil.rmtree(extract_dir) - os.replace(tmp_dir, extract_dir) - finally: - if tmp_dir.exists(): - shutil.rmtree(tmp_dir, ignore_errors=True) - - logger.info("Extracted plug-in wheel '%s' to %s", wheel_path.name, extract_dir) - return extract_dir - - @staticmethod - def _wheel_fingerprint(*, wheel_path: Path) -> str: - """ - Return a cheap change-detecting fingerprint (size and mtime) for a wheel. - - Args: - wheel_path: Path to the plug-in wheel. - - Returns: - str: A ``":"`` fingerprint that changes whenever the wheel is - rebuilt or replaced. - """ - stat = wheel_path.stat() - return f"{stat.st_size}:{stat.st_mtime_ns}" - - @classmethod - def _cached_fingerprint(cls, *, extract_dir: Path) -> str | None: - """ - Return the fingerprint recorded for a cached extraction, or ``None`` if absent. - - A missing marker (e.g. an extraction from before fingerprinting) reads as ``None`` - so the cache is treated as stale and re-extracted rather than trusted blindly. - - Args: - extract_dir: The cached extraction directory. - - Returns: - str | None: The recorded fingerprint, or ``None`` when no valid marker exists. - """ - marker = extract_dir / cls._FINGERPRINT_FILE - try: - return marker.read_text(encoding="utf-8").strip() - except OSError: - return None - - @staticmethod - def _plugin_base_dir() -> Path: - """ - Resolve the base directory for plug-in extractions. - - Uses ``PLUGIN_DIR`` when set, otherwise ``/.plugin``. - - Returns: - Path: The resolved plug-in base directory. - """ - override = os.getenv("PLUGIN_DIR") - if override: - return Path(override).expanduser().resolve() - from pyrit.common import path - - return Path(path.HOME_PATH, ".plugin").resolve() - - @staticmethod - def _resolve_package_name(*, extract_dir: Path, explicit_package: str | None) -> str: - """ - Determine the plug-in's top-level import package. - - Resolution order: the spec's ``package`` (when set), then ``*.dist-info/top_level.txt``, - then the single importable top-level directory in the extraction. - - Args: - extract_dir: The directory the wheel was extracted to. - explicit_package: The package name from the plug-in's config, if any. - - Returns: - str: The top-level package name to import. - - Raises: - ValueError: If the package cannot be unambiguously determined. - """ - if explicit_package: - return explicit_package - - for dist_info in sorted(extract_dir.glob("*.dist-info")): - top_level = dist_info / "top_level.txt" - if top_level.is_file(): - for line in top_level.read_text(encoding="utf-8").splitlines(): - name = line.strip() - if name: - return name - - candidates = sorted( - child.name - for child in extract_dir.iterdir() - if child.is_dir() - and not child.name.endswith(".dist-info") - and not child.name.endswith(".data") - and (child / "__init__.py").is_file() - ) - if len(candidates) == 1: - return candidates[0] - if not candidates: - raise ValueError( - f"Could not find an importable top-level package in {extract_dir}. " - "Set the plug-in's 'package' in its .pyrit_conf entry." - ) - raise ValueError( - f"Found multiple top-level packages in {extract_dir}: {candidates}. " - "Set the plug-in's 'package' in its .pyrit_conf entry to disambiguate." - ) - @staticmethod def _verify_module_location(*, module: ModuleType, extract_dir: Path, package_name: str) -> None: """ diff --git a/tests/unit/setup/test_plugin_formats.py b/tests/unit/setup/test_plugin_formats.py new file mode 100644 index 0000000000..242ab7c67e --- /dev/null +++ b/tests/unit/setup/test_plugin_formats.py @@ -0,0 +1,116 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import zipfile +from pathlib import Path + +import pytest + +from pyrit.exceptions import PluginWheelNotFoundError +from pyrit.setup import PluginFormat, PluginSpec +from pyrit.setup.plugin_formats import WheelPluginFormat + + +def _build_wheel(tmp_path: Path, *, package: str = "sample_plugin") -> Path: + wheel = tmp_path / f"{package}-1.0.0-py3-none-any.whl" + dist_info = f"{package}-1.0.0.dist-info" + with zipfile.ZipFile(wheel, "w") as archive: + archive.writestr(f"{package}/__init__.py", "") + archive.writestr(f"{dist_info}/top_level.txt", package) + archive.writestr( + f"{dist_info}/METADATA", + "\n".join( + [ + "Metadata-Version: 2.1", + f"Name: {package}", + "Version: 1.0.0", + "Requires-Dist: pyrit==0.14.0", + ] + ), + ) + return wheel + + +async def test_wheel_prepare_returns_common_artifact(tmp_path: Path) -> None: + wheel = _build_wheel(tmp_path) + spec = PluginSpec( + name="sample", + format=PluginFormat.WHEEL, + wheel=wheel, + package="sample_plugin", + ) + + prepared = await WheelPluginFormat(base_dir=tmp_path / ".plugin").prepare_async(spec=spec) + + assert prepared.spec is spec + assert prepared.import_root.is_dir() + assert prepared.entry_modules == ("sample_plugin",) + assert prepared.owned_module_prefixes == ("sample_plugin",) + assert prepared.artifact_fingerprint + assert prepared.declared_pyrit_version == "0.14.0" + assert (prepared.import_root / "sample_plugin" / "__init__.py").is_file() + + +async def test_wheel_prepare_resolves_package_from_metadata(tmp_path: Path) -> None: + wheel = _build_wheel(tmp_path, package="metadata_plugin") + spec = PluginSpec(name="metadata", format=PluginFormat.WHEEL, wheel=wheel) + + prepared = await WheelPluginFormat(base_dir=tmp_path / ".plugin").prepare_async(spec=spec) + + assert prepared.entry_modules == ("metadata_plugin",) + + +async def test_wheel_prepare_rejects_missing_wheel(tmp_path: Path) -> None: + spec = PluginSpec( + name="missing", + format=PluginFormat.WHEEL, + wheel=tmp_path / "missing.whl", + ) + + with pytest.raises(PluginWheelNotFoundError, match="existing"): + await WheelPluginFormat(base_dir=tmp_path / ".plugin").prepare_async(spec=spec) + + +async def test_wheel_prepare_rejects_source_spec(tmp_path: Path) -> None: + spec = PluginSpec( + name="source", + format=PluginFormat.SOURCE, + source=tmp_path / "plugin.py", + ) + + with pytest.raises(ValueError, match="wheel-format"): + await WheelPluginFormat(base_dir=tmp_path / ".plugin").prepare_async(spec=spec) + + +async def test_wheel_prepare_reuses_unchanged_extraction(tmp_path: Path) -> None: + wheel = _build_wheel(tmp_path) + spec = PluginSpec(name="sample", format=PluginFormat.WHEEL, wheel=wheel) + adapter = WheelPluginFormat(base_dir=tmp_path / ".plugin") + + first = await adapter.prepare_async(spec=spec) + marker = first.import_root / "cache_marker.txt" + marker.write_text("kept", encoding="utf-8") + second = await adapter.prepare_async(spec=spec) + + assert first.import_root == second.import_root + assert marker.is_file() + + +def test_wheel_package_resolution_prefers_explicit(tmp_path: Path) -> None: + assert ( + WheelPluginFormat._resolve_package( + extract_dir=tmp_path, + explicit_package="explicit_pkg", + ) + == "explicit_pkg" + ) + + +def test_wheel_package_resolution_rejects_ambiguous_directory(tmp_path: Path) -> None: + for name in ("pkg_a", "pkg_b"): + package_dir = tmp_path / name + package_dir.mkdir() + (package_dir / "__init__.py").write_text("", encoding="utf-8") + + with pytest.raises(ValueError, match="multiple"): + WheelPluginFormat._resolve_package(extract_dir=tmp_path, explicit_package=None) diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index 0186f3d989..dbf68dff00 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -33,7 +33,6 @@ from pyrit.setup.plugin_loader import ( PluginImportError, PluginInitializer, - PluginLoader, PluginLoadError, PluginRegisteredNothingError, PluginWheelNotFoundError, @@ -662,79 +661,6 @@ class _PreexistingScenario(RapidResponse): assert registry._classes[wheel.scenario_name] is _PreexistingScenario -# --------------------------------------------------------------------------- -# Extraction cache -# --------------------------------------------------------------------------- - - -def test_extract_wheel_reuses_cached_extraction(tmp_path: Path) -> None: - """A second extraction of an unchanged wheel reuses the cached directory.""" - wheel = build_mock_wheel(tmp_path) - plugin_dir = tmp_path / ".plugin" - - with plugin_env(PLUGIN_DIR=str(plugin_dir)): - initializer = PluginLoader(spec=_spec(wheel)) - first = initializer._extract_wheel(wheel_path=wheel.path) - marker = first / "cache_marker.txt" - marker.write_text("kept", encoding="utf-8") - - second = initializer._extract_wheel(wheel_path=wheel.path) - - assert first == second - assert marker.is_file() # not wiped -> cached, not re-extracted - - -# --------------------------------------------------------------------------- -# Package name resolution -# --------------------------------------------------------------------------- - - -def test_resolve_package_name_prefers_explicit(tmp_path: Path) -> None: - """An explicit package name takes precedence over inference.""" - (tmp_path / "some_pkg").mkdir() - (tmp_path / "some_pkg" / "__init__.py").write_text("", encoding="utf-8") - - assert PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package="explicit_pkg") == "explicit_pkg" - - -def test_resolve_package_name_infers_single_package(tmp_path: Path) -> None: - """The single importable top-level directory is inferred when no package/top_level.txt exists.""" - (tmp_path / "the_pkg").mkdir() - (tmp_path / "the_pkg" / "__init__.py").write_text("", encoding="utf-8") - (tmp_path / "the_pkg-0.0.1.dist-info").mkdir() - - assert PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) == "the_pkg" - - -def test_resolve_package_name_uses_top_level_txt(tmp_path: Path) -> None: - """top_level.txt is consulted before directory inference.""" - (tmp_path / "pkg_a").mkdir() - (tmp_path / "pkg_a" / "__init__.py").write_text("", encoding="utf-8") - (tmp_path / "pkg_b").mkdir() - (tmp_path / "pkg_b" / "__init__.py").write_text("", encoding="utf-8") - distinfo = tmp_path / "thing-0.0.1.dist-info" - distinfo.mkdir() - (distinfo / "top_level.txt").write_text("pkg_b\n", encoding="utf-8") - - assert PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) == "pkg_b" - - -def test_resolve_package_name_none_raises(tmp_path: Path) -> None: - """No importable package raises a clear error pointing at the plug-in's config.""" - with pytest.raises(ValueError, match="package"): - PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) - - -def test_resolve_package_name_multiple_raises(tmp_path: Path) -> None: - """Multiple top-level packages require an explicit package to disambiguate.""" - for name in ("pkg_a", "pkg_b"): - (tmp_path / name).mkdir() - (tmp_path / name / "__init__.py").write_text("", encoding="utf-8") - - with pytest.raises(ValueError, match="disambiguate"): - PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) - - # --------------------------------------------------------------------------- # Failure modes # --------------------------------------------------------------------------- From 6010b12d02debcb9ac8cf4f6d755d8ab0ff10c48 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 17:21:21 -0700 Subject: [PATCH 19/34] feat(setup): prepare source-file plug-ins Add a non-executing source adapter for one importable Python file with canonical ownership metadata and content fingerprinting behind the common PreparedPlugin contract. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/setup/plugin_formats.py | 50 +++++++++++++++++++++- tests/unit/setup/test_plugin_formats.py | 57 ++++++++++++++++++++++++- 2 files changed, 104 insertions(+), 3 deletions(-) diff --git a/pyrit/setup/plugin_formats.py b/pyrit/setup/plugin_formats.py index 675ec34aab..4c9bffdbd0 100644 --- a/pyrit/setup/plugin_formats.py +++ b/pyrit/setup/plugin_formats.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from pathlib import Path -from pyrit.exceptions import PluginWheelNotFoundError +from pyrit.exceptions import PluginSourceNotFoundError, PluginWheelNotFoundError from pyrit.setup.plugin_spec import PluginFormat, PluginSpec @@ -29,6 +29,54 @@ class PreparedPlugin: declared_pyrit_version: str | None = None +class SourcePluginFormat: + """Prepare a server-local Python source file for plug-in discovery.""" + + async def prepare_async(self, *, spec: PluginSpec) -> PreparedPlugin: + """ + Validate one source-file plug-in without importing or executing it. + + Args: + spec (PluginSpec): The normalized source declaration. + + Returns: + PreparedPlugin: The prepared source import root and ownership metadata. + + Raises: + ValueError: If the spec is not source format. + PluginSourceNotFoundError: If the source is absent or not a Python file. + """ + return await asyncio.to_thread(self._prepare, spec=spec) + + def _prepare(self, *, spec: PluginSpec) -> PreparedPlugin: + if spec.format is not PluginFormat.SOURCE or spec.source is None: + raise ValueError("SourcePluginFormat requires a source-format PluginSpec.") + + source = spec.source.expanduser().resolve() + if not source.is_file(): + raise PluginSourceNotFoundError(f"Plug-in source does not point to an existing Python file: {source}") + if source.suffix != ".py" or not source.stem.isidentifier(): + raise PluginSourceNotFoundError( + f"Plug-in source must be a Python file with an importable filename: {source}" + ) + + return PreparedPlugin( + spec=spec, + import_root=source.parent, + entry_modules=(source.stem,), + owned_module_prefixes=(source.stem,), + artifact_fingerprint=self._fingerprint(source=source), + ) + + @staticmethod + def _fingerprint(*, source: Path) -> str: + digest = hashlib.sha256() + with source.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + class WheelPluginFormat: """Prepare a pre-built Python wheel for plug-in discovery.""" diff --git a/tests/unit/setup/test_plugin_formats.py b/tests/unit/setup/test_plugin_formats.py index 242ab7c67e..e736e53568 100644 --- a/tests/unit/setup/test_plugin_formats.py +++ b/tests/unit/setup/test_plugin_formats.py @@ -6,9 +6,9 @@ import pytest -from pyrit.exceptions import PluginWheelNotFoundError +from pyrit.exceptions import PluginSourceNotFoundError, PluginWheelNotFoundError from pyrit.setup import PluginFormat, PluginSpec -from pyrit.setup.plugin_formats import WheelPluginFormat +from pyrit.setup.plugin_formats import SourcePluginFormat, WheelPluginFormat def _build_wheel(tmp_path: Path, *, package: str = "sample_plugin") -> Path: @@ -114,3 +114,56 @@ def test_wheel_package_resolution_rejects_ambiguous_directory(tmp_path: Path) -> with pytest.raises(ValueError, match="multiple"): WheelPluginFormat._resolve_package(extract_dir=tmp_path, explicit_package=None) + + +async def test_source_file_prepare_returns_common_artifact_without_execution(tmp_path: Path) -> None: + source = tmp_path / "operation_foobar.py" + marker = tmp_path / "executed.txt" + source.write_text( + f"from pathlib import Path\nPath({str(marker)!r}).write_text('executed')\n", + encoding="utf-8", + ) + spec = PluginSpec(name="operation_foobar", format=PluginFormat.SOURCE, source=source) + + prepared = await SourcePluginFormat().prepare_async(spec=spec) + + assert prepared.spec is spec + assert prepared.import_root == tmp_path.resolve() + assert prepared.entry_modules == ("operation_foobar",) + assert prepared.owned_module_prefixes == ("operation_foobar",) + assert prepared.artifact_fingerprint + assert prepared.declared_pyrit_version is None + assert not marker.exists() + + +async def test_source_file_prepare_rejects_missing_file(tmp_path: Path) -> None: + spec = PluginSpec( + name="missing", + format=PluginFormat.SOURCE, + source=tmp_path / "missing.py", + ) + + with pytest.raises(PluginSourceNotFoundError, match="existing"): + await SourcePluginFormat().prepare_async(spec=spec) + + +async def test_source_file_prepare_rejects_non_python_file(tmp_path: Path) -> None: + source = tmp_path / "plugin.txt" + source.write_text("not Python", encoding="utf-8") + spec = PluginSpec(name="plugin", format=PluginFormat.SOURCE, source=source) + + with pytest.raises(PluginSourceNotFoundError, match="Python"): + await SourcePluginFormat().prepare_async(spec=spec) + + +async def test_source_file_fingerprint_changes_with_content(tmp_path: Path) -> None: + source = tmp_path / "plugin.py" + source.write_text("VALUE = 1\n", encoding="utf-8") + spec = PluginSpec(name="plugin", format=PluginFormat.SOURCE, source=source) + adapter = SourcePluginFormat() + + first = await adapter.prepare_async(spec=spec) + source.write_text("VALUE = 2\n", encoding="utf-8") + second = await adapter.prepare_async(spec=spec) + + assert first.artifact_fingerprint != second.artifact_fingerprint From c5d05072dea588fbfba0faa957ac0db2ed86a793 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 17:26:22 -0700 Subject: [PATCH 20/34] feat(setup): prepare source-package plug-ins Extend source format preparation to proper Python packages, validate nested entry modules and ownership, reject loose directories, and fingerprint package trees deterministically. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/setup/plugin_formats.py | 56 +++++++++++++++++- tests/unit/setup/test_plugin_formats.py | 77 +++++++++++++++++++++++++ 2 files changed, 131 insertions(+), 2 deletions(-) diff --git a/pyrit/setup/plugin_formats.py b/pyrit/setup/plugin_formats.py index 4c9bffdbd0..9bf6e553a8 100644 --- a/pyrit/setup/plugin_formats.py +++ b/pyrit/setup/plugin_formats.py @@ -53,12 +53,21 @@ def _prepare(self, *, spec: PluginSpec) -> PreparedPlugin: raise ValueError("SourcePluginFormat requires a source-format PluginSpec.") source = spec.source.expanduser().resolve() - if not source.is_file(): - raise PluginSourceNotFoundError(f"Plug-in source does not point to an existing Python file: {source}") + if source.is_file(): + return self._prepare_file(spec=spec, source=source) + if source.is_dir(): + return self._prepare_package(spec=spec, source=source) + raise PluginSourceNotFoundError(f"Plug-in source does not point to an existing path: {source}") + + def _prepare_file(self, *, spec: PluginSpec, source: Path) -> PreparedPlugin: if source.suffix != ".py" or not source.stem.isidentifier(): raise PluginSourceNotFoundError( f"Plug-in source must be a Python file with an importable filename: {source}" ) + if spec.package is not None and spec.package != source.stem: + raise ValueError( + f"Single-file plug-in package must match the source module '{source.stem}', got '{spec.package}'." + ) return PreparedPlugin( spec=spec, @@ -68,6 +77,36 @@ def _prepare(self, *, spec: PluginSpec) -> PreparedPlugin: artifact_fingerprint=self._fingerprint(source=source), ) + def _prepare_package(self, *, spec: PluginSpec, source: Path) -> PreparedPlugin: + package_name = source.name + if not package_name.isidentifier() or not (source / "__init__.py").is_file(): + raise PluginSourceNotFoundError( + f"Plug-in source directory must be an importable package containing __init__.py: {source}" + ) + + entry_module = spec.package or package_name + if entry_module != package_name and not entry_module.startswith(f"{package_name}."): + raise ValueError(f"Plug-in entry module '{entry_module}' must be inside source package '{package_name}'.") + self._validate_entry_module(source=source, entry_module=entry_module) + + return PreparedPlugin( + spec=spec, + import_root=source.parent, + entry_modules=(entry_module,), + owned_module_prefixes=(package_name,), + artifact_fingerprint=self._fingerprint_package(source=source), + ) + + @staticmethod + def _validate_entry_module(*, source: Path, entry_module: str) -> None: + relative_parts = entry_module.split(".")[1:] + if not relative_parts: + return + module_path = source.joinpath(*relative_parts) + if module_path.with_suffix(".py").is_file() or (module_path / "__init__.py").is_file(): + return + raise ValueError(f"Plug-in entry module '{entry_module}' does not exist inside {source}.") + @staticmethod def _fingerprint(*, source: Path) -> str: digest = hashlib.sha256() @@ -76,6 +115,19 @@ def _fingerprint(*, source: Path) -> str: digest.update(chunk) return digest.hexdigest() + @classmethod + def _fingerprint_package(cls, *, source: Path) -> str: + digest = hashlib.sha256() + for path in sorted(source.rglob("*")): + if not path.is_file() or "__pycache__" in path.parts or path.suffix == ".pyc": + continue + resolved = path.resolve() + if not resolved.is_relative_to(source): + raise PluginSourceNotFoundError(f"Plug-in package file escapes the source root: {path}") + digest.update(path.relative_to(source).as_posix().encode("utf-8")) + digest.update(cls._fingerprint(source=path).encode("ascii")) + return digest.hexdigest() + class WheelPluginFormat: """Prepare a pre-built Python wheel for plug-in discovery.""" diff --git a/tests/unit/setup/test_plugin_formats.py b/tests/unit/setup/test_plugin_formats.py index e736e53568..4747402e5e 100644 --- a/tests/unit/setup/test_plugin_formats.py +++ b/tests/unit/setup/test_plugin_formats.py @@ -167,3 +167,80 @@ async def test_source_file_fingerprint_changes_with_content(tmp_path: Path) -> N second = await adapter.prepare_async(spec=spec) assert first.artifact_fingerprint != second.artifact_fingerprint + + +async def test_source_package_prepare_returns_package_ownership(tmp_path: Path) -> None: + package = tmp_path / "operation_plugin" + package.mkdir() + (package / "__init__.py").write_text("", encoding="utf-8") + (package / "scenarios.py").write_text("VALUE = 1\n", encoding="utf-8") + spec = PluginSpec( + name="operation", + format=PluginFormat.SOURCE, + source=package, + ) + + prepared = await SourcePluginFormat().prepare_async(spec=spec) + + assert prepared.import_root == tmp_path.resolve() + assert prepared.entry_modules == ("operation_plugin",) + assert prepared.owned_module_prefixes == ("operation_plugin",) + + +async def test_source_package_accepts_explicit_nested_entry_module(tmp_path: Path) -> None: + package = tmp_path / "operation_plugin" + package.mkdir() + (package / "__init__.py").write_text("", encoding="utf-8") + (package / "entry.py").write_text("", encoding="utf-8") + spec = PluginSpec( + name="operation", + format=PluginFormat.SOURCE, + source=package, + package="operation_plugin.entry", + ) + + prepared = await SourcePluginFormat().prepare_async(spec=spec) + + assert prepared.entry_modules == ("operation_plugin.entry",) + assert prepared.owned_module_prefixes == ("operation_plugin",) + + +async def test_source_package_rejects_loose_directory(tmp_path: Path) -> None: + source = tmp_path / "loose" + source.mkdir() + (source / "scenario.py").write_text("", encoding="utf-8") + spec = PluginSpec(name="loose", format=PluginFormat.SOURCE, source=source) + + with pytest.raises(PluginSourceNotFoundError, match="__init__.py"): + await SourcePluginFormat().prepare_async(spec=spec) + + +async def test_source_package_rejects_foreign_entry_module(tmp_path: Path) -> None: + package = tmp_path / "operation_plugin" + package.mkdir() + (package / "__init__.py").write_text("", encoding="utf-8") + spec = PluginSpec( + name="operation", + format=PluginFormat.SOURCE, + source=package, + package="other_plugin", + ) + + with pytest.raises(ValueError, match="inside"): + await SourcePluginFormat().prepare_async(spec=spec) + + +async def test_source_package_fingerprint_changes_with_nested_content(tmp_path: Path) -> None: + package = tmp_path / "operation_plugin" + package.mkdir() + (package / "__init__.py").write_text("", encoding="utf-8") + nested = package / "scenario.py" + nested.write_text("VALUE = 1\n", encoding="utf-8") + spec = PluginSpec(name="operation", format=PluginFormat.SOURCE, source=package) + adapter = SourcePluginFormat() + + first = await adapter.prepare_async(spec=spec) + nested.write_text("VALUE = 2\n", encoding="utf-8") + second = await adapter.prepare_async(spec=spec) + + assert first.artifact_fingerprint != second.artifact_fingerprint From 8ead0c5333d58fe5f3daf5866af79c0236403f41 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 17:33:51 -0700 Subject: [PATCH 21/34] feat(setup): discover source plug-in scenarios Import prepared source or wheel modules through a shared ownership-checked path and discover concrete module-owned Scenario subclasses with deterministic source-relative dotted names and ambiguity errors. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/setup/plugin_discovery.py | 173 ++++++++++++++++++++++ tests/unit/setup/test_plugin_discovery.py | 105 +++++++++++++ 2 files changed, 278 insertions(+) create mode 100644 pyrit/setup/plugin_discovery.py create mode 100644 tests/unit/setup/test_plugin_discovery.py diff --git a/pyrit/setup/plugin_discovery.py b/pyrit/setup/plugin_discovery.py new file mode 100644 index 0000000000..1910489811 --- /dev/null +++ b/pyrit/setup/plugin_discovery.py @@ -0,0 +1,173 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Import prepared plug-ins and discover owned PyRIT contributions.""" + +from __future__ import annotations + +import asyncio +import importlib +import inspect +import pkgutil +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from pyrit.exceptions import PluginDiscoveryError, PluginImportError +from pyrit.models import class_name_to_snake_case +from pyrit.scenario import Scenario + +if TYPE_CHECKING: + from types import ModuleType + + from pyrit.setup.plugin_formats import PreparedPlugin + + +@dataclass(frozen=True) +class ImportedPlugin: + """The imported modules owned by one prepared plug-in.""" + + prepared: PreparedPlugin + modules: tuple[ModuleType, ...] + + +@dataclass(frozen=True) +class ScenarioContribution: + """A discovered scenario class and its proposed registry name.""" + + scenario_class: type[Scenario] + registry_name: str + + +async def import_plugin_async(*, prepared: PreparedPlugin) -> ImportedPlugin: + """ + Import a prepared plug-in and all submodules under package entry modules. + + Args: + prepared (PreparedPlugin): The prepared source or wheel import root. + + Returns: + ImportedPlugin: The imported, ownership-validated modules. + + Raises: + PluginImportError: If import or ownership validation fails. + """ + return await asyncio.to_thread(_import_plugin, prepared=prepared) + + +def _import_plugin(*, prepared: PreparedPlugin) -> ImportedPlugin: + import_root = str(prepared.import_root) + if import_root not in sys.path: + sys.path.insert(0, import_root) + + try: + for entry_module in prepared.entry_modules: + module = importlib.import_module(entry_module) + _verify_module_location(module=module, import_root=prepared.import_root) + _import_submodules(module=module) + except Exception as exc: + raise PluginImportError( + f"Could not import plug-in '{prepared.spec.name}' from {prepared.import_root}: {exc}" + ) from exc + + modules = tuple( + module + for name, module in sorted(sys.modules.items()) + if module is not None and _name_owned_by_any(name=name, prefixes=prepared.owned_module_prefixes) + ) + for module in modules: + _verify_module_location(module=module, import_root=prepared.import_root) + return ImportedPlugin(prepared=prepared, modules=modules) + + +def _import_submodules(*, module: ModuleType) -> None: + module_path = getattr(module, "__path__", None) + if not module_path: + return + for item in pkgutil.walk_packages(module_path, prefix=f"{module.__name__}."): + importlib.import_module(item.name) + + +def _verify_module_location(*, module: ModuleType, import_root: Path) -> None: + root = import_root.resolve() + raw_locations = list(getattr(module, "__path__", []) or []) + module_file = getattr(module, "__file__", None) + if module_file: + raw_locations.append(module_file) + locations = [Path(location).resolve() for location in raw_locations if location] + if locations and not any(location.is_relative_to(root) for location in locations): + raise ValueError(f"Imported module '{module.__name__}' resolved outside plug-in root {root}.") + + +def discover_scenarios(*, imported: ImportedPlugin) -> list[ScenarioContribution]: + """ + Discover concrete, module-owned Scenario subclasses. + + Args: + imported (ImportedPlugin): The imported plug-in modules. + + Returns: + list[ScenarioContribution]: Deterministically ordered scenario contributions. + + Raises: + PluginDiscoveryError: If one module defines multiple implicit scenarios. + """ + contributions: list[ScenarioContribution] = [] + for module in imported.modules: + classes = _module_defined_subclasses(module=module, base_class=Scenario) + if len(classes) > 1: + raise PluginDiscoveryError( + f"Plug-in module '{module.__name__}' defines multiple Scenario classes. " + "Provide explicit registry names through the plug-in manifest." + ) + if classes: + contributions.append( + ScenarioContribution( + scenario_class=classes[0], + registry_name=_scenario_registry_name(imported=imported, module=module, cls=classes[0]), + ) + ) + return sorted(contributions, key=lambda item: item.registry_name) + + +def _module_defined_subclasses(*, module: ModuleType, base_class: type[Scenario]) -> list[type[Scenario]]: + return [ + candidate + for _, candidate in inspect.getmembers(module, inspect.isclass) + if candidate is not base_class + and issubclass(candidate, base_class) + and not inspect.isabstract(candidate) + and candidate.__module__ == module.__name__ + ] + + +def _scenario_registry_name( + *, + imported: ImportedPlugin, + module: ModuleType, + cls: type[Scenario], +) -> str: + prefix = next( + ( + owned + for owned in sorted(imported.prepared.owned_module_prefixes, key=len, reverse=True) + if module.__name__ == owned or module.__name__.startswith(f"{owned}.") + ), + "", + ) + relative = module.__name__[len(prefix) :].lstrip(".") if prefix else module.__name__ + for marker in ("scenario.scenarios.", "scenarios."): + if marker in relative: + relative = relative.split(marker, 1)[1] + break + if relative: + return relative + source = imported.prepared.spec.source + if source is not None and source.is_file(): + return source.stem + return class_name_to_snake_case(cls.__name__, suffix="Scenario") + + +def _name_owned_by_any(*, name: str, prefixes: tuple[str, ...]) -> bool: + return any(name == prefix or name.startswith(f"{prefix}.") for prefix in prefixes) diff --git a/tests/unit/setup/test_plugin_discovery.py b/tests/unit/setup/test_plugin_discovery.py new file mode 100644 index 0000000000..3cb5070475 --- /dev/null +++ b/tests/unit/setup/test_plugin_discovery.py @@ -0,0 +1,105 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import sys +from collections.abc import Iterator +from pathlib import Path + +import pytest + +from pyrit.exceptions import PluginDiscoveryError +from pyrit.setup import PluginFormat, PluginSpec +from pyrit.setup.plugin_discovery import discover_scenarios, import_plugin_async +from pyrit.setup.plugin_formats import SourcePluginFormat + + +@pytest.fixture +def restore_import_state() -> Iterator[None]: + path_snapshot = list(sys.path) + module_snapshot = set(sys.modules) + yield + sys.path[:] = path_snapshot + for name in set(sys.modules) - module_snapshot: + del sys.modules[name] + + +async def _prepare_and_import_source(source: Path, *, name: str): + spec = PluginSpec(name=name, format=PluginFormat.SOURCE, source=source) + prepared = await SourcePluginFormat().prepare_async(spec=spec) + return await import_plugin_async(prepared=prepared) + + +@pytest.mark.usefixtures("restore_import_state") +async def test_discover_scenario_from_single_file(tmp_path: Path) -> None: + source = tmp_path / "private_scenario.py" + source.write_text( + """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + +class PrivateScenario(RapidResponse): + pass +""", + encoding="utf-8", + ) + + imported = await _prepare_and_import_source(source, name="private") + contributions = discover_scenarios(imported=imported) + + assert len(contributions) == 1 + assert contributions[0].scenario_class.__name__ == "PrivateScenario" + assert contributions[0].registry_name == "private_scenario" + + +@pytest.mark.usefixtures("restore_import_state") +async def test_discover_scenario_ignores_imported_foreign_class(tmp_path: Path) -> None: + source = tmp_path / "foreign_only.py" + source.write_text( + "from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse\n", + encoding="utf-8", + ) + + imported = await _prepare_and_import_source(source, name="foreign_only") + + assert discover_scenarios(imported=imported) == [] + + +@pytest.mark.usefixtures("restore_import_state") +async def test_discover_scenario_rejects_multiple_classes_in_one_module(tmp_path: Path) -> None: + source = tmp_path / "ambiguous.py" + source.write_text( + """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + +class FirstScenario(RapidResponse): + pass + +class SecondScenario(RapidResponse): + pass +""", + encoding="utf-8", + ) + + imported = await _prepare_and_import_source(source, name="ambiguous") + + with pytest.raises(PluginDiscoveryError, match="multiple"): + discover_scenarios(imported=imported) + + +@pytest.mark.usefixtures("restore_import_state") +async def test_discover_scenario_uses_source_relative_dotted_name(tmp_path: Path) -> None: + package = tmp_path / "operation_plugin" + scenarios = package / "scenarios" / "image" + scenarios.mkdir(parents=True) + for init_file in (package / "__init__.py", package / "scenarios" / "__init__.py", scenarios / "__init__.py"): + init_file.write_text("", encoding="utf-8") + (scenarios / "abuse.py").write_text( + """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + +class AbuseScenario(RapidResponse): + pass +""", + encoding="utf-8", + ) + + imported = await _prepare_and_import_source(package, name="operation") + contributions = discover_scenarios(imported=imported) + + assert [item.registry_name for item in contributions] == ["image.abuse"] From 9f0c0c83355a87886d464153006d88b5c0feba05 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 18:10:01 -0700 Subject: [PATCH 22/34] feat(setup): discover plug-in attack techniques Discover explicit, conventional, and module-global AttackTechniqueFactory contributions without invoking arbitrary helpers, require scenario applicability, validate identifiers, and reject duplicate technique names. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/setup/plugin_discovery.py | 119 ++++++++++++++++- tests/unit/setup/test_plugin_discovery.py | 155 +++++++++++++++++++++- 2 files changed, 271 insertions(+), 3 deletions(-) diff --git a/pyrit/setup/plugin_discovery.py b/pyrit/setup/plugin_discovery.py index 1910489811..9767977eaa 100644 --- a/pyrit/setup/plugin_discovery.py +++ b/pyrit/setup/plugin_discovery.py @@ -12,13 +12,14 @@ import sys from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from pyrit.exceptions import PluginDiscoveryError, PluginImportError from pyrit.models import class_name_to_snake_case -from pyrit.scenario import Scenario +from pyrit.scenario import AttackTechniqueFactory, Scenario if TYPE_CHECKING: + from collections.abc import Callable, Sequence from types import ModuleType from pyrit.setup.plugin_formats import PreparedPlugin @@ -40,6 +41,14 @@ class ScenarioContribution: registry_name: str +@dataclass(frozen=True) +class TechniqueContribution: + """A configured attack technique and the scenarios that expose it.""" + + factory: AttackTechniqueFactory + scenario_names: frozenset[str] + + async def import_plugin_async(*, prepared: PreparedPlugin) -> ImportedPlugin: """ Import a prepared plug-in and all submodules under package entry modules. @@ -131,6 +140,112 @@ def discover_scenarios(*, imported: ImportedPlugin) -> list[ScenarioContribution return sorted(contributions, key=lambda item: item.registry_name) +def discover_techniques(*, imported: ImportedPlugin) -> list[TechniqueContribution]: + """ + Discover configured attack-technique factories from owned modules. + + Args: + imported (ImportedPlugin): The imported plug-in modules. + + Returns: + list[TechniqueContribution]: Deterministically ordered technique contributions. + + Raises: + PluginDiscoveryError: If a factory export is malformed, lacks applicability, + or duplicates another contribution name. + """ + contributions: list[TechniqueContribution] = [] + seen_objects: set[int] = set() + for module in imported.modules: + explicit = [ + value + for value in vars(module).values() + if isinstance(value, TechniqueContribution) and id(value.factory) not in seen_objects + ] + contributions.extend(explicit) + seen_objects.update(id(item.factory) for item in explicit) + + factory_builder = getattr(module, "get_technique_factories", None) + if callable(factory_builder) and getattr(factory_builder, "__module__", None) == module.__name__: + contributions.extend( + _contributions_from_factory_builder( + module=module, + factory_builder=factory_builder, + seen_objects=seen_objects, + ) + ) + + for value in vars(module).values(): + if isinstance(value, AttackTechniqueFactory) and id(value) not in seen_objects: + contributions.append(_contribution_from_factory(factory=value, module_name=module.__name__)) + seen_objects.add(id(value)) + + by_name: dict[str, TechniqueContribution] = {} + for contribution in contributions: + if not contribution.scenario_names: + raise PluginDiscoveryError( + f"Attack technique '{contribution.factory.name}' must declare at least one applicable scenario." + ) + existing = by_name.get(contribution.factory.name) + if existing is not None and existing.factory is not contribution.factory: + raise PluginDiscoveryError( + f"Plug-in discovered duplicate attack technique name '{contribution.factory.name}'." + ) + contribution.factory.get_identifier() + by_name[contribution.factory.name] = contribution + return [by_name[name] for name in sorted(by_name)] + + +def _contributions_from_factory_builder( + *, + module: ModuleType, + factory_builder: Callable[[], object], + seen_objects: set[int], +) -> list[TechniqueContribution]: + signature = inspect.signature(factory_builder) + required = [ + parameter + for parameter in signature.parameters.values() + if parameter.default is inspect.Parameter.empty + and parameter.kind + not in { + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + } + ] + if required: + raise PluginDiscoveryError( + f"Plug-in function '{module.__name__}.get_technique_factories' must take no required arguments." + ) + result = factory_builder() + if not isinstance(result, (list, tuple)) or not all(isinstance(item, AttackTechniqueFactory) for item in result): + raise PluginDiscoveryError( + f"Plug-in function '{module.__name__}.get_technique_factories' must return " + "AttackTechniqueFactory instances." + ) + factories = cast("Sequence[AttackTechniqueFactory]", result) + contributions = [ + _contribution_from_factory(factory=factory, module_name=module.__name__) + for factory in factories + if id(factory) not in seen_objects + ] + seen_objects.update(id(item.factory) for item in contributions) + return contributions + + +def _contribution_from_factory(*, factory: AttackTechniqueFactory, module_name: str) -> TechniqueContribution: + prefix = "scenario:" + scenario_names = frozenset( + tag.removeprefix(prefix) for tag in factory.strategy_tags if tag.startswith(prefix) and tag != prefix + ) + if not scenario_names: + raise PluginDiscoveryError( + f"Attack technique '{factory.name}' from '{module_name}' must declare an applicable scenario " + f"using a '{prefix}' tag or an explicit TechniqueContribution." + ) + return TechniqueContribution(factory=factory, scenario_names=scenario_names) + + def _module_defined_subclasses(*, module: ModuleType, base_class: type[Scenario]) -> list[type[Scenario]]: return [ candidate diff --git a/tests/unit/setup/test_plugin_discovery.py b/tests/unit/setup/test_plugin_discovery.py index 3cb5070475..2d582ba87f 100644 --- a/tests/unit/setup/test_plugin_discovery.py +++ b/tests/unit/setup/test_plugin_discovery.py @@ -9,7 +9,11 @@ from pyrit.exceptions import PluginDiscoveryError from pyrit.setup import PluginFormat, PluginSpec -from pyrit.setup.plugin_discovery import discover_scenarios, import_plugin_async +from pyrit.setup.plugin_discovery import ( + discover_scenarios, + discover_techniques, + import_plugin_async, +) from pyrit.setup.plugin_formats import SourcePluginFormat @@ -103,3 +107,152 @@ class AbuseScenario(RapidResponse): contributions = discover_scenarios(imported=imported) assert [item.registry_name for item in contributions] == ["image.abuse"] + + +@pytest.mark.usefixtures("restore_import_state") +async def test_discover_technique_from_conventional_factory_function(tmp_path: Path) -> None: + source = tmp_path / "private_technique.py" + source.write_text( + """from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory + +def get_technique_factories(): + return [ + AttackTechniqueFactory( + name="operation_foobar", + attack_class=PromptSendingAttack, + strategy_tags=["single_turn", "scenario:airt.rapid_response"], + ) + ] +""", + encoding="utf-8", + ) + + imported = await _prepare_and_import_source(source, name="private_technique") + contributions = discover_techniques(imported=imported) + + assert len(contributions) == 1 + assert contributions[0].factory.name == "operation_foobar" + assert contributions[0].scenario_names == frozenset({"airt.rapid_response"}) + + +@pytest.mark.usefixtures("restore_import_state") +async def test_discover_technique_from_explicit_contribution(tmp_path: Path) -> None: + source = tmp_path / "explicit_technique.py" + source.write_text( + """from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory +from pyrit.setup.plugin_discovery import TechniqueContribution + +OPERATION = TechniqueContribution( + factory=AttackTechniqueFactory( + name="operation_explicit", + attack_class=PromptSendingAttack, + strategy_tags=["single_turn"], + ), + scenario_names=frozenset({"airt.rapid_response"}), +) +""", + encoding="utf-8", + ) + + imported = await _prepare_and_import_source(source, name="explicit_technique") + contributions = discover_techniques(imported=imported) + + assert [item.factory.name for item in contributions] == ["operation_explicit"] + + +@pytest.mark.usefixtures("restore_import_state") +async def test_discover_technique_from_module_global_factory(tmp_path: Path) -> None: + source = tmp_path / "global_technique.py" + source.write_text( + """from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory + +OPERATION = AttackTechniqueFactory( + name="operation_global", + attack_class=PromptSendingAttack, + strategy_tags=["single_turn", "scenario:airt.rapid_response"], +) +""", + encoding="utf-8", + ) + + imported = await _prepare_and_import_source(source, name="global_technique") + contributions = discover_techniques(imported=imported) + + assert [item.factory.name for item in contributions] == ["operation_global"] + + +@pytest.mark.usefixtures("restore_import_state") +async def test_discover_technique_does_not_call_arbitrary_helpers(tmp_path: Path) -> None: + source = tmp_path / "safe_discovery.py" + marker = tmp_path / "helper_called.txt" + source.write_text( + f"""from pathlib import Path +from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory + +def build_something(): + Path({str(marker)!r}).write_text("called") + return AttackTechniqueFactory(name="hidden", attack_class=PromptSendingAttack) + +OPERATION = AttackTechniqueFactory( + name="visible", + attack_class=PromptSendingAttack, + strategy_tags=["scenario:airt.rapid_response"], +) +""", + encoding="utf-8", + ) + + imported = await _prepare_and_import_source(source, name="safe_discovery") + contributions = discover_techniques(imported=imported) + + assert [item.factory.name for item in contributions] == ["visible"] + assert not marker.exists() + + +@pytest.mark.usefixtures("restore_import_state") +async def test_discover_technique_requires_scenario_applicability(tmp_path: Path) -> None: + source = tmp_path / "missing_applicability.py" + source.write_text( + """from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory + +OPERATION = AttackTechniqueFactory(name="missing_scope", attack_class=PromptSendingAttack) +""", + encoding="utf-8", + ) + + imported = await _prepare_and_import_source(source, name="missing_applicability") + + with pytest.raises(PluginDiscoveryError, match="applicable scenario"): + discover_techniques(imported=imported) + + +@pytest.mark.usefixtures("restore_import_state") +async def test_discover_technique_rejects_duplicate_names(tmp_path: Path) -> None: + source = tmp_path / "duplicate_techniques.py" + source.write_text( + """from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory + +FIRST = AttackTechniqueFactory( + name="duplicate", + attack_class=PromptSendingAttack, + strategy_tags=["scenario:airt.rapid_response"], +) +SECOND = AttackTechniqueFactory( + name="duplicate", + attack_class=PromptSendingAttack, + strategy_tags=["scenario:airt.rapid_response"], +) +""", + encoding="utf-8", + ) + + imported = await _prepare_and_import_source(source, name="duplicate_techniques") + + with pytest.raises(PluginDiscoveryError, match="duplicate"): + discover_techniques(imported=imported) From 153b2c0e62df87969503a6f531f83066ec1e6b6a Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 19:41:57 -0700 Subject: [PATCH 23/34] feat(setup): register plug-in components transactionally Replace plug-in-authored bootstrap behavior with framework-owned source/wheel discovery, strict extend-only scenario and technique registrars, latent built-in collision checks, unsupported provider cleanup, and fail-closed registry/module rollback. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- .../components/attack_technique_registry.py | 37 +- .../registry/components/scenario_registry.py | 18 + pyrit/setup/plugin_loader.py | 627 ++++++------------ tests/unit/setup/test_plugin_loader.py | 212 +----- tests/unit/setup/test_plugin_registration.py | 155 +++++ 5 files changed, 424 insertions(+), 625 deletions(-) create mode 100644 tests/unit/setup/test_plugin_registration.py diff --git a/pyrit/registry/components/attack_technique_registry.py b/pyrit/registry/components/attack_technique_registry.py index 2d3ca12c24..2880ec4f87 100644 --- a/pyrit/registry/components/attack_technique_registry.py +++ b/pyrit/registry/components/attack_technique_registry.py @@ -113,6 +113,7 @@ def register_technique( name: str, factory: AttackTechniqueFactory, tags: dict[str, str] | list[str] | None = None, + metadata: dict[str, object] | None = None, ) -> None: """ Register an attack technique factory. @@ -123,10 +124,44 @@ def register_technique( tags (dict[str, str] | list[str] | None): Optional tags for categorisation. Accepts a ``dict[str, str]`` or a ``list[str]`` (each string becomes a key with value ``""``). + metadata (dict[str, object] | None): Optional non-selection metadata. """ - self.instances.register(factory, name=name, tags=tags) + self.instances.register(factory, name=name, tags=tags, metadata=metadata) logger.debug(f"Registered attack technique factory: {name} ({factory.attack_class.__name__})") + def register_contributed_factory( + self, + *, + factory: AttackTechniqueFactory, + plugin_name: str, + scenario_names: frozenset[str], + ) -> None: + """ + Register one extend-only plug-in technique contribution. + + Args: + factory (AttackTechniqueFactory): The configured attack technique. + plugin_name (str): The owning plug-in name. + scenario_names (frozenset[str]): Scenarios that expose the technique. + + Raises: + PluginCollisionError: If the technique name is already registered. + """ + from pyrit.exceptions import PluginCollisionError + + if factory.name in self.instances: + raise PluginCollisionError(f"Attack technique name '{factory.name}' is already registered.") + self.register_technique( + name=factory.name, + factory=factory, + tags=factory.strategy_tags, + metadata={ + "plugin_name": plugin_name, + "scenario_names": sorted(scenario_names), + "identifier_hash": factory.get_identifier().hash, + }, + ) + def get_factories(self) -> dict[str, AttackTechniqueFactory]: """ Return all registered factories as a name→factory dict. diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py index af4284506d..526dd7dc87 100644 --- a/pyrit/registry/components/scenario_registry.py +++ b/pyrit/registry/components/scenario_registry.py @@ -138,6 +138,24 @@ def _external_registry_name(self, cls: type[Scenario], *, package_name: str) -> break return relative or class_name_to_snake_case(cls.__name__, suffix="Scenario") + def register_contributed_scenario(self, *, scenario_class: type[Scenario], name: str) -> None: + """ + Register one extend-only plug-in scenario contribution. + + Args: + scenario_class (type[Scenario]): The concrete scenario class. + name (str): The scanner-facing dotted registry name. + + Raises: + PluginCollisionError: If the name belongs to a built-in or prior contribution. + """ + from pyrit.exceptions import PluginCollisionError + + self._ensure_discovered() + if name in self._classes: + raise PluginCollisionError(f"Scenario name '{name}' is already registered.") + self.register_class(scenario_class, name=name) + def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: """ Build metadata for a Scenario class. diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 09915393a6..9128f374bc 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -1,111 +1,64 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -""" -Load non-disclosable PyRIT plug-ins from pre-built wheels at initialization time. - -A plug-in is a pure-Python wheel that ships dataset providers and/or scenarios that -must not live in the public PyRIT repo. The loader **extracts** the wheel (stdlib -``zipfile`` — never ``pip``/``.venv``) into ``.plugin//``, prepends that directory -to ``sys.path``, imports the package (so ``SeedDatasetProvider`` subclasses self-register), -and runs the plug-in's bootstrap (a top-level ``register()`` callable or a shipped -``PyRITInitializer`` subclass). The plug-in's own ``Scenario`` subclasses are then -auto-registered by type (scoped to the plug-in package), so a plug-in's scenarios are -discovered without the bootstrap having to register each one explicitly. - -Plug-ins are declared in ``.pyrit_conf`` under the dedicated ``plugins`` key (a list, so -several plug-ins load identically to one). ``load_plugins_if_configured_async`` is invoked -as a guaranteed-first phase inside ``initialize_pyrit_async`` — after central memory is set -and **before** any configured initializer runs — so plug-in datasets and scenarios are -registered before ``LoadDefaultDatasets`` / ``PreloadScenarioMetadata`` read the registry. -Because ``plugins`` is its own always-first phase (not one of the ordered ``initializers``), -this ordering is true by construction. It is a no-op when no plug-ins are configured. -""" +"""Activate private scenarios and attack techniques from configured plug-ins.""" from __future__ import annotations -import importlib -import inspect import logging -import pkgutil import sys -from pathlib import Path -from typing import TYPE_CHECKING +from dataclasses import dataclass +from typing import TYPE_CHECKING, cast -from pyrit.exceptions.exception_classes import ( - PluginImportError, +from pyrit.datasets import SeedDatasetProvider +from pyrit.exceptions import ( + PluginCollisionError, PluginLoadError, PluginRegisteredNothingError, - PluginWheelNotFoundError, + PluginValidationError, ) -from pyrit.setup.plugin_formats import WheelPluginFormat +from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry +from pyrit.setup.plugin_discovery import ( + ScenarioContribution, + TechniqueContribution, + discover_scenarios, + discover_techniques, + import_plugin_async, +) +from pyrit.setup.plugin_formats import PreparedPlugin, SourcePluginFormat, WheelPluginFormat +from pyrit.setup.plugin_spec import PluginFormat from pyrit.setup.pyrit_initializer import PyRITInitializer if TYPE_CHECKING: - from collections.abc import Mapping, Sequence - from types import ModuleType + from collections.abc import Sequence - from pyrit.registry import ScenarioRegistry + from pyrit.models import ComponentIdentifier + from pyrit.registry import ScenarioMetadata + from pyrit.registry.instance_registry import DefaultInstanceRegistry, RegistryEntry + from pyrit.scenario import AttackTechniqueFactory, Scenario from pyrit.setup.plugin_spec import PluginSpec logger = logging.getLogger(__name__) +_REMEDIATION = "Fix or remove the plug-in entry from .pyrit_conf, then restart PyRIT." + async def load_plugins_if_configured_async(*, plugins: Sequence[PluginSpec]) -> None: """ - Load every configured plug-in, in order. A no-op when ``plugins`` is empty. - - Convenience entry point invoked by ``initialize_pyrit_async`` after memory is set and - before the configured initializers run. Plug-ins are loaded in list order, so one - plug-in behaves identically to several — the single-plug-in case is just a - one-element list. + Activate the configured V1 plug-in. Args: - plugins: The plug-in specs to load, in order. + plugins: The normalized plug-in declarations. Raises: - PluginLoadError: If a plug-in fails to load. + PluginLoadError: If activation fails. ValueError: If more than one plug-in is supplied. """ if not plugins: - logger.debug("No plug-ins configured; plug-in loading is a no-op.") return if len(plugins) > 1: raise ValueError("V1 supports one plug-in at a time; plug-in composition is not supported.") - - for spec in plugins: - await PluginLoader(spec=spec).load_async() - - -def _name_owned_by(module_name: str, package_name: str) -> bool: - """ - Return whether a module name belongs to the given plug-in package. - - Args: - module_name: A dotted module name. - package_name: The plug-in's top-level package name. - - Returns: - bool: True if ``module_name`` is the package or one of its submodules. - """ - return module_name == package_name or module_name.startswith(f"{package_name}.") - - -def _module_owned_by(cls: type, package_name: str) -> bool: - """ - Return whether ``cls`` is defined within the given plug-in package. - - Args: - cls: The class to check. - package_name: The plug-in's top-level package name. - - Returns: - bool: True if the class's module is the package or one of its submodules. - """ - return _name_owned_by(cls.__module__ or "", package_name) - - -_REMEDIATION = "Fix or remove the plug-in entry from .pyrit_conf, then restart PyRIT." + await PluginLoader(spec=plugins[0]).load_async() class PluginInitializer(PyRITInitializer): @@ -136,401 +89,209 @@ async def initialize_async(self) -> None: await load_plugins_if_configured_async(plugins=self._plugins) -class PluginLoader: - """ - Extract and register a single PyRIT plug-in wheel described by a ``PluginSpec``. +@dataclass +class _RegistrySnapshot: + scenario_classes: dict[str, type[Scenario]] + scenario_metadata: dict[str, ScenarioMetadata] | None + scenario_discovered: bool + technique_entries: dict[str, RegistryEntry[AttackTechniqueFactory]] + technique_metadata: list[ComponentIdentifier] | None + providers: dict[str, type[SeedDatasetProvider]] + sys_path: list[str] + module_names: set[str] - The wheel is extracted to ``.plugin//`` (never installed), imported, and its - bootstrap is run so its datasets and scenarios register like built-ins. Fails closed - by default. - """ + @classmethod + def capture(cls) -> _RegistrySnapshot: + scenario_registry = ScenarioRegistry.get_registry_singleton() + technique_registry = AttackTechniqueRegistry.get_registry_singleton() + technique_instances = cast( + "DefaultInstanceRegistry[AttackTechniqueFactory]", + technique_registry.instances, + ) + return cls( + scenario_classes=dict(scenario_registry._classes), + scenario_metadata=( + dict(scenario_registry._metadata_cache) if scenario_registry._metadata_cache is not None else None + ), + scenario_discovered=scenario_registry._discovered, + technique_entries=dict(technique_instances._registry_items), + technique_metadata=( + list(technique_instances._metadata_cache) if technique_instances._metadata_cache is not None else None + ), + providers=dict(SeedDatasetProvider._registry), + sys_path=list(sys.path), + module_names=set(sys.modules), + ) + + def restore(self, *, owned_prefixes: tuple[str, ...] = ()) -> None: + scenario_registry = ScenarioRegistry.get_registry_singleton() + scenario_registry._classes = dict(self.scenario_classes) + scenario_registry._metadata_cache = dict(self.scenario_metadata) if self.scenario_metadata is not None else None + scenario_registry._discovered = self.scenario_discovered + + technique_instances = cast( + "DefaultInstanceRegistry[AttackTechniqueFactory]", + AttackTechniqueRegistry.get_registry_singleton().instances, + ) + technique_instances._registry_items = dict(self.technique_entries) + technique_instances._metadata_cache = ( + list(self.technique_metadata) if self.technique_metadata is not None else None + ) + + SeedDatasetProvider._registry.clear() + SeedDatasetProvider._registry.update(self.providers) + sys.path[:] = self.sys_path + for name in list(sys.modules): + if name not in self.module_names and _name_owned_by_any(name=name, prefixes=owned_prefixes): + del sys.modules[name] + + def restore_unsupported_provider_side_effects(self) -> None: + """Remove provider registrations because datasets are outside the V1 contract.""" + SeedDatasetProvider._registry.clear() + SeedDatasetProvider._registry.update(self.providers) + + +class PluginLoader: + """Prepare, discover, validate, and transactionally register one plug-in.""" def __init__(self, *, spec: PluginSpec) -> None: """ - Initialize the loader for one plug-in. + Initialize the loader. Args: - spec: The plug-in to load (wheel path plus optional package name). + spec (PluginSpec): The normalized source or wheel declaration. """ self._spec = spec async def load_async(self) -> None: """ - Load the plug-in described by this loader's ``PluginSpec``. + Activate the configured plug-in. Raises: - PluginWheelNotFoundError: If the spec is not a wheel-format plug-in. - PluginLoadError: If the plug-in fails to load. + PluginLoadError: If preparation, discovery, validation, or registration fails. """ - wheel = self._spec.wheel - if wheel is None: - raise PluginWheelNotFoundError("The current loader only accepts wheel-format plug-ins.") + snapshot = _RegistrySnapshot.capture() + prepared: PreparedPlugin | None = None try: - await self._load_plugin_async() + prepared = await self._prepare_async() + imported = await import_plugin_async(prepared=prepared) + self._reject_import_time_registration(snapshot=snapshot) + scenarios = discover_scenarios(imported=imported) + techniques = discover_techniques(imported=imported) + self._validate_names(scenarios=scenarios, techniques=techniques) + self._register(scenarios=scenarios, techniques=techniques) + snapshot.restore_unsupported_provider_side_effects() + logger.info( + "Loaded plug-in '%s': %d scenario(s), %d attack technique(s).", + self._spec.name, + len(scenarios), + len(techniques), + ) except Exception as exc: - message = f"Failed to load plug-in from wheel '{wheel}': {exc} {_REMEDIATION}" - # Preserve the specific failure type so callers can distinguish modes, while - # always surfacing the remediation guidance. Unknown errors (e.g. a raising - # bootstrap or an unsafe archive) become a plain PluginLoadError. + prefixes = prepared.owned_module_prefixes if prepared else () + snapshot.restore(owned_prefixes=prefixes) + message = f"Failed to load plug-in '{self._spec.name}': {exc} {_REMEDIATION}" if isinstance(exc, PluginLoadError): raise type(exc)(message) from exc raise PluginLoadError(message) from exc - async def _load_plugin_async(self) -> None: - """ - Extract, import, bootstrap, and verify a single plug-in wheel. - - Global state (``sys.path``, imported plug-in modules, and the provider/scenario - registries) is rolled back if the load fails, so a failed or accepted-failure load - leaves no partial trace. - - Raises: - PluginWheelNotFoundError: If the configured wheel is absent or invalid. - PluginImportError: If the plug-in package cannot be imported. - PluginRegisteredNothingError: If the plug-in registered no datasets or scenarios. - """ - prepared = await WheelPluginFormat().prepare_async(spec=self._spec) - extract_dir = prepared.import_root - package_name = prepared.entry_modules[0] - - from pyrit.setup.plugin_compat import warn_on_version_drift - - warn_on_version_drift(extract_dir=extract_dir) - - from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider - from pyrit.registry import ScenarioRegistry + async def _prepare_async(self) -> PreparedPlugin: + if self._spec.format is PluginFormat.SOURCE: + return await SourcePluginFormat().prepare_async(spec=self._spec) + if self._spec.format is PluginFormat.WHEEL: + prepared = await WheelPluginFormat().prepare_async(spec=self._spec) + self._warn_on_version_drift(prepared=prepared) + return prepared + raise PluginValidationError(f"Unsupported plug-in format: {self._spec.format}") + @staticmethod + def _reject_import_time_registration(*, snapshot: _RegistrySnapshot) -> None: scenario_registry = ScenarioRegistry.get_registry_singleton() - provider_snapshot = dict(SeedDatasetProvider._registry) - scenario_snapshot = dict(scenario_registry._classes) - modules_snapshot = {name for name in sys.modules if _name_owned_by(name, package_name)} - syspath_entry = str(extract_dir) - added_to_syspath = syspath_entry not in sys.path - if added_to_syspath: - sys.path.insert(0, syspath_entry) - - try: - logger.info("Importing plug-in package '%s'", package_name) - try: - module = importlib.import_module(package_name) - self._verify_module_location(module=module, extract_dir=extract_dir, package_name=package_name) - self._import_submodules(module=module, package_name=package_name) - except Exception as exc: - raise PluginImportError(f"Could not import plug-in package '{package_name}': {exc}") from exc - - await self._run_bootstrap_async(package_name=package_name, module=module) - - # Register the plug-in's own Scenario subclasses the same way built-ins are - # discovered (by type, scoped to the plug-in package), so plug-in scenarios are - # picked up without the bootstrap having to register each one explicitly. Classes - # the bootstrap already registered are left untouched. - registered_scenarios = scenario_registry.register_external_subclasses(package_name=package_name) - if registered_scenarios: - logger.info("Auto-registered %d plug-in scenario(s) from '%s'.", registered_scenarios, package_name) - - provider_count, scenario_count = self._count_registered( - package_name=package_name, scenario_registry=scenario_registry - ) - if not provider_count and not scenario_count: - raise PluginRegisteredNothingError( - f"Plug-in package '{package_name}' imported successfully but registered no datasets or " - "scenarios. The wheel is likely mis-packaged (imports cleanly yet loads nothing)." - ) - - dataset_collisions = self._warn_on_dataset_name_collisions(package_name=package_name) - except Exception: - self._rollback( - package_name=package_name, - syspath_entry=syspath_entry if added_to_syspath else None, - modules_snapshot=modules_snapshot, - provider_snapshot=provider_snapshot, - scenario_registry=scenario_registry, - scenario_snapshot=scenario_snapshot, - ) - raise - - logger.info( - "Loaded plug-in '%s': %d dataset provider(s), %d scenario(s) registered; %d dataset name collision(s)%s.", - package_name, - provider_count, - scenario_count, - len(dataset_collisions), - " (see PLUGIN DATASET SHADOWED warnings above)" if dataset_collisions else "", + technique_instances = cast( + "DefaultInstanceRegistry[AttackTechniqueFactory]", + AttackTechniqueRegistry.get_registry_singleton().instances, ) - - @staticmethod - def _verify_module_location(*, module: ModuleType, extract_dir: Path, package_name: str) -> None: - """ - Verify the imported package resolves inside the extraction directory. - - Guards against an installed package of the same name shadowing the extracted - plug-in — a silent failure where import succeeds but the wheel's code/data is - ignored. - - Args: - module: The imported plug-in package module. - extract_dir: The directory the wheel was extracted to. - package_name: The plug-in's top-level package name. - - Raises: - ValueError: If the imported package resolves outside ``extract_dir``. - """ - extract_resolved = extract_dir.resolve() - raw_locations = list(getattr(module, "__path__", []) or []) - module_file = getattr(module, "__file__", None) - if module_file: - raw_locations.append(module_file) - - locations = [Path(location).resolve() for location in raw_locations if location] - if not locations: - return - if not any(location.is_relative_to(extract_resolved) for location in locations): - raise ValueError( - f"Imported package '{package_name}' resolved to {locations[0]} which is outside the " - f"plug-in extraction directory {extract_resolved}. An installed package with the same " - "name is likely shadowing the plug-in; set the plug-in's 'package' in .pyrit_conf or " - "resolve the name conflict." + technique_entries = technique_instances._registry_items + if scenario_registry._classes != snapshot.scenario_classes or technique_entries != snapshot.technique_entries: + raise PluginValidationError( + "Plug-in source registered components during import. V1 plug-ins must expose definitions " + "and let the framework register them transactionally." ) @staticmethod - def _import_submodules(*, module: ModuleType, package_name: str) -> None: - """ - Import every submodule of the plug-in package. - - Ensures dataset providers self-register and bootstrap initializers become - discoverable even when the package ``__init__`` does not import them. Import - errors surface (plug-in dependencies must be pre-satisfied — fail loud). - - Args: - module: The imported plug-in package module. - package_name: The plug-in's top-level package name. - """ - module_path = getattr(module, "__path__", None) - if not module_path: - return # Single-module plug-in (not a package); nothing to walk. - - def _raise_on_error(name: str) -> None: - raise ImportError(f"Failed to import plug-in submodule '{name}'") - - for submodule in pkgutil.walk_packages(module_path, prefix=f"{package_name}.", onerror=_raise_on_error): - importlib.import_module(submodule.name) - - async def _run_bootstrap_async(self, *, package_name: str, module: ModuleType) -> None: - """ - Run the plug-in's bootstrap so its scenarios register. - - Prefers a top-level ``register()`` callable on the package, then any - ``PyRITInitializer`` subclass defined within the package. If neither exists the - plug-in is assumed to register everything on import (datasets-only plug-ins). + def _validate_names( + *, + scenarios: Sequence[ScenarioContribution], + techniques: Sequence[TechniqueContribution], + ) -> None: + if not scenarios and not techniques: + raise PluginRegisteredNothingError("Plug-in contributed no scenarios or attack techniques.") - Args: - package_name: The plug-in's top-level package name. - module: The imported plug-in package module. - """ - register = getattr(module, "register", None) - if callable(register): - logger.info("Running plug-in bootstrap register() from '%s'", package_name) - result = register() - if inspect.isawaitable(result): - await result - return + scenario_registry = ScenarioRegistry.get_registry_singleton() + builtin_scenarios = set(scenario_registry.get_class_names()) + collisions = sorted(item.registry_name for item in scenarios if item.registry_name in builtin_scenarios) + if collisions: + raise PluginCollisionError(f"Scenario name collision(s): {collisions}.") - initializer_classes = self._find_plugin_initializers(package_name=package_name) - if initializer_classes: - for initializer_class in initializer_classes: - logger.info("Running plug-in bootstrap initializer %s", initializer_class.__name__) - await initializer_class().initialize_async() - return + from pyrit.setup.initializers.techniques import build_technique_factories - logger.info( - "Plug-in '%s' exposes no register() or PyRITInitializer bootstrap; relying on " - "import-time registration only.", - package_name, + builtin_techniques = {factory.name for factory in build_technique_factories()} + registered_techniques = set(AttackTechniqueRegistry.get_registry_singleton().get_factories()) + technique_collisions = sorted( + item.factory.name for item in techniques if item.factory.name in builtin_techniques | registered_techniques ) + if technique_collisions: + raise PluginCollisionError(f"Attack technique name collision(s): {technique_collisions}.") - @staticmethod - def _find_plugin_initializers(*, package_name: str) -> list[type[PyRITInitializer]]: - """ - Find concrete ``PyRITInitializer`` subclasses defined within the plug-in package. - - Args: - package_name: The plug-in's top-level package name. + def _register( + self, + *, + scenarios: Sequence[ScenarioContribution], + techniques: Sequence[TechniqueContribution], + ) -> None: + technique_registry = AttackTechniqueRegistry.get_registry_singleton() + for contribution in techniques: + technique_registry.register_contributed_factory( + factory=contribution.factory, + plugin_name=self._spec.name or "", + scenario_names=contribution.scenario_names, + ) - Returns: - list[type[PyRITInitializer]]: Bootstrap initializer classes owned by the plug-in. - """ - prefix = f"{package_name}." - found: list[type[PyRITInitializer]] = [] - seen: set[type[PyRITInitializer]] = set() - - stack: list[type[PyRITInitializer]] = list(PyRITInitializer.__subclasses__()) - while stack: - cls = stack.pop() - if cls in seen: - continue - seen.add(cls) - stack.extend(cls.__subclasses__()) - - module_name = cls.__module__ or "" - if inspect.isabstract(cls): - continue - if module_name == package_name or module_name.startswith(prefix): - found.append(cls) - return found + scenario_registry = ScenarioRegistry.get_registry_singleton() + for contribution in scenarios: + scenario_registry.register_contributed_scenario( + scenario_class=contribution.scenario_class, + name=contribution.registry_name, + ) @staticmethod - def _count_registered(*, package_name: str, scenario_registry: ScenarioRegistry) -> tuple[int, int]: - """ - Count providers and scenarios registered by the plug-in package. - - Both are counted by matching each registered class's module against the plug-in - package, so the check is precise to this plug-in and safe to re-run. - - Args: - package_name: The plug-in's top-level package name. - scenario_registry: The scenario registry singleton the bootstrap registered into. - - Returns: - tuple[int, int]: (dataset provider count, scenario count) owned by the plug-in. - """ - from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider + def _warn_on_version_drift(*, prepared: PreparedPlugin) -> None: + declared = prepared.declared_pyrit_version + if not declared: + return + import pyrit - provider_count = sum( - 1 for cls in SeedDatasetProvider.get_all_providers().values() if _module_owned_by(cls, package_name) + running = getattr(pyrit, "__version__", "") or "" + if _major_minor(declared) == _major_minor(running) and _major_minor(declared) is not None: + return + logger.warning( + "PLUGIN VERSION DRIFT: plug-in '%s' declares pyrit %s but pyrit %s is running. " + "Compatibility is the artifact author's responsibility.", + prepared.spec.name, + declared, + running or "(unknown)", ) - # Read the raw class catalog directly: this snapshot must not trigger built-in - # discovery, and the plug-in's register_class writes straight into it. - scenario_count = sum(1 for cls in scenario_registry._classes.values() if _module_owned_by(cls, package_name)) - - return provider_count, scenario_count - @staticmethod - def _warn_on_dataset_name_collisions(*, package_name: str) -> list[str]: - """ - Warn loudly when a plug-in dataset name collides with an existing dataset name. - - The dataset resolver treats central memory as authoritative and only consults a - provider when memory has no seeds for that ``dataset_name``. Once a same-named - dataset is in memory, a scan uses it and never consults the plug-in's provider, so - the plug-in's copy is silently bypassed. Any collision with **another** registered - provider's ``dataset_name`` is surfaced prominently at load time so the mismatch is - never silent. - - This compares the **provider registry**, not live memory, on purpose. At this phase - memory is not populated yet, and a live-memory check would false-positive on the - plug-in's own datasets persisted from a prior run (the seed rows carry no trustworthy - source, so "already in memory" cannot be told apart from "this plug-in loaded it last - run" — it is fundamentally undecidable at load time). The registry check is a - **conservative proxy** for the shadowing that ``LoadDefaultDatasets`` will cause by - loading provider datasets into memory: if the operator's config does not run - ``load_default_datasets`` (or loads only a tag subset), a built-in name may not actually - land in memory and this warning can fire without real shadowing. Over-warning is the - safe direction — do NOT "fix" this into a memory check (it reintroduces the false - positives). Governing principle: a guard's value is its precision — a check that cries - wolf on legitimate plug-in data every run desensitizes operators and defeats itself for - the real collision, so false-positive-free with a documented gap beats high-recall-but- - noisy. Hard enforcement that can tell a real mismatch from a harmless name coincidence - belongs to the scenario's declared required-dataset-names / expected-source mechanism - (which knows the operator's intent), not this loader. - - Args: - package_name: The plug-in's top-level package name. - - Returns: - list[str]: The sorted colliding dataset names (empty when there are none). - """ - from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider - - def _safe_name(provider_class: type[SeedDatasetProvider]) -> str | None: - try: - return provider_class().dataset_name - except Exception: - return None - - providers = SeedDatasetProvider.get_all_providers() - - # Map dataset_name -> owning provider class name(s), split into the plug-in's own - # providers vs. everything else. "Other" deliberately EXCLUDES the plug-in's own - # providers, so a plug-in shipping multiple datasets (or a re-run) never self-flags. - plugin_owned: dict[str, str] = {} - other_owned: dict[str, list[str]] = {} - for class_name, provider_class in providers.items(): - name = _safe_name(provider_class) - if name is None: - continue - if _module_owned_by(provider_class, package_name): - plugin_owned.setdefault(name, class_name) - else: - other_owned.setdefault(name, []).append(class_name) - - collisions = sorted(plugin_owned.keys() & other_owned.keys()) - for name in collisions: - logger.warning( - "PLUGIN DATASET SHADOWED: plug-in '%s' provider %s registers dataset_name '%s', which is " - "already provided by %s. Central memory is authoritative, so a scan will use the existing " - "dataset and the plug-in's copy will NOT take effect. Rename the plug-in dataset to a unique name.", - package_name, - plugin_owned[name], - name, - ", ".join(sorted(other_owned[name])), - ) - return collisions +def _major_minor(version: str) -> tuple[int, int] | None: + parts = version.split(".") + if len(parts) < 2: + return None + try: + return int(parts[0]), int(parts[1]) + except ValueError: + return None - @staticmethod - def _rollback( - *, - package_name: str, - syspath_entry: str | None, - modules_snapshot: set[str], - provider_snapshot: Mapping[str, type], - scenario_registry: ScenarioRegistry, - scenario_snapshot: Mapping[str, type], - ) -> None: - """ - Undo the partial global-state changes made while loading a plug-in. - Removes the plug-in's ``sys.path`` entry, the modules it newly imported, and the - provider/scenario registrations it added, and **restores any entries the plug-in - overwrote**. Both registries key by a name the plug-in does not control - (``SeedDatasetProvider`` by ``cls.__name__``; the scenario catalog by registry - name) and assign unconditionally, so a plug-in whose provider/scenario name - collides with an existing one silently replaces it; rollback must put the original - back, not just drop the new key. State present before the load — including modules - that already existed and built-ins discovered meanwhile — is preserved. - - Args: - package_name: The plug-in's top-level package name. - syspath_entry: The ``sys.path`` entry to remove, or None if it was already present. - modules_snapshot: Package-owned module names present before the load. - provider_snapshot: Provider registry contents captured before the load. - scenario_registry: The scenario registry singleton to clean up. - scenario_snapshot: Scenario catalog contents captured before the load. - """ - from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider - - if syspath_entry and syspath_entry in sys.path: - sys.path.remove(syspath_entry) - - for name in [m for m in sys.modules if _name_owned_by(m, package_name) and m not in modules_snapshot]: - del sys.modules[name] - - # For every entry the plug-in now owns: restore the pre-load value if the key - # existed before (the plug-in overwrote it), otherwise drop the key it added. - for key in list(SeedDatasetProvider._registry): - if _module_owned_by(SeedDatasetProvider._registry[key], package_name): - if key in provider_snapshot: - SeedDatasetProvider._registry[key] = provider_snapshot[key] # type: ignore[ty:invalid-assignment] - else: - del SeedDatasetProvider._registry[key] - - changed_scenarios = False - for name in list(scenario_registry._classes): - if _module_owned_by(scenario_registry._classes[name], package_name): - if name in scenario_snapshot: - scenario_registry._classes[name] = scenario_snapshot[name] # type: ignore[ty:invalid-assignment] - else: - del scenario_registry._classes[name] - changed_scenarios = True - if changed_scenarios: - scenario_registry._metadata_cache = None +def _name_owned_by_any(*, name: str, prefixes: tuple[str, ...]) -> bool: + return any(name == prefix or name.startswith(f"{prefix}.") for prefix in prefixes) diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index dbf68dff00..de4ff00553 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -11,7 +11,6 @@ """ import inspect -import logging import os import sys import textwrap @@ -25,17 +24,18 @@ import pytest from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider +from pyrit.exceptions import ( + PluginImportError, + PluginLoadError, + PluginRegisteredNothingError, + PluginWheelNotFoundError, +) from pyrit.memory import CentralMemory -from pyrit.models import SeedDataset from pyrit.registry import ScenarioRegistry from pyrit.setup import PluginSpec from pyrit.setup.initialization import IN_MEMORY, initialize_pyrit_async from pyrit.setup.plugin_loader import ( - PluginImportError, PluginInitializer, - PluginLoadError, - PluginRegisteredNothingError, - PluginWheelNotFoundError, load_plugins_if_configured_async, ) @@ -375,56 +375,11 @@ async def test_no_op_when_no_plugins() -> None: assert sys.path == path_before -# --------------------------------------------------------------------------- -# Silent-failure trap: extraction, not zipimport -# --------------------------------------------------------------------------- - - -def test_raw_wheel_on_syspath_loses_datasets(tmp_path: Path) -> None: - """A raw .whl on sys.path imports but __file__-relative datasets vanish (regression guard).""" - wheel = build_mock_wheel(tmp_path, bootstrap="none", include_scenario=False) - - sys.path.insert(0, str(wheel.path)) - module = __import__(wheel.package) - paths_module = __import__(f"{wheel.package}.paths", fromlist=["MOCK_DATASETS_PATH"]) - - assert ".whl" in (module.__file__ or "") - assert not paths_module.MOCK_DATASETS_PATH.exists() - assert list(paths_module.MOCK_DATASETS_PATH.glob("**/*.yaml")) == [] - - -def test_extracted_wheel_loads_datasets(tmp_path: Path) -> None: - """Extracting the wheel to disk makes __file__-relative datasets resolve and load.""" - wheel = build_mock_wheel(tmp_path, bootstrap="none", include_scenario=False) - extract_dir = tmp_path / "extracted" - extract_dir.mkdir() - with zipfile.ZipFile(wheel.path) as archive: - archive.extractall(extract_dir) - - sys.path.insert(0, str(extract_dir)) - paths_module = __import__(f"{wheel.package}.paths", fromlist=["MOCK_DATASETS_PATH"]) - - yamls = list(paths_module.MOCK_DATASETS_PATH.glob("**/*.yaml")) - assert len(yamls) == 1 - - dataset = SeedDataset.from_yaml_file(yamls[0]) - assert len(dataset.seeds) == 3 - - # --------------------------------------------------------------------------- # Loading via the initializer # --------------------------------------------------------------------------- -async def test_load_registers_provider_on_import(tmp_path: Path) -> None: - """Importing the plug-in package self-registers its SeedDatasetProvider.""" - wheel = build_mock_wheel(tmp_path) - - await load_plugin(wheel, tmp_path / ".plugin") - - assert "MockProvider" in SeedDatasetProvider.get_all_providers() - - async def test_load_extracts_to_plugin_dir(tmp_path: Path) -> None: """The wheel is extracted (not installed) under the configured plug-in dir.""" wheel = build_mock_wheel(tmp_path) @@ -444,24 +399,23 @@ async def test_scenario_registration_survives_discovery(tmp_path: Path) -> None: await load_plugin(wheel, tmp_path / ".plugin") registry = ScenarioRegistry.get_registry_singleton() - assert registry._discovered is False # register_class must not trigger discovery - - names = registry.get_class_names() # triggers built-in discovery - assert wheel.scenario_name in names + names = registry.get_class_names() + assert "scenario" in names assert "airt.rapid_response" in names mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario - assert registry.get_class(wheel.scenario_name) is mock_scenario + assert registry.get_class("scenario") is mock_scenario -async def test_register_callable_bootstrap(tmp_path: Path) -> None: - """A plug-in exposing a top-level register() callable is bootstrapped too.""" +async def test_register_callable_bootstrap_is_not_executed(tmp_path: Path) -> None: + """V1 ignores plug-in-authored lifecycle hooks and discovers definitions itself.""" wheel = build_mock_wheel(tmp_path, bootstrap="register") await load_plugin(wheel, tmp_path / ".plugin") names = ScenarioRegistry.get_registry_singleton().get_class_names() - assert wheel.scenario_name in names + assert "scenario" in names + assert wheel.scenario_name not in names async def test_ordering_scenario_visible_to_preload(tmp_path: Path) -> None: @@ -473,7 +427,7 @@ async def test_ordering_scenario_visible_to_preload(tmp_path: Path) -> None: # get_class_names() is exactly what PreloadScenarioMetadata iterates; the plug-in # scenario being present proves it registered before that read would happen. names = ScenarioRegistry.get_registry_singleton().get_class_names() - assert wheel.scenario_name in names + assert "scenario" in names async def test_plugin_scenario_auto_registered_without_bootstrap(tmp_path: Path) -> None: @@ -487,11 +441,11 @@ async def test_plugin_scenario_auto_registered_without_bootstrap(tmp_path: Path) registry = ScenarioRegistry.get_registry_singleton() mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario assert mock_scenario in registry._classes.values() - assert registry._discovered is False # auto-registration must not trigger built-in discovery + assert registry._discovered is True -async def test_bootstrap_registration_not_duplicated_by_auto_register(tmp_path: Path) -> None: - """A scenario the bootstrap registers is not also re-registered under a fallback name.""" +async def test_framework_owns_registration_when_bootstrap_exists(tmp_path: Path) -> None: + """The framework owns the scenario name even when a register hook exists.""" wheel = build_mock_wheel(tmp_path, bootstrap="register") await load_plugin(wheel, tmp_path / ".plugin") @@ -499,26 +453,16 @@ async def test_bootstrap_registration_not_duplicated_by_auto_register(tmp_path: registry = ScenarioRegistry.get_registry_singleton() mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario registered_names = [name for name, cls in registry._classes.items() if cls is mock_scenario] - assert registered_names == [wheel.scenario_name] - - -async def test_datasets_only_plugin_loads_without_bootstrap(tmp_path: Path) -> None: - """A datasets-only plug-in (no bootstrap, no scenario) loads via import-time registration.""" - wheel = build_mock_wheel(tmp_path, bootstrap="none", include_scenario=False) - - await load_plugin(wheel, tmp_path / ".plugin") - - assert "MockProvider" in SeedDatasetProvider.get_all_providers() + assert registered_names == ["scenario"] async def test_submodule_walk_discovers_unwired_components(tmp_path: Path) -> None: - """Provider + bootstrap register even when __init__.py does not import them.""" + """Scenario discovery walks submodules even when ``__init__.py`` does not import them.""" wheel = build_mock_wheel(tmp_path, bootstrap="initializer", wire_init=False) await load_plugin(wheel, tmp_path / ".plugin") - assert "MockProvider" in SeedDatasetProvider.get_all_providers() - assert wheel.scenario_name in ScenarioRegistry.get_registry_singleton().get_class_names() + assert "scenario" in ScenarioRegistry.get_registry_singleton().get_class_names() async def test_shadowing_installed_package_is_rejected(tmp_path: Path) -> None: @@ -530,62 +474,6 @@ async def test_shadowing_installed_package_is_rejected(tmp_path: Path) -> None: await load_plugin(wheel, tmp_path / ".plugin", package="json") -# --------------------------------------------------------------------------- -# Dataset name collision (memory-authoritative resolver guard) -# --------------------------------------------------------------------------- - - -async def test_colliding_dataset_name_warns_loudly(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: - """A plug-in dataset_name that collides with an existing provider's name warns at load.""" - wheel = build_mock_wheel(tmp_path) - colliding_name = wheel.dataset_name - - class CollidingProvider(SeedDatasetProvider): - """Non-plug-in provider that already claims the plug-in's dataset name.""" - - @property - def dataset_name(self) -> str: - return colliding_name - - async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: - raise NotImplementedError - - with caplog.at_level(logging.WARNING, logger="pyrit.setup.plugin_loader"): - await load_plugin(wheel, tmp_path / ".plugin") - - messages = [record.getMessage() for record in caplog.records] - # Un-missable: greppable prefix, names the colliding dataset and BOTH providers. - assert any( - "PLUGIN DATASET SHADOWED:" in message - and colliding_name in message - and "MockProvider" in message - and "CollidingProvider" in message - for message in messages - ) - - -async def test_unique_dataset_name_does_not_warn(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: - """A plug-in whose dataset name is unique produces no collision warning.""" - wheel = build_mock_wheel(tmp_path) - - with caplog.at_level(logging.WARNING, logger="pyrit.setup.plugin_loader"): - await load_plugin(wheel, tmp_path / ".plugin") - - assert not any("PLUGIN DATASET SHADOWED:" in record.getMessage() for record in caplog.records) - - -async def test_reload_does_not_self_flag(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: - """Loading the same plug-in twice must not flag its own provider as a collision.""" - wheel = build_mock_wheel(tmp_path) - plugin_dir = tmp_path / ".plugin" - - await load_plugin(wheel, plugin_dir) - with caplog.at_level(logging.WARNING, logger="pyrit.setup.plugin_loader"): - await load_plugin(wheel, plugin_dir) - - assert not any("PLUGIN DATASET SHADOWED:" in record.getMessage() for record in caplog.records) - - # --------------------------------------------------------------------------- # Rollback on failure # --------------------------------------------------------------------------- @@ -603,64 +491,6 @@ async def test_failed_load_rolls_back_syspath(tmp_path: Path) -> None: assert extract_dir not in sys.path -async def test_failing_bootstrap_rolls_back_partial_registration(tmp_path: Path) -> None: - """A bootstrap that registers then raises has its registration rolled back.""" - wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") - plugin_dir = tmp_path / ".plugin" - - with pytest.raises(PluginLoadError): - await load_plugin(wheel, plugin_dir) - - registry = ScenarioRegistry.get_registry_singleton() - assert wheel.scenario_name not in registry._classes - assert "MockProvider" not in SeedDatasetProvider.get_all_providers() - assert str(plugin_dir / wheel.path.stem) not in sys.path - - -async def test_rollback_restores_overwritten_provider(tmp_path: Path) -> None: - """A failed load restores a provider entry the plug-in overwrote (name collision).""" - wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") - - # SeedDatasetProvider keys by class name and the mock provider is "MockProvider"; - # occupy that key so the plug-in's import overwrites it. - class _PreexistingProvider(SeedDatasetProvider): - should_register = False - - @property - def dataset_name(self) -> str: - return "preexisting" - - async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: - raise NotImplementedError - - SeedDatasetProvider._registry["MockProvider"] = _PreexistingProvider - - with pytest.raises(PluginLoadError): - await load_plugin(wheel, tmp_path / ".plugin") - - # The original provider is restored, not deleted or left replaced by the plug-in's. - assert SeedDatasetProvider._registry["MockProvider"] is _PreexistingProvider - - -async def test_rollback_restores_overwritten_scenario(tmp_path: Path) -> None: - """A failed load restores a scenario entry the plug-in overwrote (name collision).""" - from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse - - wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") - - class _PreexistingScenario(RapidResponse): - """Sentinel scenario occupying the plug-in's registry name.""" - - registry = ScenarioRegistry.get_registry_singleton() - registry.register_class(_PreexistingScenario, name=wheel.scenario_name) - - with pytest.raises(PluginLoadError): - await load_plugin(wheel, tmp_path / ".plugin") - - # The original scenario is restored, not deleted or left replaced by the plug-in's. - assert registry._classes[wheel.scenario_name] is _PreexistingScenario - - # --------------------------------------------------------------------------- # Failure modes # --------------------------------------------------------------------------- @@ -677,7 +507,7 @@ async def test_empty_wheel_is_loud(tmp_path: Path) -> None: """A wheel that imports cleanly but registers nothing fails loudly.""" wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) - with pytest.raises(PluginRegisteredNothingError, match="registered no datasets or scenarios"): + with pytest.raises(PluginRegisteredNothingError, match="no scenarios or attack techniques"): await load_plugin(wheel, tmp_path / ".plugin") diff --git a/tests/unit/setup/test_plugin_registration.py b/tests/unit/setup/test_plugin_registration.py new file mode 100644 index 0000000000..bf828262cf --- /dev/null +++ b/tests/unit/setup/test_plugin_registration.py @@ -0,0 +1,155 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import sys +from collections.abc import Iterator +from pathlib import Path + +import pytest + +from pyrit.datasets import SeedDatasetProvider +from pyrit.exceptions import PluginCollisionError +from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry +from pyrit.setup import PluginFormat, PluginSpec +from pyrit.setup.plugin_loader import PluginInitializer + + +@pytest.fixture(autouse=True) +def reset_plugin_registration_state() -> Iterator[None]: + path_snapshot = list(sys.path) + module_snapshot = set(sys.modules) + provider_snapshot = dict(SeedDatasetProvider._registry) + AttackTechniqueRegistry.reset_registry_singleton() + ScenarioRegistry.reset_registry_singleton() + yield + AttackTechniqueRegistry.reset_registry_singleton() + ScenarioRegistry.reset_registry_singleton() + SeedDatasetProvider._registry.clear() + SeedDatasetProvider._registry.update(provider_snapshot) + sys.path[:] = path_snapshot + for name in set(sys.modules) - module_snapshot: + del sys.modules[name] + + +def _source_spec(source: Path, *, name: str = "operation") -> PluginSpec: + return PluginSpec(name=name, format=PluginFormat.SOURCE, source=source) + + +async def test_plugin_initializer_registers_source_technique(tmp_path: Path) -> None: + source = tmp_path / "operation_technique.py" + source.write_text( + """from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory + +OPERATION = AttackTechniqueFactory( + name="operation_foobar", + attack_class=PromptSendingAttack, + strategy_tags=["single_turn", "scenario:airt.rapid_response"], +) +""", + encoding="utf-8", + ) + + await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() + + registry = AttackTechniqueRegistry.get_registry_singleton() + assert registry.get_factories()["operation_foobar"] is not None + entry = registry.instances.get_entry("operation_foobar") + assert entry is not None + assert entry.metadata["plugin_name"] == "operation" + assert entry.metadata["scenario_names"] == ["airt.rapid_response"] + + +async def test_plugin_initializer_registers_source_scenario(tmp_path: Path) -> None: + source = tmp_path / "private_scenario.py" + source.write_text( + """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + +class PrivateScenario(RapidResponse): + pass +""", + encoding="utf-8", + ) + + await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() + + scenario_class = ScenarioRegistry.get_registry_singleton().get_class("private_scenario") + assert scenario_class.__name__ == "PrivateScenario" + + +async def test_plugin_initializer_rejects_latent_builtin_technique_name(tmp_path: Path) -> None: + source = tmp_path / "collision.py" + source.write_text( + """from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory + +OPERATION = AttackTechniqueFactory( + name="role_play", + attack_class=PromptSendingAttack, + strategy_tags=["scenario:airt.rapid_response"], +) +""", + encoding="utf-8", + ) + + with pytest.raises(PluginCollisionError, match="role_play"): + await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() + + assert "role_play" not in AttackTechniqueRegistry.get_registry_singleton().get_factories() + + +async def test_plugin_initializer_rolls_back_scenario_when_technique_collides(tmp_path: Path) -> None: + source = tmp_path / "partial.py" + source.write_text( + """from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory +from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + +class PartialScenario(RapidResponse): + pass + +OPERATION = AttackTechniqueFactory( + name="role_play", + attack_class=PromptSendingAttack, + strategy_tags=["scenario:airt.rapid_response"], +) +""", + encoding="utf-8", + ) + + with pytest.raises(PluginCollisionError): + await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() + + assert "partial" not in ScenarioRegistry.get_registry_singleton()._classes + assert str(tmp_path.resolve()) not in sys.path + assert "partial" not in sys.modules + + +async def test_plugin_initializer_removes_unsupported_provider_side_effect(tmp_path: Path) -> None: + source = tmp_path / "provider_side_effect.py" + source.write_text( + """from pyrit.datasets import SeedDatasetProvider +from pyrit.executor.attack import PromptSendingAttack +from pyrit.models import SeedDataset +from pyrit.scenario import AttackTechniqueFactory + +class PrivateProvider(SeedDatasetProvider): + @property + def dataset_name(self): + return "private" + + async def fetch_dataset_async(self, *, cache=True): + return SeedDataset(seeds=[], dataset_name="private") + +OPERATION = AttackTechniqueFactory( + name="provider_safe", + attack_class=PromptSendingAttack, + strategy_tags=["scenario:airt.rapid_response"], +) +""", + encoding="utf-8", + ) + + await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() + + assert "PrivateProvider" not in SeedDatasetProvider.get_all_providers() From f94cea0c96f9b2be387e248a7702b22bea98cc48 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 19:48:56 -0700 Subject: [PATCH 24/34] feat(scenario): expose applicable plug-in techniques Store scenario applicability on contributed factories and have RapidResponse union its built-in core catalog with private techniques explicitly targeting airt.rapid_response, without requiring private factories to claim the core tag. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- .../components/attack_technique_registry.py | 26 ++++++++++++++ .../scenario/scenarios/airt/rapid_response.py | 11 ++++-- .../test_attack_technique_registry.py | 36 +++++++++++++++++++ .../unit/scenario/airt/test_rapid_response.py | 19 ++++++++++ 4 files changed, 90 insertions(+), 2 deletions(-) diff --git a/pyrit/registry/components/attack_technique_registry.py b/pyrit/registry/components/attack_technique_registry.py index 2880ec4f87..d69b814014 100644 --- a/pyrit/registry/components/attack_technique_registry.py +++ b/pyrit/registry/components/attack_technique_registry.py @@ -202,6 +202,32 @@ def get_factories_or_raise(self) -> dict[str, AttackTechniqueFactory]: ) return factories + def get_factories_for_scenario( + self, + *, + scenario_name: str, + base_query: TagQuery | None = None, + ) -> dict[str, AttackTechniqueFactory]: + """ + Return built-in-query matches plus contributions applicable to a scenario. + + Args: + scenario_name (str): The scanner-facing scenario registry name. + base_query (TagQuery | None): Optional query selecting the scenario's + built-in technique pool. + + Returns: + dict[str, AttackTechniqueFactory]: Matching factories, sorted by name. + """ + factories: dict[str, AttackTechniqueFactory] = {} + for entry in self.instances.get_all_instances(): + built_in_match = base_query is not None and base_query.matches(set(entry.instance.strategy_tags)) + applicable_scenarios = entry.metadata.get("scenario_names", []) + contributed_match = scenario_name in applicable_scenarios + if built_in_match or contributed_match: + factories[entry.name] = entry.instance + return factories + @property def scorer_override_policy(self) -> ScorerOverridePolicy: """The policy applied when a scenario scorer is incompatible with an attack's annotation.""" diff --git a/pyrit/scenario/scenarios/airt/rapid_response.py b/pyrit/scenario/scenarios/airt/rapid_response.py index 652a1e9711..efc8481d3b 100644 --- a/pyrit/scenario/scenarios/airt/rapid_response.py +++ b/pyrit/scenario/scenarios/airt/rapid_response.py @@ -45,11 +45,18 @@ def _build_rapid_response_strategy() -> type[ScenarioStrategy]: from pyrit.registry.tag_query import TagQuery registry = AttackTechniqueRegistry.get_registry_singleton() - factories = list(registry.get_factories_or_raise().values()) + factories = list( + registry.get_factories_for_scenario( + scenario_name="airt.rapid_response", + base_query=TagQuery.all("core"), + ).values() + ) + if not factories: + registry.get_factories_or_raise() return AttackTechniqueRegistry.build_strategy_class_from_factories( # type: ignore[ty:invalid-return-type] class_name="RapidResponseStrategy", - factories=TagQuery.all("core").filter(factories), + factories=factories, aggregate_tags={ "default": TagQuery.any_of("default"), "single_turn": TagQuery.any_of("single_turn"), diff --git a/tests/unit/registry/test_attack_technique_registry.py b/tests/unit/registry/test_attack_technique_registry.py index 11696155ff..b268d020b4 100644 --- a/tests/unit/registry/test_attack_technique_registry.py +++ b/tests/unit/registry/test_attack_technique_registry.py @@ -13,6 +13,7 @@ from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry +from pyrit.registry.tag_query import TagQuery from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory, ScorerOverridePolicy from pyrit.setup.initializers.techniques import build_technique_factories @@ -120,6 +121,41 @@ def test_register_multiple_techniques(self): assert len(self.registry.instances) == 2 assert self.registry.instances.get_names() == ["stub_20", "stub_5"] + def test_get_factories_for_scenario_unions_base_query_and_applicability(self): + core = AttackTechniqueFactory( + name="core_technique", + attack_class=_StubAttack, + strategy_tags=["core", "single_turn"], + ) + private = AttackTechniqueFactory( + name="private_technique", + attack_class=_StubAttack, + strategy_tags=["single_turn"], + ) + unrelated = AttackTechniqueFactory( + name="unrelated", + attack_class=_StubAttack, + strategy_tags=["single_turn"], + ) + self.registry.register_technique(name=core.name, factory=core, tags=core.strategy_tags) + self.registry.register_contributed_factory( + factory=private, + plugin_name="operation", + scenario_names=frozenset({"airt.rapid_response"}), + ) + self.registry.register_contributed_factory( + factory=unrelated, + plugin_name="operation", + scenario_names=frozenset({"airt.cyber"}), + ) + + factories = self.registry.get_factories_for_scenario( + scenario_name="airt.rapid_response", + base_query=TagQuery.all("core"), + ) + assert list(factories) == ["core_technique", "private_technique"] + assert list(factories) == ["core_technique", "private_technique"] + class TestAttackTechniqueRegistryMetadata: """Tests for metadata / list_metadata on the registry.""" diff --git a/tests/unit/scenario/airt/test_rapid_response.py b/tests/unit/scenario/airt/test_rapid_response.py index b129d8b518..134d31f2fa 100644 --- a/tests/unit/scenario/airt/test_rapid_response.py +++ b/tests/unit/scenario/airt/test_rapid_response.py @@ -161,6 +161,25 @@ def test_get_strategy_class(self, mock_objective_scorer): ): assert RapidResponse()._strategy_class is strat + def test_strategy_class_includes_applicable_contributed_technique(self): + from pyrit.scenario.scenarios.airt.rapid_response import _build_rapid_response_strategy + + factory = AttackTechniqueFactory( + name="operation_foobar", + attack_class=PromptSendingAttack, + strategy_tags=["single_turn"], + ) + AttackTechniqueRegistry.get_registry_singleton().register_contributed_factory( + factory=factory, + plugin_name="operation", + scenario_names=frozenset({"airt.rapid_response"}), + ) + _build_rapid_response_strategy.cache_clear() + + strategy_class = _build_rapid_response_strategy() + + assert strategy_class("operation_foobar").value == "operation_foobar" + def test_get_default_strategy_returns_default(self, mock_objective_scorer): strat = _strategy_class() with patch( From 11e40cb7655bc6f5d9d1e25546314b3040ab0d68 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 19:57:26 -0700 Subject: [PATCH 25/34] test(setup): exercise source plug-ins through scanner Migrate wheel integration to the ConfigurationLoader-owned privileged path and complete source scenario and RapidResponse technique tests through a freshly launched stock pyrit_scan backend. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- .../setup/test_plugin_loader_integration.py | 189 ++++++++++++++++-- 1 file changed, 176 insertions(+), 13 deletions(-) diff --git a/tests/integration/setup/test_plugin_loader_integration.py b/tests/integration/setup/test_plugin_loader_integration.py index 9045d62616..230c952060 100644 --- a/tests/integration/setup/test_plugin_loader_integration.py +++ b/tests/integration/setup/test_plugin_loader_integration.py @@ -29,6 +29,8 @@ import inspect import os +import re +import subprocess import sys import textwrap import uuid @@ -43,17 +45,17 @@ from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider from pyrit.models import ScenarioResult, ScenarioRunState from pyrit.prompt_target import OpenAIChatTarget -from pyrit.registry import ScenarioRegistry +from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry from pyrit.registry.discovery import discover_in_directory from pyrit.scenario.core import DatasetAttackConfiguration, Scenario from pyrit.score import SelfAskTrueFalseScorer, TrueFalseQuestionPaths -from pyrit.setup import IN_MEMORY, PluginSpec, initialize_pyrit_async +from pyrit.setup import ConfigurationLoader, PluginFormat, PluginSpec # (module stem, class name, registry name) for the scenarios the self-contained wheel ships. _MOCK_SCENARIOS = [ - ("alpha", "MockAlphaScenario", "airt.mock_alpha"), - ("beta", "MockBetaScenario", "airt.mock_beta"), - ("gamma", "MockGammaScenario", "airt.mock_gamma"), + ("alpha", "MockAlphaScenario", "alpha"), + ("beta", "MockBetaScenario", "beta"), + ("gamma", "MockGammaScenario", "gamma"), ] # Substring of the ValueError ``Scenario.run_async`` raises when an atomic attack's @@ -63,6 +65,19 @@ _OBJECTIVES_INCOMPLETE_MARKER = "objectives incomplete" +async def _initialize_plugin_async(*, spec: PluginSpec, plugin_dir: Path | None) -> None: + """Initialize PyRIT through the config-owned privileged plug-in path.""" + config = ConfigurationLoader( + memory_db_type="in_memory", + initializers=["technique"], + env_files=[], + silent=True, + plugins=[spec.to_config()], + ) + with _plugin_dir_env(plugin_dir=plugin_dir): + await config.initialize_pyrit_async() + + def _build_scenario_plugin_wheel(dest_dir: Path, *, package: str) -> Path: """ Build a wheel whose package ships several scenarios and a ``register()`` bootstrap. @@ -293,6 +308,7 @@ def plugin_sandbox() -> Iterator[None]: del sys.modules[name] SeedDatasetProvider._registry.clear() SeedDatasetProvider._registry.update(provider_snapshot) + AttackTechniqueRegistry.reset_registry_singleton() ScenarioRegistry.reset_registry_singleton() @@ -309,8 +325,15 @@ async def test_built_wheel_scenarios_are_discovered( registry = ScenarioRegistry.get_registry_singleton() before = set(registry.get_class_names()) - with _plugin_dir_env(plugin_dir=plugin_dir): - await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel)]) + await _initialize_plugin_async( + spec=PluginSpec( + name=package, + format=PluginFormat.WHEEL, + wheel=wheel, + package=package, + ), + plugin_dir=plugin_dir, + ) registry = ScenarioRegistry.get_registry_singleton() after = set(registry.get_class_names()) @@ -348,8 +371,15 @@ async def test_injected_wheel_scenarios_are_discovered( if not scenario_dirs_env and not package: pytest.skip("Set PLUGIN_TEST_SCENARIO_DIRS or PLUGIN_TEST_PACKAGE to define the expected scenarios.") - with _plugin_dir_env(plugin_dir=None): - await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel, package=package)]) + await _initialize_plugin_async( + spec=PluginSpec( + name=package or "injected_plugin", + format=PluginFormat.WHEEL, + wheel=wheel, + package=package, + ), + plugin_dir=None, + ) found = _registered_scenario_class_names() @@ -379,8 +409,15 @@ async def test_injected_wheel_scenarios_instantiate(plugin_sandbox: None) -> Non wheel = Path(wheel_env).expanduser() assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" - with _plugin_dir_env(plugin_dir=None): - await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel, package=package)]) + await _initialize_plugin_async( + spec=PluginSpec( + name=package, + format=PluginFormat.WHEEL, + wheel=wheel, + package=package, + ), + plugin_dir=None, + ) classes = _scenario_classes_under_package(package) assert classes, f"No plug-in scenarios registered under package {package!r}." @@ -413,8 +450,15 @@ async def test_injected_wheel_scenario_executes(plugin_sandbox: None) -> None: wheel = Path(wheel_env).expanduser() assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" - with _plugin_dir_env(plugin_dir=None): - await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel, package=package)]) + await _initialize_plugin_async( + spec=PluginSpec( + name=package, + format=PluginFormat.WHEEL, + wheel=wheel, + package=package, + ), + plugin_dir=None, + ) endpoint = os.getenv("ADVERSARIAL_CHAT_ENDPOINT") if not endpoint: @@ -436,3 +480,122 @@ async def test_injected_wheel_scenario_executes(plugin_sandbox: None) -> None: if result is not None: assert result.scenario_run_state == ScenarioRunState.COMPLETED assert result.attack_results, "Completed scenario run produced no attack results." + + +# --------------------------------------------------------------------------- +# Scanner integration +# --------------------------------------------------------------------------- + + +def _write_scanner_config(*, tmp_path: Path, source: Path, plugin_name: str) -> Path: + config = tmp_path / f"{plugin_name}.pyrit_conf" + config.write_text( + textwrap.dedent( + f"""\ + memory_db_type: in_memory + initializers: + - target + - scorer + - technique + plugins: + - name: {plugin_name} + format: source + source: {source.as_posix()} + """ + ), + encoding="utf-8", + ) + return config + + +def _run_scanner_with_fresh_backend(*, config: Path) -> subprocess.CompletedProcess[str]: + base = [sys.executable, "-m", "pyrit.cli.pyrit_scan", "--config-file", str(config)] + subprocess.run([*base, "--stop-server"], capture_output=True, text=True, check=False, timeout=30) + try: + return subprocess.run( + [*base, "--start-server", "--list-scenarios"], + capture_output=True, + text=True, + check=True, + timeout=120, + ) + finally: + subprocess.run([*base, "--stop-server"], capture_output=True, text=True, check=False, timeout=30) + + +@pytest.mark.run_only_if_all_tests +def test_private_scenario_registers_from_scanner(tmp_path: Path) -> None: + """A source Scenario appears through the stock scanner catalog.""" + source = tmp_path / "private_scanner_scenario.py" + source.write_text( + """from pyrit.models import SeedObjective +from pyrit.scenario import DatasetAttackConfiguration, Scenario, ScenarioStrategy +from pyrit.score import SubStringScorer + +class PrivateStrategy(ScenarioStrategy): + ALL = ("all", {"all"}) + DIRECT = ("direct", {"direct"}) + +class PrivateScannerScenario(Scenario): + VERSION = 1 + + def __init__(self, *, scenario_result_id=None): + super().__init__( + version=self.VERSION, + strategy_class=PrivateStrategy, + default_strategy=PrivateStrategy.ALL, + default_dataset_config=DatasetAttackConfiguration( + seeds=[SeedObjective(value="integration objective")] + ), + objective_scorer=SubStringScorer(substring="integration"), + scenario_result_id=scenario_result_id, + ) + + async def _build_atomic_attacks_async(self, *, context): + return [] +""", + encoding="utf-8", + ) + config = _write_scanner_config( + tmp_path=tmp_path, + source=source, + plugin_name="private_scanner_scenario", + ) + + proc = _run_scanner_with_fresh_backend(config=config) + + assert "private_scanner_scenario" in proc.stdout + assert "PrivateScannerScenario" in proc.stdout + + +@pytest.mark.run_only_if_all_tests +def test_private_attack_technique_registers_from_scanner(tmp_path: Path) -> None: + """A source technique applicable to RapidResponse appears in scanner metadata.""" + source = tmp_path / "private_scanner_technique.py" + source.write_text( + """from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory + +PRIVATE = AttackTechniqueFactory( + name="private_scanner_technique", + attack_class=PromptSendingAttack, + strategy_tags=["single_turn", "scenario:airt.rapid_response"], +) +""", + encoding="utf-8", + ) + config = _write_scanner_config( + tmp_path=tmp_path, + source=source, + plugin_name="private_scanner_technique", + ) + + proc = _run_scanner_with_fresh_backend(config=config) + + match = re.search( + r"\n airt\.rapid_response\n(?P.*?)(?=\n [a-z][a-z0-9_.]*\n|\n={80})", + proc.stdout, + flags=re.DOTALL, + ) + assert match is not None, proc.stdout + assert "private_scanner_technique" in match.group("body") From ddf27ac769dacfe0182207d95368b3188a9e7798 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 20:10:36 -0700 Subject: [PATCH 26/34] docs(setup): explain private scenario plug-in usage Document when to use config-driven plug-ins instead of direct Python composition, add source/wheel setup and troubleshooting guides, update .pyrit_conf examples, and surface the guidance in pyrit_scan --help without adding scanner activation logic. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- .pyrit_conf_example | 36 ++-- doc/getting_started/README.md | 1 + doc/getting_started/plugins.md | 173 ++++++++++++++++++ doc/getting_started/pyrit_conf.md | 35 +++- doc/getting_started/troubleshooting/README.md | 2 + .../troubleshooting/plugins.md | 113 ++++++++++++ doc/myst.yml | 2 + pyrit/cli/pyrit_scan.py | 7 + tests/unit/cli/test_pyrit_scan.py | 5 +- 9 files changed, 356 insertions(+), 18 deletions(-) create mode 100644 doc/getting_started/plugins.md create mode 100644 doc/getting_started/troubleshooting/plugins.md diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 0522e20cab..aa3511ffaf 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -59,26 +59,32 @@ initializers: # Plug-ins # -------- -# Non-disclosable plug-ins (extra scenarios + datasets) shipped as pre-built wheels and -# loaded at initialization. Plug-ins load as a guaranteed-first phase — after memory is set -# and BEFORE the initializers above — so their scenarios/datasets register before anything -# reads the registry. They are a dedicated phase, not an entry in `initializers:`, so there -# is no ordering to get right. Loaded in list order; one plug-in behaves identically to -# several. A configless setup simply has no plug-ins. -# -# Each entry can be: -# - a wheel path (package auto-detected): - /abs/path/to/my_plugin.whl -# - a {package: wheel} pair (names the package): - my_plugin: /abs/path/to/my_plugin.whl -# - an explicit mapping: - wheel: /abs/path/to/my_plugin.whl -# package: my_plugin +# Private scenarios and attack techniques loaded from source or a pre-built wheel. +# ConfigurationLoader automatically prepends a privileged PluginInitializer ahead of +# the initializers above; do NOT list it under `initializers:`. +# +# V1 accepts one plug-in. Exactly one of `source` or `wheel` is required. +# Relative paths resolve against this configuration file. # # WARNING: loading a plug-in executes third-party code in the PyRIT process. Whoever can # write this file can run code on the host. Treat it as sensitive. # -# Example: +# Source example: +# plugins: +# - name: operation_foobar +# format: source +# source: /opt/pyrit/operation_foobar.py +# +# Wheel example: # plugins: -# - my_first_plugin: /abs/path/to/first-0.0.0-py3-none-any.whl -# - my_second_plugin: /abs/path/to/second-0.0.0-py3-none-any.whl +# - name: partner_scenarios +# format: wheel +# wheel: /opt/pyrit/partner_scenarios-1.2.0-py3-none-any.whl +# package: partner_scenarios +# +# Help and troubleshooting: +# - doc/getting_started/plugins.md +# - doc/getting_started/troubleshooting/plugins.md # Default Scenario # ---------------- diff --git a/doc/getting_started/README.md b/doc/getting_started/README.md index b5795569d7..af16d905dd 100644 --- a/doc/getting_started/README.md +++ b/doc/getting_started/README.md @@ -28,4 +28,5 @@ Once you're set up: - 🔌 [Targets](../code/targets/0_prompt_targets.md) — Connect to different AI systems - 📦 [Scenarios](../code/scenarios/0_scenarios.ipynb) — Run standardized evaluation scenarios +- 🧩 [Private Scenario Plug-Ins](./plugins.md) — Add private scenarios and attack techniques to stock PyRIT - 🖥️ [Scanner](../scanner/0_scanner.md) — Use `pyrit_scan` for automated assessments diff --git a/doc/getting_started/plugins.md b/doc/getting_started/plugins.md new file mode 100644 index 0000000000..bc9417e5be --- /dev/null +++ b/doc/getting_started/plugins.md @@ -0,0 +1,173 @@ +# Private Scenarios and Attack Techniques + +PyRIT plug-ins make private scenarios and attack techniques behave like built-in +components in the standard backend, GUI, and `pyrit_scan` catalog. + +## Plug-In or Python Dependency? + +Use ordinary Python composition when you own the application: + +```python +import pyrit + +# Build and run your custom workflow directly. +``` + +That is the simplest option for notebooks, services, and tools that already control +their Python entry point. + +Use a plug-in when operators need private components inside a **stock PyRIT +installation**: + +- keep the shipped `pyrit_scan` command and backend; +- discover private components through normal catalog APIs; +- avoid a private PyRIT fork or wrapper CLI; +- activate components from `.pyrit_conf`; +- deploy an artifact/config instead of application glue. + +## V1 Scope + +V1 plug-ins contribute: + +- concrete `Scenario` subclasses; +- configured `AttackTechniqueFactory` instances. + +They do not contribute custom initializers or lifecycle hooks. PyRIT's privileged +`PluginInitializer` owns loading, discovery, validation, registration, and rollback. + +Only one plug-in is supported at a time. + +## Source Plug-In + +Source is useful for live operations where scenarios or techniques already exist on +the backend filesystem. + +Supported shapes: + +- one importable `.py` file; +- one Python package directory containing `__init__.py`. + +### Private attack technique + +```python +# /opt/pyrit/operation_foobar.py +from pyrit.executor.attack import PromptSendingAttack +from pyrit.scenario import AttackTechniqueFactory + + +OPERATION_FOOBAR = AttackTechniqueFactory( + name="operation_foobar", + attack_class=PromptSendingAttack, + strategy_tags=[ + "single_turn", + "scenario:airt.rapid_response", + ], +) +``` + +The `scenario:` tag declares where a directly discovered factory is +available. Construction tooling may replace this bridge with an explicit contribution +manifest in the future. + +Configure it: + +```yaml +plugins: + - name: operation_foobar + format: source + source: /opt/pyrit/operation_foobar.py +``` + +Then restart the backend and run: + +```powershell +pyrit_scan airt.rapid_response ` + --target openai_chat ` + --strategies operation_foobar +``` + +### Private scenario + +A source scenario follows the normal `Scenario` contract: + +- concrete subclass; +- keyword-only constructor; +- no-argument instantiable for catalog metadata; +- runtime dependencies resolved lazily; +- `_build_atomic_attacks_async` implemented. + +The default registry name comes from the source-relative module path. For example: + +```text +/opt/pyrit/private_scenarios/image/abuse.py -> image.abuse +``` + +A package directory must contain `__init__.py`; a directory of unrelated loose files +is rejected. + +## Wheel Plug-In + +Wheels are appropriate for durable distribution of multi-file private scenarios and +resources. The wheel must already be built and compatible with the running PyRIT. + +```yaml +plugins: + - name: partner_scenarios + format: wheel + wheel: /opt/pyrit/partner_scenarios-1.2.0-py3-none-any.whl + package: partner_scenarios +``` + +PyRIT safely extracts the wheel to `.plugin/`; it does not run `pip`, install into +`.venv`, or resolve dependencies. + +If wheel metadata declares another PyRIT version, loading emits an advisory warning. +PyRIT does not rewrite incompatible APIs. Rebuild the wheel against the running +version when validation fails. + +## Initialization Behavior + +`ConfigurationLoader` converts the config entry into a privileged +`PluginInitializer`. Operators do not add this initializer to `initializers:`. + +Runtime order: + +1. load environment; +2. create CentralMemory; +3. activate the configured plug-in; +4. execute user-configured initializers; +5. serve catalog and scenario requests. + +This ordering ensures scenario metadata and technique strategies cannot be built from +a partial catalog. + +For direct Python use: + +```python +from pyrit.setup import initialize_from_config_async + +await initialize_from_config_async("/path/to/.pyrit_conf") +``` + +Low-level `initialize_pyrit_async()` does not read `.pyrit_conf` automatically. + +## Trust Boundary + +Source and wheel plug-ins execute Python with backend permissions. Treat write access +to the artifact and `.pyrit_conf` as code-execution authority. + +Dependencies must already be installed in the backend environment. Plug-ins share one +interpreter and do not receive dependency isolation. + +## Updating a Plug-In + +V1 does not hot reload. + +```powershell +pyrit_scan --stop-server +pyrit_scan --start-server --config-file /path/to/.pyrit_conf +``` + +For direct Python use, initialize in a fresh process after changing the artifact. + +See [Plug-In Troubleshooting](./troubleshooting/plugins.md) for common failures. diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 47565222e0..0f40bb4f48 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -151,6 +151,33 @@ initialization_scripts: - ./local_initializer.py ``` +### `plugins` + +One private scenario/attack-technique plug-in to activate before the configured +initializers. The plug-in can be a source file/package or a pre-built wheel: + +```yaml +plugins: + - name: operation_foobar + format: source + source: /opt/pyrit/operation_foobar.py +``` + +```yaml +plugins: + - name: partner_scenarios + format: wheel + wheel: /opt/pyrit/partner_scenarios-1.2.0-py3-none-any.whl + package: partner_scenarios +``` + +`ConfigurationLoader` automatically injects a privileged plug-in initializer first; +do not add it to `initializers:`. V1 is fail-closed, supports one plug-in, and requires +a process/backend restart after changes. + +See [Private Scenarios and Attack Techniques](./plugins.md) and +[Plug-In Troubleshooting](./troubleshooting/plugins.md). + ### `env_files` Environment file paths to load during initialization. Later files override values from earlier files. @@ -196,9 +223,13 @@ The 3-layer model above determines **which config values are selected**. Once re 1. Environment files are loaded 2. Default values are reset 3. Memory database is configured (from `memory_db_type`) -4. Initializers are executed in listed order +4. The privileged plug-in initializer runs when `plugins:` is configured +5. User-configured initializers are executed in listed order -Because initializers run last, they can modify anything set up in earlier steps — including environment variables and the memory instance. In practice, built-in initializers like `target` and `scorer` only call `set_default_value` and `set_global_variable` and do not touch memory or environment variables. However, a custom initializer could override those if needed. When this happens, the initializer's changes take effect because it runs after the other settings have been applied. +Because initializers run after environment and memory setup, they can use those +prerequisites. The plug-in initializer is framework-controlled and always precedes +the listed initializers so scenario/technique registries are complete before catalog +consumers run. ## Usage diff --git a/doc/getting_started/troubleshooting/README.md b/doc/getting_started/troubleshooting/README.md index c146eabe15..8eb820b821 100644 --- a/doc/getting_started/troubleshooting/README.md +++ b/doc/getting_started/troubleshooting/README.md @@ -14,6 +14,8 @@ Common issues and advanced setup guides. - **[Azure SQL Database Setup](./azure_sql_db.md)** — Setting up Azure SQL as your memory backend, including Entra ID authentication and user permissions. +- **[Private Scenario Plug-Ins](./plugins.md)** — Configuration, discovery, dependency, collision, and restart troubleshooting. + ## Model Deployment - **[Deploy HuggingFace Models on Azure ML](./deploy_hf_model_aml.ipynb)** — Deploy HuggingFace models to Azure ML endpoints. diff --git a/doc/getting_started/troubleshooting/plugins.md b/doc/getting_started/troubleshooting/plugins.md new file mode 100644 index 0000000000..08bba8453a --- /dev/null +++ b/doc/getting_started/troubleshooting/plugins.md @@ -0,0 +1,113 @@ +# Plug-In Troubleshooting + +## The plug-in is missing from `pyrit_scan` + +Restart the backend after changing `.pyrit_conf` or the artifact: + +```powershell +pyrit_scan --stop-server +pyrit_scan --start-server --config-file /path/to/.pyrit_conf +``` + +`pyrit_scan` is a thin client. An already-running backend keeps the plug-ins loaded +from its startup config. + +## Configuration is rejected + +V1 accepts one explicit entry: + +```yaml +plugins: + - name: operation_foobar + format: source + source: /absolute/path/operation_foobar.py +``` + +or: + +```yaml +plugins: + - name: partner_scenarios + format: wheel + wheel: /absolute/path/partner.whl + package: partner_scenarios +``` + +Check that: + +- `name` is lowercase snake case; +- `format` is `source` or `wheel`; +- exactly one matching artifact field is present; +- only one plug-in is configured. + +Relative paths resolve against the configuration file. + +## Source path is rejected + +A source path must be: + +- one `.py` file with an importable filename; or +- one package directory containing `__init__.py`. + +Loose directories containing unrelated Python files are not supported. + +## Import fails + +Plug-in dependencies are not installed automatically. Install them into the backend +environment and restart. + +Also verify: + +- package imports work from the configured source root; +- the configured package name matches the wheel/source package; +- another installed package is not shadowing the plug-in. + +## Scenario is rejected + +The scenario must: + +- inherit from `Scenario`; +- be concrete; +- use a keyword-only constructor; +- construct with no arguments for catalog metadata; +- implement `_build_atomic_attacks_async`. + +Plug-in scenarios are inspected before user-configured target/scorer/technique +initializers run. Keep import and no-argument construction side-effect-light and defer +runtime dependency resolution. + +## Technique is rejected + +Expose factories through either: + +- a module-owned `get_technique_factories()` returning + `AttackTechniqueFactory` instances; or +- module-global `AttackTechniqueFactory` instances. + +Directly discovered factories must identify an applicable scenario, for example: + +```python +strategy_tags=["single_turn", "scenario:airt.rapid_response"] +``` + +Do not register the factory during module import. The framework owns registration and +rollback. + +## Name collision + +V1 is extend-only. A private scenario or technique cannot replace a built-in or +existing registry name. Rename the contribution and restart. + +## Version drift warning + +The warning is advisory. PyRIT validates the live scenario/factory contract but does +not patch incompatible code. Rebuild or update the artifact against the installed +PyRIT version. + +## Partial state after failure + +Plug-in activation is fail-closed and transactional for supported scenario/technique +registries. If initialization fails, fix the reported stage and restart the process. + +Continuing production work in the same process after failed initialization is not +supported. diff --git a/doc/myst.yml b/doc/myst.yml index 08902dba7e..27c9898434 100644 --- a/doc/myst.yml +++ b/doc/myst.yml @@ -51,6 +51,7 @@ project: children: - file: getting_started/populating_secrets.md - file: getting_started/pyrit_conf.md + - file: getting_started/plugins.md - file: getting_started/troubleshooting/README.md children: - file: getting_started/troubleshooting/installation.md @@ -60,6 +61,7 @@ project: - file: getting_started/troubleshooting/local_dev.md - file: getting_started/troubleshooting/jupyter_setup.md - file: getting_started/troubleshooting/azure_sql_db.md + - file: getting_started/troubleshooting/plugins.md - file: getting_started/troubleshooting/deployment.md children: - file: getting_started/troubleshooting/deploy_hf_model_aml.ipynb diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index 7f92b85b1b..c45120b65e 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -84,6 +84,13 @@ def _print_cli_exception(*, exc: BaseException) -> None: Requires a running PyRIT backend server. Use --start-server to launch one, or connect to an existing server with --server-url. +Private scenario plug-ins: + Use a plug-in when private scenarios or attack techniques must behave like + built-ins in the stock scanner/backend/catalog. If you own a custom Python + application, importing PyRIT as a dependency and composing it directly is + usually simpler. Configure trusted source or wheel artifacts in .pyrit_conf. + See doc/getting_started/plugins.md. + Examples: # Start the backend server pyrit_scan --start-server diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index ae15a9ade8..c3e51fea32 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -146,10 +146,13 @@ def test_parse_args_invalid_max_retries(self): with pytest.raises(SystemExit): pyrit_scan.parse_args(["test_scenario", "--max-retries", "-1"]) - def test_parse_args_help_flag(self): + def test_parse_args_help_flag(self, capsys): with pytest.raises(SystemExit) as exc_info: pyrit_scan.parse_args(["--help"]) assert exc_info.value.code == 0 + help_text = capsys.readouterr().out + assert "Private scenario plug-ins" in help_text + assert "doc/getting_started/plugins.md" in help_text def test_parse_args_with_target(self): args = pyrit_scan.parse_args(["test_scenario", "--target", "my_target"]) From 34daddf2133b23ca382fcba46b3035beaa092331 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 20:23:53 -0700 Subject: [PATCH 27/34] refactor(setup): validate plug-in scenarios during activation Seed the privileged core technique prerequisite, validate contributed scenario metadata transactionally before startup completes, remove runtime API mutation shims, and align wheel tests with framework-owned lifecycle behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/setup/plugin_compat.py | 255 ------------------- pyrit/setup/plugin_loader.py | 9 + tests/unit/setup/test_plugin_loader.py | 5 +- tests/unit/setup/test_plugin_registration.py | 36 ++- 4 files changed, 44 insertions(+), 261 deletions(-) delete mode 100644 pyrit/setup/plugin_compat.py diff --git a/pyrit/setup/plugin_compat.py b/pyrit/setup/plugin_compat.py deleted file mode 100644 index e418a8b45b..0000000000 --- a/pyrit/setup/plugin_compat.py +++ /dev/null @@ -1,255 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Best-effort compatibility shims for plug-ins built against a slightly different PyRIT. - -A plug-in wheel is authored against one PyRIT version, then loaded into whatever PyRIT the -operator is running. When the host API has drifted, small mechanical differences — most -commonly a renamed extension-point method — can leave the plug-in's scenarios abstract and -therefore undiscoverable, even though the plug-in's own logic is unchanged. Rather than -force every plug-in author to re-release on each host bump, the loader applies narrow, loud -heuristics that bridge known mechanical renames so minor drift does not block loading. - -The shims are deliberately conservative. A scenario is bridged only when it is abstract -*solely* because of a recognized rename and a usable predecessor method is present; a class -abstract for any other reason is left alone (and fails loudly downstream). Every bridge and -every detected version mismatch logs a loud, greppable warning so the drift is visible and -the operator knows that rebuilding the plug-in against the running PyRIT removes the shim. -This module owns drift-bridging so the loader stays focused on extract/import/register. -""" - -from __future__ import annotations - -import inspect -import logging -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from collections.abc import Callable, Coroutine - from pathlib import Path - from typing import Any - -logger = logging.getLogger(__name__) - -# Renamed scenario extension points: each maps a current (host) abstract method to the -# older predecessor a plug-in may still define. A predecessor took no ``context`` (it read -# ``self._*`` state the base still populates before the build call), so the bridge can -# delegate to it and ignore ``context``. Add an entry here when a future rename needs the -# same mechanical bridge. -_SCENARIO_METHOD_RENAMES: dict[str, str] = { - "_build_atomic_attacks_async": "_get_atomic_attacks_async", -} - - -def warn_on_version_drift(*, extract_dir: Path) -> str | None: - """ - Log a loud warning when a plug-in was built against a different PyRIT minor version. - - Reads the plug-in's ``Requires-Dist: pyrit`` pin from its wheel ``METADATA`` and - compares it to the running ``pyrit`` version. A mismatch is surfaced but never fatal: - the loader tolerates slight drift and lets the import/registration path (plus the - scenario shims below) decide whether the plug-in actually works. Absence of a pin, or - an unparseable one, is silently ignored — this is a best-effort signal, not a gate. - - Args: - extract_dir: The directory the plug-in wheel was extracted to. - - Returns: - str | None: The plug-in's declared PyRIT pin when a drift warning was emitted, - else ``None``. - """ - declared = _read_required_pyrit_version(extract_dir=extract_dir) - if declared is None: - return None - - import pyrit - - running = getattr(pyrit, "__version__", "") or "" - if _same_minor(declared, running): - return None - - logger.warning( - "PLUGIN VERSION DRIFT: plug-in was built against pyrit %s but pyrit %s is running. " - "Proceeding; compatibility shims will bridge known mechanical differences, but " - "rebuild the plug-in against the running pyrit if scenarios fail to load.", - declared, - running or "(unknown)", - ) - return declared - - -def bridge_scenario_extension_points(*, package_name: str) -> list[str]: - """ - Make a plug-in's scenarios concrete by bridging renamed extension-point methods. - - Enumerates concrete-intent ``Scenario`` subclasses owned by ``package_name`` that are - still abstract, and for each one whose only unimplemented abstract methods are known - renames with a usable predecessor, injects a thin adapter so the class satisfies the - current contract. Classes abstract for any other reason are left untouched. - - Args: - package_name: The plug-in's top-level package name. - - Returns: - list[str]: Human-readable descriptions of the bridges applied, for logging. - """ - from pyrit.scenario.core import Scenario - - prefix = f"{package_name}." - applied: list[str] = [] - - for cls in _iter_owned_subclasses(base=Scenario, package_name=package_name, prefix=prefix): - if not inspect.isabstract(cls): - continue - - missing = set(cls.__abstractmethods__) # type: ignore[ty:unresolved-attribute] - # Only bridge when every missing method is a known rename we can satisfy from a - # predecessor the class actually provides. Otherwise the class is abstract for a - # reason we must not paper over, so leave it (it will fail loudly downstream). - if not missing or not missing.issubset(_SCENARIO_METHOD_RENAMES.keys()): - continue - if not all(_has_concrete_method(cls, _SCENARIO_METHOD_RENAMES[name]) for name in missing): - continue - - for new_name in missing: - predecessor = _SCENARIO_METHOD_RENAMES[new_name] - setattr(cls, new_name, _make_rename_bridge(predecessor_name=predecessor)) - applied.append(f"{cls.__name__}.{predecessor} -> {new_name}") - logger.warning( - "PLUGIN COMPAT SHIM: bridged %s.%s to the renamed %s (plug-in built against " - "an older pyrit). Rebuild the plug-in against the running pyrit to remove this shim.", - cls.__name__, - predecessor, - new_name, - ) - cls.__abstractmethods__ = frozenset(cls.__abstractmethods__ - missing) # type: ignore[ty:unresolved-attribute] - - return applied - - -def _make_rename_bridge(*, predecessor_name: str) -> Callable[..., Coroutine[Any, Any, Any]]: - """ - Build an adapter that satisfies a renamed async extension point via its predecessor. - - Args: - predecessor_name: The older method name the plug-in still defines. - - Returns: - Callable[..., Coroutine[Any, Any, Any]]: An async adapter that ignores the new - ``context`` keyword and delegates to the predecessor (which reads ``self._*`` - state the base populates before the build call). - """ - - async def _bridged_build_async(self, **_kwargs: Any) -> Any: # noqa: ANN001 - return await getattr(self, predecessor_name)() - - _bridged_build_async.__name__ = predecessor_name - _bridged_build_async.__qualname__ = predecessor_name - return _bridged_build_async - - -def _has_concrete_method(cls: type, name: str) -> bool: - """ - Return whether ``cls`` provides a concrete (non-abstract) method of the given name. - - Args: - cls: The class to inspect. - name: The method name to look for. - - Returns: - bool: True when the method exists and is not itself abstract. - """ - method = getattr(cls, name, None) - return callable(method) and name not in getattr(cls, "__abstractmethods__", frozenset()) - - -def _iter_owned_subclasses(*, base: type, package_name: str, prefix: str) -> list[type]: - """ - Return every subclass of ``base`` whose module is owned by the plug-in package. - - Args: - base: The base class to enumerate subclasses of. - package_name: The plug-in's top-level package name. - prefix: ``f"{package_name}."`` (passed in to avoid recomputation). - - Returns: - list[type]: The owned subclasses currently loaded in memory. - """ - seen: set[int] = set() - owned: list[type] = [] - stack = list(base.__subclasses__()) - while stack: - cls = stack.pop() - if id(cls) in seen: - continue - seen.add(id(cls)) - stack.extend(cls.__subclasses__()) - module = cls.__module__ or "" - if module == package_name or module.startswith(prefix): - owned.append(cls) - return owned - - -def _read_required_pyrit_version(*, extract_dir: Path) -> str | None: - """ - Return the pinned ``pyrit`` version from the plug-in's wheel ``METADATA``, if any. - - Args: - extract_dir: The directory the plug-in wheel was extracted to. - - Returns: - str | None: The pinned version string (e.g. ``"0.14.0"``) or ``None`` when no - parseable ``Requires-Dist: pyrit`` pin is present. - """ - import re - - for metadata in extract_dir.glob("*.dist-info/METADATA"): - try: - text = metadata.read_text(encoding="utf-8", errors="replace") - except OSError: - continue - for line in text.splitlines(): - if not line.lower().startswith("requires-dist:"): - continue - if not re.match(r"requires-dist:\s*pyrit\b", line, flags=re.IGNORECASE): - continue - match = re.search(r"==\s*([0-9][0-9A-Za-z.\-]*)", line) - if match: - return match.group(1) - return None - - -def _same_minor(left: str, right: str) -> bool: - """ - Return whether two version strings share the same ``major.minor`` prefix. - - Args: - left: A version string (e.g. ``"0.14.0"``). - right: A version string (e.g. ``"0.15.0.dev0"``). - - Returns: - bool: True when both parse to the same ``(major, minor)``; False otherwise. An - unparseable side compares unequal so drift is surfaced rather than hidden. - """ - return _major_minor(left) == _major_minor(right) and _major_minor(left) is not None - - -def _major_minor(version: str) -> tuple[int, int] | None: - """ - Parse the leading ``major.minor`` from a version string. - - Args: - version: A version string. - - Returns: - tuple[int, int] | None: The ``(major, minor)`` pair, or ``None`` when the string - does not begin with two integer components. - """ - parts = version.split(".") - if len(parts) < 2: - return None - try: - return int(parts[0]), int(parts[1]) - except ValueError: - return None diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 9128f374bc..c9c1015997 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -249,7 +249,10 @@ def _register( scenarios: Sequence[ScenarioContribution], techniques: Sequence[TechniqueContribution], ) -> None: + from pyrit.setup.initializers.techniques import build_technique_factories + technique_registry = AttackTechniqueRegistry.get_registry_singleton() + technique_registry.register_from_factories(build_technique_factories(groups=["core"])) for contribution in techniques: technique_registry.register_contributed_factory( factory=contribution.factory, @@ -263,6 +266,12 @@ def _register( scenario_class=contribution.scenario_class, name=contribution.registry_name, ) + try: + scenario_registry.get_class_metadata(contribution.scenario_class) + except Exception as exc: + raise PluginValidationError( + f"Scenario '{contribution.registry_name}' could not build registry metadata: {exc}" + ) from exc @staticmethod def _warn_on_version_drift(*, prepared: PreparedPlugin) -> None: diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index de4ff00553..4ec019090a 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -5,9 +5,8 @@ Unit tests for the PyRIT plug-in loader. These tests build a **mock plug-in wheel** at test time (no dependency on any real -plug-in) and exercise the full consumer mechanism: extract -> sys.path -> -import -> bootstrap -> assert-loaded, plus the accept-load-failures/fail-closed policy and the -silent-failure guards called out in the design brief. +plug-in) and exercise the V1 consumer mechanism: prepare, import, discover, validate, +register, and fail-closed rollback. """ import inspect diff --git a/tests/unit/setup/test_plugin_registration.py b/tests/unit/setup/test_plugin_registration.py index bf828262cf..8476a5321c 100644 --- a/tests/unit/setup/test_plugin_registration.py +++ b/tests/unit/setup/test_plugin_registration.py @@ -4,11 +4,12 @@ import sys from collections.abc import Iterator from pathlib import Path +from unittest.mock import patch import pytest from pyrit.datasets import SeedDatasetProvider -from pyrit.exceptions import PluginCollisionError +from pyrit.exceptions import PluginCollisionError, PluginValidationError from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry from pyrit.setup import PluginFormat, PluginSpec from pyrit.setup.plugin_loader import PluginInitializer @@ -60,7 +61,10 @@ async def test_plugin_initializer_registers_source_technique(tmp_path: Path) -> assert entry.metadata["scenario_names"] == ["airt.rapid_response"] -async def test_plugin_initializer_registers_source_scenario(tmp_path: Path) -> None: +async def test_plugin_initializer_registers_source_scenario( + tmp_path: Path, + patch_central_database: None, +) -> None: source = tmp_path / "private_scenario.py" source.write_text( """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse @@ -71,7 +75,15 @@ class PrivateScenario(RapidResponse): encoding="utf-8", ) - await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() + with patch.dict( + "os.environ", + { + "OPENAI_CHAT_ENDPOINT": "https://example.test", + "OPENAI_CHAT_KEY": "test-key", + "OPENAI_CHAT_MODEL": "test-model", + }, + ): + await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() scenario_class = ScenarioRegistry.get_registry_singleton().get_class("private_scenario") assert scenario_class.__name__ == "PrivateScenario" @@ -153,3 +165,21 @@ async def fetch_dataset_async(self, *, cache=True): await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() assert "PrivateProvider" not in SeedDatasetProvider.get_all_providers() + + +async def test_plugin_initializer_rejects_scenario_that_cannot_build_metadata(tmp_path: Path) -> None: + source = tmp_path / "invalid_scenario.py" + source.write_text( + """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + +class InvalidScenario(RapidResponse): + def __init__(self, *, required_value): + super().__init__() +""", + encoding="utf-8", + ) + + with pytest.raises(PluginValidationError, match="metadata"): + await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() + + assert "invalid_scenario" not in ScenarioRegistry.get_registry_singleton()._classes From d8377e740a6f36664b0dbf2258d7bb9f3bf7a2b7 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 13 Jul 2026 20:36:52 -0700 Subject: [PATCH 28/34] test(setup): make mock plug-in scenario self-contained Use a side-effect-light scenario with inline objective and scorer so plug-in metadata validation is deterministic in isolated and parallel test processes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- tests/unit/setup/test_plugin_loader.py | 28 ++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index 4ec019090a..89f5c59e3e 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -167,11 +167,35 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: (pkg / "scenario.py").write_text( textwrap.dedent( """\ - from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + from pyrit.models import SeedObjective + from pyrit.scenario import DatasetAttackConfiguration, Scenario, ScenarioStrategy + from pyrit.score import SubStringScorer - class MockScenario(RapidResponse): + class MockStrategy(ScenarioStrategy): + ALL = ("all", {"all"}) + DIRECT = ("direct", {"direct"}) + + + class MockScenario(Scenario): \"\"\"Mock plugin scenario for registration test.\"\"\" + + VERSION = 1 + + def __init__(self, *, scenario_result_id=None): + super().__init__( + version=self.VERSION, + strategy_class=MockStrategy, + default_strategy=MockStrategy.ALL, + default_dataset_config=DatasetAttackConfiguration( + seeds=[SeedObjective(value="mock objective")] + ), + objective_scorer=SubStringScorer(substring="mock"), + scenario_result_id=scenario_result_id, + ) + + async def _build_atomic_attacks_async(self, *, context): + return [] """ ), encoding="utf-8", From 26267c78f49b084f1addccb2cfa62173515243ec Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 14 Jul 2026 13:08:17 -0700 Subject: [PATCH 29/34] refactor(scenario): expose all registered RapidResponse techniques The plug-in PR had RapidResponse filter its technique pool to the built-in 'core' tag via get_factories_for_scenario(base_query=TagQuery.all('core')), with a fallback whose return value was discarded. Restore the builder to the upstream form (aligning with #2186): read every registered factory via get_factories_or_raise and build the technique enum with no core filter. This lets an initializer-registered private technique (any tag) surface through --techniques with no scenario-routing bridge, and removes a latent empty-enum path. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/scenario/scenarios/airt/rapid_response.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/pyrit/scenario/scenarios/airt/rapid_response.py b/pyrit/scenario/scenarios/airt/rapid_response.py index 19f3810900..e88558a291 100644 --- a/pyrit/scenario/scenarios/airt/rapid_response.py +++ b/pyrit/scenario/scenarios/airt/rapid_response.py @@ -35,8 +35,9 @@ def _build_rapid_response_technique() -> type[ScenarioTechnique]: """ Build the RapidResponse technique class dynamically from the registered factories. - Reads the singleton ``AttackTechniqueRegistry`` and filters to factories - tagged ``core``. + Reads the singleton ``AttackTechniqueRegistry`` and exposes every registered + technique. Which techniques are available is decided by what initializers + (including a plug-in initializer) have registered before this runs. Returns: type[ScenarioTechnique]: The dynamically generated technique enum class. @@ -45,14 +46,7 @@ def _build_rapid_response_technique() -> type[ScenarioTechnique]: from pyrit.registry.tag_query import TagQuery registry = AttackTechniqueRegistry.get_registry_singleton() - factories = list( - registry.get_factories_for_scenario( - scenario_name="airt.rapid_response", - base_query=TagQuery.all("core"), - ).values() - ) - if not factories: - registry.get_factories_or_raise() + factories = list(registry.get_factories_or_raise().values()) return AttackTechniqueRegistry.build_technique_class_from_factories( # type: ignore[ty:invalid-return-type] class_name="RapidResponseTechnique", From 3fbdd0484c5b149ea9a4ec3f5a5c7e7451613ad8 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 14 Jul 2026 13:08:41 -0700 Subject: [PATCH 30/34] refactor(registry): drop plug-in contributed and external-subclass APIs Remove the registration surface the old discovery-based plug-in system needed, now that a plug-in points at a PyRITInitializer that registers components itself via the ordinary register_from_factories / register_class paths: - AttackTechniqueRegistry.register_contributed_factory and get_factories_for_scenario (the scenario: tag bridge consumer; no remaining callers after the RapidResponse restore). - ScenarioRegistry.register_contributed_scenario and _external_registry_name. - Registry.register_external_subclasses, _external_registry_name, and the external_package branch of _register_subclasses_in_package (dead code with no callers). V1 is intentionally extend-only-by-trust: the contributed-* methods were the only built-in name-collision guards, so a plug-in initializer may now shadow a built-in name. This is accepted for the trusted-artifact model and revisited post-MVP. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- .../components/attack_technique_registry.py | 59 ----------------- .../registry/components/scenario_registry.py | 47 ------------- pyrit/registry/registry.py | 66 ++----------------- 3 files changed, 5 insertions(+), 167 deletions(-) diff --git a/pyrit/registry/components/attack_technique_registry.py b/pyrit/registry/components/attack_technique_registry.py index bf828a7747..3bc00f55dd 100644 --- a/pyrit/registry/components/attack_technique_registry.py +++ b/pyrit/registry/components/attack_technique_registry.py @@ -129,39 +129,6 @@ def register_technique( self.instances.register(factory, name=name, tags=tags, metadata=metadata) logger.debug(f"Registered attack technique factory: {name} ({factory.attack_class.__name__})") - def register_contributed_factory( - self, - *, - factory: AttackTechniqueFactory, - plugin_name: str, - scenario_names: frozenset[str], - ) -> None: - """ - Register one extend-only plug-in technique contribution. - - Args: - factory (AttackTechniqueFactory): The configured attack technique. - plugin_name (str): The owning plug-in name. - scenario_names (frozenset[str]): Scenarios that expose the technique. - - Raises: - PluginCollisionError: If the technique name is already registered. - """ - from pyrit.exceptions import PluginCollisionError - - if factory.name in self.instances: - raise PluginCollisionError(f"Attack technique name '{factory.name}' is already registered.") - self.register_technique( - name=factory.name, - factory=factory, - tags=factory.technique_tags, - metadata={ - "plugin_name": plugin_name, - "scenario_names": sorted(scenario_names), - "identifier_hash": factory.get_identifier().hash, - }, - ) - def get_factories(self) -> dict[str, AttackTechniqueFactory]: """ Return all registered factories as a name→factory dict. @@ -202,32 +169,6 @@ def get_factories_or_raise(self) -> dict[str, AttackTechniqueFactory]: ) return factories - def get_factories_for_scenario( - self, - *, - scenario_name: str, - base_query: TagQuery | None = None, - ) -> dict[str, AttackTechniqueFactory]: - """ - Return built-in-query matches plus contributions applicable to a scenario. - - Args: - scenario_name (str): The scanner-facing scenario registry name. - base_query (TagQuery | None): Optional query selecting the scenario's - built-in technique pool. - - Returns: - dict[str, AttackTechniqueFactory]: Matching factories, sorted by name. - """ - factories: dict[str, AttackTechniqueFactory] = {} - for entry in self.instances.get_all_instances(): - built_in_match = base_query is not None and base_query.matches(set(entry.instance.technique_tags)) - applicable_scenarios = entry.metadata.get("scenario_names", []) - contributed_match = scenario_name in applicable_scenarios - if built_in_match or contributed_match: - factories[entry.name] = entry.instance - return factories - @property def scorer_override_policy(self) -> ScorerOverridePolicy: """The policy applied when a scenario scorer is incompatible with an attack's annotation.""" diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py index 9bd66fae4b..1551148085 100644 --- a/pyrit/registry/components/scenario_registry.py +++ b/pyrit/registry/components/scenario_registry.py @@ -109,53 +109,6 @@ def _get_registry_name(self, cls: type[Scenario]) -> str: return relative return class_name_to_snake_case(cls.__name__, suffix="Scenario") - def _external_registry_name(self, cls: type[Scenario], *, package_name: str) -> str: - """ - Key a plug-in scenario by its module path within the plug-in package. - - Mirrors the built-in dotted-name scheme (module path relative to the scenarios - package) so plug-in scenarios read like built-ins and same-named scenarios in - different submodules do not collide on a bare class name. The name is the class's - module relative to the plug-in's top-level ``package_name`` with a leading - ``scenario.scenarios.`` / ``scenarios.`` segment stripped to match built-in style; - it falls back to the suffix-stripped snake_case class name when no module context - is available. - - Args: - cls (type[Scenario]): The plug-in scenario class. - package_name (str): The plug-in's top-level package name. - - Returns: - str: The dotted registry name for the plug-in scenario. - """ - module = cls.__module__ or "" - prefix = f"{package_name}." - relative = module[len(prefix) :] if module.startswith(prefix) else module - for marker in ("scenario.scenarios.", "scenarios."): - index = relative.find(marker) - if index != -1: - relative = relative[index + len(marker) :] - break - return relative or class_name_to_snake_case(cls.__name__, suffix="Scenario") - - def register_contributed_scenario(self, *, scenario_class: type[Scenario], name: str) -> None: - """ - Register one extend-only plug-in scenario contribution. - - Args: - scenario_class (type[Scenario]): The concrete scenario class. - name (str): The scanner-facing dotted registry name. - - Raises: - PluginCollisionError: If the name belongs to a built-in or prior contribution. - """ - from pyrit.exceptions import PluginCollisionError - - self._ensure_discovered() - if name in self._classes: - raise PluginCollisionError(f"Scenario name '{name}' is already registered.") - self.register_class(scenario_class, name=name) - def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: """ Build metadata for a Scenario class. diff --git a/pyrit/registry/registry.py b/pyrit/registry/registry.py index 9946f2e276..3c89249796 100644 --- a/pyrit/registry/registry.py +++ b/pyrit/registry/registry.py @@ -283,70 +283,31 @@ def _discover(self) -> None: logger.debug(f"Skipping lazily-exported '{exported_name}': {exc}") self._register_subclasses_in_package(package_name=package_name) - def register_external_subclasses(self, *, package_name: str) -> int: - """ - Register concrete base-type subclasses that live under an external package. - - This is how a loaded plug-in's components (e.g. scenarios shipped in a plug-in - wheel) enter the registry: the plug-in's modules are imported elsewhere, then - this enumerates the base type's in-memory subclasses under ``package_name`` and - registers each, alongside the built-in discovery package. Built-in discovery is - **not** triggered, so lazy discovery is preserved, and classes already registered - (e.g. by the plug-in's own bootstrap) are left untouched. - - Args: - package_name (str): The external (plug-in) top-level package name. - - Returns: - int: The number of classes newly registered from the package. - """ - return self._register_subclasses_in_package( - package_name=package_name, skip_registered_classes=True, external_package=package_name - ) - - def _register_subclasses_in_package( - self, *, package_name: str, skip_registered_classes: bool = False, external_package: str | None = None - ) -> int: + def _register_subclasses_in_package(self, *, package_name: str) -> int: """ Register every concrete base-type subclass whose module lives under a package. - Shared by built-in discovery (``_discover``) and external plug-in - registration (``register_external_subclasses``). Enumerates the in-memory - subclasses of ``_base_type()``, keeps those whose module is ``package_name`` or a - submodule of it, and registers each. Built-in classes are named by - ``_get_registry_name``; when ``external_package`` is set, names come from - ``_external_registry_name`` so a plug-in's classes are keyed by their module path - within the plug-in (mirroring built-ins), which keeps same-named classes in - different submodules from colliding. When ``skip_registered_classes`` is set, - classes already in the catalog (by identity) are skipped so an explicit bootstrap - registration is never shadowed by a fallback-named duplicate. + Used by built-in discovery (``_discover``): enumerates the in-memory subclasses of + ``_base_type()``, keeps those whose module is ``package_name`` or a submodule of + it, names each via ``_get_registry_name``, and registers it. Args: package_name (str): The package whose subclasses to register. - skip_registered_classes (bool): Skip classes already registered by identity. - external_package (str | None): When set, the plug-in package the classes belong - to; names are derived relative to it via ``_external_registry_name``. Returns: int: The number of classes newly registered. """ base = self._base_type() package_prefix = f"{package_name}." - already_registered = set(self._classes.values()) if skip_registered_classes else set() count = 0 for cls in self._iter_concrete_subclasses(base): module = cls.__module__ or "" if module != package_name and not module.startswith(package_prefix): continue - if skip_registered_classes and cls in already_registered: - continue if (cls.__doc__ or "").strip().startswith("Deprecated alias"): logger.debug(f"Skipping deprecated alias: {cls.__name__}") continue - if external_package is not None: - name = self._external_registry_name(cls, package_name=external_package) - else: - name = self._get_registry_name(cls) + name = self._get_registry_name(cls) existing = self._classes.get(name) if existing is not None and existing is not cls: logger.warning( @@ -359,23 +320,6 @@ def _register_subclasses_in_package( logger.debug(f"Registered {base.__name__} class: {name} ({cls.__name__})") return count - def _external_registry_name(self, cls: type[T], *, package_name: str) -> str: - """ - Return the catalog name for a class discovered in an external (plug-in) package. - - Defaults to the same name built-in discovery would use. Registries that key - built-ins by module path (e.g. scenarios) override this so plug-in classes get an - equally path-based, collision-resistant name relative to the plug-in package. - - Args: - cls (type[T]): The external class being registered. - package_name (str): The plug-in's top-level package name. - - Returns: - str: The catalog name to register the class under. - """ - return self._get_registry_name(cls) - @staticmethod def _iter_concrete_subclasses(base: type[T]) -> list[type[T]]: """ From a9764ac5eaaf404f07002f2977905b9a350b21de Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 14 Jul 2026 13:08:57 -0700 Subject: [PATCH 31/34] refactor(setup): point plug-ins at a PyRITInitializer, drop discovery and wheels Reframe a plug-in as a config-declared pointer to a private PyRITInitializer that lives outside the PyRIT tree, run privileged-first. The framework only anchors the source root on sys.path, imports the dotted initializer, runs it, and fails closed; the initializer owns all registration (scenarios, techniques, datasets, targets). - PluginSpec is now {name, source, initializer} (source-only; no format/wheel/package). - PluginInitializer resolves and runs the dotted initializer; no component discovery, no transactional snapshot/rollback, no version-drift machinery. - Delete plugin_discovery.py and plugin_formats.py (wheel extraction, source/scenario discovery, the scenario: tag bridge). - Prune plug-in exceptions to PluginLoadError + PluginSourceNotFoundError. - Drop the PluginFormat export and refresh the ConfigurationLoader plugins docstring. This is a thin convenience + ordering layer over --initialization-scripts: config-declared, run-first, and package-root-anchored so a packaged initializer with intra-package imports loads. Wheels, discovery, and versioning are out of scope for the MVP. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- pyrit/exceptions/__init__.py | 12 - pyrit/exceptions/exception_classes.py | 24 -- pyrit/setup/__init__.py | 3 +- pyrit/setup/configuration_loader.py | 4 +- pyrit/setup/plugin_discovery.py | 288 ------------------------ pyrit/setup/plugin_formats.py | 265 ---------------------- pyrit/setup/plugin_loader.py | 311 +++++--------------------- pyrit/setup/plugin_spec.py | 106 ++------- 8 files changed, 83 insertions(+), 930 deletions(-) delete mode 100644 pyrit/setup/plugin_discovery.py delete mode 100644 pyrit/setup/plugin_formats.py diff --git a/pyrit/exceptions/__init__.py b/pyrit/exceptions/__init__.py index 1070cde32a..abb5552ccc 100644 --- a/pyrit/exceptions/__init__.py +++ b/pyrit/exceptions/__init__.py @@ -10,14 +10,8 @@ ExperimentalWarning, InvalidJsonException, MissingPromptPlaceholderException, - PluginCollisionError, - PluginDiscoveryError, - PluginImportError, PluginLoadError, - PluginRegisteredNothingError, PluginSourceNotFoundError, - PluginValidationError, - PluginWheelNotFoundError, PyritException, RateLimitException, ScorerLLMResponseBlockedException, @@ -61,14 +55,8 @@ "handle_bad_request_exception", "InvalidJsonException", "MissingPromptPlaceholderException", - "PluginCollisionError", - "PluginDiscoveryError", - "PluginImportError", "PluginLoadError", - "PluginRegisteredNothingError", "PluginSourceNotFoundError", - "PluginValidationError", - "PluginWheelNotFoundError", "PyritException", "pyrit_custom_result_retry", "pyrit_json_retry", diff --git a/pyrit/exceptions/exception_classes.py b/pyrit/exceptions/exception_classes.py index d45dfbd2fa..e1bdf5e40b 100644 --- a/pyrit/exceptions/exception_classes.py +++ b/pyrit/exceptions/exception_classes.py @@ -265,34 +265,10 @@ class PluginLoadError(RuntimeError): """Base error raised when a configured plug-in fails to load.""" -class PluginCollisionError(PluginLoadError): - """A plug-in contribution conflicts with an existing scenario or attack technique.""" - - -class PluginDiscoveryError(PluginLoadError): - """A plug-in imported but its scenario or attack-technique contributions could not be discovered.""" - - -class PluginImportError(PluginLoadError): - """A plug-in source module or wheel package could not be imported.""" - - -class PluginRegisteredNothingError(PluginLoadError): - """The plug-in imported cleanly but contributed no scenarios or attack techniques.""" - - class PluginSourceNotFoundError(PluginLoadError): """The configured plug-in source path does not point to a readable Python file or package.""" -class PluginValidationError(PluginLoadError): - """A discovered plug-in scenario or attack technique violates the current PyRIT contract.""" - - -class PluginWheelNotFoundError(PluginLoadError): - """The configured plug-in wheel path does not point to a readable ``.whl`` file.""" - - def pyrit_custom_result_retry( retry_function: Callable[..., bool], retry_max_num_attempts: int | None = None ) -> Callable[..., Any]: diff --git a/pyrit/setup/__init__.py b/pyrit/setup/__init__.py index a94f4b3ec4..9260523b0d 100644 --- a/pyrit/setup/__init__.py +++ b/pyrit/setup/__init__.py @@ -11,7 +11,7 @@ MemoryDatabaseType, initialize_pyrit_async, ) -from pyrit.setup.plugin_spec import PluginFormat, PluginSpec +from pyrit.setup.plugin_spec import PluginSpec __all__ = [ "AZURE_SQL", @@ -20,7 +20,6 @@ "initialize_pyrit_async", "initialize_from_config_async", "MemoryDatabaseType", - "PluginFormat", "ConfigurationLoader", "PluginSpec", ] diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 99d678f2bc..299ebe8d1b 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -111,8 +111,8 @@ class ConfigurationLoader(YamlLoadable): operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. plugins: List of plug-ins to load as the guaranteed-first initialization phase. - V1 accepts at most one explicit mapping with ``name``, ``format``, exactly - one matching ``source`` or ``wheel`` path, and an optional ``package``. + V1 accepts at most one explicit mapping with ``name``, a ``source`` path, and + a dotted ``initializer`` (module.Class) pointing at a ``PyRITInitializer``. Empty means no plug-ins. Example YAML configuration: diff --git a/pyrit/setup/plugin_discovery.py b/pyrit/setup/plugin_discovery.py deleted file mode 100644 index e9b45a3e15..0000000000 --- a/pyrit/setup/plugin_discovery.py +++ /dev/null @@ -1,288 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""Import prepared plug-ins and discover owned PyRIT contributions.""" - -from __future__ import annotations - -import asyncio -import importlib -import inspect -import pkgutil -import sys -from dataclasses import dataclass -from pathlib import Path -from typing import TYPE_CHECKING, cast - -from pyrit.exceptions import PluginDiscoveryError, PluginImportError -from pyrit.models import class_name_to_snake_case -from pyrit.scenario import AttackTechniqueFactory, Scenario - -if TYPE_CHECKING: - from collections.abc import Callable, Sequence - from types import ModuleType - - from pyrit.setup.plugin_formats import PreparedPlugin - - -@dataclass(frozen=True) -class ImportedPlugin: - """The imported modules owned by one prepared plug-in.""" - - prepared: PreparedPlugin - modules: tuple[ModuleType, ...] - - -@dataclass(frozen=True) -class ScenarioContribution: - """A discovered scenario class and its proposed registry name.""" - - scenario_class: type[Scenario] - registry_name: str - - -@dataclass(frozen=True) -class TechniqueContribution: - """A configured attack technique and the scenarios that expose it.""" - - factory: AttackTechniqueFactory - scenario_names: frozenset[str] - - -async def import_plugin_async(*, prepared: PreparedPlugin) -> ImportedPlugin: - """ - Import a prepared plug-in and all submodules under package entry modules. - - Args: - prepared (PreparedPlugin): The prepared source or wheel import root. - - Returns: - ImportedPlugin: The imported, ownership-validated modules. - - Raises: - PluginImportError: If import or ownership validation fails. - """ - return await asyncio.to_thread(_import_plugin, prepared=prepared) - - -def _import_plugin(*, prepared: PreparedPlugin) -> ImportedPlugin: - import_root = str(prepared.import_root) - if import_root not in sys.path: - sys.path.insert(0, import_root) - - try: - for entry_module in prepared.entry_modules: - module = importlib.import_module(entry_module) - _verify_module_location(module=module, import_root=prepared.import_root) - _import_submodules(module=module) - except Exception as exc: - raise PluginImportError( - f"Could not import plug-in '{prepared.spec.name}' from {prepared.import_root}: {exc}" - ) from exc - - modules = tuple( - module - for name, module in sorted(sys.modules.items()) - if module is not None and _name_owned_by_any(name=name, prefixes=prepared.owned_module_prefixes) - ) - for module in modules: - _verify_module_location(module=module, import_root=prepared.import_root) - return ImportedPlugin(prepared=prepared, modules=modules) - - -def _import_submodules(*, module: ModuleType) -> None: - module_path = getattr(module, "__path__", None) - if not module_path: - return - for item in pkgutil.walk_packages(module_path, prefix=f"{module.__name__}."): - importlib.import_module(item.name) - - -def _verify_module_location(*, module: ModuleType, import_root: Path) -> None: - root = import_root.resolve() - raw_locations = list(getattr(module, "__path__", []) or []) - module_file = getattr(module, "__file__", None) - if module_file: - raw_locations.append(module_file) - locations = [Path(location).resolve() for location in raw_locations if location] - if locations and not any(location.is_relative_to(root) for location in locations): - raise ValueError(f"Imported module '{module.__name__}' resolved outside plug-in root {root}.") - - -def discover_scenarios(*, imported: ImportedPlugin) -> list[ScenarioContribution]: - """ - Discover concrete, module-owned Scenario subclasses. - - Args: - imported (ImportedPlugin): The imported plug-in modules. - - Returns: - list[ScenarioContribution]: Deterministically ordered scenario contributions. - - Raises: - PluginDiscoveryError: If one module defines multiple implicit scenarios. - """ - contributions: list[ScenarioContribution] = [] - for module in imported.modules: - classes = _module_defined_subclasses(module=module, base_class=Scenario) - if len(classes) > 1: - raise PluginDiscoveryError( - f"Plug-in module '{module.__name__}' defines multiple Scenario classes. " - "Move each Scenario class into a separate module." - ) - if classes: - contributions.append( - ScenarioContribution( - scenario_class=classes[0], - registry_name=_scenario_registry_name(imported=imported, module=module, cls=classes[0]), - ) - ) - return sorted(contributions, key=lambda item: item.registry_name) - - -def discover_techniques(*, imported: ImportedPlugin) -> list[TechniqueContribution]: - """ - Discover configured attack-technique factories from owned modules. - - Args: - imported (ImportedPlugin): The imported plug-in modules. - - Returns: - list[TechniqueContribution]: Deterministically ordered technique contributions. - - Raises: - PluginDiscoveryError: If a factory export is malformed, lacks applicability, - or duplicates another contribution name. - """ - contributions: list[TechniqueContribution] = [] - seen_objects: set[int] = set() - for module in imported.modules: - explicit = [ - value - for value in vars(module).values() - if isinstance(value, TechniqueContribution) and id(value.factory) not in seen_objects - ] - contributions.extend(explicit) - seen_objects.update(id(item.factory) for item in explicit) - - factory_builder = getattr(module, "get_technique_factories", None) - if callable(factory_builder) and getattr(factory_builder, "__module__", None) == module.__name__: - contributions.extend( - _contributions_from_factory_builder( - module=module, - factory_builder=factory_builder, - seen_objects=seen_objects, - ) - ) - - for value in vars(module).values(): - if isinstance(value, AttackTechniqueFactory) and id(value) not in seen_objects: - contributions.append(_contribution_from_factory(factory=value, module_name=module.__name__)) - seen_objects.add(id(value)) - - by_name: dict[str, TechniqueContribution] = {} - for contribution in contributions: - if not contribution.scenario_names: - raise PluginDiscoveryError( - f"Attack technique '{contribution.factory.name}' must declare at least one applicable scenario." - ) - existing = by_name.get(contribution.factory.name) - if existing is not None and existing.factory is not contribution.factory: - raise PluginDiscoveryError( - f"Plug-in discovered duplicate attack technique name '{contribution.factory.name}'." - ) - contribution.factory.get_identifier() - by_name[contribution.factory.name] = contribution - return [by_name[name] for name in sorted(by_name)] - - -def _contributions_from_factory_builder( - *, - module: ModuleType, - factory_builder: Callable[[], object], - seen_objects: set[int], -) -> list[TechniqueContribution]: - signature = inspect.signature(factory_builder) - required = [ - parameter - for parameter in signature.parameters.values() - if parameter.default is inspect.Parameter.empty - and parameter.kind - not in { - inspect.Parameter.VAR_POSITIONAL, - inspect.Parameter.VAR_KEYWORD, - } - ] - if required: - raise PluginDiscoveryError( - f"Plug-in function '{module.__name__}.get_technique_factories' must take no required arguments." - ) - result = factory_builder() - if not isinstance(result, (list, tuple)) or not all(isinstance(item, AttackTechniqueFactory) for item in result): - raise PluginDiscoveryError( - f"Plug-in function '{module.__name__}.get_technique_factories' must return " - "AttackTechniqueFactory instances." - ) - factories = cast("Sequence[AttackTechniqueFactory]", result) - contributions = [ - _contribution_from_factory(factory=factory, module_name=module.__name__) - for factory in factories - if id(factory) not in seen_objects - ] - seen_objects.update(id(item.factory) for item in contributions) - return contributions - - -def _contribution_from_factory(*, factory: AttackTechniqueFactory, module_name: str) -> TechniqueContribution: - prefix = "scenario:" - scenario_names = frozenset( - tag.removeprefix(prefix) for tag in factory.technique_tags if tag.startswith(prefix) and tag != prefix - ) - if not scenario_names: - raise PluginDiscoveryError( - f"Attack technique '{factory.name}' from '{module_name}' must declare an applicable scenario " - f"using a '{prefix}' tag or an explicit TechniqueContribution." - ) - return TechniqueContribution(factory=factory, scenario_names=scenario_names) - - -def _module_defined_subclasses(*, module: ModuleType, base_class: type[Scenario]) -> list[type[Scenario]]: - return [ - candidate - for _, candidate in inspect.getmembers(module, inspect.isclass) - if candidate is not base_class - and issubclass(candidate, base_class) - and not inspect.isabstract(candidate) - and candidate.__module__ == module.__name__ - ] - - -def _scenario_registry_name( - *, - imported: ImportedPlugin, - module: ModuleType, - cls: type[Scenario], -) -> str: - prefix = next( - ( - owned - for owned in sorted(imported.prepared.owned_module_prefixes, key=len, reverse=True) - if module.__name__ == owned or module.__name__.startswith(f"{owned}.") - ), - "", - ) - relative = module.__name__[len(prefix) :].lstrip(".") if prefix else module.__name__ - for marker in ("scenario.scenarios.", "scenarios."): - if marker in relative: - relative = relative.split(marker, 1)[1] - break - if relative: - return relative - source = imported.prepared.spec.source - if source is not None and source.is_file(): - return source.stem - return class_name_to_snake_case(cls.__name__, suffix="Scenario") - - -def _name_owned_by_any(*, name: str, prefixes: tuple[str, ...]) -> bool: - return any(name == prefix or name.startswith(f"{prefix}.") for prefix in prefixes) diff --git a/pyrit/setup/plugin_formats.py b/pyrit/setup/plugin_formats.py deleted file mode 100644 index 9bf6e553a8..0000000000 --- a/pyrit/setup/plugin_formats.py +++ /dev/null @@ -1,265 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""Artifact-format adapters for scenario plug-in consumption.""" - -from __future__ import annotations - -import asyncio -import hashlib -import os -import shutil -import tempfile -from dataclasses import dataclass -from pathlib import Path - -from pyrit.exceptions import PluginSourceNotFoundError, PluginWheelNotFoundError -from pyrit.setup.plugin_spec import PluginFormat, PluginSpec - - -@dataclass(frozen=True) -class PreparedPlugin: - """A prepared import root shared by every plug-in artifact format.""" - - spec: PluginSpec - import_root: Path - entry_modules: tuple[str, ...] - owned_module_prefixes: tuple[str, ...] - artifact_fingerprint: str - declared_pyrit_version: str | None = None - - -class SourcePluginFormat: - """Prepare a server-local Python source file for plug-in discovery.""" - - async def prepare_async(self, *, spec: PluginSpec) -> PreparedPlugin: - """ - Validate one source-file plug-in without importing or executing it. - - Args: - spec (PluginSpec): The normalized source declaration. - - Returns: - PreparedPlugin: The prepared source import root and ownership metadata. - - Raises: - ValueError: If the spec is not source format. - PluginSourceNotFoundError: If the source is absent or not a Python file. - """ - return await asyncio.to_thread(self._prepare, spec=spec) - - def _prepare(self, *, spec: PluginSpec) -> PreparedPlugin: - if spec.format is not PluginFormat.SOURCE or spec.source is None: - raise ValueError("SourcePluginFormat requires a source-format PluginSpec.") - - source = spec.source.expanduser().resolve() - if source.is_file(): - return self._prepare_file(spec=spec, source=source) - if source.is_dir(): - return self._prepare_package(spec=spec, source=source) - raise PluginSourceNotFoundError(f"Plug-in source does not point to an existing path: {source}") - - def _prepare_file(self, *, spec: PluginSpec, source: Path) -> PreparedPlugin: - if source.suffix != ".py" or not source.stem.isidentifier(): - raise PluginSourceNotFoundError( - f"Plug-in source must be a Python file with an importable filename: {source}" - ) - if spec.package is not None and spec.package != source.stem: - raise ValueError( - f"Single-file plug-in package must match the source module '{source.stem}', got '{spec.package}'." - ) - - return PreparedPlugin( - spec=spec, - import_root=source.parent, - entry_modules=(source.stem,), - owned_module_prefixes=(source.stem,), - artifact_fingerprint=self._fingerprint(source=source), - ) - - def _prepare_package(self, *, spec: PluginSpec, source: Path) -> PreparedPlugin: - package_name = source.name - if not package_name.isidentifier() or not (source / "__init__.py").is_file(): - raise PluginSourceNotFoundError( - f"Plug-in source directory must be an importable package containing __init__.py: {source}" - ) - - entry_module = spec.package or package_name - if entry_module != package_name and not entry_module.startswith(f"{package_name}."): - raise ValueError(f"Plug-in entry module '{entry_module}' must be inside source package '{package_name}'.") - self._validate_entry_module(source=source, entry_module=entry_module) - - return PreparedPlugin( - spec=spec, - import_root=source.parent, - entry_modules=(entry_module,), - owned_module_prefixes=(package_name,), - artifact_fingerprint=self._fingerprint_package(source=source), - ) - - @staticmethod - def _validate_entry_module(*, source: Path, entry_module: str) -> None: - relative_parts = entry_module.split(".")[1:] - if not relative_parts: - return - module_path = source.joinpath(*relative_parts) - if module_path.with_suffix(".py").is_file() or (module_path / "__init__.py").is_file(): - return - raise ValueError(f"Plug-in entry module '{entry_module}' does not exist inside {source}.") - - @staticmethod - def _fingerprint(*, source: Path) -> str: - digest = hashlib.sha256() - with source.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - @classmethod - def _fingerprint_package(cls, *, source: Path) -> str: - digest = hashlib.sha256() - for path in sorted(source.rglob("*")): - if not path.is_file() or "__pycache__" in path.parts or path.suffix == ".pyc": - continue - resolved = path.resolve() - if not resolved.is_relative_to(source): - raise PluginSourceNotFoundError(f"Plug-in package file escapes the source root: {path}") - digest.update(path.relative_to(source).as_posix().encode("utf-8")) - digest.update(cls._fingerprint(source=path).encode("ascii")) - return digest.hexdigest() - - -class WheelPluginFormat: - """Prepare a pre-built Python wheel for plug-in discovery.""" - - _FINGERPRINT_FILE = ".plugin_wheel_fingerprint" - - def __init__(self, *, base_dir: Path | None = None) -> None: - """ - Initialize the wheel adapter. - - Args: - base_dir (Path | None): Extraction root. Defaults to ``PLUGIN_DIR`` or - PyRIT's ``.plugin`` directory. - """ - self._base_dir = base_dir - - async def prepare_async(self, *, spec: PluginSpec) -> PreparedPlugin: - """ - Validate and extract one wheel-format plug-in. - - Args: - spec (PluginSpec): The normalized wheel declaration. - - Returns: - PreparedPlugin: The prepared wheel import root and ownership metadata. - - Raises: - ValueError: If the spec is not wheel format or package discovery is ambiguous. - PluginWheelNotFoundError: If the wheel path is absent or invalid. - """ - return await asyncio.to_thread(self._prepare, spec=spec) - - def _prepare(self, *, spec: PluginSpec) -> PreparedPlugin: - if spec.format is not PluginFormat.WHEEL or spec.wheel is None: - raise ValueError("WheelPluginFormat requires a wheel-format PluginSpec.") - wheel = spec.wheel.expanduser() - if not wheel.is_file(): - raise PluginWheelNotFoundError(f"Plug-in wheel does not point to an existing file: {wheel}") - if wheel.suffix != ".whl": - raise PluginWheelNotFoundError(f"Plug-in wheel must be a .whl file, got: {wheel}") - - fingerprint = self._fingerprint(wheel=wheel) - extract_dir = self._extract(wheel=wheel, fingerprint=fingerprint) - package = self._resolve_package(extract_dir=extract_dir, explicit_package=spec.package) - return PreparedPlugin( - spec=spec, - import_root=extract_dir, - entry_modules=(package,), - owned_module_prefixes=(package,), - artifact_fingerprint=fingerprint, - declared_pyrit_version=self._read_required_pyrit_version(extract_dir=extract_dir), - ) - - def _extract(self, *, wheel: Path, fingerprint: str) -> Path: - from pyrit.common.safe_extract import safe_extract_zip - - base_dir = self._resolve_base_dir() - base_dir.mkdir(parents=True, exist_ok=True) - extract_dir = base_dir / wheel.stem - marker = extract_dir / self._FINGERPRINT_FILE - if extract_dir.is_dir() and marker.is_file(): - try: - if marker.read_text(encoding="utf-8").strip() == fingerprint: - return extract_dir - except OSError: - pass - - temp_dir = Path(tempfile.mkdtemp(prefix=f".{wheel.stem}.tmp-", dir=base_dir)) - try: - safe_extract_zip(source=wheel, dest_dir=temp_dir) - (temp_dir / self._FINGERPRINT_FILE).write_text(fingerprint, encoding="utf-8") - if extract_dir.exists(): - shutil.rmtree(extract_dir) - os.replace(temp_dir, extract_dir) - finally: - if temp_dir.exists(): - shutil.rmtree(temp_dir, ignore_errors=True) - return extract_dir - - def _resolve_base_dir(self) -> Path: - if self._base_dir is not None: - return self._base_dir.expanduser().resolve() - override = os.getenv("PLUGIN_DIR") - if override: - return Path(override).expanduser().resolve() - from pyrit.common import path - - return (Path(path.HOME_PATH) / ".plugin").resolve() - - @staticmethod - def _fingerprint(*, wheel: Path) -> str: - digest = hashlib.sha256() - with wheel.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - @staticmethod - def _resolve_package(*, extract_dir: Path, explicit_package: str | None) -> str: - if explicit_package: - return explicit_package - for dist_info in sorted(extract_dir.glob("*.dist-info")): - top_level = dist_info / "top_level.txt" - if top_level.is_file(): - names = [line.strip() for line in top_level.read_text(encoding="utf-8").splitlines() if line.strip()] - if len(names) == 1: - return names[0] - - candidates = sorted( - child.name - for child in extract_dir.iterdir() - if child.is_dir() and not child.name.endswith((".dist-info", ".data")) and (child / "__init__.py").is_file() - ) - if len(candidates) == 1: - return candidates[0] - if not candidates: - raise ValueError(f"Could not find an importable top-level package in {extract_dir}.") - raise ValueError(f"Found multiple top-level packages in {extract_dir}: {candidates}.") - - @staticmethod - def _read_required_pyrit_version(*, extract_dir: Path) -> str | None: - import re - - for metadata in extract_dir.glob("*.dist-info/METADATA"): - try: - text = metadata.read_text(encoding="utf-8", errors="replace") - except OSError: - continue - for line in text.splitlines(): - if not re.match(r"Requires-Dist:\s*pyrit\b", line, flags=re.IGNORECASE): - continue - match = re.search(r"==\s*([0-9][0-9A-Za-z.\-]*)", line) - if match: - return match.group(1) - return None diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index c9c1015997..10c891a67a 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -1,41 +1,31 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -"""Activate private scenarios and attack techniques from configured plug-ins.""" +""" +Activate a configured plug-in by running its ``PyRITInitializer`` first. + +A plug-in is a config-declared pointer to a ``PyRITInitializer`` that lives outside the +PyRIT tree. ``ConfigurationLoader`` prepends this privileged initializer ahead of the +user-configured ones, so the plug-in's components exist before any later initializer or +lazy catalog consumer runs. The plug-in initializer itself owns registration of every +scenario, attack technique, dataset, and default target it wants discoverable; PyRIT does +no component discovery here. +""" from __future__ import annotations +import asyncio +import importlib import logging import sys -from dataclasses import dataclass -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING -from pyrit.datasets import SeedDatasetProvider -from pyrit.exceptions import ( - PluginCollisionError, - PluginLoadError, - PluginRegisteredNothingError, - PluginValidationError, -) -from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry -from pyrit.setup.plugin_discovery import ( - ScenarioContribution, - TechniqueContribution, - discover_scenarios, - discover_techniques, - import_plugin_async, -) -from pyrit.setup.plugin_formats import PreparedPlugin, SourcePluginFormat, WheelPluginFormat -from pyrit.setup.plugin_spec import PluginFormat +from pyrit.exceptions import PluginLoadError, PluginSourceNotFoundError from pyrit.setup.pyrit_initializer import PyRITInitializer if TYPE_CHECKING: from collections.abc import Sequence - from pyrit.models import ComponentIdentifier - from pyrit.registry import ScenarioMetadata - from pyrit.registry.instance_registry import DefaultInstanceRegistry, RegistryEntry - from pyrit.scenario import AttackTechniqueFactory, Scenario from pyrit.setup.plugin_spec import PluginSpec logger = logging.getLogger(__name__) @@ -43,26 +33,8 @@ _REMEDIATION = "Fix or remove the plug-in entry from .pyrit_conf, then restart PyRIT." -async def load_plugins_if_configured_async(*, plugins: Sequence[PluginSpec]) -> None: - """ - Activate the configured V1 plug-in. - - Args: - plugins: The normalized plug-in declarations. - - Raises: - PluginLoadError: If activation fails. - ValueError: If more than one plug-in is supplied. - """ - if not plugins: - return - if len(plugins) > 1: - raise ValueError("V1 supports one plug-in at a time; plug-in composition is not supported.") - await PluginLoader(spec=plugins[0]).load_async() - - class PluginInitializer(PyRITInitializer): - """Privileged config-owned initializer that activates one scenario plug-in.""" + """Privileged, config-owned initializer that runs one plug-in's ``PyRITInitializer`` first.""" def __init__(self, *, plugins: Sequence[PluginSpec]) -> None: """ @@ -72,235 +44,68 @@ def __init__(self, *, plugins: Sequence[PluginSpec]) -> None: plugins (Sequence[PluginSpec]): The normalized plug-in declarations. Raises: - ValueError: If V1 receives anything other than one plug-in. + ValueError: If anything other than exactly one plug-in is supplied (V1). """ super().__init__() if len(plugins) != 1: raise ValueError("PluginInitializer requires exactly one plug-in in V1.") - self._plugins = tuple(plugins) + self._spec = plugins[0] @property def plugins(self) -> list[PluginSpec]: """The normalized plug-in declarations.""" - return list(self._plugins) + return [self._spec] async def initialize_async(self) -> None: - """Activate the configured plug-in before user-configured initializers.""" - await load_plugins_if_configured_async(plugins=self._plugins) - - -@dataclass -class _RegistrySnapshot: - scenario_classes: dict[str, type[Scenario]] - scenario_metadata: dict[str, ScenarioMetadata] | None - scenario_discovered: bool - technique_entries: dict[str, RegistryEntry[AttackTechniqueFactory]] - technique_metadata: list[ComponentIdentifier] | None - providers: dict[str, type[SeedDatasetProvider]] - sys_path: list[str] - module_names: set[str] - - @classmethod - def capture(cls) -> _RegistrySnapshot: - scenario_registry = ScenarioRegistry.get_registry_singleton() - technique_registry = AttackTechniqueRegistry.get_registry_singleton() - technique_instances = cast( - "DefaultInstanceRegistry[AttackTechniqueFactory]", - technique_registry.instances, - ) - return cls( - scenario_classes=dict(scenario_registry._classes), - scenario_metadata=( - dict(scenario_registry._metadata_cache) if scenario_registry._metadata_cache is not None else None - ), - scenario_discovered=scenario_registry._discovered, - technique_entries=dict(technique_instances._registry_items), - technique_metadata=( - list(technique_instances._metadata_cache) if technique_instances._metadata_cache is not None else None - ), - providers=dict(SeedDatasetProvider._registry), - sys_path=list(sys.path), - module_names=set(sys.modules), - ) - - def restore(self, *, owned_prefixes: tuple[str, ...] = ()) -> None: - scenario_registry = ScenarioRegistry.get_registry_singleton() - scenario_registry._classes = dict(self.scenario_classes) - scenario_registry._metadata_cache = dict(self.scenario_metadata) if self.scenario_metadata is not None else None - scenario_registry._discovered = self.scenario_discovered - - technique_instances = cast( - "DefaultInstanceRegistry[AttackTechniqueFactory]", - AttackTechniqueRegistry.get_registry_singleton().instances, - ) - technique_instances._registry_items = dict(self.technique_entries) - technique_instances._metadata_cache = ( - list(self.technique_metadata) if self.technique_metadata is not None else None - ) - - SeedDatasetProvider._registry.clear() - SeedDatasetProvider._registry.update(self.providers) - sys.path[:] = self.sys_path - for name in list(sys.modules): - if name not in self.module_names and _name_owned_by_any(name=name, prefixes=owned_prefixes): - del sys.modules[name] - - def restore_unsupported_provider_side_effects(self) -> None: - """Remove provider registrations because datasets are outside the V1 contract.""" - SeedDatasetProvider._registry.clear() - SeedDatasetProvider._registry.update(self.providers) - - -class PluginLoader: - """Prepare, discover, validate, and transactionally register one plug-in.""" - - def __init__(self, *, spec: PluginSpec) -> None: - """ - Initialize the loader. - - Args: - spec (PluginSpec): The normalized source or wheel declaration. - """ - self._spec = spec - - async def load_async(self) -> None: """ - Activate the configured plug-in. + Run the configured plug-in's initializer before user-configured initializers. Raises: - PluginLoadError: If preparation, discovery, validation, or registration fails. + PluginLoadError: If the source cannot be prepared, the initializer cannot be + resolved, or the initializer raises. Fails closed — a restart is required + after fixing the configuration or artifact. """ - snapshot = _RegistrySnapshot.capture() - prepared: PreparedPlugin | None = None + initializer = await asyncio.to_thread(self._resolve_initializer) try: - prepared = await self._prepare_async() - imported = await import_plugin_async(prepared=prepared) - self._reject_import_time_registration(snapshot=snapshot) - scenarios = discover_scenarios(imported=imported) - techniques = discover_techniques(imported=imported) - self._validate_names(scenarios=scenarios, techniques=techniques) - self._register(scenarios=scenarios, techniques=techniques) - snapshot.restore_unsupported_provider_side_effects() - logger.info( - "Loaded plug-in '%s': %d scenario(s), %d attack technique(s).", - self._spec.name, - len(scenarios), - len(techniques), - ) + await initializer.initialize_async() except Exception as exc: - prefixes = prepared.owned_module_prefixes if prepared else () - snapshot.restore(owned_prefixes=prefixes) - message = f"Failed to load plug-in '{self._spec.name}': {exc} {_REMEDIATION}" - if isinstance(exc, PluginLoadError): - raise type(exc)(message) from exc - raise PluginLoadError(message) from exc - - async def _prepare_async(self) -> PreparedPlugin: - if self._spec.format is PluginFormat.SOURCE: - return await SourcePluginFormat().prepare_async(spec=self._spec) - if self._spec.format is PluginFormat.WHEEL: - prepared = await WheelPluginFormat().prepare_async(spec=self._spec) - self._warn_on_version_drift(prepared=prepared) - return prepared - raise PluginValidationError(f"Unsupported plug-in format: {self._spec.format}") - - @staticmethod - def _reject_import_time_registration(*, snapshot: _RegistrySnapshot) -> None: - scenario_registry = ScenarioRegistry.get_registry_singleton() - technique_instances = cast( - "DefaultInstanceRegistry[AttackTechniqueFactory]", - AttackTechniqueRegistry.get_registry_singleton().instances, - ) - technique_entries = technique_instances._registry_items - if scenario_registry._classes != snapshot.scenario_classes or technique_entries != snapshot.technique_entries: - raise PluginValidationError( - "Plug-in source registered components during import. V1 plug-ins must expose definitions " - "and let the framework register them transactionally." - ) + raise PluginLoadError(f"Plug-in '{self._spec.name}' initializer failed: {exc} {_REMEDIATION}") from exc + logger.info("Activated plug-in '%s' via %s.", self._spec.name, self._spec.initializer) - @staticmethod - def _validate_names( - *, - scenarios: Sequence[ScenarioContribution], - techniques: Sequence[TechniqueContribution], - ) -> None: - if not scenarios and not techniques: - raise PluginRegisteredNothingError("Plug-in contributed no scenarios or attack techniques.") - - scenario_registry = ScenarioRegistry.get_registry_singleton() - builtin_scenarios = set(scenario_registry.get_class_names()) - collisions = sorted(item.registry_name for item in scenarios if item.registry_name in builtin_scenarios) - if collisions: - raise PluginCollisionError(f"Scenario name collision(s): {collisions}.") + def _resolve_initializer(self) -> PyRITInitializer: + """ + Anchor the source root on ``sys.path`` and import the configured initializer. - from pyrit.setup.initializers.techniques import build_technique_factories + Returns: + PyRITInitializer: A fresh instance of the resolved initializer subclass. - builtin_techniques = {factory.name for factory in build_technique_factories()} - registered_techniques = set(AttackTechniqueRegistry.get_registry_singleton().get_factories()) - technique_collisions = sorted( - item.factory.name for item in techniques if item.factory.name in builtin_techniques | registered_techniques - ) - if technique_collisions: - raise PluginCollisionError(f"Attack technique name collision(s): {technique_collisions}.") + Raises: + PluginSourceNotFoundError: If the source path does not exist. + PluginLoadError: If the dotted initializer cannot be imported or is not a + ``PyRITInitializer`` subclass. + """ + spec = self._spec + if not spec.source.exists(): + raise PluginSourceNotFoundError(f"Plug-in '{spec.name}' source path does not exist: {spec.source}") - def _register( - self, - *, - scenarios: Sequence[ScenarioContribution], - techniques: Sequence[TechniqueContribution], - ) -> None: - from pyrit.setup.initializers.techniques import build_technique_factories + root = spec.source if spec.source.is_dir() else spec.source.parent + root_str = str(root) + if root_str not in sys.path: + sys.path.insert(0, root_str) - technique_registry = AttackTechniqueRegistry.get_registry_singleton() - technique_registry.register_from_factories(build_technique_factories(groups=["core"])) - for contribution in techniques: - technique_registry.register_contributed_factory( - factory=contribution.factory, - plugin_name=self._spec.name or "", - scenario_names=contribution.scenario_names, - ) - - scenario_registry = ScenarioRegistry.get_registry_singleton() - for contribution in scenarios: - scenario_registry.register_contributed_scenario( - scenario_class=contribution.scenario_class, - name=contribution.registry_name, + module_path, _, class_name = spec.initializer.rpartition(".") + try: + module = importlib.import_module(module_path) + initializer_cls = getattr(module, class_name) + except Exception as exc: + raise PluginLoadError( + f"Could not import plug-in '{spec.name}' initializer '{spec.initializer}' from {root}: " + f"{exc} {_REMEDIATION}" + ) from exc + + if not (isinstance(initializer_cls, type) and issubclass(initializer_cls, PyRITInitializer)): + raise PluginLoadError( + f"Plug-in '{spec.name}' initializer '{spec.initializer}' is not a PyRITInitializer subclass. " + f"{_REMEDIATION}" ) - try: - scenario_registry.get_class_metadata(contribution.scenario_class) - except Exception as exc: - raise PluginValidationError( - f"Scenario '{contribution.registry_name}' could not build registry metadata: {exc}" - ) from exc - - @staticmethod - def _warn_on_version_drift(*, prepared: PreparedPlugin) -> None: - declared = prepared.declared_pyrit_version - if not declared: - return - import pyrit - - running = getattr(pyrit, "__version__", "") or "" - if _major_minor(declared) == _major_minor(running) and _major_minor(declared) is not None: - return - logger.warning( - "PLUGIN VERSION DRIFT: plug-in '%s' declares pyrit %s but pyrit %s is running. " - "Compatibility is the artifact author's responsibility.", - prepared.spec.name, - declared, - running or "(unknown)", - ) - - -def _major_minor(version: str) -> tuple[int, int] | None: - parts = version.split(".") - if len(parts) < 2: - return None - try: - return int(parts[0]), int(parts[1]) - except ValueError: - return None - - -def _name_owned_by_any(*, name: str, prefixes: tuple[str, ...]) -> bool: - return any(name == prefix or name.startswith(f"{prefix}.") for prefix in prefixes) + return initializer_cls() diff --git a/pyrit/setup/plugin_spec.py b/pyrit/setup/plugin_spec.py index 2a8c21029a..32998864e4 100644 --- a/pyrit/setup/plugin_spec.py +++ b/pyrit/setup/plugin_spec.py @@ -1,63 +1,24 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -"""Static plug-in configuration models shared by setup entry points.""" +"""Static plug-in configuration model shared by setup entry points.""" from __future__ import annotations from dataclasses import dataclass -from enum import Enum from pathlib import Path from typing import Any from pyrit.models import validate_registry_name -class PluginFormat(str, Enum): - """The supported plug-in artifact formats.""" - - SOURCE = "source" - WHEEL = "wheel" - - @dataclass(frozen=True) class PluginSpec: - """A normalized plug-in artifact declaration.""" + """A normalized plug-in declaration: a source root plus a dotted initializer.""" - wheel: Path | None = None - package: str | None = None - name: str | None = None - format: PluginFormat | None = None - source: Path | None = None - - def __post_init__(self) -> None: - """ - Infer legacy direct-construction fields while the loader is refactored. - - Raises: - ValueError: If the artifact fields are missing, conflicting, or inconsistent - with the explicit format. - """ - if (self.source is None) == (self.wheel is None): - raise ValueError("PluginSpec requires exactly one of 'source' or 'wheel'.") - - inferred_format = PluginFormat.SOURCE if self.source is not None else PluginFormat.WHEEL - if self.format is not None and self.format is not inferred_format: - raise ValueError(f"Plug-in format '{self.format.value}' does not match its artifact field.") - object.__setattr__(self, "format", inferred_format) - - if self.name is None: - path = self.source or self.wheel - assert path is not None - inferred_name = self.package.split(".", 1)[0] if self.package else path.stem.replace("-", "_") - object.__setattr__(self, "name", inferred_name) - - @property - def artifact_path(self) -> Path: - """The normalized source or wheel artifact path.""" - path = self.source or self.wheel - assert path is not None - return path + name: str + source: Path + initializer: str @classmethod def from_config(cls, entry: dict[str, Any], *, base_dir: Path | None = None) -> PluginSpec: @@ -66,60 +27,44 @@ def from_config(cls, entry: dict[str, Any], *, base_dir: Path | None = None) -> Args: entry (dict[str, Any]): The YAML plug-in mapping. - base_dir (Path | None): Directory used to resolve relative artifact paths. + base_dir (Path | None): Directory used to resolve a relative ``source`` path. Returns: PluginSpec: The normalized plug-in declaration. Raises: - ValueError: If the mapping is malformed or inconsistent. + ValueError: If the mapping is malformed. """ if not isinstance(entry, dict): raise ValueError(f"Plug-in entry must be a mapping, got {type(entry).__name__}.") - allowed = {"name", "format", "source", "wheel", "package"} + allowed = {"name", "source", "initializer"} unexpected = set(entry) - allowed if unexpected: - raise ValueError(f"Plug-in entry has unexpected key(s): {sorted(unexpected)}.") + raise ValueError(f"Plug-in entry has unexpected key(s): {sorted(unexpected)}. Allowed: {sorted(allowed)}.") name = entry.get("name") if not isinstance(name, str): raise ValueError("Plug-in entry requires a string 'name'.") validate_registry_name(name) - try: - plugin_format = PluginFormat(entry.get("format")) - except (TypeError, ValueError) as exc: - raise ValueError("Plug-in entry 'format' must be 'source' or 'wheel'.") from exc - source = entry.get("source") - wheel = entry.get("wheel") - if (source is None) == (wheel is None): - raise ValueError("Plug-in entry requires exactly one of 'source' or 'wheel'.") - artifact_key = plugin_format.value - artifact = entry.get(artifact_key) - other_key = PluginFormat.WHEEL.value if plugin_format is PluginFormat.SOURCE else PluginFormat.SOURCE.value - if not isinstance(artifact, str) or entry.get(other_key) is not None: - raise ValueError(f"Plug-in format '{plugin_format.value}' requires only the '{artifact_key}' field.") - - package = entry.get("package") - if package is not None and ( - not isinstance(package, str) or not all(part.isidentifier() for part in package.split(".")) - ): - raise ValueError("Plug-in 'package' must be a dotted Python identifier.") - - path = Path(artifact).expanduser() + if not isinstance(source, str) or not source: + raise ValueError("Plug-in entry requires a non-empty string 'source' path.") + + initializer = entry.get("initializer") + if not isinstance(initializer, str) or "." not in initializer: + raise ValueError( + "Plug-in entry requires a dotted 'initializer' path to a PyRITInitializer subclass " + "(e.g. 'my_package.setup.MyInitializer')." + ) + + path = Path(source).expanduser() if not path.is_absolute(): path = (base_dir or Path.cwd()) / path path = path.resolve() - return cls( - name=name, - format=plugin_format, - source=path if plugin_format is PluginFormat.SOURCE else None, - wheel=path if plugin_format is PluginFormat.WHEEL else None, - package=package, - ) + return cls(name=name, source=path, initializer=initializer) def to_config(self) -> dict[str, str]: """ @@ -128,11 +73,4 @@ def to_config(self) -> dict[str, str]: Returns: dict[str, str]: The normalized configuration mapping. """ - config = { - "name": self.name or "", - "format": self.format.value if self.format else "", - self.format.value if self.format else "source": str(self.artifact_path), - } - if self.package: - config["package"] = self.package - return config + return {"name": self.name, "source": str(self.source), "initializer": self.initializer} From d2885063b917aea43f7225ff8acee3e64aed0723 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 14 Jul 2026 13:09:15 -0700 Subject: [PATCH 32/34] docs: retarget plug-in config and docs to the source-initializer model Remove wheel/discovery/scenario-tag content and describe the source-only, initializer-pointing plug-in end to end. - .pyrit_conf_example and doc/getting_started/pyrit_conf.md: {name, source, initializer} schema. - doc/getting_started/plugins.md: initializer contract, the plug-in-vs-initialization-scripts decision, and how the initializer registers scenarios/techniques/datasets/targets (including gray-area content at different abstraction levels). - doc/getting_started/troubleshooting/plugins.md: source-path, import, and non-initializer failures. - Remove the .plugin/ extraction dir and its .gitignore/.gitkeep entries; drop a stray .env_example blank line. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- .env_example | 1 - .gitignore | 5 - .plugin/.gitkeep | 2 - .pyrit_conf_example | 35 ++- doc/getting_started/plugins.md | 215 ++++++------------ doc/getting_started/pyrit_conf.md | 26 +-- doc/getting_started/troubleshooting/README.md | 2 +- .../troubleshooting/plugins.md | 117 +++------- 8 files changed, 136 insertions(+), 267 deletions(-) delete mode 100644 .plugin/.gitkeep diff --git a/.env_example b/.env_example index 0bec2cbb6e..0d70a48cfd 100644 --- a/.env_example +++ b/.env_example @@ -250,7 +250,6 @@ AZURE_OPENAI_VIDEO_MODEL="sora-2" AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" OPENAI_VIDEO_ENDPOINT = ${AZURE_OPENAI_VIDEO_ENDPOINT} - OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" diff --git a/.gitignore b/.gitignore index 6908547442..f92df639f0 100644 --- a/.gitignore +++ b/.gitignore @@ -5,11 +5,6 @@ dbdata/ eval/ default_memory.json.memory -# Plug-in extraction working directory (see DefaultPluginInitializer / PLUGIN_WHEEL). -# The directory is kept; extracted plug-in artifacts are ignored. -.plugin/* -!.plugin/.gitkeep - # Frontend build artifacts copied to backend for packaging pyrit/backend/frontend/ diff --git a/.plugin/.gitkeep b/.plugin/.gitkeep deleted file mode 100644 index 32d9e609e5..0000000000 --- a/.plugin/.gitkeep +++ /dev/null @@ -1,2 +0,0 @@ -# Keeps the .plugin/ directory in version control while its contents stay ignored. -# PyRIT extracts plug-in wheels here at runtime (override with PLUGIN_DIR). diff --git a/.pyrit_conf_example b/.pyrit_conf_example index aa3511ffaf..5a072c4dda 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -59,28 +59,27 @@ initializers: # Plug-ins # -------- -# Private scenarios and attack techniques loaded from source or a pre-built wheel. -# ConfigurationLoader automatically prepends a privileged PluginInitializer ahead of -# the initializers above; do NOT list it under `initializers:`. -# -# V1 accepts one plug-in. Exactly one of `source` or `wheel` is required. -# Relative paths resolve against this configuration file. +# A plug-in points PyRIT at a private ``PyRITInitializer`` that lives outside the public +# tree (for example, a team's internal red-teaming package). ConfigurationLoader runs it +# as a guaranteed-first, privileged initializer -- before the initializers above -- so the +# scenarios, attack techniques, datasets, and targets it registers exist before anything +# else consumes the registries. Do NOT list it under `initializers:`. +# +# The plug-in's initializer owns all registration; PyRIT discovers nothing on its own. +# V1 accepts one plug-in with three fields: +# - name: an operator label (used in logs/errors) +# - source: a directory placed on sys.path so `import ` resolves +# - initializer: dotted 'module.Class' path to a PyRITInitializer subclass +# Relative source paths resolve against this configuration file. # # WARNING: loading a plug-in executes third-party code in the PyRIT process. Whoever can -# write this file can run code on the host. Treat it as sensitive. -# -# Source example: -# plugins: -# - name: operation_foobar -# format: source -# source: /opt/pyrit/operation_foobar.py +# write this file (or the source) can run code on the host. Treat it as sensitive. # -# Wheel example: +# Example: # plugins: -# - name: partner_scenarios -# format: wheel -# wheel: /opt/pyrit/partner_scenarios-1.2.0-py3-none-any.whl -# package: partner_scenarios +# - name: rapid_response +# source: /repos/pyrit-internal +# initializer: pyrit_internal.setup.initializers.RapidResponseInitializer # # Help and troubleshooting: # - doc/getting_started/plugins.md diff --git a/doc/getting_started/plugins.md b/doc/getting_started/plugins.md index 556d3c59d1..d00817468c 100644 --- a/doc/getting_started/plugins.md +++ b/doc/getting_started/plugins.md @@ -1,173 +1,106 @@ # Private Scenarios and Attack Techniques -PyRIT plug-ins make private scenarios and attack techniques behave like built-in -components in the standard backend, GUI, and `pyrit_scan` catalog. +PyRIT plug-ins let an operator activate private scenarios and attack techniques from a +**stock PyRIT installation** — without forking PyRIT or passing per-run CLI flags. -## Plug-In or Python Dependency? +## Plug-in or `--initialization-scripts`? -Use ordinary Python composition when you own the application: +A plug-in is a thin layer over the existing initializer path: -```python -import pyrit - -# Build and run your custom workflow directly. -``` - -That is the simplest option for notebooks, services, and tools that already control -their Python entry point. - -Use a plug-in when operators need private components inside a **stock PyRIT -installation**: - -- keep the shipped `pyrit_scan` command and backend; -- discover private components through normal catalog APIs; -- avoid a private PyRIT fork or wrapper CLI; -- activate components from `.pyrit_conf`; -- deploy an artifact/config instead of application glue. - -## V1 Scope - -V1 plug-ins contribute: - -- concrete `Scenario` subclasses; -- configured `AttackTechniqueFactory` instances. - -They do not contribute custom initializers or lifecycle hooks. PyRIT's privileged -`PluginInitializer` owns loading, discovery, validation, registration, and rollback. - -Only one plug-in is supported at a time. - -## Source Plug-In +- `--initialization-scripts ./my_init.py` runs a custom `PyRITInitializer` from a loose + file. It works for a self-contained script but is per-run and breaks for a packaged + initializer that imports from its own package. +- A `plugins:` entry in `.pyrit_conf` names a private initializer once, runs it + **first** (before other initializers and before catalog warming), and anchors the + package root on `sys.path` so a packaged initializer with intra-package imports loads + correctly. -Source is useful for live operations where scenarios or techniques already exist on -the backend filesystem. +Use a plug-in when a private, packaged initializer (for example a team's internal +red-teaming package) needs to behave like a built-in inside the standard backend, GUI, +and `pyrit_scan` catalog. -Supported shapes: +## What a plug-in is -- one importable `.py` file; -- one Python package directory containing `__init__.py`. - -### Private attack technique - -```python -# /opt/pyrit/operation_foobar.py -from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory - - -OPERATION_FOOBAR = AttackTechniqueFactory( - name="operation_foobar", - attack_class=PromptSendingAttack, - technique_tags=[ - "single_turn", - "scenario:airt.rapid_response", - ], -) -``` - -The `scenario:` tag declares where a directly discovered factory is -available. Construction tooling may replace this bridge with an explicit contribution -manifest in the future. - -Configure it: +A plug-in is a config entry that points at a concrete `PyRITInitializer` reachable by a +dotted path from a source root: ```yaml +# .pyrit_conf plugins: - - name: operation_foobar - format: source - source: /opt/pyrit/operation_foobar.py + - name: rapid_response + source: /repos/pyrit-internal + initializer: pyrit_internal.setup.initializers.RapidResponseInitializer ``` -Then restart the backend and run: +- `name` — an operator label used in logs and errors. +- `source` — a directory placed on `sys.path` so `import ` resolves. A + relative path resolves against the config file. +- `initializer` — a dotted `module.Class` path to a concrete `PyRITInitializer`. -```powershell -pyrit_scan airt.rapid_response ` - --target openai_chat ` - --techniques operation_foobar -``` +`ConfigurationLoader` prepends a privileged initializer that anchors `source`, imports +the initializer, and runs it before any user-configured initializer. Do **not** list it +under `initializers:`. -### Private scenario +## The initializer owns registration -A source scenario follows the normal `Scenario` contract: +PyRIT discovers nothing on its own. The plug-in's initializer registers everything it +wants discoverable, at whatever level of abstraction fits: -- concrete subclass; -- keyword-only constructor; -- no-argument instantiable for catalog metadata; -- runtime dependencies resolved lazily; -- `_build_atomic_attacks_async` implemented. +- **Attack techniques** — + `AttackTechniqueRegistry.get_registry_singleton().register_from_factories([...])`; + selectable via `--techniques`. +- **Scenarios** — + `ScenarioRegistry.get_registry_singleton().register_class(MyScenario, name="airt_internal.violence")`; + runnable via `pyrit_scan airt_internal.violence`. +- **Datasets** — register providers and load them into memory so private seeds stay in + the operator's database and are never published. +- **Default targets** — `set_default_value(...)`. -The default registry name comes from the source-relative module path. For example: +This is why the plug-in fits gray-area content: sometimes only the *dataset* must stay +private, sometimes a *technique*, and sometimes an entire *scenario* — the initializer +decides. -```text -/opt/pyrit/private_scenarios/image/abuse.py -> image.abuse -``` - -A package directory must contain `__init__.py`; a directory of unrelated loose files -is rejected. - -## Wheel Plug-In - -Wheels are appropriate for durable distribution of multi-file private scenarios and -resources. The wheel must already be built and compatible with the running PyRIT. - -```yaml -plugins: - - name: partner_scenarios - format: wheel - wheel: /opt/pyrit/partner_scenarios-1.2.0-py3-none-any.whl - package: partner_scenarios -``` - -PyRIT safely extracts the wheel to `.plugin/`; it does not run `pip`, install into -`.venv`, or resolve dependencies. - -If wheel metadata declares another PyRIT version, loading emits an advisory warning. -PyRIT does not rewrite incompatible APIs. Rebuild the wheel against the running -version when validation fails. - -## Initialization Behavior - -`ConfigurationLoader` converts the config entry into a privileged -`PluginInitializer`. Operators do not add this initializer to `initializers:`. - -Runtime order: - -1. load environment; -2. create CentralMemory; -3. activate the configured plug-in; -4. execute user-configured initializers; -5. serve catalog and scenario requests. - -This ordering ensures scenario metadata and technique catalogs cannot be built from -a partial catalog. - -For direct Python use: +### Example initializer ```python -from pyrit.setup import initialize_from_config_async - -await initialize_from_config_async("/path/to/.pyrit_conf") -``` - -Low-level `initialize_pyrit_async()` does not read `.pyrit_conf` automatically. +from pyrit.executor.attack import PromptSendingAttack +from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry +from pyrit.scenario.core import AttackTechniqueFactory +from pyrit.setup.pyrit_initializer import PyRITInitializer -## Trust Boundary +from my_package.scenarios import Violence -Source and wheel plug-ins execute Python with backend permissions. Treat write access -to the artifact and `.pyrit_conf` as code-execution authority. -Dependencies must already be installed in the backend environment. Plug-ins share one -interpreter and do not receive dependency isolation. +class MyInitializer(PyRITInitializer): + """Register a private technique and a private scenario.""" -## Updating a Plug-In + async def initialize_async(self) -> None: + AttackTechniqueRegistry.get_registry_singleton().register_from_factories( + [AttackTechniqueFactory(name="operation_foobar", attack_class=PromptSendingAttack)] + ) + ScenarioRegistry.get_registry_singleton().register_class(Violence, name="airt_internal.violence") +``` -V1 does not hot reload. +## Usage ```powershell -pyrit_scan --stop-server -pyrit_scan --start-server --config-file /path/to/.pyrit_conf +# A private technique through a public scenario +pyrit_scan airt.rapid_response --target openai_chat --techniques operation_foobar + +# A private scenario +pyrit_scan airt_internal.violence --target openai_chat ``` -For direct Python use, initialize in a fresh process after changing the artifact. +## Behavior and limits + +- The plug-in initializer runs **first**, so lazy catalog/metadata consumers see a + complete registry. +- Loading executes third-party Python with backend permissions; whoever can write the + config or the source can run code on the host. Treat the config as sensitive. +- Dependencies must already be installed in the backend environment. +- V1 is **fail-closed** and supports **one** plug-in. A failed load aborts + initialization — fix the config or source and restart. +- Plug-ins activate only at process/backend startup. Restart after changing the config + or the source; there is no hot reload. See [Plug-In Troubleshooting](./troubleshooting/plugins.md) for common failures. diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 0f40bb4f48..ef04b93397 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -153,27 +153,21 @@ initialization_scripts: ### `plugins` -One private scenario/attack-technique plug-in to activate before the configured -initializers. The plug-in can be a source file/package or a pre-built wheel: +One private plug-in to activate before the configured initializers. A plug-in points at +a private ``PyRITInitializer`` reached by a dotted path from a source root; the +initializer registers whatever private scenarios, attack techniques, datasets, and +targets it wants discoverable: ```yaml plugins: - - name: operation_foobar - format: source - source: /opt/pyrit/operation_foobar.py + - name: rapid_response + source: /repos/pyrit-internal + initializer: pyrit_internal.setup.initializers.RapidResponseInitializer ``` -```yaml -plugins: - - name: partner_scenarios - format: wheel - wheel: /opt/pyrit/partner_scenarios-1.2.0-py3-none-any.whl - package: partner_scenarios -``` - -`ConfigurationLoader` automatically injects a privileged plug-in initializer first; -do not add it to `initializers:`. V1 is fail-closed, supports one plug-in, and requires -a process/backend restart after changes. +`ConfigurationLoader` automatically injects this as a privileged initializer that runs +first; do not add it to `initializers:`. V1 is fail-closed, supports one plug-in, and +requires a process/backend restart after changes. See [Private Scenarios and Attack Techniques](./plugins.md) and [Plug-In Troubleshooting](./troubleshooting/plugins.md). diff --git a/doc/getting_started/troubleshooting/README.md b/doc/getting_started/troubleshooting/README.md index 8eb820b821..c91af61d85 100644 --- a/doc/getting_started/troubleshooting/README.md +++ b/doc/getting_started/troubleshooting/README.md @@ -14,7 +14,7 @@ Common issues and advanced setup guides. - **[Azure SQL Database Setup](./azure_sql_db.md)** — Setting up Azure SQL as your memory backend, including Entra ID authentication and user permissions. -- **[Private Scenario Plug-Ins](./plugins.md)** — Configuration, discovery, dependency, collision, and restart troubleshooting. +- **[Private Scenario Plug-Ins](./plugins.md)** — Configuration, initializer resolution, import, and restart troubleshooting. ## Model Deployment diff --git a/doc/getting_started/troubleshooting/plugins.md b/doc/getting_started/troubleshooting/plugins.md index fe89ce91a5..c6b2f92693 100644 --- a/doc/getting_started/troubleshooting/plugins.md +++ b/doc/getting_started/troubleshooting/plugins.md @@ -1,113 +1,64 @@ # Plug-In Troubleshooting -## The plug-in is missing from `pyrit_scan` +## The plug-in did not activate -Restart the backend after changing `.pyrit_conf` or the artifact: +Restart the backend after changing `.pyrit_conf` or the source: ```powershell pyrit_scan --stop-server pyrit_scan --start-server --config-file /path/to/.pyrit_conf ``` -`pyrit_scan` is a thin client. An already-running backend keeps the plug-ins loaded -from its startup config. +`pyrit_scan` is a thin client. An already-running backend keeps the plug-in loaded from +its startup config. ## Configuration is rejected -V1 accepts one explicit entry: +V1 accepts exactly one entry with three fields: ```yaml plugins: - - name: operation_foobar - format: source - source: /absolute/path/operation_foobar.py -``` - -or: - -```yaml -plugins: - - name: partner_scenarios - format: wheel - wheel: /absolute/path/partner.whl - package: partner_scenarios + - name: rapid_response + source: /repos/pyrit-internal + initializer: pyrit_internal.setup.initializers.RapidResponseInitializer ``` Check that: -- `name` is lowercase snake case; -- `format` is `source` or `wheel`; -- exactly one matching artifact field is present; +- `name` is a valid lowercase snake_case registry name; +- `source` is present (relative paths resolve against the config file); +- `initializer` is a dotted `module.Class` path; - only one plug-in is configured. -Relative paths resolve against the configuration file. - -## Source path is rejected - -A source path must be: - -- one `.py` file with an importable filename; or -- one package directory containing `__init__.py`. - -Loose directories containing unrelated Python files are not supported. - -## Import fails - -Plug-in dependencies are not installed automatically. Install them into the backend -environment and restart. +## Source path does not exist -Also verify: - -- package imports work from the configured source root; -- the configured package name matches the wheel/source package; -- another installed package is not shadowing the plug-in. - -## Scenario is rejected - -The scenario must: - -- inherit from `Scenario`; -- be concrete; -- use a keyword-only constructor; -- construct with no arguments for catalog metadata; -- implement `_build_atomic_attacks_async`. - -Plug-in scenarios are inspected before user-configured target/scorer/technique -initializers run. Keep import and no-argument construction side-effect-light and defer -runtime dependency resolution. - -## Technique is rejected - -Expose factories through either: - -- a module-owned `get_technique_factories()` returning - `AttackTechniqueFactory` instances; or -- module-global `AttackTechniqueFactory` instances. - -Directly discovered factories must identify an applicable scenario, for example: - -```python -technique_tags=["single_turn", "scenario:airt.rapid_response"] -``` +The loader fails closed with a "source path does not exist" error. Point `source` at the +directory that contains your package (the parent of the top-level package) so +`import ` resolves once that directory is on `sys.path`. -Do not register the factory during module import. The framework owns registration and -rollback. +## Initializer cannot be imported -## Name collision +- Confirm `initializer` names a real `module.Class` importable from `source`. +- Install the plug-in's dependencies into the backend Python environment. +- A packaged initializer must import from its own package (for example + `from my_package... import ...`); pointing `source` at the package root makes that + work. -V1 is extend-only. A private scenario or technique cannot replace a built-in or -existing registry name. Rename the contribution and restart. +## Target is not a PyRITInitializer -## Version drift warning +`initializer` must resolve to a concrete subclass of `PyRITInitializer`. A plug-in +contributes an initializer, not loose scenario or technique objects. -The warning is advisory. PyRIT validates the live scenario/factory contract but does -not patch incompatible code. Rebuild or update the artifact against the installed -PyRIT version. +## Nothing shows up in the catalog -## Partial state after failure +PyRIT does not discover components — the initializer must register them. For scanner +discovery, the initializer must call +`ScenarioRegistry.get_registry_singleton().register_class(...)` for scenarios and +`AttackTechniqueRegistry.get_registry_singleton().register_from_factories(...)` for +techniques. Datasets must be registered as providers and loaded into memory. -Plug-in activation is fail-closed and transactional for supported scenario/technique -registries. If initialization fails, fix the reported stage and restart the process. +## Partial state after a failed load -Continuing production work in the same process after failed initialization is not -supported. +V1 is fail-closed. A failed plug-in load aborts initialization; fix the reported stage +and restart the process. Continuing in the same process after a failed initialization is +not supported. From 19bdbf9032b1022d971963a92cc33997a82530e1 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 14 Jul 2026 13:09:29 -0700 Subject: [PATCH 33/34] test(setup): rewrite plug-in tests around the initializer path Replace the discovery/wheel/contributed-registration test surface with tests for the initializer-pointing model. - test_plugin_spec.py: the {name, source, initializer} schema and its validation. - test_plugin_loader.py: PluginInitializer anchors the source root, runs the dotted initializer, and fails closed on a missing source, an unresolvable/non-PyRITInitializer target, or an initializer that raises. - test_plugin_loader_integration.py: a source package whose initializer registers a real private scenario + technique becomes discoverable (gated). - Update test_configuration_loader plug-in cases to the new schema. - Delete test_plugin_discovery/test_plugin_formats/test_plugin_registration. - Drop the contributed-technique cases from test_rapid_response and test_attack_technique_registry, and the removed-exception cases from test_exceptions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- .../setup/test_plugin_loader_integration.py | 677 +++--------------- tests/unit/exceptions/test_exceptions.py | 29 +- .../test_attack_technique_registry.py | 35 - .../unit/scenario/airt/test_rapid_response.py | 19 - tests/unit/setup/test_configuration_loader.py | 29 +- tests/unit/setup/test_plugin_discovery.py | 258 ------- tests/unit/setup/test_plugin_formats.py | 246 ------- tests/unit/setup/test_plugin_loader.py | 633 ++-------------- tests/unit/setup/test_plugin_registration.py | 185 ----- tests/unit/setup/test_plugin_spec.py | 80 +-- 10 files changed, 239 insertions(+), 1952 deletions(-) delete mode 100644 tests/unit/setup/test_plugin_discovery.py delete mode 100644 tests/unit/setup/test_plugin_formats.py delete mode 100644 tests/unit/setup/test_plugin_registration.py diff --git a/tests/integration/setup/test_plugin_loader_integration.py b/tests/integration/setup/test_plugin_loader_integration.py index 1a1299f036..b43794374f 100644 --- a/tests/integration/setup/test_plugin_loader_integration.py +++ b/tests/integration/setup/test_plugin_loader_integration.py @@ -1,601 +1,152 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -""" -Integration test: plug-in scenarios load, register, instantiate, and execute after init. - -This exercises the full public loading path end-to-end: point ``PLUGIN_WHEEL`` at a -built wheel, run ``initialize_pyrit_async``, and assert the wheel's scenarios are -registered in ``ScenarioRegistry``, construct cleanly, and drive through the public -execution pipeline. - -Two cases: +"""Integration test for the initializer-pointing plug-in path. -* A self-contained case builds a small wheel at test time and verifies its scenarios - are discovered. This always runs and guards the mechanism in public CI. -* An injected case loads a wheel supplied out-of-band via environment variables and - verifies all of its scenarios are discovered, instantiate, and that at least one - executes. It is skipped unless those variables are set, so the committed test depends - on no specific external package. Point it at a real scenario wheel (for example in a - downstream/private CI job) to guarantee every scenario that wheel ships is picked up:: - - PLUGIN_TEST_WHEEL=/path/to/plugin.whl - PLUGIN_TEST_PACKAGE=the_plugin_package # enables package enumeration + instantiation - PLUGIN_TEST_SCENARIO_DIRS=/path/to/scenarios # optional; os.pathsep-separated - PLUGIN_TEST_EXEC_SCENARIO=the.registry.name # optional; enables the execution case - ADVERSARIAL_CHAT_ENDPOINT=... # execution target + scorer endpoint - ADVERSARIAL_CHAT_MODEL=... # optional model name +Exercises the full source-plug-in flow against real registries (no mocks): a source +package whose ``PyRITInitializer`` registers a private scenario and a private attack +technique becomes discoverable through the standard registries after the privileged +plug-in phase runs. """ -import inspect -import os -import re -import subprocess import sys import textwrap -import uuid -import zipfile -from collections.abc import Callable, Iterator -from contextlib import contextmanager from pathlib import Path -from typing import Any import pytest -from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider -from pyrit.models import ScenarioResult, ScenarioRunState -from pyrit.prompt_target import OpenAIChatTarget from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry -from pyrit.registry.discovery import discover_in_directory -from pyrit.scenario.core import DatasetAttackConfiguration, Scenario -from pyrit.score import SelfAskTrueFalseScorer, TrueFalseQuestionPaths -from pyrit.setup import ConfigurationLoader, PluginFormat, PluginSpec - -# (module stem, class name, registry name) for the scenarios the self-contained wheel ships. -_MOCK_SCENARIOS = [ - ("alpha", "MockAlphaScenario", "alpha"), - ("beta", "MockBetaScenario", "beta"), - ("gamma", "MockGammaScenario", "gamma"), -] - -# Substring of the ValueError ``Scenario.run_async`` raises when an atomic attack's -# objective did not complete (for example a target refusal). The attack still ran through -# the public pipeline, so this outcome proves execution while a drift regression -- which -# surfaces a different exception type before or during the run -- still fails loudly. -_OBJECTIVES_INCOMPLETE_MARKER = "objectives incomplete" - - -async def _initialize_plugin_async(*, spec: PluginSpec, plugin_dir: Path | None) -> None: - """Initialize PyRIT through the config-owned privileged plug-in path.""" - config = ConfigurationLoader( - memory_db_type="in_memory", - initializers=["technique"], - env_files=[], - silent=True, - plugins=[spec.to_config()], - ) - with _plugin_dir_env(plugin_dir=plugin_dir): - await config.initialize_pyrit_async() +from pyrit.setup import PluginSpec +from pyrit.setup.plugin_loader import PluginInitializer +_PLUGIN_PACKAGE = "airt_internal_plugin" -def _build_scenario_plugin_wheel(dest_dir: Path, *, package: str) -> Path: - """ - Build a wheel whose package ships several scenarios and a ``register()`` bootstrap. - Args: - dest_dir: Directory to write the source tree and .whl into. - package: Top-level package name for the wheel. - - Returns: - Path: The built wheel. - """ - src = dest_dir / f"{package}_src" - pkg = src / package - scenarios_pkg = pkg / "scenarios" - scenarios_pkg.mkdir(parents=True, exist_ok=True) - - (pkg / "__init__.py").write_text("from .bootstrap import register # noqa: F401\n", encoding="utf-8") - (scenarios_pkg / "__init__.py").write_text("", encoding="utf-8") - - for stem, class_name, _registry_name in _MOCK_SCENARIOS: - (scenarios_pkg / f"{stem}.py").write_text( - textwrap.dedent( - f"""\ - from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse - - - class {class_name}(RapidResponse): - \"\"\"Mock plug-in scenario {stem}.\"\"\" - """ - ), - encoding="utf-8", - ) - - imports = "\n".join(f"from .scenarios.{stem} import {class_name}" for stem, class_name, _ in _MOCK_SCENARIOS) - registrations = "\n".join( - f' registry.register_class({class_name}, name="{registry_name}")' - for _stem, class_name, registry_name in _MOCK_SCENARIOS +def _write_source_plugin(root: Path) -> None: + """Write a source package whose initializer registers a scenario and a technique.""" + pkg = root / _PLUGIN_PACKAGE + pkg.mkdir(parents=True, exist_ok=True) + (pkg / "__init__.py").write_text("", encoding="utf-8") + (pkg / "scenario.py").write_text( + textwrap.dedent( + ''' + from functools import cache + + from pyrit.executor.attack import PromptSendingAttack + from pyrit.models import SeedObjective + from pyrit.registry import AttackTechniqueRegistry + from pyrit.scenario.core import ( + AttackTechniqueFactory, + DatasetAttackConfiguration, + Scenario, + ) + from pyrit.score import SubStringScorer + + + @cache + def _build_private_technique(): + return AttackTechniqueRegistry.build_technique_class_from_factories( + class_name="PrivateTechnique", + factories=[ + AttackTechniqueFactory(name="private_prompt_sending", attack_class=PromptSendingAttack) + ], + aggregate_tags={}, + ) + + + class PrivateScenario(Scenario): + """Private scenario for the plug-in integration test.""" + + VERSION = 1 + + def __init__(self, *, scenario_result_id=None): + technique_class = _build_private_technique() + super().__init__( + version=self.VERSION, + technique_class=technique_class, + default_technique=technique_class.ALL, + default_dataset_config=DatasetAttackConfiguration( + seeds=[SeedObjective(value="integration objective")] + ), + objective_scorer=SubStringScorer(substring="integration"), + scenario_result_id=scenario_result_id, + ) + + async def _build_atomic_attacks_async(self, *, context): + return [] + ''' + ), + encoding="utf-8", ) (pkg / "bootstrap.py").write_text( textwrap.dedent( - """\ - from pyrit.registry import ScenarioRegistry - - {imports} - - - def register() -> None: - registry = ScenarioRegistry.get_registry_singleton() - {registrations} - """ - ).format(imports=imports, registrations=registrations), + f''' + from pyrit.executor.attack import PromptSendingAttack + from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry + from pyrit.scenario.core import AttackTechniqueFactory + from pyrit.setup.pyrit_initializer import PyRITInitializer + + from {_PLUGIN_PACKAGE}.scenario import PrivateScenario + + + class PrivateInitializer(PyRITInitializer): + """Register the private scenario and a private attack technique.""" + + async def initialize_async(self) -> None: + AttackTechniqueRegistry.get_registry_singleton().register_from_factories( + [AttackTechniqueFactory(name="operation_foobar", attack_class=PromptSendingAttack)] + ) + ScenarioRegistry.get_registry_singleton().register_class( + PrivateScenario, name="airt_internal.private" + ) + ''' + ), encoding="utf-8", ) - distinfo = src / f"{package}-0.0.1.dist-info" - distinfo.mkdir(parents=True, exist_ok=True) - (distinfo / "METADATA").write_text(f"Metadata-Version: 2.1\nName: {package}\nVersion: 0.0.1\n", encoding="utf-8") - (distinfo / "WHEEL").write_text( - "Wheel-Version: 1.0\nGenerator: test\nRoot-Is-Purelib: true\nTag: py3-none-any\n", encoding="utf-8" - ) - (distinfo / "RECORD").write_text("", encoding="utf-8") - - wheel = dest_dir / f"{package}-0.0.1-py3-none-any.whl" - with zipfile.ZipFile(wheel, "w", zipfile.ZIP_DEFLATED) as archive: - for root, _, files in os.walk(src): - for file_name in files: - file_path = Path(root) / file_name - archive.write(file_path, str(file_path.relative_to(src))) - return wheel - - -def _registered_scenario_class_names() -> set[str]: - """Return the class names of every scenario currently registered in ScenarioRegistry.""" - registry = ScenarioRegistry.get_registry_singleton() - return {registry.get_class(name).__name__ for name in registry.get_class_names()} - - -def all_scenario_class_names(from_dirs: list[Path]) -> set[str]: - """ - Walk source directories and collect the class name of every ``Scenario`` subclass. - - This is the expected-set builder: given the scenario source directories to cover, it - enumerates the scenario classes defined there so the test can assert each is - discovered after initialization. - - Args: - from_dirs: Directories to walk recursively for ``Scenario`` subclasses. - - Returns: - set[str]: The scenario class names found across the directories. - """ - names: set[str] = set() - for directory in from_dirs: - for _stem, _path, cls in discover_in_directory(directory=directory, base_class=Scenario, recursive=True): - names.add(cls.__name__) - return names - - -def _scenario_classes_under_package(package_prefix: str) -> dict[str, type[Scenario]]: - """ - Return registered concrete ``Scenario`` subclasses owned by a package, keyed by registry name. - - Uses the post-load registry (reliable after the plug-in has been imported and - bootstrapped), which sidesteps the standalone-import problems a filesystem walk can - hit for a package that uses relative imports. - - Args: - package_prefix: The plug-in's top-level package name. - - Returns: - dict[str, type[Scenario]]: Registry name -> scenario class for that package. - """ - prefix = f"{package_prefix}." - registry = ScenarioRegistry.get_registry_singleton() - found: dict[str, type[Scenario]] = {} - for name in registry.get_class_names(): - cls = registry.get_class(name) - module = cls.__module__ or "" - if not inspect.isabstract(cls) and (module == package_prefix or module.startswith(prefix)): - found[name] = cls - return found - - -async def _execute_scenario_async( - *, scenario_cls: type[Scenario], endpoint: str, model: str | None -) -> ScenarioResult | None: - """ - Drive one plug-in scenario through the public initialize/run pipeline. - - Constructs an objective target and scorer from the supplied endpoint, initializes the - scenario, and runs it. ``initialize_async`` must succeed and build at least one atomic - attack; that is where dataset-config resolution drift would surface. - - Args: - scenario_cls: The plug-in scenario class to execute. - endpoint: Chat endpoint backing both the objective target and the scorer. - model: Optional model name for the endpoint. - - Returns: - ScenarioResult | None: The completed result, or ``None`` when the run finished with - incomplete objectives (the attack ran end-to-end but the objective was not achieved). - """ - target = OpenAIChatTarget(endpoint=endpoint, model_name=model) - scorer = SelfAskTrueFalseScorer( - chat_target=target, - true_false_question_path=TrueFalseQuestionPaths.TASK_ACHIEVED_REFINED.value, - ) - - try: - scenario = scenario_cls(objective_scorer=scorer, fast_mode=True) # type: ignore[ty:missing-argument, ty:unknown-argument] - except TypeError: - scenario = scenario_cls(objective_scorer=scorer) # type: ignore[ty:missing-argument] - - dataset_config: DatasetAttackConfiguration | None = None - required_datasets = getattr(scenario_cls, "required_datasets", None) - if callable(required_datasets): - names = list(required_datasets()) - if names: - dataset_config = DatasetAttackConfiguration(dataset_names=names, max_dataset_size=1) - - args: dict[str, Any] = {"objective_target": target, "max_concurrency": 1} - if dataset_config is not None: - args["dataset_config"] = dataset_config - scenario.set_params_from_args(args=args) - await scenario.initialize_async() - assert scenario.atomic_attack_count >= 1, "Scenario built no atomic attacks during initialization." - - try: - return await scenario.run_async() - except ValueError as exc: - if _OBJECTIVES_INCOMPLETE_MARKER in str(exc): - return None - raise - - -@contextmanager -def _plugin_dir_env(*, plugin_dir: Path | None) -> Iterator[None]: - """Set PLUGIN_DIR (extraction dir) for the duration of a load, then restore it.""" - values: dict[str, str] = {} - if plugin_dir is not None: - values["PLUGIN_DIR"] = str(plugin_dir) - - saved: dict[str, str | None] = {key: os.environ.get(key) for key in values} - os.environ.update(values) - try: - yield - finally: - for key, previous in saved.items(): - if previous is None: - os.environ.pop(key, None) - else: - os.environ[key] = previous - - -@pytest.fixture -def plugin_dir(tmp_path: Path) -> Path: - """The directory the loader extracts wheels into for a test.""" - return tmp_path / ".plugin" - - -@pytest.fixture -def build_mock_wheel(tmp_path: Path) -> Callable[[str], Path]: - """A builder for the self-contained scenario wheel, parameterized by package name.""" - - def build(package: str) -> Path: - return _build_scenario_plugin_wheel(tmp_path, package=package) - - return build - - -@pytest.fixture -def all_scenarios() -> Callable[[list[Path]], set[str]]: - """The expected-set builder: scenario class names defined under the given source dirs.""" - return all_scenario_class_names - - -@pytest.fixture -def plugin_sandbox() -> Iterator[None]: - """Snapshot and restore global import + registry state so a load does not leak.""" - sys_path_snapshot = list(sys.path) - modules_snapshot = set(sys.modules) - provider_snapshot = dict(SeedDatasetProvider._registry) +@pytest.fixture(autouse=True) +def _isolate_registries_and_imports(): + original_path = list(sys.path) + original_modules = set(sys.modules) + AttackTechniqueRegistry.reset_registry_singleton() + ScenarioRegistry.reset_registry_singleton() yield - - sys.path[:] = sys_path_snapshot - for name in set(sys.modules) - modules_snapshot: - del sys.modules[name] - SeedDatasetProvider._registry.clear() - SeedDatasetProvider._registry.update(provider_snapshot) AttackTechniqueRegistry.reset_registry_singleton() ScenarioRegistry.reset_registry_singleton() + sys.path[:] = original_path + for name in list(sys.modules): + if name not in original_modules and name.split(".", 1)[0] == _PLUGIN_PACKAGE: + del sys.modules[name] @pytest.mark.run_only_if_all_tests -async def test_built_wheel_scenarios_are_discovered( - plugin_sandbox: None, - build_mock_wheel: Callable[[str], Path], - plugin_dir: Path, -) -> None: - """A built wheel's scenarios are all registered in ScenarioRegistry after init.""" - package = f"mock_scenario_plugin_{uuid.uuid4().hex[:8]}" - wheel = build_mock_wheel(package) - - registry = ScenarioRegistry.get_registry_singleton() - before = set(registry.get_class_names()) - - await _initialize_plugin_async( - spec=PluginSpec( - name=package, - format=PluginFormat.WHEEL, - wheel=wheel, - package=package, - ), - plugin_dir=plugin_dir, +async def test_source_plugin_initializer_registers_real_components(tmp_path: Path) -> None: + _write_source_plugin(tmp_path) + spec = PluginSpec( + name="airt_internal", + source=tmp_path, + initializer=f"{_PLUGIN_PACKAGE}.bootstrap.PrivateInitializer", ) - registry = ScenarioRegistry.get_registry_singleton() - after = set(registry.get_class_names()) + await PluginInitializer(plugins=[spec]).initialize_async() - expected_registry_names = {registry_name for _stem, _cls, registry_name in _MOCK_SCENARIOS} - # The plug-in adds exactly its scenarios and removes none of the built-ins. - assert after - before == expected_registry_names - assert before <= after - - # And the scenario classes themselves resolve to the plug-in's classes. - expected_class_names = {class_name for _stem, class_name, _ in _MOCK_SCENARIOS} - assert expected_class_names <= _registered_scenario_class_names() - - -@pytest.mark.run_only_if_all_tests -async def test_injected_wheel_scenarios_are_discovered( - plugin_sandbox: None, - all_scenarios: Callable[[list[Path]], set[str]], -) -> None: - """Every scenario in an out-of-band wheel is discovered after initialization. - - Skipped unless ``PLUGIN_TEST_WHEEL`` is set, so the committed test names no external - package. Provide ``PLUGIN_TEST_SCENARIO_DIRS`` (preferred) or ``PLUGIN_TEST_PACKAGE`` - to tell the test which scenarios to expect. - """ - wheel_env = os.getenv("PLUGIN_TEST_WHEEL") - if not wheel_env: - pytest.skip("PLUGIN_TEST_WHEEL is not set; skipping injected-wheel scenario discovery test.") - - wheel = Path(wheel_env).expanduser() - assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" - - scenario_dirs_env = os.getenv("PLUGIN_TEST_SCENARIO_DIRS") - package = os.getenv("PLUGIN_TEST_PACKAGE") - if not scenario_dirs_env and not package: - pytest.skip("Set PLUGIN_TEST_SCENARIO_DIRS or PLUGIN_TEST_PACKAGE to define the expected scenarios.") - - await _initialize_plugin_async( - spec=PluginSpec( - name=package or "injected_plugin", - format=PluginFormat.WHEEL, - wheel=wheel, - package=package, - ), - plugin_dir=None, - ) - - found = _registered_scenario_class_names() - - if scenario_dirs_env: - dirs = [Path(p) for p in scenario_dirs_env.split(os.pathsep) if p] - expected = all_scenarios(dirs) - else: - assert package is not None - expected = {cls.__name__ for cls in _scenario_classes_under_package(package).values()} - - assert expected, "No plug-in scenarios were found to verify; check the injected wheel/scenario source." - missing = expected - found - assert not missing, f"Plug-in scenarios not discovered by ScenarioRegistry: {sorted(missing)}" - - -@pytest.mark.run_only_if_all_tests -async def test_injected_wheel_scenarios_instantiate(plugin_sandbox: None) -> None: - """Every scenario in an out-of-band wheel constructs cleanly after initialization. - - Skipped unless ``PLUGIN_TEST_WHEEL`` and ``PLUGIN_TEST_PACKAGE`` are set. - """ - wheel_env = os.getenv("PLUGIN_TEST_WHEEL") - package = os.getenv("PLUGIN_TEST_PACKAGE") - if not wheel_env or not package: - pytest.skip("Set PLUGIN_TEST_WHEEL and PLUGIN_TEST_PACKAGE to run the instantiation test.") - - wheel = Path(wheel_env).expanduser() - assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" - - await _initialize_plugin_async( - spec=PluginSpec( - name=package, - format=PluginFormat.WHEEL, - wheel=wheel, - package=package, - ), - plugin_dir=None, - ) - - classes = _scenario_classes_under_package(package) - assert classes, f"No plug-in scenarios registered under package {package!r}." - - failures: dict[str, str] = {} - for name, cls in sorted(classes.items()): - try: - instance = cls() # type: ignore[ty:missing-argument] - assert isinstance(instance, Scenario) - except Exception as exc: # noqa: BLE001 - failures[name] = f"{type(exc).__name__}: {exc}" - - assert not failures, f"Plug-in scenarios failed to instantiate: {failures}" - - -@pytest.mark.run_only_if_all_tests -async def test_injected_wheel_scenario_executes(plugin_sandbox: None) -> None: - """One named scenario from an out-of-band wheel runs through the public pipeline. - - Skipped unless ``PLUGIN_TEST_WHEEL``, ``PLUGIN_TEST_PACKAGE``, and - ``PLUGIN_TEST_EXEC_SCENARIO`` are set, and an ``ADVERSARIAL_CHAT_ENDPOINT`` is - configured (it may come from the loaded ``.env``, so it is checked after init). - """ - wheel_env = os.getenv("PLUGIN_TEST_WHEEL") - package = os.getenv("PLUGIN_TEST_PACKAGE") - exec_scenario = os.getenv("PLUGIN_TEST_EXEC_SCENARIO") - if not wheel_env or not package or not exec_scenario: - pytest.skip("Set PLUGIN_TEST_WHEEL, PLUGIN_TEST_PACKAGE, and PLUGIN_TEST_EXEC_SCENARIO to run the exec test.") - - wheel = Path(wheel_env).expanduser() - assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" - - await _initialize_plugin_async( - spec=PluginSpec( - name=package, - format=PluginFormat.WHEEL, - wheel=wheel, - package=package, - ), - plugin_dir=None, - ) - - endpoint = os.getenv("ADVERSARIAL_CHAT_ENDPOINT") - if not endpoint: - pytest.skip("ADVERSARIAL_CHAT_ENDPOINT is not configured; skipping plug-in scenario execution.") - - registry = ScenarioRegistry.get_registry_singleton() - assert exec_scenario in registry.get_class_names(), ( - f"PLUGIN_TEST_EXEC_SCENARIO {exec_scenario!r} is not registered; " - f"available: {sorted(registry.get_class_names())}" - ) - scenario_cls = registry.get_class(exec_scenario) - - result = await _execute_scenario_async( - scenario_cls=scenario_cls, - endpoint=endpoint, - model=os.getenv("ADVERSARIAL_CHAT_MODEL"), - ) - - if result is not None: - assert result.scenario_run_state == ScenarioRunState.COMPLETED - assert result.attack_results, "Completed scenario run produced no attack results." - - -# --------------------------------------------------------------------------- -# Scanner integration -# --------------------------------------------------------------------------- - - -def _write_scanner_config(*, tmp_path: Path, source: Path, plugin_name: str) -> Path: - config = tmp_path / f"{plugin_name}.pyrit_conf" - config.write_text( - textwrap.dedent( - f"""\ - memory_db_type: in_memory - initializers: - - target - - scorer - - technique - plugins: - - name: {plugin_name} - format: source - source: {source.as_posix()} - """ - ), - encoding="utf-8", - ) - return config - - -def _run_scanner_with_fresh_backend(*, config: Path) -> subprocess.CompletedProcess[str]: - base = [sys.executable, "-m", "pyrit.cli.pyrit_scan", "--config-file", str(config)] - subprocess.run([*base, "--stop-server"], capture_output=True, text=True, check=False, timeout=30) - try: - return subprocess.run( - [*base, "--start-server", "--list-scenarios"], - capture_output=True, - text=True, - check=True, - timeout=120, - ) - finally: - subprocess.run([*base, "--stop-server"], capture_output=True, text=True, check=False, timeout=30) + scenario_registry = ScenarioRegistry.get_registry_singleton() + assert "airt_internal.private" in scenario_registry.get_class_names() + assert scenario_registry.get_class("airt_internal.private").__name__ == "PrivateScenario" + assert "operation_foobar" in AttackTechniqueRegistry.get_registry_singleton().get_factories() @pytest.mark.run_only_if_all_tests -def test_private_scenario_registers_from_scanner(tmp_path: Path) -> None: - """A source Scenario appears through the stock scanner catalog.""" - source = tmp_path / "private_scanner_scenario.py" - source.write_text( - """from pyrit.models import SeedObjective -from pyrit.scenario import DatasetAttackConfiguration, Scenario, ScenarioTechnique -from pyrit.score import SubStringScorer +async def test_source_plugin_missing_initializer_fails_closed(tmp_path: Path) -> None: + from pyrit.exceptions import PluginLoadError -class PrivateTechnique(ScenarioTechnique): - ALL = ("all", {"all"}) - DIRECT = ("direct", {"direct"}) - -class PrivateScannerScenario(Scenario): - VERSION = 1 - - def __init__(self, *, scenario_result_id=None): - super().__init__( - version=self.VERSION, - technique_class=PrivateTechnique, - default_technique=PrivateTechnique.ALL, - default_dataset_config=DatasetAttackConfiguration( - seeds=[SeedObjective(value="integration objective")] - ), - objective_scorer=SubStringScorer(substring="integration"), - scenario_result_id=scenario_result_id, - ) - - async def _build_atomic_attacks_async(self, *, context): - return [] -""", - encoding="utf-8", - ) - config = _write_scanner_config( - tmp_path=tmp_path, - source=source, - plugin_name="private_scanner_scenario", - ) - - proc = _run_scanner_with_fresh_backend(config=config) - - assert "private_scanner_scenario" in proc.stdout - assert "PrivateScannerScenario" in proc.stdout - - -@pytest.mark.run_only_if_all_tests -def test_private_attack_technique_registers_from_scanner(tmp_path: Path) -> None: - """A source technique applicable to RapidResponse appears in scanner metadata.""" - source = tmp_path / "private_scanner_technique.py" - source.write_text( - """from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory - -PRIVATE = AttackTechniqueFactory( - name="private_scanner_technique", - attack_class=PromptSendingAttack, - technique_tags=["single_turn", "scenario:airt.rapid_response"], -) -""", - encoding="utf-8", - ) - config = _write_scanner_config( - tmp_path=tmp_path, - source=source, - plugin_name="private_scanner_technique", + _write_source_plugin(tmp_path) + spec = PluginSpec( + name="airt_internal", + source=tmp_path, + initializer=f"{_PLUGIN_PACKAGE}.bootstrap.DoesNotExist", ) - proc = _run_scanner_with_fresh_backend(config=config) - - match = re.search( - r"\n airt\.rapid_response\n(?P.*?)(?=\n [a-z][a-z0-9_.]*\n|\n={80})", - proc.stdout, - flags=re.DOTALL, - ) - assert match is not None, proc.stdout - assert "private_scanner_technique" in match.group("body") + with pytest.raises(PluginLoadError): + await PluginInitializer(plugins=[spec]).initialize_async() diff --git a/tests/unit/exceptions/test_exceptions.py b/tests/unit/exceptions/test_exceptions.py index e6e35542ec..374781517b 100644 --- a/tests/unit/exceptions/test_exceptions.py +++ b/tests/unit/exceptions/test_exceptions.py @@ -15,14 +15,8 @@ EmptyResponseException, InvalidJsonException, MissingPromptPlaceholderException, - PluginCollisionError, - PluginDiscoveryError, - PluginImportError, PluginLoadError, - PluginRegisteredNothingError, PluginSourceNotFoundError, - PluginValidationError, - PluginWheelNotFoundError, PyritException, RateLimitException, handle_bad_request_exception, @@ -74,26 +68,9 @@ def test_invalid_json_exception_initialization(): assert str(ex) == "Status Code: 500, Message: Invalid JSON Response" -@pytest.mark.parametrize( - "error_type", - [ - PluginSourceNotFoundError, - PluginWheelNotFoundError, - PluginImportError, - PluginDiscoveryError, - PluginValidationError, - PluginCollisionError, - PluginRegisteredNothingError, - ], -) -def test_plugin_specific_errors_inherit_plugin_load_error(error_type: type[PluginLoadError]) -> None: - """Every specific plug-in failure must be catchable through ``PluginLoadError``.""" - assert issubclass(error_type, PluginLoadError) - - -def test_plugin_registered_nothing_error_describes_supported_components() -> None: - """The empty-contribution error must describe the V1 scenario/technique scope.""" - assert "scenarios or attack techniques" in (PluginRegisteredNothingError.__doc__ or "") +def test_plugin_source_not_found_error_inherits_plugin_load_error() -> None: + """The specific plug-in failure must be catchable through ``PluginLoadError``.""" + assert issubclass(PluginSourceNotFoundError, PluginLoadError) def test_bad_request_exception_process_exception(caplog): diff --git a/tests/unit/registry/test_attack_technique_registry.py b/tests/unit/registry/test_attack_technique_registry.py index 243155fa26..d0a109d365 100644 --- a/tests/unit/registry/test_attack_technique_registry.py +++ b/tests/unit/registry/test_attack_technique_registry.py @@ -13,7 +13,6 @@ from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry -from pyrit.registry.tag_query import TagQuery from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory, ScorerOverridePolicy from pyrit.setup.initializers.techniques import build_technique_factories @@ -121,40 +120,6 @@ def test_register_multiple_techniques(self): assert len(self.registry.instances) == 2 assert self.registry.instances.get_names() == ["stub_20", "stub_5"] - def test_get_factories_for_scenario_unions_base_query_and_applicability(self): - core = AttackTechniqueFactory( - name="core_technique", - attack_class=_StubAttack, - technique_tags=["core", "single_turn"], - ) - private = AttackTechniqueFactory( - name="private_technique", - attack_class=_StubAttack, - technique_tags=["single_turn"], - ) - unrelated = AttackTechniqueFactory( - name="unrelated", - attack_class=_StubAttack, - technique_tags=["single_turn"], - ) - self.registry.register_technique(name=core.name, factory=core, tags=core.technique_tags) - self.registry.register_contributed_factory( - factory=private, - plugin_name="operation", - scenario_names=frozenset({"airt.rapid_response"}), - ) - self.registry.register_contributed_factory( - factory=unrelated, - plugin_name="operation", - scenario_names=frozenset({"airt.cyber"}), - ) - - factories = self.registry.get_factories_for_scenario( - scenario_name="airt.rapid_response", - base_query=TagQuery.all("core"), - ) - assert list(factories) == ["core_technique", "private_technique"] - class TestAttackTechniqueRegistryMetadata: """Tests for metadata / list_metadata on the registry.""" diff --git a/tests/unit/scenario/airt/test_rapid_response.py b/tests/unit/scenario/airt/test_rapid_response.py index 838099b5d5..4a04ef3bcf 100644 --- a/tests/unit/scenario/airt/test_rapid_response.py +++ b/tests/unit/scenario/airt/test_rapid_response.py @@ -165,25 +165,6 @@ def test_get_technique_class(self, mock_objective_scorer): ): assert RapidResponse()._technique_class is strat - def test_technique_class_includes_applicable_contributed_technique(self): - from pyrit.scenario.scenarios.airt.rapid_response import _build_rapid_response_technique - - factory = AttackTechniqueFactory( - name="operation_foobar", - attack_class=PromptSendingAttack, - technique_tags=["single_turn"], - ) - AttackTechniqueRegistry.get_registry_singleton().register_contributed_factory( - factory=factory, - plugin_name="operation", - scenario_names=frozenset({"airt.rapid_response"}), - ) - _build_rapid_response_technique.cache_clear() - - technique_class = _build_rapid_response_technique() - - assert technique_class("operation_foobar").value == "operation_foobar" - def test_get_default_technique_returns_default(self, mock_objective_scorer): strat = _technique_class() with patch( diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 91723fe174..6c97a2218d 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -726,22 +726,23 @@ def test_explicit_source_mapping(self, tmp_path: pathlib.Path): plugins=[ { "name": "operation_foobar", - "format": "source", - "source": str(tmp_path / "operation_foobar.py"), + "source": str(tmp_path / "operation_foobar"), + "initializer": "operation_foobar.setup.OperationInitializer", } ] ) spec = config.resolve_plugins()[0] assert spec.name == "operation_foobar" - assert spec.source == (tmp_path / "operation_foobar.py").resolve() + assert spec.source == (tmp_path / "operation_foobar").resolve() + assert spec.initializer == "operation_foobar.setup.OperationInitializer" def test_multiple_plugins_rejected(self, tmp_path: pathlib.Path): with pytest.raises(ValueError, match="one plug-in"): ConfigurationLoader( plugins=[ - {"name": "first", "format": "source", "source": str(tmp_path / "first.py")}, - {"name": "second", "format": "source", "source": str(tmp_path / "second.py")}, + {"name": "first", "source": str(tmp_path / "first"), "initializer": "first.Init"}, + {"name": "second", "source": str(tmp_path / "second"), "initializer": "second.Init"}, ] ) @@ -757,14 +758,13 @@ def test_from_dict_with_plugins(self, tmp_path: pathlib.Path): "plugins": [ { "name": "partner", - "format": "wheel", - "wheel": str(tmp_path / "plugin.whl"), - "package": "partner.plugin", + "source": str(tmp_path / "partner"), + "initializer": "partner.setup.PartnerInitializer", } ] } ) - assert config.resolve_plugins()[0].package == "partner.plugin" + assert config.resolve_plugins()[0].initializer == "partner.setup.PartnerInitializer" def test_yaml_plugin_path_resolves_relative_to_config(self, tmp_path: pathlib.Path): config_dir = tmp_path / "config" @@ -773,15 +773,15 @@ def test_yaml_plugin_path_resolves_relative_to_config(self, tmp_path: pathlib.Pa config_path.write_text( """plugins: - name: operation_foobar - format: source - source: ../plugins/operation_foobar.py + source: ../plugins/operation_foobar + initializer: operation_foobar.setup.OperationInitializer """, encoding="utf-8", ) config = ConfigurationLoader.from_yaml_file(config_path) - assert config.resolve_plugins()[0].source == (tmp_path / "plugins" / "operation_foobar.py").resolve() + assert config.resolve_plugins()[0].source == (tmp_path / "plugins" / "operation_foobar").resolve() async def test_initialize_prepends_privileged_plugin_initializer(self, tmp_path: pathlib.Path): from pyrit.setup.plugin_loader import PluginInitializer @@ -792,9 +792,8 @@ async def test_initialize_prepends_privileged_plugin_initializer(self, tmp_path: plugins=[ { "name": "partner", - "format": "wheel", - "wheel": str(tmp_path / "partner.whl"), - "package": "partner.plugin", + "source": str(tmp_path / "partner"), + "initializer": "partner.setup.PartnerInitializer", } ], ) diff --git a/tests/unit/setup/test_plugin_discovery.py b/tests/unit/setup/test_plugin_discovery.py deleted file mode 100644 index 39c08cd6ea..0000000000 --- a/tests/unit/setup/test_plugin_discovery.py +++ /dev/null @@ -1,258 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import sys -from collections.abc import Iterator -from pathlib import Path - -import pytest - -from pyrit.exceptions import PluginDiscoveryError -from pyrit.setup import PluginFormat, PluginSpec -from pyrit.setup.plugin_discovery import ( - discover_scenarios, - discover_techniques, - import_plugin_async, -) -from pyrit.setup.plugin_formats import SourcePluginFormat - - -@pytest.fixture -def restore_import_state() -> Iterator[None]: - path_snapshot = list(sys.path) - module_snapshot = set(sys.modules) - yield - sys.path[:] = path_snapshot - for name in set(sys.modules) - module_snapshot: - del sys.modules[name] - - -async def _prepare_and_import_source(source: Path, *, name: str): - spec = PluginSpec(name=name, format=PluginFormat.SOURCE, source=source) - prepared = await SourcePluginFormat().prepare_async(spec=spec) - return await import_plugin_async(prepared=prepared) - - -@pytest.mark.usefixtures("restore_import_state") -async def test_discover_scenario_from_single_file(tmp_path: Path) -> None: - source = tmp_path / "private_scenario.py" - source.write_text( - """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse - -class PrivateScenario(RapidResponse): - pass -""", - encoding="utf-8", - ) - - imported = await _prepare_and_import_source(source, name="private") - contributions = discover_scenarios(imported=imported) - - assert len(contributions) == 1 - assert contributions[0].scenario_class.__name__ == "PrivateScenario" - assert contributions[0].registry_name == "private_scenario" - - -@pytest.mark.usefixtures("restore_import_state") -async def test_discover_scenario_ignores_imported_foreign_class(tmp_path: Path) -> None: - source = tmp_path / "foreign_only.py" - source.write_text( - "from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse\n", - encoding="utf-8", - ) - - imported = await _prepare_and_import_source(source, name="foreign_only") - - assert discover_scenarios(imported=imported) == [] - - -@pytest.mark.usefixtures("restore_import_state") -async def test_discover_scenario_rejects_multiple_classes_in_one_module(tmp_path: Path) -> None: - source = tmp_path / "ambiguous.py" - source.write_text( - """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse - -class FirstScenario(RapidResponse): - pass - -class SecondScenario(RapidResponse): - pass -""", - encoding="utf-8", - ) - - imported = await _prepare_and_import_source(source, name="ambiguous") - - with pytest.raises(PluginDiscoveryError, match="multiple"): - discover_scenarios(imported=imported) - - -@pytest.mark.usefixtures("restore_import_state") -async def test_discover_scenario_uses_source_relative_dotted_name(tmp_path: Path) -> None: - package = tmp_path / "operation_plugin" - scenarios = package / "scenarios" / "image" - scenarios.mkdir(parents=True) - for init_file in (package / "__init__.py", package / "scenarios" / "__init__.py", scenarios / "__init__.py"): - init_file.write_text("", encoding="utf-8") - (scenarios / "abuse.py").write_text( - """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse - -class AbuseScenario(RapidResponse): - pass -""", - encoding="utf-8", - ) - - imported = await _prepare_and_import_source(package, name="operation") - contributions = discover_scenarios(imported=imported) - - assert [item.registry_name for item in contributions] == ["image.abuse"] - - -@pytest.mark.usefixtures("restore_import_state") -async def test_discover_technique_from_conventional_factory_function(tmp_path: Path) -> None: - source = tmp_path / "private_technique.py" - source.write_text( - """from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory - -def get_technique_factories(): - return [ - AttackTechniqueFactory( - name="operation_foobar", - attack_class=PromptSendingAttack, - technique_tags=["single_turn", "scenario:airt.rapid_response"], - ) - ] -""", - encoding="utf-8", - ) - - imported = await _prepare_and_import_source(source, name="private_technique") - contributions = discover_techniques(imported=imported) - - assert len(contributions) == 1 - assert contributions[0].factory.name == "operation_foobar" - assert contributions[0].scenario_names == frozenset({"airt.rapid_response"}) - - -@pytest.mark.usefixtures("restore_import_state") -async def test_discover_technique_from_explicit_contribution(tmp_path: Path) -> None: - source = tmp_path / "explicit_technique.py" - source.write_text( - """from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory -from pyrit.setup.plugin_discovery import TechniqueContribution - -OPERATION = TechniqueContribution( - factory=AttackTechniqueFactory( - name="operation_explicit", - attack_class=PromptSendingAttack, - technique_tags=["single_turn"], - ), - scenario_names=frozenset({"airt.rapid_response"}), -) -""", - encoding="utf-8", - ) - - imported = await _prepare_and_import_source(source, name="explicit_technique") - contributions = discover_techniques(imported=imported) - - assert [item.factory.name for item in contributions] == ["operation_explicit"] - - -@pytest.mark.usefixtures("restore_import_state") -async def test_discover_technique_from_module_global_factory(tmp_path: Path) -> None: - source = tmp_path / "global_technique.py" - source.write_text( - """from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory - -OPERATION = AttackTechniqueFactory( - name="operation_global", - attack_class=PromptSendingAttack, - technique_tags=["single_turn", "scenario:airt.rapid_response"], -) -""", - encoding="utf-8", - ) - - imported = await _prepare_and_import_source(source, name="global_technique") - contributions = discover_techniques(imported=imported) - - assert [item.factory.name for item in contributions] == ["operation_global"] - - -@pytest.mark.usefixtures("restore_import_state") -async def test_discover_technique_does_not_call_arbitrary_helpers(tmp_path: Path) -> None: - source = tmp_path / "safe_discovery.py" - marker = tmp_path / "helper_called.txt" - source.write_text( - f"""from pathlib import Path -from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory - -def build_something(): - Path({str(marker)!r}).write_text("called") - return AttackTechniqueFactory(name="hidden", attack_class=PromptSendingAttack) - -OPERATION = AttackTechniqueFactory( - name="visible", - attack_class=PromptSendingAttack, - technique_tags=["scenario:airt.rapid_response"], -) -""", - encoding="utf-8", - ) - - imported = await _prepare_and_import_source(source, name="safe_discovery") - contributions = discover_techniques(imported=imported) - - assert [item.factory.name for item in contributions] == ["visible"] - assert not marker.exists() - - -@pytest.mark.usefixtures("restore_import_state") -async def test_discover_technique_requires_scenario_applicability(tmp_path: Path) -> None: - source = tmp_path / "missing_applicability.py" - source.write_text( - """from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory - -OPERATION = AttackTechniqueFactory(name="missing_scope", attack_class=PromptSendingAttack) -""", - encoding="utf-8", - ) - - imported = await _prepare_and_import_source(source, name="missing_applicability") - - with pytest.raises(PluginDiscoveryError, match="applicable scenario"): - discover_techniques(imported=imported) - - -@pytest.mark.usefixtures("restore_import_state") -async def test_discover_technique_rejects_duplicate_names(tmp_path: Path) -> None: - source = tmp_path / "duplicate_techniques.py" - source.write_text( - """from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory - -FIRST = AttackTechniqueFactory( - name="duplicate", - attack_class=PromptSendingAttack, - technique_tags=["scenario:airt.rapid_response"], -) -SECOND = AttackTechniqueFactory( - name="duplicate", - attack_class=PromptSendingAttack, - technique_tags=["scenario:airt.rapid_response"], -) -""", - encoding="utf-8", - ) - - imported = await _prepare_and_import_source(source, name="duplicate_techniques") - - with pytest.raises(PluginDiscoveryError, match="duplicate"): - discover_techniques(imported=imported) diff --git a/tests/unit/setup/test_plugin_formats.py b/tests/unit/setup/test_plugin_formats.py deleted file mode 100644 index 4747402e5e..0000000000 --- a/tests/unit/setup/test_plugin_formats.py +++ /dev/null @@ -1,246 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import zipfile -from pathlib import Path - -import pytest - -from pyrit.exceptions import PluginSourceNotFoundError, PluginWheelNotFoundError -from pyrit.setup import PluginFormat, PluginSpec -from pyrit.setup.plugin_formats import SourcePluginFormat, WheelPluginFormat - - -def _build_wheel(tmp_path: Path, *, package: str = "sample_plugin") -> Path: - wheel = tmp_path / f"{package}-1.0.0-py3-none-any.whl" - dist_info = f"{package}-1.0.0.dist-info" - with zipfile.ZipFile(wheel, "w") as archive: - archive.writestr(f"{package}/__init__.py", "") - archive.writestr(f"{dist_info}/top_level.txt", package) - archive.writestr( - f"{dist_info}/METADATA", - "\n".join( - [ - "Metadata-Version: 2.1", - f"Name: {package}", - "Version: 1.0.0", - "Requires-Dist: pyrit==0.14.0", - ] - ), - ) - return wheel - - -async def test_wheel_prepare_returns_common_artifact(tmp_path: Path) -> None: - wheel = _build_wheel(tmp_path) - spec = PluginSpec( - name="sample", - format=PluginFormat.WHEEL, - wheel=wheel, - package="sample_plugin", - ) - - prepared = await WheelPluginFormat(base_dir=tmp_path / ".plugin").prepare_async(spec=spec) - - assert prepared.spec is spec - assert prepared.import_root.is_dir() - assert prepared.entry_modules == ("sample_plugin",) - assert prepared.owned_module_prefixes == ("sample_plugin",) - assert prepared.artifact_fingerprint - assert prepared.declared_pyrit_version == "0.14.0" - assert (prepared.import_root / "sample_plugin" / "__init__.py").is_file() - - -async def test_wheel_prepare_resolves_package_from_metadata(tmp_path: Path) -> None: - wheel = _build_wheel(tmp_path, package="metadata_plugin") - spec = PluginSpec(name="metadata", format=PluginFormat.WHEEL, wheel=wheel) - - prepared = await WheelPluginFormat(base_dir=tmp_path / ".plugin").prepare_async(spec=spec) - - assert prepared.entry_modules == ("metadata_plugin",) - - -async def test_wheel_prepare_rejects_missing_wheel(tmp_path: Path) -> None: - spec = PluginSpec( - name="missing", - format=PluginFormat.WHEEL, - wheel=tmp_path / "missing.whl", - ) - - with pytest.raises(PluginWheelNotFoundError, match="existing"): - await WheelPluginFormat(base_dir=tmp_path / ".plugin").prepare_async(spec=spec) - - -async def test_wheel_prepare_rejects_source_spec(tmp_path: Path) -> None: - spec = PluginSpec( - name="source", - format=PluginFormat.SOURCE, - source=tmp_path / "plugin.py", - ) - - with pytest.raises(ValueError, match="wheel-format"): - await WheelPluginFormat(base_dir=tmp_path / ".plugin").prepare_async(spec=spec) - - -async def test_wheel_prepare_reuses_unchanged_extraction(tmp_path: Path) -> None: - wheel = _build_wheel(tmp_path) - spec = PluginSpec(name="sample", format=PluginFormat.WHEEL, wheel=wheel) - adapter = WheelPluginFormat(base_dir=tmp_path / ".plugin") - - first = await adapter.prepare_async(spec=spec) - marker = first.import_root / "cache_marker.txt" - marker.write_text("kept", encoding="utf-8") - second = await adapter.prepare_async(spec=spec) - - assert first.import_root == second.import_root - assert marker.is_file() - - -def test_wheel_package_resolution_prefers_explicit(tmp_path: Path) -> None: - assert ( - WheelPluginFormat._resolve_package( - extract_dir=tmp_path, - explicit_package="explicit_pkg", - ) - == "explicit_pkg" - ) - - -def test_wheel_package_resolution_rejects_ambiguous_directory(tmp_path: Path) -> None: - for name in ("pkg_a", "pkg_b"): - package_dir = tmp_path / name - package_dir.mkdir() - (package_dir / "__init__.py").write_text("", encoding="utf-8") - - with pytest.raises(ValueError, match="multiple"): - WheelPluginFormat._resolve_package(extract_dir=tmp_path, explicit_package=None) - - -async def test_source_file_prepare_returns_common_artifact_without_execution(tmp_path: Path) -> None: - source = tmp_path / "operation_foobar.py" - marker = tmp_path / "executed.txt" - source.write_text( - f"from pathlib import Path\nPath({str(marker)!r}).write_text('executed')\n", - encoding="utf-8", - ) - spec = PluginSpec(name="operation_foobar", format=PluginFormat.SOURCE, source=source) - - prepared = await SourcePluginFormat().prepare_async(spec=spec) - - assert prepared.spec is spec - assert prepared.import_root == tmp_path.resolve() - assert prepared.entry_modules == ("operation_foobar",) - assert prepared.owned_module_prefixes == ("operation_foobar",) - assert prepared.artifact_fingerprint - assert prepared.declared_pyrit_version is None - assert not marker.exists() - - -async def test_source_file_prepare_rejects_missing_file(tmp_path: Path) -> None: - spec = PluginSpec( - name="missing", - format=PluginFormat.SOURCE, - source=tmp_path / "missing.py", - ) - - with pytest.raises(PluginSourceNotFoundError, match="existing"): - await SourcePluginFormat().prepare_async(spec=spec) - - -async def test_source_file_prepare_rejects_non_python_file(tmp_path: Path) -> None: - source = tmp_path / "plugin.txt" - source.write_text("not Python", encoding="utf-8") - spec = PluginSpec(name="plugin", format=PluginFormat.SOURCE, source=source) - - with pytest.raises(PluginSourceNotFoundError, match="Python"): - await SourcePluginFormat().prepare_async(spec=spec) - - -async def test_source_file_fingerprint_changes_with_content(tmp_path: Path) -> None: - source = tmp_path / "plugin.py" - source.write_text("VALUE = 1\n", encoding="utf-8") - spec = PluginSpec(name="plugin", format=PluginFormat.SOURCE, source=source) - adapter = SourcePluginFormat() - - first = await adapter.prepare_async(spec=spec) - source.write_text("VALUE = 2\n", encoding="utf-8") - second = await adapter.prepare_async(spec=spec) - - assert first.artifact_fingerprint != second.artifact_fingerprint - - -async def test_source_package_prepare_returns_package_ownership(tmp_path: Path) -> None: - package = tmp_path / "operation_plugin" - package.mkdir() - (package / "__init__.py").write_text("", encoding="utf-8") - (package / "scenarios.py").write_text("VALUE = 1\n", encoding="utf-8") - spec = PluginSpec( - name="operation", - format=PluginFormat.SOURCE, - source=package, - ) - - prepared = await SourcePluginFormat().prepare_async(spec=spec) - - assert prepared.import_root == tmp_path.resolve() - assert prepared.entry_modules == ("operation_plugin",) - assert prepared.owned_module_prefixes == ("operation_plugin",) - - -async def test_source_package_accepts_explicit_nested_entry_module(tmp_path: Path) -> None: - package = tmp_path / "operation_plugin" - package.mkdir() - (package / "__init__.py").write_text("", encoding="utf-8") - (package / "entry.py").write_text("", encoding="utf-8") - spec = PluginSpec( - name="operation", - format=PluginFormat.SOURCE, - source=package, - package="operation_plugin.entry", - ) - - prepared = await SourcePluginFormat().prepare_async(spec=spec) - - assert prepared.entry_modules == ("operation_plugin.entry",) - assert prepared.owned_module_prefixes == ("operation_plugin",) - - -async def test_source_package_rejects_loose_directory(tmp_path: Path) -> None: - source = tmp_path / "loose" - source.mkdir() - (source / "scenario.py").write_text("", encoding="utf-8") - spec = PluginSpec(name="loose", format=PluginFormat.SOURCE, source=source) - - with pytest.raises(PluginSourceNotFoundError, match="__init__.py"): - await SourcePluginFormat().prepare_async(spec=spec) - - -async def test_source_package_rejects_foreign_entry_module(tmp_path: Path) -> None: - package = tmp_path / "operation_plugin" - package.mkdir() - (package / "__init__.py").write_text("", encoding="utf-8") - spec = PluginSpec( - name="operation", - format=PluginFormat.SOURCE, - source=package, - package="other_plugin", - ) - - with pytest.raises(ValueError, match="inside"): - await SourcePluginFormat().prepare_async(spec=spec) - - -async def test_source_package_fingerprint_changes_with_nested_content(tmp_path: Path) -> None: - package = tmp_path / "operation_plugin" - package.mkdir() - (package / "__init__.py").write_text("", encoding="utf-8") - nested = package / "scenario.py" - nested.write_text("VALUE = 1\n", encoding="utf-8") - spec = PluginSpec(name="operation", format=PluginFormat.SOURCE, source=package) - adapter = SourcePluginFormat() - - first = await adapter.prepare_async(spec=spec) - nested.write_text("VALUE = 2\n", encoding="utf-8") - second = await adapter.prepare_async(spec=spec) - - assert first.artifact_fingerprint != second.artifact_fingerprint diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py index 8b508e2afe..5d87b3f741 100644 --- a/tests/unit/setup/test_plugin_loader.py +++ b/tests/unit/setup/test_plugin_loader.py @@ -1,603 +1,120 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -""" -Unit tests for the PyRIT plug-in loader. +"""Tests for the initializer-pointing ``PluginInitializer``. -These tests build a **mock plug-in wheel** at test time (no dependency on any real -plug-in) and exercise the V1 consumer mechanism: prepare, import, discover, validate, -register, and fail-closed rollback. +A plug-in resolves to a ``PyRITInitializer`` reached by dotted path from a source root. +The loader anchors the root on ``sys.path``, imports the initializer, and runs it. The +initializer itself owns all registration; the loader only orchestrates and fails closed. """ -import inspect -import os import sys import textwrap -import uuid -import zipfile -from collections.abc import Iterator -from contextlib import contextmanager from pathlib import Path -from unittest.mock import AsyncMock, MagicMock, patch import pytest -from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider -from pyrit.exceptions import ( - PluginImportError, - PluginLoadError, - PluginRegisteredNothingError, - PluginWheelNotFoundError, -) -from pyrit.memory import CentralMemory -from pyrit.registry import ScenarioRegistry +from pyrit.exceptions import PluginLoadError, PluginSourceNotFoundError +from pyrit.registry import AttackTechniqueRegistry from pyrit.setup import PluginSpec -from pyrit.setup.initialization import IN_MEMORY, initialize_pyrit_async -from pyrit.setup.plugin_loader import ( - PluginInitializer, - load_plugins_if_configured_async, -) - -# --------------------------------------------------------------------------- -# Mock-wheel builder -# --------------------------------------------------------------------------- - - -class MockWheel: - """Handle describing a built mock plug-in wheel.""" - - def __init__(self, *, path: Path, package: str, scenario_name: str, dataset_name: str) -> None: - self.path = path - self.package = package - self.scenario_name = scenario_name - self.dataset_name = dataset_name - - -def _unique_package_name() -> str: - """Return a unique, import-safe mock package name.""" - return f"mock_plugin_{uuid.uuid4().hex[:8]}" - - -def build_mock_wheel( - dest_dir: Path, - *, - bootstrap: str = "initializer", - include_provider: bool = True, - include_scenario: bool = True, - wire_init: bool = True, - package_name: str | None = None, -) -> MockWheel: - """ - Build a mock plug-in wheel in ``dest_dir`` and return a handle to it. +from pyrit.setup.plugin_loader import PluginInitializer - Args: - dest_dir: Directory to write the wheel source tree and .whl into. - bootstrap: Bootstrap style: "initializer" (a PyRITInitializer subclass), - "register" (a top-level register() callable), or "none". - include_provider: Whether to ship a self-registering SeedDatasetProvider. - include_scenario: Whether to ship a Scenario subclass. - wire_init: Whether __init__.py imports the submodules. When False, submodules are - shipped but not imported by __init__, exercising the loader's submodule walk. - package_name: Optional explicit package name; a unique one is generated otherwise. +_TEST_PACKAGE_ROOTS = {"mock_plugin", "notinit"} - Returns: - MockWheel: The built wheel handle (path + package/scenario/dataset names). - """ - package_name = package_name or _unique_package_name() - scenario_name = f"airt.{package_name}" - dataset_name = f"{package_name}_dataset" - src = dest_dir / f"{package_name}_src" - pkg = src / package_name +def _write_plugin_package(root: Path, *, technique_name: str) -> None: + """Write a source package whose initializer registers one private technique.""" + pkg = root / "mock_plugin" pkg.mkdir(parents=True, exist_ok=True) + (pkg / "__init__.py").write_text("", encoding="utf-8") + (pkg / "bootstrap.py").write_text( + textwrap.dedent( + f""" + from pyrit.executor.attack import PromptSendingAttack + from pyrit.registry import AttackTechniqueRegistry + from pyrit.scenario.core import AttackTechniqueFactory + from pyrit.setup.pyrit_initializer import PyRITInitializer - imports = [] - if include_provider: - imports.append("provider") - if include_scenario: - imports.append("scenario") - if bootstrap in ("initializer", "initializer_raises", "register"): - imports.append("bootstrap") - - init_lines = [] - if wire_init and imports: - init_lines.append(f"from . import {', '.join(imports)} # noqa: F401") - if wire_init and bootstrap == "register": - init_lines.append("from .bootstrap import register # noqa: F401") - (pkg / "__init__.py").write_text(("\n".join(init_lines) + "\n") if init_lines else "", encoding="utf-8") - # __file__-relative dataset path (as a real plug-in ships). Only resolves on real disk. - (pkg / "paths.py").write_text( - textwrap.dedent( - """\ - from pathlib import Path + class MockInitializer(PyRITInitializer): + \"\"\"Register one private attack technique.\"\"\" - MOCK_ROOT = Path(__file__, "..").resolve() - MOCK_DATASETS_PATH = Path(MOCK_ROOT, "datasets").resolve() + async def initialize_async(self) -> None: + AttackTechniqueRegistry.get_registry_singleton().register_from_factories( + [AttackTechniqueFactory(name="{technique_name}", attack_class=PromptSendingAttack)] + ) """ ), encoding="utf-8", ) - if include_provider: - datasets = pkg / "datasets" - datasets.mkdir(parents=True, exist_ok=True) - (pkg / "provider.py").write_text( - textwrap.dedent( - f"""\ - from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider - from pyrit.models.seeds.seed_dataset import SeedDataset - - from .paths import MOCK_DATASETS_PATH - - - class MockProvider(SeedDatasetProvider): - @property - def dataset_name(self) -> str: - return "{dataset_name}" - - async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: - return SeedDataset.from_yaml_file(MOCK_DATASETS_PATH / "seed.yaml") - """ - ), - encoding="utf-8", - ) - (datasets / "seed.yaml").write_text( - textwrap.dedent( - f"""\ - dataset_name: {dataset_name} - harm_categories: - - mock - data_type: text - description: mock dataset for plugin test - authors: - - tester - groups: - - test - seeds: - - value: mock prompt one - - value: mock prompt two - - value: mock prompt three - """ - ), - encoding="utf-8", - ) - - if include_scenario: - (pkg / "scenario.py").write_text( - textwrap.dedent( - """\ - from pyrit.models import SeedObjective - from pyrit.scenario import DatasetAttackConfiguration, Scenario, ScenarioTechnique - from pyrit.score import SubStringScorer - - - class MockTechnique(ScenarioTechnique): - ALL = ("all", {"all"}) - DIRECT = ("direct", {"direct"}) - - - class MockScenario(Scenario): - \"\"\"Mock plugin scenario for registration test.\"\"\" - - VERSION = 1 - - def __init__(self, *, scenario_result_id=None): - super().__init__( - version=self.VERSION, - technique_class=MockTechnique, - default_technique=MockTechnique.ALL, - default_dataset_config=DatasetAttackConfiguration( - seeds=[SeedObjective(value="mock objective")] - ), - objective_scorer=SubStringScorer(substring="mock"), - scenario_result_id=scenario_result_id, - ) - - async def _build_atomic_attacks_async(self, *, context): - return [] - """ - ), - encoding="utf-8", - ) - - if bootstrap == "initializer": - (pkg / "bootstrap.py").write_text( - textwrap.dedent( - f"""\ - from pyrit.registry import ScenarioRegistry - from pyrit.setup.pyrit_initializer import PyRITInitializer - - from .scenario import MockScenario - - - class MockBootstrapInitializer(PyRITInitializer): - \"\"\"Register the mock plugin scenario.\"\"\" - - async def initialize_async(self) -> None: - ScenarioRegistry.get_registry_singleton().register_class( - MockScenario, name="{scenario_name}" - ) - """ - ), - encoding="utf-8", - ) - elif bootstrap == "initializer_raises": - (pkg / "bootstrap.py").write_text( - textwrap.dedent( - f"""\ - from pyrit.registry import ScenarioRegistry - from pyrit.setup.pyrit_initializer import PyRITInitializer - - from .scenario import MockScenario - - - class MockBootstrapInitializer(PyRITInitializer): - \"\"\"Register the scenario, then fail to test rollback.\"\"\" - - async def initialize_async(self) -> None: - ScenarioRegistry.get_registry_singleton().register_class( - MockScenario, name="{scenario_name}" - ) - raise RuntimeError("bootstrap failed after registering") - """ - ), - encoding="utf-8", - ) - elif bootstrap == "register": - (pkg / "bootstrap.py").write_text( - textwrap.dedent( - f"""\ - from pyrit.registry import ScenarioRegistry - - from .scenario import MockScenario - - - def register() -> None: - ScenarioRegistry.get_registry_singleton().register_class( - MockScenario, name="{scenario_name}" - ) - """ - ), - encoding="utf-8", - ) - - # Minimal dist-info without top_level.txt so package name inference is exercised. - distinfo = src / f"{package_name}-0.0.1.dist-info" - distinfo.mkdir(parents=True, exist_ok=True) - (distinfo / "METADATA").write_text( - f"Metadata-Version: 2.1\nName: {package_name}\nVersion: 0.0.1\n", encoding="utf-8" - ) - (distinfo / "WHEEL").write_text( - "Wheel-Version: 1.0\nGenerator: test\nRoot-Is-Purelib: true\nTag: py3-none-any\n", encoding="utf-8" - ) - (distinfo / "RECORD").write_text("", encoding="utf-8") - - wheel = dest_dir / f"{package_name}-0.0.1-py3-none-any.whl" - with zipfile.ZipFile(wheel, "w", zipfile.ZIP_DEFLATED) as archive: - for root, _, files in os.walk(src): - for file_name in files: - file_path = Path(root) / file_name - archive.write(file_path, str(file_path.relative_to(src))) - - return MockWheel(path=wheel, package=package_name, scenario_name=scenario_name, dataset_name=dataset_name) - - -# --------------------------------------------------------------------------- -# Fixtures / helpers -# --------------------------------------------------------------------------- - @pytest.fixture(autouse=True) -def plugin_sandbox() -> Iterator[None]: - """Snapshot and restore global import + registry state around each test.""" - sys_path_snapshot = list(sys.path) - provider_snapshot = dict(SeedDatasetProvider._registry) - +def _isolate_import_state(): + original_path = list(sys.path) + original_modules = set(sys.modules) + AttackTechniqueRegistry.reset_registry_singleton() yield + AttackTechniqueRegistry.reset_registry_singleton() + sys.path[:] = original_path + for name in list(sys.modules): + if name not in original_modules and name.split(".", 1)[0] in _TEST_PACKAGE_ROOTS: + del sys.modules[name] - sys.path[:] = sys_path_snapshot - # Only drop the mock plug-in modules this suite imports; leave real pyrit modules - # in place so re-imports don't create duplicate class objects for other tests. - for name in [m for m in sys.modules if m == "mock_plugin" or m.startswith("mock_plugin_")]: - del sys.modules[name] - SeedDatasetProvider._registry.clear() - SeedDatasetProvider._registry.update(provider_snapshot) - ScenarioRegistry.reset_registry_singleton() - - -@contextmanager -def plugin_env(**overrides: str) -> Iterator[None]: - """Patch os.environ so only the given PLUGIN_* overrides are present.""" - with patch.dict(os.environ, overrides, clear=False): - for key in ("PLUGIN_DIR",): - if key not in overrides: - os.environ.pop(key, None) - yield - - -def _spec(wheel: MockWheel, *, package: str | None = None) -> PluginSpec: - """Build a PluginSpec from a mock wheel.""" - return PluginSpec(wheel=wheel.path, package=package) - - -async def load_plugin( - wheel: MockWheel, - plugin_dir: Path, - *, - package: str | None = None, - extra_env: dict[str, str] | None = None, -) -> None: - """Run the plug-in loader against a single mock wheel with an isolated env.""" - with plugin_env(PLUGIN_DIR=str(plugin_dir), **(extra_env or {})): - await load_plugins_if_configured_async(plugins=[_spec(wheel, package=package)]) - - -# --------------------------------------------------------------------------- -# Privileged PluginInitializer -# --------------------------------------------------------------------------- - - -def test_core_initialize_has_no_plugin_parameters() -> None: - parameters = inspect.signature(initialize_pyrit_async).parameters - assert "plugins" not in parameters - assert "plugin_accept_load_failures" not in parameters - - -async def test_plugin_initializer_delegates_to_loader() -> None: - load_plugins_mock = AsyncMock() - spec = PluginSpec(wheel=Path("plugin.whl")) - initializer = PluginInitializer(plugins=[spec]) - - with patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", load_plugins_mock): - await initializer.initialize_async() - - load_plugins_mock.assert_awaited_once() - assert list(load_plugins_mock.await_args.kwargs["plugins"]) == [spec] - - -async def test_plugin_initializer_runs_after_memory_setup() -> None: - memory = MagicMock() - initializer = PluginInitializer(plugins=[PluginSpec(wheel=Path("plugin.whl"))]) - - async def _assert_memory_ready() -> None: - assert CentralMemory.get_memory_instance() is memory - - with ( - patch("pyrit.setup.initialization.SQLiteMemory", return_value=memory), - patch.object(initializer, "initialize_async", side_effect=_assert_memory_ready) as initialize_mock, - ): - await initialize_pyrit_async(IN_MEMORY, initializers=[initializer], env_files=[], silent=True) - - initialize_mock.assert_awaited_once() - - -def test_plugin_initializer_is_not_user_discoverable() -> None: - from pyrit.registry import InitializerRegistry - InitializerRegistry.reset_registry_singleton() - registry = InitializerRegistry.get_registry_singleton() +async def test_runs_configured_initializer(tmp_path: Path) -> None: + _write_plugin_package(tmp_path, technique_name="operation_alpha") + spec = PluginSpec(name="alpha", source=tmp_path, initializer="mock_plugin.bootstrap.MockInitializer") - assert "plugin" not in registry.get_class_names() + await PluginInitializer(plugins=[spec]).initialize_async() + assert "operation_alpha" in AttackTechniqueRegistry.get_registry_singleton().get_factories() -# --------------------------------------------------------------------------- -# No-op behavior -# --------------------------------------------------------------------------- +async def test_missing_source_fails_closed(tmp_path: Path) -> None: + spec = PluginSpec(name="p", source=tmp_path / "nope", initializer="mock_plugin.bootstrap.MockInitializer") + with pytest.raises(PluginSourceNotFoundError): + await PluginInitializer(plugins=[spec]).initialize_async() -async def test_no_op_when_no_plugins() -> None: - """With an empty plug-in list the loader does nothing and registers nothing.""" - providers_before = dict(SeedDatasetProvider.get_all_providers()) - path_before = list(sys.path) - - with plugin_env(): - await load_plugins_if_configured_async(plugins=[]) - - assert SeedDatasetProvider.get_all_providers() == providers_before - assert sys.path == path_before - - -# --------------------------------------------------------------------------- -# Loading via the initializer -# --------------------------------------------------------------------------- - - -async def test_load_extracts_to_plugin_dir(tmp_path: Path) -> None: - """The wheel is extracted (not installed) under the configured plug-in dir.""" - wheel = build_mock_wheel(tmp_path) - plugin_dir = tmp_path / ".plugin" - - await load_plugin(wheel, plugin_dir) - - extract_dir = plugin_dir / wheel.path.stem - assert (extract_dir / wheel.package / "__init__.py").is_file() - assert (extract_dir / wheel.package / "datasets" / "seed.yaml").is_file() - - -async def test_scenario_registration_survives_discovery(tmp_path: Path) -> None: - """A plug-in scenario registered before discovery coexists with built-ins afterwards.""" - wheel = build_mock_wheel(tmp_path, bootstrap="initializer") - - await load_plugin(wheel, tmp_path / ".plugin") - - registry = ScenarioRegistry.get_registry_singleton() - names = registry.get_class_names() - assert "scenario" in names - assert "airt.rapid_response" in names - - mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario - assert registry.get_class("scenario") is mock_scenario - - -async def test_register_callable_bootstrap_is_not_executed(tmp_path: Path) -> None: - """V1 ignores plug-in-authored lifecycle hooks and discovers definitions itself.""" - wheel = build_mock_wheel(tmp_path, bootstrap="register") - - await load_plugin(wheel, tmp_path / ".plugin") - - names = ScenarioRegistry.get_registry_singleton().get_class_names() - assert "scenario" in names - assert wheel.scenario_name not in names - - -async def test_ordering_scenario_visible_to_preload(tmp_path: Path) -> None: - """The plug-in scenario is registered before a later PreloadScenarioMetadata read.""" - wheel = build_mock_wheel(tmp_path, bootstrap="initializer") - - await load_plugin(wheel, tmp_path / ".plugin") - - # get_class_names() is exactly what PreloadScenarioMetadata iterates; the plug-in - # scenario being present proves it registered before that read would happen. - names = ScenarioRegistry.get_registry_singleton().get_class_names() - assert "scenario" in names - - -async def test_plugin_scenario_auto_registered_without_bootstrap(tmp_path: Path) -> None: - """A plug-in's Scenario subclass is auto-registered by discovery even with no bootstrap.""" - wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False) - - await load_plugin(wheel, tmp_path / ".plugin") - - # The scenario is picked up purely by type-scoped discovery of the plug-in package, - # so a scenario-only plug-in with no register()/initializer still loads. - registry = ScenarioRegistry.get_registry_singleton() - mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario - assert mock_scenario in registry._classes.values() - assert registry._discovered is True - - -async def test_framework_owns_registration_when_bootstrap_exists(tmp_path: Path) -> None: - """The framework owns the scenario name even when a register hook exists.""" - wheel = build_mock_wheel(tmp_path, bootstrap="register") - - await load_plugin(wheel, tmp_path / ".plugin") - - registry = ScenarioRegistry.get_registry_singleton() - mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario - registered_names = [name for name, cls in registry._classes.items() if cls is mock_scenario] - assert registered_names == ["scenario"] - - -async def test_submodule_walk_discovers_unwired_components(tmp_path: Path) -> None: - """Scenario discovery walks submodules even when ``__init__.py`` does not import them.""" - wheel = build_mock_wheel(tmp_path, bootstrap="initializer", wire_init=False) - - await load_plugin(wheel, tmp_path / ".plugin") - - assert "scenario" in ScenarioRegistry.get_registry_singleton().get_class_names() - - -async def test_shadowing_installed_package_is_rejected(tmp_path: Path) -> None: - """An installed package of the same name shadowing the plug-in fails loudly.""" - wheel = build_mock_wheel(tmp_path) - - # The spec's package points at a stdlib package that imports from outside the extraction dir. - with pytest.raises(PluginImportError, match="shadowing"): - await load_plugin(wheel, tmp_path / ".plugin", package="json") - - -# --------------------------------------------------------------------------- -# Rollback on failure -# --------------------------------------------------------------------------- - - -async def test_failed_load_rolls_back_syspath(tmp_path: Path) -> None: - """A failed load removes its own sys.path entry (fail-closed leaves no trace).""" - wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) - plugin_dir = tmp_path / ".plugin" +async def test_unresolvable_initializer_fails_closed(tmp_path: Path) -> None: + _write_plugin_package(tmp_path, technique_name="operation_beta") + spec = PluginSpec(name="p", source=tmp_path, initializer="mock_plugin.bootstrap.DoesNotExist") with pytest.raises(PluginLoadError): - await load_plugin(wheel, plugin_dir) - - extract_dir = str(plugin_dir / wheel.path.stem) - assert extract_dir not in sys.path - - -# --------------------------------------------------------------------------- -# Failure modes -# --------------------------------------------------------------------------- - - -async def test_missing_wheel_fails_closed() -> None: - """A configured-but-missing wheel raises PluginWheelNotFoundError by default (fail-closed).""" - with plugin_env(): - with pytest.raises(PluginWheelNotFoundError, match="Failed to load plug-in"): - await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=Path("does_not_exist.whl"))]) - - -async def test_empty_wheel_is_loud(tmp_path: Path) -> None: - """A wheel that imports cleanly but registers nothing fails loudly.""" - wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) - - with pytest.raises(PluginRegisteredNothingError, match="no scenarios or attack techniques"): - await load_plugin(wheel, tmp_path / ".plugin") - - -async def test_non_whl_path_fails_closed(tmp_path: Path) -> None: - """A wheel path that is not a .whl file fails closed with PluginWheelNotFoundError.""" - not_a_wheel = tmp_path / "plugin.zip" - not_a_wheel.write_text("not a wheel", encoding="utf-8") - - with plugin_env(): - with pytest.raises(PluginWheelNotFoundError, match="Failed to load plug-in"): - await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=not_a_wheel)]) + await PluginInitializer(plugins=[spec]).initialize_async() -def test_error_subclasses_are_plugin_load_errors() -> None: - """All specific plug-in errors subclass PluginLoadError so one except still catches them.""" - for error_cls in (PluginWheelNotFoundError, PluginImportError, PluginRegisteredNothingError): - assert issubclass(error_cls, PluginLoadError) +async def test_non_initializer_target_fails_closed(tmp_path: Path) -> None: + pkg = tmp_path / "notinit" + pkg.mkdir() + (pkg / "__init__.py").write_text("class NotAnInitializer:\n pass\n", encoding="utf-8") + spec = PluginSpec(name="p", source=tmp_path, initializer="notinit.NotAnInitializer") + with pytest.raises(PluginLoadError, match="not a PyRITInitializer"): + await PluginInitializer(plugins=[spec]).initialize_async() -async def test_wheel_with_path_traversal_member_fails_closed(tmp_path: Path) -> None: - """A wheel containing a path-traversal member is rejected during safe extraction.""" - malicious = tmp_path / "evil-0.0.1-py3-none-any.whl" - with zipfile.ZipFile(malicious, "w") as archive: - archive.writestr("evil_pkg/__init__.py", "") - archive.writestr("../escape.py", "compromised = True") - - with plugin_env(PLUGIN_DIR=str(tmp_path / ".plugin")): - with pytest.raises(PluginLoadError, match="Failed to load plug-in"): - await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=malicious)]) - - # The traversal target was not written outside the extraction directory. - assert not (tmp_path / "escape.py").exists() - - -# --------------------------------------------------------------------------- -# No-arg-instantiable contract -# --------------------------------------------------------------------------- - - -def test_non_no_arg_scenario_fails_metadata_cleanly() -> None: - """A registered scenario that is not no-arg instantiable fails metadata build clearly.""" - from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse - - class BadScenario(RapidResponse): - """Scenario that violates the no-arg-instantiable contract.""" - - def __init__(self, *, required_value: str) -> None: - super().__init__() - self._required_value = required_value - - registry = ScenarioRegistry() - registry.register_class(BadScenario, name="airt.bad") # signature-only validation passes - - with pytest.raises(TypeError, match="no arguments"): - registry._build_metadata("airt.bad", BadScenario) - +async def test_initializer_that_raises_fails_closed(tmp_path: Path) -> None: + pkg = tmp_path / "notinit" + pkg.mkdir() + (pkg / "__init__.py").write_text( + textwrap.dedent( + """ + from pyrit.setup.pyrit_initializer import PyRITInitializer -# --------------------------------------------------------------------------- -# Multiple plug-ins -# --------------------------------------------------------------------------- + class Boom(PyRITInitializer): + async def initialize_async(self) -> None: + raise RuntimeError("boom") + """ + ), + encoding="utf-8", + ) + spec = PluginSpec(name="p", source=tmp_path, initializer="notinit.Boom") + with pytest.raises(PluginLoadError, match="initializer failed"): + await PluginInitializer(plugins=[spec]).initialize_async() -async def test_multiple_plugins_rejected(tmp_path: Path) -> None: - """V1 rejects plug-in composition.""" - wheel_a = build_mock_wheel(tmp_path / "a", package_name="mock_plugin_alpha") - wheel_b = build_mock_wheel(tmp_path / "b", package_name="mock_plugin_beta") - with plugin_env(PLUGIN_DIR=str(tmp_path / ".plugin")): - with pytest.raises(ValueError, match="one plug-in"): - await load_plugins_if_configured_async(plugins=[_spec(wheel_a), _spec(wheel_b)]) +def test_requires_exactly_one_plugin(tmp_path: Path) -> None: + spec = PluginSpec(name="p", source=tmp_path, initializer="m.C") + with pytest.raises(ValueError): + PluginInitializer(plugins=[spec, spec]) diff --git a/tests/unit/setup/test_plugin_registration.py b/tests/unit/setup/test_plugin_registration.py deleted file mode 100644 index 615bb5b58f..0000000000 --- a/tests/unit/setup/test_plugin_registration.py +++ /dev/null @@ -1,185 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import sys -from collections.abc import Iterator -from pathlib import Path -from unittest.mock import patch - -import pytest - -from pyrit.datasets import SeedDatasetProvider -from pyrit.exceptions import PluginCollisionError, PluginValidationError -from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry -from pyrit.setup import PluginFormat, PluginSpec -from pyrit.setup.plugin_loader import PluginInitializer - - -@pytest.fixture(autouse=True) -def reset_plugin_registration_state() -> Iterator[None]: - path_snapshot = list(sys.path) - module_snapshot = set(sys.modules) - provider_snapshot = dict(SeedDatasetProvider._registry) - AttackTechniqueRegistry.reset_registry_singleton() - ScenarioRegistry.reset_registry_singleton() - yield - AttackTechniqueRegistry.reset_registry_singleton() - ScenarioRegistry.reset_registry_singleton() - SeedDatasetProvider._registry.clear() - SeedDatasetProvider._registry.update(provider_snapshot) - sys.path[:] = path_snapshot - for name in set(sys.modules) - module_snapshot: - del sys.modules[name] - - -def _source_spec(source: Path, *, name: str = "operation") -> PluginSpec: - return PluginSpec(name=name, format=PluginFormat.SOURCE, source=source) - - -async def test_plugin_initializer_registers_source_technique(tmp_path: Path) -> None: - source = tmp_path / "operation_technique.py" - source.write_text( - """from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory - -OPERATION = AttackTechniqueFactory( - name="operation_foobar", - attack_class=PromptSendingAttack, - technique_tags=["single_turn", "scenario:airt.rapid_response"], -) -""", - encoding="utf-8", - ) - - await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() - - registry = AttackTechniqueRegistry.get_registry_singleton() - assert registry.get_factories()["operation_foobar"] is not None - entry = registry.instances.get_entry("operation_foobar") - assert entry is not None - assert entry.metadata["plugin_name"] == "operation" - assert entry.metadata["scenario_names"] == ["airt.rapid_response"] - - -async def test_plugin_initializer_registers_source_scenario( - tmp_path: Path, - patch_central_database: None, -) -> None: - source = tmp_path / "private_scenario.py" - source.write_text( - """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse - -class PrivateScenario(RapidResponse): - pass -""", - encoding="utf-8", - ) - - with patch.dict( - "os.environ", - { - "OPENAI_CHAT_ENDPOINT": "https://example.test", - "OPENAI_CHAT_KEY": "test-key", - "OPENAI_CHAT_MODEL": "test-model", - }, - ): - await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() - - scenario_class = ScenarioRegistry.get_registry_singleton().get_class("private_scenario") - assert scenario_class.__name__ == "PrivateScenario" - - -async def test_plugin_initializer_rejects_latent_builtin_technique_name(tmp_path: Path) -> None: - source = tmp_path / "collision.py" - source.write_text( - """from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory - -OPERATION = AttackTechniqueFactory( - name="role_play_movie_script", - attack_class=PromptSendingAttack, - technique_tags=["scenario:airt.rapid_response"], -) -""", - encoding="utf-8", - ) - - with pytest.raises(PluginCollisionError, match="role_play_movie_script"): - await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() - - assert "role_play_movie_script" not in AttackTechniqueRegistry.get_registry_singleton().get_factories() - - -async def test_plugin_initializer_rolls_back_scenario_when_technique_collides(tmp_path: Path) -> None: - source = tmp_path / "partial.py" - source.write_text( - """from pyrit.executor.attack import PromptSendingAttack -from pyrit.scenario import AttackTechniqueFactory -from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse - -class PartialScenario(RapidResponse): - pass - -OPERATION = AttackTechniqueFactory( - name="role_play_movie_script", - attack_class=PromptSendingAttack, - technique_tags=["scenario:airt.rapid_response"], -) -""", - encoding="utf-8", - ) - - with pytest.raises(PluginCollisionError): - await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() - - assert "partial" not in ScenarioRegistry.get_registry_singleton()._classes - assert str(tmp_path.resolve()) not in sys.path - assert "partial" not in sys.modules - - -async def test_plugin_initializer_removes_unsupported_provider_side_effect(tmp_path: Path) -> None: - source = tmp_path / "provider_side_effect.py" - source.write_text( - """from pyrit.datasets import SeedDatasetProvider -from pyrit.executor.attack import PromptSendingAttack -from pyrit.models import SeedDataset -from pyrit.scenario import AttackTechniqueFactory - -class PrivateProvider(SeedDatasetProvider): - @property - def dataset_name(self): - return "private" - - async def fetch_dataset_async(self, *, cache=True): - return SeedDataset(seeds=[], dataset_name="private") - -OPERATION = AttackTechniqueFactory( - name="provider_safe", - attack_class=PromptSendingAttack, - technique_tags=["scenario:airt.rapid_response"], -) -""", - encoding="utf-8", - ) - - await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() - - assert "PrivateProvider" not in SeedDatasetProvider.get_all_providers() - - -async def test_plugin_initializer_rejects_scenario_that_cannot_build_metadata(tmp_path: Path) -> None: - source = tmp_path / "invalid_scenario.py" - source.write_text( - """from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse - -class InvalidScenario(RapidResponse): - def __init__(self, *, required_value): - super().__init__() -""", - encoding="utf-8", - ) - - with pytest.raises(PluginValidationError, match="metadata"): - await PluginInitializer(plugins=[_source_spec(source)]).initialize_async() - - assert "invalid_scenario" not in ScenarioRegistry.get_registry_singleton()._classes diff --git a/tests/unit/setup/test_plugin_spec.py b/tests/unit/setup/test_plugin_spec.py index a4d33ca0f0..9a9cfb73c6 100644 --- a/tests/unit/setup/test_plugin_spec.py +++ b/tests/unit/setup/test_plugin_spec.py @@ -1,69 +1,55 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +"""Tests for the source-only, initializer-pointing ``PluginSpec`` schema.""" + from pathlib import Path import pytest -from pyrit.setup import PluginFormat, PluginSpec - +from pyrit.setup import PluginSpec -def test_source_spec_resolves_relative_path(tmp_path: Path) -> None: - spec = PluginSpec.from_config( - { - "name": "operation_foobar", - "format": "source", - "source": "plugins/operation_foobar.py", - }, - base_dir=tmp_path, - ) - assert spec.name == "operation_foobar" - assert spec.format is PluginFormat.SOURCE - assert spec.source == (tmp_path / "plugins" / "operation_foobar.py").resolve() - assert spec.wheel is None - assert spec.artifact_path == spec.source +def test_from_config_normalizes_source_and_initializer(tmp_path: Path) -> None: + entry = { + "name": "rapid_response", + "source": "pkg_root", + "initializer": "pyrit_internal.setup.RapidResponseInitializer", + } + spec = PluginSpec.from_config(entry, base_dir=tmp_path) + assert spec.name == "rapid_response" + assert spec.source == (tmp_path / "pkg_root").resolve() + assert spec.initializer == "pyrit_internal.setup.RapidResponseInitializer" -def test_wheel_spec_resolves_package_and_path(tmp_path: Path) -> None: +def test_from_config_keeps_absolute_source(tmp_path: Path) -> None: + absolute = (tmp_path / "artifacts").resolve() spec = PluginSpec.from_config( - { - "name": "partner_scenarios", - "format": "wheel", - "wheel": "plugins/partner.whl", - "package": "partner_scenarios.plugin", - }, - base_dir=tmp_path, + {"name": "p", "source": str(absolute), "initializer": "m.C"}, + base_dir=tmp_path / "elsewhere", ) + assert spec.source == absolute - assert spec.name == "partner_scenarios" - assert spec.format is PluginFormat.WHEEL - assert spec.wheel == (tmp_path / "plugins" / "partner.whl").resolve() - assert spec.source is None - assert spec.package == "partner_scenarios.plugin" - assert spec.artifact_path == spec.wheel + +def test_to_config_round_trips() -> None: + spec = PluginSpec(name="p", source=Path("/opt/plugin").resolve(), initializer="m.C") + assert spec.to_config() == { + "name": "p", + "source": str(Path("/opt/plugin").resolve()), + "initializer": "m.C", + } @pytest.mark.parametrize( "entry", [ - {}, - {"name": "x", "format": "source"}, - {"name": "x", "format": "wheel"}, - {"name": "x", "format": "source", "source": "x.py", "wheel": "x.whl"}, - {"name": "x", "format": "source", "wheel": "x.whl"}, - {"name": "x", "format": "wheel", "source": "x.py"}, - {"name": "Bad-Name", "format": "source", "source": "x.py"}, - {"name": "x", "format": "wheel", "wheel": "x.whl", "package": "bad-package"}, - {"name": "x", "format": "archive", "wheel": "x.whl"}, - {"name": "x", "format": "source", "source": "x.py", "unexpected": True}, + {"source": "s", "initializer": "m.C"}, # missing name + {"name": "p", "initializer": "m.C"}, # missing source + {"name": "p", "source": "s"}, # missing initializer + {"name": "p", "source": "s", "initializer": "no_dot"}, # non-dotted initializer + {"name": "p", "source": "s", "initializer": "m.C", "wheel": "x"}, # unexpected key ], ) -def test_plugin_spec_rejects_invalid_config(entry: dict[str, object], tmp_path: Path) -> None: +def test_from_config_rejects_malformed_entries(entry: dict) -> None: with pytest.raises(ValueError): - PluginSpec.from_config(entry, base_dir=tmp_path) - - -def test_plugin_spec_rejects_legacy_shorthand() -> None: - with pytest.raises(ValueError, match="mapping"): - PluginSpec.from_config("plugin.whl") # type: ignore[arg-type] + PluginSpec.from_config(entry) From 5b3e996e423f7652da83b3d0f9d868f06e2975eb Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 14 Jul 2026 13:51:14 -0700 Subject: [PATCH 34/34] docs: address review comments; drop version and internal-repo references Resolve the inline review TODOs and genericize the plug-in docs/config for a shipped feature: - Explain, in plain language, what 'source' should point at and why the package root goes on sys.path (and the consequences of getting it wrong). - Replace the 'do not list it under initializers:' instruction with what actually happens: the privileged initializer is not a registered name, so listing it fails; the framework injects it. - Reframe the 'gray-area' guidance from the operator's perspective (contribute publicly vs keep tracked vs keep private; sometimes only the dataset/technique/scenario is sensitive). - Remove all 'V1' version references from docs, config example, and code messages/docstrings. - Replace every pyrit-internal / RapidResponseInitializer example with a generic my_redteam package, and rename the integration-test namespace accordingly. - Fix a stale 'source or wheel' docstring left in ConfigurationLoader._normalize_plugins. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 153315ea-4360-489a-89c5-630a41cf3fd2 --- .pyrit_conf_example | 11 ++-- doc/getting_started/plugins.md | 58 +++++++++++++------ doc/getting_started/pyrit_conf.md | 12 ++-- .../troubleshooting/plugins.md | 17 +++--- pyrit/setup/configuration_loader.py | 10 ++-- pyrit/setup/plugin_loader.py | 4 +- .../setup/test_plugin_loader_integration.py | 12 ++-- tests/unit/setup/test_plugin_spec.py | 4 +- 8 files changed, 77 insertions(+), 51 deletions(-) diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 5a072c4dda..2d1166fcad 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -66,9 +66,10 @@ initializers: # else consumes the registries. Do NOT list it under `initializers:`. # # The plug-in's initializer owns all registration; PyRIT discovers nothing on its own. -# V1 accepts one plug-in with three fields: +# A plug-in is one entry with three fields: # - name: an operator label (used in logs/errors) -# - source: a directory placed on sys.path so `import ` resolves +# - source: the folder that contains your package, placed on sys.path so +# `import ` resolves # - initializer: dotted 'module.Class' path to a PyRITInitializer subclass # Relative source paths resolve against this configuration file. # @@ -77,9 +78,9 @@ initializers: # # Example: # plugins: -# - name: rapid_response -# source: /repos/pyrit-internal -# initializer: pyrit_internal.setup.initializers.RapidResponseInitializer +# - name: my_redteam +# source: /repos/my-redteam +# initializer: my_redteam.setup.MyInitializer # # Help and troubleshooting: # - doc/getting_started/plugins.md diff --git a/doc/getting_started/plugins.md b/doc/getting_started/plugins.md index d00817468c..29dfef7ae6 100644 --- a/doc/getting_started/plugins.md +++ b/doc/getting_started/plugins.md @@ -27,19 +27,33 @@ dotted path from a source root: ```yaml # .pyrit_conf plugins: - - name: rapid_response - source: /repos/pyrit-internal - initializer: pyrit_internal.setup.initializers.RapidResponseInitializer + - name: my_redteam + source: /repos/my-redteam + initializer: my_redteam.setup.MyInitializer ``` - `name` — an operator label used in logs and errors. -- `source` — a directory placed on `sys.path` so `import ` resolves. A - relative path resolves against the config file. +- `source` — the folder placed on `sys.path` so your package can be imported (see below). - `initializer` — a dotted `module.Class` path to a concrete `PyRITInitializer`. -`ConfigurationLoader` prepends a privileged initializer that anchors `source`, imports -the initializer, and runs it before any user-configured initializer. Do **not** list it -under `initializers:`. +### What `source` should point at (and why it matters) + +`source` should be the folder that **contains** your package — the directory you would be +sitting in for `import my_redteam` to succeed at a Python prompt. PyRIT adds that folder +to Python's import search path (`sys.path`) before importing your initializer. + +You need this because a real private package doesn't live next to PyRIT; its modules +import from *each other* (for example `from my_redteam.datasets import load`). If Python +can't find the package root, those imports fail and the plug-in won't load. In plain +terms: point `source` at the folder above your package, not at the package folder itself +and not at a single file buried inside it. If you point at the wrong place, loading fails +closed with an import error naming what could not be found. + +`ConfigurationLoader` runs the plug-in as a privileged initializer, always **first** — +before your other initializers and before anything reads the scenario/technique catalog. +You do not (and cannot) add it to `initializers:`: it isn't a registered initializer name, +so listing it there fails with an "initializer not found" error. The framework constructs +and runs it for you. ## The initializer owns registration @@ -50,15 +64,25 @@ wants discoverable, at whatever level of abstraction fits: `AttackTechniqueRegistry.get_registry_singleton().register_from_factories([...])`; selectable via `--techniques`. - **Scenarios** — - `ScenarioRegistry.get_registry_singleton().register_class(MyScenario, name="airt_internal.violence")`; - runnable via `pyrit_scan airt_internal.violence`. + `ScenarioRegistry.get_registry_singleton().register_class(MyScenario, name="my_redteam.violence")`; + runnable via `pyrit_scan my_redteam.violence`. - **Datasets** — register providers and load them into memory so private seeds stay in the operator's database and are never published. - **Default targets** — `set_default_value(...)`. -This is why the plug-in fits gray-area content: sometimes only the *dataset* must stay -private, sometimes a *technique*, and sometimes an entire *scenario* — the initializer -decides. +### Deciding what to keep private + +"Private" is rarely all-or-nothing. If you build a custom scenario or technique, you have +to decide whether to contribute it publicly, keep it in your own tracked repo, or keep it +fully private — and often only *part* of it is sensitive. A plug-in lets you keep exactly +the sensitive layer private while everything else stays public and works out of the box: + +- sometimes only the **dataset** (the prompts/objectives) is sensitive; +- sometimes it's a niche **technique**; +- sometimes an entire **scenario** should not be exposed at all. + +Because the initializer registers each level independently, you can publish the generic +parts and register only the sensitive parts from your private package. ### Example initializer @@ -68,7 +92,7 @@ from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry from pyrit.scenario.core import AttackTechniqueFactory from pyrit.setup.pyrit_initializer import PyRITInitializer -from my_package.scenarios import Violence +from my_redteam.scenarios import Violence class MyInitializer(PyRITInitializer): @@ -78,7 +102,7 @@ class MyInitializer(PyRITInitializer): AttackTechniqueRegistry.get_registry_singleton().register_from_factories( [AttackTechniqueFactory(name="operation_foobar", attack_class=PromptSendingAttack)] ) - ScenarioRegistry.get_registry_singleton().register_class(Violence, name="airt_internal.violence") + ScenarioRegistry.get_registry_singleton().register_class(Violence, name="my_redteam.violence") ``` ## Usage @@ -88,7 +112,7 @@ class MyInitializer(PyRITInitializer): pyrit_scan airt.rapid_response --target openai_chat --techniques operation_foobar # A private scenario -pyrit_scan airt_internal.violence --target openai_chat +pyrit_scan my_redteam.violence --target openai_chat ``` ## Behavior and limits @@ -98,7 +122,7 @@ pyrit_scan airt_internal.violence --target openai_chat - Loading executes third-party Python with backend permissions; whoever can write the config or the source can run code on the host. Treat the config as sensitive. - Dependencies must already be installed in the backend environment. -- V1 is **fail-closed** and supports **one** plug-in. A failed load aborts +- The plug-in path is **fail-closed** and supports **one** plug-in. A failed load aborts initialization — fix the config or source and restart. - Plug-ins activate only at process/backend startup. Restart after changing the config or the source; there is no hot reload. diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index ef04b93397..44bd4c2a3d 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -51,6 +51,7 @@ flowchart LR ### Using .env.local for Overrides You can use `~/.pyrit/.env.local` to override values in `~/.pyrit/.env` without modifying the base file. This is useful for: + - Testing different targets - Using personal credentials instead of shared ones - Switching between configurations quickly @@ -160,14 +161,15 @@ targets it wants discoverable: ```yaml plugins: - - name: rapid_response - source: /repos/pyrit-internal - initializer: pyrit_internal.setup.initializers.RapidResponseInitializer + - name: my_redteam + source: /repos/my-redteam + initializer: my_redteam.setup.MyInitializer ``` `ConfigurationLoader` automatically injects this as a privileged initializer that runs -first; do not add it to `initializers:`. V1 is fail-closed, supports one plug-in, and -requires a process/backend restart after changes. +first; you cannot add it to `initializers:` yourself (it is not a registered initializer +name). It is fail-closed, supports one plug-in, and requires a process/backend restart +after changes. See [Private Scenarios and Attack Techniques](./plugins.md) and [Plug-In Troubleshooting](./troubleshooting/plugins.md). diff --git a/doc/getting_started/troubleshooting/plugins.md b/doc/getting_started/troubleshooting/plugins.md index c6b2f92693..94e5a11ee6 100644 --- a/doc/getting_started/troubleshooting/plugins.md +++ b/doc/getting_started/troubleshooting/plugins.md @@ -14,13 +14,13 @@ its startup config. ## Configuration is rejected -V1 accepts exactly one entry with three fields: +A plug-in is one entry with three fields: ```yaml plugins: - - name: rapid_response - source: /repos/pyrit-internal - initializer: pyrit_internal.setup.initializers.RapidResponseInitializer + - name: my_redteam + source: /repos/my-redteam + initializer: my_redteam.setup.MyInitializer ``` Check that: @@ -41,8 +41,7 @@ directory that contains your package (the parent of the top-level package) so - Confirm `initializer` names a real `module.Class` importable from `source`. - Install the plug-in's dependencies into the backend Python environment. - A packaged initializer must import from its own package (for example - `from my_package... import ...`); pointing `source` at the package root makes that - work. + `from my_redteam... import ...`); pointing `source` at the package root makes that work. ## Target is not a PyRITInitializer @@ -59,6 +58,6 @@ techniques. Datasets must be registered as providers and loaded into memory. ## Partial state after a failed load -V1 is fail-closed. A failed plug-in load aborts initialization; fix the reported stage -and restart the process. Continuing in the same process after a failed initialization is -not supported. +Plug-in loading is fail-closed. A failed load aborts initialization; fix the reported +stage and restart the process. Continuing in the same process after a failed +initialization is not supported. diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 299ebe8d1b..d63986344d 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -111,7 +111,7 @@ class ConfigurationLoader(YamlLoadable): operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. plugins: List of plug-ins to load as the guaranteed-first initialization phase. - V1 accepts at most one explicit mapping with ``name``, a ``source`` path, and + Accepts at most one explicit mapping with ``name``, a ``source`` path, and a dotted ``initializer`` (module.Class) pointing at a ``PyRITInitializer``. Empty means no plug-ins. @@ -233,9 +233,9 @@ def _normalize_plugins(self) -> None: """ Normalize ``plugins`` entries to ``PluginSpec`` instances. - V1 accepts at most one explicit mapping that selects the ``source`` or ``wheel`` - format. Validating here (before any other normalization) fails fast on a - malformed plug-in entry. + Accepts at most one explicit mapping (``name`` + ``source`` + dotted + ``initializer``). Validating here (before any other normalization) fails fast on + a malformed plug-in entry. Raises: ValueError: If a plug-in entry has an unsupported shape. @@ -243,7 +243,7 @@ def _normalize_plugins(self) -> None: from pyrit.setup.plugin_spec import PluginSpec if len(self.plugins) > 1: - raise ValueError("V1 supports one plug-in at a time; plug-in composition is not supported.") + raise ValueError("Only one plug-in is supported at a time; plug-in composition is not supported.") self._plugin_specs = [PluginSpec.from_config(entry, base_dir=self._plugin_base_dir) for entry in self.plugins] def _normalize_scenario(self) -> None: diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py index 10c891a67a..31fc2e65a5 100644 --- a/pyrit/setup/plugin_loader.py +++ b/pyrit/setup/plugin_loader.py @@ -44,11 +44,11 @@ def __init__(self, *, plugins: Sequence[PluginSpec]) -> None: plugins (Sequence[PluginSpec]): The normalized plug-in declarations. Raises: - ValueError: If anything other than exactly one plug-in is supplied (V1). + ValueError: If anything other than exactly one plug-in is supplied. """ super().__init__() if len(plugins) != 1: - raise ValueError("PluginInitializer requires exactly one plug-in in V1.") + raise ValueError("PluginInitializer requires exactly one plug-in.") self._spec = plugins[0] @property diff --git a/tests/integration/setup/test_plugin_loader_integration.py b/tests/integration/setup/test_plugin_loader_integration.py index b43794374f..10b01e92ff 100644 --- a/tests/integration/setup/test_plugin_loader_integration.py +++ b/tests/integration/setup/test_plugin_loader_integration.py @@ -19,7 +19,7 @@ from pyrit.setup import PluginSpec from pyrit.setup.plugin_loader import PluginInitializer -_PLUGIN_PACKAGE = "airt_internal_plugin" +_PLUGIN_PACKAGE = "my_redteam_plugin" def _write_source_plugin(root: Path) -> None: @@ -97,7 +97,7 @@ async def initialize_async(self) -> None: [AttackTechniqueFactory(name="operation_foobar", attack_class=PromptSendingAttack)] ) ScenarioRegistry.get_registry_singleton().register_class( - PrivateScenario, name="airt_internal.private" + PrivateScenario, name="my_redteam.private" ) ''' ), @@ -124,7 +124,7 @@ def _isolate_registries_and_imports(): async def test_source_plugin_initializer_registers_real_components(tmp_path: Path) -> None: _write_source_plugin(tmp_path) spec = PluginSpec( - name="airt_internal", + name="my_redteam", source=tmp_path, initializer=f"{_PLUGIN_PACKAGE}.bootstrap.PrivateInitializer", ) @@ -132,8 +132,8 @@ async def test_source_plugin_initializer_registers_real_components(tmp_path: Pat await PluginInitializer(plugins=[spec]).initialize_async() scenario_registry = ScenarioRegistry.get_registry_singleton() - assert "airt_internal.private" in scenario_registry.get_class_names() - assert scenario_registry.get_class("airt_internal.private").__name__ == "PrivateScenario" + assert "my_redteam.private" in scenario_registry.get_class_names() + assert scenario_registry.get_class("my_redteam.private").__name__ == "PrivateScenario" assert "operation_foobar" in AttackTechniqueRegistry.get_registry_singleton().get_factories() @@ -143,7 +143,7 @@ async def test_source_plugin_missing_initializer_fails_closed(tmp_path: Path) -> _write_source_plugin(tmp_path) spec = PluginSpec( - name="airt_internal", + name="my_redteam", source=tmp_path, initializer=f"{_PLUGIN_PACKAGE}.bootstrap.DoesNotExist", ) diff --git a/tests/unit/setup/test_plugin_spec.py b/tests/unit/setup/test_plugin_spec.py index 9a9cfb73c6..7df4ce607b 100644 --- a/tests/unit/setup/test_plugin_spec.py +++ b/tests/unit/setup/test_plugin_spec.py @@ -14,12 +14,12 @@ def test_from_config_normalizes_source_and_initializer(tmp_path: Path) -> None: entry = { "name": "rapid_response", "source": "pkg_root", - "initializer": "pyrit_internal.setup.RapidResponseInitializer", + "initializer": "my_redteam.setup.MyInitializer", } spec = PluginSpec.from_config(entry, base_dir=tmp_path) assert spec.name == "rapid_response" assert spec.source == (tmp_path / "pkg_root").resolve() - assert spec.initializer == "pyrit_internal.setup.RapidResponseInitializer" + assert spec.initializer == "my_redteam.setup.MyInitializer" def test_from_config_keeps_absolute_source(tmp_path: Path) -> None: