diff --git a/.pyrit_conf_example b/.pyrit_conf_example index bbf1e98903..2d1166fcad 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,35 @@ initializers: - name: technique - name: load_default_datasets +# Plug-ins +# -------- +# A plug-in points PyRIT at a private ``PyRITInitializer`` that lives outside the public +# tree (for example, a team's internal red-teaming package). ConfigurationLoader runs it +# as a guaranteed-first, privileged initializer -- before the initializers above -- so the +# scenarios, attack techniques, datasets, and targets it registers exist before anything +# else consumes the registries. Do NOT list it under `initializers:`. +# +# The plug-in's initializer owns all registration; PyRIT discovers nothing on its own. +# A plug-in is one entry with three fields: +# - name: an operator label (used in logs/errors) +# - source: the folder that contains your package, placed on sys.path so +# `import ` resolves +# - initializer: dotted 'module.Class' path to a PyRITInitializer subclass +# Relative source paths resolve against this configuration file. +# +# WARNING: loading a plug-in executes third-party code in the PyRIT process. Whoever can +# write this file (or the source) can run code on the host. Treat it as sensitive. +# +# Example: +# plugins: +# - name: my_redteam +# source: /repos/my-redteam +# initializer: my_redteam.setup.MyInitializer +# +# Help and troubleshooting: +# - doc/getting_started/plugins.md +# - doc/getting_started/troubleshooting/plugins.md + # Default Scenario # ---------------- # Optional default scenario to run when invoking `pyrit_scan` without a diff --git a/doc/getting_started/README.md b/doc/getting_started/README.md index b5795569d7..af16d905dd 100644 --- a/doc/getting_started/README.md +++ b/doc/getting_started/README.md @@ -28,4 +28,5 @@ Once you're set up: - 🔌 [Targets](../code/targets/0_prompt_targets.md) — Connect to different AI systems - 📦 [Scenarios](../code/scenarios/0_scenarios.ipynb) — Run standardized evaluation scenarios +- 🧩 [Private Scenario Plug-Ins](./plugins.md) — Add private scenarios and attack techniques to stock PyRIT - 🖥️ [Scanner](../scanner/0_scanner.md) — Use `pyrit_scan` for automated assessments diff --git a/doc/getting_started/plugins.md b/doc/getting_started/plugins.md new file mode 100644 index 0000000000..29dfef7ae6 --- /dev/null +++ b/doc/getting_started/plugins.md @@ -0,0 +1,130 @@ +# Private Scenarios and Attack Techniques + +PyRIT plug-ins let an operator activate private scenarios and attack techniques from a +**stock PyRIT installation** — without forking PyRIT or passing per-run CLI flags. + +## Plug-in or `--initialization-scripts`? + +A plug-in is a thin layer over the existing initializer path: + +- `--initialization-scripts ./my_init.py` runs a custom `PyRITInitializer` from a loose + file. It works for a self-contained script but is per-run and breaks for a packaged + initializer that imports from its own package. +- A `plugins:` entry in `.pyrit_conf` names a private initializer once, runs it + **first** (before other initializers and before catalog warming), and anchors the + package root on `sys.path` so a packaged initializer with intra-package imports loads + correctly. + +Use a plug-in when a private, packaged initializer (for example a team's internal +red-teaming package) needs to behave like a built-in inside the standard backend, GUI, +and `pyrit_scan` catalog. + +## What a plug-in is + +A plug-in is a config entry that points at a concrete `PyRITInitializer` reachable by a +dotted path from a source root: + +```yaml +# .pyrit_conf +plugins: + - name: my_redteam + source: /repos/my-redteam + initializer: my_redteam.setup.MyInitializer +``` + +- `name` — an operator label used in logs and errors. +- `source` — the folder placed on `sys.path` so your package can be imported (see below). +- `initializer` — a dotted `module.Class` path to a concrete `PyRITInitializer`. + +### What `source` should point at (and why it matters) + +`source` should be the folder that **contains** your package — the directory you would be +sitting in for `import my_redteam` to succeed at a Python prompt. PyRIT adds that folder +to Python's import search path (`sys.path`) before importing your initializer. + +You need this because a real private package doesn't live next to PyRIT; its modules +import from *each other* (for example `from my_redteam.datasets import load`). If Python +can't find the package root, those imports fail and the plug-in won't load. In plain +terms: point `source` at the folder above your package, not at the package folder itself +and not at a single file buried inside it. If you point at the wrong place, loading fails +closed with an import error naming what could not be found. + +`ConfigurationLoader` runs the plug-in as a privileged initializer, always **first** — +before your other initializers and before anything reads the scenario/technique catalog. +You do not (and cannot) add it to `initializers:`: it isn't a registered initializer name, +so listing it there fails with an "initializer not found" error. The framework constructs +and runs it for you. + +## The initializer owns registration + +PyRIT discovers nothing on its own. The plug-in's initializer registers everything it +wants discoverable, at whatever level of abstraction fits: + +- **Attack techniques** — + `AttackTechniqueRegistry.get_registry_singleton().register_from_factories([...])`; + selectable via `--techniques`. +- **Scenarios** — + `ScenarioRegistry.get_registry_singleton().register_class(MyScenario, name="my_redteam.violence")`; + runnable via `pyrit_scan my_redteam.violence`. +- **Datasets** — register providers and load them into memory so private seeds stay in + the operator's database and are never published. +- **Default targets** — `set_default_value(...)`. + +### Deciding what to keep private + +"Private" is rarely all-or-nothing. If you build a custom scenario or technique, you have +to decide whether to contribute it publicly, keep it in your own tracked repo, or keep it +fully private — and often only *part* of it is sensitive. A plug-in lets you keep exactly +the sensitive layer private while everything else stays public and works out of the box: + +- sometimes only the **dataset** (the prompts/objectives) is sensitive; +- sometimes it's a niche **technique**; +- sometimes an entire **scenario** should not be exposed at all. + +Because the initializer registers each level independently, you can publish the generic +parts and register only the sensitive parts from your private package. + +### Example initializer + +```python +from pyrit.executor.attack import PromptSendingAttack +from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry +from pyrit.scenario.core import AttackTechniqueFactory +from pyrit.setup.pyrit_initializer import PyRITInitializer + +from my_redteam.scenarios import Violence + + +class MyInitializer(PyRITInitializer): + """Register a private technique and a private scenario.""" + + async def initialize_async(self) -> None: + AttackTechniqueRegistry.get_registry_singleton().register_from_factories( + [AttackTechniqueFactory(name="operation_foobar", attack_class=PromptSendingAttack)] + ) + ScenarioRegistry.get_registry_singleton().register_class(Violence, name="my_redteam.violence") +``` + +## Usage + +```powershell +# A private technique through a public scenario +pyrit_scan airt.rapid_response --target openai_chat --techniques operation_foobar + +# A private scenario +pyrit_scan my_redteam.violence --target openai_chat +``` + +## Behavior and limits + +- The plug-in initializer runs **first**, so lazy catalog/metadata consumers see a + complete registry. +- Loading executes third-party Python with backend permissions; whoever can write the + config or the source can run code on the host. Treat the config as sensitive. +- Dependencies must already be installed in the backend environment. +- The plug-in path is **fail-closed** and supports **one** plug-in. A failed load aborts + initialization — fix the config or source and restart. +- Plug-ins activate only at process/backend startup. Restart after changing the config + or the source; there is no hot reload. + +See [Plug-In Troubleshooting](./troubleshooting/plugins.md) for common failures. diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 47565222e0..44bd4c2a3d 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -51,6 +51,7 @@ flowchart LR ### Using .env.local for Overrides You can use `~/.pyrit/.env.local` to override values in `~/.pyrit/.env` without modifying the base file. This is useful for: + - Testing different targets - Using personal credentials instead of shared ones - Switching between configurations quickly @@ -151,6 +152,28 @@ initialization_scripts: - ./local_initializer.py ``` +### `plugins` + +One private plug-in to activate before the configured initializers. A plug-in points at +a private ``PyRITInitializer`` reached by a dotted path from a source root; the +initializer registers whatever private scenarios, attack techniques, datasets, and +targets it wants discoverable: + +```yaml +plugins: + - name: my_redteam + source: /repos/my-redteam + initializer: my_redteam.setup.MyInitializer +``` + +`ConfigurationLoader` automatically injects this as a privileged initializer that runs +first; you cannot add it to `initializers:` yourself (it is not a registered initializer +name). It is fail-closed, supports one plug-in, and requires a process/backend restart +after changes. + +See [Private Scenarios and Attack Techniques](./plugins.md) and +[Plug-In Troubleshooting](./troubleshooting/plugins.md). + ### `env_files` Environment file paths to load during initialization. Later files override values from earlier files. @@ -196,9 +219,13 @@ The 3-layer model above determines **which config values are selected**. Once re 1. Environment files are loaded 2. Default values are reset 3. Memory database is configured (from `memory_db_type`) -4. Initializers are executed in listed order +4. The privileged plug-in initializer runs when `plugins:` is configured +5. User-configured initializers are executed in listed order -Because initializers run last, they can modify anything set up in earlier steps — including environment variables and the memory instance. In practice, built-in initializers like `target` and `scorer` only call `set_default_value` and `set_global_variable` and do not touch memory or environment variables. However, a custom initializer could override those if needed. When this happens, the initializer's changes take effect because it runs after the other settings have been applied. +Because initializers run after environment and memory setup, they can use those +prerequisites. The plug-in initializer is framework-controlled and always precedes +the listed initializers so scenario/technique registries are complete before catalog +consumers run. ## Usage diff --git a/doc/getting_started/troubleshooting/README.md b/doc/getting_started/troubleshooting/README.md index c146eabe15..c91af61d85 100644 --- a/doc/getting_started/troubleshooting/README.md +++ b/doc/getting_started/troubleshooting/README.md @@ -14,6 +14,8 @@ Common issues and advanced setup guides. - **[Azure SQL Database Setup](./azure_sql_db.md)** — Setting up Azure SQL as your memory backend, including Entra ID authentication and user permissions. +- **[Private Scenario Plug-Ins](./plugins.md)** — Configuration, initializer resolution, import, and restart troubleshooting. + ## Model Deployment - **[Deploy HuggingFace Models on Azure ML](./deploy_hf_model_aml.ipynb)** — Deploy HuggingFace models to Azure ML endpoints. diff --git a/doc/getting_started/troubleshooting/plugins.md b/doc/getting_started/troubleshooting/plugins.md new file mode 100644 index 0000000000..94e5a11ee6 --- /dev/null +++ b/doc/getting_started/troubleshooting/plugins.md @@ -0,0 +1,63 @@ +# Plug-In Troubleshooting + +## The plug-in did not activate + +Restart the backend after changing `.pyrit_conf` or the source: + +```powershell +pyrit_scan --stop-server +pyrit_scan --start-server --config-file /path/to/.pyrit_conf +``` + +`pyrit_scan` is a thin client. An already-running backend keeps the plug-in loaded from +its startup config. + +## Configuration is rejected + +A plug-in is one entry with three fields: + +```yaml +plugins: + - name: my_redteam + source: /repos/my-redteam + initializer: my_redteam.setup.MyInitializer +``` + +Check that: + +- `name` is a valid lowercase snake_case registry name; +- `source` is present (relative paths resolve against the config file); +- `initializer` is a dotted `module.Class` path; +- only one plug-in is configured. + +## Source path does not exist + +The loader fails closed with a "source path does not exist" error. Point `source` at the +directory that contains your package (the parent of the top-level package) so +`import ` resolves once that directory is on `sys.path`. + +## Initializer cannot be imported + +- Confirm `initializer` names a real `module.Class` importable from `source`. +- Install the plug-in's dependencies into the backend Python environment. +- A packaged initializer must import from its own package (for example + `from my_redteam... import ...`); pointing `source` at the package root makes that work. + +## Target is not a PyRITInitializer + +`initializer` must resolve to a concrete subclass of `PyRITInitializer`. A plug-in +contributes an initializer, not loose scenario or technique objects. + +## Nothing shows up in the catalog + +PyRIT does not discover components — the initializer must register them. For scanner +discovery, the initializer must call +`ScenarioRegistry.get_registry_singleton().register_class(...)` for scenarios and +`AttackTechniqueRegistry.get_registry_singleton().register_from_factories(...)` for +techniques. Datasets must be registered as providers and loaded into memory. + +## Partial state after a failed load + +Plug-in loading is fail-closed. A failed load aborts initialization; fix the reported +stage and restart the process. Continuing in the same process after a failed +initialization is not supported. diff --git a/doc/myst.yml b/doc/myst.yml index c0f22d68ff..45a5de823d 100644 --- a/doc/myst.yml +++ b/doc/myst.yml @@ -51,6 +51,7 @@ project: children: - file: getting_started/populating_secrets.md - file: getting_started/pyrit_conf.md + - file: getting_started/plugins.md - file: getting_started/troubleshooting/README.md children: - file: getting_started/troubleshooting/installation.md @@ -60,6 +61,7 @@ project: - file: getting_started/troubleshooting/local_dev.md - file: getting_started/troubleshooting/jupyter_setup.md - file: getting_started/troubleshooting/azure_sql_db.md + - file: getting_started/troubleshooting/plugins.md - file: getting_started/troubleshooting/deployment.md children: - file: getting_started/troubleshooting/deploy_hf_model_aml.ipynb diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index 8c10d020cd..3d601c5391 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -84,6 +84,13 @@ def _print_cli_exception(*, exc: BaseException) -> None: Requires a running PyRIT backend server. Use --start-server to launch one, or connect to an existing server with --server-url. +Private scenario plug-ins: + Use a plug-in when private scenarios or attack techniques must behave like + built-ins in the stock scanner/backend/catalog. If you own a custom Python + application, importing PyRIT as a dependency and composing it directly is + usually simpler. Configure trusted source or wheel artifacts in .pyrit_conf. + See doc/getting_started/plugins.md. + Examples: # Start the backend server pyrit_scan --start-server diff --git a/pyrit/exceptions/__init__.py b/pyrit/exceptions/__init__.py index 6d889ab95d..abb5552ccc 100644 --- a/pyrit/exceptions/__init__.py +++ b/pyrit/exceptions/__init__.py @@ -10,6 +10,8 @@ ExperimentalWarning, InvalidJsonException, MissingPromptPlaceholderException, + PluginLoadError, + PluginSourceNotFoundError, PyritException, RateLimitException, ScorerLLMResponseBlockedException, @@ -53,6 +55,8 @@ "handle_bad_request_exception", "InvalidJsonException", "MissingPromptPlaceholderException", + "PluginLoadError", + "PluginSourceNotFoundError", "PyritException", "pyrit_custom_result_retry", "pyrit_json_retry", diff --git a/pyrit/exceptions/exception_classes.py b/pyrit/exceptions/exception_classes.py index 04043a2381..e1bdf5e40b 100644 --- a/pyrit/exceptions/exception_classes.py +++ b/pyrit/exceptions/exception_classes.py @@ -261,6 +261,14 @@ class ExperimentalWarning(FutureWarning): """ +class PluginLoadError(RuntimeError): + """Base error raised when a configured plug-in fails to load.""" + + +class PluginSourceNotFoundError(PluginLoadError): + """The configured plug-in source path does not point to a readable Python file or package.""" + + def pyrit_custom_result_retry( retry_function: Callable[..., bool], retry_max_num_attempts: int | None = None ) -> Callable[..., Any]: diff --git a/pyrit/registry/components/attack_technique_registry.py b/pyrit/registry/components/attack_technique_registry.py index ffa0bbbca0..3bc00f55dd 100644 --- a/pyrit/registry/components/attack_technique_registry.py +++ b/pyrit/registry/components/attack_technique_registry.py @@ -113,6 +113,7 @@ def register_technique( name: str, factory: AttackTechniqueFactory, tags: dict[str, str] | list[str] | None = None, + metadata: dict[str, object] | None = None, ) -> None: """ Register an attack technique factory. @@ -123,8 +124,9 @@ def register_technique( tags (dict[str, str] | list[str] | None): Optional tags for categorisation. Accepts a ``dict[str, str]`` or a ``list[str]`` (each string becomes a key with value ``""``). + metadata (dict[str, object] | None): Optional non-selection metadata. """ - self.instances.register(factory, name=name, tags=tags) + self.instances.register(factory, name=name, tags=tags, metadata=metadata) logger.debug(f"Registered attack technique factory: {name} ({factory.attack_class.__name__})") def get_factories(self) -> dict[str, AttackTechniqueFactory]: diff --git a/pyrit/registry/registry.py b/pyrit/registry/registry.py index f79f1f20c7..3c89249796 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,6 +281,25 @@ 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_subclasses_in_package(self, *, package_name: str) -> int: + """ + Register every concrete base-type subclass whose module lives under a package. + + Used by built-in discovery (``_discover``): enumerates the in-memory subclasses of + ``_base_type()``, keeps those whose module is ``package_name`` or a submodule of + it, names each via ``_get_registry_name``, and registers it. + + Args: + package_name (str): The package whose subclasses to register. + + Returns: + int: The number of classes newly registered. + """ + base = self._base_type() + package_prefix = f"{package_name}." + count = 0 for cls in self._iter_concrete_subclasses(base): module = cls.__module__ or "" if module != package_name and not module.startswith(package_prefix): @@ -299,7 +316,9 @@ 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 @staticmethod def _iter_concrete_subclasses(base: type[T]) -> list[type[T]]: diff --git a/pyrit/scenario/scenarios/airt/rapid_response.py b/pyrit/scenario/scenarios/airt/rapid_response.py index 8362d53595..e88558a291 100644 --- a/pyrit/scenario/scenarios/airt/rapid_response.py +++ b/pyrit/scenario/scenarios/airt/rapid_response.py @@ -35,8 +35,9 @@ def _build_rapid_response_technique() -> type[ScenarioTechnique]: """ Build the RapidResponse technique class dynamically from the registered factories. - Reads the singleton ``AttackTechniqueRegistry`` and filters to factories - tagged ``core``. + Reads the singleton ``AttackTechniqueRegistry`` and exposes every registered + technique. Which techniques are available is decided by what initializers + (including a plug-in initializer) have registered before this runs. Returns: type[ScenarioTechnique]: The dynamically generated technique enum class. @@ -50,7 +51,6 @@ def _build_rapid_response_technique() -> type[ScenarioTechnique]: return AttackTechniqueRegistry.build_technique_class_from_factories( # type: ignore[ty:invalid-return-type] class_name="RapidResponseTechnique", factories=factories, - available=TagQuery.all("core"), aggregate_tags={ "single_turn": TagQuery.any_of("single_turn"), "multi_turn": TagQuery.any_of("multi_turn"), diff --git a/pyrit/setup/__init__.py b/pyrit/setup/__init__.py index 4cac6e1470..9260523b0d 100644 --- a/pyrit/setup/__init__.py +++ b/pyrit/setup/__init__.py @@ -11,6 +11,7 @@ MemoryDatabaseType, initialize_pyrit_async, ) +from pyrit.setup.plugin_spec 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..d63986344d 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -13,7 +13,10 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar +import yaml + from pyrit.common.path import DEFAULT_CONFIG_PATH +from pyrit.common.utils import verify_and_resolve_path from pyrit.common.yaml_loadable import YamlLoadable from pyrit.models import class_name_to_snake_case from pyrit.setup.initialization import ( @@ -24,6 +27,7 @@ ) if TYPE_CHECKING: + from pyrit.setup.plugin_spec import PluginSpec from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -106,6 +110,10 @@ 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. + Accepts at most one explicit mapping with ``name``, a ``source`` path, and + a dotted ``initializer`` (module.Class) pointing at a ``PyRITInitializer``. + Empty means no plug-ins. Example YAML configuration: memory_db_type: sqlite @@ -149,10 +157,15 @@ class ConfigurationLoader(YamlLoadable): max_concurrent_scenario_runs: int = 3 allow_custom_initializers: bool = False server: dict[str, Any] | None = None + plugins: list[dict[str, Any]] = field(default_factory=list) extensions: dict[str, Any] = field(default_factory=dict) + _plugin_base_dir: pathlib.Path = field(default_factory=pathlib.Path.cwd, repr=False, compare=False) 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 +229,23 @@ 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. + + Accepts at most one explicit mapping (``name`` + ``source`` + dotted + ``initializer``). Validating here (before any other normalization) fails fast on + a malformed plug-in entry. + + Raises: + ValueError: If a plug-in entry has an unsupported shape. + """ + from pyrit.setup.plugin_spec import PluginSpec + + if len(self.plugins) > 1: + raise ValueError("Only one plug-in is supported at a time; plug-in composition is not supported.") + self._plugin_specs = [PluginSpec.from_config(entry, base_dir=self._plugin_base_dir) for entry in self.plugins] + def _normalize_scenario(self) -> None: """ Normalize the optional ``scenario`` block to a ``ScenarioConfig``. @@ -288,12 +318,13 @@ def scenario_config(self) -> ScenarioConfig | None: return self._scenario_config @classmethod - def from_dict(cls, data: dict[str, Any]) -> "ConfigurationLoader": + def from_dict(cls, data: dict[str, Any], *, plugin_base_dir: pathlib.Path | None = None) -> "ConfigurationLoader": """ Create a ConfigurationLoader from a dictionary. Args: data: Dictionary containing configuration values. + plugin_base_dir: Directory used to resolve relative plug-in paths. Returns: A new ConfigurationLoader instance. @@ -303,7 +334,7 @@ def from_dict(cls, data: dict[str, Any]) -> "ConfigurationLoader": """ # Filter out None values only - empty lists are meaningful ("load nothing") filtered_data = {k: v for k, v in data.items() if v is not None} - known_fields = set(cls.__dataclass_fields__.keys()) + known_fields = {name for name in cls.__dataclass_fields__ if not name.startswith("_")} known_data = {k: v for k, v in filtered_data.items() if k in known_fields and k != "extensions"} extra_data = {k: v for k, v in filtered_data.items() if k not in known_fields} if "extensions" in filtered_data: @@ -311,7 +342,37 @@ def from_dict(cls, data: dict[str, Any]) -> "ConfigurationLoader": if not isinstance(extensions, dict): raise ValueError(f"ConfigurationLoader.extensions must be a dict. Got: {type(extensions).__name__}") extra_data = {**extra_data, **extensions} - return cls(**known_data, extensions=extra_data) + return cls( + **known_data, + extensions=extra_data, + _plugin_base_dir=(plugin_base_dir or pathlib.Path.cwd()).resolve(), + ) + + @classmethod + def from_yaml_file(cls, file: pathlib.Path | str) -> "ConfigurationLoader": + """ + Load configuration while retaining the file directory for plug-in paths. + + Args: + file: The YAML configuration path. + + Returns: + ConfigurationLoader: The parsed and normalized configuration. + + Raises: + FileNotFoundError: If the path does not exist. + ValueError: If the YAML is empty, malformed, or has a non-mapping root. + """ + resolved = verify_and_resolve_path(file) + try: + yaml_data = yaml.safe_load(resolved.read_text(encoding="utf-8")) + except yaml.YAMLError as exc: + raise ValueError(f"Invalid YAML file '{resolved}': {exc}") from exc + if yaml_data is None: + raise ValueError(f"YAML file '{resolved}' is empty.") + if not isinstance(yaml_data, dict): + raise ValueError(f"YAML file '{resolved}' must contain a top-level mapping.") + return cls.from_dict(yaml_data, plugin_base_dir=resolved.parent) @staticmethod def load_with_overrides( @@ -362,6 +423,7 @@ def load_with_overrides( "env_files": None, # None = use defaults "env_akv_ref": None, "silent": False, + "plugins": [], } # 1. Try loading default config file if it exists @@ -381,6 +443,7 @@ 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"] = [spec.to_config() for spec in default_config.resolve_plugins()] if default_config.operator: config_data["operator"] = default_config.operator if default_config.operation: @@ -407,6 +470,7 @@ 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"] = [spec.to_config() for spec in explicit_config.resolve_plugins()] if explicit_config.operator: config_data["operator"] = explicit_config.operator if explicit_config.operation: @@ -491,6 +555,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. @@ -560,9 +633,14 @@ async def initialize_pyrit_async(self) -> None: Raises: ValueError: If configuration is invalid or initializers cannot be resolved. """ - resolved_initializers = self.resolve_initializers() + resolved_initializers = list(self.resolve_initializers()) resolved_scripts = self.resolve_initialization_scripts() resolved_env_files = self.resolve_env_files() + resolved_plugins = self.resolve_plugins() + if resolved_plugins: + from pyrit.setup.plugin_loader import PluginInitializer + + resolved_initializers.insert(0, PluginInitializer(plugins=resolved_plugins)) # Map snake_case memory_db_type to internal constant internal_memory_db_type = self._MEMORY_DB_TYPE_MAP[self.memory_db_type] diff --git a/pyrit/setup/plugin_loader.py b/pyrit/setup/plugin_loader.py new file mode 100644 index 0000000000..31fc2e65a5 --- /dev/null +++ b/pyrit/setup/plugin_loader.py @@ -0,0 +1,111 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Activate a configured plug-in by running its ``PyRITInitializer`` first. + +A plug-in is a config-declared pointer to a ``PyRITInitializer`` that lives outside the +PyRIT tree. ``ConfigurationLoader`` prepends this privileged initializer ahead of the +user-configured ones, so the plug-in's components exist before any later initializer or +lazy catalog consumer runs. The plug-in initializer itself owns registration of every +scenario, attack technique, dataset, and default target it wants discoverable; PyRIT does +no component discovery here. +""" + +from __future__ import annotations + +import asyncio +import importlib +import logging +import sys +from typing import TYPE_CHECKING + +from pyrit.exceptions import PluginLoadError, PluginSourceNotFoundError +from pyrit.setup.pyrit_initializer import PyRITInitializer + +if TYPE_CHECKING: + from collections.abc import Sequence + + from pyrit.setup.plugin_spec import PluginSpec + +logger = logging.getLogger(__name__) + +_REMEDIATION = "Fix or remove the plug-in entry from .pyrit_conf, then restart PyRIT." + + +class PluginInitializer(PyRITInitializer): + """Privileged, config-owned initializer that runs one plug-in's ``PyRITInitializer`` first.""" + + def __init__(self, *, plugins: Sequence[PluginSpec]) -> None: + """ + Initialize the privileged plug-in loader. + + Args: + plugins (Sequence[PluginSpec]): The normalized plug-in declarations. + + Raises: + ValueError: If anything other than exactly one plug-in is supplied. + """ + super().__init__() + if len(plugins) != 1: + raise ValueError("PluginInitializer requires exactly one plug-in.") + self._spec = plugins[0] + + @property + def plugins(self) -> list[PluginSpec]: + """The normalized plug-in declarations.""" + return [self._spec] + + async def initialize_async(self) -> None: + """ + Run the configured plug-in's initializer before user-configured initializers. + + Raises: + PluginLoadError: If the source cannot be prepared, the initializer cannot be + resolved, or the initializer raises. Fails closed — a restart is required + after fixing the configuration or artifact. + """ + initializer = await asyncio.to_thread(self._resolve_initializer) + try: + await initializer.initialize_async() + except Exception as exc: + raise PluginLoadError(f"Plug-in '{self._spec.name}' initializer failed: {exc} {_REMEDIATION}") from exc + logger.info("Activated plug-in '%s' via %s.", self._spec.name, self._spec.initializer) + + def _resolve_initializer(self) -> PyRITInitializer: + """ + Anchor the source root on ``sys.path`` and import the configured initializer. + + Returns: + PyRITInitializer: A fresh instance of the resolved initializer subclass. + + Raises: + PluginSourceNotFoundError: If the source path does not exist. + PluginLoadError: If the dotted initializer cannot be imported or is not a + ``PyRITInitializer`` subclass. + """ + spec = self._spec + if not spec.source.exists(): + raise PluginSourceNotFoundError(f"Plug-in '{spec.name}' source path does not exist: {spec.source}") + + root = spec.source if spec.source.is_dir() else spec.source.parent + root_str = str(root) + if root_str not in sys.path: + sys.path.insert(0, root_str) + + module_path, _, class_name = spec.initializer.rpartition(".") + try: + module = importlib.import_module(module_path) + initializer_cls = getattr(module, class_name) + except Exception as exc: + raise PluginLoadError( + f"Could not import plug-in '{spec.name}' initializer '{spec.initializer}' from {root}: " + f"{exc} {_REMEDIATION}" + ) from exc + + if not (isinstance(initializer_cls, type) and issubclass(initializer_cls, PyRITInitializer)): + raise PluginLoadError( + f"Plug-in '{spec.name}' initializer '{spec.initializer}' is not a PyRITInitializer subclass. " + f"{_REMEDIATION}" + ) + return initializer_cls() diff --git a/pyrit/setup/plugin_spec.py b/pyrit/setup/plugin_spec.py new file mode 100644 index 0000000000..32998864e4 --- /dev/null +++ b/pyrit/setup/plugin_spec.py @@ -0,0 +1,76 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Static plug-in configuration model shared by setup entry points.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from pyrit.models import validate_registry_name + + +@dataclass(frozen=True) +class PluginSpec: + """A normalized plug-in declaration: a source root plus a dotted initializer.""" + + name: str + source: Path + initializer: str + + @classmethod + def from_config(cls, entry: dict[str, Any], *, base_dir: Path | None = None) -> PluginSpec: + """ + Normalize one explicit ``.pyrit_conf`` plug-in entry. + + Args: + entry (dict[str, Any]): The YAML plug-in mapping. + base_dir (Path | None): Directory used to resolve a relative ``source`` path. + + Returns: + PluginSpec: The normalized plug-in declaration. + + Raises: + ValueError: If the mapping is malformed. + """ + if not isinstance(entry, dict): + raise ValueError(f"Plug-in entry must be a mapping, got {type(entry).__name__}.") + + allowed = {"name", "source", "initializer"} + unexpected = set(entry) - allowed + if unexpected: + raise ValueError(f"Plug-in entry has unexpected key(s): {sorted(unexpected)}. Allowed: {sorted(allowed)}.") + + name = entry.get("name") + if not isinstance(name, str): + raise ValueError("Plug-in entry requires a string 'name'.") + validate_registry_name(name) + + source = entry.get("source") + if not isinstance(source, str) or not source: + raise ValueError("Plug-in entry requires a non-empty string 'source' path.") + + initializer = entry.get("initializer") + if not isinstance(initializer, str) or "." not in initializer: + raise ValueError( + "Plug-in entry requires a dotted 'initializer' path to a PyRITInitializer subclass " + "(e.g. 'my_package.setup.MyInitializer')." + ) + + path = Path(source).expanduser() + if not path.is_absolute(): + path = (base_dir or Path.cwd()) / path + path = path.resolve() + + return cls(name=name, source=path, initializer=initializer) + + def to_config(self) -> dict[str, str]: + """ + Serialize this spec to the explicit YAML-style mapping. + + Returns: + dict[str, str]: The normalized configuration mapping. + """ + return {"name": self.name, "source": str(self.source), "initializer": self.initializer} 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..10b01e92ff --- /dev/null +++ b/tests/integration/setup/test_plugin_loader_integration.py @@ -0,0 +1,152 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Integration test for the initializer-pointing plug-in path. + +Exercises the full source-plug-in flow against real registries (no mocks): a source +package whose ``PyRITInitializer`` registers a private scenario and a private attack +technique becomes discoverable through the standard registries after the privileged +plug-in phase runs. +""" + +import sys +import textwrap +from pathlib import Path + +import pytest + +from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry +from pyrit.setup import PluginSpec +from pyrit.setup.plugin_loader import PluginInitializer + +_PLUGIN_PACKAGE = "my_redteam_plugin" + + +def _write_source_plugin(root: Path) -> None: + """Write a source package whose initializer registers a scenario and a technique.""" + pkg = root / _PLUGIN_PACKAGE + pkg.mkdir(parents=True, exist_ok=True) + (pkg / "__init__.py").write_text("", encoding="utf-8") + (pkg / "scenario.py").write_text( + textwrap.dedent( + ''' + from functools import cache + + from pyrit.executor.attack import PromptSendingAttack + from pyrit.models import SeedObjective + from pyrit.registry import AttackTechniqueRegistry + from pyrit.scenario.core import ( + AttackTechniqueFactory, + DatasetAttackConfiguration, + Scenario, + ) + from pyrit.score import SubStringScorer + + + @cache + def _build_private_technique(): + return AttackTechniqueRegistry.build_technique_class_from_factories( + class_name="PrivateTechnique", + factories=[ + AttackTechniqueFactory(name="private_prompt_sending", attack_class=PromptSendingAttack) + ], + aggregate_tags={}, + ) + + + class PrivateScenario(Scenario): + """Private scenario for the plug-in integration test.""" + + VERSION = 1 + + def __init__(self, *, scenario_result_id=None): + technique_class = _build_private_technique() + super().__init__( + version=self.VERSION, + technique_class=technique_class, + default_technique=technique_class.ALL, + default_dataset_config=DatasetAttackConfiguration( + seeds=[SeedObjective(value="integration objective")] + ), + objective_scorer=SubStringScorer(substring="integration"), + scenario_result_id=scenario_result_id, + ) + + async def _build_atomic_attacks_async(self, *, context): + return [] + ''' + ), + encoding="utf-8", + ) + (pkg / "bootstrap.py").write_text( + textwrap.dedent( + f''' + from pyrit.executor.attack import PromptSendingAttack + from pyrit.registry import AttackTechniqueRegistry, ScenarioRegistry + from pyrit.scenario.core import AttackTechniqueFactory + from pyrit.setup.pyrit_initializer import PyRITInitializer + + from {_PLUGIN_PACKAGE}.scenario import PrivateScenario + + + class PrivateInitializer(PyRITInitializer): + """Register the private scenario and a private attack technique.""" + + async def initialize_async(self) -> None: + AttackTechniqueRegistry.get_registry_singleton().register_from_factories( + [AttackTechniqueFactory(name="operation_foobar", attack_class=PromptSendingAttack)] + ) + ScenarioRegistry.get_registry_singleton().register_class( + PrivateScenario, name="my_redteam.private" + ) + ''' + ), + encoding="utf-8", + ) + + +@pytest.fixture(autouse=True) +def _isolate_registries_and_imports(): + original_path = list(sys.path) + original_modules = set(sys.modules) + AttackTechniqueRegistry.reset_registry_singleton() + ScenarioRegistry.reset_registry_singleton() + yield + AttackTechniqueRegistry.reset_registry_singleton() + ScenarioRegistry.reset_registry_singleton() + sys.path[:] = original_path + for name in list(sys.modules): + if name not in original_modules and name.split(".", 1)[0] == _PLUGIN_PACKAGE: + del sys.modules[name] + + +@pytest.mark.run_only_if_all_tests +async def test_source_plugin_initializer_registers_real_components(tmp_path: Path) -> None: + _write_source_plugin(tmp_path) + spec = PluginSpec( + name="my_redteam", + source=tmp_path, + initializer=f"{_PLUGIN_PACKAGE}.bootstrap.PrivateInitializer", + ) + + await PluginInitializer(plugins=[spec]).initialize_async() + + scenario_registry = ScenarioRegistry.get_registry_singleton() + assert "my_redteam.private" in scenario_registry.get_class_names() + assert scenario_registry.get_class("my_redteam.private").__name__ == "PrivateScenario" + assert "operation_foobar" in AttackTechniqueRegistry.get_registry_singleton().get_factories() + + +@pytest.mark.run_only_if_all_tests +async def test_source_plugin_missing_initializer_fails_closed(tmp_path: Path) -> None: + from pyrit.exceptions import PluginLoadError + + _write_source_plugin(tmp_path) + spec = PluginSpec( + name="my_redteam", + source=tmp_path, + initializer=f"{_PLUGIN_PACKAGE}.bootstrap.DoesNotExist", + ) + + with pytest.raises(PluginLoadError): + await PluginInitializer(plugins=[spec]).initialize_async() diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index cb196b60ce..112adce262 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -146,10 +146,13 @@ def test_parse_args_invalid_max_retries(self): with pytest.raises(SystemExit): pyrit_scan.parse_args(["test_scenario", "--max-retries", "-1"]) - def test_parse_args_help_flag(self): + def test_parse_args_help_flag(self, capsys): with pytest.raises(SystemExit) as exc_info: pyrit_scan.parse_args(["--help"]) assert exc_info.value.code == 0 + help_text = capsys.readouterr().out + assert "Private scenario plug-ins" in help_text + assert "doc/getting_started/plugins.md" in help_text def test_parse_args_with_target(self): args = pyrit_scan.parse_args(["test_scenario", "--target", "my_target"]) diff --git a/tests/unit/exceptions/test_exceptions.py b/tests/unit/exceptions/test_exceptions.py index e228efed32..374781517b 100644 --- a/tests/unit/exceptions/test_exceptions.py +++ b/tests/unit/exceptions/test_exceptions.py @@ -15,6 +15,8 @@ EmptyResponseException, InvalidJsonException, MissingPromptPlaceholderException, + PluginLoadError, + PluginSourceNotFoundError, PyritException, RateLimitException, handle_bad_request_exception, @@ -66,6 +68,11 @@ def test_invalid_json_exception_initialization(): assert str(ex) == "Status Code: 500, Message: Invalid JSON Response" +def test_plugin_source_not_found_error_inherits_plugin_load_error() -> None: + """The specific plug-in failure must be catchable through ``PluginLoadError``.""" + assert issubclass(PluginSourceNotFoundError, PluginLoadError) + + def test_bad_request_exception_process_exception(caplog): ex = BadRequestException() with caplog.at_level(logging.ERROR): diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 3c33fac5dc..6c97a2218d 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -710,3 +710,114 @@ 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`` plug-in list.""" + + 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() == [] + + def test_explicit_source_mapping(self, tmp_path: pathlib.Path): + config = ConfigurationLoader( + plugins=[ + { + "name": "operation_foobar", + "source": str(tmp_path / "operation_foobar"), + "initializer": "operation_foobar.setup.OperationInitializer", + } + ] + ) + + spec = config.resolve_plugins()[0] + assert spec.name == "operation_foobar" + assert spec.source == (tmp_path / "operation_foobar").resolve() + assert spec.initializer == "operation_foobar.setup.OperationInitializer" + + def test_multiple_plugins_rejected(self, tmp_path: pathlib.Path): + with pytest.raises(ValueError, match="one plug-in"): + ConfigurationLoader( + plugins=[ + {"name": "first", "source": str(tmp_path / "first"), "initializer": "first.Init"}, + {"name": "second", "source": str(tmp_path / "second"), "initializer": "second.Init"}, + ] + ) + + 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, tmp_path: pathlib.Path): + """``from_dict`` wires the explicit plug-in mapping onto the loader.""" + config = ConfigurationLoader.from_dict( + { + "plugins": [ + { + "name": "partner", + "source": str(tmp_path / "partner"), + "initializer": "partner.setup.PartnerInitializer", + } + ] + } + ) + assert config.resolve_plugins()[0].initializer == "partner.setup.PartnerInitializer" + + def test_yaml_plugin_path_resolves_relative_to_config(self, tmp_path: pathlib.Path): + config_dir = tmp_path / "config" + config_dir.mkdir() + config_path = config_dir / ".pyrit_conf" + config_path.write_text( + """plugins: + - name: operation_foobar + source: ../plugins/operation_foobar + initializer: operation_foobar.setup.OperationInitializer +""", + encoding="utf-8", + ) + + config = ConfigurationLoader.from_yaml_file(config_path) + + assert config.resolve_plugins()[0].source == (tmp_path / "plugins" / "operation_foobar").resolve() + + async def test_initialize_prepends_privileged_plugin_initializer(self, tmp_path: pathlib.Path): + from pyrit.setup.plugin_loader import PluginInitializer + + user_initializer = mock.MagicMock() + config = ConfigurationLoader( + memory_db_type="in_memory", + plugins=[ + { + "name": "partner", + "source": str(tmp_path / "partner"), + "initializer": "partner.setup.PartnerInitializer", + } + ], + ) + + with ( + mock.patch.object(config, "resolve_initializers", return_value=[user_initializer]), + mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") as core_init, + ): + await config.initialize_pyrit_async() + + initializers = core_init.call_args.kwargs["initializers"] + assert isinstance(initializers[0], PluginInitializer) + assert initializers[0].plugins == config.resolve_plugins() + assert initializers[1] is user_initializer + assert "plugins" not in core_init.call_args.kwargs + + async def test_initialize_without_plugins_uses_only_configured_initializers(self): + user_initializer = mock.MagicMock() + config = ConfigurationLoader(memory_db_type="in_memory") + + with ( + mock.patch.object(config, "resolve_initializers", return_value=[user_initializer]), + mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") as core_init, + ): + await config.initialize_pyrit_async() + + assert core_init.call_args.kwargs["initializers"] == [user_initializer] diff --git a/tests/unit/setup/test_plugin_loader.py b/tests/unit/setup/test_plugin_loader.py new file mode 100644 index 0000000000..5d87b3f741 --- /dev/null +++ b/tests/unit/setup/test_plugin_loader.py @@ -0,0 +1,120 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for the initializer-pointing ``PluginInitializer``. + +A plug-in resolves to a ``PyRITInitializer`` reached by dotted path from a source root. +The loader anchors the root on ``sys.path``, imports the initializer, and runs it. The +initializer itself owns all registration; the loader only orchestrates and fails closed. +""" + +import sys +import textwrap +from pathlib import Path + +import pytest + +from pyrit.exceptions import PluginLoadError, PluginSourceNotFoundError +from pyrit.registry import AttackTechniqueRegistry +from pyrit.setup import PluginSpec +from pyrit.setup.plugin_loader import PluginInitializer + +_TEST_PACKAGE_ROOTS = {"mock_plugin", "notinit"} + + +def _write_plugin_package(root: Path, *, technique_name: str) -> None: + """Write a source package whose initializer registers one private technique.""" + pkg = root / "mock_plugin" + pkg.mkdir(parents=True, exist_ok=True) + (pkg / "__init__.py").write_text("", encoding="utf-8") + (pkg / "bootstrap.py").write_text( + textwrap.dedent( + f""" + from pyrit.executor.attack import PromptSendingAttack + from pyrit.registry import AttackTechniqueRegistry + from pyrit.scenario.core import AttackTechniqueFactory + from pyrit.setup.pyrit_initializer import PyRITInitializer + + + class MockInitializer(PyRITInitializer): + \"\"\"Register one private attack technique.\"\"\" + + async def initialize_async(self) -> None: + AttackTechniqueRegistry.get_registry_singleton().register_from_factories( + [AttackTechniqueFactory(name="{technique_name}", attack_class=PromptSendingAttack)] + ) + """ + ), + encoding="utf-8", + ) + + +@pytest.fixture(autouse=True) +def _isolate_import_state(): + original_path = list(sys.path) + original_modules = set(sys.modules) + AttackTechniqueRegistry.reset_registry_singleton() + yield + AttackTechniqueRegistry.reset_registry_singleton() + sys.path[:] = original_path + for name in list(sys.modules): + if name not in original_modules and name.split(".", 1)[0] in _TEST_PACKAGE_ROOTS: + del sys.modules[name] + + +async def test_runs_configured_initializer(tmp_path: Path) -> None: + _write_plugin_package(tmp_path, technique_name="operation_alpha") + spec = PluginSpec(name="alpha", source=tmp_path, initializer="mock_plugin.bootstrap.MockInitializer") + + await PluginInitializer(plugins=[spec]).initialize_async() + + assert "operation_alpha" in AttackTechniqueRegistry.get_registry_singleton().get_factories() + + +async def test_missing_source_fails_closed(tmp_path: Path) -> None: + spec = PluginSpec(name="p", source=tmp_path / "nope", initializer="mock_plugin.bootstrap.MockInitializer") + with pytest.raises(PluginSourceNotFoundError): + await PluginInitializer(plugins=[spec]).initialize_async() + + +async def test_unresolvable_initializer_fails_closed(tmp_path: Path) -> None: + _write_plugin_package(tmp_path, technique_name="operation_beta") + spec = PluginSpec(name="p", source=tmp_path, initializer="mock_plugin.bootstrap.DoesNotExist") + with pytest.raises(PluginLoadError): + await PluginInitializer(plugins=[spec]).initialize_async() + + +async def test_non_initializer_target_fails_closed(tmp_path: Path) -> None: + pkg = tmp_path / "notinit" + pkg.mkdir() + (pkg / "__init__.py").write_text("class NotAnInitializer:\n pass\n", encoding="utf-8") + spec = PluginSpec(name="p", source=tmp_path, initializer="notinit.NotAnInitializer") + with pytest.raises(PluginLoadError, match="not a PyRITInitializer"): + await PluginInitializer(plugins=[spec]).initialize_async() + + +async def test_initializer_that_raises_fails_closed(tmp_path: Path) -> None: + pkg = tmp_path / "notinit" + pkg.mkdir() + (pkg / "__init__.py").write_text( + textwrap.dedent( + """ + from pyrit.setup.pyrit_initializer import PyRITInitializer + + + class Boom(PyRITInitializer): + async def initialize_async(self) -> None: + raise RuntimeError("boom") + """ + ), + encoding="utf-8", + ) + spec = PluginSpec(name="p", source=tmp_path, initializer="notinit.Boom") + with pytest.raises(PluginLoadError, match="initializer failed"): + await PluginInitializer(plugins=[spec]).initialize_async() + + +def test_requires_exactly_one_plugin(tmp_path: Path) -> None: + spec = PluginSpec(name="p", source=tmp_path, initializer="m.C") + with pytest.raises(ValueError): + PluginInitializer(plugins=[spec, spec]) diff --git a/tests/unit/setup/test_plugin_spec.py b/tests/unit/setup/test_plugin_spec.py new file mode 100644 index 0000000000..7df4ce607b --- /dev/null +++ b/tests/unit/setup/test_plugin_spec.py @@ -0,0 +1,55 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for the source-only, initializer-pointing ``PluginSpec`` schema.""" + +from pathlib import Path + +import pytest + +from pyrit.setup import PluginSpec + + +def test_from_config_normalizes_source_and_initializer(tmp_path: Path) -> None: + entry = { + "name": "rapid_response", + "source": "pkg_root", + "initializer": "my_redteam.setup.MyInitializer", + } + spec = PluginSpec.from_config(entry, base_dir=tmp_path) + assert spec.name == "rapid_response" + assert spec.source == (tmp_path / "pkg_root").resolve() + assert spec.initializer == "my_redteam.setup.MyInitializer" + + +def test_from_config_keeps_absolute_source(tmp_path: Path) -> None: + absolute = (tmp_path / "artifacts").resolve() + spec = PluginSpec.from_config( + {"name": "p", "source": str(absolute), "initializer": "m.C"}, + base_dir=tmp_path / "elsewhere", + ) + assert spec.source == absolute + + +def test_to_config_round_trips() -> None: + spec = PluginSpec(name="p", source=Path("/opt/plugin").resolve(), initializer="m.C") + assert spec.to_config() == { + "name": "p", + "source": str(Path("/opt/plugin").resolve()), + "initializer": "m.C", + } + + +@pytest.mark.parametrize( + "entry", + [ + {"source": "s", "initializer": "m.C"}, # missing name + {"name": "p", "initializer": "m.C"}, # missing source + {"name": "p", "source": "s"}, # missing initializer + {"name": "p", "source": "s", "initializer": "no_dot"}, # non-dotted initializer + {"name": "p", "source": "s", "initializer": "m.C", "wheel": "x"}, # unexpected key + ], +) +def test_from_config_rejects_malformed_entries(entry: dict) -> None: + with pytest.raises(ValueError): + PluginSpec.from_config(entry)