From 9f409abce269546ede61594857f50bb13246d3ea Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 7 Jul 2026 16:06:08 -0700 Subject: [PATCH 1/8] 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 2/8] 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 3/8] 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 4/8] 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 5/8] 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 6/8] 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 7/8] 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 8/8] 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():