[FEAT] Resolve Key Vault-Backend Environment References and Update .env_example

.pyrit_conf_example

@@ -79,33 +79,20 @@ operation: op_trash_panda
 #   - /path/to/my_custom_initializer.py
 #   - ./local_initializer.py
 
-# Environment Files
-# -----------------
-# List of .env file paths to load during initialization.
-# Later files override values from earlier files.
-#
-# Behavior:
-# - Omit this field (or set to null): Load default .env and .env.local from ~/.pyrit/ if they exist
-# - Set to []: Explicitly load NO environment files
-# - Set to list of paths: Load only the specified files
-#
-# Example:
-# env_files:
-#   - /path/to/.env
-#   - /path/to/.env.local
-
-# Azure Key Vault Environment References
-# ---------------------------------------
-# List of AKV secret URLs to load during initialization.
-# Each secret's value must be the full contents of a .env file.
-# Loaded after env_files, so AKV secrets take precedence.
-# Authentication uses DefaultAzureCredential (managed identity, Azure CLI, etc.).
-#
-# Requires: pip install azure-keyvault-secrets
-#
-# Example:
+# Environment Configuration
+# -------------------------
+# Azure Key Vault is recommended for shared and deployed configurations.
+# See doc/getting_started/pyrit_conf.md for loading order, references, and migration guidance.
 # env_akv_ref:
 #   - https://my-vault.vault.azure.net/secrets/my-pyrit-env
+# env_akv_strict: true
+# env_akv_write_env: false  # Opt in to writing ~/.pyrit/.env for inspection.
+
+# Local dotenv files remain supported. Explicit paths load after Key Vault and override it.
+# Omit env_files to load ~/.pyrit/.env and ~/.pyrit/.env.local, or use [] for no local files.
+# env_files:
+#   - /path/to/.env.local
+# env_akv_strict: false
 
 # Max Concurrent Scenario Runs
 # ----------------------------

build_scripts/env_local_integration_test

@@ -7,12 +7,12 @@ OPENAI_CHAT_ENDPOINT=${AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT}
 OPENAI_CHAT_KEY=${AZURE_OPENAI_INTEGRATION_TEST_KEY}
 OPENAI_CHAT_MODEL=${AZURE_OPENAI_INTEGRATION_TEST_MODEL}
 
-OPENAI_IMAGE_ENDPOINT=${OPENAI_IMAGE_ENDPOINT2}
-OPENAI_IMAGE_API_KEY=${OPENAI_IMAGE_API_KEY2}
-OPENAI_IMAGE_MODEL=${OPENAI_IMAGE_MODEL2}
+OPENAI_IMAGE_ENDPOINT=${AZURE_OPENAI_IMAGE_ENDPOINT2}
+OPENAI_IMAGE_API_KEY=${AZURE_OPENAI_IMAGE_API_KEY2}
+OPENAI_IMAGE_MODEL=${AZURE_OPENAI_IMAGE_MODEL2}
 
-OPENAI_TTS_ENDPOINT=${OPENAI_TTS_ENDPOINT2}
-OPENAI_TTS_KEY=${OPENAI_TTS_KEY2}
+OPENAI_TTS_ENDPOINT=${AZURE_OPENAI_TTS_ENDPOINT2}
+OPENAI_TTS_KEY=${AZURE_OPENAI_TTS_KEY2}
 
 AZURE_SQL_DB_CONNECTION_STRING=${AZURE_SQL_DB_CONNECTION_STRING_TEST}
 AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL=${AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST}

doc/code/executor/gcg/1_gcg_azure_ml.ipynb

@@ -57,7 +57,7 @@
    "source": [
     "import os\n",
     "\n",
-    "from pyrit.setup.initialization import _load_environment_files\n",
+    "from pyrit.setup.akv_initialization import _load_environment_files\n",
     "\n",
     "_load_environment_files(env_files=None)\n",
     "\n",

doc/code/executor/gcg/1_gcg_azure_ml.py

@@ -29,7 +29,7 @@
 # %%
 import os
 
-from pyrit.setup.initialization import _load_environment_files
+from pyrit.setup.akv_initialization import _load_environment_files
 
 _load_environment_files(env_files=None)
 

doc/getting_started/pyrit_conf.md

@@ -10,47 +10,50 @@ cp .pyrit_conf_example ~/.pyrit/.pyrit_conf
 cp .env_example ~/.pyrit/.env
 ```
 
-Then edit both files for your environment. The `.pyrit_conf` tells PyRIT _how_ to initialize; the `.env` tells it _where_ your AI targets are.
+Then edit both files for your environment. The `.pyrit_conf` tells PyRIT _how_ to initialize; the `.env` tells it _where_ your targets are.
 
 ## File Location
 
 The default configuration file path is:
 
-```
+```text
 ~/.pyrit/.pyrit_conf
 ```
 
 PyRIT looks for this file automatically on startup (via the CLI, shell, or `ConfigurationLoader`). If the file does not exist, PyRIT falls back to built-in defaults.
 
-## Setting Up Secrets (.env files)
+## Environment Configuration
 
-The `.pyrit_conf` file works hand-in-hand with `.env` files for your API credentials. See [Populating Secrets](./populating_secrets.md) for provider-specific examples of what to put in your `.env` file.
+```{important}
+Azure Key Vault is the recommended place for shared, CI/CD, and deployed PyRIT configuration. It avoids keeping credentials in a local `.env` file while preserving standard dotenv syntax. Existing `.env` configurations remain supported for backward compatibility and local development.
+```
 
-### Environment Variable Precedence
+See [Populating Secrets](./populating_secrets.md) for provider-specific variable examples.
 
-When PyRIT initializes, environment variables are loaded in a specific order. **Later sources override earlier ones:**
+### Loading Order
 
-```{mermaid}
-flowchart LR
-    A["1. System Environment"] --> B{"env_files in .pyrit_conf?"}
-    B -->|No| C["2. ~/.pyrit/.env"]
-    C --> D["3. ~/.pyrit/.env.local"]
-    B -->|Yes| E["2. Your specified files (in order)"]
-```
+PyRIT loads environment sources in this order. Each later source overrides matching values from earlier sources:
 
-**Default behavior** (no `env_files` field in `.pyrit_conf`):
+1. Existing process environment variables.
+2. Key Vault bootstrap documents from `env_akv_ref`, in list order.
+3. Local dotenv files:
+   - If `env_files` is configured, those files load in list order.
+   - Otherwise, `~/.pyrit/.env` loads if present, followed by `~/.pyrit/.env.local`.
 
-| Priority | Source | Description |
-|----------|--------|-------------|
-| Lowest | System environment variables | Always loaded as the baseline |
-| Medium | `~/.pyrit/.env` | Default config file (loaded if it exists) |
-| Highest | `~/.pyrit/.env.local` | Local overrides (loaded if it exists) |
+For a Key Vault-only setup, explicitly disable local dotenv loading:
 
-**Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. Default paths are completely ignored.
+```yaml
+env_akv_ref:
+  - https://my-vault.vault.azure.net/secrets/my-pyrit-env
+env_files: []
+```
+
+If Key Vault and local files are both configured, PyRIT warns that local values may override the fetched configuration. Remove stale local files when Key Vault should be authoritative.
 
 ### 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
@@ -107,7 +110,7 @@ Use `pyrit list initializers` in the CLI to see all registered initializers. See
 Most users should enable the following initializers. These are what the `.pyrit_conf_example` ships with and are required for features like `pyrit_scan` and automated scenarios.
 
 | Initializer | What It Registers | When You Need It |
-|---|---|---|
+| --- | --- | --- |
 | `target` | Prompt targets (OpenAI, Azure, AML, etc.) into the `TargetRegistry` | **Required for `pyrit_scan`** and any registry-based workflows |
 | `scorer` | Scorers (refusal, content safety, harm-category, Likert, etc.) into the `ScorerRegistry` | **Required for automated scoring** and `pyrit_scan` evaluations |
 | `technique` | Attack techniques into the `AttackTechniqueRegistry` | **Required for `pyrit_scan` scenarios** that select techniques |
@@ -157,20 +160,93 @@ initialization_scripts:
 
 ### `env_files`
 
-Environment file paths to load during initialization. Later files override values from earlier files.
+Optional local dotenv paths. Key Vault is recommended for shared or deployed configuration; use local files for backward compatibility and deliberate local overrides.
 
-| Value             | Behavior                                                             |
-| ----------------- | -------------------------------------------------------------------- |
-| Omitted or `null` | Load default `~/.pyrit/.env` and `~/.pyrit/.env.local` if they exist |
-| `[]` (empty list) | Load **no** environment files                                        |
-| List of paths     | Load **only** the specified files (defaults are skipped)             |
+| Value             | Behavior                                                 |
+| ----------------- | -------------------------------------------------------- |
+| Omitted or `null` | Load default `~/.pyrit/.env` and `~/.pyrit/.env.local`   |
+| `[]` (empty list) | Load **no** environment files                            |
+| List of paths     | Load **only** the specified files (defaults are skipped) |
 
 ```yaml
 env_files:
   - /path/to/.env
   - /path/to/.env.local
 ```
 
+Local environment files use standard python-dotenv parsing and `${NAME}` interpolation. Interpolation follows assignment and file load order. The default `.env.local` can reference a value loaded earlier from `.env`, for example `FOOBAR=${OPENAI_CHAT_ENDPOINT}`. A `.env` value cannot reference a variable introduced only by the later `.env.local`; values are not resolved retroactively. Explicit `env_files` follow the order in which they are listed.
+
+`env_akv_strict` does not apply to local files: malformed local lines retain python-dotenv's existing permissive skip-and-warn behavior. Local `kv:`, `akv:`, `azure_key_vault:`, and `env_akv_ref:` values remain literal; child-secret resolution is limited to Key Vault bootstrap documents. PyRIT does not define `env:` or `literal:` interpolation syntax. Use standard `${NAME}` interpolation instead.
+
+Environment loading preserves the historical non-transactional dotenv behavior. Each bootstrap document and local file updates `os.environ` as it loads. If a later source or child-secret lookup fails, assignments made by earlier sources remain in the process environment.
+
+When `env_akv_ref` is not configured, an empty `env_files` list or missing default files leaves existing process environment variables unchanged and initialization continues.
+
+### `env_akv_ref`
+
+Ordered Azure Key Vault secret URLs used to obtain bootstrap environment documents. This is the recommended configuration path. Each secret value must contain dotenv-formatted entries. Authentication uses `DefaultAzureCredential`.
+
+```yaml
+env_akv_ref:
+  - https://my-vault.vault.azure.net/secrets/shared-pyrit-env
+  - https://my-vault.vault.azure.net/secrets/team-pyrit-env
+```
+
+Bootstrap documents load in list order with `override=True`; local environment files load afterward. Each document uses native dotenv interpolation against the process environment and assignments already parsed. A bootstrap document can mix literal values, `${NAME}` interpolation, and complete-value references to scalar secrets in the same vault:
+
+```dotenv
+OPENAI_CHAT_ENDPOINT="https://example.openai.azure.com/openai/v1"
+OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key"
+PINNED_OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version-id"
+OPENAI_CHAT_MODEL="${PYRIT_OPENAI_CHAT_MODEL}"
+```
+
+Resolution is limited to one child-secret lookup:
+
+1. PyRIT validates and loads the bootstrap dotenv document.
+2. For each complete-value Key Vault reference in that document, PyRIT fetches the same-vault scalar secret and replaces the environment value.
+
+For example, if `OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key"`, the value of the `openai-chat-key` secret becomes `OPENAI_CHAT_KEY` verbatim. If that secret happens to contain `kv:another-secret`, the final environment value is the string `kv:another-secret`; PyRIT does not fetch `another-secret`.
+
+References must occupy the entire value. `kv:` is the canonical Key Vault prefix; `akv:`, `azure_key_vault:`, and `env_akv_ref:` are accepted aliases.
+
+A Key Vault reference must use a full HTTPS secret URL from the bootstrap document's vault. Supported vault DNS suffixes are `.vault.azure.net`, `.vault.azure.cn`, and `.vault.usgovcloudapi.net`. An unversioned URL reads the latest secret version at initialization. Include the version in the URL to pin it. Short names, malformed paths, arbitrary hosts, and cross-vault child references are rejected before a client is created.
+
+PyRIT does not cache referenced secrets. Each `kv:` occurrence in a bootstrap document performs a Key Vault read during initialization, including repeated references to the same URI. A later bootstrap or local file may override a reference after it has already been fetched.
+
+```dotenv
+LATEST_KEY_URI="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key"
+PINNED_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version-id"
+```
+
+Bootstrap documents stay in memory by default. They load before explicit `env_files` or the default `~/.pyrit/.env` and `~/.pyrit/.env.local`, allowing intentional local overrides.
+
+### `env_akv_strict`
+
+Controls validation only of the Key Vault bootstrap document and defaults to `true`. It does not change parsing of `.env`, `.env.local`, or explicit `env_files`.
+
+```yaml
+env_akv_strict: false
+```
+
+In strict mode, any malformed dotenv line or variable without an equals sign stops that bootstrap document before it mutates the environment. Empty assignments such as `OPTIONAL_VALUE=` remain valid and set the variable to an empty string. A referenced Key Vault secret whose value is an empty string is also valid. A missing value represented by `None` is treated as an error. With `env_akv_strict: false`, PyRIT emits a warning containing only malformed line numbers and valueless variable names, skips those entries, and loads the valid assignments. Secret values are never included in the warning.
+
+Non-strict mode does not suppress Key Vault or reference failures. Missing secrets, invalid `kv:` URLs, and bootstrap documents with no valid assignments still stop initialization. Because loading is non-transactional, values from earlier bootstrap documents remain if a later document fails, and raw values from the current document may remain if a child-secret lookup fails.
+
+Key Vault clients use an explicit Azure retry policy with up to three retries and exponential backoff. Bootstrap parsing, invalid or missing secrets, authentication, authorization, and Azure transport failures are raised as `KeyVaultInitializationException` with the original exception preserved as the cause. The exception remains `ValueError`-compatible for callers migrating from the previous contract.
+
+### `env_akv_write_env`
+
+Defaults to `false`. Set it to `true` to write the fetched bootstrap document to `~/.pyrit/.env` for inspecting configured targets and aliases:
+
+```yaml
+env_akv_write_env: true
+```
+
+The written file contains the bootstrap text before child `kv:` references are resolved. This makes target configuration readable without writing referenced child-secret values. However, any literal secret already present in the bootstrap document is written as-is, so treat the file as sensitive.
+
+Writing is opt-in and overwrites an existing `~/.pyrit/.env`. PyRIT does not load the generated `.env` during that same initialization, because its unresolved `kv:` references would otherwise replace resolved values. `.env.local` and other explicit local files still load afterward. The generated file is not a secure backup and should be removed when debugging is complete.
+
 ### `silent`
 
 If `true`, suppresses print statements during initialization. Useful for non-interactive environments or when embedding PyRIT in other tools. Defaults to `false`.
@@ -180,7 +256,7 @@ If `true`, suppresses print statements during initialization. Useful for non-int
 Client settings for connecting to or launching a PyRIT backend.
 
 | Field | Description | Default |
-|---|---|---|
+| --- | --- | --- |
 | `url` | Backend URL used when `--server-url` is omitted | `http://localhost:8000` |
 | `startup_timeout` | Seconds `pyrit_scan --start-server` waits for a healthy backend before terminating the spawned process | `120` |
 
@@ -216,7 +292,7 @@ This means you can set sensible defaults in `~/.pyrit/.pyrit_conf` and override
 
 The 3-layer model above determines **which config values are selected**. Once resolved, the values are applied in a fixed runtime order:
 
-1. Environment files are loaded
+1. Configured AKV bootstrap documents load in order, followed by selected environment files
 2. Default values are reset
 3. Memory database is configured (from `memory_db_type`)
 4. Initializers are executed in listed order
@@ -292,12 +368,14 @@ initializers:
 # initialization_scripts:
 #   - /path/to/my_custom_initializer.py
 
-# Environment files (optional)
-# Omit or set to null to use defaults (~/.pyrit/.env, ~/.pyrit/.env.local)
-# Set to [] to load no env files
-# env_files:
-#   - /path/to/.env
-#   - /path/to/.env.local
+# Recommended: ordered Azure Key Vault bootstrap environment documents
+# env_akv_ref:
+#   - https://my-vault.vault.azure.net/secrets/my-pyrit-env
+# env_akv_strict: true
+# env_akv_write_env: false  # Opt in to writing ~/.pyrit/.env for inspection
+
+# Recommended with Key Vault: disable local dotenv overrides
+# env_files: []
 
 # Suppress initialization messages
 silent: false

infra/env.demo.template

@@ -36,16 +36,16 @@ AZURE_CONTENT_SAFETY_API_ENDPOINT=https://YOUR_CONTENT_SAFETY.cognitiveservices.
 AZURE_CONTENT_SAFETY_API_KEY=
 
 # ─── Image Target (optional — for image generation demos) ───
-# OPENAI_IMAGE_ENDPOINT1=https://YOUR_IMAGE_ENDPOINT.openai.azure.com/openai/v1
-# OPENAI_IMAGE_API_KEY1=
-# OPENAI_IMAGE_MODEL1=dall-e-3
-# OPENAI_IMAGE_UNDERLYING_MODEL1=dall-e-3
+# AZURE_OPENAI_IMAGE_ENDPOINT1=https://YOUR_IMAGE_ENDPOINT.openai.azure.com/openai/v1
+# AZURE_OPENAI_IMAGE_API_KEY1=
+# AZURE_OPENAI_IMAGE_MODEL1=dall-e-3
+# AZURE_OPENAI_IMAGE_UNDERLYING_MODEL1=dall-e-3
 
 # ─── TTS Target (optional — for text-to-speech demos) ───
-# OPENAI_TTS_ENDPOINT1=https://YOUR_TTS_ENDPOINT.openai.azure.com/openai/v1
-# OPENAI_TTS_KEY1=
-# OPENAI_TTS_MODEL1=tts-1
-# OPENAI_TTS_UNDERLYING_MODEL1=tts-1
+# AZURE_OPENAI_TTS_ENDPOINT1=https://YOUR_TTS_ENDPOINT.openai.azure.com/openai/v1
+# AZURE_OPENAI_TTS_KEY1=
+# AZURE_OPENAI_TTS_MODEL1=tts-1
+# AZURE_OPENAI_TTS_UNDERLYING_MODEL1=tts-1
 
 # ─── Video Target (optional — for video generation demos) ───
 # AZURE_OPENAI_VIDEO_ENDPOINT=https://YOUR_VIDEO_ENDPOINT.openai.azure.com/openai/v1

pyrit/exceptions/__init__.py

@@ -9,6 +9,7 @@
     EmptyResponseException,
     ExperimentalWarning,
     InvalidJsonException,
+    KeyVaultInitializationException,
     MissingPromptPlaceholderException,
     PyritException,
     RateLimitException,
@@ -53,6 +54,7 @@
     "get_retry_max_num_attempts",
     "handle_bad_request_exception",
     "InvalidJsonException",
+    "KeyVaultInitializationException",
     "MissingPromptPlaceholderException",
     "PyritException",
     "pyrit_custom_result_retry",

pyrit/exceptions/exception_classes.py

@@ -190,6 +190,25 @@ def __init__(self, *, status_code: int = 500, message: str = "Server Error", bod
         self.body = body
 
 
+class KeyVaultInitializationException(PyritException, ValueError):  # noqa: N818
+    """Exception raised when Key Vault-backed environment initialization fails."""
+
+    def __init__(
+        self,
+        *,
+        status_code: int = 500,
+        message: str = "Key Vault environment initialization failed",
+    ) -> None:
+        """
+        Initialize a Key Vault initialization exception.
+
+        Args:
+            status_code (int): HTTP-style status code associated with the failure.
+            message (str): Human-readable failure description.
+        """
+        super().__init__(status_code=status_code, message=message)
+
+
 class EmptyResponseException(BadRequestException):
     """Exception class for empty response errors."""
 

pyrit/executor/promptgen/gcg/experiments/run.py

@@ -27,7 +27,7 @@
 from pyrit.executor.promptgen.gcg.config import GCGConfig, GCGDataConfig, GCGOutputConfig
 from pyrit.executor.promptgen.gcg.data import load_goals_and_targets
 from pyrit.executor.promptgen.gcg.generator import GCGGenerator
-from pyrit.setup.initialization import _load_environment_files
+from pyrit.setup.akv_initialization import _load_environment_files
 
 
 def _parse_arguments() -> argparse.Namespace: