diff --git a/.env_example b/.env_example index 0d70a48cfd..0bec2cbb6e 100644 --- a/.env_example +++ b/.env_example @@ -250,6 +250,7 @@ AZURE_OPENAI_VIDEO_MODEL="sora-2" AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" OPENAI_VIDEO_ENDPOINT = ${AZURE_OPENAI_VIDEO_ENDPOINT} + OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" diff --git a/.gitignore b/.gitignore index 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..e14a0c599d 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 are configured below via the dedicated `plugins:` key, NOT as an entry +# in this `initializers:` list. Plug-ins load as a guaranteed-first phase (before these +# initializers), so there is no ordering to get right here. +# # Each initializer can be specified as: # - A simple string (name only) # - A dictionary with 'name' and optional 'args' for parameters @@ -53,6 +57,36 @@ initializers: - name: technique - name: load_default_datasets +# Plug-ins +# -------- +# Non-disclosable plug-ins (extra scenarios + datasets) shipped as pre-built wheels and +# loaded at initialization. Plug-ins load as a guaranteed-first phase — after memory is set +# and BEFORE the initializers above — so their scenarios/datasets register before anything +# reads the registry. They are a dedicated phase, not an entry in `initializers:`, so there +# is no ordering to get right. Loaded in list order; one plug-in behaves identically to +# several. A configless setup simply has no plug-ins. +# +# Each entry can be: +# - a wheel path (package auto-detected): - /abs/path/to/my_plugin.whl +# - a {package: wheel} pair (names the package): - my_plugin: /abs/path/to/my_plugin.whl +# - an explicit mapping: - wheel: /abs/path/to/my_plugin.whl +# package: my_plugin +# +# WARNING: loading a plug-in executes third-party code in the PyRIT process. Whoever can +# write this file can run code on the host. Treat it as sensitive. +# +# Example: +# plugins: +# - my_first_plugin: /abs/path/to/first-0.0.0-py3-none-any.whl +# - my_second_plugin: /abs/path/to/second-0.0.0-py3-none-any.whl + +# Accept plug-in load failures +# ---------------------------- +# When true, a plug-in that fails to load is skipped with a warning instead of aborting +# initialization. Defaults to false (fail-closed). Precedence: this value wins when set; +# the PLUGIN_ACCEPT_LOAD_FAILURES env var is used only when this is omitted. +# plugin_accept_load_failures: false + # Default Scenario # ---------------- # Optional default scenario to run when invoking `pyrit_scan` without a diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py index f22686e755..af4284506d 100644 --- a/pyrit/registry/components/scenario_registry.py +++ b/pyrit/registry/components/scenario_registry.py @@ -109,6 +109,35 @@ def _get_registry_name(self, cls: type[Scenario]) -> str: return relative return class_name_to_snake_case(cls.__name__, suffix="Scenario") + def _external_registry_name(self, cls: type[Scenario], *, package_name: str) -> str: + """ + Key a plug-in scenario by its module path within the plug-in package. + + Mirrors the built-in dotted-name scheme (module path relative to the scenarios + package) so plug-in scenarios read like built-ins and same-named scenarios in + different submodules do not collide on a bare class name. The name is the class's + module relative to the plug-in's top-level ``package_name`` with a leading + ``scenario.scenarios.`` / ``scenarios.`` segment stripped to match built-in style; + it falls back to the suffix-stripped snake_case class name when no module context + is available. + + Args: + cls (type[Scenario]): The plug-in scenario class. + package_name (str): The plug-in's top-level package name. + + Returns: + str: The dotted registry name for the plug-in scenario. + """ + module = cls.__module__ or "" + prefix = f"{package_name}." + relative = module[len(prefix) :] if module.startswith(prefix) else module + for marker in ("scenario.scenarios.", "scenarios."): + index = relative.find(marker) + if index != -1: + relative = relative[index + len(marker) :] + break + return relative or class_name_to_snake_case(cls.__name__, suffix="Scenario") + def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: """ Build metadata for a Scenario class. diff --git a/pyrit/registry/registry.py b/pyrit/registry/registry.py index e159016a0b..578c815334 100644 --- a/pyrit/registry/registry.py +++ b/pyrit/registry/registry.py @@ -272,9 +272,7 @@ def _discover(self) -> None: ``_base_type``/``_discovery_package``. """ package = self._discovery_package() - base = self._base_type() package_name = package.__name__ - package_prefix = f"{package_name}." # Materialize lazily-exported classes so they are loaded before enumeration. # A lazy import backed by an optional dependency may fail; skip it rather # than fail the whole discovery (the class cannot be built without the dep). @@ -283,14 +281,72 @@ def _discover(self) -> None: getattr(package, exported_name) except Exception as exc: logger.debug(f"Skipping lazily-exported '{exported_name}': {exc}") + self._register_subclasses_in_package(package_name=package_name) + + def register_external_subclasses(self, *, package_name: str) -> int: + """ + Register concrete base-type subclasses that live under an external package. + + This is how a loaded plug-in's components (e.g. scenarios shipped in a plug-in + wheel) enter the registry: the plug-in's modules are imported elsewhere, then + this enumerates the base type's in-memory subclasses under ``package_name`` and + registers each, alongside the built-in discovery package. Built-in discovery is + **not** triggered, so lazy discovery is preserved, and classes already registered + (e.g. by the plug-in's own bootstrap) are left untouched. + + Args: + package_name (str): The external (plug-in) top-level package name. + + Returns: + int: The number of classes newly registered from the package. + """ + return self._register_subclasses_in_package( + package_name=package_name, skip_registered_classes=True, external_package=package_name + ) + + def _register_subclasses_in_package( + self, *, package_name: str, skip_registered_classes: bool = False, external_package: str | None = None + ) -> int: + """ + Register every concrete base-type subclass whose module lives under a package. + + Shared by built-in discovery (``_discover``) and external plug-in + registration (``register_external_subclasses``). Enumerates the in-memory + subclasses of ``_base_type()``, keeps those whose module is ``package_name`` or a + submodule of it, and registers each. Built-in classes are named by + ``_get_registry_name``; when ``external_package`` is set, names come from + ``_external_registry_name`` so a plug-in's classes are keyed by their module path + within the plug-in (mirroring built-ins), which keeps same-named classes in + different submodules from colliding. When ``skip_registered_classes`` is set, + classes already in the catalog (by identity) are skipped so an explicit bootstrap + registration is never shadowed by a fallback-named duplicate. + + Args: + package_name (str): The package whose subclasses to register. + skip_registered_classes (bool): Skip classes already registered by identity. + external_package (str | None): When set, the plug-in package the classes belong + to; names are derived relative to it via ``_external_registry_name``. + + Returns: + int: The number of classes newly registered. + """ + base = self._base_type() + package_prefix = f"{package_name}." + already_registered = set(self._classes.values()) if skip_registered_classes else set() + count = 0 for cls in self._iter_concrete_subclasses(base): module = cls.__module__ or "" if module != package_name and not module.startswith(package_prefix): continue + if skip_registered_classes and cls in already_registered: + continue if (cls.__doc__ or "").strip().startswith("Deprecated alias"): logger.debug(f"Skipping deprecated alias: {cls.__name__}") continue - name = self._get_registry_name(cls) + if external_package is not None: + name = self._external_registry_name(cls, package_name=external_package) + else: + name = self._get_registry_name(cls) existing = self._classes.get(name) if existing is not None and existing is not cls: logger.warning( @@ -299,7 +355,26 @@ def _discover(self) -> None: ) continue self.register_class(cls, name=name) + count += 1 logger.debug(f"Registered {base.__name__} class: {name} ({cls.__name__})") + return count + + def _external_registry_name(self, cls: type[T], *, package_name: str) -> str: + """ + Return the catalog name for a class discovered in an external (plug-in) package. + + Defaults to the same name built-in discovery would use. Registries that key + built-ins by module path (e.g. scenarios) override this so plug-in classes get an + equally path-based, collision-resistant name relative to the plug-in package. + + Args: + cls (type[T]): The external class being registered. + package_name (str): The plug-in's top-level package name. + + Returns: + str: The catalog name to register the class under. + """ + return self._get_registry_name(cls) @staticmethod def _iter_concrete_subclasses(base: type[T]) -> list[type[T]]: diff --git a/pyrit/setup/__init__.py b/pyrit/setup/__init__.py index 4cac6e1470..1f16003a42 100644 --- a/pyrit/setup/__init__.py +++ b/pyrit/setup/__init__.py @@ -11,6 +11,7 @@ MemoryDatabaseType, initialize_pyrit_async, ) +from pyrit.setup.plugin_loader import PluginSpec __all__ = [ "AZURE_SQL", @@ -20,4 +21,5 @@ "initialize_from_config_async", "MemoryDatabaseType", "ConfigurationLoader", + "PluginSpec", ] diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 416421cb91..6870def2fa 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -24,6 +24,7 @@ ) if TYPE_CHECKING: + from pyrit.setup.plugin_loader import PluginSpec from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -106,6 +107,11 @@ class ConfigurationLoader(YamlLoadable): silent: Whether to suppress initialization messages. operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. + plugins: List of plug-ins to load as the guaranteed-first initialization phase. Each + entry is a wheel path string, a ``{package: wheel}`` mapping, or a + ``{wheel: ..., package: ...}`` mapping. Empty means no plug-ins. + plugin_accept_load_failures: When True, a plug-in that fails to load is skipped with a + warning instead of raising. None defers to ``PLUGIN_ACCEPT_LOAD_FAILURES`` (else fail-closed). Example YAML configuration: memory_db_type: sqlite @@ -149,10 +155,15 @@ class ConfigurationLoader(YamlLoadable): max_concurrent_scenario_runs: int = 3 allow_custom_initializers: bool = False server: dict[str, Any] | None = None + plugins: list[str | dict[str, Any]] = field(default_factory=list) + plugin_accept_load_failures: bool | None = None extensions: dict[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: """Validate and normalize the configuration after loading.""" + # Normalize plug-ins first: they load as the guaranteed-first phase of + # initialization, so a malformed plug-in entry should fail fast before anything else. + self._normalize_plugins() self._normalize_memory_db_type() self._normalize_initializers() self._normalize_scenario() @@ -216,6 +227,22 @@ def _normalize_initializers(self) -> None: raise ValueError(f"Initializer entry must be a string or dict, got: {type(entry).__name__}") self._initializer_configs = normalized + def _normalize_plugins(self) -> None: + """ + Normalize ``plugins`` entries to ``PluginSpec`` instances. + + Each entry is a wheel path string, a single-key ``{package: wheel}`` mapping, or an + explicit ``{wheel: ..., package: ...}`` mapping. Validating here (before any other + normalization) fails fast on a malformed plug-in entry, since plug-ins load as the + guaranteed-first phase of initialization. + + Raises: + ValueError: If a plug-in entry has an unsupported shape. + """ + from pyrit.setup.plugin_loader import PluginSpec + + self._plugin_specs = [PluginSpec.from_config(entry) for entry in self.plugins] + def _normalize_scenario(self) -> None: """ Normalize the optional ``scenario`` block to a ``ScenarioConfig``. @@ -362,6 +389,8 @@ def load_with_overrides( "env_files": None, # None = use defaults "env_akv_ref": None, "silent": False, + "plugins": [], + "plugin_accept_load_failures": None, } # 1. Try loading default config file if it exists @@ -381,6 +410,8 @@ def load_with_overrides( config_data["env_files"] = default_config.env_files config_data["env_akv_ref"] = default_config.env_akv_ref config_data["silent"] = default_config.silent + config_data["plugins"] = list(default_config.plugins) + config_data["plugin_accept_load_failures"] = default_config.plugin_accept_load_failures if default_config.operator: config_data["operator"] = default_config.operator if default_config.operation: @@ -407,6 +438,8 @@ def load_with_overrides( config_data["env_files"] = explicit_config.env_files config_data["env_akv_ref"] = explicit_config.env_akv_ref config_data["silent"] = explicit_config.silent + config_data["plugins"] = list(explicit_config.plugins) + config_data["plugin_accept_load_failures"] = explicit_config.plugin_accept_load_failures if explicit_config.operator: config_data["operator"] = explicit_config.operator if explicit_config.operation: @@ -491,6 +524,15 @@ def resolve_initializers(self) -> Sequence["PyRITInitializer"]: return resolved + def resolve_plugins(self) -> list["PluginSpec"]: + """ + Resolve the configured ``plugins`` entries to ``PluginSpec`` instances. + + Returns: + list[PluginSpec]: The plug-ins to load, in configured order (empty if none). + """ + return list(self._plugin_specs) + def resolve_initialization_scripts(self) -> Sequence[pathlib.Path] | None: """ Resolve initialization script paths. @@ -563,6 +605,7 @@ async def initialize_pyrit_async(self) -> None: resolved_initializers = self.resolve_initializers() resolved_scripts = self.resolve_initialization_scripts() resolved_env_files = self.resolve_env_files() + resolved_plugins = self.resolve_plugins() # Map snake_case memory_db_type to internal constant internal_memory_db_type = self._MEMORY_DB_TYPE_MAP[self.memory_db_type] @@ -574,6 +617,8 @@ async def initialize_pyrit_async(self) -> None: env_files=resolved_env_files, env_akv_ref=self.env_akv_ref, silent=self.silent, + plugins=resolved_plugins if resolved_plugins else None, + plugin_accept_load_failures=self.plugin_accept_load_failures, ) diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 43e7bfd4e0..efb711316f 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -13,6 +13,7 @@ from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory if TYPE_CHECKING: + from pyrit.setup.plugin_loader import PluginSpec from pyrit.setup.pyrit_initializer import PyRITInitializer logger = logging.getLogger(__name__) @@ -195,6 +196,30 @@ async def _execute_initializers_async(*, initializers: Sequence["PyRITInitialize raise +def _verify_plugin_prerequisites() -> None: + """ + Verify the initialization phases plug-ins depend on completed successfully. + + Plug-in loading runs a third party's bootstrap, which typically constructs targets and + scorers (needing loaded environment variables) and reads or writes central memory. This + is the health gate for that phase: plug-ins are loaded only if the prior initialization + phases were green. Central memory being set is the concrete, verifiable signal that the + environment and memory phases completed — the linear init flow would have raised earlier + otherwise, so a violation here means a broken initialization state, not a plug-in fault, + and is surfaced loudly regardless of ``plugin_accept_load_failures``. + + Raises: + RuntimeError: If central memory is not set, indicating a prior phase did not complete. + """ + try: + CentralMemory.get_memory_instance() + except Exception as exc: + raise RuntimeError( + "Cannot load plug-ins: central memory is not initialized. Plug-in loading requires " + "the environment and memory phases of initialize_pyrit_async to have completed first." + ) from exc + + async def initialize_pyrit_async( memory_db_type: MemoryDatabaseType | str, *, @@ -203,6 +228,8 @@ async def initialize_pyrit_async( env_files: Sequence[pathlib.Path] | None = None, env_akv_ref: Sequence[str] | None = None, silent: bool = False, + plugins: Sequence["PluginSpec"] | None = None, + plugin_accept_load_failures: bool | None = None, **memory_instance_kwargs: Any, ) -> None: """ @@ -224,6 +251,15 @@ async def initialize_pyrit_async( so local files take precedence over AKV. Requires ``azure-keyvault-secrets``. silent (bool): If True, suppresses print statements about environment file loading and schema migration. Defaults to False. + plugins (Sequence[PluginSpec] | None): Optional plug-ins to load as a guaranteed-first phase, + after central memory is set and before the configured initializers run. Each plug-in is a + pre-built wheel shipping datasets/scenarios that register like built-ins. Plug-ins are + loaded in list order, so one plug-in behaves identically to several. Typically populated + from the ``plugins`` key of ``.pyrit_conf``; direct callers pass ``PluginSpec`` instances. + plugin_accept_load_failures (bool | None): Overrides ``PLUGIN_ACCEPT_LOAD_FAILURES`` for plug-in + loading. When True, a plug-in that fails to load is skipped with a warning instead of + raising. Defaults to None (use the ``PLUGIN_ACCEPT_LOAD_FAILURES`` environment variable, + else fail-closed). **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. Raises: @@ -259,6 +295,20 @@ async def initialize_pyrit_async( CentralMemory.set_memory_instance(memory) + # Load configured plug-ins as a guaranteed-first phase: after central memory is set (a + # plug-in bootstrap may use it) and BEFORE any configured initializer runs, so plug-in + # datasets/scenarios are registered before LoadDefaultDatasets and PreloadScenarioMetadata + # read the registry. Because `plugins` is its own always-first phase (not one of the + # ordered `initializers`), this ordering is true by construction. Plug-ins are loaded only + # after a health check confirms the initialization phases they depend on completed. No-op + # when no plug-ins are configured. + if plugins: + _verify_plugin_prerequisites() + + from pyrit.setup.plugin_loader import load_plugins_if_configured_async + + await load_plugins_if_configured_async(plugins=plugins, accept_load_failures=plugin_accept_load_failures) + # Combine directly provided initializers with those loaded from scripts all_initializers = list(initializers) if initializers else [] diff --git a/pyrit/setup/plugin_compat.py b/pyrit/setup/plugin_compat.py new file mode 100644 index 0000000000..e418a8b45b --- /dev/null +++ b/pyrit/setup/plugin_compat.py @@ -0,0 +1,255 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Best-effort compatibility shims for plug-ins built against a slightly different PyRIT. + +A plug-in wheel is authored against one PyRIT version, then loaded into whatever PyRIT the +operator is running. When the host API has drifted, small mechanical differences — most +commonly a renamed extension-point method — can leave the plug-in's scenarios abstract and +therefore undiscoverable, even though the plug-in's own logic is unchanged. Rather than +force every plug-in author to re-release on each host bump, the loader applies narrow, loud +heuristics that bridge known mechanical renames so minor drift does not block loading. + +The shims are deliberately conservative. A scenario is bridged only when it is abstract +*solely* because of a recognized rename and a usable predecessor method is present; a class +abstract for any other reason is left alone (and fails loudly downstream). Every bridge and +every detected version mismatch logs a loud, greppable warning so the drift is visible and +the operator knows that rebuilding the plug-in against the running PyRIT removes the shim. +This module owns drift-bridging so the loader stays focused on extract/import/register. +""" + +from __future__ import annotations + +import inspect +import logging +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine + from pathlib import Path + from typing import Any + +logger = logging.getLogger(__name__) + +# Renamed scenario extension points: each maps a current (host) abstract method to the +# older predecessor a plug-in may still define. A predecessor took no ``context`` (it read +# ``self._*`` state the base still populates before the build call), so the bridge can +# delegate to it and ignore ``context``. Add an entry here when a future rename needs the +# same mechanical bridge. +_SCENARIO_METHOD_RENAMES: dict[str, str] = { + "_build_atomic_attacks_async": "_get_atomic_attacks_async", +} + + +def warn_on_version_drift(*, extract_dir: Path) -> str | None: + """ + Log a loud warning when a plug-in was built against a different PyRIT minor version. + + Reads the plug-in's ``Requires-Dist: pyrit`` pin from its wheel ``METADATA`` and + compares it to the running ``pyrit`` version. A mismatch is surfaced but never fatal: + the loader tolerates slight drift and lets the import/registration path (plus the + scenario shims below) decide whether the plug-in actually works. Absence of a pin, or + an unparseable one, is silently ignored — this is a best-effort signal, not a gate. + + Args: + extract_dir: The directory the plug-in wheel was extracted to. + + Returns: + str | None: The plug-in's declared PyRIT pin when a drift warning was emitted, + else ``None``. + """ + declared = _read_required_pyrit_version(extract_dir=extract_dir) + if declared is None: + return None + + import pyrit + + running = getattr(pyrit, "__version__", "") or "" + if _same_minor(declared, running): + return None + + logger.warning( + "PLUGIN VERSION DRIFT: plug-in was built against pyrit %s but pyrit %s is running. " + "Proceeding; compatibility shims will bridge known mechanical differences, but " + "rebuild the plug-in against the running pyrit if scenarios fail to load.", + declared, + running or "(unknown)", + ) + return declared + + +def bridge_scenario_extension_points(*, package_name: str) -> list[str]: + """ + Make a plug-in's scenarios concrete by bridging renamed extension-point methods. + + Enumerates concrete-intent ``Scenario`` subclasses owned by ``package_name`` that are + still abstract, and for each one whose only unimplemented abstract methods are known + renames with a usable predecessor, injects a thin adapter so the class satisfies the + current contract. Classes abstract for any other reason are left untouched. + + Args: + package_name: The plug-in's top-level package name. + + Returns: + list[str]: Human-readable descriptions of the bridges applied, for logging. + """ + from pyrit.scenario.core import Scenario + + prefix = f"{package_name}." + applied: list[str] = [] + + for cls in _iter_owned_subclasses(base=Scenario, package_name=package_name, prefix=prefix): + if not inspect.isabstract(cls): + continue + + missing = set(cls.__abstractmethods__) # type: ignore[ty:unresolved-attribute] + # Only bridge when every missing method is a known rename we can satisfy from a + # predecessor the class actually provides. Otherwise the class is abstract for a + # reason we must not paper over, so leave it (it will fail loudly downstream). + if not missing or not missing.issubset(_SCENARIO_METHOD_RENAMES.keys()): + continue + if not all(_has_concrete_method(cls, _SCENARIO_METHOD_RENAMES[name]) for name in missing): + continue + + for new_name in missing: + predecessor = _SCENARIO_METHOD_RENAMES[new_name] + setattr(cls, new_name, _make_rename_bridge(predecessor_name=predecessor)) + applied.append(f"{cls.__name__}.{predecessor} -> {new_name}") + logger.warning( + "PLUGIN COMPAT SHIM: bridged %s.%s to the renamed %s (plug-in built against " + "an older pyrit). Rebuild the plug-in against the running pyrit to remove this shim.", + cls.__name__, + predecessor, + new_name, + ) + cls.__abstractmethods__ = frozenset(cls.__abstractmethods__ - missing) # type: ignore[ty:unresolved-attribute] + + return applied + + +def _make_rename_bridge(*, predecessor_name: str) -> Callable[..., Coroutine[Any, Any, Any]]: + """ + Build an adapter that satisfies a renamed async extension point via its predecessor. + + Args: + predecessor_name: The older method name the plug-in still defines. + + Returns: + Callable[..., Coroutine[Any, Any, Any]]: An async adapter that ignores the new + ``context`` keyword and delegates to the predecessor (which reads ``self._*`` + state the base populates before the build call). + """ + + async def _bridged_build_async(self, **_kwargs: Any) -> Any: # noqa: ANN001 + return await getattr(self, predecessor_name)() + + _bridged_build_async.__name__ = predecessor_name + _bridged_build_async.__qualname__ = predecessor_name + return _bridged_build_async + + +def _has_concrete_method(cls: type, name: str) -> bool: + """ + Return whether ``cls`` provides a concrete (non-abstract) method of the given name. + + Args: + cls: The class to inspect. + name: The method name to look for. + + Returns: + bool: True when the method exists and is not itself abstract. + """ + method = getattr(cls, name, None) + return callable(method) and name not in getattr(cls, "__abstractmethods__", frozenset()) + + +def _iter_owned_subclasses(*, base: type, package_name: str, prefix: str) -> list[type]: + """ + Return every subclass of ``base`` whose module is owned by the plug-in package. + + Args: + base: The base class to enumerate subclasses of. + package_name: The plug-in's top-level package name. + prefix: ``f"{package_name}."`` (passed in to avoid recomputation). + + Returns: + list[type]: The owned subclasses currently loaded in memory. + """ + seen: set[int] = set() + owned: list[type] = [] + stack = list(base.__subclasses__()) + while stack: + cls = stack.pop() + if id(cls) in seen: + continue + seen.add(id(cls)) + stack.extend(cls.__subclasses__()) + module = cls.__module__ or "" + if module == package_name or module.startswith(prefix): + owned.append(cls) + return owned + + +def _read_required_pyrit_version(*, extract_dir: Path) -> str | None: + """ + Return the pinned ``pyrit`` version from the plug-in's wheel ``METADATA``, if any. + + Args: + extract_dir: The directory the plug-in wheel was extracted to. + + Returns: + str | None: The pinned version string (e.g. ``"0.14.0"``) or ``None`` when no + parseable ``Requires-Dist: pyrit`` pin is present. + """ + import re + + for metadata in extract_dir.glob("*.dist-info/METADATA"): + try: + text = metadata.read_text(encoding="utf-8", errors="replace") + except OSError: + continue + for line in text.splitlines(): + if not line.lower().startswith("requires-dist:"): + continue + if not re.match(r"requires-dist:\s*pyrit\b", line, flags=re.IGNORECASE): + continue + match = re.search(r"==\s*([0-9][0-9A-Za-z.\-]*)", line) + if match: + return match.group(1) + return None + + +def _same_minor(left: str, right: str) -> bool: + """ + Return whether two version strings share the same ``major.minor`` prefix. + + Args: + left: A version string (e.g. ``"0.14.0"``). + right: A version string (e.g. ``"0.15.0.dev0"``). + + Returns: + bool: True when both parse to the same ``(major, minor)``; False otherwise. An + unparseable side compares unequal so drift is surfaced rather than hidden. + """ + return _major_minor(left) == _major_minor(right) and _major_minor(left) is not None + + +def _major_minor(version: str) -> tuple[int, int] | None: + """ + Parse the leading ``major.minor`` from a version string. + + Args: + version: A version string. + + Returns: + tuple[int, int] | None: The ``(major, minor)`` pair, or ``None`` when the string + does not begin with two integer components. + """ + parts = version.split(".") + if len(parts) < 2: + return None + try: + return int(parts[0]), int(parts[1]) + except ValueError: + return None diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py new file mode 100644 index 0000000000..61866b97d0 --- /dev/null +++ b/pyrit/setup/plugin_loader.py @@ -0,0 +1,831 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Load non-disclosable PyRIT plug-ins from pre-built wheels at initialization time. + +A plug-in is a pure-Python wheel that ships dataset providers and/or scenarios that +must not live in the public PyRIT repo. The loader **extracts** the wheel (stdlib +``zipfile`` — never ``pip``/``.venv``) into ``.plugin//``, prepends that directory +to ``sys.path``, imports the package (so ``SeedDatasetProvider`` subclasses self-register), +and runs the plug-in's bootstrap (a top-level ``register()`` callable or a shipped +``PyRITInitializer`` subclass). The plug-in's own ``Scenario`` subclasses are then +auto-registered by type (scoped to the plug-in package), so a plug-in's scenarios are +discovered without the bootstrap having to register each one explicitly. + +Plug-ins are declared in ``.pyrit_conf`` under the dedicated ``plugins`` key (a list, so +several plug-ins load identically to one). ``load_plugins_if_configured_async`` is invoked +as a guaranteed-first phase inside ``initialize_pyrit_async`` — after central memory is set +and **before** any configured initializer runs — so plug-in datasets and scenarios are +registered before ``LoadDefaultDatasets`` / ``PreloadScenarioMetadata`` read the registry. +Because ``plugins`` is its own always-first phase (not one of the ordered ``initializers``), +this ordering is true by construction. It is a no-op when no plug-ins are configured. +""" + +from __future__ import annotations + +import asyncio +import importlib +import inspect +import logging +import os +import pkgutil +import shutil +import sys +import tempfile +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from pyrit.setup.pyrit_initializer import PyRITInitializer + +if TYPE_CHECKING: + from collections.abc import Sequence + from types import ModuleType + + from pyrit.registry import ScenarioRegistry + +logger = logging.getLogger(__name__) + +_TRUE_TOKENS = frozenset({"1", "true", "yes", "on"}) + + +@dataclass(frozen=True) +class PluginSpec: + """ + A single plug-in to load: a pre-built wheel plus its optional top-level package name. + + Attributes: + wheel (Path): Filesystem path to the pre-built plug-in wheel (``.whl``). + package (str | None): The wheel's top-level import package. When ``None`` it is + auto-detected from the wheel; set it to disambiguate a multi-package wheel. + """ + + wheel: Path + package: str | None = None + + @classmethod + def from_config(cls, entry: str | Mapping[str, object]) -> PluginSpec: + """ + Build a ``PluginSpec`` from a ``.pyrit_conf`` ``plugins`` list entry. + + Accepted shapes: + + * A bare wheel path string (package auto-detected):: + + - /abs/path/to/plugin.whl + + * A single-key ``{package: wheel}`` mapping (concise, names the package):: + + - my_plugin: /abs/path/to/plugin.whl + + * An explicit ``{wheel: ..., package: ...}`` mapping (``package`` optional):: + + - wheel: /abs/path/to/plugin.whl + package: my_plugin + + Args: + entry: One ``plugins`` list item from the parsed configuration. + + Returns: + PluginSpec: The normalized spec. + + Raises: + ValueError: If the entry shape is not one of the accepted forms. + """ + if isinstance(entry, str): + return cls(wheel=Path(entry).expanduser()) + + if isinstance(entry, Mapping): + mapping = dict(entry) + if "wheel" in mapping: + extra_keys = set(mapping) - {"wheel", "package"} + if extra_keys: + raise ValueError( + f"Plug-in mapping has unexpected key(s) {sorted(extra_keys)}; only 'wheel' and " + f"'package' are allowed. Got: {entry}" + ) + wheel = mapping["wheel"] + package = mapping.get("package") + if not isinstance(wheel, str) or (package is not None and not isinstance(package, str)): + raise ValueError(f"Plug-in entry 'wheel'/'package' must be strings. Got: {entry}") + return cls(wheel=Path(wheel).expanduser(), package=package) + if len(mapping) == 1 and "package" not in mapping: + ((package, wheel),) = mapping.items() + if not isinstance(package, str) or not isinstance(wheel, str): + raise ValueError(f"Plug-in entry must map a package name to a wheel path. Got: {entry}") + return cls(wheel=Path(wheel).expanduser(), package=package) + raise ValueError( + f"Plug-in mapping must be a single {{package_name: wheel}} pair or carry a 'wheel' key. Got: {entry}" + ) + + raise ValueError(f"Plug-in entry must be a wheel path string or mapping, got: {type(entry).__name__}") + + +async def load_plugins_if_configured_async( + *, plugins: Sequence[PluginSpec], accept_load_failures: bool | None = None +) -> None: + """ + Load every configured plug-in, in order. A no-op when ``plugins`` is empty. + + Convenience entry point invoked by ``initialize_pyrit_async`` after memory is set and + before the configured initializers run. Plug-ins are loaded in list order, so one + plug-in behaves identically to several — the single-plug-in case is just a + one-element list. + + Args: + plugins: The plug-in specs to load, in order. + accept_load_failures: If provided, overrides the ``PLUGIN_ACCEPT_LOAD_FAILURES`` + environment variable. When True, a plug-in that fails to load is skipped with + a warning and the next plug-in still loads. + + Raises: + PluginLoadError: If a plug-in fails to load and load failures are not accepted. + """ + if not plugins: + logger.debug("No plug-ins configured; plug-in loading is a no-op.") + return + + for spec in plugins: + await PluginLoader(spec=spec, accept_load_failures=accept_load_failures).load_async() + + +def _name_owned_by(module_name: str, package_name: str) -> bool: + """ + Return whether a module name belongs to the given plug-in package. + + Args: + module_name: A dotted module name. + package_name: The plug-in's top-level package name. + + Returns: + bool: True if ``module_name`` is the package or one of its submodules. + """ + return module_name == package_name or module_name.startswith(f"{package_name}.") + + +def _module_owned_by(cls: type, package_name: str) -> bool: + """ + Return whether ``cls`` is defined within the given plug-in package. + + Args: + cls: The class to check. + package_name: The plug-in's top-level package name. + + Returns: + bool: True if the class's module is the package or one of its submodules. + """ + return _name_owned_by(cls.__module__ or "", package_name) + + +_REMEDIATION = ( + "Remove the plug-in configuration, or accept load failures (PLUGIN_ACCEPT_LOAD_FAILURES=true, or the " + "initialize_pyrit_async(plugin_accept_load_failures=True) parameter) to continue without it." +) + + +class PluginLoadError(RuntimeError): + """Base error raised when a configured plug-in fails to load and load failures are not accepted.""" + + +class PluginWheelNotFoundError(PluginLoadError): + """The configured plug-in wheel path does not point to a readable ``.whl`` file.""" + + +class PluginImportError(PluginLoadError): + """The plug-in package could not be imported (missing dependency, or an installed package shadows it).""" + + +class PluginRegisteredNothingError(PluginLoadError): + """The plug-in imported cleanly but registered no datasets or scenarios.""" + + +class PluginLoader: + """ + Extract and register a single PyRIT plug-in wheel described by a ``PluginSpec``. + + The wheel is extracted to ``.plugin//`` (never installed), imported, and its + bootstrap is run so its datasets and scenarios register like built-ins. Fails closed + by default; set ``accept_load_failures`` (constructor / ``initialize_pyrit_async`` + param) or ``PLUGIN_ACCEPT_LOAD_FAILURES`` to continue without the plug-in when it + cannot be loaded. + """ + + _FINGERPRINT_FILE = ".plugin_wheel_fingerprint" + + def __init__(self, *, spec: PluginSpec, accept_load_failures: bool | None = None) -> None: + """ + Initialize the loader for one plug-in. + + Args: + spec: The plug-in to load (wheel path plus optional package name). + accept_load_failures: If provided, overrides ``PLUGIN_ACCEPT_LOAD_FAILURES``. + When True, a plug-in that fails to load is skipped with a warning instead + of raising. + """ + self._spec = spec + self._explicit_accept_load_failures = accept_load_failures + + async def load_async(self) -> None: + """ + Load the plug-in described by this loader's ``PluginSpec``. + + Raises: + PluginLoadError: If the plug-in fails to load and load failures are not accepted. + """ + wheel = self._spec.wheel + accept_load_failures = self._resolve_accept_load_failures() + + try: + await self._load_plugin_async(wheel_path=wheel.expanduser()) + except Exception as exc: + if accept_load_failures: + logger.warning( + "Plug-in wheel '%s' failed to load; load failures are accepted so continuing without it: %s", + wheel, + exc, + ) + return + message = f"Failed to load plug-in from wheel '{wheel}': {exc} {_REMEDIATION}" + # Preserve the specific failure type so callers can distinguish modes, while + # always surfacing the remediation guidance. Unknown errors (e.g. a raising + # 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: + """ + Extract, import, bootstrap, and verify a single plug-in wheel. + + Global state (``sys.path``, imported plug-in modules, and the provider/scenario + registries) is rolled back if the load fails, so a failed or accepted-failure load + leaves no partial trace. + + Args: + wheel_path: Path to the pre-built plug-in wheel on disk. + + Raises: + 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 PluginWheelNotFoundError(f"Plug-in wheel does not point to an existing file: {wheel_path}") + if wheel_path.suffix != ".whl": + raise PluginWheelNotFoundError(f"Plug-in wheel must be a .whl file, got: {wheel_path}") + + # Wheel extraction and directory scanning are blocking filesystem work; run them + # off the event loop so init does not stall unrelated async tasks. + extract_dir = await asyncio.to_thread(self._extract_wheel, wheel_path=wheel_path) + package_name = await asyncio.to_thread( + self._resolve_package_name, extract_dir=extract_dir, explicit_package=self._spec.package + ) + + from pyrit.setup.plugin_compat import bridge_scenario_extension_points, warn_on_version_drift + + # Surface any host/plug-in version drift before importing, so a subsequent import or + # registration failure is read in the light of the mismatch (tolerate slight drift). + warn_on_version_drift(extract_dir=extract_dir) + + from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider + from pyrit.registry import ScenarioRegistry + + scenario_registry = ScenarioRegistry.get_registry_singleton() + provider_snapshot = dict(SeedDatasetProvider._registry) + scenario_snapshot = dict(scenario_registry._classes) + modules_snapshot = {name for name in sys.modules if _name_owned_by(name, package_name)} + syspath_entry = str(extract_dir) + added_to_syspath = syspath_entry not in sys.path + if added_to_syspath: + sys.path.insert(0, syspath_entry) + + try: + logger.info("Importing plug-in package '%s'", package_name) + try: + module = importlib.import_module(package_name) + self._verify_module_location(module=module, extract_dir=extract_dir, package_name=package_name) + self._import_submodules(module=module, package_name=package_name) + except Exception as exc: + raise PluginImportError(f"Could not import plug-in package '{package_name}': {exc}") from exc + + # Bridge known renamed extension points so scenarios built against an older + # PyRIT are made concrete (and thus discoverable) before bootstrap/registration. + bridged = bridge_scenario_extension_points(package_name=package_name) + if bridged: + logger.warning( + "Applied %d plug-in compatibility shim(s) for '%s': %s", + len(bridged), + package_name, + ", ".join(bridged), + ) + + await self._run_bootstrap_async(package_name=package_name, module=module) + + # Register the plug-in's own Scenario subclasses the same way built-ins are + # discovered (by type, scoped to the plug-in package), so plug-in scenarios are + # picked up without the bootstrap having to register each one explicitly. Classes + # the bootstrap already registered are left untouched. + registered_scenarios = scenario_registry.register_external_subclasses(package_name=package_name) + if registered_scenarios: + logger.info("Auto-registered %d plug-in scenario(s) from '%s'.", registered_scenarios, package_name) + + provider_count, scenario_count = self._count_registered( + package_name=package_name, scenario_registry=scenario_registry + ) + if not provider_count and not scenario_count: + raise PluginRegisteredNothingError( + f"Plug-in package '{package_name}' imported successfully but registered no datasets or " + "scenarios. The wheel is likely mis-packaged (imports cleanly yet loads nothing)." + ) + + dataset_collisions = self._warn_on_dataset_name_collisions(package_name=package_name) + except Exception: + self._rollback( + package_name=package_name, + syspath_entry=syspath_entry if added_to_syspath else None, + modules_snapshot=modules_snapshot, + provider_snapshot=provider_snapshot, + scenario_registry=scenario_registry, + scenario_snapshot=scenario_snapshot, + ) + raise + + logger.info( + "Loaded plug-in '%s': %d dataset provider(s), %d scenario(s) registered; %d dataset name collision(s)%s.", + package_name, + provider_count, + scenario_count, + len(dataset_collisions), + " (see PLUGIN DATASET SHADOWED warnings above)" if dataset_collisions else "", + ) + + def _extract_wheel(self, *, wheel_path: Path) -> Path: + """ + Extract the wheel into ``.plugin//``, reusing a fresh cached extraction. + + A cached extraction is reused only when its recorded fingerprint (wheel size and + modification time) matches the wheel on disk. A wheel rebuilt at the same path + (common during plug-in development, where the version-stamped filename is stable) + has a newer fingerprint, so the stale extraction is discarded and re-extracted + rather than silently reused — the exact silent-staleness trap the plug-in design + guards against. Extraction is atomic: the wheel is unpacked into a temporary + sibling directory and moved into place only on success, so a crash mid-extraction + never leaves a partial tree that would later be treated as a valid cache. + ``safe_extract_zip`` validates every member first (path traversal, symlinks, and + size / entry-count / compression caps) so a tampered wheel cannot escape the + extraction directory or exhaust disk. + + Args: + wheel_path: Path to the plug-in wheel. + + Returns: + Path: The directory the wheel was extracted to. + """ + from pyrit.common.safe_extract import safe_extract_zip + + base_dir = self._plugin_base_dir() + base_dir.mkdir(parents=True, exist_ok=True) + + extract_dir = base_dir / wheel_path.stem + fingerprint = self._wheel_fingerprint(wheel_path=wheel_path) + if extract_dir.is_dir() and any(extract_dir.iterdir()): + if self._cached_fingerprint(extract_dir=extract_dir) == fingerprint: + logger.info("Reusing cached plug-in extraction at %s", extract_dir) + return extract_dir + logger.warning( + "PLUGIN CACHE STALE: extraction at %s does not match the current wheel " + "(rebuilt or replaced); re-extracting '%s'.", + extract_dir, + wheel_path.name, + ) + + # Unique per-extraction temp dir so concurrent loads of the same wheel in one + # process cannot collide on a shared path (mkdtemp is atomic and 0700). safe_extract + # into it, then atomically move into place. + tmp_dir = Path(tempfile.mkdtemp(prefix=f".{wheel_path.stem}.tmp-", dir=base_dir)) + try: + safe_extract_zip(source=wheel_path, dest_dir=tmp_dir) + (tmp_dir / self._FINGERPRINT_FILE).write_text(fingerprint, encoding="utf-8") + if extract_dir.exists(): + shutil.rmtree(extract_dir) + os.replace(tmp_dir, extract_dir) + finally: + if tmp_dir.exists(): + shutil.rmtree(tmp_dir, ignore_errors=True) + + logger.info("Extracted plug-in wheel '%s' to %s", wheel_path.name, extract_dir) + return extract_dir + + @staticmethod + def _wheel_fingerprint(*, wheel_path: Path) -> str: + """ + Return a cheap change-detecting fingerprint (size and mtime) for a wheel. + + Args: + wheel_path: Path to the plug-in wheel. + + Returns: + str: A ``":"`` fingerprint that changes whenever the wheel is + rebuilt or replaced. + """ + stat = wheel_path.stat() + return f"{stat.st_size}:{stat.st_mtime_ns}" + + @classmethod + def _cached_fingerprint(cls, *, extract_dir: Path) -> str | None: + """ + Return the fingerprint recorded for a cached extraction, or ``None`` if absent. + + A missing marker (e.g. an extraction from before fingerprinting) reads as ``None`` + so the cache is treated as stale and re-extracted rather than trusted blindly. + + Args: + extract_dir: The cached extraction directory. + + Returns: + str | None: The recorded fingerprint, or ``None`` when no valid marker exists. + """ + marker = extract_dir / cls._FINGERPRINT_FILE + try: + return marker.read_text(encoding="utf-8").strip() + except OSError: + return None + + @staticmethod + def _plugin_base_dir() -> Path: + """ + Resolve the base directory for plug-in extractions. + + Uses ``PLUGIN_DIR`` when set, otherwise ``/.plugin``. + + Returns: + Path: The resolved plug-in base directory. + """ + override = os.getenv("PLUGIN_DIR") + if override: + return Path(override).expanduser().resolve() + from pyrit.common import path + + return Path(path.HOME_PATH, ".plugin").resolve() + + @staticmethod + def _resolve_package_name(*, extract_dir: Path, explicit_package: str | None) -> str: + """ + Determine the plug-in's top-level import package. + + Resolution order: the spec's ``package`` (when set), then ``*.dist-info/top_level.txt``, + then the single importable top-level directory in the extraction. + + Args: + extract_dir: The directory the wheel was extracted to. + explicit_package: The package name from the plug-in's config, if any. + + Returns: + str: The top-level package name to import. + + Raises: + ValueError: If the package cannot be unambiguously determined. + """ + if explicit_package: + return explicit_package + + for dist_info in sorted(extract_dir.glob("*.dist-info")): + top_level = dist_info / "top_level.txt" + if top_level.is_file(): + for line in top_level.read_text(encoding="utf-8").splitlines(): + name = line.strip() + if name: + return name + + candidates = sorted( + child.name + for child in extract_dir.iterdir() + if child.is_dir() + and not child.name.endswith(".dist-info") + and not child.name.endswith(".data") + and (child / "__init__.py").is_file() + ) + if len(candidates) == 1: + return candidates[0] + if not candidates: + raise ValueError( + f"Could not find an importable top-level package in {extract_dir}. " + "Set the plug-in's 'package' in its .pyrit_conf entry." + ) + raise ValueError( + f"Found multiple top-level packages in {extract_dir}: {candidates}. " + "Set the plug-in's 'package' in its .pyrit_conf entry to disambiguate." + ) + + @staticmethod + def _verify_module_location(*, module: ModuleType, extract_dir: Path, package_name: str) -> None: + """ + Verify the imported package resolves inside the extraction directory. + + Guards against an installed package of the same name shadowing the extracted + plug-in — a silent failure where import succeeds but the wheel's code/data is + ignored. + + Args: + module: The imported plug-in package module. + extract_dir: The directory the wheel was extracted to. + package_name: The plug-in's top-level package name. + + Raises: + ValueError: If the imported package resolves outside ``extract_dir``. + """ + extract_resolved = extract_dir.resolve() + raw_locations = list(getattr(module, "__path__", []) or []) + module_file = getattr(module, "__file__", None) + if module_file: + raw_locations.append(module_file) + + locations = [Path(location).resolve() for location in raw_locations if location] + if not locations: + return + if not any(location.is_relative_to(extract_resolved) for location in locations): + raise ValueError( + f"Imported package '{package_name}' resolved to {locations[0]} which is outside the " + f"plug-in extraction directory {extract_resolved}. An installed package with the same " + "name is likely shadowing the plug-in; set the plug-in's 'package' in .pyrit_conf or " + "resolve the name conflict." + ) + + @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_ACCEPT_LOAD_FAILURES``. + + 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 the + provider/scenario registrations it added, and **restores any entries the plug-in + overwrote**. Both registries key by a name the plug-in does not control + (``SeedDatasetProvider`` by ``cls.__name__``; the scenario catalog by registry + name) and assign unconditionally, so a plug-in whose provider/scenario name + collides with an existing one silently replaces it; rollback must put the original + back, not just drop the new key. State present before the load — including modules + that already existed and built-ins discovered meanwhile — is preserved. + + Args: + package_name: The plug-in's top-level package name. + syspath_entry: The ``sys.path`` entry to remove, or None if it was already present. + modules_snapshot: Package-owned module names present before the load. + provider_snapshot: Provider registry contents captured before the load. + scenario_registry: The scenario registry singleton to clean up. + scenario_snapshot: Scenario catalog contents captured before the load. + """ + from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider + + if syspath_entry and syspath_entry in sys.path: + sys.path.remove(syspath_entry) + + for name in [m for m in sys.modules if _name_owned_by(m, package_name) and m not in modules_snapshot]: + del sys.modules[name] + + # For every entry the plug-in now owns: restore the pre-load value if the key + # existed before (the plug-in overwrote it), otherwise drop the key it added. + for key in list(SeedDatasetProvider._registry): + if _module_owned_by(SeedDatasetProvider._registry[key], package_name): + if key in provider_snapshot: + SeedDatasetProvider._registry[key] = provider_snapshot[key] # type: ignore[ty:invalid-assignment] + else: + del SeedDatasetProvider._registry[key] + + changed_scenarios = False + for name in list(scenario_registry._classes): + if _module_owned_by(scenario_registry._classes[name], package_name): + if name in scenario_snapshot: + scenario_registry._classes[name] = scenario_snapshot[name] # type: ignore[ty:invalid-assignment] + else: + del scenario_registry._classes[name] + changed_scenarios = True + if changed_scenarios: + scenario_registry._metadata_cache = None + + def _resolve_accept_load_failures(self) -> bool: + """ + Resolve the accept-load-failures setting from the explicit value or the environment. + + Precedence: the explicit ``accept_load_failures`` passed to the constructor (e.g. from + ``initialize_pyrit_async``), then the ``PLUGIN_ACCEPT_LOAD_FAILURES`` environment + variable, otherwise fail-closed. + + Returns: + bool: True if a failed plug-in load should be skipped with a warning. + """ + if self._explicit_accept_load_failures is not None: + return self._explicit_accept_load_failures + + env_value = os.getenv("PLUGIN_ACCEPT_LOAD_FAILURES") + if env_value is not None: + return self._coerce_bool(env_value) + + return False + + @staticmethod + def _coerce_bool(value: str) -> bool: + """ + Interpret a string as a boolean flag. + + Args: + value: The raw string value. + + Returns: + bool: True for common truthy tokens (1/true/yes/on, case-insensitive). + """ + return str(value).strip().lower() in _TRUE_TOKENS diff --git a/tests/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..9045d62616 --- /dev/null +++ b/tests/integration/setup/test_plugin_loader_integration.py @@ -0,0 +1,438 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Integration test: plug-in scenarios load, register, instantiate, and execute after init. + +This exercises the full public loading path end-to-end: point ``PLUGIN_WHEEL`` at a +built wheel, run ``initialize_pyrit_async``, and assert the wheel's scenarios are +registered in ``ScenarioRegistry``, construct cleanly, and drive through the public +execution pipeline. + +Two cases: + +* A self-contained case builds a small wheel at test time and verifies its scenarios + are discovered. This always runs and guards the mechanism in public CI. +* An injected case loads a wheel supplied out-of-band via environment variables and + verifies all of its scenarios are discovered, instantiate, and that at least one + executes. It is skipped unless those variables are set, so the committed test depends + on no specific external package. Point it at a real scenario wheel (for example in a + downstream/private CI job) to guarantee every scenario that wheel ships is picked up:: + + PLUGIN_TEST_WHEEL=/path/to/plugin.whl + PLUGIN_TEST_PACKAGE=the_plugin_package # enables package enumeration + instantiation + PLUGIN_TEST_SCENARIO_DIRS=/path/to/scenarios # optional; os.pathsep-separated + PLUGIN_TEST_EXEC_SCENARIO=the.registry.name # optional; enables the execution case + ADVERSARIAL_CHAT_ENDPOINT=... # execution target + scorer endpoint + ADVERSARIAL_CHAT_MODEL=... # optional model name +""" + +import inspect +import os +import sys +import textwrap +import uuid +import zipfile +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import pytest + +from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider +from pyrit.models import ScenarioResult, ScenarioRunState +from pyrit.prompt_target import OpenAIChatTarget +from pyrit.registry import ScenarioRegistry +from pyrit.registry.discovery import discover_in_directory +from pyrit.scenario.core import DatasetAttackConfiguration, Scenario +from pyrit.score import SelfAskTrueFalseScorer, TrueFalseQuestionPaths +from pyrit.setup import IN_MEMORY, PluginSpec, initialize_pyrit_async + +# (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"), +] + +# Substring of the ValueError ``Scenario.run_async`` raises when an atomic attack's +# objective did not complete (for example a target refusal). The attack still ran through +# the public pipeline, so this outcome proves execution while a drift regression -- which +# surfaces a different exception type before or during the run -- still fails loudly. +_OBJECTIVES_INCOMPLETE_MARKER = "objectives incomplete" + + +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_classes_under_package(package_prefix: str) -> dict[str, type[Scenario]]: + """ + Return registered concrete ``Scenario`` subclasses owned by a package, keyed by registry name. + + Uses the post-load registry (reliable after the plug-in has been imported and + bootstrapped), which sidesteps the standalone-import problems a filesystem walk can + hit for a package that uses relative imports. + + Args: + package_prefix: The plug-in's top-level package name. + + Returns: + dict[str, type[Scenario]]: Registry name -> scenario class for that package. + """ + prefix = f"{package_prefix}." + registry = ScenarioRegistry.get_registry_singleton() + found: dict[str, type[Scenario]] = {} + for name in registry.get_class_names(): + cls = registry.get_class(name) + module = cls.__module__ or "" + if not inspect.isabstract(cls) and (module == package_prefix or module.startswith(prefix)): + found[name] = cls + return found + + +async def _execute_scenario_async( + *, scenario_cls: type[Scenario], endpoint: str, model: str | None +) -> ScenarioResult | None: + """ + Drive one plug-in scenario through the public initialize/run pipeline. + + Constructs an objective target and scorer from the supplied endpoint, initializes the + scenario, and runs it. ``initialize_async`` must succeed and build at least one atomic + attack; that is where dataset-config resolution drift would surface. + + Args: + scenario_cls: The plug-in scenario class to execute. + endpoint: Chat endpoint backing both the objective target and the scorer. + model: Optional model name for the endpoint. + + Returns: + ScenarioResult | None: The completed result, or ``None`` when the run finished with + incomplete objectives (the attack ran end-to-end but the objective was not achieved). + """ + target = OpenAIChatTarget(endpoint=endpoint, model_name=model) + scorer = SelfAskTrueFalseScorer( + chat_target=target, + true_false_question_path=TrueFalseQuestionPaths.TASK_ACHIEVED_REFINED.value, + ) + + try: + scenario = scenario_cls(objective_scorer=scorer, fast_mode=True) # type: ignore[ty:missing-argument, ty:unknown-argument] + except TypeError: + scenario = scenario_cls(objective_scorer=scorer) # type: ignore[ty:missing-argument] + + dataset_config: DatasetAttackConfiguration | None = None + required_datasets = getattr(scenario_cls, "required_datasets", None) + if callable(required_datasets): + names = list(required_datasets()) + if names: + dataset_config = DatasetAttackConfiguration(dataset_names=names, max_dataset_size=1) + + args: dict[str, Any] = {"objective_target": target, "max_concurrency": 1} + if dataset_config is not None: + args["dataset_config"] = dataset_config + scenario.set_params_from_args(args=args) + await scenario.initialize_async() + assert scenario.atomic_attack_count >= 1, "Scenario built no atomic attacks during initialization." + + try: + return await scenario.run_async() + except ValueError as exc: + if _OBJECTIVES_INCOMPLETE_MARKER in str(exc): + return None + raise + + +@contextmanager +def _plugin_dir_env(*, plugin_dir: Path | None) -> Iterator[None]: + """Set PLUGIN_DIR (extraction dir) for the duration of a load, then restore it.""" + values: dict[str, str] = {} + if plugin_dir is not None: + values["PLUGIN_DIR"] = str(plugin_dir) + + saved: dict[str, str | None] = {key: os.environ.get(key) for key in values} + os.environ.update(values) + try: + yield + finally: + for key, previous in saved.items(): + if previous is None: + os.environ.pop(key, None) + else: + os.environ[key] = previous + + +@pytest.fixture +def plugin_dir(tmp_path: Path) -> Path: + """The directory the loader extracts wheels into for a test.""" + return tmp_path / ".plugin" + + +@pytest.fixture +def build_mock_wheel(tmp_path: Path) -> Callable[[str], Path]: + """A builder for the self-contained scenario wheel, parameterized by package name.""" + + def build(package: str) -> Path: + return _build_scenario_plugin_wheel(tmp_path, package=package) + + return build + + +@pytest.fixture +def all_scenarios() -> Callable[[list[Path]], set[str]]: + """The expected-set builder: scenario class names defined under the given source dirs.""" + return all_scenario_class_names + + +@pytest.fixture +def plugin_sandbox() -> Iterator[None]: + """Snapshot and restore global import + registry state so a load does not leak.""" + sys_path_snapshot = list(sys.path) + modules_snapshot = set(sys.modules) + provider_snapshot = dict(SeedDatasetProvider._registry) + + 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( + plugin_sandbox: None, + build_mock_wheel: Callable[[str], Path], + plugin_dir: Path, +) -> None: + """A built wheel's scenarios are all registered in ScenarioRegistry after init.""" + package = f"mock_scenario_plugin_{uuid.uuid4().hex[:8]}" + wheel = build_mock_wheel(package) + + registry = ScenarioRegistry.get_registry_singleton() + before = set(registry.get_class_names()) + + with _plugin_dir_env(plugin_dir=plugin_dir): + await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel)]) + + 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, + all_scenarios: Callable[[list[Path]], set[str]], +) -> None: + """Every scenario in an out-of-band wheel is discovered after initialization. + + Skipped unless ``PLUGIN_TEST_WHEEL`` is set, so the committed test names no external + package. Provide ``PLUGIN_TEST_SCENARIO_DIRS`` (preferred) or ``PLUGIN_TEST_PACKAGE`` + to tell the test which scenarios to expect. + """ + wheel_env = os.getenv("PLUGIN_TEST_WHEEL") + if not wheel_env: + pytest.skip("PLUGIN_TEST_WHEEL is not set; skipping injected-wheel scenario discovery test.") + + wheel = Path(wheel_env).expanduser() + assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" + + scenario_dirs_env = os.getenv("PLUGIN_TEST_SCENARIO_DIRS") + package = os.getenv("PLUGIN_TEST_PACKAGE") + if not scenario_dirs_env and not package: + pytest.skip("Set PLUGIN_TEST_SCENARIO_DIRS or PLUGIN_TEST_PACKAGE to define the expected scenarios.") + + with _plugin_dir_env(plugin_dir=None): + await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel, package=package)]) + + found = _registered_scenario_class_names() + + if scenario_dirs_env: + dirs = [Path(p) for p in scenario_dirs_env.split(os.pathsep) if p] + expected = all_scenarios(dirs) + else: + assert package is not None + expected = {cls.__name__ for cls in _scenario_classes_under_package(package).values()} + + assert expected, "No plug-in scenarios were found to verify; check the injected wheel/scenario source." + missing = expected - found + assert not missing, f"Plug-in scenarios not discovered by ScenarioRegistry: {sorted(missing)}" + + +@pytest.mark.run_only_if_all_tests +async def test_injected_wheel_scenarios_instantiate(plugin_sandbox: None) -> None: + """Every scenario in an out-of-band wheel constructs cleanly after initialization. + + Skipped unless ``PLUGIN_TEST_WHEEL`` and ``PLUGIN_TEST_PACKAGE`` are set. + """ + wheel_env = os.getenv("PLUGIN_TEST_WHEEL") + package = os.getenv("PLUGIN_TEST_PACKAGE") + if not wheel_env or not package: + pytest.skip("Set PLUGIN_TEST_WHEEL and PLUGIN_TEST_PACKAGE to run the instantiation test.") + + wheel = Path(wheel_env).expanduser() + assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" + + with _plugin_dir_env(plugin_dir=None): + await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel, package=package)]) + + classes = _scenario_classes_under_package(package) + assert classes, f"No plug-in scenarios registered under package {package!r}." + + failures: dict[str, str] = {} + for name, cls in sorted(classes.items()): + try: + instance = cls() # type: ignore[ty:missing-argument] + assert isinstance(instance, Scenario) + except Exception as exc: # noqa: BLE001 + failures[name] = f"{type(exc).__name__}: {exc}" + + assert not failures, f"Plug-in scenarios failed to instantiate: {failures}" + + +@pytest.mark.run_only_if_all_tests +async def test_injected_wheel_scenario_executes(plugin_sandbox: None) -> None: + """One named scenario from an out-of-band wheel runs through the public pipeline. + + Skipped unless ``PLUGIN_TEST_WHEEL``, ``PLUGIN_TEST_PACKAGE``, and + ``PLUGIN_TEST_EXEC_SCENARIO`` are set, and an ``ADVERSARIAL_CHAT_ENDPOINT`` is + configured (it may come from the loaded ``.env``, so it is checked after init). + """ + wheel_env = os.getenv("PLUGIN_TEST_WHEEL") + package = os.getenv("PLUGIN_TEST_PACKAGE") + exec_scenario = os.getenv("PLUGIN_TEST_EXEC_SCENARIO") + if not wheel_env or not package or not exec_scenario: + pytest.skip("Set PLUGIN_TEST_WHEEL, PLUGIN_TEST_PACKAGE, and PLUGIN_TEST_EXEC_SCENARIO to run the exec test.") + + wheel = Path(wheel_env).expanduser() + assert wheel.is_file(), f"PLUGIN_TEST_WHEEL does not exist: {wheel}" + + with _plugin_dir_env(plugin_dir=None): + await initialize_pyrit_async(IN_MEMORY, plugins=[PluginSpec(wheel=wheel, package=package)]) + + endpoint = os.getenv("ADVERSARIAL_CHAT_ENDPOINT") + if not endpoint: + pytest.skip("ADVERSARIAL_CHAT_ENDPOINT is not configured; skipping plug-in scenario execution.") + + registry = ScenarioRegistry.get_registry_singleton() + assert exec_scenario in registry.get_class_names(), ( + f"PLUGIN_TEST_EXEC_SCENARIO {exec_scenario!r} is not registered; " + f"available: {sorted(registry.get_class_names())}" + ) + scenario_cls = registry.get_class(exec_scenario) + + result = await _execute_scenario_async( + scenario_cls=scenario_cls, + endpoint=endpoint, + model=os.getenv("ADVERSARIAL_CHAT_MODEL"), + ) + + if result is not None: + assert result.scenario_run_state == ScenarioRunState.COMPLETED + assert result.attack_results, "Completed scenario run produced no attack results." diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 3c33fac5dc..804a67be6b 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -710,3 +710,82 @@ def test_server_url_non_string_raises(self): def test_server_non_dict_raises(self): with pytest.raises(ValueError, match="Server entry must be a dict"): ConfigurationLoader(server="http://oops:8000") # type: ignore[arg-type] + + +class TestConfigurationLoaderPlugins: + """Tests for the .pyrit_conf `plugins` list and `plugin_accept_load_failures`.""" + + def test_no_plugins_by_default(self): + """A config with no plugins resolves to an empty plug-in list.""" + config = ConfigurationLoader() + assert config.plugins == [] + assert config.resolve_plugins() == [] + assert config.plugin_accept_load_failures is None + + def test_bare_string_entry(self): + """A bare wheel-path string resolves to a spec with no explicit package.""" + from pyrit.setup.plugin_loader import PluginSpec + + config = ConfigurationLoader(plugins=["/abs/path/plugin.whl"]) + assert config.resolve_plugins() == [PluginSpec(wheel=pathlib.Path("/abs/path/plugin.whl"))] + + def test_concise_package_wheel_pair(self): + """A {package: wheel} mapping names the package.""" + from pyrit.setup.plugin_loader import PluginSpec + + config = ConfigurationLoader(plugins=[{"my_pkg": "/abs/path/plugin.whl"}]) + assert config.resolve_plugins() == [PluginSpec(wheel=pathlib.Path("/abs/path/plugin.whl"), package="my_pkg")] + + def test_explicit_mapping_entry(self): + """An explicit {wheel, package} mapping resolves both fields.""" + from pyrit.setup.plugin_loader import PluginSpec + + config = ConfigurationLoader(plugins=[{"wheel": "/abs/path/plugin.whl", "package": "my_pkg"}]) + assert config.resolve_plugins() == [PluginSpec(wheel=pathlib.Path("/abs/path/plugin.whl"), package="my_pkg")] + + def test_multiple_plugins_preserve_order(self): + """Several plug-ins resolve in configured order.""" + config = ConfigurationLoader(plugins=["/a/first.whl", {"second_pkg": "/b/second.whl"}]) + specs = config.resolve_plugins() + assert [s.wheel.name for s in specs] == ["first.whl", "second.whl"] + assert specs[1].package == "second_pkg" + + def test_malformed_plugin_entry_fails_fast(self): + """A malformed plug-in entry raises during construction (before other setup).""" + with pytest.raises(ValueError): + ConfigurationLoader(plugins=[123]) # type: ignore[list-item] + + def test_from_dict_with_plugins(self): + """from_dict wires the plugins list and accept-load-failures flag onto the loader.""" + config = ConfigurationLoader.from_dict( + {"plugins": [{"pkg": "/x/plugin.whl"}], "plugin_accept_load_failures": True} + ) + assert config.plugin_accept_load_failures is True + assert config.resolve_plugins()[0].package == "pkg" + + async def test_initialize_forwards_plugins_and_flag(self): + """ConfigurationLoader.initialize_pyrit_async forwards resolved plugins and the flag to core.""" + from pyrit.setup.plugin_loader import PluginSpec + + config = ConfigurationLoader( + memory_db_type="in_memory", + plugins=[{"pkg": "/x/plugin.whl"}], + plugin_accept_load_failures=True, + ) + + with mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") as core_init: + await config.initialize_pyrit_async() + + _, kwargs = core_init.call_args + assert kwargs["plugins"] == [PluginSpec(wheel=pathlib.Path("/x/plugin.whl"), package="pkg")] + assert kwargs["plugin_accept_load_failures"] is True + + async def test_initialize_passes_none_plugins_when_empty(self): + """With no plug-ins configured, core is called with plugins=None (no-op path).""" + config = ConfigurationLoader(memory_db_type="in_memory") + + with mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") as core_init: + await config.initialize_pyrit_async() + + _, kwargs = core_init.call_args + assert kwargs["plugins"] is None diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py new file mode 100644 index 0000000000..9d93cd5362 --- /dev/null +++ b/tests/unit/setup/test_plugin_loader.py @@ -0,0 +1,986 @@ +# 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 accept-load-failures/fail-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, _verify_plugin_prerequisites, initialize_pyrit_async +from pyrit.setup.plugin_loader import ( + PluginImportError, + PluginLoader, + PluginLoadError, + PluginRegisteredNothingError, + PluginSpec, + PluginWheelNotFoundError, + load_plugins_if_configured_async, +) + +# --------------------------------------------------------------------------- +# Mock-wheel builder +# --------------------------------------------------------------------------- + + +class MockWheel: + """Handle describing a built mock plug-in wheel.""" + + def __init__(self, *, path: Path, package: str, scenario_name: str, dataset_name: str) -> None: + self.path = path + self.package = package + self.scenario_name = scenario_name + self.dataset_name = dataset_name + + +def _unique_package_name() -> str: + """Return a unique, import-safe mock package name.""" + return f"mock_plugin_{uuid.uuid4().hex[:8]}" + + +def build_mock_wheel( + dest_dir: Path, + *, + bootstrap: str = "initializer", + include_provider: bool = True, + include_scenario: bool = True, + wire_init: bool = True, + package_name: str | None = None, +) -> MockWheel: + """ + Build a mock plug-in wheel in ``dest_dir`` and return a handle to it. + + 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_DIR", "PLUGIN_ACCEPT_LOAD_FAILURES"): + if key not in overrides: + os.environ.pop(key, None) + yield + + +def _spec(wheel: MockWheel, *, package: str | None = None) -> PluginSpec: + """Build a PluginSpec from a mock wheel.""" + return PluginSpec(wheel=wheel.path, package=package) + + +def _dummy_loader(*, accept_load_failures: bool | None = None) -> PluginLoader: + """Build a PluginLoader with a placeholder spec for testing policy resolution.""" + return PluginLoader(spec=PluginSpec(wheel=Path("dummy.whl")), accept_load_failures=accept_load_failures) + + +async def load_plugin( + wheel: MockWheel, + plugin_dir: Path, + *, + accept_load_failures: bool | None = None, + package: str | None = None, + extra_env: dict[str, str] | None = None, +) -> None: + """Run the plug-in loader against a single mock wheel with an isolated env.""" + with plugin_env(PLUGIN_DIR=str(plugin_dir), **(extra_env or {})): + await load_plugins_if_configured_async( + plugins=[_spec(wheel, package=package)], accept_load_failures=accept_load_failures + ) + + +# --------------------------------------------------------------------------- +# Loader phase inside initialize_pyrit_async +# --------------------------------------------------------------------------- + + +def test_verify_plugin_prerequisites_passes_when_memory_set() -> None: + """The health check passes when central memory is initialized.""" + with patch.object(CentralMemory, "get_memory_instance", return_value=MagicMock()): + _verify_plugin_prerequisites() # must not raise + + +def test_verify_plugin_prerequisites_raises_when_memory_unset() -> None: + """The health check fails loudly when a prerequisite phase (memory) did not complete.""" + with patch.object(CentralMemory, "get_memory_instance", side_effect=ValueError("not set")): + with pytest.raises(RuntimeError, match="central memory is not initialized"): + _verify_plugin_prerequisites() + + +async def test_plugin_phase_runs_after_memory_before_initializers() -> None: + """initialize_pyrit_async loads plug-ins after memory is set, before initializers.""" + manager = MagicMock() + manager.attach_mock(MagicMock(), "set_memory") + manager.attach_mock(AsyncMock(), "load_plugins") + manager.attach_mock(AsyncMock(), "execute") + + with ( + patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), + patch.object(CentralMemory, "set_memory_instance", manager.set_memory), + patch("pyrit.setup.initialization._verify_plugin_prerequisites"), + patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", manager.load_plugins), + patch("pyrit.setup.initialization._execute_initializers_async", manager.execute), + ): + await initialize_pyrit_async( + IN_MEMORY, + initializers=[MagicMock()], + env_files=[], + silent=True, + plugins=[PluginSpec(wheel=Path("plugin.whl"))], + ) + + order = [call[0] for call in manager.mock_calls if call[0] in {"set_memory", "load_plugins", "execute"}] + assert order.index("set_memory") < order.index("load_plugins") < order.index("execute") + + +async def test_plugin_phase_forwards_accept_load_failures_param() -> None: + """initialize_pyrit_async forwards plugins and plugin_accept_load_failures to the loader.""" + load_plugins_mock = AsyncMock() + spec = PluginSpec(wheel=Path("plugin.whl")) + with ( + patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), + patch.object(CentralMemory, "set_memory_instance"), + patch("pyrit.setup.initialization._verify_plugin_prerequisites"), + patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", load_plugins_mock), + ): + await initialize_pyrit_async( + IN_MEMORY, env_files=[], silent=True, plugins=[spec], plugin_accept_load_failures=True + ) + + load_plugins_mock.assert_awaited_once_with(plugins=[spec], accept_load_failures=True) + + +async def test_plugin_phase_skipped_when_no_plugins() -> None: + """initialize_pyrit_async does not touch the plug-in loader when no plug-ins are configured.""" + load_plugins_mock = AsyncMock() + with ( + patch("pyrit.setup.initialization.SQLiteMemory", return_value=MagicMock()), + patch.object(CentralMemory, "set_memory_instance"), + patch("pyrit.setup.plugin_loader.load_plugins_if_configured_async", load_plugins_mock), + ): + await initialize_pyrit_async(IN_MEMORY, env_files=[], silent=True) + + load_plugins_mock.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# No-op behavior +# --------------------------------------------------------------------------- + + +async def test_no_op_when_no_plugins() -> None: + """With an empty plug-in list the loader does nothing and registers nothing.""" + providers_before = dict(SeedDatasetProvider.get_all_providers()) + path_before = list(sys.path) + + with plugin_env(): + await load_plugins_if_configured_async(plugins=[]) + + assert SeedDatasetProvider.get_all_providers() == providers_before + assert sys.path == path_before + + +# --------------------------------------------------------------------------- +# 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_plugin_scenario_auto_registered_without_bootstrap(tmp_path: Path) -> None: + """A plug-in's Scenario subclass is auto-registered by discovery even with no bootstrap.""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False) + + await load_plugin(wheel, tmp_path / ".plugin") + + # The scenario is picked up purely by type-scoped discovery of the plug-in package, + # so a scenario-only plug-in with no register()/initializer still loads. + registry = ScenarioRegistry.get_registry_singleton() + mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario + assert mock_scenario in registry._classes.values() + assert registry._discovered is False # auto-registration must not trigger built-in discovery + + +async def test_bootstrap_registration_not_duplicated_by_auto_register(tmp_path: Path) -> None: + """A scenario the bootstrap registers is not also re-registered under a fallback name.""" + wheel = build_mock_wheel(tmp_path, bootstrap="register") + + await load_plugin(wheel, tmp_path / ".plugin") + + registry = ScenarioRegistry.get_registry_singleton() + mock_scenario = sys.modules[f"{wheel.package}.scenario"].MockScenario + registered_names = [name for name, cls in registry._classes.items() if cls is mock_scenario] + assert registered_names == [wheel.scenario_name] + + +async def test_datasets_only_plugin_loads_without_bootstrap(tmp_path: Path) -> None: + """A datasets-only plug-in (no bootstrap, no scenario) loads via import-time registration.""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_scenario=False) + + 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) + + # The spec's package points at a stdlib package that imports from outside the extraction dir. + with pytest.raises(PluginImportError, match="shadowing"): + await load_plugin(wheel, tmp_path / ".plugin", package="json") + + +# --------------------------------------------------------------------------- +# 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 pytest.raises(PluginLoadError): + await load_plugin(wheel, plugin_dir) + + 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 pytest.raises(PluginLoadError): + await load_plugin(wheel, plugin_dir) + + registry = ScenarioRegistry.get_registry_singleton() + assert wheel.scenario_name not in registry._classes + assert "MockProvider" not in SeedDatasetProvider.get_all_providers() + assert str(plugin_dir / wheel.path.stem) not in sys.path + + +async def test_accept_load_failures_rolls_back_partial_registration(tmp_path: Path) -> None: + """When load failures are accepted, a partially-registered failed plug-in is still fully rolled back.""" + wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") + plugin_dir = tmp_path / ".plugin" + + await load_plugin(wheel, plugin_dir, accept_load_failures=True) # must not raise + + registry = ScenarioRegistry.get_registry_singleton() + assert wheel.scenario_name not in registry._classes + + +async def test_rollback_restores_overwritten_provider(tmp_path: Path) -> None: + """A failed load restores a provider entry the plug-in overwrote (name collision).""" + wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") + + # SeedDatasetProvider keys by class name and the mock provider is "MockProvider"; + # occupy that key so the plug-in's import overwrites it. + class _PreexistingProvider(SeedDatasetProvider): + should_register = False + + @property + def dataset_name(self) -> str: + return "preexisting" + + async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: + raise NotImplementedError + + SeedDatasetProvider._registry["MockProvider"] = _PreexistingProvider + + with pytest.raises(PluginLoadError): + await load_plugin(wheel, tmp_path / ".plugin") + + # The original provider is restored, not deleted or left replaced by the plug-in's. + assert SeedDatasetProvider._registry["MockProvider"] is _PreexistingProvider + + +async def test_rollback_restores_overwritten_scenario(tmp_path: Path) -> None: + """A failed load restores a scenario entry the plug-in overwrote (name collision).""" + from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + + wheel = build_mock_wheel(tmp_path, bootstrap="initializer_raises") + + class _PreexistingScenario(RapidResponse): + """Sentinel scenario occupying the plug-in's registry name.""" + + registry = ScenarioRegistry.get_registry_singleton() + registry.register_class(_PreexistingScenario, name=wheel.scenario_name) + + with pytest.raises(PluginLoadError): + await load_plugin(wheel, tmp_path / ".plugin") + + # The original scenario is restored, not deleted or left replaced by the plug-in's. + assert registry._classes[wheel.scenario_name] is _PreexistingScenario + + +# --------------------------------------------------------------------------- +# Extraction cache +# --------------------------------------------------------------------------- + + +def test_extract_wheel_reuses_cached_extraction(tmp_path: Path) -> None: + """A second extraction of an unchanged wheel reuses the cached directory.""" + wheel = build_mock_wheel(tmp_path) + plugin_dir = tmp_path / ".plugin" + + with plugin_env(PLUGIN_DIR=str(plugin_dir)): + initializer = PluginLoader(spec=_spec(wheel)) + first = initializer._extract_wheel(wheel_path=wheel.path) + marker = first / "cache_marker.txt" + marker.write_text("kept", encoding="utf-8") + + second = initializer._extract_wheel(wheel_path=wheel.path) + + assert first == second + assert marker.is_file() # not wiped -> cached, not re-extracted + + +# --------------------------------------------------------------------------- +# Package name resolution +# --------------------------------------------------------------------------- + + +def test_resolve_package_name_prefers_explicit(tmp_path: Path) -> None: + """An explicit package name takes precedence over inference.""" + (tmp_path / "some_pkg").mkdir() + (tmp_path / "some_pkg" / "__init__.py").write_text("", encoding="utf-8") + + assert PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package="explicit_pkg") == "explicit_pkg" + + +def test_resolve_package_name_infers_single_package(tmp_path: Path) -> None: + """The single importable top-level directory is inferred when no package/top_level.txt exists.""" + (tmp_path / "the_pkg").mkdir() + (tmp_path / "the_pkg" / "__init__.py").write_text("", encoding="utf-8") + (tmp_path / "the_pkg-0.0.1.dist-info").mkdir() + + assert PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) == "the_pkg" + + +def test_resolve_package_name_uses_top_level_txt(tmp_path: Path) -> None: + """top_level.txt is consulted before directory inference.""" + (tmp_path / "pkg_a").mkdir() + (tmp_path / "pkg_a" / "__init__.py").write_text("", encoding="utf-8") + (tmp_path / "pkg_b").mkdir() + (tmp_path / "pkg_b" / "__init__.py").write_text("", encoding="utf-8") + distinfo = tmp_path / "thing-0.0.1.dist-info" + distinfo.mkdir() + (distinfo / "top_level.txt").write_text("pkg_b\n", encoding="utf-8") + + assert PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) == "pkg_b" + + +def test_resolve_package_name_none_raises(tmp_path: Path) -> None: + """No importable package raises a clear error pointing at the plug-in's config.""" + with pytest.raises(ValueError, match="package"): + PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) + + +def test_resolve_package_name_multiple_raises(tmp_path: Path) -> None: + """Multiple top-level packages require an explicit package to disambiguate.""" + for name in ("pkg_a", "pkg_b"): + (tmp_path / name).mkdir() + (tmp_path / name / "__init__.py").write_text("", encoding="utf-8") + + with pytest.raises(ValueError, match="disambiguate"): + PluginLoader._resolve_package_name(extract_dir=tmp_path, explicit_package=None) + + +# --------------------------------------------------------------------------- +# Failure modes +# --------------------------------------------------------------------------- + + +async def test_missing_wheel_fails_closed() -> None: + """A configured-but-missing wheel raises PluginWheelNotFoundError by default (fail-closed).""" + with plugin_env(): + with pytest.raises(PluginWheelNotFoundError, match="Failed to load plug-in"): + await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=Path("does_not_exist.whl"))]) + + +async def test_missing_wheel_accept_load_failures_param_proceeds() -> None: + """Accepting load failures via the explicit param skips a broken plug-in with a warning.""" + with plugin_env(): + await load_plugins_if_configured_async( + plugins=[PluginSpec(wheel=Path("does_not_exist.whl"))], accept_load_failures=True + ) # must not raise + + +async def test_missing_wheel_accept_load_failures_env_proceeds() -> None: + """Accepting load failures via PLUGIN_ACCEPT_LOAD_FAILURES env skips a broken plug-in with a warning.""" + with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): + await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=Path("does_not_exist.whl"))]) # must not raise + + +async def test_empty_wheel_is_loud(tmp_path: Path) -> None: + """A wheel that imports cleanly but registers nothing fails loudly.""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) + + with pytest.raises(PluginRegisteredNothingError, match="registered no datasets or scenarios"): + await load_plugin(wheel, tmp_path / ".plugin") + + +async def test_empty_wheel_accept_load_failures_proceeds(tmp_path: Path) -> None: + """An empty wheel with load failures accepted proceeds instead of raising.""" + wheel = build_mock_wheel(tmp_path, bootstrap="none", include_provider=False, include_scenario=False) + + await load_plugin(wheel, tmp_path / ".plugin", accept_load_failures=True) # must not raise + + +async def test_non_whl_path_fails_closed(tmp_path: Path) -> None: + """A wheel path that is not a .whl file fails closed with PluginWheelNotFoundError.""" + not_a_wheel = tmp_path / "plugin.zip" + not_a_wheel.write_text("not a wheel", encoding="utf-8") + + with plugin_env(): + with pytest.raises(PluginWheelNotFoundError, match="Failed to load plug-in"): + await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=not_a_wheel)]) + + +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" + with zipfile.ZipFile(malicious, "w") as archive: + archive.writestr("evil_pkg/__init__.py", "") + archive.writestr("../escape.py", "compromised = True") + + with plugin_env(PLUGIN_DIR=str(tmp_path / ".plugin")): + with pytest.raises(PluginLoadError, match="Failed to load plug-in"): + await load_plugins_if_configured_async(plugins=[PluginSpec(wheel=malicious)]) + + # The traversal target was not written outside the extraction directory. + assert not (tmp_path / "escape.py").exists() + + +# --------------------------------------------------------------------------- +# No-arg-instantiable contract +# --------------------------------------------------------------------------- + + +def test_non_no_arg_scenario_fails_metadata_cleanly() -> None: + """A registered scenario that is not no-arg instantiable fails metadata build clearly.""" + from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse + + class BadScenario(RapidResponse): + """Scenario that violates the no-arg-instantiable contract.""" + + def __init__(self, *, required_value: str) -> None: + super().__init__() + self._required_value = required_value + + registry = ScenarioRegistry() + registry.register_class(BadScenario, name="airt.bad") # signature-only validation passes + + with pytest.raises(TypeError, match="no arguments"): + registry._build_metadata("airt.bad", BadScenario) + + +# --------------------------------------------------------------------------- +# accept_load_failures resolution +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "value,expected", + [("true", True), ("True", True), ("1", True), ("yes", True), ("on", True), ("false", False), ("0", False)], +) +def test_resolve_accept_load_failures_from_env_tokens(value: str, expected: bool) -> None: + """accept_load_failures resolves from PLUGIN_ACCEPT_LOAD_FAILURES across truthy/falsey tokens.""" + with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES=value): + assert _dummy_loader()._resolve_accept_load_failures() is expected + + +def test_resolve_accept_load_failures_explicit_true() -> None: + """An explicit accept_load_failures=True resolves to True.""" + with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="false"): + assert _dummy_loader(accept_load_failures=True)._resolve_accept_load_failures() is True + + +def test_resolve_accept_load_failures_explicit_overrides_env() -> None: + """An explicit accept_load_failures value takes precedence over the env var.""" + with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): + assert _dummy_loader(accept_load_failures=False)._resolve_accept_load_failures() is False + + +def test_resolve_accept_load_failures_from_env_when_no_explicit() -> None: + """accept_load_failures falls back to PLUGIN_ACCEPT_LOAD_FAILURES when no explicit value is set.""" + with plugin_env(PLUGIN_ACCEPT_LOAD_FAILURES="true"): + assert _dummy_loader()._resolve_accept_load_failures() is True + + +def test_resolve_accept_load_failures_defaults_false() -> None: + """accept_load_failures defaults to False (fail-closed).""" + with plugin_env(): + assert _dummy_loader()._resolve_accept_load_failures() is False + + +# --------------------------------------------------------------------------- +# PluginSpec.from_config parsing +# --------------------------------------------------------------------------- + + +def test_plugin_spec_from_config_bare_string() -> None: + """A bare wheel-path string parses to a spec with no explicit package.""" + spec = PluginSpec.from_config("/abs/path/plugin.whl") + assert spec.wheel == Path("/abs/path/plugin.whl") + assert spec.package is None + + +def test_plugin_spec_from_config_single_key_pair() -> None: + """A {package: wheel} mapping names the package.""" + spec = PluginSpec.from_config({"my_pkg": "/abs/path/plugin.whl"}) + assert spec.wheel == Path("/abs/path/plugin.whl") + assert spec.package == "my_pkg" + + +def test_plugin_spec_from_config_explicit_mapping() -> None: + """An explicit {wheel, package} mapping parses both fields; package is optional.""" + spec = PluginSpec.from_config({"wheel": "/abs/path/plugin.whl", "package": "my_pkg"}) + assert spec == PluginSpec(wheel=Path("/abs/path/plugin.whl"), package="my_pkg") + assert PluginSpec.from_config({"wheel": "/abs/path/plugin.whl"}).package is None + + +@pytest.mark.parametrize( + "entry", [123, {"a": "1", "b": "2"}, {"package": "no_wheel"}, {"wheel": "/x.whl", "packge": "typo"}] +) +def test_plugin_spec_from_config_rejects_bad_shapes(entry: object) -> None: + """Unsupported entry shapes (including explicit mappings with unexpected keys) raise ValueError.""" + with pytest.raises(ValueError): + PluginSpec.from_config(entry) # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- +# Multiple plug-ins +# --------------------------------------------------------------------------- + + +async def test_multiple_plugins_all_load(tmp_path: Path) -> None: + """Every configured plug-in loads; a single plug-in is just the one-element case.""" + wheel_a = build_mock_wheel(tmp_path / "a", package_name="mock_plugin_alpha") + wheel_b = build_mock_wheel(tmp_path / "b", package_name="mock_plugin_beta") + + with plugin_env(PLUGIN_DIR=str(tmp_path / ".plugin")): + await load_plugins_if_configured_async(plugins=[_spec(wheel_a), _spec(wheel_b)]) + + names = ScenarioRegistry.get_registry_singleton().get_class_names() + assert wheel_a.scenario_name in names + assert wheel_b.scenario_name in names + + +async def test_plugins_load_in_configured_order(tmp_path: Path) -> None: + """Plug-ins load in list order (later entries load after earlier ones).""" + wheel_a = build_mock_wheel(tmp_path / "a", package_name="mock_plugin_first") + wheel_b = build_mock_wheel(tmp_path / "b", package_name="mock_plugin_second") + + loaded: list[str] = [] + original = PluginLoader.load_async + + async def _record(self: PluginLoader) -> None: + loaded.append(self._spec.package or self._spec.wheel.stem) + await original(self) + + with plugin_env(PLUGIN_DIR=str(tmp_path / ".plugin")): + with patch.object(PluginLoader, "load_async", _record): + await load_plugins_if_configured_async( + plugins=[_spec(wheel_a, package="mock_plugin_first"), _spec(wheel_b, package="mock_plugin_second")] + ) + + assert loaded == ["mock_plugin_first", "mock_plugin_second"]