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..3e5d9b3c07 --- /dev/null +++ b/pyrit/setup/plugin_loader.py @@ -0,0 +1,629 @@ +# 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: + 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) + 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