[FEAT] Resolve Key Vault-Backend Environment References

[Victor Valbuena (ValbuenaVC)]

Description

This PR adds Azure Key Vault-backed environment bootstrapping while retaining python-dotenv as the owner of dotenv parsing and ${NAME} interpolation. Previous iterations of this PR updated .env_example and added tests for drift across the PyRIT repo, .env_example, and the key vault; these have been removed and deferred.

Canonically, .pyrit_conf should now look like this:

env_akv_ref:
  - https://my-vault.vault.azure.net/secrets/pyrit-env
env_akv_strict: true

Each env_akv_ref URL identifies a secret containing a dotenv document. Documents load in list order and may contain literals, interpolation, and complete-value references to scalar Key Vault secrets.

Loading Behavior

Environment sources follow the existing precedence rules:

1. Process environment values are retained.
2. Key Vault bootstrap documents fill missing values in list order.
3. Ordinary env_files fill missing values in list order.
4. Files named .env.local override existing values.

When Key Vault is configured without explicit env_files, auto-discovered ~/.pyrit/.env is ignored and emits a security warning. ~/.pyrit/.env.local remains available for deliberate local overrides.

PyRIT gathers all source candidates before resolving references, so only winning Key Vault references are fetched. Process-origin values remain opaque and are never interpreted as references. Historical non-transactional dotenv behavior is preserved.

When PYTHON_DOTENV_DISABLED is enabled, all Key Vault and dotenv sources are skipped without I/O. Existing process environment values remain unchanged.

Key Vault References

Bootstrap documents and local dotenv files support complete-value references:

ENDPOINT="https://example.openai.azure.com/openai/v1"
MODEL="${PYRIT_MODEL}"
API_KEY="kv:https://my-vault.vault.azure.net/secrets/api-key"
PINNED_KEY="kv:https://my-vault.vault.azure.net/secrets/api-key/version-id"

kv: is canonical. akv:, azure_key_vault:, and env_akv_ref: remain compatibility aliases.

References require full HTTPS secret URLs. References originating from a bootstrap document must remain in that document’s vault, keeping the document within one trust boundary. Explicit local files may reference any validated supported Azure Key Vault URL.

Supported vault DNS suffixes are:

• .vault.azure.net
• .vault.azure.cn
• .vault.usgovcloudapi.net

Resolution is one hop: fetched child-secret values are terminal and are not interpreted as additional references.

Validation And Failures

Strict mode rejects malformed bootstrap entries and malformed references. Non-strict mode warns and falls through to the next applicable source candidate. Authentication, authorization, transport, missing-secret, and missing-value failures always stop initialization.

Key Vault clients use asynchronous retries with exponential backoff. KeyVaultInitializationException preserves upstream HTTP status codes and the original exception. Failures without an HTTP response use PyRIT’s generic status 500 fallback.

Debug Export

Runtime initialization never writes secrets to disk. An explicit helper exports AKV-only configuration for debugging:

python -m build_scripts.export_akv_environment `
  --secret-url https://my-vault.vault.azure.net/secrets/pyrit-env

The helper writes ~/.pyrit/.env_akv, which is not auto-loaded and is Git-ignored. Publication is owner-restricted where supported, atomic, and refuses to overwrite an existing path.

Validation

• Setup, exporter, and import-hygiene tests: 506 passed, 2 skipped
• Environment-loading tests: 83 passed
• Ruff lint and formatting: passed
• ty type checking: passed
• Documentation validation: passed
• Async-suffix and reST-role checks: passed
• Git whitespace checks: passed

[Copilot]

Pull request overview

Adds Azure Key Vault-backed environment bootstrapping with recursive reference resolution, precedence handling, warnings, and documentation.

Changes:

• Resolves env:, kv:, aliases, and escaped literals.
• Adds environment-source validation and AKV/local-file precedence.
• Expands tests and configuration documentation.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
pyrit/setup/initialization.py	Implements AKV loading and reference resolution.
tests/unit/setup/test_initialization.py	Tests environment initialization behavior.
doc/getting_started/pyrit_conf.md	Documents loading precedence and AKV references.
.pyrit_conf_example	Updates example AKV configuration guidance.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

Suppressed comments (5)

tests/unit/setup/test_initialization.py:380

• This patch target is no longer called by initialize_pyrit_async, leaving the output assertion dependent on any real default environment files. Patch _resolve_environment_files instead so unrelated local files cannot add output or trigger reference resolution.

    @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=True)

pyrit/setup/initialization.py:537

• Direct callers using the former list-shaped env_akv_ref reach .strip() here and get AttributeError, rather than the deliberate ValueError used by ConfigurationLoader. Validate the runtime type before calling string methods so this public API rejects legacy values consistently.

        if not env_akv_ref.strip():

pyrit/setup/initialization.py:439

• This merge is case-sensitive even on Windows. For example, ambient Path=old plus a winning PATH=new leaves both keys, so env:Path returns the ambient exact match and violates the documented merged-source precedence. Normalize keys on Windows while applying values last.

    reference_environment = {**ambient_environment, **values}

tests/unit/setup/test_initialization.py:372

• initialize_pyrit_async no longer calls _load_environment_files, so this patch is inert and the test can read real ~/.pyrit files (and even resolve their Key Vault references). Patch the resolver now used by initialization to keep the unit test isolated.

This issue also appears on line 380 of the same file.

    @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=True)

doc/getting_started/pyrit_conf.md:174

• This row contradicts both the implementation and the earlier AKV precedence section: when env_files is omitted, initialization loads both .env and .env.local after the bootstrap. Remove the claim that only .env.local is loaded.

| Omitted or `null` | Load default `~/.pyrit/.env` and `~/.pyrit/.env.local`, or only `.env.local` after an AKV root |

[Bashir Partovi (bashirpartovi)]

Thanks Victor Valbuena (@ValbuenaVC) for working through all the earlier feedback on this. I pulled the latest changes, read through the existing review threads, and reran the focused setup/configuration tests locally. The source ordering is much clearer now, and keeping python-dotenv responsible for parsing was the right simplification.

I found a few cases that I think still need to be fixed before this merges. Most of them sit around the new debug-export path or configuration migration, so the current tests and green CI do not exercise them.

First, env_akv_strict and env_akv_write_env are annotated as booleans but are not validated at runtime. A config like this is accepted:

env_akv_write_env: "false"
env_akv_strict: "false"

Both values remain strings, and because a non-empty string is treated as True, the first one actually enables writing plaintext secrets to ~/.pyrit/.env. I think these fields should reject anything that is not a real bool, both in ConfigurationLoader and at the direct initialization boundary. This one feels especially important because it can turn on a sensitive feature when the user explicitly wrote false.

There are also two data-integrity problems in the generated debug file:

1. _serialize_terminal_dotenv_value() does not escape backslashes before putting a secret in a single-quoted dotenv value. In a round-trip through python-dotenv, two consecutive backslashes become one and four become two. That means some passwords, tokens, paths, or other arbitrary secret values will be silently changed in the file.

2. References created through interpolation are resolved at runtime but are not rewritten in the generated file. For example:

   A=kv:https://vault.vault.azure.net/secrets/key
   B=${A}
   A=literal

   At runtime, B becomes the fetched secret. In the generated file, B=${A} is retained, so reloading the file gives B the original kv: URI instead. The renderer currently checks whether the raw assignment text looks like a reference, while resolution checks the interpolated value. I think both runtime application and debug rendering need to consume the same assignment-level resolution result.

The no-clobber guarantee for ~/.pyrit/.env also has a race. The existence check happens before the Key Vault network calls, but _write_akv_env_file() later publishes with unconditional os.replace(). I simulated another process creating .env during the fetch, and the newly created file was overwritten. A second exists() check would still race; the final publish needs an atomic create-if-absent operation. The early check can stay for fast feedback, but the writer needs to enforce the guarantee authoritatively.

I also found a non-strict fallback case that leaves a Key Vault URI in the environment instead of resolving it. If an ordinary file provides a valid kv: reference and a later .env.local overrides it with malformed kv:short, non-strict mode restores the earlier URI and then immediately continues. My probe ended with the literal earlier kv: URI in os.environ and made zero child-secret reads. Since the documented behavior is to skip the malformed assignment, I think the loader needs an ordered candidate chain per variable, rather than a single scalar fallback, so it can discard the malformed winner and fully resolve the next candidate.

The image/TTS environment-variable changes look breaking as well. On main, PyRIT accepts the documented OPENAI_IMAGE_*1/2 and OPENAI_TTS_*1/2 variables. This branch removes those inputs completely and replaces them with AZURE_OPENAI_*1/2, without aliases or a migration warning. Existing configurations will silently stop registering those targets.

There is a provider mismatch in the new names too: openai_image_platform and openai_tts_platform now read Azure-prefixed variables whose example values are Azure endpoints. The second TTS example is https://xxxxx.openai.azure.com/v1, which is also missing the /openai/v1 path that PyRIT recommends and passes through unchanged. I would preserve the old variables as deprecated aliases and keep the *_platform registry entries backed by actual platform OpenAI settings. A larger provider-explicit rename would be safer as a separate migration with compatibility tests.

Two smaller contract gaps are worth covering in the same pass:

• With PYTHON_DOTENV_DISABLED=true, load_dotenv() returns false and the AKV loader exits before resolving child references, but debug mode can still write that raw document and call it fully resolved. Either dotenv-disabled mode should suppress this publication, the option combination should be rejected, or output resolution should be separated from environment mutation.
• test_env_example_names_are_referenced_in_repository is one-way: it begins with names already present in .env_example and looks for weak textual references elsewhere. It cannot catch a code-required name missing from the example, removed compatibility names, duplicate assignments, provider mismatches, or invalid endpoint families. Deriving required names from TARGET_CONFIGS, comparing both directions with an explicit allowlist, and checking duplicates/provider URLs structurally would make this a real drift test.

I don't think the documented non-transactional updates, one-hop terminal references, sequential precedence, or lack of child-secret caching are problems by themselves. Those are clear choices now. The recurring issue is that precedence is represented by mutating os.environ plus one fallback value per variable, which loses information needed by fallback resolution and debug rendering. A small internal ordered assignment/candidate model would address both without changing the public API or replacing python-dotenv.

I also reproduced the boolean, backslash, interpolation, no-clobber, and fallback cases directly against the latest changes.

[Behnam (behnam-o)]

I think this PR can be shortened by quite a lot, and I'd keep the functionality simple, and changes minimal (i.e. just allow pyrit to understand when a value in an .env file or secret - the ones we load today- starts with akv: , and try resolving that value when setting the env variable)

I the other bits (write back an akv to a file, allow-same-akv-ref, etc. ) adding the kind of convenience that could cause more confusion.

[Behnam (behnam-o)]

Thanks for addressing the comments. As we discussed offline, I think you should consider limiting the changes to the minimum possible that achieve the desired feature (i.e. instead of blindly setting env vars from .env documents, try to resolve those that start with akv: ) - as implemented right now, it does achieve it, but with quite a lot of overhead and fixes that don't necessarily belong to this particular feature (for example, updating the parse_url function, supporting array of akv env documents instead of one, etc. that I consider parallel fixes/features)

[Victor Valbuena (ValbuenaVC)]

Thanks for addressing the comments. As we discussed offline, I think you should consider limiting the changes to the minimum possible that achieve the desired feature (i.e. instead of blindly setting env vars from .env documents, try to resolve those that start with akv: ) - as implemented right now, it does achieve it, but with quite a lot of overhead and fixes that don't necessarily belong to this particular feature (for example, updating the parse_url function, supporting array of akv env documents instead of one, etc. that I consider parallel fixes/features)

Agreed, and thanks for the feedback + approval! There are some sub-features that I think should accompany this change given that it's setting a new default for environment variables in the library (so updating things like URL parsing are appropriate I feel, despite being somewhat different concerns). But I'll minimize the blast radius as much as I can. You're right that it's a bit too broad, and the changes should be narrower.