FEAT Plug-in mechanism for private scenarios

.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 = ""

.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/
 

.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).

.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

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.

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]]:

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",
 ]

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,
         )