[DRAFT] FEAT Plug-in mechanism for private scenarios#2131
Open
ValbuenaVC wants to merge 2 commits into
Open
Conversation
Adds a plug-in mechanism so an operator can run non-disclosable scenarios (and self-contained datasets) from a stock public PyRIT install, from a pre-built wheel referenced by PLUGIN_WHEEL, without those components living in the public repo. Plug-in loading runs as a guaranteed-first phase inside initialize_pyrit_async: after central memory is set and before the configured initializers, so plug-in datasets/scenarios register before LoadDefaultDatasets and PreloadScenarioMetadata read the registry. Ordering is true by construction, independent of .pyrit_conf list position. The wheel is extracted (stdlib zipfile - never pip/.venv) to .plugin/<name>/, prepended to sys.path, imported (dataset providers self-register), and its bootstrap (a top-level register() callable or a shipped PyRITInitializer subclass) is run; the loader then asserts the plug-in registered something. No-op when PLUGIN_WHEEL is unset. Fails closed by default; fail-open via PLUGIN_FAIL_OPEN or initialize_pyrit_async(plugin_fail_open=True). Guards against silent-failure modes: extraction (not zipimport) so __file__-relative datasets resolve, atomic extraction, submodule discovery so components register even if the package __init__ does not import them, a package-shadowing guard when an installed package of the same name is importable, a loud warning when a plug-in dataset name collides with an existing name (the resolver is memory-authoritative, so the plug-in copy would otherwise be silently ignored), and rollback of sys.path / sys.modules / registries on failure. Wiring: .gitignore keeps .plugin/ (ignores its contents); .env_example documents the PLUGIN_* variables. Tests build a mock wheel (no dependency on any real plug-in) covering extraction, registration, ordering, collisions, and failure modes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a wheel-based plug-in loading mechanism to PyRIT initialization so operators can ship private scenarios (and dataset providers) outside the public repo, while ensuring plug-ins load before other initializers consume registries.
Changes:
- Introduces
pyrit.setup.plugin_loaderto extract a configured.whlinto.plugin/, import it, run a bootstrap (register()orPyRITInitializer), and validate registration with rollback-on-failure. - Hooks plug-in loading into
initialize_pyrit_asyncas a guaranteed-first phase (after memory is set, before other initializers), with aplugin_fail_openoverride. - Adds extensive unit coverage for loader behavior/guards and documents operator configuration via
.env_exampleand.pyrit_conf_example; ensures.plugin/is kept but ignored in git.
Reviewed changes
Copilot reviewed 6 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
tests/unit/setup/test_plugin_loader.py |
New unit test suite that builds a mock wheel and exercises loader spine + guardrails (fail-open/closed, rollback, collisions, cache, package inference). |
pyrit/setup/plugin_loader.py |
New plug-in loader implementation: extraction, import/submodule discovery, bootstrap execution, collision warnings, and rollback. |
pyrit/setup/initialization.py |
Adds the plug-in load phase into initialize_pyrit_async and exposes plugin_fail_open. |
.pyrit_conf_example |
Documents that plug-ins are not configured via .pyrit_conf and are loaded automatically during initialization. |
.plugin/.gitkeep |
Keeps .plugin/ directory present in the repo while allowing extracted artifacts to be ignored. |
.gitignore |
Ignores extracted plug-in artifacts under .plugin/ while retaining .gitkeep. |
.env_example |
Documents PLUGIN_WHEEL, PLUGIN_PACKAGE, PLUGIN_FAIL_OPEN, and PLUGIN_DIR configuration and trust boundary. |
Comment on lines
+250
to
+251
| with zipfile.ZipFile(wheel_path) as wheel_zip: | ||
| wheel_zip.extractall(tmp_dir) |
Comment on lines
+170
to
+171
| extract_dir = self._extract_wheel(wheel_path=wheel_path) | ||
| package_name = self._resolve_package_name(extract_dir=extract_dir) |
Comment on lines
+581
to
+584
| for key, cls in list(SeedDatasetProvider._registry.items()): | ||
| if key not in provider_snapshot and _module_owned_by(cls, package_name): | ||
| del SeedDatasetProvider._registry[key] | ||
|
|
Comment on lines
+585
to
+591
| removed_scenario = False | ||
| for name, cls in list(scenario_registry._classes.items()): | ||
| if name not in scenario_snapshot and _module_owned_by(cls, package_name): | ||
| del scenario_registry._classes[name] | ||
| removed_scenario = True | ||
| if removed_scenario: | ||
| scenario_registry._metadata_cache = None |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Operators who use PyRIT extensively sometimes need to run bespoke scenario logic that cannot be published to the public repo and cannot be represented as database rows. Today there is no delivery path for that code into a stock public install. This PR adds a minimal plug-in mechanism so a private scenario can be shipped as a pre-built Python wheel and loaded at initialization, without those components living in the public repository. Datasets come along as a secondary benefit (a plug-in's dataset providers self-register on import), but the primary purpose is distributing private scenario code.
TL; DR How it works. Plug-in loading runs as a guaranteed-first phase inside
initialize_pyrit_async, after central memory is set and before the configured initializers run. Because it is a fixed phase rather than a user-ordered.pyrit_confinitializer, "plug-ins register before anything reads the registry" is true by construction. This deliberately avoids putting a load-bearing ordering requirement into user-editable config (a silent-failure trap: a misordered conf would let scenario metadata warm before plug-in scenarios registered). SetPLUGIN_WHEELto a wheel on disk and the loader extracts it (stdlibzipfile, never pip or.venv) into.plugin/<name>/, prepends that dir tosys.path, imports the package (dataset providers self-register), runs the plug-in's bootstrap (a top-levelregister()callable or a shippedPyRITInitializersubclass that registers scenarios), and asserts the plug-in registered something. It is a no-op whenPLUGIN_WHEELis unset.Optional config:
PLUGIN_PACKAGE(auto-detected top-level package),PLUGIN_FAIL_OPEN(orinitialize_pyrit_async(plugin_fail_open=True)), andPLUGIN_DIR. It fails closed by default; fail-open skips a broken plug-in with a warning.Why extraction, not install. Plug-in seed files load via
__file__-relative paths; a raw.whlonsys.path(zipimport) makes those resolve inside the zip and the datasets silently vanish. Runtime pip/setuptools are also not guaranteed present (uv venvs ship without them). A wheel is just a zip, so stdlib extraction works on any environment.Silent-failure guards (the non-obvious part). Each guard exists because it is a way a plug-in could load "successfully" yet be wrong: extraction rather than zipimport; atomic extraction (temp dir then move) so a crash never leaves a poisoned cache; submodule discovery so providers and the bootstrap register even if the package
__init__does not import them; a package-shadowing guard that fails loudly if an installed package of the same name shadows the wheel; fail-loud when a wheel imports cleanly but registers nothing; and full rollback ofsys.path, imported modules, and registries on any failure, so a failed (or fail-open) load leaves no trace.Dataset guardrail (worth a careful look). The dataset resolver treats central memory as authoritative and only consults a provider when memory has no seeds for that name, so a plug-in cannot override an existing same-named dataset (a scan would silently use the existing copy). To surface this, the loader warns loudly at load time (greppable
PLUGIN DATASET SHADOWED:prefix, naming both providers) when a plug-indataset_namecollides with an existing registered provider's name. The check compares the provider registry rather than live memory on purpose: at that 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 source is undecidable at load time). Precedence for names that exist only in shared storage without a provider is out of scope here and belongs to a separate dataset-management effort.Trust boundary. Running a plug-in bootstrap executes third-party code in the PyRIT process, so whoever can set
PLUGIN_*or write the backend's.envcan run code on the host. The config is the trust boundary and should be treated as sensitive.Scoped out of this PR (natural follow-ups): custom attacks/converters/targets as plug-ins, a live-iteration/editable-install workflow, nested dependency resolution, multi-plug-in composition, and
.pyrit_confcomposition.Tests and Documentation
Tests:
tests/unit/setup/test_plugin_loader.pybuilds a mock plug-in wheel at test time (no dependency on any real plug-in) and covers the full spine plus every guard: extract-not-zipimport, provider self-registration on import (including via submodule discovery), scenarioregister_classsurviving lazy discovery and coexisting with built-ins, the no-arg-instantiable contract, the phase running after memory is set and before the configured initializers, package-name resolution, atomic-extraction cache reuse, the package-shadowing guard, the dataset-name collision warning (and no false positive on the plug-in's own reload), and all failure modes (fail-closed, fail-open, loud-on-empty, rollback). 43 test cases; the fulltests/unit/setup/suite passes (323). ruff, ty, and the async-suffix check are clean.Documentation: operator-facing configuration is documented inline in
.env_example(thePLUGIN_*variables) and.pyrit_conf_example(a note that plug-ins are not configured there). No Jupyter notebooks changed, so JupyText was not run.