From 9f409abce269546ede61594857f50bb13246d3ea Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 7 Jul 2026 16:06:08 -0700 Subject: [PATCH 01/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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)