From 6941013c37fad8bdd672ebc5156a9cc81b6b501e Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 11 Aug 2026 11:03:39 -0400 Subject: [PATCH 01/16] FEAT Add Key Vault environment resolution --- .pyrit_conf_example | 12 +- doc/getting_started/pyrit_conf.md | 55 ++++- pyrit/setup/initialization.py | 254 ++++++++++++++++++++++-- tests/unit/setup/test_initialization.py | 248 +++++++++++++++++------ 4 files changed, 473 insertions(+), 96 deletions(-) diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 523ccc28e1..e20faecf34 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -97,9 +97,17 @@ operation: op_trash_panda # 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. +# The first secret's value must be the full contents of a .env file. +# Values in that document may reference ambient variables with env:NAME or +# scalar secrets in the same vault with kv:SECRET_NAME. +# Additional entries are currently ignored pending support for labeled references. +# The AKV document replaces the default ~/.pyrit/.env source. Explicit env_files +# and the default ~/.pyrit/.env.local are loaded afterward and take precedence. +# PyRIT emits a warning when these local files coexist with env_akv_ref so stale +# configuration cannot silently mask or be mistaken for the Key Vault document. # Authentication uses DefaultAzureCredential (managed identity, Azure CLI, etc.). +# If env_akv_ref is omitted, at least one configured or default environment file +# must exist. System environment variables remain available but are not a source by themselves. # # Requires: pip install azure-keyvault-secrets # diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 153c93b9bf..d60f55967c 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -32,25 +32,33 @@ When PyRIT initializes, environment variables are loaded in a specific 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)"] + A["1. System Environment"] --> B{"env_akv_ref configured?"} + B -->|Yes| C["2. First AKV secret"] + B -->|No| D["2. ~/.pyrit/.env"] + C --> E["3. Explicit env_files or ~/.pyrit/.env.local"] + D --> F["3. ~/.pyrit/.env.local"] ``` -**Default behavior** (no `env_files` field in `.pyrit_conf`): +System environment variables are always the baseline, but initialization requires either an AKV root or at least one environment file. A system-environment-only configuration is not considered a complete source. + +**Default file behavior** (no `env_akv_ref` or `env_files` field in `.pyrit_conf`): | 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) | -**Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. Default paths are completely ignored. +**AKV behavior** (with `env_akv_ref`): The first referenced secret replaces `~/.pyrit/.env` as the root document. `~/.pyrit/.env.local` is loaded afterward if present. Additional AKV entries are currently ignored pending support for labeled references. + +PyRIT emits a warning when `env_akv_ref` is selected and default or explicit environment files coexist with it. The warning distinguishes files that are ignored from files that load afterward and override Key Vault values, making stale migration files visible at startup. + +**Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. They override the AKV root when both fields are configured, and default paths are completely ignored. ### 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 +115,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 | @@ -161,7 +169,7 @@ Environment file paths to load during initialization. Later files override value | Value | Behavior | | ----------------- | -------------------------------------------------------------------- | -| Omitted or `null` | Load default `~/.pyrit/.env` and `~/.pyrit/.env.local` if they exist | +| Omitted or `null` | Load default `~/.pyrit/.env` and `~/.pyrit/.env.local`, or only `.env.local` after an AKV root | | `[]` (empty list) | Load **no** environment files | | List of paths | Load **only** the specified files (defaults are skipped) | @@ -171,6 +179,29 @@ env_files: - /path/to/.env.local ``` +When `env_akv_ref` is not configured, an empty list or missing default files causes initialization to fail because no environment source is available. + +### `env_akv_ref` + +Azure Key Vault secret URLs used to obtain the root environment document. The first URL is used; its secret value must contain dotenv-formatted entries. Authentication uses `DefaultAzureCredential`. + +```yaml +env_akv_ref: + - https://my-vault.vault.azure.net/secrets/my-pyrit-env +``` + +The root document can mix literal values with references to ambient environment variables and scalar secrets in the same vault: + +```dotenv +OPENAI_CHAT_ENDPOINT="https://example.openai.azure.com/openai/v1" +OPENAI_CHAT_KEY="kv:openai-chat-key" +OPENAI_CHAT_MODEL="env:OPENAI_CHAT_MODEL" +``` + +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 referenced secret can point to another supported reference, subject to bounded depth and cycle detection. Prefix a value with `literal:` when its actual content starts with a reserved reference prefix. + +The AKV document is loaded before explicit `env_files` or the default `~/.pyrit/.env.local`, allowing local values to override shared configuration without writing the fetched document to disk. + ### `silent` If `true`, suppresses print statements during initialization. Useful for non-interactive environments or when embedding PyRIT in other tools. Defaults to `false`. @@ -197,7 +228,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. The AKV root or environment files are loaded, followed by local overrides 2. Default values are reset 3. Memory database is configured (from `memory_db_type`) 4. Initializers are executed in listed order @@ -280,6 +311,10 @@ initializers: # - /path/to/.env # - /path/to/.env.local +# Optional Azure Key Vault root environment document +# env_akv_ref: +# - https://my-vault.vault.azure.net/secrets/my-pyrit-env + # Suppress initialization messages silent: false ``` diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index eb0cf04ff8..b3d0dfa9a2 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -1,7 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio import io import logging +import os import pathlib from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Literal, get_args @@ -13,6 +15,8 @@ from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory if TYPE_CHECKING: + from azure.keyvault.secrets.aio import SecretClient + from pyrit.setup.pyrit_initializer import PyRITInitializer logger = logging.getLogger(__name__) @@ -22,8 +26,17 @@ AZURE_SQL = "AzureSQL" MemoryDatabaseType = Literal["InMemory", "SQLite", "AzureSQL"] +_AKV_REFERENCE_PREFIXES = frozenset({"akv", "kv", "azure_key_vault", "env_akv_ref"}) +_MAX_REFERENCE_DEPTH = 10 +_MAX_UNIQUE_SECRET_REFERENCES = 100 -def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: bool = False) -> None: + +def _load_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool = False, + include_default_base: bool = True, +) -> bool: """ Load environment files in the order they are provided. Later files override values from earlier files. @@ -33,6 +46,11 @@ def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: .env and .env.local from PyRIT home directory (only if they exist). silent: If True, suppresses print statements about environment file loading. Defaults to False. + include_default_base: If False and env_files is None, skips the default + .env file while still loading .env.local. Defaults to True. + + Returns: + True if at least one environment file was loaded, otherwise False. Raises: ValueError: If any provided env_files do not exist. @@ -51,7 +69,7 @@ def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" - if base_file.exists(): + if include_default_base and base_file.exists(): default_files.append(base_file) if local_file.exists(): default_files.append(local_file) @@ -63,7 +81,7 @@ def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: ) else: _print_msg( - "No default environment files found. Using system environment variables only.", + "No default environment files found.", quiet=silent, log=True, ) @@ -75,6 +93,8 @@ def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: if not silent: _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) + return bool(env_files) + def _print_msg(message: str, quiet: bool, log: bool) -> None: """ @@ -91,6 +111,41 @@ def _print_msg(message: str, quiet: bool, log: bool) -> None: logger.info(message) +def _warn_about_akv_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool = False, +) -> None: + """Warn when local environment files coexist with an AKV environment source.""" + base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" + local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" + messages: list[str] = [] + + if base_file.exists(): + messages.append(f"{base_file} exists and will be ignored because Key Vault supplies the base environment") + + if local_file.exists(): + if env_files is None: + messages.append(f"{local_file} will load after Key Vault and override matching values") + else: + messages.append(f"{local_file} exists but will be ignored because env_files was explicitly configured") + + if env_files: + messages.append(f"explicit env_files will load after Key Vault and override matching values: {list(env_files)}") + + if not messages: + return + + message = ( + "env_akv_ref is configured, but local environment files were also found:\n- " + + "\n- ".join(messages) + + "\nConfirm that this precedence is intentional." + ) + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + + def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: """ Parse an AKV secret URL into vault URL, secret name, and optional version. @@ -120,38 +175,172 @@ def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: async def _load_env_from_akv_async(*, secret_urls: Sequence[str], silent: bool = False) -> None: """ - Load environment variables from Azure Key Vault secrets. + Load environment variables from an Azure Key Vault secret. - Each secret's value is treated as the full contents of a ``.env`` file and - parsed accordingly. Later secrets override values from earlier ones. + The first secret URL identifies the root environment document. Additional + URLs are ignored until the configuration supports labeled references. Authentication uses ``DefaultAzureCredential``, which silently tries managed identity, Azure CLI, VS Code credentials, etc., and falls back to interactive browser authentication when running locally. Args: - secret_urls (Sequence[str]): Sequence of AKV secret URLs to load, each in - the format ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. + secret_urls (Sequence[str]): Sequence of AKV secret URLs. The first URL + must use the format ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. silent (bool): If True, suppresses print statements. Defaults to False. Raises: ImportError: If ``azure-keyvault-secrets`` is not installed. - ValueError: If a secret URL is malformed. + ValueError: If no root secret is configured, the root URL is malformed, + or the environment document cannot be fully resolved. """ if not secret_urls: - return + raise ValueError("At least one env_akv_ref URL is required to load an environment document.") + from azure.identity.aio import DefaultAzureCredential from azure.keyvault.secrets.aio import SecretClient - credential = DefaultAzureCredential() - for secret_url in secret_urls: - _print_msg(f"Loading environment from AKV secret: {secret_url}", quiet=silent, log=True) - vault_url, secret_name, secret_version = _parse_akv_secret_url(secret_url) - client = SecretClient(vault_url=vault_url, credential=credential) - secret = await client.get_secret(secret_name, version=secret_version) - if secret.value: - dotenv.load_dotenv(stream=io.StringIO(secret.value), override=True) - _print_msg(f"Loaded environment from AKV secret: {secret_url}", quiet=silent, log=True) + secret_url = secret_urls[0] + if len(secret_urls) > 1: + _print_msg( + "Multiple env_akv_ref values were provided; using the first as the root environment document.", + quiet=silent, + log=True, + ) + + _print_msg(f"Loading environment from AKV secret: {secret_url}", quiet=silent, log=True) + vault_url, secret_name, secret_version = _parse_akv_secret_url(secret_url) + ambient_environment = dict(os.environ) + async with DefaultAzureCredential() as credential: + async with SecretClient(vault_url=vault_url, credential=credential) as client: + secret = await client.get_secret(secret_name, version=secret_version) + + if not secret.value: + raise ValueError(f"AKV environment secret has no value: {secret_url}") + + parsed_environment = dotenv.dotenv_values(stream=io.StringIO(secret.value), interpolate=True) + if not parsed_environment: + raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") + + missing_values = [name for name, value in parsed_environment.items() if value is None] + if missing_values: + raise ValueError( + "AKV environment document contains variables without values: " + ", ".join(missing_values) + ) + + resolved_secrets: dict[str, str] = {} + resolved_environment: dict[str, str] = {} + for variable_name, value in parsed_environment.items(): + if value is None: + continue + resolved_environment[variable_name] = await _resolve_environment_value_async( + value=value, + variable_name=variable_name, + secret_client=client, + ambient_environment=ambient_environment, + resolved_secrets=resolved_secrets, + ) + + os.environ.update(resolved_environment) + + _print_msg(f"Loaded environment from AKV secret: {secret_url}", quiet=silent, log=True) + + +def _parse_environment_value_reference(value: str) -> tuple[str, str] | None: + """ + Parse an exact whole-value environment or Key Vault reference. + + Returns: + The normalized reference type and target, or None for a literal value. + """ + prefix, separator, target = value.partition(":") + if not separator: + return None + if prefix == "env": + return "env", target.strip() + if prefix in _AKV_REFERENCE_PREFIXES: + return "akv", target.strip() + if prefix == "literal": + return "literal", target + return None + + +def _validate_akv_secret_name(*, secret_name: str, variable_name: str) -> None: + if not secret_name or len(secret_name) > 127 or any(not char.isalnum() and char != "-" for char in secret_name): + raise ValueError( + f"Invalid same-vault secret name '{secret_name}' referenced by environment variable '{variable_name}'. " + "Secret names must contain only letters, numbers, and hyphens." + ) + + +async def _resolve_environment_value_async( + *, + value: str, + variable_name: str, + secret_client: "SecretClient", + ambient_environment: dict[str, str], + resolved_secrets: dict[str, str], + reference_path: tuple[str, ...] = (), +) -> str: + reference = _parse_environment_value_reference(value) + if reference is None: + return value + + reference_type, target = reference + if reference_type == "literal": + return target + if not target: + raise ValueError(f"Empty {reference_type} reference for environment variable '{variable_name}'.") + if len(reference_path) >= _MAX_REFERENCE_DEPTH: + raise ValueError( + f"Environment reference depth exceeded {_MAX_REFERENCE_DEPTH} while resolving '{variable_name}'." + ) + + normalized_target = target if reference_type == "env" else target.casefold() + reference_token = f"{reference_type}:{normalized_target}" + if reference_token in reference_path: + cycle = " -> ".join((*reference_path, reference_token)) + raise ValueError(f"Environment reference cycle detected while resolving '{variable_name}': {cycle}") + next_path = (*reference_path, reference_token) + + if reference_type == "env": + if target not in ambient_environment: + raise ValueError( + f"Environment variable '{target}' referenced by '{variable_name}' " + "is not set in the ambient environment." + ) + return await _resolve_environment_value_async( + value=ambient_environment[target], + variable_name=variable_name, + secret_client=secret_client, + ambient_environment=ambient_environment, + resolved_secrets=resolved_secrets, + reference_path=next_path, + ) + + _validate_akv_secret_name(secret_name=target, variable_name=variable_name) + secret_cache_key = target.casefold() + if secret_cache_key in resolved_secrets: + return resolved_secrets[secret_cache_key] + if len(resolved_secrets) >= _MAX_UNIQUE_SECRET_REFERENCES: + raise ValueError( + f"Environment secret reference limit of {_MAX_UNIQUE_SECRET_REFERENCES} exceeded while resolving " + f"'{variable_name}'." + ) + + secret = await secret_client.get_secret(target) + if secret.value is None: + raise ValueError(f"AKV secret '{target}' referenced by environment variable '{variable_name}' has no value.") + resolved_value = await _resolve_environment_value_async( + value=secret.value, + variable_name=variable_name, + secret_client=secret_client, + ambient_environment=ambient_environment, + resolved_secrets=resolved_secrets, + reference_path=next_path, + ) + resolved_secrets[secret_cache_key] = resolved_value + return resolved_value async def _execute_initializers_async(*, initializers: Sequence["PyRITInitializer"]) -> None: @@ -237,12 +426,35 @@ async def initialize_pyrit_async( **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. Raises: - ValueError: If an unsupported memory_db_type is provided or if env_files contains non-existent files. + ValueError: If an unsupported memory_db_type is provided, env_files contains non-existent files, + or neither env_akv_ref nor an environment file is available. """ if env_akv_ref: + await asyncio.to_thread( + _warn_about_akv_environment_files, + env_files=env_files, + silent=silent, + ) await _load_env_from_akv_async(secret_urls=env_akv_ref, silent=silent) - _load_environment_files(env_files=env_files, silent=silent) + # PR review decision: .env.local and explicit files currently override the Key Vault document. + # The default .env is always skipped because Key Vault supplies the base environment. + await asyncio.to_thread( + _load_environment_files, + env_files=env_files, + silent=silent, + include_default_base=False, + ) + else: + loaded_local_file = await asyncio.to_thread( + _load_environment_files, + env_files=env_files, + silent=silent, + ) + if not loaded_local_file: + raise ValueError( + "No environment source found. Configure env_akv_ref or provide at least one .env or .env.local file." + ) # Reset all default values before executing initialization scripts # This ensures a clean state for each initialization diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index b919df4338..d65e0f4c6e 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -3,7 +3,6 @@ import os import pathlib -import sys import tempfile import types from unittest import mock @@ -14,7 +13,12 @@ from pyrit.common.singleton import Singleton from pyrit.registry import InitializerRegistry from pyrit.setup import IN_MEMORY, initialize_pyrit_async -from pyrit.setup.initialization import _load_env_from_akv_async, _load_environment_files, _parse_akv_secret_url +from pyrit.setup.initialization import ( + _load_env_from_akv_async, + _load_environment_files, + _parse_akv_secret_url, + _warn_about_akv_environment_files, +) class TestLoadInitializersFromScripts: @@ -125,7 +129,9 @@ def setup_method(self) -> None: @mock.patch("pyrit.setup.initialization._load_environment_files") async def test_initialize_basic(self, mock_load_env, mock_set_memory): """Test basic initialization.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY) + mock_load_env.return_value = True + + await initialize_pyrit_async(memory_db_type=IN_MEMORY, load_defaults=False) mock_load_env.assert_called_once() mock_set_memory.assert_called_once() @@ -161,10 +167,11 @@ async def initialize_async(self) -> None: finally: os.unlink(script_path) - async def test_invalid_memory_type_raises_error(self): + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=True) + async def test_invalid_memory_type_raises_error(self, mock_load_env): """Test that invalid memory type raises ValueError.""" with pytest.raises(ValueError, match="is not a supported type"): - await initialize_pyrit_async(memory_db_type="InvalidType") # type: ignore[arg-type] + await initialize_pyrit_async(memory_db_type="InvalidType", load_defaults=False) # type: ignore[arg-type] @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") @mock.patch("pyrit.setup.initialization._load_environment_files") @@ -173,7 +180,7 @@ async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, m """Test that env_akv_ref triggers AKV env loading.""" refs = ["https://vault.vault.azure.net/secrets/test-secret"] - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) mock_load_akv.assert_awaited_once() assert mock_load_akv.await_args.kwargs["secret_urls"] == refs @@ -188,7 +195,9 @@ async def test_initialize_with_empty_env_akv_ref_does_not_load_akv( self, mock_load_akv, mock_load_env, mock_set_memory ): """Test that empty env_akv_ref does not invoke AKV loading.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=[]) + mock_load_env.return_value = True + + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=[], load_defaults=False) mock_load_akv.assert_not_called() mock_load_env.assert_called_once() @@ -199,23 +208,36 @@ async def test_initialize_loads_akv_before_env_files(self, mock_set_memory): """Test that AKV refs are loaded before env_files so env_files can override values.""" call_order: list[str] = [] + def _record_warning(*, env_files, silent=False): + call_order.append("warning") + async def _record_akv_call(*, secret_urls, silent=False): call_order.append("akv") - def _record_env_file_call(*, env_files, silent=False): + def _record_env_file_call(*, env_files, silent=False, include_default_base=True): call_order.append("env_files") + assert include_default_base is False + return True refs = ["https://vault.vault.azure.net/secrets/test-secret"] with ( + mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files", side_effect=_record_warning), mock.patch("pyrit.setup.initialization._load_env_from_akv_async", side_effect=_record_akv_call), mock.patch("pyrit.setup.initialization._load_environment_files", side_effect=_record_env_file_call), ): - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) - assert call_order == ["akv", "env_files"] + assert call_order == ["warning", "akv", "env_files"] mock_set_memory.assert_called_once() + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_without_environment_source_raises(self, mock_set_memory): + with pytest.raises(ValueError, match="No environment source found"): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[], load_defaults=False) + + mock_set_memory.assert_not_called() + @pytest.fixture def reset_memory_singletons(): @@ -237,16 +259,18 @@ def setup_method(self) -> None: """Clear default values before each test.""" reset_default_values() - async def test_initialize_silent_produces_no_output(self, capsys): + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=True) + async def test_initialize_silent_produces_no_output(self, mock_load_env, capsys): """initialize_pyrit_async with silent=True must not print anything to stdout.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=True) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=True, load_defaults=False) captured = capsys.readouterr() assert captured.out == "" - async def test_initialize_not_silent_prints_migration_message(self, capsys): + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=True) + async def test_initialize_not_silent_prints_migration_message(self, mock_load_env, capsys): """Without silent, the Alembic schema-check message is printed and tagged as Alembic output.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=False) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=False, load_defaults=False) captured = capsys.readouterr() assert "[pyrit:alembic] No new upgrade operations detected." in captured.out @@ -273,9 +297,10 @@ async def test_loads_default_env_files_when_none_provided(self, mock_config_path mock_config_path.__truediv__ = lambda self, other: temp_path / other # Call the function with None (default behavior) - _load_environment_files(env_files=None) + loaded = _load_environment_files(env_files=None) # Verify both files were loaded + assert loaded is True assert mock_load_dotenv.call_count == 2 calls = [call[0][0] for call in mock_load_dotenv.call_args_list] assert env_file in calls @@ -294,12 +319,76 @@ async def test_only_loads_existing_default_files(self, mock_config_path, mock_lo mock_config_path.__truediv__ = lambda self, other: temp_path / other - _load_environment_files(env_files=None) + loaded = _load_environment_files(env_files=None) # Verify only one file was loaded + assert loaded is True assert mock_load_dotenv.call_count == 1 assert mock_load_dotenv.call_args[0][0] == env_file + @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_excludes_default_env_when_loading_local_override(self, mock_config_path, mock_load_dotenv): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text("VAR=base") + env_local_file.write_text("VAR=local") + + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + loaded = _load_environment_files(env_files=None, include_default_base=False) + + assert loaded is True + mock_load_dotenv.assert_called_once() + assert mock_load_dotenv.call_args.args[0] == env_local_file + + @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_returns_false_when_no_default_files_exist(self, mock_config_path, mock_load_dotenv): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + loaded = _load_environment_files(env_files=None) + + assert loaded is False + mock_load_dotenv.assert_not_called() + + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_warns_when_default_files_coexist_with_akv(self, mock_config_path, caplog, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text("VAR=base") + env_local_file.write_text("VAR=local") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with caplog.at_level("WARNING", logger="pyrit.setup.initialization"): + _warn_about_akv_environment_files(env_files=None) + + output = capsys.readouterr().out + assert output.startswith("WARNING: env_akv_ref is configured") + assert f"{env_file} exists and will be ignored" in output + assert f"{env_local_file} will load after Key Vault and override matching values" in output + assert "Confirm that this precedence is intentional." in output + assert caplog.records[0].levelname == "WARNING" + + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_akv_environment_file_warning_respects_silent(self, mock_config_path, caplog, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VAR=base") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with caplog.at_level("WARNING", logger="pyrit.setup.initialization"): + _warn_about_akv_environment_files(env_files=None, silent=True) + + assert capsys.readouterr().out == "" + assert "will be ignored because Key Vault supplies the base environment" in caplog.text + @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") async def test_loads_custom_env_files_in_order(self, mock_load_dotenv): """Test that custom env_files are loaded in the order provided.""" @@ -338,7 +427,7 @@ async def test_initialize_pyrit_with_custom_env_files(self, mock_set_memory): env_file.write_text("CUSTOM_VAR=custom_value") # Should not raise an error - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[env_file]) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[env_file], load_defaults=False) mock_set_memory.assert_called_once() @@ -371,7 +460,7 @@ async def test_custom_env_files_override_default_behavior(self, mock_set_memory, mock_home_path.__truediv__ = lambda self, other: temp_path / other # Pass custom env_files - should NOT load defaults - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[custom_env]) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[custom_env], load_defaults=False) # Verify only custom file was loaded, not the default ones assert mock_load_dotenv.call_count == 1 @@ -403,61 +492,94 @@ def test_parse_akv_secret_url_invalid_raises(self): with pytest.raises(ValueError, match="Invalid AKV secret URL"): _parse_akv_secret_url("https://myvault.vault.azure.net/not-secrets/my-secret") - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") - async def test_load_env_from_akv_async_empty_urls_noop(self, mock_load_dotenv): - await _load_env_from_akv_async(secret_urls=[]) - mock_load_dotenv.assert_not_called() - - async def test_load_env_from_akv_async_loads_secret_content(self): - class FakeCredential: - pass - - client_calls: list[tuple[str, object, object]] = [] + async def test_load_env_from_akv_async_empty_urls_raises(self): + with pytest.raises(ValueError, match="At least one env_akv_ref URL is required"): + await _load_env_from_akv_async(secret_urls=[]) + + @pytest.mark.parametrize( + "secret_urls", + [ + ["https://myvault.vault.azure.net/secrets/my-secret/v1"], + [ + "https://myvault.vault.azure.net/secrets/my-secret/v1", + "https://myvault.vault.azure.net/secrets/ignored/v2", + ], + ], + ) + async def test_load_env_from_akv_async_loads_first_secret_content(self, secret_urls): + credential = mock.MagicMock() + credential.__aenter__ = mock.AsyncMock(return_value=credential) + credential.__aexit__ = mock.AsyncMock(return_value=None) + client = mock.MagicMock() + client.__aenter__ = mock.AsyncMock(return_value=client) + client.__aexit__ = mock.AsyncMock(return_value=None) + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="AKV_VAR=from_secret\n")) - class FakeSecretClient: - def __init__(self, *, vault_url, credential): - client_calls.append(("init", vault_url, credential)) + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential) as mock_credential_cls, + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, + mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") as mock_load_dotenv, + mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, + ): + mock_load_dotenv.return_value = True + await _load_env_from_akv_async(secret_urls=secret_urls, silent=True) - async def get_secret(self, name, version=None): - client_calls.append(("get_secret", name, version)) - return types.SimpleNamespace(value="AKV_VAR=from_secret\n") + mock_credential_cls.assert_called_once_with() + mock_client_cls.assert_called_once_with(vault_url="https://myvault.vault.azure.net", credential=credential) + client.get_secret.assert_awaited_once_with("my-secret", version="v1") + credential.__aenter__.assert_awaited_once() + credential.__aexit__.assert_awaited_once() + client.__aenter__.assert_awaited_once() + client.__aexit__.assert_awaited_once() - azure_module = types.ModuleType("azure") - identity_module = types.ModuleType("azure.identity") - identity_aio_module = types.ModuleType("azure.identity.aio") - keyvault_module = types.ModuleType("azure.keyvault") - keyvault_secrets_module = types.ModuleType("azure.keyvault.secrets") - keyvault_secrets_aio_module = types.ModuleType("azure.keyvault.secrets.aio") + stream = mock_load_dotenv.call_args.kwargs["stream"] + assert stream.getvalue() == "AKV_VAR=from_secret\n" + assert mock_load_dotenv.call_args.kwargs["override"] is True + assert mock_print_msg.call_count == 2 + (len(secret_urls) > 1) - identity_aio_module.DefaultAzureCredential = FakeCredential - keyvault_secrets_aio_module.SecretClient = FakeSecretClient + async def test_load_env_from_akv_async_empty_secret_raises(self): + credential = mock.MagicMock() + credential.__aenter__ = mock.AsyncMock(return_value=credential) + credential.__aexit__ = mock.AsyncMock(return_value=None) + client = mock.MagicMock() + client.__aenter__ = mock.AsyncMock(return_value=client) + client.__aexit__ = mock.AsyncMock(return_value=None) + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=None)) with ( - mock.patch.dict( - sys.modules, - { - "azure": azure_module, - "azure.identity": identity_module, - "azure.identity.aio": identity_aio_module, - "azure.keyvault": keyvault_module, - "azure.keyvault.secrets": keyvault_secrets_module, - "azure.keyvault.secrets.aio": keyvault_secrets_aio_module, - }, - ), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") as mock_load_dotenv, - mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, + pytest.raises(ValueError, match="has no value"), ): await _load_env_from_akv_async( - secret_urls=["https://myvault.vault.azure.net/secrets/my-secret/v1"], + secret_urls=["https://myvault.vault.azure.net/secrets/my-secret"], silent=True, ) - assert client_calls[0][0] == "init" - assert client_calls[0][1] == "https://myvault.vault.azure.net" - assert isinstance(client_calls[0][2], FakeCredential) - assert client_calls[1] == ("get_secret", "my-secret", "v1") + mock_load_dotenv.assert_not_called() + credential.__aexit__.assert_awaited_once() + client.__aexit__.assert_awaited_once() + + async def test_load_env_from_akv_async_without_entries_raises(self): + credential = mock.MagicMock() + credential.__aenter__ = mock.AsyncMock(return_value=credential) + credential.__aexit__ = mock.AsyncMock(return_value=None) + client = mock.MagicMock() + client.__aenter__ = mock.AsyncMock(return_value=client) + client.__aexit__ = mock.AsyncMock(return_value=None) + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="# comments only\n")) - stream = mock_load_dotenv.call_args.kwargs["stream"] - assert stream.getvalue() == "AKV_VAR=from_secret\n" - assert mock_load_dotenv.call_args.kwargs["override"] is True - assert mock_print_msg.call_count == 2 + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + mock.patch("pyrit.setup.initialization.dotenv.load_dotenv", return_value=False), + pytest.raises(ValueError, match="contains no environment entries"), + ): + await _load_env_from_akv_async( + secret_urls=["https://myvault.vault.azure.net/secrets/my-secret"], + silent=True, + ) + + credential.__aexit__.assert_awaited_once() + client.__aexit__.assert_awaited_once() From eafe42c5552e1de4677f7996b1c961dc90bab0bb Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 12 Aug 2026 10:20:37 -0400 Subject: [PATCH 02/16] FEAT: Removed recursion for KV lookups --- .pyrit_conf_example | 3 +- doc/getting_started/pyrit_conf.md | 22 ++++- pyrit/setup/initialization.py | 68 ++++++--------- tests/unit/setup/test_initialization.py | 106 +++++++++++++++--------- 4 files changed, 112 insertions(+), 87 deletions(-) diff --git a/.pyrit_conf_example b/.pyrit_conf_example index e20faecf34..c1227bf57a 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -100,7 +100,8 @@ operation: op_trash_panda # The first secret's value must be the full contents of a .env file. # Values in that document may reference ambient variables with env:NAME or # scalar secrets in the same vault with kv:SECRET_NAME. -# Additional entries are currently ignored pending support for labeled references. +# Referenced values are terminal scalars; they are not parsed for more references. +# If multiple URLs are listed, only the first is used. # The AKV document replaces the default ~/.pyrit/.env source. Explicit env_files # and the default ~/.pyrit/.env.local are loaded afterward and take precedence. # PyRIT emits a warning when these local files coexist with env_akv_ref so stale diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index d60f55967c..226cc3f80b 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -49,7 +49,7 @@ System environment variables are always the baseline, but initialization require | Medium | `~/.pyrit/.env` | Default config file (loaded if it exists) | | Highest | `~/.pyrit/.env.local` | Local overrides (loaded if it exists) | -**AKV behavior** (with `env_akv_ref`): The first referenced secret replaces `~/.pyrit/.env` as the root document. `~/.pyrit/.env.local` is loaded afterward if present. Additional AKV entries are currently ignored pending support for labeled references. +**AKV behavior** (with `env_akv_ref`): The first referenced secret replaces `~/.pyrit/.env` as the root document. `~/.pyrit/.env.local` is loaded afterward if present. If multiple AKV URLs are configured, only the first is used. PyRIT emits a warning when `env_akv_ref` is selected and default or explicit environment files coexist with it. The warning distinguishes files that are ignored from files that load afterward and override Key Vault values, making stale migration files visible at startup. @@ -195,10 +195,26 @@ The root document can mix literal values with references to ambient environment ```dotenv OPENAI_CHAT_ENDPOINT="https://example.openai.azure.com/openai/v1" OPENAI_CHAT_KEY="kv:openai-chat-key" -OPENAI_CHAT_MODEL="env:OPENAI_CHAT_MODEL" +OPENAI_CHAT_MODEL="env:PYRIT_OPENAI_CHAT_MODEL" ``` -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 referenced secret can point to another supported reference, subject to bounded depth and cycle detection. Prefix a value with `literal:` when its actual content starts with a reserved reference prefix. +Resolution is deliberately limited to two levels: + +1. PyRIT fetches the first `env_akv_ref` secret and parses it as the bootstrap dotenv document. +2. For each reference in that document, PyRIT either copies one ambient `env:` value or fetches one scalar secret from the same vault. The resulting value is final and is not parsed as another reference. + +For example, if `OPENAI_CHAT_KEY="kv: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. + +`literal:` is an escape hatch for a bootstrap value that begins with a reserved reference prefix. PyRIT removes `literal:` and returns the remainder without interpreting it as a reference. Quoting does not provide this escape because dotenv removes quotes while parsing. Values fetched from child secrets are already terminal and do not need this escape. + +```dotenv +REFERENCE="kv:openai-chat-key" +LITERAL_VALUE="literal:kv:not-a-secret-name" +``` + +Here, `REFERENCE` retrieves `openai-chat-key`, while `LITERAL_VALUE` becomes the string `kv:not-a-secret-name`. The AKV document is loaded before explicit `env_files` or the default `~/.pyrit/.env.local`, allowing local values to override shared configuration without writing the fetched document to disk. diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index b3d0dfa9a2..41fb313292 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -27,8 +27,6 @@ MemoryDatabaseType = Literal["InMemory", "SQLite", "AzureSQL"] _AKV_REFERENCE_PREFIXES = frozenset({"akv", "kv", "azure_key_vault", "env_akv_ref"}) -_MAX_REFERENCE_DEPTH = 10 -_MAX_UNIQUE_SECRET_REFERENCES = 100 def _load_environment_files( @@ -177,8 +175,9 @@ async def _load_env_from_akv_async(*, secret_urls: Sequence[str], silent: bool = """ Load environment variables from an Azure Key Vault secret. - The first secret URL identifies the root environment document. Additional - URLs are ignored until the configuration supports labeled references. + The first secret URL identifies the bootstrap environment document. Values + in that document may directly reference scalar secrets in the same vault. + Additional root URLs are ignored. Authentication uses ``DefaultAzureCredential``, which silently tries managed identity, Azure CLI, VS Code credentials, etc., and falls back to interactive @@ -192,7 +191,7 @@ async def _load_env_from_akv_async(*, secret_urls: Sequence[str], silent: bool = Raises: ImportError: If ``azure-keyvault-secrets`` is not installed. ValueError: If no root secret is configured, the root URL is malformed, - or the environment document cannot be fully resolved. + or the bootstrap environment document cannot be fully resolved. """ if not secret_urls: raise ValueError("At least one env_akv_ref URL is required to load an environment document.") @@ -280,8 +279,23 @@ async def _resolve_environment_value_async( secret_client: "SecretClient", ambient_environment: dict[str, str], resolved_secrets: dict[str, str], - reference_path: tuple[str, ...] = (), ) -> str: + """ + Resolve one value from the bootstrap environment document. + + Args: + value (str): The parsed bootstrap value. + variable_name (str): The environment variable receiving the resolved value. + secret_client (SecretClient): The client for the bootstrap document's vault. + ambient_environment (dict[str, str]): Snapshot used for ``env:`` references. + resolved_secrets (dict[str, str]): Same-vault scalar cache keyed by secret name. + + Returns: + str: The literal, ambient, or same-vault scalar value. + + Raises: + ValueError: If a reference is empty or cannot resolve to a value. + """ reference = _parse_environment_value_reference(value) if reference is None: return value @@ -291,17 +305,6 @@ async def _resolve_environment_value_async( return target if not target: raise ValueError(f"Empty {reference_type} reference for environment variable '{variable_name}'.") - if len(reference_path) >= _MAX_REFERENCE_DEPTH: - raise ValueError( - f"Environment reference depth exceeded {_MAX_REFERENCE_DEPTH} while resolving '{variable_name}'." - ) - - normalized_target = target if reference_type == "env" else target.casefold() - reference_token = f"{reference_type}:{normalized_target}" - if reference_token in reference_path: - cycle = " -> ".join((*reference_path, reference_token)) - raise ValueError(f"Environment reference cycle detected while resolving '{variable_name}': {cycle}") - next_path = (*reference_path, reference_token) if reference_type == "env": if target not in ambient_environment: @@ -309,38 +312,18 @@ async def _resolve_environment_value_async( f"Environment variable '{target}' referenced by '{variable_name}' " "is not set in the ambient environment." ) - return await _resolve_environment_value_async( - value=ambient_environment[target], - variable_name=variable_name, - secret_client=secret_client, - ambient_environment=ambient_environment, - resolved_secrets=resolved_secrets, - reference_path=next_path, - ) + return ambient_environment[target] _validate_akv_secret_name(secret_name=target, variable_name=variable_name) secret_cache_key = target.casefold() if secret_cache_key in resolved_secrets: return resolved_secrets[secret_cache_key] - if len(resolved_secrets) >= _MAX_UNIQUE_SECRET_REFERENCES: - raise ValueError( - f"Environment secret reference limit of {_MAX_UNIQUE_SECRET_REFERENCES} exceeded while resolving " - f"'{variable_name}'." - ) secret = await secret_client.get_secret(target) if secret.value is None: raise ValueError(f"AKV secret '{target}' referenced by environment variable '{variable_name}' has no value.") - resolved_value = await _resolve_environment_value_async( - value=secret.value, - variable_name=variable_name, - secret_client=secret_client, - ambient_environment=ambient_environment, - resolved_secrets=resolved_secrets, - reference_path=next_path, - ) - resolved_secrets[secret_cache_key] = resolved_value - return resolved_value + resolved_secrets[secret_cache_key] = secret.value + return secret.value async def _execute_initializers_async(*, initializers: Sequence["PyRITInitializer"]) -> None: @@ -419,8 +402,9 @@ async def initialize_pyrit_async( in order. If not provided, will load default .env and .env.local files from PyRIT home if they exist. All paths must be valid pathlib.Path objects. env_akv_ref (Sequence[str] | None): Optional sequence of Azure Key Vault secret URLs to load. - Each secret's value must be the full contents of a .env file. Loaded before ``env_files`` - so local files take precedence over AKV. Requires ``azure-keyvault-secrets``. + The first secret's value must contain the bootstrap .env document; additional URLs are ignored. + Loaded before ``env_files`` so local files take precedence over AKV. Requires + ``azure-keyvault-secrets``. silent (bool): If True, suppresses print statements about environment file loading and schema migration. Defaults to False. **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index d65e0f4c6e..52ce0a92ef 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -467,6 +467,16 @@ async def test_custom_env_files_override_default_behavior(self, mock_set_memory, assert mock_load_dotenv.call_args[0][0] == custom_env +def _create_mock_akv_clients() -> tuple[mock.MagicMock, mock.MagicMock]: + credential = mock.MagicMock() + credential.__aenter__ = mock.AsyncMock(return_value=credential) + credential.__aexit__ = mock.AsyncMock(return_value=None) + client = mock.MagicMock() + client.__aenter__ = mock.AsyncMock(return_value=client) + client.__aexit__ = mock.AsyncMock(return_value=None) + return credential, client + + class TestAkvEnvironmentLoading: """Tests for AKV URL parsing and env loading helpers.""" @@ -496,60 +506,59 @@ async def test_load_env_from_akv_async_empty_urls_raises(self): with pytest.raises(ValueError, match="At least one env_akv_ref URL is required"): await _load_env_from_akv_async(secret_urls=[]) - @pytest.mark.parametrize( - "secret_urls", - [ - ["https://myvault.vault.azure.net/secrets/my-secret/v1"], - [ - "https://myvault.vault.azure.net/secrets/my-secret/v1", - "https://myvault.vault.azure.net/secrets/ignored/v2", - ], - ], - ) - async def test_load_env_from_akv_async_loads_first_secret_content(self, secret_urls): - credential = mock.MagicMock() - credential.__aenter__ = mock.AsyncMock(return_value=credential) - credential.__aexit__ = mock.AsyncMock(return_value=None) - client = mock.MagicMock() - client.__aenter__ = mock.AsyncMock(return_value=client) - client.__aexit__ = mock.AsyncMock(return_value=None) - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="AKV_VAR=from_secret\n")) + async def test_load_env_from_akv_async_uses_first_root_and_resolves_one_level(self): + credential, client = _create_mock_akv_clients() + root_document = ( + "DIRECT=from-bootstrap\n" + "FROM_ENV=env:SOURCE_VALUE\n" + "FROM_KV=kv:api-key\n" + "DUPLICATE_KV=akv:API-KEY\n" + "ESCAPED=literal:kv:not-a-secret" + ) + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value=root_document), + types.SimpleNamespace(value="env:not-resolved-again"), + ] + ) + secret_urls = [ + "https://myvault.vault.azure.net/secrets/bootstrap/v1", + "https://myvault.vault.azure.net/secrets/ignored/v2", + ] with ( + mock.patch.dict(os.environ, {"SOURCE_VALUE": "kv:not-fetched"}, clear=False), mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential) as mock_credential_cls, mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, - mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") as mock_load_dotenv, mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, ): - mock_load_dotenv.return_value = True await _load_env_from_akv_async(secret_urls=secret_urls, silent=True) + assert os.environ["DIRECT"] == "from-bootstrap" + assert os.environ["FROM_ENV"] == "kv:not-fetched" + assert os.environ["FROM_KV"] == "env:not-resolved-again" + assert os.environ["DUPLICATE_KV"] == "env:not-resolved-again" + assert os.environ["ESCAPED"] == "kv:not-a-secret" + mock_credential_cls.assert_called_once_with() mock_client_cls.assert_called_once_with(vault_url="https://myvault.vault.azure.net", credential=credential) - client.get_secret.assert_awaited_once_with("my-secret", version="v1") + assert client.get_secret.await_args_list == [ + mock.call("bootstrap", version="v1"), + mock.call("api-key"), + ] credential.__aenter__.assert_awaited_once() credential.__aexit__.assert_awaited_once() client.__aenter__.assert_awaited_once() client.__aexit__.assert_awaited_once() - - stream = mock_load_dotenv.call_args.kwargs["stream"] - assert stream.getvalue() == "AKV_VAR=from_secret\n" - assert mock_load_dotenv.call_args.kwargs["override"] is True - assert mock_print_msg.call_count == 2 + (len(secret_urls) > 1) + assert mock_print_msg.call_count == 3 async def test_load_env_from_akv_async_empty_secret_raises(self): - credential = mock.MagicMock() - credential.__aenter__ = mock.AsyncMock(return_value=credential) - credential.__aexit__ = mock.AsyncMock(return_value=None) - client = mock.MagicMock() - client.__aenter__ = mock.AsyncMock(return_value=client) - client.__aexit__ = mock.AsyncMock(return_value=None) + credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=None)) with ( mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") as mock_load_dotenv, pytest.raises(ValueError, match="has no value"), ): await _load_env_from_akv_async( @@ -557,23 +566,16 @@ async def test_load_env_from_akv_async_empty_secret_raises(self): silent=True, ) - mock_load_dotenv.assert_not_called() credential.__aexit__.assert_awaited_once() client.__aexit__.assert_awaited_once() async def test_load_env_from_akv_async_without_entries_raises(self): - credential = mock.MagicMock() - credential.__aenter__ = mock.AsyncMock(return_value=credential) - credential.__aexit__ = mock.AsyncMock(return_value=None) - client = mock.MagicMock() - client.__aenter__ = mock.AsyncMock(return_value=client) - client.__aexit__ = mock.AsyncMock(return_value=None) + credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="# comments only\n")) with ( mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - mock.patch("pyrit.setup.initialization.dotenv.load_dotenv", return_value=False), pytest.raises(ValueError, match="contains no environment entries"), ): await _load_env_from_akv_async( @@ -583,3 +585,25 @@ async def test_load_env_from_akv_async_without_entries_raises(self): credential.__aexit__.assert_awaited_once() client.__aexit__.assert_awaited_once() + + async def test_load_env_from_akv_async_failure_does_not_partially_update_environment(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value="GOOD=resolved\nBAD=kv:missing-value"), + types.SimpleNamespace(value=None), + ] + ) + + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="has no value"), + ): + await _load_env_from_akv_async( + secret_urls=["https://myvault.vault.azure.net/secrets/bootstrap"], + silent=True, + ) + + assert "GOOD" not in os.environ From d14821b624c163780727a8715baf99fb14418258 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 12 Aug 2026 11:53:39 -0400 Subject: [PATCH 03/16] FEAT: Added strict mode for KV --- .pyrit_conf_example | 7 + doc/getting_started/pyrit_conf.md | 24 ++- pyrit/setup/configuration_loader.py | 9 + pyrit/setup/initialization.py | 158 +++++++++++++---- tests/unit/setup/test_configuration_loader.py | 7 +- tests/unit/setup/test_initialization.py | 162 +++++++++++++++--- 6 files changed, 312 insertions(+), 55 deletions(-) diff --git a/.pyrit_conf_example b/.pyrit_conf_example index b9b2528a11..de75f08995 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -100,6 +100,9 @@ operation: op_trash_panda # The first secret's value must be the full contents of a .env file. # Values in that document may reference ambient variables with env:NAME or # scalar secrets in the same vault with kv:SECRET_NAME. +# Full same-vault URIs are also accepted, including a version to pin a secret: +# kv:https://my-vault.vault.azure.net/secrets/SECRET_NAME/SECRET_VERSION +# Cross-vault child references are rejected. # Referenced values are terminal scalars; they are not parsed for more references. # If multiple URLs are listed, only the first is used. # The AKV document replaces the default ~/.pyrit/.env source. Explicit env_files @@ -115,6 +118,10 @@ operation: op_trash_panda # Example: # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env +# +# Strict validation is enabled by default. Set this to false to skip malformed +# or valueless bootstrap entries with a warning while loading valid entries. +# env_akv_strict: false # Max Concurrent Scenario Runs # ---------------------------- diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 35d219fa8c..a332c79456 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -195,6 +195,7 @@ The root document can mix literal values with references to ambient environment ```dotenv OPENAI_CHAT_ENDPOINT="https://example.openai.azure.com/openai/v1" OPENAI_CHAT_KEY="kv:openai-chat-key" +PINNED_OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version-id" OPENAI_CHAT_MODEL="env:PYRIT_OPENAI_CHAT_MODEL" ``` @@ -207,6 +208,14 @@ For example, if `OPENAI_CHAT_KEY="kv:openai-chat-key"`, the value of the `openai 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 may use a secret name or a full secret URI from the bootstrap document's vault. A name or unversioned URI reads the latest secret version at initialization. Include the version in the URI to pin it. Cross-vault child references are rejected. + +```dotenv +LATEST_KEY="kv:openai-chat-key" +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" +``` + `literal:` is an escape hatch for a bootstrap value that begins with a reserved reference prefix. PyRIT removes `literal:` and returns the remainder without interpreting it as a reference. Quoting does not provide this escape because dotenv removes quotes while parsing. Values fetched from child secrets are already terminal and do not need this escape. ```dotenv @@ -218,6 +227,18 @@ Here, `REFERENCE` retrieves `openai-chat-key`, while `LITERAL_VALUE` becomes the The AKV document is loaded before explicit `env_files` or the default `~/.pyrit/.env.local`, allowing local values to override shared configuration without writing the fetched document to disk. +### `env_akv_strict` + +Controls validation of the Key Vault bootstrap document and defaults to `true`. + +```yaml +env_akv_strict: false +``` + +In strict mode, any malformed dotenv line or variable without an equals sign stops initialization. Empty assignments such as `OPTIONAL_VALUE=` remain valid. 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:` names, unresolved `env:` references, and a bootstrap document with no valid assignments still stop initialization. + ### `silent` If `true`, suppresses print statements during initialization. Useful for non-interactive environments or when embedding PyRIT in other tools. Defaults to `false`. @@ -227,7 +248,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` | @@ -349,6 +370,7 @@ initializers: # Optional Azure Key Vault root environment document # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env +# env_akv_strict: false # Optional; defaults to true # Suppress initialization messages silent: false diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 26a29c45b9..f34ab08b71 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -96,6 +96,8 @@ class ConfigurationLoader(YamlLoadable): None means "use defaults", [] means "load nothing". env_files: List of environment file paths to load. None means "use defaults (.env, .env.local)", [] means "load nothing". + env_akv_strict: Whether malformed or valueless entries in a Key Vault + bootstrap document should fail initialization. silent: Whether to suppress initialization messages. operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. @@ -135,6 +137,7 @@ class ConfigurationLoader(YamlLoadable): initialization_scripts: list[str] | None = None env_files: list[str] | None = None env_akv_ref: list[str] | None = None + env_akv_strict: bool = True silent: bool = False operator: str | None = None operation: str | None = None @@ -401,6 +404,7 @@ def load_with_overrides( initialization_scripts: Sequence[str] | None = None, env_files: Sequence[str] | None = None, env_akv_ref: Sequence[str] | None = None, + env_akv_strict: bool | None = None, ) -> "ConfigurationLoader": """ Load configuration with optional overrides. @@ -417,6 +421,7 @@ def load_with_overrides( initialization_scripts: Override for initialization script paths. env_files: Override for environment file paths. env_akv_ref: Override for Azure Key Vault secret URLs. + env_akv_strict: Override for strict Key Vault bootstrap validation. Returns: A merged ConfigurationLoader instance. @@ -479,6 +484,9 @@ def to_init_data(config: ConfigurationLoader) -> dict[str, Any]: if env_akv_ref is not None: config_data["env_akv_ref"] = list(env_akv_ref) + if env_akv_strict is not None: + config_data["env_akv_strict"] = env_akv_strict + return cls.from_dict(config_data) @classmethod @@ -614,6 +622,7 @@ async def initialize_pyrit_async(self) -> None: initializers=resolved_initializers if resolved_initializers else None, env_files=resolved_env_files, env_akv_ref=self.env_akv_ref, + env_akv_strict=self.env_akv_strict, silent=self.silent, ) diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 41fb313292..bcebc26761 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Literal, get_args import dotenv +from dotenv.parser import parse_stream from pyrit.common import path from pyrit.common.apply_defaults import reset_default_values @@ -171,42 +172,84 @@ def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: return vault_url, secret_name, secret_version -async def _load_env_from_akv_async(*, secret_urls: Sequence[str], silent: bool = False) -> None: +def _validate_dotenv_document( + document: str, + *, + strict: bool = True, + silent: bool = False, +) -> str: + """ + Validate that every dotenv binding uses ``NAME=VALUE`` syntax. + + Args: + document (str): The dotenv document to validate. + strict (bool): If True, reject any invalid entry. If False, warn and + allow python-dotenv to skip invalid entries. Defaults to True. + silent (bool): If True, suppress the console warning. Defaults to False. + + Returns: + str: The original document, or a sanitized document when strict is False. + + Raises: + ValueError: If strict is True and the document contains invalid entries. + """ + bindings = list(parse_stream(io.StringIO(document))) + malformed_lines = [str(binding.original.line) for binding in bindings if binding.error] + valueless_names = [binding.key for binding in bindings if binding.key is not None and binding.value is None] + issues: list[str] = [] + if malformed_lines: + issues.append("malformed entries at lines: " + ", ".join(malformed_lines)) + if valueless_names: + issues.append("variables without values: " + ", ".join(valueless_names)) + if not issues: + return document + + details = "; ".join(issues) + if strict: + raise ValueError("AKV environment document contains " + details) + + message = "AKV environment document contains invalid entries that will be skipped: " + details + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + return "".join( + binding.original.string + for binding in bindings + if not binding.error and not (binding.key is not None and binding.value is None) + ) + + +async def _load_env_from_akv_async( + *, + secret_url: str, + strict: bool = True, + silent: bool = False, +) -> None: """ Load environment variables from an Azure Key Vault secret. - The first secret URL identifies the bootstrap environment document. Values - in that document may directly reference scalar secrets in the same vault. - Additional root URLs are ignored. + The secret URL identifies the bootstrap environment document. Values in + that document may directly reference scalar secrets in the same vault. Authentication uses ``DefaultAzureCredential``, which silently tries managed identity, Azure CLI, VS Code credentials, etc., and falls back to interactive browser authentication when running locally. Args: - secret_urls (Sequence[str]): Sequence of AKV secret URLs. The first URL - must use the format ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. + secret_url (str): AKV secret URL in the format + ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. + strict (bool): If True, reject malformed or valueless dotenv entries. + If False, warn and skip those entries. Defaults to True. silent (bool): If True, suppresses print statements. Defaults to False. Raises: ImportError: If ``azure-keyvault-secrets`` is not installed. - ValueError: If no root secret is configured, the root URL is malformed, - or the bootstrap environment document cannot be fully resolved. + ValueError: If the root URL is malformed or the bootstrap environment + document cannot be fully resolved. """ - if not secret_urls: - raise ValueError("At least one env_akv_ref URL is required to load an environment document.") - from azure.identity.aio import DefaultAzureCredential from azure.keyvault.secrets.aio import SecretClient - secret_url = secret_urls[0] - if len(secret_urls) > 1: - _print_msg( - "Multiple env_akv_ref values were provided; using the first as the root environment document.", - quiet=silent, - log=True, - ) - _print_msg(f"Loading environment from AKV secret: {secret_url}", quiet=silent, log=True) vault_url, secret_name, secret_version = _parse_akv_secret_url(secret_url) ambient_environment = dict(os.environ) @@ -217,16 +260,11 @@ async def _load_env_from_akv_async(*, secret_urls: Sequence[str], silent: bool = if not secret.value: raise ValueError(f"AKV environment secret has no value: {secret_url}") - parsed_environment = dotenv.dotenv_values(stream=io.StringIO(secret.value), interpolate=True) + validated_document = _validate_dotenv_document(secret.value, strict=strict, silent=silent) + parsed_environment = dotenv.dotenv_values(stream=io.StringIO(validated_document), interpolate=True) if not parsed_environment: raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") - missing_values = [name for name, value in parsed_environment.items() if value is None] - if missing_values: - raise ValueError( - "AKV environment document contains variables without values: " + ", ".join(missing_values) - ) - resolved_secrets: dict[str, str] = {} resolved_environment: dict[str, str] = {} for variable_name, value in parsed_environment.items(): @@ -236,6 +274,7 @@ async def _load_env_from_akv_async(*, secret_urls: Sequence[str], silent: bool = value=value, variable_name=variable_name, secret_client=client, + vault_url=vault_url, ambient_environment=ambient_environment, resolved_secrets=resolved_secrets, ) @@ -272,11 +311,47 @@ def _validate_akv_secret_name(*, secret_name: str, variable_name: str) -> None: ) +def _resolve_akv_secret_reference( + *, + target: str, + variable_name: str, + vault_url: str, +) -> tuple[str, str | None, str]: + """ + Resolve a same-vault secret name or full secret URI. + + Args: + target (str): A secret name or full Key Vault secret URI. + variable_name (str): The environment variable receiving the secret. + vault_url (str): The bootstrap document's vault URL. + + Returns: + tuple[str, str | None, str]: Secret name, optional version, and cache key. + + Raises: + ValueError: If the target is invalid or references another vault. + """ + secret_name = target + secret_version: str | None = None + if target.casefold().startswith("https://"): + referenced_vault_url, secret_name, secret_version = _parse_akv_secret_url(target) + if referenced_vault_url.rstrip("/").casefold() != vault_url.rstrip("/").casefold(): + raise ValueError( + f"Cross-vault AKV reference for environment variable '{variable_name}' is not supported. " + f"Expected vault '{vault_url}', got '{referenced_vault_url}'." + ) + + _validate_akv_secret_name(secret_name=secret_name, variable_name=variable_name) + cache_key = f"{secret_name.casefold()}|{secret_version or ''}" + return secret_name, secret_version, cache_key + + async def _resolve_environment_value_async( *, value: str, variable_name: str, secret_client: "SecretClient", + vault_url: str, ambient_environment: dict[str, str], resolved_secrets: dict[str, str], ) -> str: @@ -287,6 +362,7 @@ async def _resolve_environment_value_async( value (str): The parsed bootstrap value. variable_name (str): The environment variable receiving the resolved value. secret_client (SecretClient): The client for the bootstrap document's vault. + vault_url (str): The bootstrap document's vault URL. ambient_environment (dict[str, str]): Snapshot used for ``env:`` references. resolved_secrets (dict[str, str]): Same-vault scalar cache keyed by secret name. @@ -314,14 +390,19 @@ async def _resolve_environment_value_async( ) return ambient_environment[target] - _validate_akv_secret_name(secret_name=target, variable_name=variable_name) - secret_cache_key = target.casefold() + secret_name, secret_version, secret_cache_key = _resolve_akv_secret_reference( + target=target, + variable_name=variable_name, + vault_url=vault_url, + ) if secret_cache_key in resolved_secrets: return resolved_secrets[secret_cache_key] - secret = await secret_client.get_secret(target) + secret = await secret_client.get_secret(secret_name, version=secret_version) if secret.value is None: - raise ValueError(f"AKV secret '{target}' referenced by environment variable '{variable_name}' has no value.") + raise ValueError( + f"AKV secret '{secret_name}' referenced by environment variable '{variable_name}' has no value." + ) resolved_secrets[secret_cache_key] = secret.value return secret.value @@ -375,6 +456,7 @@ async def initialize_pyrit_async( load_defaults: bool = True, env_files: Sequence[pathlib.Path] | None = None, env_akv_ref: Sequence[str] | None = None, + env_akv_strict: bool = True, silent: bool = False, **memory_instance_kwargs: Any, ) -> None: @@ -405,6 +487,8 @@ async def initialize_pyrit_async( The first secret's value must contain the bootstrap .env document; additional URLs are ignored. Loaded before ``env_files`` so local files take precedence over AKV. Requires ``azure-keyvault-secrets``. + env_akv_strict (bool): If True, reject malformed or valueless entries in the Key Vault + bootstrap document. If False, warn and skip those entries. Defaults to True. silent (bool): If True, suppresses print statements about environment file loading and schema migration. Defaults to False. **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. @@ -419,10 +503,18 @@ async def initialize_pyrit_async( env_files=env_files, silent=silent, ) - await _load_env_from_akv_async(secret_urls=env_akv_ref, silent=silent) + if len(env_akv_ref) > 1: + _print_msg( + "Multiple env_akv_ref values were provided; using the first as the root environment document.", + quiet=silent, + log=True, + ) + await _load_env_from_akv_async( + secret_url=env_akv_ref[0], + strict=env_akv_strict, + silent=silent, + ) - # PR review decision: .env.local and explicit files currently override the Key Vault document. - # The default .env is always skipped because Key Vault supplies the base environment. await asyncio.to_thread( _load_environment_files, env_files=env_files, diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 99bd2c5fbc..e36d14ecf7 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -42,6 +42,7 @@ def test_default_values(self): assert config.initialization_scripts is None # None means "use defaults" assert config.env_files is None # None means "use defaults" assert config.env_akv_ref is None + assert config.env_akv_strict is True assert config.silent is False def test_valid_memory_db_types_snake_case(self): @@ -147,6 +148,7 @@ def test_from_dict_with_all_fields(self): "initialization_scripts": ["/path/to/script.py"], "env_files": ["/path/to/.env"], "env_akv_ref": ["https://vault.vault.azure.net/secrets/one"], + "env_akv_strict": False, "silent": True, } config = ConfigurationLoader.from_dict(data) @@ -155,6 +157,7 @@ def test_from_dict_with_all_fields(self): assert config.initialization_scripts == ["/path/to/script.py"] assert config.env_files == ["/path/to/.env"] assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] + assert config.env_akv_strict is False assert config.silent is True def test_from_dict_filters_none_values(self): @@ -334,6 +337,7 @@ async def test_initialize_pyrit_async_basic(self, mock_init): assert call_kwargs["initializers"] is None assert call_kwargs["env_files"] is None assert call_kwargs["env_akv_ref"] is None + assert call_kwargs["env_akv_strict"] is True assert call_kwargs["silent"] is False @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") @@ -343,13 +347,14 @@ async def test_initialize_pyrit_async_with_env_akv_ref(self, mock_init): "https://vault.vault.azure.net/secrets/first", "https://vault.vault.azure.net/secrets/second/version", ] - config = ConfigurationLoader(memory_db_type="in_memory", env_akv_ref=refs) + config = ConfigurationLoader(memory_db_type="in_memory", env_akv_ref=refs, env_akv_strict=False) await config.initialize_pyrit_async() mock_init.assert_called_once() call_kwargs = mock_init.call_args.kwargs assert call_kwargs["env_akv_ref"] == refs + assert call_kwargs["env_akv_strict"] is False @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") @mock.patch("pyrit.registry.InitializerRegistry") diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 52ce0a92ef..a2ad00f212 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -17,6 +17,7 @@ _load_env_from_akv_async, _load_environment_files, _parse_akv_secret_url, + _parse_environment_value_reference, _warn_about_akv_environment_files, ) @@ -177,13 +178,17 @@ async def test_invalid_memory_type_raises_error(self, mock_load_env): @mock.patch("pyrit.setup.initialization._load_environment_files") @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, mock_set_memory): - """Test that env_akv_ref triggers AKV env loading.""" - refs = ["https://vault.vault.azure.net/secrets/test-secret"] + """Test that env_akv_ref loads only its first entry.""" + refs = [ + "https://vault.vault.azure.net/secrets/test-secret", + "https://vault.vault.azure.net/secrets/ignored", + ] await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) mock_load_akv.assert_awaited_once() - assert mock_load_akv.await_args.kwargs["secret_urls"] == refs + assert mock_load_akv.await_args.kwargs["secret_url"] == refs[0] + assert mock_load_akv.await_args.kwargs["strict"] is True assert mock_load_akv.await_args.kwargs["silent"] is False mock_load_env.assert_called_once() mock_set_memory.assert_called_once() @@ -211,7 +216,7 @@ async def test_initialize_loads_akv_before_env_files(self, mock_set_memory): def _record_warning(*, env_files, silent=False): call_order.append("warning") - async def _record_akv_call(*, secret_urls, silent=False): + async def _record_akv_call(*, secret_url, strict=True, silent=False): call_order.append("akv") def _record_env_file_call(*, env_files, silent=False, include_default_base=True): @@ -480,6 +485,15 @@ def _create_mock_akv_clients() -> tuple[mock.MagicMock, mock.MagicMock]: class TestAkvEnvironmentLoading: """Tests for AKV URL parsing and env loading helpers.""" + @pytest.mark.parametrize("prefix", ["kv", "akv", "azure_key_vault", "env_akv_ref"]) + def test_parse_environment_value_reference_accepts_akv_aliases(self, prefix): + assert _parse_environment_value_reference(f"{prefix}:api-key") == ("akv", "api-key") + + def test_parse_environment_value_reference_rejects_azure_app_service_syntax(self): + value = "@Microsoft.KeyVault(SecretUri=https://myvault.vault.azure.net/secrets/api-key)" + + assert _parse_environment_value_reference(value) is None + def test_parse_akv_secret_url_with_version(self): url = "https://myvault.vault.azure.net/secrets/my-secret/abc123" @@ -502,29 +516,24 @@ def test_parse_akv_secret_url_invalid_raises(self): with pytest.raises(ValueError, match="Invalid AKV secret URL"): _parse_akv_secret_url("https://myvault.vault.azure.net/not-secrets/my-secret") - async def test_load_env_from_akv_async_empty_urls_raises(self): - with pytest.raises(ValueError, match="At least one env_akv_ref URL is required"): - await _load_env_from_akv_async(secret_urls=[]) - - async def test_load_env_from_akv_async_uses_first_root_and_resolves_one_level(self): + async def test_load_env_from_akv_async_resolves_one_level(self): credential, client = _create_mock_akv_clients() root_document = ( "DIRECT=from-bootstrap\n" "FROM_ENV=env:SOURCE_VALUE\n" "FROM_KV=kv:api-key\n" - "DUPLICATE_KV=akv:API-KEY\n" + "DUPLICATE_KV=akv:https://MYVAULT.vault.azure.net/secrets/API-KEY\n" + "PINNED_KV=kv:https://myvault.vault.azure.net/secrets/api-key/version-2\n" "ESCAPED=literal:kv:not-a-secret" ) client.get_secret = mock.AsyncMock( side_effect=[ types.SimpleNamespace(value=root_document), types.SimpleNamespace(value="env:not-resolved-again"), + types.SimpleNamespace(value="pinned-secret-value"), ] ) - secret_urls = [ - "https://myvault.vault.azure.net/secrets/bootstrap/v1", - "https://myvault.vault.azure.net/secrets/ignored/v2", - ] + secret_url = "https://myvault.vault.azure.net/secrets/bootstrap/v1" with ( mock.patch.dict(os.environ, {"SOURCE_VALUE": "kv:not-fetched"}, clear=False), @@ -532,25 +541,50 @@ async def test_load_env_from_akv_async_uses_first_root_and_resolves_one_level(se mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, ): - await _load_env_from_akv_async(secret_urls=secret_urls, silent=True) + await _load_env_from_akv_async(secret_url=secret_url, silent=True) assert os.environ["DIRECT"] == "from-bootstrap" assert os.environ["FROM_ENV"] == "kv:not-fetched" assert os.environ["FROM_KV"] == "env:not-resolved-again" assert os.environ["DUPLICATE_KV"] == "env:not-resolved-again" + assert os.environ["PINNED_KV"] == "pinned-secret-value" assert os.environ["ESCAPED"] == "kv:not-a-secret" mock_credential_cls.assert_called_once_with() mock_client_cls.assert_called_once_with(vault_url="https://myvault.vault.azure.net", credential=credential) assert client.get_secret.await_args_list == [ mock.call("bootstrap", version="v1"), - mock.call("api-key"), + mock.call("api-key", version=None), + mock.call("api-key", version="version-2"), ] credential.__aenter__.assert_awaited_once() credential.__aexit__.assert_awaited_once() client.__aenter__.assert_awaited_once() client.__aexit__.assert_awaited_once() - assert mock_print_msg.call_count == 3 + assert mock_print_msg.call_count == 2 + + async def test_load_env_from_akv_async_rejects_cross_vault_reference(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + return_value=types.SimpleNamespace( + value="API_KEY=kv:https://other-vault.vault.azure.net/secrets/api-key/version-1" + ) + ) + + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="Cross-vault AKV reference"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert "API_KEY" not in os.environ + + client.get_secret.assert_awaited_once_with("bootstrap", version=None) async def test_load_env_from_akv_async_empty_secret_raises(self): credential, client = _create_mock_akv_clients() @@ -562,7 +596,7 @@ async def test_load_env_from_akv_async_empty_secret_raises(self): pytest.raises(ValueError, match="has no value"), ): await _load_env_from_akv_async( - secret_urls=["https://myvault.vault.azure.net/secrets/my-secret"], + secret_url="https://myvault.vault.azure.net/secrets/my-secret", silent=True, ) @@ -579,13 +613,101 @@ async def test_load_env_from_akv_async_without_entries_raises(self): pytest.raises(ValueError, match="contains no environment entries"), ): await _load_env_from_akv_async( - secret_urls=["https://myvault.vault.azure.net/secrets/my-secret"], + secret_url="https://myvault.vault.azure.net/secrets/my-secret", silent=True, ) credential.__aexit__.assert_awaited_once() client.__aexit__.assert_awaited_once() + @pytest.mark.parametrize( + ("document", "error"), + [ + ("GOOD=resolved\n=malformed\nOTHER=resolved", "malformed entries at lines: 2"), + ("MISSING_VALUE\n", "variables without values: MISSING_VALUE"), + ], + ) + async def test_load_env_from_akv_async_rejects_non_assignments(self, document, error): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) + + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match=error), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert "GOOD" not in os.environ + assert "OTHER" not in os.environ + + async def test_load_env_from_akv_async_allows_empty_assignment(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="EMPTY=")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert os.environ["EMPTY"] == "" + + async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entries(self, caplog, capsys): + credential, client = _create_mock_akv_clients() + document = "GOOD=resolved\n=malformed\nMISSING_VALUE\nOTHER=also-resolved" + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + caplog.at_level("WARNING", logger="pyrit.setup.initialization"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=False, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" + assert "MISSING_VALUE" not in os.environ + + output = capsys.readouterr().out + assert "WARNING: AKV environment document contains invalid entries that will be skipped" in output + assert "malformed entries at lines: 2" in output + assert "variables without values: MISSING_VALUE" in output + assert "GOOD" not in caplog.text + assert "resolved" not in caplog.text + + async def test_load_env_from_akv_async_non_strict_silent_logs_warning(self, caplog, capsys): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="GOOD=resolved\nMISSING_VALUE")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + caplog.at_level("WARNING", logger="pyrit.setup.initialization"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=True, + ) + + assert capsys.readouterr().out == "" + assert "variables without values: MISSING_VALUE" in caplog.text + async def test_load_env_from_akv_async_failure_does_not_partially_update_environment(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock( @@ -602,7 +724,7 @@ async def test_load_env_from_akv_async_failure_does_not_partially_update_environ pytest.raises(ValueError, match="has no value"), ): await _load_env_from_akv_async( - secret_urls=["https://myvault.vault.azure.net/secrets/bootstrap"], + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", silent=True, ) From de5be15f3268d9a36a60089f1c93015e1c9af30a Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 12 Aug 2026 12:05:37 -0400 Subject: [PATCH 04/16] FIX: Restore ambient-only setup path --- .pyrit_conf_example | 8 ++++++-- doc/getting_started/pyrit_conf.md | 8 +++++--- pyrit/setup/initialization.py | 15 ++++++-------- tests/unit/setup/test_initialization.py | 27 ++++++++++++++++++++----- 4 files changed, 39 insertions(+), 19 deletions(-) diff --git a/.pyrit_conf_example b/.pyrit_conf_example index de75f08995..b3ebfb8053 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -88,6 +88,8 @@ operation: op_trash_panda # - 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 +# - PyRIT reference prefixes are not resolved in local files; standard dotenv +# interpolation such as DERIVED=${BASE} remains enabled. # # Example: # env_files: @@ -109,9 +111,11 @@ operation: op_trash_panda # and the default ~/.pyrit/.env.local are loaded afterward and take precedence. # PyRIT emits a warning when these local files coexist with env_akv_ref so stale # configuration cannot silently mask or be mistaken for the Key Vault document. +# When migrating, remove or clear ~/.pyrit/.env and ~/.pyrit/.env.local, remove +# explicit env_files if Key Vault should be authoritative, and restart PyRIT. # Authentication uses DefaultAzureCredential (managed identity, Azure CLI, etc.). -# If env_akv_ref is omitted, at least one configured or default environment file -# must exist. System environment variables remain available but are not a source by themselves. +# If env_akv_ref and local files are omitted, PyRIT uses existing process +# environment variables and continues initialization. # # Requires: pip install azure-keyvault-secrets # diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index a332c79456..9b4dd84a6f 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -39,7 +39,7 @@ flowchart LR D --> F["3. ~/.pyrit/.env.local"] ``` -System environment variables are always the baseline, but initialization requires either an AKV root or at least one environment file. A system-environment-only configuration is not considered a complete source. +System environment variables are always the baseline. If no AKV root or environment file is available, PyRIT continues initialization using the existing process environment only. **Default file behavior** (no `env_akv_ref` or `env_files` field in `.pyrit_conf`): @@ -51,7 +51,7 @@ System environment variables are always the baseline, but initialization require **AKV behavior** (with `env_akv_ref`): The first referenced secret replaces `~/.pyrit/.env` as the root document. `~/.pyrit/.env.local` is loaded afterward if present. If multiple AKV URLs are configured, only the first is used. -PyRIT emits a warning when `env_akv_ref` is selected and default or explicit environment files coexist with it. The warning distinguishes files that are ignored from files that load afterward and override Key Vault values, making stale migration files visible at startup. +PyRIT emits a warning when `env_akv_ref` is selected and default or explicit environment files coexist with it. The warning distinguishes files that are ignored from files that load afterward and override Key Vault values, making stale migration files visible at startup. When migrating to Key Vault, clear or remove `~/.pyrit/.env` and `~/.pyrit/.env.local`, remove explicit `env_files` when Key Vault should be the only source, and restart PyRIT so values already present in the process environment cannot mask the Key Vault configuration. **Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. They override the AKV root when both fields are configured, and default paths are completely ignored. @@ -179,7 +179,9 @@ env_files: - /path/to/.env.local ``` -When `env_akv_ref` is not configured, an empty list or missing default files causes initialization to fail because no environment source is available. +Local environment files use standard dotenv behavior. PyRIT does not interpret `kv:`, `akv:`, `env:`, or `literal:` prefixes in `.env`, `.env.local`, or explicit `env_files`; those strings remain literal values. Standard dotenv interpolation such as `DERIVED=${BASE}` remains enabled. + +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` diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index bcebc26761..9b4ca03175 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -80,7 +80,7 @@ def _load_environment_files( ) else: _print_msg( - "No default environment files found.", + "No default environment files found. Using system environment variables only.", quiet=silent, log=True, ) @@ -138,7 +138,9 @@ def _warn_about_akv_environment_files( message = ( "env_akv_ref is configured, but local environment files were also found:\n- " + "\n- ".join(messages) - + "\nConfirm that this precedence is intentional." + + "\nWhen migrating to Key Vault, clear or remove ~/.pyrit/.env and ~/.pyrit/.env.local, " + "remove explicit env_files when Key Vault should be the only source, and restart PyRIT so stale " + "process values cannot mask Key Vault configuration." ) if not silent: print(f"WARNING: {message}") @@ -494,8 +496,7 @@ async def initialize_pyrit_async( **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. Raises: - ValueError: If an unsupported memory_db_type is provided, env_files contains non-existent files, - or neither env_akv_ref nor an environment file is available. + ValueError: If an unsupported memory_db_type is provided or env_files contains non-existent files. """ if env_akv_ref: await asyncio.to_thread( @@ -522,15 +523,11 @@ async def initialize_pyrit_async( include_default_base=False, ) else: - loaded_local_file = await asyncio.to_thread( + await asyncio.to_thread( _load_environment_files, env_files=env_files, silent=silent, ) - if not loaded_local_file: - raise ValueError( - "No environment source found. Configure env_akv_ref or provide at least one .env or .env.local file." - ) # Reset all default values before executing initialization scripts # This ensures a clean state for each initialization diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index a2ad00f212..408d493878 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -237,11 +237,10 @@ def _record_env_file_call(*, env_files, silent=False, include_default_base=True) mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_without_environment_source_raises(self, mock_set_memory): - with pytest.raises(ValueError, match="No environment source found"): - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[], load_defaults=False) + async def test_initialize_without_environment_file_uses_system_environment(self, mock_set_memory): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[], load_defaults=False) - mock_set_memory.assert_not_called() + mock_set_memory.assert_called_once() @pytest.fixture @@ -378,7 +377,9 @@ def test_warns_when_default_files_coexist_with_akv(self, mock_config_path, caplo assert output.startswith("WARNING: env_akv_ref is configured") assert f"{env_file} exists and will be ignored" in output assert f"{env_local_file} will load after Key Vault and override matching values" in output - assert "Confirm that this precedence is intentional." in output + assert "clear or remove ~/.pyrit/.env and ~/.pyrit/.env.local" in output + assert "remove explicit env_files when Key Vault should be the only source" in output + assert "restart PyRIT" in output assert caplog.records[0].levelname == "WARNING" @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") @@ -393,6 +394,7 @@ def test_akv_environment_file_warning_respects_silent(self, mock_config_path, ca assert capsys.readouterr().out == "" assert "will be ignored because Key Vault supplies the base environment" in caplog.text + assert "restart PyRIT" in caplog.text @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") async def test_loads_custom_env_files_in_order(self, mock_load_dotenv): @@ -416,6 +418,21 @@ async def test_loads_custom_env_files_in_order(self, mock_load_dotenv): call_args = [call[0][0] for call in mock_load_dotenv.call_args_list] assert call_args == [env1, env2, env3] + async def test_local_environment_files_keep_pyrit_references_literal(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text( + "BASE_VALUE=base\nKV_REFERENCE=kv:api-key\nENV_REFERENCE=env:SOURCE_VALUE\nINTERPOLATED=${BASE_VALUE}" + ) + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert os.environ["KV_REFERENCE"] == "kv:api-key" + assert os.environ["ENV_REFERENCE"] == "env:SOURCE_VALUE" + assert os.environ["INTERPOLATED"] == "base" + async def test_raises_error_for_nonexistent_env_file(self): """Test that ValueError is raised for non-existent env file.""" nonexistent = pathlib.Path("/nonexistent/path/.env") From dd309263fea15a08bb2c965297f785ffea91cf7c Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 12 Aug 2026 13:30:09 -0400 Subject: [PATCH 05/16] FEAT: Changed precedence for AKV secrets. No longer raises on both .env and AKV. --- .pyrit_conf_example | 8 +- doc/getting_started/pyrit_conf.md | 6 +- pyrit/setup/initialization.py | 142 ++++++++++++++++--- tests/unit/setup/test_initialization.py | 176 +++++++++++++++++------- 4 files changed, 260 insertions(+), 72 deletions(-) diff --git a/.pyrit_conf_example b/.pyrit_conf_example index b3ebfb8053..afe2325c3f 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -90,6 +90,8 @@ operation: op_trash_panda # - Set to list of paths: Load only the specified files # - PyRIT reference prefixes are not resolved in local files; standard dotenv # interpolation such as DERIVED=${BASE} remains enabled. +# - During PyRIT initialization, selected environment sources are staged and +# committed together only after every source loads successfully. # # Example: # env_files: @@ -123,8 +125,10 @@ operation: op_trash_panda # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env # -# Strict validation is enabled by default. Set this to false to skip malformed -# or valueless bootstrap entries with a warning while loading valid entries. +# Strict validation applies only to the Key Vault bootstrap and is enabled by +# default. Set this to false to skip malformed or valueless bootstrap entries +# with a warning while loading valid entries. Local files retain standard +# python-dotenv parsing regardless of this setting. # env_akv_strict: false # Max Concurrent Scenario Runs diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 9b4dd84a6f..43a5a5e908 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -179,7 +179,9 @@ env_files: - /path/to/.env.local ``` -Local environment files use standard dotenv behavior. PyRIT does not interpret `kv:`, `akv:`, `env:`, or `literal:` prefixes in `.env`, `.env.local`, or explicit `env_files`; those strings remain literal values. Standard dotenv interpolation such as `DERIVED=${BASE}` remains enabled. +Local environment files use standard dotenv behavior. PyRIT does not interpret `kv:`, `akv:`, `env:`, or `literal:` prefixes in `.env`, `.env.local`, or explicit `env_files`; those strings remain literal values. Standard dotenv interpolation such as `DERIVED=${BASE}` remains enabled. `env_akv_strict` does not apply to local files: malformed local lines retain python-dotenv's existing permissive skip-and-warn behavior. + +During `initialize_pyrit_async`, PyRIT stages the Key Vault mapping and every selected local file before updating `os.environ`. Later files can interpolate and override earlier staged values. If any selected source fails to load or resolve, none of the staged environment values are committed. Memory setup and initializers run after this environment commit and are outside this transaction. 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. @@ -231,7 +233,7 @@ The AKV document is loaded before explicit `env_files` or the default `~/.pyrit/ ### `env_akv_strict` -Controls validation of the Key Vault bootstrap document and defaults to `true`. +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 diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 9b4ca03175..8729a715bc 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -5,11 +5,12 @@ import logging import os import pathlib -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Literal, get_args import dotenv from dotenv.parser import parse_stream +from dotenv.variables import parse_variables from pyrit.common import path from pyrit.common.apply_defaults import reset_default_values @@ -54,7 +55,79 @@ def _load_environment_files( Raises: ValueError: If any provided env_files do not exist. """ - # Validate env_files exist if they were provided + selected_files = _select_environment_files( + env_files=env_files, + silent=silent, + include_default_base=include_default_base, + ) + for env_file in selected_files: + dotenv.load_dotenv(env_file, override=True, interpolate=True) + if not silent: + _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) + + return bool(selected_files) + + +def _resolve_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + base_environment: Mapping[str, str], + silent: bool = False, + include_default_base: bool = True, +) -> tuple[dict[str, str], bool]: + """ + Resolve environment files without mutating ``os.environ``. + + Args: + env_files: Optional sequence of environment file paths. If None, resolves + default files from the PyRIT configuration directory. + base_environment: Environment visible to interpolation before file values. + silent: If True, suppresses loading messages. Defaults to False. + include_default_base: If False and env_files is None, skips the default + .env file while still resolving .env.local. Defaults to True. + + Returns: + tuple[dict[str, str], bool]: Resolved values and whether any file was selected. + + Raises: + ValueError: If any explicitly provided environment file does not exist. + """ + selected_files = _select_environment_files( + env_files=env_files, + silent=silent, + include_default_base=include_default_base, + ) + if _dotenv_loading_disabled(): + return {}, bool(selected_files) + + staged_environment = dict(base_environment) + resolved_environment: dict[str, str] = {} + for env_file in selected_files: + raw_values = dotenv.dotenv_values(dotenv_path=env_file, interpolate=False) + file_values = _interpolate_dotenv_values(values=raw_values, base_environment=staged_environment) + staged_environment.update(file_values) + resolved_environment.update(file_values) + if not silent: + _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) + + return resolved_environment, bool(selected_files) + + +def _select_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool, + include_default_base: bool, +) -> list[pathlib.Path]: + """ + Select and validate environment files without reading their contents. + + Returns: + list[pathlib.Path]: Environment files in load order. + + Raises: + ValueError: If an explicitly provided environment file does not exist. + """ if env_files is not None: if not silent: _print_msg(f"Loading custom environment files: {[str(f) for f in env_files]}", quiet=silent, log=True) @@ -87,12 +160,37 @@ def _load_environment_files( env_files = default_files - for env_file in env_files: - dotenv.load_dotenv(env_file, override=True, interpolate=True) - if not silent: - _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) + return list(env_files) - return bool(env_files) + +def _interpolate_dotenv_values( + *, + values: Mapping[str, str | None], + base_environment: Mapping[str, str], +) -> dict[str, str]: + """ + Resolve dotenv interpolation against a staged environment mapping. + + Returns: + dict[str, str]: Interpolated assignments, excluding valueless entries. + """ + visible_environment: dict[str, str | None] = dict(base_environment) + resolved_values: dict[str, str] = {} + for name, value in values.items(): + if value is None: + visible_environment[name] = None + continue + + resolved_value = "".join(atom.resolve(visible_environment) for atom in parse_variables(value)) + visible_environment[name] = resolved_value + resolved_values[name] = resolved_value + + return resolved_values + + +def _dotenv_loading_disabled() -> bool: + value = os.environ.get("PYTHON_DOTENV_DISABLED", "") + return value.casefold() in {"1", "true", "t", "yes", "y"} def _print_msg(message: str, quiet: bool, log: bool) -> None: @@ -226,7 +324,7 @@ async def _load_env_from_akv_async( secret_url: str, strict: bool = True, silent: bool = False, -) -> None: +) -> dict[str, str]: """ Load environment variables from an Azure Key Vault secret. @@ -244,6 +342,9 @@ async def _load_env_from_akv_async( If False, warn and skip those entries. Defaults to True. silent (bool): If True, suppresses print statements. Defaults to False. + Returns: + dict[str, str]: The fully resolved Key Vault environment mapping. + Raises: ImportError: If ``azure-keyvault-secrets`` is not installed. ValueError: If the root URL is malformed or the bootstrap environment @@ -281,9 +382,7 @@ async def _load_env_from_akv_async( resolved_secrets=resolved_secrets, ) - os.environ.update(resolved_environment) - - _print_msg(f"Loaded environment from AKV secret: {secret_url}", quiet=silent, log=True) + return resolved_environment def _parse_environment_value_reference(value: str) -> tuple[str, str] | None: @@ -498,6 +597,8 @@ async def initialize_pyrit_async( Raises: ValueError: If an unsupported memory_db_type is provided or env_files contains non-existent files. """ + base_environment = dict(os.environ) + environment_updates: dict[str, str] = {} if env_akv_ref: await asyncio.to_thread( _warn_about_akv_environment_files, @@ -510,24 +611,31 @@ async def initialize_pyrit_async( quiet=silent, log=True, ) - await _load_env_from_akv_async( + akv_environment = await _load_env_from_akv_async( secret_url=env_akv_ref[0], strict=env_akv_strict, silent=silent, ) - - await asyncio.to_thread( - _load_environment_files, + staged_environment = {**base_environment, **akv_environment} + local_environment, _ = await asyncio.to_thread( + _resolve_environment_files, env_files=env_files, + base_environment=staged_environment, silent=silent, include_default_base=False, ) + environment_updates.update(akv_environment) + environment_updates.update(local_environment) else: - await asyncio.to_thread( - _load_environment_files, + local_environment, _ = await asyncio.to_thread( + _resolve_environment_files, env_files=env_files, + base_environment=base_environment, silent=silent, ) + environment_updates.update(local_environment) + + os.environ.update(environment_updates) # Reset all default values before executing initialization scripts # This ensures a clean state for each initialization diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 408d493878..7ff8bc8bab 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -127,19 +127,17 @@ def setup_method(self) -> None: reset_default_values() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files") - async def test_initialize_basic(self, mock_load_env, mock_set_memory): + @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) + async def test_initialize_basic(self, mock_resolve_env, mock_set_memory): """Test basic initialization.""" - mock_load_env.return_value = True - await initialize_pyrit_async(memory_db_type=IN_MEMORY, load_defaults=False) - mock_load_env.assert_called_once() + mock_resolve_env.assert_called_once() mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files") - async def test_initialize_with_script(self, mock_load_env, mock_set_memory): + @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) + async def test_initialize_with_script(self, mock_resolve_env, mock_set_memory): """Test initialization with a script.""" with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f: f.write( @@ -163,49 +161,49 @@ async def initialize_async(self) -> None: try: await initialize_pyrit_async(memory_db_type=IN_MEMORY, initialization_scripts=[script_path]) - mock_load_env.assert_called_once() + mock_resolve_env.assert_called_once() mock_set_memory.assert_called_once() finally: os.unlink(script_path) - @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=True) - async def test_invalid_memory_type_raises_error(self, mock_load_env): + @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) + async def test_invalid_memory_type_raises_error(self, mock_resolve_env): """Test that invalid memory type raises ValueError.""" with pytest.raises(ValueError, match="is not a supported type"): await initialize_pyrit_async(memory_db_type="InvalidType", load_defaults=False) # type: ignore[arg-type] @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files") + @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) - async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, mock_set_memory): + async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_resolve_env, mock_set_memory): """Test that env_akv_ref loads only its first entry.""" refs = [ "https://vault.vault.azure.net/secrets/test-secret", "https://vault.vault.azure.net/secrets/ignored", ] + mock_load_akv.return_value = {} + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) mock_load_akv.assert_awaited_once() assert mock_load_akv.await_args.kwargs["secret_url"] == refs[0] assert mock_load_akv.await_args.kwargs["strict"] is True assert mock_load_akv.await_args.kwargs["silent"] is False - mock_load_env.assert_called_once() + mock_resolve_env.assert_called_once() mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files") + @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) async def test_initialize_with_empty_env_akv_ref_does_not_load_akv( - self, mock_load_akv, mock_load_env, mock_set_memory + self, mock_load_akv, mock_resolve_env, mock_set_memory ): """Test that empty env_akv_ref does not invoke AKV loading.""" - mock_load_env.return_value = True - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=[], load_defaults=False) mock_load_akv.assert_not_called() - mock_load_env.assert_called_once() + mock_resolve_env.assert_called_once() mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") @@ -218,24 +216,84 @@ def _record_warning(*, env_files, silent=False): async def _record_akv_call(*, secret_url, strict=True, silent=False): call_order.append("akv") + return {"FROM_AKV": "shared"} - def _record_env_file_call(*, env_files, silent=False, include_default_base=True): + def _record_env_file_call(*, env_files, base_environment, silent=False, include_default_base=True): call_order.append("env_files") assert include_default_base is False - return True + assert base_environment["FROM_AKV"] == "shared" + return {"FROM_LOCAL": "override"}, True refs = ["https://vault.vault.azure.net/secrets/test-secret"] - with ( - mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files", side_effect=_record_warning), - mock.patch("pyrit.setup.initialization._load_env_from_akv_async", side_effect=_record_akv_call), - mock.patch("pyrit.setup.initialization._load_environment_files", side_effect=_record_env_file_call), - ): - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files", side_effect=_record_warning), + mock.patch("pyrit.setup.initialization._load_env_from_akv_async", side_effect=_record_akv_call), + mock.patch("pyrit.setup.initialization._resolve_environment_files", side_effect=_record_env_file_call), + ): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) + + assert os.environ == {"FROM_AKV": "shared", "FROM_LOCAL": "override"} assert call_order == ["warning", "akv", "env_files"] mock_set_memory.assert_called_once() + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_akv_and_local_files_are_applied_atomically(self, mock_set_memory): + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + nonexistent = pathlib.Path("/nonexistent/.env") + + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + return_value={"FROM_AKV": "resolved"}, + ), + pytest.raises(ValueError, match="Environment file not found"), + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=refs, + env_files=[nonexistent], + load_defaults=False, + ) + + assert "FROM_AKV" not in os.environ + + mock_set_memory.assert_not_called() + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_stages_local_overrides_on_akv_environment(self, mock_set_memory): + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + with tempfile.TemporaryDirectory() as temp_dir: + local_file = pathlib.Path(temp_dir) / ".env.local" + local_file.write_text("DERIVED=${BASE}\nBASE=local") + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + return_value={"BASE": "akv", "ONLY_AKV": "shared"}, + ), + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=refs, + env_files=[local_file], + load_defaults=False, + ) + + assert os.environ["BASE"] == "local" + assert os.environ["DERIVED"] == "akv" + assert os.environ["ONLY_AKV"] == "shared" + + mock_set_memory.assert_called_once() + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") async def test_initialize_without_environment_file_uses_system_environment(self, mock_set_memory): await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[], load_defaults=False) @@ -433,6 +491,23 @@ async def test_local_environment_files_keep_pyrit_references_literal(self): assert os.environ["ENV_REFERENCE"] == "env:SOURCE_VALUE" assert os.environ["INTERPOLATED"] == "base" + async def test_env_akv_strict_does_not_validate_local_environment_files(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("GOOD=resolved\n=malformed\nOTHER=also-resolved") + + with mock.patch.dict(os.environ, {}, clear=True): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_files=[env_file], + env_akv_strict=True, + load_defaults=False, + silent=True, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" + async def test_raises_error_for_nonexistent_env_file(self): """Test that ValueError is raised for non-existent env file.""" nonexistent = pathlib.Path("/nonexistent/path/.env") @@ -461,10 +536,8 @@ async def test_initialize_pyrit_raises_for_nonexistent_env_file(self, mock_set_m with pytest.raises(ValueError, match="Environment file not found"): await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[nonexistent]) - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") - @mock.patch("pyrit.setup.initialization.path.HOME_PATH") @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_custom_env_files_override_default_behavior(self, mock_set_memory, mock_home_path, mock_load_dotenv): + async def test_custom_env_files_override_default_behavior(self, mock_set_memory): """Test that passing custom env_files prevents loading default files.""" with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -479,14 +552,12 @@ async def test_custom_env_files_override_default_behavior(self, mock_set_memory, custom_env = temp_path / ".env.custom" custom_env.write_text("CUSTOM=value") - mock_home_path.__truediv__ = lambda self, other: temp_path / other + with mock.patch.dict(os.environ, {}, clear=True): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[custom_env], load_defaults=False) - # Pass custom env_files - should NOT load defaults - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[custom_env], load_defaults=False) - - # Verify only custom file was loaded, not the default ones - assert mock_load_dotenv.call_count == 1 - assert mock_load_dotenv.call_args[0][0] == custom_env + assert os.environ["CUSTOM"] == "value" + assert "DEFAULT" not in os.environ + assert "DEFAULT_LOCAL" not in os.environ def _create_mock_akv_clients() -> tuple[mock.MagicMock, mock.MagicMock]: @@ -558,14 +629,17 @@ async def test_load_env_from_akv_async_resolves_one_level(self): mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, ): - await _load_env_from_akv_async(secret_url=secret_url, silent=True) - - assert os.environ["DIRECT"] == "from-bootstrap" - assert os.environ["FROM_ENV"] == "kv:not-fetched" - assert os.environ["FROM_KV"] == "env:not-resolved-again" - assert os.environ["DUPLICATE_KV"] == "env:not-resolved-again" - assert os.environ["PINNED_KV"] == "pinned-secret-value" - assert os.environ["ESCAPED"] == "kv:not-a-secret" + resolved_environment = await _load_env_from_akv_async(secret_url=secret_url, silent=True) + + assert resolved_environment == { + "DIRECT": "from-bootstrap", + "FROM_ENV": "kv:not-fetched", + "FROM_KV": "env:not-resolved-again", + "DUPLICATE_KV": "env:not-resolved-again", + "PINNED_KV": "pinned-secret-value", + "ESCAPED": "kv:not-a-secret", + } + assert "DIRECT" not in os.environ mock_credential_cls.assert_called_once_with() mock_client_cls.assert_called_once_with(vault_url="https://myvault.vault.azure.net", credential=credential) @@ -578,7 +652,7 @@ async def test_load_env_from_akv_async_resolves_one_level(self): credential.__aexit__.assert_awaited_once() client.__aenter__.assert_awaited_once() client.__aexit__.assert_awaited_once() - assert mock_print_msg.call_count == 2 + mock_print_msg.assert_called_once() async def test_load_env_from_akv_async_rejects_cross_vault_reference(self): credential, client = _create_mock_akv_clients() @@ -671,12 +745,13 @@ async def test_load_env_from_akv_async_allows_empty_assignment(self): mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), ): - await _load_env_from_akv_async( + resolved_environment = await _load_env_from_akv_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", silent=True, ) - assert os.environ["EMPTY"] == "" + assert resolved_environment["EMPTY"] == "" + assert "EMPTY" not in os.environ async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entries(self, caplog, capsys): credential, client = _create_mock_akv_clients() @@ -689,15 +764,14 @@ async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entrie mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), caplog.at_level("WARNING", logger="pyrit.setup.initialization"), ): - await _load_env_from_akv_async( + resolved_environment = await _load_env_from_akv_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", strict=False, silent=False, ) - assert os.environ["GOOD"] == "resolved" - assert os.environ["OTHER"] == "also-resolved" - assert "MISSING_VALUE" not in os.environ + assert resolved_environment == {"GOOD": "resolved", "OTHER": "also-resolved"} + assert "GOOD" not in os.environ output = capsys.readouterr().out assert "WARNING: AKV environment document contains invalid entries that will be skipped" in output From be956df5874d2e5d839f047d53c203b796b087ce Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 12 Aug 2026 17:34:43 -0400 Subject: [PATCH 06/16] FEAT: Refactored precedence order and simplified environment variable resolution --- .pyrit_conf_example | 31 +- doc/getting_started/pyrit_conf.md | 53 +- pyrit/exceptions/__init__.py | 2 + pyrit/exceptions/exception_classes.py | 19 + pyrit/setup/configuration_loader.py | 27 +- pyrit/setup/initialization.py | 435 ++++++++++---- tests/unit/exceptions/test_exceptions.py | 9 + tests/unit/setup/test_configuration_loader.py | 45 +- tests/unit/setup/test_initialization.py | 554 +++++++++++++----- 9 files changed, 830 insertions(+), 345 deletions(-) diff --git a/.pyrit_conf_example b/.pyrit_conf_example index afe2325c3f..3ed13b71b9 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -88,8 +88,11 @@ operation: op_trash_panda # - 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 -# - PyRIT reference prefixes are not resolved in local files; standard dotenv -# interpolation such as DERIVED=${BASE} remains enabled. +# - Local files retain standard dotenv parsing and interpolation. After source +# precedence is applied, PyRIT resolves complete-value kv:/env: references in +# winning values from any source. Overridden references are not fetched. +# - Interpolation follows load order: .env.local can reference .env, but .env +# cannot see variables introduced only by the later .env.local. # - During PyRIT initialization, selected environment sources are staged and # committed together only after every source loads successfully. # @@ -100,35 +103,39 @@ operation: op_trash_panda # Azure Key Vault Environment References # --------------------------------------- -# List of AKV secret URLs to load during initialization. -# The first secret's value must be the full contents of a .env file. -# Values in that document may reference ambient variables with env:NAME or -# scalar secrets in the same vault with kv:SECRET_NAME. -# Full same-vault URIs are also accepted, including a version to pin a secret: +# AKV secret URL whose value is the bootstrap .env document. +# Winning values may reference another merged environment key with env:NAME, +# falling back to the existing process environment, or reference a scalar +# secret in the same vault using a full URL: +# kv:https://my-vault.vault.azure.net/secrets/SECRET_NAME +# Include a version to pin a secret: # kv:https://my-vault.vault.azure.net/secrets/SECRET_NAME/SECRET_VERSION +# Short secret names such as kv:SECRET_NAME are rejected. # Cross-vault child references are rejected. +# Referenced secrets are not cached; each kv: occurrence performs a vault read. # Referenced values are terminal scalars; they are not parsed for more references. -# If multiple URLs are listed, only the first is used. -# The AKV document replaces the default ~/.pyrit/.env source. Explicit env_files -# and the default ~/.pyrit/.env.local are loaded afterward and take precedence. +# Source precedence is AKV bootstrap -> ~/.pyrit/.env -> ~/.pyrit/.env.local. +# Explicit env_files replace the default files and load after the AKV bootstrap. # PyRIT emits a warning when these local files coexist with env_akv_ref so stale # configuration cannot silently mask or be mistaken for the Key Vault document. # When migrating, remove or clear ~/.pyrit/.env and ~/.pyrit/.env.local, remove # explicit env_files if Key Vault should be authoritative, and restart PyRIT. # Authentication uses DefaultAzureCredential (managed identity, Azure CLI, etc.). +# Key Vault operations use up to three retries with exponential backoff and +# raise KeyVaultInitializationException on bootstrap or secret-resolution failure. # If env_akv_ref and local files are omitted, PyRIT uses existing process # environment variables and continues initialization. # # Requires: pip install azure-keyvault-secrets # # Example: -# env_akv_ref: -# - https://my-vault.vault.azure.net/secrets/my-pyrit-env +# env_akv_ref: https://my-vault.vault.azure.net/secrets/my-pyrit-env # # Strict validation applies only to the Key Vault bootstrap and is enabled by # default. Set this to false to skip malformed or valueless bootstrap entries # with a warning while loading valid entries. Local files retain standard # python-dotenv parsing regardless of this setting. +# Empty assignments (NAME=) and child secrets containing an empty string are valid. # env_akv_strict: false # Max Concurrent Scenario Runs diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 43a5a5e908..baa0023961 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -32,11 +32,13 @@ When PyRIT initializes, environment variables are loaded in a specific order. ** ```{mermaid} flowchart LR - A["1. System Environment"] --> B{"env_akv_ref configured?"} - B -->|Yes| C["2. First AKV secret"] - B -->|No| D["2. ~/.pyrit/.env"] - C --> E["3. Explicit env_files or ~/.pyrit/.env.local"] - D --> F["3. ~/.pyrit/.env.local"] + A["System environment"] --> B{"env_akv_ref configured?"} + B -->|Yes| C["AKV bootstrap"] + B -->|No| D{"Explicit env_files?"} + C --> D + D -->|Yes| E["Explicit files in order"] + D -->|No| F["~/.pyrit/.env"] + F --> G["~/.pyrit/.env.local"] ``` System environment variables are always the baseline. If no AKV root or environment file is available, PyRIT continues initialization using the existing process environment only. @@ -49,7 +51,7 @@ System environment variables are always the baseline. If no AKV root or environm | Medium | `~/.pyrit/.env` | Default config file (loaded if it exists) | | Highest | `~/.pyrit/.env.local` | Local overrides (loaded if it exists) | -**AKV behavior** (with `env_akv_ref`): The first referenced secret replaces `~/.pyrit/.env` as the root document. `~/.pyrit/.env.local` is loaded afterward if present. If multiple AKV URLs are configured, only the first is used. +**AKV behavior** (with `env_akv_ref`): The referenced secret is the lowest-priority file source. Unless custom `env_files` are configured, `~/.pyrit/.env` loads afterward and `~/.pyrit/.env.local` loads last; either may override matching Key Vault values. PyRIT emits a warning when `env_akv_ref` is selected and default or explicit environment files coexist with it. The warning distinguishes files that are ignored from files that load afterward and override Key Vault values, making stale migration files visible at startup. When migrating to Key Vault, clear or remove `~/.pyrit/.env` and `~/.pyrit/.env.local`, remove explicit `env_files` when Key Vault should be the only source, and restart PyRIT so values already present in the process environment cannot mask the Key Vault configuration. @@ -179,43 +181,49 @@ env_files: - /path/to/.env.local ``` -Local environment files use standard dotenv behavior. PyRIT does not interpret `kv:`, `akv:`, `env:`, or `literal:` prefixes in `.env`, `.env.local`, or explicit `env_files`; those strings remain literal values. Standard dotenv interpolation such as `DERIVED=${BASE}` remains enabled. `env_akv_strict` does not apply to local files: malformed local lines retain python-dotenv's existing permissive skip-and-warn behavior. +Local environment files use standard dotenv parsing and interpolation. `env_akv_strict` does not apply to them: malformed local lines retain python-dotenv's existing permissive skip-and-warn behavior. -During `initialize_pyrit_async`, PyRIT stages the Key Vault mapping and every selected local file before updating `os.environ`. Later files can interpolate and override earlier staged values. If any selected source fails to load or resolve, none of the staged environment values are committed. Memory setup and initializers run after this environment commit and are outside this transaction. +During `initialize_pyrit_async`, PyRIT first applies source precedence across the optional Key Vault bootstrap, `.env`, and `.env.local` or explicit `env_files`. It then resolves complete-value `kv:`, `akv:`, `azure_key_vault:`, `env_akv_ref:`, `env:`, and `literal:` references in the winning values, regardless of which source declared them. References overridden by a later source are never fetched. A local file can therefore use a full Key Vault URL even when no bootstrap document is configured. + +An `env:NAME` alias first reads the winning `NAME` value from the merged sources. If no source declares `NAME`, it falls back to the process environment captured before initialization. Merged values take precedence over ambient values with the same name. Alias resolution is one hop. Direct self-reference such as `MODEL="env:MODEL"` is rejected; use a distinct source variable such as `MODEL="env:PYRIT_MODEL"`. + +Interpolation follows 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. + +PyRIT stages the Key Vault mapping and every selected local file before updating `os.environ`. Later files can interpolate and override earlier staged values. If any selected source fails to load or resolve, none of the staged environment values are committed. Memory setup and initializers run after this environment commit and are outside this transaction. 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` -Azure Key Vault secret URLs used to obtain the root environment document. The first URL is used; its secret value must contain dotenv-formatted entries. Authentication uses `DefaultAzureCredential`. +Azure Key Vault secret URL used to obtain the root environment document. Its value must contain dotenv-formatted entries. Authentication uses `DefaultAzureCredential`. ```yaml -env_akv_ref: - - https://my-vault.vault.azure.net/secrets/my-pyrit-env +env_akv_ref: https://my-vault.vault.azure.net/secrets/my-pyrit-env ``` -The root document can mix literal values with references to ambient environment variables and scalar secrets in the same vault: +The root document can mix literal values with references to merged or ambient environment variables and scalar secrets in the same vault: ```dotenv OPENAI_CHAT_ENDPOINT="https://example.openai.azure.com/openai/v1" -OPENAI_CHAT_KEY="kv:openai-chat-key" +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="env:PYRIT_OPENAI_CHAT_MODEL" ``` Resolution is deliberately limited to two levels: -1. PyRIT fetches the first `env_akv_ref` secret and parses it as the bootstrap dotenv document. -2. For each reference in that document, PyRIT either copies one ambient `env:` value or fetches one scalar secret from the same vault. The resulting value is final and is not parsed as another reference. +1. PyRIT fetches the `env_akv_ref` secret and parses it as the bootstrap dotenv document. +2. After all sources are merged, PyRIT either copies one merged-or-ambient `env:` value or fetches one scalar secret from the same vault. The resulting value is final and is not parsed as another reference. -For example, if `OPENAI_CHAT_KEY="kv: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`. +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 may use a secret name or a full secret URI from the bootstrap document's vault. A name or unversioned URI reads the latest secret version at initialization. Include the version in the URI to pin it. Cross-vault child references are rejected. +A Key Vault reference must use a full secret URL from the bootstrap document's vault. An unversioned URL reads the latest secret version at initialization. Include the version in the URL to pin it. Short names and cross-vault child references are rejected. + +PyRIT does not cache referenced secrets. Each `kv:` occurrence performs a Key Vault read during initialization, including repeated references to the same URI. ```dotenv -LATEST_KEY="kv:openai-chat-key" 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" ``` @@ -223,7 +231,7 @@ PINNED_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version- `literal:` is an escape hatch for a bootstrap value that begins with a reserved reference prefix. PyRIT removes `literal:` and returns the remainder without interpreting it as a reference. Quoting does not provide this escape because dotenv removes quotes while parsing. Values fetched from child secrets are already terminal and do not need this escape. ```dotenv -REFERENCE="kv:openai-chat-key" +REFERENCE="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key" LITERAL_VALUE="literal:kv:not-a-secret-name" ``` @@ -239,10 +247,12 @@ Controls validation only of the Key Vault bootstrap document and defaults to `tr env_akv_strict: false ``` -In strict mode, any malformed dotenv line or variable without an equals sign stops initialization. Empty assignments such as `OPTIONAL_VALUE=` remain valid. 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. +In strict mode, any malformed dotenv line or variable without an equals sign stops initialization. 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:` names, unresolved `env:` references, and a bootstrap document with no valid assignments still stop initialization. +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. + ### `silent` If `true`, suppresses print statements during initialization. Useful for non-interactive environments or when embedding PyRIT in other tools. Defaults to `false`. @@ -372,8 +382,7 @@ initializers: # - /path/to/.env.local # Optional Azure Key Vault root environment document -# env_akv_ref: -# - https://my-vault.vault.azure.net/secrets/my-pyrit-env +# env_akv_ref: https://my-vault.vault.azure.net/secrets/my-pyrit-env # env_akv_strict: false # Optional; defaults to true # Suppress initialization messages diff --git a/pyrit/exceptions/__init__.py b/pyrit/exceptions/__init__.py index 9e8a074b67..c4be6078ce 100644 --- a/pyrit/exceptions/__init__.py +++ b/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", diff --git a/pyrit/exceptions/exception_classes.py b/pyrit/exceptions/exception_classes.py index b2aa780083..d6edf92999 100644 --- a/pyrit/exceptions/exception_classes.py +++ b/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.""" diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index f34ab08b71..0201f3fe82 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -136,7 +136,7 @@ class ConfigurationLoader(YamlLoadable): initializers: list[str | dict[str, Any]] = field(default_factory=list) initialization_scripts: list[str] | None = None env_files: list[str] | None = None - env_akv_ref: list[str] | None = None + env_akv_ref: str | None = None env_akv_strict: bool = True silent: bool = False operator: str | None = None @@ -150,8 +150,21 @@ def __post_init__(self) -> None: """Validate and normalize the configuration after loading.""" self._normalize_memory_db_type() self._normalize_initializers() + self._validate_env_akv_ref() self._normalize_server() + def _validate_env_akv_ref(self) -> None: + """ + Validate the Key Vault bootstrap secret reference. + + Raises: + ValueError: If env_akv_ref is not one non-empty string. + """ + if self.env_akv_ref is None: + return + if not isinstance(self.env_akv_ref, str) or not self.env_akv_ref.strip(): + raise ValueError("env_akv_ref must be one non-empty Azure Key Vault secret URL.") + def _normalize_memory_db_type(self) -> None: """ Normalize and validate memory_db_type. @@ -403,7 +416,7 @@ def load_with_overrides( initializers: Sequence[str | dict[str, Any]] | None = None, initialization_scripts: Sequence[str] | None = None, env_files: Sequence[str] | None = None, - env_akv_ref: Sequence[str] | None = None, + env_akv_ref: str | None = None, env_akv_strict: bool | None = None, ) -> "ConfigurationLoader": """ @@ -420,7 +433,7 @@ def load_with_overrides( initializers: Override for initializer list. initialization_scripts: Override for initialization script paths. env_files: Override for environment file paths. - env_akv_ref: Override for Azure Key Vault secret URLs. + env_akv_ref: Override for the Azure Key Vault bootstrap secret URL. env_akv_strict: Override for strict Key Vault bootstrap validation. Returns: @@ -482,7 +495,7 @@ def to_init_data(config: ConfigurationLoader) -> dict[str, Any]: config_data["env_files"] = list(env_files) if env_akv_ref is not None: - config_data["env_akv_ref"] = list(env_akv_ref) + config_data["env_akv_ref"] = env_akv_ref if env_akv_strict is not None: config_data["env_akv_strict"] = env_akv_strict @@ -588,12 +601,12 @@ def resolve_env_files(self) -> Sequence[pathlib.Path] | None: return resolved - def resolve_env_akv_ref(self) -> list[str] | None: + def resolve_env_akv_ref(self) -> str | None: """ - Return the list of AKV secret URLs, or ``None`` when not configured. + Return the AKV bootstrap secret URL, or ``None`` when not configured. Returns: - list[str] | None: The configured AKV secret URLs, or ``None``. + str | None: The configured AKV bootstrap secret URL, or ``None``. """ return self.env_akv_ref diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 8729a715bc..b192548430 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -14,6 +14,7 @@ from pyrit.common import path from pyrit.common.apply_defaults import reset_default_values +from pyrit.exceptions import KeyVaultInitializationException from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory if TYPE_CHECKING: @@ -29,6 +30,8 @@ MemoryDatabaseType = Literal["InMemory", "SQLite", "AzureSQL"] _AKV_REFERENCE_PREFIXES = frozenset({"akv", "kv", "azure_key_vault", "env_akv_ref"}) +_AKV_RETRY_TOTAL = 3 +_AKV_RETRY_BACKOFF_FACTOR = 0.8 def _load_environment_files( @@ -55,17 +58,14 @@ def _load_environment_files( Raises: ValueError: If any provided env_files do not exist. """ - selected_files = _select_environment_files( + resolved_environment, files_selected = _resolve_environment_files( env_files=env_files, + base_environment=dict(os.environ), silent=silent, include_default_base=include_default_base, ) - for env_file in selected_files: - dotenv.load_dotenv(env_file, override=True, interpolate=True) - if not silent: - _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) - - return bool(selected_files) + os.environ.update(resolved_environment) + return files_selected def _resolve_environment_files( @@ -219,7 +219,10 @@ def _warn_about_akv_environment_files( messages: list[str] = [] if base_file.exists(): - messages.append(f"{base_file} exists and will be ignored because Key Vault supplies the base environment") + if env_files is None: + messages.append(f"{base_file} will load after Key Vault and override matching values") + else: + messages.append(f"{base_file} exists but will be ignored because env_files was explicitly configured") if local_file.exists(): if env_files is None: @@ -272,6 +275,40 @@ def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: return vault_url, secret_name, secret_version +def _create_akv_secret_client(*, vault_url: str, credential: Any) -> "SecretClient": + """ + Create an asynchronous Key Vault client with an explicit retry policy. + + Returns: + SecretClient: Configured asynchronous secret client. + """ + from azure.core.pipeline.policies import AsyncRetryPolicy + from azure.keyvault.secrets.aio import SecretClient + + retry_policy = AsyncRetryPolicy( + retry_total=_AKV_RETRY_TOTAL, + retry_connect=_AKV_RETRY_TOTAL, + retry_read=_AKV_RETRY_TOTAL, + retry_status=_AKV_RETRY_TOTAL, + retry_backoff_factor=_AKV_RETRY_BACKOFF_FACTOR, + ) + return SecretClient(vault_url=vault_url, credential=credential, retry_policy=retry_policy) + + +def _key_vault_initialization_error(*, message: str, error: Exception) -> KeyVaultInitializationException: + """ + Create a contextual Key Vault exception without losing the original cause. + + Returns: + KeyVaultInitializationException: Wrapped contextual exception. + """ + status_code = getattr(error, "status_code", None) + return KeyVaultInitializationException( + status_code=status_code if isinstance(status_code, int) else 500, + message=f"{message}: {error}", + ) + + def _validate_dotenv_document( document: str, *, @@ -324,12 +361,11 @@ async def _load_env_from_akv_async( secret_url: str, strict: bool = True, silent: bool = False, -) -> dict[str, str]: +) -> tuple[dict[str, str], str]: """ - Load environment variables from an Azure Key Vault secret. + Load a bootstrap environment document from an Azure Key Vault secret. - The secret URL identifies the bootstrap environment document. Values in - that document may directly reference scalar secrets in the same vault. + References remain unresolved until all environment sources are merged. Authentication uses ``DefaultAzureCredential``, which silently tries managed identity, Azure CLI, VS Code credentials, etc., and falls back to interactive @@ -343,46 +379,188 @@ async def _load_env_from_akv_async( silent (bool): If True, suppresses print statements. Defaults to False. Returns: - dict[str, str]: The fully resolved Key Vault environment mapping. + tuple[dict[str, str], str]: Parsed bootstrap values and the vault URL. Raises: ImportError: If ``azure-keyvault-secrets`` is not installed. - ValueError: If the root URL is malformed or the bootstrap environment + KeyVaultInitializationException: If the root URL is malformed or the bootstrap environment document cannot be fully resolved. + ValueError: Compatibility base of ``KeyVaultInitializationException``. """ from azure.identity.aio import DefaultAzureCredential - from azure.keyvault.secrets.aio import SecretClient - _print_msg(f"Loading environment from AKV secret: {secret_url}", quiet=silent, log=True) - vault_url, secret_name, secret_version = _parse_akv_secret_url(secret_url) - ambient_environment = dict(os.environ) - async with DefaultAzureCredential() as credential: - async with SecretClient(vault_url=vault_url, credential=credential) as client: - secret = await client.get_secret(secret_name, version=secret_version) - - if not secret.value: - raise ValueError(f"AKV environment secret has no value: {secret_url}") - - validated_document = _validate_dotenv_document(secret.value, strict=strict, silent=silent) - parsed_environment = dotenv.dotenv_values(stream=io.StringIO(validated_document), interpolate=True) - if not parsed_environment: - raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") - - resolved_secrets: dict[str, str] = {} - resolved_environment: dict[str, str] = {} - for variable_name, value in parsed_environment.items(): - if value is None: - continue - resolved_environment[variable_name] = await _resolve_environment_value_async( - value=value, + try: + _print_msg(f"Loading environment from AKV secret: {secret_url}", quiet=silent, log=True) + vault_url, secret_name, secret_version = _parse_akv_secret_url(secret_url) + async with DefaultAzureCredential() as credential: + async with _create_akv_secret_client(vault_url=vault_url, credential=credential) as client: + secret = await client.get_secret(secret_name, version=secret_version) + + if not secret.value: + raise ValueError(f"AKV environment secret has no value: {secret_url}") + + validated_document = _validate_dotenv_document(secret.value, strict=strict, silent=silent) + parsed_environment = dotenv.dotenv_values(stream=io.StringIO(validated_document), interpolate=True) + if not parsed_environment: + raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") + + return {name: value for name, value in parsed_environment.items() if value is not None}, vault_url + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to load Key Vault bootstrap secret '{secret_url}'", + error=error, + ) + raise wrapped_error from error + + +async def _resolve_environment_references_async( + *, + values: Mapping[str, str], + ambient_environment: Mapping[str, str], + bootstrap_vault_url: str | None = None, +) -> dict[str, str]: + """ + Resolve references after all environment sources have been merged. + + Args: + values (Mapping[str, str]): Winning values after source precedence. + ambient_environment (Mapping[str, str]): Process environment visible to ``env:`` references. + bootstrap_vault_url (str | None): Vault URL imposed by the bootstrap document. + + Returns: + dict[str, str]: Values with complete-value references resolved. + + Raises: + KeyVaultInitializationException: If a Key Vault reference is invalid or uses another vault. + ValueError: If an environment reference cannot be resolved. + """ + reference_environment = {**ambient_environment, **values} + vault_url = bootstrap_vault_url + reference_variable_name = "" + try: + for variable_name, value in values.items(): + reference_variable_name = variable_name + reference = _parse_environment_value_reference(value) + if reference is None or reference[0] != "akv": + continue + target = reference[1] + if not target.casefold().startswith("https://"): + _resolve_akv_secret_reference( + target=target, variable_name=variable_name, - secret_client=client, - vault_url=vault_url, - ambient_environment=ambient_environment, - resolved_secrets=resolved_secrets, + vault_url=vault_url or "", ) + referenced_vault_url, _, _ = _parse_akv_secret_url(target) + if vault_url is None: + vault_url = referenced_vault_url + _resolve_akv_secret_reference( + target=target, + variable_name=variable_name, + vault_url=vault_url, + ) + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Invalid Key Vault reference for environment variable '{reference_variable_name}'", + error=error, + ) + raise wrapped_error from error + + if vault_url is None: + return { + name: await _resolve_environment_value_async( + value=value, + variable_name=name, + secret_client=None, + vault_url=None, + reference_environment=reference_environment, + ) + for name, value in values.items() + } + + from azure.core.exceptions import AzureError + from azure.identity.aio import DefaultAzureCredential - return resolved_environment + try: + async with DefaultAzureCredential() as credential: + async with _create_akv_secret_client(vault_url=vault_url, credential=credential) as client: + return { + name: await _resolve_environment_value_async( + value=value, + variable_name=name, + secret_client=client, + vault_url=vault_url, + reference_environment=reference_environment, + ) + for name, value in values.items() + } + except KeyVaultInitializationException: + raise + except AzureError as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to connect to Key Vault '{vault_url}'", + error=error, + ) + raise wrapped_error from error + + +async def _prepare_environment_updates_async( + *, + env_akv_ref: str | None, + env_files: Sequence[pathlib.Path] | None, + env_akv_strict: bool, + silent: bool, +) -> dict[str, str]: + """ + Stage all environment sources and resolve references before committing. + + Args: + env_akv_ref (str | None): Optional Key Vault bootstrap secret URL. + env_files (Sequence[pathlib.Path] | None): Optional ordered local environment files. + env_akv_strict (bool): Whether bootstrap dotenv validation is strict. + silent (bool): Whether initialization messages are suppressed. + + Returns: + dict[str, str]: Fully resolved environment updates. + + Raises: + ValueError: If a configured source or reference is invalid. + """ + ambient_environment = dict(os.environ) + merged_values: dict[str, str] = {} + bootstrap_vault_url: str | None = None + + if env_akv_ref is not None: + if not env_akv_ref.strip(): + raise ValueError("env_akv_ref must be a non-empty Azure Key Vault secret URL.") + await asyncio.to_thread( + _warn_about_akv_environment_files, + env_files=env_files, + silent=silent, + ) + bootstrap_values, bootstrap_vault_url = await _load_env_from_akv_async( + secret_url=env_akv_ref, + strict=env_akv_strict, + silent=silent, + ) + merged_values.update(bootstrap_values) + + local_values, _ = await asyncio.to_thread( + _resolve_environment_files, + env_files=env_files, + base_environment={**ambient_environment, **merged_values}, + silent=silent, + ) + merged_values.update(local_values) + + return await _resolve_environment_references_async( + values=merged_values, + ambient_environment=ambient_environment, + bootstrap_vault_url=bootstrap_vault_url, + ) def _parse_environment_value_reference(value: str) -> tuple[str, str] | None: @@ -404,6 +582,31 @@ def _parse_environment_value_reference(value: str) -> tuple[str, str] | None: return None +def _lookup_environment_value(*, environment: Mapping[str, str], name: str) -> str | None: + """ + Look up an environment value using platform-appropriate name semantics. + + Returns: + str | None: The matched value, or None when no name matches. + """ + if name in environment: + return environment[name] + if os.name == "nt": + folded_name = name.casefold() + return next((value for key, value in environment.items() if key.casefold() == folded_name), None) + return None + + +def _environment_names_equal(*, left: str, right: str) -> bool: + """ + Compare environment variable names using platform semantics. + + Returns: + bool: True when the names identify the same environment variable. + """ + return left.casefold() == right.casefold() if os.name == "nt" else left == right + + def _validate_akv_secret_name(*, secret_name: str, variable_name: str) -> None: if not secret_name or len(secret_name) > 127 or any(not char.isalnum() and char != "-" for char in secret_name): raise ValueError( @@ -417,44 +620,45 @@ def _resolve_akv_secret_reference( target: str, variable_name: str, vault_url: str, -) -> tuple[str, str | None, str]: +) -> tuple[str, str | None]: """ - Resolve a same-vault secret name or full secret URI. + Resolve a full same-vault secret URI. Args: - target (str): A secret name or full Key Vault secret URI. + target (str): Full Key Vault secret URI. variable_name (str): The environment variable receiving the secret. vault_url (str): The bootstrap document's vault URL. Returns: - tuple[str, str | None, str]: Secret name, optional version, and cache key. + tuple[str, str | None]: Secret name and optional version. Raises: - ValueError: If the target is invalid or references another vault. + ValueError: If the target is not a full URI, is invalid, or references another vault. """ - secret_name = target - secret_version: str | None = None - if target.casefold().startswith("https://"): - referenced_vault_url, secret_name, secret_version = _parse_akv_secret_url(target) - if referenced_vault_url.rstrip("/").casefold() != vault_url.rstrip("/").casefold(): - raise ValueError( - f"Cross-vault AKV reference for environment variable '{variable_name}' is not supported. " - f"Expected vault '{vault_url}', got '{referenced_vault_url}'." - ) + if not target.casefold().startswith("https://"): + raise ValueError( + f"AKV reference for environment variable '{variable_name}' must use a full secret URL, " + "for example kv:https://my-vault.vault.azure.net/secrets/my-secret." + ) + + referenced_vault_url, secret_name, secret_version = _parse_akv_secret_url(target) + if referenced_vault_url.rstrip("/").casefold() != vault_url.rstrip("/").casefold(): + raise ValueError( + f"Cross-vault AKV reference for environment variable '{variable_name}' is not supported. " + f"Expected vault '{vault_url}', got '{referenced_vault_url}'." + ) _validate_akv_secret_name(secret_name=secret_name, variable_name=variable_name) - cache_key = f"{secret_name.casefold()}|{secret_version or ''}" - return secret_name, secret_version, cache_key + return secret_name, secret_version async def _resolve_environment_value_async( *, value: str, variable_name: str, - secret_client: "SecretClient", - vault_url: str, - ambient_environment: dict[str, str], - resolved_secrets: dict[str, str], + secret_client: "SecretClient | None", + vault_url: str | None, + reference_environment: Mapping[str, str], ) -> str: """ Resolve one value from the bootstrap environment document. @@ -462,15 +666,15 @@ async def _resolve_environment_value_async( Args: value (str): The parsed bootstrap value. variable_name (str): The environment variable receiving the resolved value. - secret_client (SecretClient): The client for the bootstrap document's vault. - vault_url (str): The bootstrap document's vault URL. - ambient_environment (dict[str, str]): Snapshot used for ``env:`` references. - resolved_secrets (dict[str, str]): Same-vault scalar cache keyed by secret name. + secret_client (SecretClient | None): Client for Key Vault references, when needed. + vault_url (str | None): The allowed Key Vault URL, when one is needed. + reference_environment (Mapping[str, str]): Merged source values with ambient fallback. Returns: str: The literal, ambient, or same-vault scalar value. Raises: + KeyVaultInitializationException: If a Key Vault reference cannot be resolved. ValueError: If a reference is empty or cannot resolve to a value. """ reference = _parse_environment_value_reference(value) @@ -484,28 +688,42 @@ async def _resolve_environment_value_async( raise ValueError(f"Empty {reference_type} reference for environment variable '{variable_name}'.") if reference_type == "env": - if target not in ambient_environment: + if _environment_names_equal(left=target, right=variable_name): + raise ValueError( + f"Environment variable '{variable_name}' cannot reference itself. Use a distinct source variable name." + ) + resolved_value = _lookup_environment_value(environment=reference_environment, name=target) + if resolved_value is None: raise ValueError( f"Environment variable '{target}' referenced by '{variable_name}' " - "is not set in the ambient environment." + "is not available in the merged environment." ) - return ambient_environment[target] + return resolved_value - secret_name, secret_version, secret_cache_key = _resolve_akv_secret_reference( - target=target, - variable_name=variable_name, - vault_url=vault_url, - ) - if secret_cache_key in resolved_secrets: - return resolved_secrets[secret_cache_key] + try: + if secret_client is None or vault_url is None: + raise ValueError(f"AKV reference for environment variable '{variable_name}' has no available vault client.") - secret = await secret_client.get_secret(secret_name, version=secret_version) - if secret.value is None: - raise ValueError( - f"AKV secret '{secret_name}' referenced by environment variable '{variable_name}' has no value." + secret_name, secret_version = _resolve_akv_secret_reference( + target=target, + variable_name=variable_name, + vault_url=vault_url, + ) + + secret = await secret_client.get_secret(secret_name, version=secret_version) + if secret.value is None: + raise ValueError( + f"AKV secret '{secret_name}' referenced by environment variable '{variable_name}' has no value." + ) + return secret.value + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to resolve Key Vault reference for environment variable '{variable_name}'", + error=error, ) - resolved_secrets[secret_cache_key] = secret.value - return secret.value + raise wrapped_error from error async def _execute_initializers_async(*, initializers: Sequence["PyRITInitializer"]) -> None: @@ -556,7 +774,7 @@ async def initialize_pyrit_async( initializers: Sequence["PyRITInitializer"] | None = None, load_defaults: bool = True, env_files: Sequence[pathlib.Path] | None = None, - env_akv_ref: Sequence[str] | None = None, + env_akv_ref: str | None = None, env_akv_strict: bool = True, silent: bool = False, **memory_instance_kwargs: Any, @@ -584,10 +802,9 @@ async def initialize_pyrit_async( env_files (Sequence[pathlib.Path] | None): Optional sequence of environment file paths to load in order. If not provided, will load default .env and .env.local files from PyRIT home if they exist. All paths must be valid pathlib.Path objects. - env_akv_ref (Sequence[str] | None): Optional sequence of Azure Key Vault secret URLs to load. - The first secret's value must contain the bootstrap .env document; additional URLs are ignored. - Loaded before ``env_files`` so local files take precedence over AKV. Requires - ``azure-keyvault-secrets``. + env_akv_ref (str | None): Optional Azure Key Vault URL whose secret value contains the + bootstrap .env document. Loaded before ``env_files`` so local files take precedence + over AKV. Requires ``azure-keyvault-secrets``. env_akv_strict (bool): If True, reject malformed or valueless entries in the Key Vault bootstrap document. If False, warn and skip those entries. Defaults to True. silent (bool): If True, suppresses print statements about environment file loading and @@ -597,44 +814,12 @@ async def initialize_pyrit_async( Raises: ValueError: If an unsupported memory_db_type is provided or env_files contains non-existent files. """ - base_environment = dict(os.environ) - environment_updates: dict[str, str] = {} - if env_akv_ref: - await asyncio.to_thread( - _warn_about_akv_environment_files, - env_files=env_files, - silent=silent, - ) - if len(env_akv_ref) > 1: - _print_msg( - "Multiple env_akv_ref values were provided; using the first as the root environment document.", - quiet=silent, - log=True, - ) - akv_environment = await _load_env_from_akv_async( - secret_url=env_akv_ref[0], - strict=env_akv_strict, - silent=silent, - ) - staged_environment = {**base_environment, **akv_environment} - local_environment, _ = await asyncio.to_thread( - _resolve_environment_files, - env_files=env_files, - base_environment=staged_environment, - silent=silent, - include_default_base=False, - ) - environment_updates.update(akv_environment) - environment_updates.update(local_environment) - else: - local_environment, _ = await asyncio.to_thread( - _resolve_environment_files, - env_files=env_files, - base_environment=base_environment, - silent=silent, - ) - environment_updates.update(local_environment) - + environment_updates = await _prepare_environment_updates_async( + env_akv_ref=env_akv_ref, + env_files=env_files, + env_akv_strict=env_akv_strict, + silent=silent, + ) os.environ.update(environment_updates) # Reset all default values before executing initialization scripts diff --git a/tests/unit/exceptions/test_exceptions.py b/tests/unit/exceptions/test_exceptions.py index e228efed32..ae30546cd7 100644 --- a/tests/unit/exceptions/test_exceptions.py +++ b/tests/unit/exceptions/test_exceptions.py @@ -14,6 +14,7 @@ BadRequestException, EmptyResponseException, InvalidJsonException, + KeyVaultInitializationException, MissingPromptPlaceholderException, PyritException, RateLimitException, @@ -59,6 +60,14 @@ def test_empty_response_exception_initialization(): assert str(ex) == "Status Code: 204, Message: No Content" +def test_key_vault_initialization_exception_is_value_error_compatible(): + ex = KeyVaultInitializationException(status_code=403, message="Key Vault access denied") + + assert isinstance(ex, ValueError) + assert ex.status_code == 403 + assert ex.message == "Key Vault access denied" + + def test_invalid_json_exception_initialization(): ex = InvalidJsonException() assert ex.status_code == 500 diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index e36d14ecf7..986c5a052b 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -147,7 +147,7 @@ def test_from_dict_with_all_fields(self): "initializers": ["simple"], "initialization_scripts": ["/path/to/script.py"], "env_files": ["/path/to/.env"], - "env_akv_ref": ["https://vault.vault.azure.net/secrets/one"], + "env_akv_ref": "https://vault.vault.azure.net/secrets/one", "env_akv_strict": False, "silent": True, } @@ -156,7 +156,7 @@ def test_from_dict_with_all_fields(self): assert config.initializers == ["simple"] assert config.initialization_scripts == ["/path/to/script.py"] assert config.env_files == ["/path/to/.env"] - assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] + assert config.env_akv_ref == "https://vault.vault.azure.net/secrets/one" assert config.env_akv_strict is False assert config.silent is True @@ -310,13 +310,15 @@ def testresolve_env_akv_ref_none_returns_none(self): assert config.resolve_env_akv_ref() is None def testresolve_env_akv_ref_returns_configured_values(self): - """Test that configured AKV references are returned unchanged.""" - refs = [ - "https://vault.vault.azure.net/secrets/first", - "https://vault.vault.azure.net/secrets/second/version", - ] - config = ConfigurationLoader(env_akv_ref=refs) - assert config.resolve_env_akv_ref() == refs + """Test that the configured AKV reference is returned unchanged.""" + ref = "https://vault.vault.azure.net/secrets/first" + config = ConfigurationLoader(env_akv_ref=ref) + assert config.resolve_env_akv_ref() == ref + + @pytest.mark.parametrize("env_akv_ref", [[], ["https://vault.vault.azure.net/secrets/one"], ""]) + def test_env_akv_ref_rejects_non_scalar_or_empty_values(self, env_akv_ref): + with pytest.raises(ValueError, match="env_akv_ref must be one non-empty"): + ConfigurationLoader(env_akv_ref=env_akv_ref) # type: ignore[arg-type] @pytest.mark.usefixtures("patch_central_database") @@ -343,17 +345,14 @@ async def test_initialize_pyrit_async_basic(self, mock_init): @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") async def test_initialize_pyrit_async_with_env_akv_ref(self, mock_init): """Test initialization forwards env_akv_ref to initialize_pyrit_async.""" - refs = [ - "https://vault.vault.azure.net/secrets/first", - "https://vault.vault.azure.net/secrets/second/version", - ] - config = ConfigurationLoader(memory_db_type="in_memory", env_akv_ref=refs, env_akv_strict=False) + ref = "https://vault.vault.azure.net/secrets/first" + config = ConfigurationLoader(memory_db_type="in_memory", env_akv_ref=ref, env_akv_strict=False) await config.initialize_pyrit_async() mock_init.assert_called_once() call_kwargs = mock_init.call_args.kwargs - assert call_kwargs["env_akv_ref"] == refs + assert call_kwargs["env_akv_ref"] == ref assert call_kwargs["env_akv_strict"] is False @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") @@ -466,13 +465,13 @@ def test_load_with_overrides_reads_env_akv_ref_from_default_config(self, mock_de mock_default_path.exists.return_value = True mock_from_yaml.return_value = ConfigurationLoader( memory_db_type="sqlite", - env_akv_ref=["https://default.vault.azure.net/secrets/from-default"], + env_akv_ref="https://default.vault.azure.net/secrets/from-default", ) config = ConfigurationLoader.load_with_overrides() mock_from_yaml.assert_called_once_with(mock_default_path) - assert config.env_akv_ref == ["https://default.vault.azure.net/secrets/from-default"] + assert config.env_akv_ref == "https://default.vault.azure.net/secrets/from-default" @mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH") def test_load_with_overrides_memory_db_type_override(self, mock_default_path): @@ -517,10 +516,10 @@ def test_load_with_overrides_env_akv_ref_override(self, mock_default_path): mock_default_path.exists.return_value = False config = ConfigurationLoader.load_with_overrides( - env_akv_ref=["https://vault.vault.azure.net/secrets/one"], + env_akv_ref="https://vault.vault.azure.net/secrets/one", ) - assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] + assert config.env_akv_ref == "https://vault.vault.azure.net/secrets/one" @mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH") def test_load_with_overrides_converts_sequence_to_list(self, mock_default_path): @@ -532,18 +531,15 @@ def test_load_with_overrides_converts_sequence_to_list(self, mock_default_path): initializers=("init1", "init2"), initialization_scripts=("/path/script1.py", "/path/script2.py"), env_files=("/path/.env",), - env_akv_ref=("https://vault.vault.azure.net/secrets/one",), ) # Verify they are stored as lists assert isinstance(config.initializers, list) assert isinstance(config.initialization_scripts, list) assert isinstance(config.env_files, list) - assert isinstance(config.env_akv_ref, list) assert config.initializers == ["init1", "init2"] assert config.initialization_scripts == ["/path/script1.py", "/path/script2.py"] assert config.env_files == ["/path/.env"] - assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] def test_load_with_overrides_explicit_config_file_not_found(self): """Test FileNotFoundError when explicit config file doesn't exist.""" @@ -566,8 +562,7 @@ def test_load_with_overrides_explicit_config_file_overrides_default(self, mock_d - /explicit/script.py env_files: - /explicit/.env -env_akv_ref: - - https://vault.vault.azure.net/secrets/explicit +env_akv_ref: https://vault.vault.azure.net/secrets/explicit """ with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: f.write(yaml_content) @@ -580,7 +575,7 @@ def test_load_with_overrides_explicit_config_file_overrides_default(self, mock_d assert config._initializer_configs[0].name == "explicit_init" assert config.initialization_scripts == ["/explicit/script.py"] assert config.env_files == ["/explicit/.env"] - assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/explicit"] + assert config.env_akv_ref == "https://vault.vault.azure.net/secrets/explicit" finally: config_path.unlink() diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 7ff8bc8bab..694dd3046b 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -8,9 +8,11 @@ from unittest import mock import pytest +from azure.core.exceptions import ResourceNotFoundError from pyrit.common.apply_defaults import reset_default_values from pyrit.common.singleton import Singleton +from pyrit.exceptions import KeyVaultInitializationException from pyrit.registry import InitializerRegistry from pyrit.setup import IN_MEMORY, initialize_pyrit_async from pyrit.setup.initialization import ( @@ -18,6 +20,8 @@ _load_environment_files, _parse_akv_secret_url, _parse_environment_value_reference, + _resolve_environment_files, + _resolve_environment_references_async, _warn_about_akv_environment_files, ) @@ -176,18 +180,15 @@ async def test_invalid_memory_type_raises_error(self, mock_resolve_env): @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_resolve_env, mock_set_memory): - """Test that env_akv_ref loads only its first entry.""" - refs = [ - "https://vault.vault.azure.net/secrets/test-secret", - "https://vault.vault.azure.net/secrets/ignored", - ] + """Test that env_akv_ref loads its bootstrap secret.""" + ref = "https://vault.vault.azure.net/secrets/test-secret" - mock_load_akv.return_value = {} + mock_load_akv.return_value = {}, "https://vault.vault.azure.net" - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=ref, load_defaults=False) mock_load_akv.assert_awaited_once() - assert mock_load_akv.await_args.kwargs["secret_url"] == refs[0] + assert mock_load_akv.await_args.kwargs["secret_url"] == ref assert mock_load_akv.await_args.kwargs["strict"] is True assert mock_load_akv.await_args.kwargs["silent"] is False mock_resolve_env.assert_called_once() @@ -196,52 +197,18 @@ async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_resolve_env @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) - async def test_initialize_with_empty_env_akv_ref_does_not_load_akv( - self, mock_load_akv, mock_resolve_env, mock_set_memory - ): - """Test that empty env_akv_ref does not invoke AKV loading.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=[], load_defaults=False) + async def test_initialize_with_empty_env_akv_ref_raises(self, mock_load_akv, mock_resolve_env, mock_set_memory): + """Test that an empty env_akv_ref is rejected.""" + with pytest.raises(ValueError, match="env_akv_ref must be a non-empty"): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref="", load_defaults=False) mock_load_akv.assert_not_called() - mock_resolve_env.assert_called_once() - mock_set_memory.assert_called_once() - - @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_loads_akv_before_env_files(self, mock_set_memory): - """Test that AKV refs are loaded before env_files so env_files can override values.""" - call_order: list[str] = [] - - def _record_warning(*, env_files, silent=False): - call_order.append("warning") - - async def _record_akv_call(*, secret_url, strict=True, silent=False): - call_order.append("akv") - return {"FROM_AKV": "shared"} - - def _record_env_file_call(*, env_files, base_environment, silent=False, include_default_base=True): - call_order.append("env_files") - assert include_default_base is False - assert base_environment["FROM_AKV"] == "shared" - return {"FROM_LOCAL": "override"}, True - - refs = ["https://vault.vault.azure.net/secrets/test-secret"] - - with mock.patch.dict(os.environ, {}, clear=True): - with ( - mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files", side_effect=_record_warning), - mock.patch("pyrit.setup.initialization._load_env_from_akv_async", side_effect=_record_akv_call), - mock.patch("pyrit.setup.initialization._resolve_environment_files", side_effect=_record_env_file_call), - ): - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) - - assert os.environ == {"FROM_AKV": "shared", "FROM_LOCAL": "override"} - - assert call_order == ["warning", "akv", "env_files"] - mock_set_memory.assert_called_once() + mock_resolve_env.assert_not_called() + mock_set_memory.assert_not_called() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") async def test_initialize_akv_and_local_files_are_applied_atomically(self, mock_set_memory): - refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + ref = "https://vault.vault.azure.net/secrets/bootstrap" nonexistent = pathlib.Path("/nonexistent/.env") with mock.patch.dict(os.environ, {}, clear=True): @@ -250,13 +217,13 @@ async def test_initialize_akv_and_local_files_are_applied_atomically(self, mock_ mock.patch( "pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock, - return_value={"FROM_AKV": "resolved"}, + return_value=({"FROM_AKV": "resolved"}, "https://vault.vault.azure.net"), ), pytest.raises(ValueError, match="Environment file not found"), ): await initialize_pyrit_async( memory_db_type=IN_MEMORY, - env_akv_ref=refs, + env_akv_ref=ref, env_files=[nonexistent], load_defaults=False, ) @@ -267,7 +234,7 @@ async def test_initialize_akv_and_local_files_are_applied_atomically(self, mock_ @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") async def test_initialize_stages_local_overrides_on_akv_environment(self, mock_set_memory): - refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + ref = "https://vault.vault.azure.net/secrets/bootstrap" with tempfile.TemporaryDirectory() as temp_dir: local_file = pathlib.Path(temp_dir) / ".env.local" local_file.write_text("DERIVED=${BASE}\nBASE=local") @@ -278,12 +245,12 @@ async def test_initialize_stages_local_overrides_on_akv_environment(self, mock_s mock.patch( "pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock, - return_value={"BASE": "akv", "ONLY_AKV": "shared"}, + return_value=({"BASE": "akv", "ONLY_AKV": "shared"}, "https://vault.vault.azure.net"), ), ): await initialize_pyrit_async( memory_db_type=IN_MEMORY, - env_akv_ref=refs, + env_akv_ref=ref, env_files=[local_file], load_defaults=False, ) @@ -294,6 +261,87 @@ async def test_initialize_stages_local_overrides_on_akv_environment(self, mock_s mock_set_memory.assert_called_once() + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_default_files_override_akv_in_order(self, mock_set_memory): + ref = "https://vault.vault.azure.net/secrets/bootstrap" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VALUE=env") + (temp_path / ".env.local").write_text("VALUE=local") + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + return_value=({"VALUE": "akv"}, "https://vault.vault.azure.net"), + ), + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=ref, + load_defaults=False, + silent=True, + ) + + assert os.environ["VALUE"] == "local" + + mock_set_memory.assert_called_once() + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_resolves_only_winning_references_after_local_override(self, mock_set_memory): + ref = "https://vault.vault.azure.net/secrets/bootstrap" + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value="bootstrap-secret-value"), + types.SimpleNamespace(value="local-secret-value"), + ] + ) + with tempfile.TemporaryDirectory() as temp_dir: + local_file = pathlib.Path(temp_dir) / ".env.local" + local_file.write_text( + "OVERRIDDEN=local\n" + "LOCAL_SECRET=kv:https://vault.vault.azure.net/secrets/local-secret\n" + "LOCAL_ENV=env:BOOTSTRAP_SOURCE" + ) + bootstrap_environment = { + "OVERRIDDEN": "kv:https://vault.vault.azure.net/secrets/unused-secret", + "BOOTSTRAP_SECRET": "kv:https://vault.vault.azure.net/secrets/bootstrap-secret", + "BOOTSTRAP_SOURCE": "bootstrap-value", + } + + with ( + mock.patch.dict(os.environ, {"SOURCE_VALUE": "ambient-value"}, clear=True), + mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + return_value=(bootstrap_environment, "https://vault.vault.azure.net"), + ), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=ref, + env_files=[local_file], + load_defaults=False, + ) + + assert os.environ["OVERRIDDEN"] == "local" + assert os.environ["BOOTSTRAP_SECRET"] == "bootstrap-secret-value" + assert os.environ["LOCAL_SECRET"] == "local-secret-value" + assert os.environ["LOCAL_ENV"] == "bootstrap-value" + + assert client.get_secret.await_args_list == [ + mock.call("bootstrap-secret", version=None), + mock.call("local-secret", version=None), + ] + mock_set_memory.assert_called_once() + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") async def test_initialize_without_environment_file_uses_system_environment(self, mock_set_memory): await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[], load_defaults=False) @@ -341,56 +389,41 @@ async def test_initialize_not_silent_prints_migration_message(self, mock_load_en class TestLoadEnvironmentFiles: """Tests for _load_environment_files function and env_files parameter in initialize_pyrit_async.""" - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_loads_default_env_files_when_none_provided(self, mock_config_path, mock_load_dotenv): + async def test_loads_default_env_files_when_none_provided(self, mock_config_path): """Test that default .env and .env.local files are loaded when env_files is None.""" - # Create temporary directory and files with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) env_file = temp_path / ".env" env_local_file = temp_path / ".env.local" - - # Create the files env_file.write_text("VAR1=value1") env_local_file.write_text("VAR2=value2") - - # Mock CONFIGURATION_DIRECTORY_PATH to point to our temp directory mock_config_path.__truediv__ = lambda self, other: temp_path / other - # Call the function with None (default behavior) - loaded = _load_environment_files(env_files=None) + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None) - # Verify both files were loaded - assert loaded is True - assert mock_load_dotenv.call_count == 2 - calls = [call[0][0] for call in mock_load_dotenv.call_args_list] - assert env_file in calls - assert env_local_file in calls + assert loaded is True + assert os.environ["VAR1"] == "value1" + assert os.environ["VAR2"] == "value2" - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_only_loads_existing_default_files(self, mock_config_path, mock_load_dotenv): + async def test_only_loads_existing_default_files(self, mock_config_path): """Test that only existing default files are loaded.""" with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) env_file = temp_path / ".env" - - # Only create .env, not .env.local env_file.write_text("VAR1=value1") - mock_config_path.__truediv__ = lambda self, other: temp_path / other - loaded = _load_environment_files(env_files=None) + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None) - # Verify only one file was loaded - assert loaded is True - assert mock_load_dotenv.call_count == 1 - assert mock_load_dotenv.call_args[0][0] == env_file + assert loaded is True + assert os.environ["VAR1"] == "value1" - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_excludes_default_env_when_loading_local_override(self, mock_config_path, mock_load_dotenv): + async def test_excludes_default_env_when_loading_local_override(self, mock_config_path): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) env_file = temp_path / ".env" @@ -400,23 +433,23 @@ async def test_excludes_default_env_when_loading_local_override(self, mock_confi mock_config_path.__truediv__ = lambda self, other: temp_path / other - loaded = _load_environment_files(env_files=None, include_default_base=False) + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None, include_default_base=False) - assert loaded is True - mock_load_dotenv.assert_called_once() - assert mock_load_dotenv.call_args.args[0] == env_local_file + assert loaded is True + assert os.environ["VAR"] == "local" - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_returns_false_when_no_default_files_exist(self, mock_config_path, mock_load_dotenv): + async def test_returns_false_when_no_default_files_exist(self, mock_config_path): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) mock_config_path.__truediv__ = lambda self, other: temp_path / other - loaded = _load_environment_files(env_files=None) + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None) - assert loaded is False - mock_load_dotenv.assert_not_called() + assert loaded is False + assert os.environ == {} @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") def test_warns_when_default_files_coexist_with_akv(self, mock_config_path, caplog, capsys): @@ -433,13 +466,32 @@ def test_warns_when_default_files_coexist_with_akv(self, mock_config_path, caplo output = capsys.readouterr().out assert output.startswith("WARNING: env_akv_ref is configured") - assert f"{env_file} exists and will be ignored" in output + assert f"{env_file} will load after Key Vault and override matching values" in output assert f"{env_local_file} will load after Key Vault and override matching values" in output assert "clear or remove ~/.pyrit/.env and ~/.pyrit/.env.local" in output assert "remove explicit env_files when Key Vault should be the only source" in output assert "restart PyRIT" in output assert caplog.records[0].levelname == "WARNING" + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_warns_when_explicit_files_replace_defaults_with_akv(self, mock_config_path, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + custom_file = temp_path / ".env.custom" + env_file.write_text("VAR=base") + env_local_file.write_text("VAR=local") + custom_file.write_text("VAR=custom") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + _warn_about_akv_environment_files(env_files=[custom_file]) + + output = capsys.readouterr().out + assert f"{env_file} exists but will be ignored because env_files was explicitly configured" in output + assert f"{env_local_file} exists but will be ignored because env_files was explicitly configured" in output + assert f"explicit env_files will load after Key Vault and override matching values: {[custom_file]}" in output + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") def test_akv_environment_file_warning_respects_silent(self, mock_config_path, caplog, capsys): with tempfile.TemporaryDirectory() as temp_dir: @@ -451,11 +503,10 @@ def test_akv_environment_file_warning_respects_silent(self, mock_config_path, ca _warn_about_akv_environment_files(env_files=None, silent=True) assert capsys.readouterr().out == "" - assert "will be ignored because Key Vault supplies the base environment" in caplog.text + assert "will load after Key Vault and override matching values" in caplog.text assert "restart PyRIT" in caplog.text - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") - async def test_loads_custom_env_files_in_order(self, mock_load_dotenv): + async def test_loads_custom_env_files_in_order(self): """Test that custom env_files are loaded in the order provided.""" with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -468,15 +519,24 @@ async def test_loads_custom_env_files_in_order(self, mock_load_dotenv): env2.write_text("VAR=prod") env3.write_text("VAR=local") - # Pass custom files - _load_environment_files(env_files=[env1, env2, env3]) + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env1, env2, env3]) + + assert loaded is True + assert os.environ["VAR"] == "local" + + async def test_load_environment_files_honors_python_dotenv_disabled(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("DISABLED_VALUE=not-loaded") - # Verify all three files were loaded in order - assert mock_load_dotenv.call_count == 3 - call_args = [call[0][0] for call in mock_load_dotenv.call_args_list] - assert call_args == [env1, env2, env3] + with mock.patch.dict(os.environ, {"PYTHON_DOTENV_DISABLED": "true"}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) - async def test_local_environment_files_keep_pyrit_references_literal(self): + assert loaded is True + assert "DISABLED_VALUE" not in os.environ + + async def test_direct_local_file_loader_keeps_pyrit_references_literal(self): with tempfile.TemporaryDirectory() as temp_dir: env_file = pathlib.Path(temp_dir) / ".env" env_file.write_text( @@ -491,6 +551,29 @@ async def test_local_environment_files_keep_pyrit_references_literal(self): assert os.environ["ENV_REFERENCE"] == "env:SOURCE_VALUE" assert os.environ["INTERPOLATED"] == "base" + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_default_local_file_can_interpolate_base_file_but_not_reverse(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text( + "OPENAI_CHAT_ENDPOINT=https://example.openai.azure.com/openai/v1\nFROM_LATER_LOCAL=${LOCAL_ONLY}" + ) + env_local_file.write_text("FOOBAR=${OPENAI_CHAT_ENDPOINT}\nLOCAL_ONLY=local") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + resolved, loaded = _resolve_environment_files( + env_files=None, + base_environment={}, + silent=True, + ) + + assert loaded is True + assert resolved["FOOBAR"] == "https://example.openai.azure.com/openai/v1" + assert resolved["FROM_LATER_LOCAL"] == "" + assert resolved["LOCAL_ONLY"] == "local" + async def test_env_akv_strict_does_not_validate_local_environment_files(self): with tempfile.TemporaryDirectory() as temp_dir: env_file = pathlib.Path(temp_dir) / ".env" @@ -508,6 +591,36 @@ async def test_env_akv_strict_does_not_validate_local_environment_files(self): assert os.environ["GOOD"] == "resolved" assert os.environ["OTHER"] == "also-resolved" + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_local_file_resolves_full_akv_reference_without_bootstrap(self, mock_set_memory): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="local-secret-value")) + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("API_KEY=kv:https://myvault.vault.azure.net/secrets/api-key") + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_files=[env_file], + load_defaults=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "local-secret-value" + + _assert_mock_akv_client_created( + mock_client_cls, + vault_url="https://myvault.vault.azure.net", + credential=credential, + ) + client.get_secret.assert_awaited_once_with("api-key", version=None) + mock_set_memory.assert_called_once() + async def test_raises_error_for_nonexistent_env_file(self): """Test that ValueError is raised for non-existent env file.""" nonexistent = pathlib.Path("/nonexistent/path/.env") @@ -570,12 +683,32 @@ def _create_mock_akv_clients() -> tuple[mock.MagicMock, mock.MagicMock]: return credential, client +def _assert_mock_akv_client_created( + mock_client_cls: mock.MagicMock, + *, + vault_url: str, + credential: mock.MagicMock, +) -> None: + mock_client_cls.assert_called_once() + call_kwargs = mock_client_cls.call_args.kwargs + assert call_kwargs["vault_url"] == vault_url + assert call_kwargs["credential"] is credential + retry_policy = call_kwargs["retry_policy"] + assert retry_policy.total_retries == 3 + assert retry_policy.connect_retries == 3 + assert retry_policy.read_retries == 3 + assert retry_policy.status_retries == 3 + assert retry_policy.backoff_factor == 0.8 + + class TestAkvEnvironmentLoading: """Tests for AKV URL parsing and env loading helpers.""" @pytest.mark.parametrize("prefix", ["kv", "akv", "azure_key_vault", "env_akv_ref"]) def test_parse_environment_value_reference_accepts_akv_aliases(self, prefix): - assert _parse_environment_value_reference(f"{prefix}:api-key") == ("akv", "api-key") + secret_url = "https://myvault.vault.azure.net/secrets/api-key" + + assert _parse_environment_value_reference(f"{prefix}:{secret_url}") == ("akv", secret_url) def test_parse_environment_value_reference_rejects_azure_app_service_syntax(self): value = "@Microsoft.KeyVault(SecretUri=https://myvault.vault.azure.net/secrets/api-key)" @@ -604,78 +737,144 @@ def test_parse_akv_secret_url_invalid_raises(self): with pytest.raises(ValueError, match="Invalid AKV secret URL"): _parse_akv_secret_url("https://myvault.vault.azure.net/not-secrets/my-secret") - async def test_load_env_from_akv_async_resolves_one_level(self): + async def test_load_env_from_akv_async_returns_unresolved_bootstrap(self): credential, client = _create_mock_akv_clients() root_document = ( "DIRECT=from-bootstrap\n" "FROM_ENV=env:SOURCE_VALUE\n" - "FROM_KV=kv:api-key\n" - "DUPLICATE_KV=akv:https://MYVAULT.vault.azure.net/secrets/API-KEY\n" + "FROM_KV=kv:https://myvault.vault.azure.net/secrets/api-key\n" "PINNED_KV=kv:https://myvault.vault.azure.net/secrets/api-key/version-2\n" "ESCAPED=literal:kv:not-a-secret" ) - client.get_secret = mock.AsyncMock( - side_effect=[ - types.SimpleNamespace(value=root_document), - types.SimpleNamespace(value="env:not-resolved-again"), - types.SimpleNamespace(value="pinned-secret-value"), - ] - ) + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=root_document)) secret_url = "https://myvault.vault.azure.net/secrets/bootstrap/v1" with ( - mock.patch.dict(os.environ, {"SOURCE_VALUE": "kv:not-fetched"}, clear=False), mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential) as mock_credential_cls, mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, ): - resolved_environment = await _load_env_from_akv_async(secret_url=secret_url, silent=True) + parsed_environment, vault_url = await _load_env_from_akv_async(secret_url=secret_url, silent=True) - assert resolved_environment == { + assert parsed_environment == { "DIRECT": "from-bootstrap", - "FROM_ENV": "kv:not-fetched", - "FROM_KV": "env:not-resolved-again", - "DUPLICATE_KV": "env:not-resolved-again", - "PINNED_KV": "pinned-secret-value", - "ESCAPED": "kv:not-a-secret", + "FROM_ENV": "env:SOURCE_VALUE", + "FROM_KV": "kv:https://myvault.vault.azure.net/secrets/api-key", + "PINNED_KV": "kv:https://myvault.vault.azure.net/secrets/api-key/version-2", + "ESCAPED": "literal:kv:not-a-secret", } - assert "DIRECT" not in os.environ + assert vault_url == "https://myvault.vault.azure.net" mock_credential_cls.assert_called_once_with() - mock_client_cls.assert_called_once_with(vault_url="https://myvault.vault.azure.net", credential=credential) - assert client.get_secret.await_args_list == [ - mock.call("bootstrap", version="v1"), - mock.call("api-key", version=None), - mock.call("api-key", version="version-2"), - ] + _assert_mock_akv_client_created( + mock_client_cls, + vault_url="https://myvault.vault.azure.net", + credential=credential, + ) + client.get_secret.assert_awaited_once_with("bootstrap", version="v1") credential.__aenter__.assert_awaited_once() credential.__aexit__.assert_awaited_once() client.__aenter__.assert_awaited_once() client.__aexit__.assert_awaited_once() mock_print_msg.assert_called_once() - async def test_load_env_from_akv_async_rejects_cross_vault_reference(self): + async def test_resolve_environment_references_async_resolves_local_values(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock( - return_value=types.SimpleNamespace( - value="API_KEY=kv:https://other-vault.vault.azure.net/secrets/api-key/version-1" + side_effect=[ + types.SimpleNamespace(value="local-secret-value"), + types.SimpleNamespace(value="pinned-secret-value"), + ] + ) + values = { + "DIRECT": "from-local", + "FROM_ENV": "env:SOURCE_VALUE", + "DECLARED": "merged-value", + "FROM_DECLARED": "env:DECLARED", + "SHADOWED": "merged-wins", + "FROM_SHADOWED": "env:SHADOWED", + "FROM_KV": "kv:https://myvault.vault.azure.net/secrets/api-key", + "PINNED_KV": "akv:https://myvault.vault.azure.net/secrets/api-key/version-2", + "ESCAPED": "literal:kv:not-a-secret", + } + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, + ): + resolved = await _resolve_environment_references_async( + values=values, + ambient_environment={"SOURCE_VALUE": "ambient-value", "SHADOWED": "ambient-loses"}, ) + + assert resolved == { + "DIRECT": "from-local", + "FROM_ENV": "ambient-value", + "DECLARED": "merged-value", + "FROM_DECLARED": "merged-value", + "SHADOWED": "merged-wins", + "FROM_SHADOWED": "merged-wins", + "FROM_KV": "local-secret-value", + "PINNED_KV": "pinned-secret-value", + "ESCAPED": "kv:not-a-secret", + } + _assert_mock_akv_client_created( + mock_client_cls, + vault_url="https://myvault.vault.azure.net", + credential=credential, ) + assert client.get_secret.await_args_list == [ + mock.call("api-key", version=None), + mock.call("api-key", version="version-2"), + ] - with mock.patch.dict(os.environ, {}, clear=True): - with ( - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - pytest.raises(ValueError, match="Cross-vault AKV reference"), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - ) + async def test_resolve_environment_references_async_rejects_self_reference(self): + with pytest.raises(ValueError, match="cannot reference itself"): + await _resolve_environment_references_async( + values={"MODEL": "env:MODEL"}, + ambient_environment={"MODEL": "ambient-model"}, + ) + + async def test_resolve_environment_references_async_preserves_windows_case_insensitive_lookup(self): + with mock.patch("pyrit.setup.initialization.os.name", "nt"): + resolved = await _resolve_environment_references_async( + values={"ALIAS": "env:Path"}, + ambient_environment={"PATH": "windows-path"}, + ) - assert "API_KEY" not in os.environ + assert resolved["ALIAS"] == "windows-path" - client.get_secret.assert_awaited_once_with("bootstrap", version=None) + async def test_resolve_environment_references_async_rejects_windows_case_variant_self_reference(self): + with ( + mock.patch("pyrit.setup.initialization.os.name", "nt"), + pytest.raises(ValueError, match="cannot reference itself"), + ): + await _resolve_environment_references_async( + values={"MODEL": "env:model"}, + ambient_environment={}, + ) + + async def test_resolve_environment_references_async_rejects_short_secret_name(self): + with pytest.raises(ValueError, match="must use a full secret URL"): + await _resolve_environment_references_async( + values={"API_KEY": "kv:api-key"}, + ambient_environment={}, + ) + + @pytest.mark.parametrize( + "reference_url", + [ + "https://other-vault.vault.azure.net/secrets/api-key", + "https://other-vault.vault.azure.net/secrets/api-key/version-1", + ], + ) + async def test_resolve_environment_references_async_rejects_cross_vault_reference(self, reference_url): + with pytest.raises(ValueError, match="Cross-vault AKV reference"): + await _resolve_environment_references_async( + values={"API_KEY": f"kv:{reference_url}"}, + ambient_environment={}, + bootstrap_vault_url="https://myvault.vault.azure.net", + ) async def test_load_env_from_akv_async_empty_secret_raises(self): credential, client = _create_mock_akv_clients() @@ -736,6 +935,39 @@ async def test_load_env_from_akv_async_rejects_non_assignments(self, document, e assert "GOOD" not in os.environ assert "OTHER" not in os.environ + async def test_load_env_from_akv_async_wraps_malformed_bootstrap(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="=malformed")) + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(KeyVaultInitializationException, match="malformed entries") as exc_info, + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert isinstance(exc_info.value.__cause__, ValueError) + + async def test_resolve_environment_references_async_wraps_missing_secret(self): + credential, client = _create_mock_akv_clients() + missing_error = ResourceNotFoundError(message="Secret was not found") + client.get_secret = mock.AsyncMock(side_effect=missing_error) + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(KeyVaultInitializationException, match="Failed to resolve Key Vault reference") as exc_info, + ): + await _resolve_environment_references_async( + values={"API_KEY": "kv:https://myvault.vault.azure.net/secrets/missing"}, + ambient_environment={}, + ) + + assert exc_info.value.__cause__ is missing_error + async def test_load_env_from_akv_async_allows_empty_assignment(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="EMPTY=")) @@ -745,7 +977,7 @@ async def test_load_env_from_akv_async_allows_empty_assignment(self): mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), ): - resolved_environment = await _load_env_from_akv_async( + resolved_environment, _ = await _load_env_from_akv_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", silent=True, ) @@ -753,6 +985,22 @@ async def test_load_env_from_akv_async_allows_empty_assignment(self): assert resolved_environment["EMPTY"] == "" assert "EMPTY" not in os.environ + async def test_resolve_environment_references_async_allows_empty_child_secret(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="")) + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + resolved_environment = await _resolve_environment_references_async( + values={"EMPTY": "kv:https://myvault.vault.azure.net/secrets/empty-secret"}, + ambient_environment={}, + ) + + assert resolved_environment["EMPTY"] == "" + client.get_secret.assert_awaited_once_with("empty-secret", version=None) + async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entries(self, caplog, capsys): credential, client = _create_mock_akv_clients() document = "GOOD=resolved\n=malformed\nMISSING_VALUE\nOTHER=also-resolved" @@ -764,7 +1012,7 @@ async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entrie mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), caplog.at_level("WARNING", logger="pyrit.setup.initialization"), ): - resolved_environment = await _load_env_from_akv_async( + resolved_environment, _ = await _load_env_from_akv_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", strict=False, silent=False, @@ -799,14 +1047,9 @@ async def test_load_env_from_akv_async_non_strict_silent_logs_warning(self, capl assert capsys.readouterr().out == "" assert "variables without values: MISSING_VALUE" in caplog.text - async def test_load_env_from_akv_async_failure_does_not_partially_update_environment(self): + async def test_resolve_environment_references_async_failure_returns_no_partial_mapping(self): credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock( - side_effect=[ - types.SimpleNamespace(value="GOOD=resolved\nBAD=kv:missing-value"), - types.SimpleNamespace(value=None), - ] - ) + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=None)) with mock.patch.dict(os.environ, {}, clear=True): with ( @@ -814,9 +1057,12 @@ async def test_load_env_from_akv_async_failure_does_not_partially_update_environ mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), pytest.raises(ValueError, match="has no value"), ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, + await _resolve_environment_references_async( + values={ + "GOOD": "resolved", + "BAD": "kv:https://myvault.vault.azure.net/secrets/missing-value", + }, + ambient_environment={}, ) assert "GOOD" not in os.environ From 8ee9c631b1a080a0001232c21326d16355b3553f Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 13 Aug 2026 12:53:35 -0400 Subject: [PATCH 07/16] FEAT Simplification refactor --- pyrit/setup/configuration_loader.py | 25 +- pyrit/setup/initialization.py | 463 +++++------------- tests/unit/setup/test_configuration_loader.py | 51 +- tests/unit/setup/test_initialization.py | 449 ++++++++--------- 4 files changed, 405 insertions(+), 583 deletions(-) diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 0201f3fe82..ecd11f0344 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -96,6 +96,7 @@ class ConfigurationLoader(YamlLoadable): None means "use defaults", [] means "load nothing". env_files: List of environment file paths to load. None means "use defaults (.env, .env.local)", [] means "load nothing". + env_akv_ref: Ordered list of Key Vault bootstrap secret URLs. env_akv_strict: Whether malformed or valueless entries in a Key Vault bootstrap document should fail initialization. silent: Whether to suppress initialization messages. @@ -136,7 +137,7 @@ class ConfigurationLoader(YamlLoadable): initializers: list[str | dict[str, Any]] = field(default_factory=list) initialization_scripts: list[str] | None = None env_files: list[str] | None = None - env_akv_ref: str | None = None + env_akv_ref: list[str] | None = None env_akv_strict: bool = True silent: bool = False operator: str | None = None @@ -158,12 +159,14 @@ def _validate_env_akv_ref(self) -> None: Validate the Key Vault bootstrap secret reference. Raises: - ValueError: If env_akv_ref is not one non-empty string. + ValueError: If env_akv_ref is not a list of non-empty strings. """ if self.env_akv_ref is None: return - if not isinstance(self.env_akv_ref, str) or not self.env_akv_ref.strip(): - raise ValueError("env_akv_ref must be one non-empty Azure Key Vault secret URL.") + if not isinstance(self.env_akv_ref, list): + raise ValueError("env_akv_ref must be a list of Azure Key Vault secret URLs.") + if any(not isinstance(secret_url, str) or not secret_url.strip() for secret_url in self.env_akv_ref): + raise ValueError("env_akv_ref must contain only non-empty Azure Key Vault secret URLs.") def _normalize_memory_db_type(self) -> None: """ @@ -416,7 +419,7 @@ def load_with_overrides( initializers: Sequence[str | dict[str, Any]] | None = None, initialization_scripts: Sequence[str] | None = None, env_files: Sequence[str] | None = None, - env_akv_ref: str | None = None, + env_akv_ref: Sequence[str] | None = None, env_akv_strict: bool | None = None, ) -> "ConfigurationLoader": """ @@ -433,7 +436,7 @@ def load_with_overrides( initializers: Override for initializer list. initialization_scripts: Override for initialization script paths. env_files: Override for environment file paths. - env_akv_ref: Override for the Azure Key Vault bootstrap secret URL. + env_akv_ref: Override for the ordered Azure Key Vault bootstrap secret URLs. env_akv_strict: Override for strict Key Vault bootstrap validation. Returns: @@ -495,7 +498,9 @@ def to_init_data(config: ConfigurationLoader) -> dict[str, Any]: config_data["env_files"] = list(env_files) if env_akv_ref is not None: - config_data["env_akv_ref"] = env_akv_ref + if isinstance(env_akv_ref, str): + raise ValueError("env_akv_ref must be a sequence of Azure Key Vault secret URLs.") + config_data["env_akv_ref"] = list(env_akv_ref) if env_akv_strict is not None: config_data["env_akv_strict"] = env_akv_strict @@ -601,12 +606,12 @@ def resolve_env_files(self) -> Sequence[pathlib.Path] | None: return resolved - def resolve_env_akv_ref(self) -> str | None: + def resolve_env_akv_ref(self) -> list[str] | None: """ - Return the AKV bootstrap secret URL, or ``None`` when not configured. + Return the AKV bootstrap secret URLs, or ``None`` when not configured. Returns: - str | None: The configured AKV bootstrap secret URL, or ``None``. + list[str] | None: The configured AKV bootstrap secret URLs, or ``None``. """ return self.env_akv_ref diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index b192548430..19ab1171e1 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -5,12 +5,12 @@ import logging import os import pathlib -from collections.abc import Mapping, Sequence +import urllib.parse +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Literal, get_args import dotenv from dotenv.parser import parse_stream -from dotenv.variables import parse_variables from pyrit.common import path from pyrit.common.apply_defaults import reset_default_values @@ -30,6 +30,7 @@ MemoryDatabaseType = Literal["InMemory", "SQLite", "AzureSQL"] _AKV_REFERENCE_PREFIXES = frozenset({"akv", "kv", "azure_key_vault", "env_akv_ref"}) +_AKV_VAULT_DNS_SUFFIXES = frozenset({"vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"}) _AKV_RETRY_TOTAL = 3 _AKV_RETRY_BACKOFF_FACTOR = 0.8 @@ -58,59 +59,17 @@ def _load_environment_files( Raises: ValueError: If any provided env_files do not exist. """ - resolved_environment, files_selected = _resolve_environment_files( - env_files=env_files, - base_environment=dict(os.environ), - silent=silent, - include_default_base=include_default_base, - ) - os.environ.update(resolved_environment) - return files_selected - - -def _resolve_environment_files( - env_files: Sequence[pathlib.Path] | None, - *, - base_environment: Mapping[str, str], - silent: bool = False, - include_default_base: bool = True, -) -> tuple[dict[str, str], bool]: - """ - Resolve environment files without mutating ``os.environ``. - - Args: - env_files: Optional sequence of environment file paths. If None, resolves - default files from the PyRIT configuration directory. - base_environment: Environment visible to interpolation before file values. - silent: If True, suppresses loading messages. Defaults to False. - include_default_base: If False and env_files is None, skips the default - .env file while still resolving .env.local. Defaults to True. - - Returns: - tuple[dict[str, str], bool]: Resolved values and whether any file was selected. - - Raises: - ValueError: If any explicitly provided environment file does not exist. - """ selected_files = _select_environment_files( env_files=env_files, silent=silent, include_default_base=include_default_base, ) - if _dotenv_loading_disabled(): - return {}, bool(selected_files) - - staged_environment = dict(base_environment) - resolved_environment: dict[str, str] = {} for env_file in selected_files: - raw_values = dotenv.dotenv_values(dotenv_path=env_file, interpolate=False) - file_values = _interpolate_dotenv_values(values=raw_values, base_environment=staged_environment) - staged_environment.update(file_values) - resolved_environment.update(file_values) + dotenv.load_dotenv(dotenv_path=env_file, override=True, interpolate=True) if not silent: _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) - return resolved_environment, bool(selected_files) + return bool(selected_files) def _select_environment_files( @@ -163,36 +122,6 @@ def _select_environment_files( return list(env_files) -def _interpolate_dotenv_values( - *, - values: Mapping[str, str | None], - base_environment: Mapping[str, str], -) -> dict[str, str]: - """ - Resolve dotenv interpolation against a staged environment mapping. - - Returns: - dict[str, str]: Interpolated assignments, excluding valueless entries. - """ - visible_environment: dict[str, str | None] = dict(base_environment) - resolved_values: dict[str, str] = {} - for name, value in values.items(): - if value is None: - visible_environment[name] = None - continue - - resolved_value = "".join(atom.resolve(visible_environment) for atom in parse_variables(value)) - visible_environment[name] = resolved_value - resolved_values[name] = resolved_value - - return resolved_values - - -def _dotenv_loading_disabled() -> bool: - value = os.environ.get("PYTHON_DOTENV_DISABLED", "") - return value.casefold() in {"1", "true", "t", "yes", "y"} - - def _print_msg(message: str, quiet: bool, log: bool) -> None: """ Print a standard initialization message unless quiet is True. @@ -262,17 +191,61 @@ def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: Raises: ValueError: If the URL does not match the expected format. """ - parts = secret_url.split("/secrets/") - if len(parts) != 2: - raise ValueError( - f"Invalid AKV secret URL: '{secret_url}'. " - "Expected format: https://{{vault}}.vault.azure.net/secrets/{{name}}[/{{version}}]" - ) - vault_url = parts[0] - name_parts = parts[1].rstrip("/").split("/") - secret_name = name_parts[0] - secret_version = name_parts[1] if len(name_parts) > 1 else None - return vault_url, secret_name, secret_version + error_message = ( + f"Invalid AKV secret URL: '{secret_url}'. Expected an HTTPS Azure Key Vault URL in the format " + "https://{vault}.{vault-dns-suffix}/secrets/{name}[/{version}]." + ) + try: + parsed_url = urllib.parse.urlsplit(secret_url) + port = parsed_url.port + except (TypeError, ValueError) as error: + raise ValueError(error_message) from error + + hostname = parsed_url.hostname + vault_name, separator, dns_suffix = hostname.partition(".") if hostname else ("", "", "") + valid_vault_name = ( + 1 <= len(vault_name) <= 63 + and all(char.isascii() and (char.isalnum() or char == "-") for char in vault_name) + ) + valid_authority = ( + parsed_url.scheme.casefold() == "https" + and parsed_url.username is None + and parsed_url.password is None + and port is None + and separator == "." + and dns_suffix in _AKV_VAULT_DNS_SUFFIXES + and valid_vault_name + ) + path_parts = parsed_url.path.split("/") + valid_path = ( + len(path_parts) in {3, 4} + and path_parts[0] == "" + and path_parts[1] == "secrets" + and all(path_parts[2:]) + ) + if not valid_authority or not valid_path or parsed_url.query or parsed_url.fragment: + raise ValueError(error_message) + + secret_name = path_parts[2] + secret_version = path_parts[3] if len(path_parts) == 4 else None + if not _is_valid_akv_identifier(secret_name) or ( + secret_version is not None and not _is_valid_akv_identifier(secret_version) + ): + raise ValueError(error_message) + + return f"https://{hostname}", secret_name, secret_version + + +def _is_valid_akv_identifier(identifier: str) -> bool: + """ + Check whether a Key Vault secret name or version uses URL-safe characters. + + Returns: + bool: True when the identifier is valid. + """ + return 1 <= len(identifier) <= 127 and all( + char.isascii() and (char.isalnum() or char == "-") for char in identifier + ) def _create_akv_secret_client(*, vault_url: str, credential: Any) -> "SecretClient": @@ -361,11 +334,12 @@ async def _load_env_from_akv_async( secret_url: str, strict: bool = True, silent: bool = False, -) -> tuple[dict[str, str], str]: +) -> None: """ - Load a bootstrap environment document from an Azure Key Vault secret. + Load a bootstrap dotenv document and resolve its same-vault secret references. - References remain unresolved until all environment sources are merged. + References are resolved once. Referenced secret values are treated as terminal + strings and are not interpreted as additional references. Authentication uses ``DefaultAzureCredential``, which silently tries managed identity, Azure CLI, VS Code credentials, etc., and falls back to interactive @@ -378,9 +352,6 @@ async def _load_env_from_akv_async( If False, warn and skip those entries. Defaults to True. silent (bool): If True, suppresses print statements. Defaults to False. - Returns: - tuple[dict[str, str], str]: Parsed bootstrap values and the vault URL. - Raises: ImportError: If ``azure-keyvault-secrets`` is not installed. KeyVaultInitializationException: If the root URL is malformed or the bootstrap environment @@ -403,212 +374,107 @@ async def _load_env_from_akv_async( parsed_environment = dotenv.dotenv_values(stream=io.StringIO(validated_document), interpolate=True) if not parsed_environment: raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") - - return {name: value for name, value in parsed_environment.items() if value is not None}, vault_url - except KeyVaultInitializationException: - raise - except Exception as error: - wrapped_error = _key_vault_initialization_error( - message=f"Failed to load Key Vault bootstrap secret '{secret_url}'", - error=error, - ) - raise wrapped_error from error - - -async def _resolve_environment_references_async( - *, - values: Mapping[str, str], - ambient_environment: Mapping[str, str], - bootstrap_vault_url: str | None = None, -) -> dict[str, str]: - """ - Resolve references after all environment sources have been merged. - - Args: - values (Mapping[str, str]): Winning values after source precedence. - ambient_environment (Mapping[str, str]): Process environment visible to ``env:`` references. - bootstrap_vault_url (str | None): Vault URL imposed by the bootstrap document. - - Returns: - dict[str, str]: Values with complete-value references resolved. - - Raises: - KeyVaultInitializationException: If a Key Vault reference is invalid or uses another vault. - ValueError: If an environment reference cannot be resolved. - """ - reference_environment = {**ambient_environment, **values} - vault_url = bootstrap_vault_url - reference_variable_name = "" - try: - for variable_name, value in values.items(): - reference_variable_name = variable_name - reference = _parse_environment_value_reference(value) - if reference is None or reference[0] != "akv": - continue - target = reference[1] - if not target.casefold().startswith("https://"): - _resolve_akv_secret_reference( - target=target, - variable_name=variable_name, - vault_url=vault_url or "", + loaded = dotenv.load_dotenv( + stream=io.StringIO(validated_document), + override=True, + interpolate=True, ) - referenced_vault_url, _, _ = _parse_akv_secret_url(target) - if vault_url is None: - vault_url = referenced_vault_url - _resolve_akv_secret_reference( - target=target, - variable_name=variable_name, - vault_url=vault_url, - ) + if not loaded: + return + + for variable_name, value in parsed_environment.items(): + if value is None: + continue + target = _parse_akv_reference(value) + if target is None: + continue + try: + referenced_name, referenced_version = _resolve_akv_secret_reference( + target=target, + variable_name=variable_name, + vault_url=vault_url, + ) + referenced_secret = await client.get_secret(referenced_name, version=referenced_version) + if referenced_secret.value is None: + raise ValueError( + f"AKV secret '{referenced_name}' referenced by environment variable " + f"'{variable_name}' has no value." + ) + os.environ[variable_name] = referenced_secret.value + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to resolve Key Vault reference for environment variable '{variable_name}'", + error=error, + ) + raise wrapped_error from error except KeyVaultInitializationException: raise except Exception as error: wrapped_error = _key_vault_initialization_error( - message=f"Invalid Key Vault reference for environment variable '{reference_variable_name}'", - error=error, - ) - raise wrapped_error from error - - if vault_url is None: - return { - name: await _resolve_environment_value_async( - value=value, - variable_name=name, - secret_client=None, - vault_url=None, - reference_environment=reference_environment, - ) - for name, value in values.items() - } - - from azure.core.exceptions import AzureError - from azure.identity.aio import DefaultAzureCredential - - try: - async with DefaultAzureCredential() as credential: - async with _create_akv_secret_client(vault_url=vault_url, credential=credential) as client: - return { - name: await _resolve_environment_value_async( - value=value, - variable_name=name, - secret_client=client, - vault_url=vault_url, - reference_environment=reference_environment, - ) - for name, value in values.items() - } - except KeyVaultInitializationException: - raise - except AzureError as error: - wrapped_error = _key_vault_initialization_error( - message=f"Failed to connect to Key Vault '{vault_url}'", + message=f"Failed to load Key Vault bootstrap secret '{secret_url}'", error=error, ) raise wrapped_error from error -async def _prepare_environment_updates_async( +async def _load_environment_async( *, - env_akv_ref: str | None, + env_akv_ref: Sequence[str] | None, env_files: Sequence[pathlib.Path] | None, env_akv_strict: bool, silent: bool, -) -> dict[str, str]: +) -> None: """ - Stage all environment sources and resolve references before committing. + Load environment sources in precedence order. Args: - env_akv_ref (str | None): Optional Key Vault bootstrap secret URL. + env_akv_ref (Sequence[str] | None): Optional ordered Key Vault bootstrap secret URLs. env_files (Sequence[pathlib.Path] | None): Optional ordered local environment files. env_akv_strict (bool): Whether bootstrap dotenv validation is strict. silent (bool): Whether initialization messages are suppressed. - Returns: - dict[str, str]: Fully resolved environment updates. - Raises: ValueError: If a configured source or reference is invalid. """ - ambient_environment = dict(os.environ) - merged_values: dict[str, str] = {} - bootstrap_vault_url: str | None = None - - if env_akv_ref is not None: - if not env_akv_ref.strip(): - raise ValueError("env_akv_ref must be a non-empty Azure Key Vault secret URL.") + if isinstance(env_akv_ref, str): + raise ValueError("env_akv_ref must be a sequence of Azure Key Vault secret URLs.") + if env_akv_ref: + if any(not isinstance(secret_url, str) or not secret_url.strip() for secret_url in env_akv_ref): + raise ValueError("env_akv_ref must contain only non-empty Azure Key Vault secret URLs.") await asyncio.to_thread( _warn_about_akv_environment_files, env_files=env_files, silent=silent, ) - bootstrap_values, bootstrap_vault_url = await _load_env_from_akv_async( - secret_url=env_akv_ref, - strict=env_akv_strict, - silent=silent, - ) - merged_values.update(bootstrap_values) + for secret_url in env_akv_ref: + await _load_env_from_akv_async( + secret_url=secret_url, + strict=env_akv_strict, + silent=silent, + ) - local_values, _ = await asyncio.to_thread( - _resolve_environment_files, + await asyncio.to_thread( + _load_environment_files, env_files=env_files, - base_environment={**ambient_environment, **merged_values}, silent=silent, ) - merged_values.update(local_values) - return await _resolve_environment_references_async( - values=merged_values, - ambient_environment=ambient_environment, - bootstrap_vault_url=bootstrap_vault_url, - ) - -def _parse_environment_value_reference(value: str) -> tuple[str, str] | None: +def _parse_akv_reference(value: str) -> str | None: """ - Parse an exact whole-value environment or Key Vault reference. + Parse an exact whole-value Key Vault reference. Returns: - The normalized reference type and target, or None for a literal value. + The referenced secret URL, or None for a literal value. """ prefix, separator, target = value.partition(":") - if not separator: - return None - if prefix == "env": - return "env", target.strip() - if prefix in _AKV_REFERENCE_PREFIXES: - return "akv", target.strip() - if prefix == "literal": - return "literal", target - return None - - -def _lookup_environment_value(*, environment: Mapping[str, str], name: str) -> str | None: - """ - Look up an environment value using platform-appropriate name semantics. - - Returns: - str | None: The matched value, or None when no name matches. - """ - if name in environment: - return environment[name] - if os.name == "nt": - folded_name = name.casefold() - return next((value for key, value in environment.items() if key.casefold() == folded_name), None) - return None - - -def _environment_names_equal(*, left: str, right: str) -> bool: - """ - Compare environment variable names using platform semantics. - - Returns: - bool: True when the names identify the same environment variable. - """ - return left.casefold() == right.casefold() if os.name == "nt" else left == right + return target.strip() if separator and prefix in _AKV_REFERENCE_PREFIXES else None def _validate_akv_secret_name(*, secret_name: str, variable_name: str) -> None: - if not secret_name or len(secret_name) > 127 or any(not char.isalnum() and char != "-" for char in secret_name): + if not _is_valid_akv_identifier(secret_name): raise ValueError( f"Invalid same-vault secret name '{secret_name}' referenced by environment variable '{variable_name}'. " "Secret names must contain only letters, numbers, and hyphens." @@ -652,80 +518,6 @@ def _resolve_akv_secret_reference( return secret_name, secret_version -async def _resolve_environment_value_async( - *, - value: str, - variable_name: str, - secret_client: "SecretClient | None", - vault_url: str | None, - reference_environment: Mapping[str, str], -) -> str: - """ - Resolve one value from the bootstrap environment document. - - Args: - value (str): The parsed bootstrap value. - variable_name (str): The environment variable receiving the resolved value. - secret_client (SecretClient | None): Client for Key Vault references, when needed. - vault_url (str | None): The allowed Key Vault URL, when one is needed. - reference_environment (Mapping[str, str]): Merged source values with ambient fallback. - - Returns: - str: The literal, ambient, or same-vault scalar value. - - Raises: - KeyVaultInitializationException: If a Key Vault reference cannot be resolved. - ValueError: If a reference is empty or cannot resolve to a value. - """ - reference = _parse_environment_value_reference(value) - if reference is None: - return value - - reference_type, target = reference - if reference_type == "literal": - return target - if not target: - raise ValueError(f"Empty {reference_type} reference for environment variable '{variable_name}'.") - - if reference_type == "env": - if _environment_names_equal(left=target, right=variable_name): - raise ValueError( - f"Environment variable '{variable_name}' cannot reference itself. Use a distinct source variable name." - ) - resolved_value = _lookup_environment_value(environment=reference_environment, name=target) - if resolved_value is None: - raise ValueError( - f"Environment variable '{target}' referenced by '{variable_name}' " - "is not available in the merged environment." - ) - return resolved_value - - try: - if secret_client is None or vault_url is None: - raise ValueError(f"AKV reference for environment variable '{variable_name}' has no available vault client.") - - secret_name, secret_version = _resolve_akv_secret_reference( - target=target, - variable_name=variable_name, - vault_url=vault_url, - ) - - secret = await secret_client.get_secret(secret_name, version=secret_version) - if secret.value is None: - raise ValueError( - f"AKV secret '{secret_name}' referenced by environment variable '{variable_name}' has no value." - ) - return secret.value - except KeyVaultInitializationException: - raise - except Exception as error: - wrapped_error = _key_vault_initialization_error( - message=f"Failed to resolve Key Vault reference for environment variable '{variable_name}'", - error=error, - ) - raise wrapped_error from error - - async def _execute_initializers_async(*, initializers: Sequence["PyRITInitializer"]) -> None: """ Execute PyRITInitializer instances in the order provided. @@ -774,7 +566,7 @@ async def initialize_pyrit_async( initializers: Sequence["PyRITInitializer"] | None = None, load_defaults: bool = True, env_files: Sequence[pathlib.Path] | None = None, - env_akv_ref: str | None = None, + env_akv_ref: Sequence[str] | None = None, env_akv_strict: bool = True, silent: bool = False, **memory_instance_kwargs: Any, @@ -802,9 +594,9 @@ async def initialize_pyrit_async( env_files (Sequence[pathlib.Path] | None): Optional sequence of environment file paths to load in order. If not provided, will load default .env and .env.local files from PyRIT home if they exist. All paths must be valid pathlib.Path objects. - env_akv_ref (str | None): Optional Azure Key Vault URL whose secret value contains the - bootstrap .env document. Loaded before ``env_files`` so local files take precedence - over AKV. Requires ``azure-keyvault-secrets``. + env_akv_ref (Sequence[str] | None): Optional ordered Azure Key Vault URLs whose secret values + contain bootstrap .env documents. Loaded before ``env_files`` so later bootstrap documents + and local files take precedence. Requires ``azure-keyvault-secrets``. env_akv_strict (bool): If True, reject malformed or valueless entries in the Key Vault bootstrap document. If False, warn and skip those entries. Defaults to True. silent (bool): If True, suppresses print statements about environment file loading and @@ -814,13 +606,12 @@ async def initialize_pyrit_async( Raises: ValueError: If an unsupported memory_db_type is provided or env_files contains non-existent files. """ - environment_updates = await _prepare_environment_updates_async( + await _load_environment_async( env_akv_ref=env_akv_ref, env_files=env_files, env_akv_strict=env_akv_strict, silent=silent, ) - os.environ.update(environment_updates) # Reset all default values before executing initialization scripts # This ensures a clean state for each initialization diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 986c5a052b..9682e9b2ae 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -147,7 +147,7 @@ def test_from_dict_with_all_fields(self): "initializers": ["simple"], "initialization_scripts": ["/path/to/script.py"], "env_files": ["/path/to/.env"], - "env_akv_ref": "https://vault.vault.azure.net/secrets/one", + "env_akv_ref": ["https://vault.vault.azure.net/secrets/one"], "env_akv_strict": False, "silent": True, } @@ -156,7 +156,7 @@ def test_from_dict_with_all_fields(self): assert config.initializers == ["simple"] assert config.initialization_scripts == ["/path/to/script.py"] assert config.env_files == ["/path/to/.env"] - assert config.env_akv_ref == "https://vault.vault.azure.net/secrets/one" + assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] assert config.env_akv_strict is False assert config.silent is True @@ -310,14 +310,20 @@ def testresolve_env_akv_ref_none_returns_none(self): assert config.resolve_env_akv_ref() is None def testresolve_env_akv_ref_returns_configured_values(self): - """Test that the configured AKV reference is returned unchanged.""" - ref = "https://vault.vault.azure.net/secrets/first" - config = ConfigurationLoader(env_akv_ref=ref) - assert config.resolve_env_akv_ref() == ref - - @pytest.mark.parametrize("env_akv_ref", [[], ["https://vault.vault.azure.net/secrets/one"], ""]) - def test_env_akv_ref_rejects_non_scalar_or_empty_values(self, env_akv_ref): - with pytest.raises(ValueError, match="env_akv_ref must be one non-empty"): + """Test that the configured AKV references are returned unchanged.""" + refs = [ + "https://vault.vault.azure.net/secrets/first", + "https://vault.vault.azure.net/secrets/second/version", + ] + config = ConfigurationLoader(env_akv_ref=refs) + assert config.resolve_env_akv_ref() == refs + + def test_env_akv_ref_allows_empty_list(self): + assert ConfigurationLoader(env_akv_ref=[]).env_akv_ref == [] + + @pytest.mark.parametrize("env_akv_ref", ["", "https://vault.vault.azure.net/secrets/one", [""], [None]]) + def test_env_akv_ref_rejects_scalar_or_invalid_entries(self, env_akv_ref): + with pytest.raises(ValueError, match="env_akv_ref must"): ConfigurationLoader(env_akv_ref=env_akv_ref) # type: ignore[arg-type] @@ -345,14 +351,17 @@ async def test_initialize_pyrit_async_basic(self, mock_init): @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") async def test_initialize_pyrit_async_with_env_akv_ref(self, mock_init): """Test initialization forwards env_akv_ref to initialize_pyrit_async.""" - ref = "https://vault.vault.azure.net/secrets/first" - config = ConfigurationLoader(memory_db_type="in_memory", env_akv_ref=ref, env_akv_strict=False) + refs = [ + "https://vault.vault.azure.net/secrets/first", + "https://vault.vault.azure.net/secrets/second/version", + ] + config = ConfigurationLoader(memory_db_type="in_memory", env_akv_ref=refs, env_akv_strict=False) await config.initialize_pyrit_async() mock_init.assert_called_once() call_kwargs = mock_init.call_args.kwargs - assert call_kwargs["env_akv_ref"] == ref + assert call_kwargs["env_akv_ref"] == refs assert call_kwargs["env_akv_strict"] is False @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") @@ -465,13 +474,13 @@ def test_load_with_overrides_reads_env_akv_ref_from_default_config(self, mock_de mock_default_path.exists.return_value = True mock_from_yaml.return_value = ConfigurationLoader( memory_db_type="sqlite", - env_akv_ref="https://default.vault.azure.net/secrets/from-default", + env_akv_ref=["https://default.vault.azure.net/secrets/from-default"], ) config = ConfigurationLoader.load_with_overrides() mock_from_yaml.assert_called_once_with(mock_default_path) - assert config.env_akv_ref == "https://default.vault.azure.net/secrets/from-default" + assert config.env_akv_ref == ["https://default.vault.azure.net/secrets/from-default"] @mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH") def test_load_with_overrides_memory_db_type_override(self, mock_default_path): @@ -516,10 +525,10 @@ def test_load_with_overrides_env_akv_ref_override(self, mock_default_path): mock_default_path.exists.return_value = False config = ConfigurationLoader.load_with_overrides( - env_akv_ref="https://vault.vault.azure.net/secrets/one", + env_akv_ref=["https://vault.vault.azure.net/secrets/one"], ) - assert config.env_akv_ref == "https://vault.vault.azure.net/secrets/one" + assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] @mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH") def test_load_with_overrides_converts_sequence_to_list(self, mock_default_path): @@ -531,15 +540,18 @@ def test_load_with_overrides_converts_sequence_to_list(self, mock_default_path): initializers=("init1", "init2"), initialization_scripts=("/path/script1.py", "/path/script2.py"), env_files=("/path/.env",), + env_akv_ref=("https://vault.vault.azure.net/secrets/one",), ) # Verify they are stored as lists assert isinstance(config.initializers, list) assert isinstance(config.initialization_scripts, list) assert isinstance(config.env_files, list) + assert isinstance(config.env_akv_ref, list) assert config.initializers == ["init1", "init2"] assert config.initialization_scripts == ["/path/script1.py", "/path/script2.py"] assert config.env_files == ["/path/.env"] + assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] def test_load_with_overrides_explicit_config_file_not_found(self): """Test FileNotFoundError when explicit config file doesn't exist.""" @@ -562,7 +574,8 @@ def test_load_with_overrides_explicit_config_file_overrides_default(self, mock_d - /explicit/script.py env_files: - /explicit/.env -env_akv_ref: https://vault.vault.azure.net/secrets/explicit +env_akv_ref: + - https://vault.vault.azure.net/secrets/explicit """ with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: f.write(yaml_content) @@ -575,7 +588,7 @@ def test_load_with_overrides_explicit_config_file_overrides_default(self, mock_d assert config._initializer_configs[0].name == "explicit_init" assert config.initialization_scripts == ["/explicit/script.py"] assert config.env_files == ["/explicit/.env"] - assert config.env_akv_ref == "https://vault.vault.azure.net/secrets/explicit" + assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/explicit"] finally: config_path.unlink() diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 694dd3046b..0e53da0deb 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -18,10 +18,8 @@ from pyrit.setup.initialization import ( _load_env_from_akv_async, _load_environment_files, + _parse_akv_reference, _parse_akv_secret_url, - _parse_environment_value_reference, - _resolve_environment_files, - _resolve_environment_references_async, _warn_about_akv_environment_files, ) @@ -131,17 +129,17 @@ def setup_method(self) -> None: reset_default_values() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) - async def test_initialize_basic(self, mock_resolve_env, mock_set_memory): + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) + async def test_initialize_basic(self, mock_load_env, mock_set_memory): """Test basic initialization.""" await initialize_pyrit_async(memory_db_type=IN_MEMORY, load_defaults=False) - mock_resolve_env.assert_called_once() + mock_load_env.assert_called_once() mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) - async def test_initialize_with_script(self, mock_resolve_env, mock_set_memory): + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) + async def test_initialize_with_script(self, mock_load_env, mock_set_memory): """Test initialization with a script.""" with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f: f.write( @@ -165,50 +163,63 @@ async def initialize_async(self) -> None: try: await initialize_pyrit_async(memory_db_type=IN_MEMORY, initialization_scripts=[script_path]) - mock_resolve_env.assert_called_once() + mock_load_env.assert_called_once() mock_set_memory.assert_called_once() finally: os.unlink(script_path) - @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) - async def test_invalid_memory_type_raises_error(self, mock_resolve_env): + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) + async def test_invalid_memory_type_raises_error(self, mock_load_env): """Test that invalid memory type raises ValueError.""" with pytest.raises(ValueError, match="is not a supported type"): await initialize_pyrit_async(memory_db_type="InvalidType", load_defaults=False) # type: ignore[arg-type] @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) - async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_resolve_env, mock_set_memory): - """Test that env_akv_ref loads its bootstrap secret.""" - ref = "https://vault.vault.azure.net/secrets/test-secret" + async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, mock_set_memory): + """Test that env_akv_ref loads bootstrap secrets in order.""" + refs = [ + "https://vault.vault.azure.net/secrets/first", + "https://vault.vault.azure.net/secrets/second/version", + ] - mock_load_akv.return_value = {}, "https://vault.vault.azure.net" + mock_load_akv.return_value = None - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=ref, load_defaults=False) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) - mock_load_akv.assert_awaited_once() - assert mock_load_akv.await_args.kwargs["secret_url"] == ref - assert mock_load_akv.await_args.kwargs["strict"] is True - assert mock_load_akv.await_args.kwargs["silent"] is False - mock_resolve_env.assert_called_once() + assert mock_load_akv.await_args_list == [ + mock.call(secret_url=refs[0], strict=True, silent=False), + mock.call(secret_url=refs[1], strict=True, silent=False), + ] + mock_load_env.assert_called_once() mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._resolve_environment_files", return_value=({}, False)) + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) - async def test_initialize_with_empty_env_akv_ref_raises(self, mock_load_akv, mock_resolve_env, mock_set_memory): - """Test that an empty env_akv_ref is rejected.""" - with pytest.raises(ValueError, match="env_akv_ref must be a non-empty"): - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref="", load_defaults=False) + async def test_initialize_with_empty_env_akv_ref_does_not_load_akv( + self, mock_load_akv, mock_load_env, mock_set_memory + ): + """Test that an empty env_akv_ref list skips AKV loading.""" + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=[], load_defaults=False) mock_load_akv.assert_not_called() - mock_resolve_env.assert_not_called() - mock_set_memory.assert_not_called() + mock_load_env.assert_called_once() + mock_set_memory.assert_called_once() + + @pytest.mark.parametrize("env_akv_ref", ["https://vault.vault.azure.net/secrets/one", [""], [None]]) + async def test_initialize_rejects_invalid_env_akv_ref(self, env_akv_ref): + with pytest.raises(ValueError, match="env_akv_ref must"): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=env_akv_ref, # type: ignore[arg-type] + load_defaults=False, + ) @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_akv_and_local_files_are_applied_atomically(self, mock_set_memory): - ref = "https://vault.vault.azure.net/secrets/bootstrap" + async def test_initialize_keeps_akv_values_when_local_file_loading_fails(self, mock_set_memory): + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] nonexistent = pathlib.Path("/nonexistent/.env") with mock.patch.dict(os.environ, {}, clear=True): @@ -217,24 +228,24 @@ async def test_initialize_akv_and_local_files_are_applied_atomically(self, mock_ mock.patch( "pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock, - return_value=({"FROM_AKV": "resolved"}, "https://vault.vault.azure.net"), + side_effect=lambda **_: os.environ.update({"FROM_AKV": "resolved"}), ), pytest.raises(ValueError, match="Environment file not found"), ): await initialize_pyrit_async( memory_db_type=IN_MEMORY, - env_akv_ref=ref, + env_akv_ref=refs, env_files=[nonexistent], load_defaults=False, ) - assert "FROM_AKV" not in os.environ + assert os.environ["FROM_AKV"] == "resolved" mock_set_memory.assert_not_called() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_stages_local_overrides_on_akv_environment(self, mock_set_memory): - ref = "https://vault.vault.azure.net/secrets/bootstrap" + async def test_initialize_loads_local_overrides_on_akv_environment(self, mock_set_memory): + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] with tempfile.TemporaryDirectory() as temp_dir: local_file = pathlib.Path(temp_dir) / ".env.local" local_file.write_text("DERIVED=${BASE}\nBASE=local") @@ -245,12 +256,12 @@ async def test_initialize_stages_local_overrides_on_akv_environment(self, mock_s mock.patch( "pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock, - return_value=({"BASE": "akv", "ONLY_AKV": "shared"}, "https://vault.vault.azure.net"), + side_effect=lambda **_: os.environ.update({"BASE": "akv", "ONLY_AKV": "shared"}), ), ): await initialize_pyrit_async( memory_db_type=IN_MEMORY, - env_akv_ref=ref, + env_akv_ref=refs, env_files=[local_file], load_defaults=False, ) @@ -263,7 +274,7 @@ async def test_initialize_stages_local_overrides_on_akv_environment(self, mock_s @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") async def test_initialize_default_files_override_akv_in_order(self, mock_set_memory): - ref = "https://vault.vault.azure.net/secrets/bootstrap" + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) (temp_path / ".env").write_text("VALUE=env") @@ -276,12 +287,12 @@ async def test_initialize_default_files_override_akv_in_order(self, mock_set_mem mock.patch( "pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock, - return_value=({"VALUE": "akv"}, "https://vault.vault.azure.net"), + side_effect=lambda **_: os.environ.update({"VALUE": "akv"}), ), ): await initialize_pyrit_async( memory_db_type=IN_MEMORY, - env_akv_ref=ref, + env_akv_ref=refs, load_defaults=False, silent=True, ) @@ -291,15 +302,8 @@ async def test_initialize_default_files_override_akv_in_order(self, mock_set_mem mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_resolves_only_winning_references_after_local_override(self, mock_set_memory): - ref = "https://vault.vault.azure.net/secrets/bootstrap" - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock( - side_effect=[ - types.SimpleNamespace(value="bootstrap-secret-value"), - types.SimpleNamespace(value="local-secret-value"), - ] - ) + async def test_initialize_resolves_bootstrap_references_before_local_overrides(self, mock_set_memory): + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] with tempfile.TemporaryDirectory() as temp_dir: local_file = pathlib.Path(temp_dir) / ".env.local" local_file.write_text( @@ -308,8 +312,8 @@ async def test_initialize_resolves_only_winning_references_after_local_override( "LOCAL_ENV=env:BOOTSTRAP_SOURCE" ) bootstrap_environment = { - "OVERRIDDEN": "kv:https://vault.vault.azure.net/secrets/unused-secret", - "BOOTSTRAP_SECRET": "kv:https://vault.vault.azure.net/secrets/bootstrap-secret", + "OVERRIDDEN": "unused-secret-value", + "BOOTSTRAP_SECRET": "bootstrap-secret-value", "BOOTSTRAP_SOURCE": "bootstrap-value", } @@ -319,27 +323,21 @@ async def test_initialize_resolves_only_winning_references_after_local_override( mock.patch( "pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock, - return_value=(bootstrap_environment, "https://vault.vault.azure.net"), + side_effect=lambda **_: os.environ.update(bootstrap_environment), ), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), ): await initialize_pyrit_async( memory_db_type=IN_MEMORY, - env_akv_ref=ref, + env_akv_ref=refs, env_files=[local_file], load_defaults=False, ) assert os.environ["OVERRIDDEN"] == "local" assert os.environ["BOOTSTRAP_SECRET"] == "bootstrap-secret-value" - assert os.environ["LOCAL_SECRET"] == "local-secret-value" - assert os.environ["LOCAL_ENV"] == "bootstrap-value" + assert os.environ["LOCAL_SECRET"] == "kv:https://vault.vault.azure.net/secrets/local-secret" + assert os.environ["LOCAL_ENV"] == "env:BOOTSTRAP_SOURCE" - assert client.get_secret.await_args_list == [ - mock.call("bootstrap-secret", version=None), - mock.call("local-secret", version=None), - ] mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") @@ -525,6 +523,19 @@ async def test_loads_custom_env_files_in_order(self): assert loaded is True assert os.environ["VAR"] == "local" + async def test_load_environment_files_interpolates_in_assignment_order(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("A=one\nB=${A}\nA=two\nC=${A}") + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert os.environ["A"] == "two" + assert os.environ["B"] == "one" + assert os.environ["C"] == "two" + async def test_load_environment_files_honors_python_dotenv_disabled(self): with tempfile.TemporaryDirectory() as temp_dir: env_file = pathlib.Path(temp_dir) / ".env" @@ -563,16 +574,13 @@ def test_default_local_file_can_interpolate_base_file_but_not_reverse(self, mock env_local_file.write_text("FOOBAR=${OPENAI_CHAT_ENDPOINT}\nLOCAL_ONLY=local") mock_config_path.__truediv__ = lambda self, other: temp_path / other - resolved, loaded = _resolve_environment_files( - env_files=None, - base_environment={}, - silent=True, - ) + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None, silent=True) - assert loaded is True - assert resolved["FOOBAR"] == "https://example.openai.azure.com/openai/v1" - assert resolved["FROM_LATER_LOCAL"] == "" - assert resolved["LOCAL_ONLY"] == "local" + assert loaded is True + assert os.environ["FOOBAR"] == "https://example.openai.azure.com/openai/v1" + assert os.environ["FROM_LATER_LOCAL"] == "" + assert os.environ["LOCAL_ONLY"] == "local" async def test_env_akv_strict_does_not_validate_local_environment_files(self): with tempfile.TemporaryDirectory() as temp_dir: @@ -592,18 +600,12 @@ async def test_env_akv_strict_does_not_validate_local_environment_files(self): assert os.environ["OTHER"] == "also-resolved" @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_local_file_resolves_full_akv_reference_without_bootstrap(self, mock_set_memory): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="local-secret-value")) + async def test_initialize_keeps_local_akv_reference_literal_without_bootstrap(self, mock_set_memory): with tempfile.TemporaryDirectory() as temp_dir: env_file = pathlib.Path(temp_dir) / ".env" env_file.write_text("API_KEY=kv:https://myvault.vault.azure.net/secrets/api-key") - with ( - mock.patch.dict(os.environ, {}, clear=True), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, - ): + with mock.patch.dict(os.environ, {}, clear=True): await initialize_pyrit_async( memory_db_type=IN_MEMORY, env_files=[env_file], @@ -611,14 +613,8 @@ async def test_initialize_local_file_resolves_full_akv_reference_without_bootstr silent=True, ) - assert os.environ["API_KEY"] == "local-secret-value" + assert os.environ["API_KEY"] == "kv:https://myvault.vault.azure.net/secrets/api-key" - _assert_mock_akv_client_created( - mock_client_cls, - vault_url="https://myvault.vault.azure.net", - credential=credential, - ) - client.get_secret.assert_awaited_once_with("api-key", version=None) mock_set_memory.assert_called_once() async def test_raises_error_for_nonexistent_env_file(self): @@ -705,15 +701,21 @@ class TestAkvEnvironmentLoading: """Tests for AKV URL parsing and env loading helpers.""" @pytest.mark.parametrize("prefix", ["kv", "akv", "azure_key_vault", "env_akv_ref"]) - def test_parse_environment_value_reference_accepts_akv_aliases(self, prefix): + def test_parse_akv_reference_accepts_aliases(self, prefix): secret_url = "https://myvault.vault.azure.net/secrets/api-key" - assert _parse_environment_value_reference(f"{prefix}:{secret_url}") == ("akv", secret_url) + assert _parse_akv_reference(f"{prefix}:{secret_url}") == secret_url - def test_parse_environment_value_reference_rejects_azure_app_service_syntax(self): - value = "@Microsoft.KeyVault(SecretUri=https://myvault.vault.azure.net/secrets/api-key)" - - assert _parse_environment_value_reference(value) is None + @pytest.mark.parametrize( + "value", + [ + "env:SOURCE_VALUE", + "literal:kv:https://myvault.vault.azure.net/secrets/api-key", + "@Microsoft.KeyVault(SecretUri=https://myvault.vault.azure.net/secrets/api-key)", + ], + ) + def test_parse_akv_reference_ignores_non_akv_syntax(self, value): + assert _parse_akv_reference(value) is None def test_parse_akv_secret_url_with_version(self): url = "https://myvault.vault.azure.net/secrets/my-secret/abc123" @@ -733,37 +735,88 @@ def test_parse_akv_secret_url_without_version(self): assert secret_name == "my-secret" assert secret_version is None - def test_parse_akv_secret_url_invalid_raises(self): + @pytest.mark.parametrize("dns_suffix", ["vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"]) + def test_parse_akv_secret_url_accepts_supported_clouds(self, dns_suffix): + url = f"https://myvault.{dns_suffix}/secrets/my-secret/version-1" + + vault_url, secret_name, secret_version = _parse_akv_secret_url(url) + + assert vault_url == f"https://myvault.{dns_suffix}" + assert secret_name == "my-secret" + assert secret_version == "version-1" + + @pytest.mark.parametrize( + "url", + [ + "http://myvault.vault.azure.net/secrets/my-secret", + "https://attacker.example/secrets/my-secret", + "https://myvault.vault.azure.net.attacker.example/secrets/my-secret", + "https://nested.myvault.vault.azure.net/secrets/my-secret", + "https://user@myvault.vault.azure.net/secrets/my-secret", + "https://myvault.vault.azure.net:443/secrets/my-secret", + "https://myvault.vault.azure.net/not-secrets/my-secret", + "https://myvault.vault.azure.net/secrets", + "https://myvault.vault.azure.net/secrets/my-secret/", + "https://myvault.vault.azure.net/secrets/my-secret/version/extra", + "https://myvault.vault.azure.net/secrets/my-secret?api-version=7.4", + "https://myvault.vault.azure.net/secrets/my-secret#fragment", + "https://myvault.vault.azure.net/secrets/my%2Fsecret", + ], + ) + def test_parse_akv_secret_url_invalid_raises(self, url): with pytest.raises(ValueError, match="Invalid AKV secret URL"): - _parse_akv_secret_url("https://myvault.vault.azure.net/not-secrets/my-secret") + _parse_akv_secret_url(url) - async def test_load_env_from_akv_async_returns_unresolved_bootstrap(self): + async def test_load_env_from_akv_async_rejects_non_azure_host_before_authentication(self): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + mock.patch("pyrit.setup.initialization._create_akv_secret_client") as mock_create_client, + pytest.raises(KeyVaultInitializationException, match="attacker.example"), + ): + await _load_env_from_akv_async( + secret_url="https://attacker.example/secrets/bootstrap", + silent=True, + ) + + mock_credential_cls.assert_not_called() + mock_create_client.assert_not_called() + + async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_secrets(self): credential, client = _create_mock_akv_clients() root_document = ( "DIRECT=from-bootstrap\n" - "FROM_ENV=env:SOURCE_VALUE\n" + "FROM_ENV=${SOURCE_VALUE}\n" "FROM_KV=kv:https://myvault.vault.azure.net/secrets/api-key\n" "PINNED_KV=kv:https://myvault.vault.azure.net/secrets/api-key/version-2\n" - "ESCAPED=literal:kv:not-a-secret" + "TERMINAL=kv:https://myvault.vault.azure.net/secrets/terminal\n" + "A=one\nB=${A}\nA=two\nC=${A}" + ) + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value=root_document), + types.SimpleNamespace(value="api-key-value"), + types.SimpleNamespace(value="pinned-key-value"), + types.SimpleNamespace(value="kv:https://myvault.vault.azure.net/secrets/not-followed"), + ] ) - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=root_document)) secret_url = "https://myvault.vault.azure.net/secrets/bootstrap/v1" with ( + mock.patch.dict(os.environ, {"SOURCE_VALUE": "ambient-value"}, clear=True), mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential) as mock_credential_cls, mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, ): - parsed_environment, vault_url = await _load_env_from_akv_async(secret_url=secret_url, silent=True) - - assert parsed_environment == { - "DIRECT": "from-bootstrap", - "FROM_ENV": "env:SOURCE_VALUE", - "FROM_KV": "kv:https://myvault.vault.azure.net/secrets/api-key", - "PINNED_KV": "kv:https://myvault.vault.azure.net/secrets/api-key/version-2", - "ESCAPED": "literal:kv:not-a-secret", - } - assert vault_url == "https://myvault.vault.azure.net" + await _load_env_from_akv_async(secret_url=secret_url, silent=True) + + assert os.environ["DIRECT"] == "from-bootstrap" + assert os.environ["FROM_ENV"] == "ambient-value" + assert os.environ["FROM_KV"] == "api-key-value" + assert os.environ["PINNED_KV"] == "pinned-key-value" + assert os.environ["TERMINAL"] == "kv:https://myvault.vault.azure.net/secrets/not-followed" + assert os.environ["A"] == "two" + assert os.environ["B"] == "one" + assert os.environ["C"] == "two" mock_credential_cls.assert_called_once_with() _assert_mock_akv_client_created( @@ -771,94 +824,31 @@ async def test_load_env_from_akv_async_returns_unresolved_bootstrap(self): vault_url="https://myvault.vault.azure.net", credential=credential, ) - client.get_secret.assert_awaited_once_with("bootstrap", version="v1") + assert client.get_secret.await_args_list == [ + mock.call("bootstrap", version="v1"), + mock.call("api-key", version=None), + mock.call("api-key", version="version-2"), + mock.call("terminal", version=None), + ] credential.__aenter__.assert_awaited_once() credential.__aexit__.assert_awaited_once() client.__aenter__.assert_awaited_once() client.__aexit__.assert_awaited_once() mock_print_msg.assert_called_once() - async def test_resolve_environment_references_async_resolves_local_values(self): + async def test_load_env_from_akv_async_rejects_short_secret_name(self): credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock( - side_effect=[ - types.SimpleNamespace(value="local-secret-value"), - types.SimpleNamespace(value="pinned-secret-value"), - ] - ) - values = { - "DIRECT": "from-local", - "FROM_ENV": "env:SOURCE_VALUE", - "DECLARED": "merged-value", - "FROM_DECLARED": "env:DECLARED", - "SHADOWED": "merged-wins", - "FROM_SHADOWED": "env:SHADOWED", - "FROM_KV": "kv:https://myvault.vault.azure.net/secrets/api-key", - "PINNED_KV": "akv:https://myvault.vault.azure.net/secrets/api-key/version-2", - "ESCAPED": "literal:kv:not-a-secret", - } + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="API_KEY=kv:api-key")) with ( + mock.patch.dict(os.environ, {}, clear=True), mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, - ): - resolved = await _resolve_environment_references_async( - values=values, - ambient_environment={"SOURCE_VALUE": "ambient-value", "SHADOWED": "ambient-loses"}, - ) - - assert resolved == { - "DIRECT": "from-local", - "FROM_ENV": "ambient-value", - "DECLARED": "merged-value", - "FROM_DECLARED": "merged-value", - "SHADOWED": "merged-wins", - "FROM_SHADOWED": "merged-wins", - "FROM_KV": "local-secret-value", - "PINNED_KV": "pinned-secret-value", - "ESCAPED": "kv:not-a-secret", - } - _assert_mock_akv_client_created( - mock_client_cls, - vault_url="https://myvault.vault.azure.net", - credential=credential, - ) - assert client.get_secret.await_args_list == [ - mock.call("api-key", version=None), - mock.call("api-key", version="version-2"), - ] - - async def test_resolve_environment_references_async_rejects_self_reference(self): - with pytest.raises(ValueError, match="cannot reference itself"): - await _resolve_environment_references_async( - values={"MODEL": "env:MODEL"}, - ambient_environment={"MODEL": "ambient-model"}, - ) - - async def test_resolve_environment_references_async_preserves_windows_case_insensitive_lookup(self): - with mock.patch("pyrit.setup.initialization.os.name", "nt"): - resolved = await _resolve_environment_references_async( - values={"ALIAS": "env:Path"}, - ambient_environment={"PATH": "windows-path"}, - ) - - assert resolved["ALIAS"] == "windows-path" - - async def test_resolve_environment_references_async_rejects_windows_case_variant_self_reference(self): - with ( - mock.patch("pyrit.setup.initialization.os.name", "nt"), - pytest.raises(ValueError, match="cannot reference itself"), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="must use a full secret URL"), ): - await _resolve_environment_references_async( - values={"MODEL": "env:model"}, - ambient_environment={}, - ) - - async def test_resolve_environment_references_async_rejects_short_secret_name(self): - with pytest.raises(ValueError, match="must use a full secret URL"): - await _resolve_environment_references_async( - values={"API_KEY": "kv:api-key"}, - ambient_environment={}, + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, ) @pytest.mark.parametrize( @@ -868,12 +858,19 @@ async def test_resolve_environment_references_async_rejects_short_secret_name(se "https://other-vault.vault.azure.net/secrets/api-key/version-1", ], ) - async def test_resolve_environment_references_async_rejects_cross_vault_reference(self, reference_url): - with pytest.raises(ValueError, match="Cross-vault AKV reference"): - await _resolve_environment_references_async( - values={"API_KEY": f"kv:{reference_url}"}, - ambient_environment={}, - bootstrap_vault_url="https://myvault.vault.azure.net", + async def test_load_env_from_akv_async_rejects_cross_vault_reference(self, reference_url): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=f"API_KEY=kv:{reference_url}")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="Cross-vault AKV reference"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, ) async def test_load_env_from_akv_async_empty_secret_raises(self): @@ -951,19 +948,25 @@ async def test_load_env_from_akv_async_wraps_malformed_bootstrap(self): assert isinstance(exc_info.value.__cause__, ValueError) - async def test_resolve_environment_references_async_wraps_missing_secret(self): + async def test_load_env_from_akv_async_wraps_missing_child_secret(self): credential, client = _create_mock_akv_clients() missing_error = ResourceNotFoundError(message="Secret was not found") - client.get_secret = mock.AsyncMock(side_effect=missing_error) + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value="API_KEY=kv:https://myvault.vault.azure.net/secrets/missing"), + missing_error, + ] + ) with ( + mock.patch.dict(os.environ, {}, clear=True), mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), pytest.raises(KeyVaultInitializationException, match="Failed to resolve Key Vault reference") as exc_info, ): - await _resolve_environment_references_async( - values={"API_KEY": "kv:https://myvault.vault.azure.net/secrets/missing"}, - ambient_environment={}, + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, ) assert exc_info.value.__cause__ is missing_error @@ -977,29 +980,34 @@ async def test_load_env_from_akv_async_allows_empty_assignment(self): mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), ): - resolved_environment, _ = await _load_env_from_akv_async( + await _load_env_from_akv_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", silent=True, ) - assert resolved_environment["EMPTY"] == "" - assert "EMPTY" not in os.environ + assert os.environ["EMPTY"] == "" - async def test_resolve_environment_references_async_allows_empty_child_secret(self): + async def test_load_env_from_akv_async_allows_empty_child_secret(self): credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="")) + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value="EMPTY=kv:https://myvault.vault.azure.net/secrets/empty-secret"), + types.SimpleNamespace(value=""), + ] + ) with ( + mock.patch.dict(os.environ, {}, clear=True), mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), ): - resolved_environment = await _resolve_environment_references_async( - values={"EMPTY": "kv:https://myvault.vault.azure.net/secrets/empty-secret"}, - ambient_environment={}, + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, ) - assert resolved_environment["EMPTY"] == "" - client.get_secret.assert_awaited_once_with("empty-secret", version=None) + assert os.environ["EMPTY"] == "" + assert client.get_secret.await_args_list[-1] == mock.call("empty-secret", version=None) async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entries(self, caplog, capsys): credential, client = _create_mock_akv_clients() @@ -1012,14 +1020,14 @@ async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entrie mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), caplog.at_level("WARNING", logger="pyrit.setup.initialization"), ): - resolved_environment, _ = await _load_env_from_akv_async( + await _load_env_from_akv_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", strict=False, silent=False, ) - assert resolved_environment == {"GOOD": "resolved", "OTHER": "also-resolved"} - assert "GOOD" not in os.environ + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" output = capsys.readouterr().out assert "WARNING: AKV environment document contains invalid entries that will be skipped" in output @@ -1047,9 +1055,16 @@ async def test_load_env_from_akv_async_non_strict_silent_logs_warning(self, capl assert capsys.readouterr().out == "" assert "variables without values: MISSING_VALUE" in caplog.text - async def test_resolve_environment_references_async_failure_returns_no_partial_mapping(self): + async def test_load_env_from_akv_async_child_failure_keeps_loaded_bootstrap_values(self): credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=None)) + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace( + value=("GOOD=resolved\nBAD=kv:https://myvault.vault.azure.net/secrets/missing-value") + ), + types.SimpleNamespace(value=None), + ] + ) with mock.patch.dict(os.environ, {}, clear=True): with ( @@ -1057,12 +1072,10 @@ async def test_resolve_environment_references_async_failure_returns_no_partial_m mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), pytest.raises(ValueError, match="has no value"), ): - await _resolve_environment_references_async( - values={ - "GOOD": "resolved", - "BAD": "kv:https://myvault.vault.azure.net/secrets/missing-value", - }, - ambient_environment={}, + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, ) - assert "GOOD" not in os.environ + assert os.environ["GOOD"] == "resolved" + assert os.environ["BAD"] == "kv:https://myvault.vault.azure.net/secrets/missing-value" From 8529c9d855b622b83aa3a70cf60e1bc5b2c0a5fa Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 13 Aug 2026 12:55:31 -0400 Subject: [PATCH 08/16] FIX: precommit --- pyrit/setup/initialization.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 19ab1171e1..5d3c713ec2 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -203,9 +203,8 @@ def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: hostname = parsed_url.hostname vault_name, separator, dns_suffix = hostname.partition(".") if hostname else ("", "", "") - valid_vault_name = ( - 1 <= len(vault_name) <= 63 - and all(char.isascii() and (char.isalnum() or char == "-") for char in vault_name) + valid_vault_name = 1 <= len(vault_name) <= 63 and all( + char.isascii() and (char.isalnum() or char == "-") for char in vault_name ) valid_authority = ( parsed_url.scheme.casefold() == "https" @@ -218,10 +217,7 @@ def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: ) path_parts = parsed_url.path.split("/") valid_path = ( - len(path_parts) in {3, 4} - and path_parts[0] == "" - and path_parts[1] == "secrets" - and all(path_parts[2:]) + len(path_parts) in {3, 4} and path_parts[0] == "" and path_parts[1] == "secrets" and all(path_parts[2:]) ) if not valid_authority or not valid_path or parsed_url.query or parsed_url.fragment: raise ValueError(error_message) From 9d98fb1f5f64cce09fec1b0f7e6ac0c12e6f756f Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 13 Aug 2026 13:04:43 -0400 Subject: [PATCH 09/16] FIX: docs --- doc/getting_started/pyrit_conf.md | 70 +++++++++++++------------------ 1 file changed, 30 insertions(+), 40 deletions(-) diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index baa0023961..283b614406 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -33,7 +33,7 @@ When PyRIT initializes, environment variables are loaded in a specific order. ** ```{mermaid} flowchart LR A["System environment"] --> B{"env_akv_ref configured?"} - B -->|Yes| C["AKV bootstrap"] + B -->|Yes| C["AKV bootstrap documents in order"] B -->|No| D{"Explicit env_files?"} C --> D D -->|Yes| E["Explicit files in order"] @@ -41,7 +41,7 @@ flowchart LR F --> G["~/.pyrit/.env.local"] ``` -System environment variables are always the baseline. If no AKV root or environment file is available, PyRIT continues initialization using the existing process environment only. +System environment variables are always the baseline. If no AKV bootstrap document or environment file is available, PyRIT continues initialization using the existing process environment only. **Default file behavior** (no `env_akv_ref` or `env_files` field in `.pyrit_conf`): @@ -51,11 +51,11 @@ System environment variables are always the baseline. If no AKV root or environm | Medium | `~/.pyrit/.env` | Default config file (loaded if it exists) | | Highest | `~/.pyrit/.env.local` | Local overrides (loaded if it exists) | -**AKV behavior** (with `env_akv_ref`): The referenced secret is the lowest-priority file source. Unless custom `env_files` are configured, `~/.pyrit/.env` loads afterward and `~/.pyrit/.env.local` loads last; either may override matching Key Vault values. +**AKV behavior** (with `env_akv_ref`): The referenced secrets load in list order before local files. Unless custom `env_files` are configured, `~/.pyrit/.env` loads afterward and `~/.pyrit/.env.local` loads last. Later bootstrap documents and local files may override earlier values. PyRIT emits a warning when `env_akv_ref` is selected and default or explicit environment files coexist with it. The warning distinguishes files that are ignored from files that load afterward and override Key Vault values, making stale migration files visible at startup. When migrating to Key Vault, clear or remove `~/.pyrit/.env` and `~/.pyrit/.env.local`, remove explicit `env_files` when Key Vault should be the only source, and restart PyRIT so values already present in the process environment cannot mask the Key Vault configuration. -**Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. They override the AKV root when both fields are configured, and default paths are completely ignored. +**Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. They override Key Vault bootstrap values when both fields are configured, and default paths are completely ignored. ### Using .env.local for Overrides @@ -169,11 +169,11 @@ initialization_scripts: Environment file paths to load during initialization. Later files override values from earlier files. -| Value | Behavior | -| ----------------- | -------------------------------------------------------------------- | -| Omitted or `null` | Load default `~/.pyrit/.env` and `~/.pyrit/.env.local`, or only `.env.local` after an AKV root | -| `[]` (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: @@ -181,63 +181,52 @@ env_files: - /path/to/.env.local ``` -Local environment files use standard dotenv parsing and interpolation. `env_akv_strict` does not apply to them: malformed local lines retain python-dotenv's existing permissive skip-and-warn behavior. +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. -During `initialize_pyrit_async`, PyRIT first applies source precedence across the optional Key Vault bootstrap, `.env`, and `.env.local` or explicit `env_files`. It then resolves complete-value `kv:`, `akv:`, `azure_key_vault:`, `env_akv_ref:`, `env:`, and `literal:` references in the winning values, regardless of which source declared them. References overridden by a later source are never fetched. A local file can therefore use a full Key Vault URL even when no bootstrap document is configured. +`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. -An `env:NAME` alias first reads the winning `NAME` value from the merged sources. If no source declares `NAME`, it falls back to the process environment captured before initialization. Merged values take precedence over ambient values with the same name. Alias resolution is one hop. Direct self-reference such as `MODEL="env:MODEL"` is rejected; use a distinct source variable such as `MODEL="env:PYRIT_MODEL"`. - -Interpolation follows 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. - -PyRIT stages the Key Vault mapping and every selected local file before updating `os.environ`. Later files can interpolate and override earlier staged values. If any selected source fails to load or resolve, none of the staged environment values are committed. Memory setup and initializers run after this environment commit and are outside this transaction. +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` -Azure Key Vault secret URL used to obtain the root environment document. Its value must contain dotenv-formatted entries. Authentication uses `DefaultAzureCredential`. +Ordered Azure Key Vault secret URLs used to obtain bootstrap environment documents. Each secret value must contain dotenv-formatted entries. Authentication uses `DefaultAzureCredential`. ```yaml -env_akv_ref: https://my-vault.vault.azure.net/secrets/my-pyrit-env +env_akv_ref: + - https://my-vault.vault.azure.net/secrets/shared-pyrit-env + - https://my-vault.vault.azure.net/secrets/team-pyrit-env ``` -The root document can mix literal values with references to merged or ambient environment variables and scalar secrets in the same vault: +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="env:PYRIT_OPENAI_CHAT_MODEL" +OPENAI_CHAT_MODEL="${PYRIT_OPENAI_CHAT_MODEL}" ``` -Resolution is deliberately limited to two levels: +Resolution is limited to one child-secret lookup: -1. PyRIT fetches the `env_akv_ref` secret and parses it as the bootstrap dotenv document. -2. After all sources are merged, PyRIT either copies one merged-or-ambient `env:` value or fetches one scalar secret from the same vault. The resulting value is final and is not parsed as another reference. +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 secret URL from the bootstrap document's vault. An unversioned URL reads the latest secret version at initialization. Include the version in the URL to pin it. Short names and cross-vault child references are rejected. +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 performs a Key Vault read during initialization, including repeated references to the same URI. +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" ``` -`literal:` is an escape hatch for a bootstrap value that begins with a reserved reference prefix. PyRIT removes `literal:` and returns the remainder without interpreting it as a reference. Quoting does not provide this escape because dotenv removes quotes while parsing. Values fetched from child secrets are already terminal and do not need this escape. - -```dotenv -REFERENCE="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key" -LITERAL_VALUE="literal:kv:not-a-secret-name" -``` - -Here, `REFERENCE` retrieves `openai-chat-key`, while `LITERAL_VALUE` becomes the string `kv:not-a-secret-name`. - -The AKV document is loaded before explicit `env_files` or the default `~/.pyrit/.env.local`, allowing local values to override shared configuration without writing the fetched document to disk. +The bootstrap documents are held in memory and never written to disk. They load before explicit `env_files` or the default `~/.pyrit/.env` and `~/.pyrit/.env.local`, allowing local values to override shared configuration. ### `env_akv_strict` @@ -247,9 +236,9 @@ Controls validation only of the Key Vault bootstrap document and defaults to `tr env_akv_strict: false ``` -In strict mode, any malformed dotenv line or variable without an equals sign stops initialization. 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. +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:` names, unresolved `env:` references, and a bootstrap document with no valid assignments still stop initialization. +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. @@ -298,7 +287,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. The AKV root or environment files are loaded, followed by local overrides +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 @@ -381,8 +370,9 @@ initializers: # - /path/to/.env # - /path/to/.env.local -# Optional Azure Key Vault root environment document -# env_akv_ref: https://my-vault.vault.azure.net/secrets/my-pyrit-env +# Optional ordered Azure Key Vault bootstrap environment documents +# env_akv_ref: +# - https://my-vault.vault.azure.net/secrets/my-pyrit-env # env_akv_strict: false # Optional; defaults to true # Suppress initialization messages From f3f7ebb9117298b55a60351c53cad450e5e6749f Mon Sep 17 00:00:00 2001 From: Victor Valbuena <50061128+ValbuenaVC@users.noreply.github.com> Date: Thu, 13 Aug 2026 15:16:35 -0400 Subject: [PATCH 10/16] Update doc/getting_started/pyrit_conf.md Co-authored-by: Justin Song --- doc/getting_started/pyrit_conf.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 283b614406..2cdb0bdd13 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -53,7 +53,7 @@ System environment variables are always the baseline. If no AKV bootstrap docume **AKV behavior** (with `env_akv_ref`): The referenced secrets load in list order before local files. Unless custom `env_files` are configured, `~/.pyrit/.env` loads afterward and `~/.pyrit/.env.local` loads last. Later bootstrap documents and local files may override earlier values. -PyRIT emits a warning when `env_akv_ref` is selected and default or explicit environment files coexist with it. The warning distinguishes files that are ignored from files that load afterward and override Key Vault values, making stale migration files visible at startup. When migrating to Key Vault, clear or remove `~/.pyrit/.env` and `~/.pyrit/.env.local`, remove explicit `env_files` when Key Vault should be the only source, and restart PyRIT so values already present in the process environment cannot mask the Key Vault configuration. +PyRIT emits a warning when `env_akv_ref` is selected and default or explicit environment files coexist with it. The warning distinguishes files that are ignored from files that load afterward and override Key Vault values, making stale migration files visible at startup. When migrating to Key Vault, clear or remove `~/.pyrit/.env` and `~/.pyrit/.env.local`, remove explicit `env_files` when Key Vault should be the only source, and re-initialize PyRIT so values already present in the process environment cannot mask the Key Vault configuration. **Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. They override Key Vault bootstrap values when both fields are configured, and default paths are completely ignored. From dc5e133405c70d9bd73accf5abad0884102e2ebe Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 13 Aug 2026 15:18:05 -0400 Subject: [PATCH 11/16] FIX: docs consistency --- .env_example | 200 ++++++++++++++++++++++++++++++-------------- .pyrit_conf_example | 24 +++--- 2 files changed, 148 insertions(+), 76 deletions(-) diff --git a/.env_example b/.env_example index 0d70a48cfd..039544d8a4 100644 --- a/.env_example +++ b/.env_example @@ -1,138 +1,176 @@ # ============================================================================ + # PyRIT Environment File Example + # ============================================================================ + # -# Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need. + +# Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need + # -# MOST USERS ONLY NEED 3 VARIABLES to get started: + +# MOST USERS ONLY NEED 3 VARIABLES to get started + # -# OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" # or any OpenAI-compatible API -# OPENAI_CHAT_KEY="your-key-here" -# OPENAI_CHAT_MODEL="gpt-4o" + +# OPENAI_CHAT_ENDPOINT="" # or any OpenAI-compatible API + +# OPENAI_CHAT_KEY="your-key-here" + +# OPENAI_CHAT_MODEL="gpt-4o" + # + # These work with OpenAI, Azure OpenAI, Ollama, Groq, OpenRouter, and any + # other OpenAI-compatible endpoint. See doc/setup/populating_secrets.md -# for provider-specific examples. + +# for provider-specific examples + # -# If you are using Entra authentication for Azure resources, + +# If you are using Entra authentication for Azure resources + # keys for those resources are not needed. PyRIT auto-detects: if an API key -# is set, it uses key auth; otherwise it falls back to Entra ID automatically. + +# is set, it uses key auth; otherwise it falls back to Entra ID automatically + # -# ============================================================================ +# ============================================================================ ################################### + # OPENAI TARGET SECRETS + # + # The below models work with OpenAIChatTarget - either pass via environment variables + # or copy to OPENAI_CHAT_ENDPOINT + ################################### -PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_CHAT_ENDPOINT="" PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" PLATFORM_OPENAI_CHAT_MODEL="gpt-4o" # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately -# Example: https://xxxx.openai.azure.com/openai/v1 -AZURE_OPENAI_GPT4O_ENDPOINT="https://xxxx.openai.azure.com/openai/v1" + +# Example: + +AZURE_OPENAI_GPT4O_ENDPOINT="" AZURE_OPENAI_GPT4O_KEY="xxxxx" AZURE_OPENAI_GPT4O_MODEL="deployment-name" -# Since Azure deployment name may be custom and differ from the actual underlying model, -# you can specify the underlying model for identifier purposes. If not specified, -# identifiers will default to the value of the standard MODEL environment variable. + +# Since Azure deployment name may be custom and differ from the actual underlying model + +# you can specify the underlying model for identifier purposes. If not specified + +# identifiers will default to the value of the standard MODEL environment variable + AZURE_OPENAI_GPT4O_UNDERLYING_MODEL="gpt-4o" -# Optional second GPT-4o endpoint (that can be used for round-robin distribution). +# Optional second GPT-4o endpoint (that can be used for round-robin distribution) + # TargetInitializer creates RoundRobinTargets that automatically group together + # targets with identical underlying model names and behavioral params, allowing -# for distribution of requests across them for rate-limit relief. -AZURE_OPENAI_GPT4O_ENDPOINT2="https://xxxx.openai.azure.com/openai/v1" + +# for distribution of requests across them for rate-limit relief + +AZURE_OPENAI_GPT4O_ENDPOINT2="" AZURE_OPENAI_GPT4O_KEY2="xxxxx" AZURE_OPENAI_GPT4O_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNDERLYING_MODEL2="gpt-4o" -AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="" AZURE_OPENAI_INTEGRATION_TEST_KEY="xxxxx" AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="" AZURE_OPENAI_GPT3_5_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT3_5_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT3_5_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT4_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT4_CHAT_ENDPOINT="" AZURE_OPENAI_GPT4_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT5_4_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT5_4_ENDPOINT="" AZURE_OPENAI_GPT5_4_KEY="xxxxx" AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" # Endpoints that host models with fewer safety mechanisms (e.g. via adversarial fine tuning -# or content filters turned off) can be defined below and used in adversarial attack testing scenarios. -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" + +# or content filters turned off) can be defined below and used in adversarial attack testing scenarios + +AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2="" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" # Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) + # Default endpoint goes here; specialized ones below -ADVERSARIAL_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" + +ADVERSARIAL_CHAT_ENDPOINT="" ADVERSARIAL_CHAT_KEY="xxxxx" ADVERSARIAL_CHAT_MODEL="deployment-name" -ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="" ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" ADVERSARIAL_CHAT_SINGLETURN_MODEL="deployment-name" -ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="" ADVERSARIAL_CHAT_MULTITURN_KEY="xxxxx" ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" -ADVERSARIAL_CHAT_REASONING_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_REASONING_ENDPOINT="" ADVERSARIAL_CHAT_REASONING_KEY="xxxxx" ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" - # Objective Scorer chat target (used in scorers in scenarios) -OBJECTIVE_SCORER_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" + +OBJECTIVE_SCORER_CHAT_ENDPOINT="" OBJECTIVE_SCORER_CHAT_KEY="xxxxx" OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" -AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="https://xxxxx.eastus2.models.ai.azure.com" +AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="" AZURE_FOUNDRY_DEEPSEEK_KEY="xxxxx" AZURE_FOUNDRY_DEEPSEEK_MODEL="" -AZURE_FOUNDRY_PHI4_ENDPOINT="https://xxxxx.models.ai.azure.com" +AZURE_FOUNDRY_PHI4_ENDPOINT="" AZURE_CHAT_PHI4_KEY="xxxxx" AZURE_CHAT_PHI4_MODEL="" -AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="https://xxxxx.services.ai.azure.com/openai/v1/" +AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="" AZURE_FOUNDRY_MISTRAL_LARGE_KEY="xxxxx" AZURE_FOUNDRY_MISTRAL_LARGE_MODEL="Mistral-Large-3" -AWS_ENDPOINT="https://bedrock-mantle.us-east-1.api.aws/v1" +AWS_ENDPOINT="" AWS_KEY="xxxxx" AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" AWS_RESPONSES_MODEL="openai.gpt-oss-120b" -GROQ_ENDPOINT="https://api.groq.com/openai/v1" +GROQ_ENDPOINT="" GROQ_KEY="gsk_xxxxxxxx" GROQ_LLAMA_MODEL="llama3-8b-8192" -OPEN_ROUTER_ENDPOINT="https://openrouter.ai/api/v1" +OPEN_ROUTER_ENDPOINT="" OPEN_ROUTER_KEY="sk-or-v1-xxxxx" OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" -OLLAMA_CHAT_ENDPOINT="http://127.0.0.1:11434/v1" +OLLAMA_CHAT_ENDPOINT="" OLLAMA_MODEL="llama2" DEFAULT_OPENAI_FRONTEND_ENDPOINT = ${AZURE_OPENAI_GPT4O_AAD_ENDPOINT} @@ -142,25 +180,30 @@ DEFAULT_OPENAI_FRONTEND_MODEL = "gpt-4o" OPENAI_CHAT_ENDPOINT=${PLATFORM_OPENAI_CHAT_ENDPOINT} OPENAI_CHAT_KEY=${PLATFORM_OPENAI_CHAT_KEY} OPENAI_CHAT_MODEL=${PLATFORM_OPENAI_CHAT_MODEL} + # The following line can be populated if using an Azure OpenAI deployment + # where the deployment name differs from the actual underlying model + OPENAI_CHAT_UNDERLYING_MODEL="" ################################## + # OPENAI RESPONSES TARGET SECRETS + ################################## -AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" -AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" +AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="" +AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="" AZURE_OPENAI_GPT5_KEY="xxxxxxx" AZURE_OPENAI_GPT5_MODEL="gpt-5" AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" -PLATFORM_OPENAI_RESPONSES_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_RESPONSES_ENDPOINT="" PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" -AZURE_OPENAI_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_RESPONSES_ENDPOINT="" AZURE_OPENAI_RESPONSES_KEY="xxxxx" AZURE_OPENAI_RESPONSES_MODEL="o4-mini" AZURE_OPENAI_RESPONSES_UNDERLYING_MODEL="o4-mini" @@ -171,10 +214,15 @@ OPENAI_RESPONSES_MODEL=${PLATFORM_OPENAI_RESPONSES_MODEL} OPENAI_RESPONSES_UNDERLYING_MODEL="" ################################## + # OPENAI REALTIME TARGET SECRETS + # + # The below models work with RealtimeTarget - either pass via environment variables + # or copy to OPENAI_REALTIME_ENDPOINT + ################################## PLATFORM_OPENAI_REALTIME_ENDPOINT="wss://api.openai.com/v1" @@ -192,18 +240,23 @@ OPENAI_REALTIME_MODEL = ${PLATFORM_OPENAI_REALTIME_MODEL} OPENAI_REALTIME_UNDERLYING_MODEL = "" ################################## + # IMAGE TARGET SECRETS + # + # The below models work with OpenAIImageTarget - either pass via environment variables + # or copy to OPENAI_IMAGE_ENDPOINT + ################################### -OPENAI_IMAGE_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" +OPENAI_IMAGE_ENDPOINT1 = "" OPENAI_IMAGE_API_KEY1 = "xxxxxx" OPENAI_IMAGE_MODEL1 = "deployment-name" OPENAI_IMAGE_UNDERLYING_MODEL1 = "dall-e-3" -OPENAI_IMAGE_ENDPOINT2 = "https://api.openai.com/v1" +OPENAI_IMAGE_ENDPOINT2 = "" OPENAI_IMAGE_API_KEY2 = "sk-xxxxx" OPENAI_IMAGE_MODEL2 = "dall-e-3" OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" @@ -213,20 +266,24 @@ OPENAI_IMAGE_API_KEY = ${OPENAI_IMAGE_API_KEY2} OPENAI_IMAGE_MODEL = ${OPENAI_IMAGE_MODEL2} OPENAI_IMAGE_UNDERLYING_MODEL = "" - ################################## + # TTS TARGET SECRETS + # + # The below models work with OpenAITTSTarget - either pass via environment variables + # or copy to OPENAI_TTS_ENDPOINT + ################################### -OPENAI_TTS_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" +OPENAI_TTS_ENDPOINT1 = "" OPENAI_TTS_KEY1 = "xxxxxxx" OPENAI_TTS_MODEL1 = "tts" OPENAI_TTS_UNDERLYING_MODEL1 = "tts" -OPENAI_TTS_ENDPOINT2 = "https://api.openai.com/v1" +OPENAI_TTS_ENDPOINT2 = "" OPENAI_TTS_KEY2 = "xxxxxx" OPENAI_TTS_MODEL2 = "tts-1" OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" @@ -237,14 +294,20 @@ OPENAI_TTS_MODEL = ${OPENAI_TTS_MODEL2} OPENAI_TTS_UNDERLYING_MODEL = "" ################################## + # VIDEO TARGET SECRETS + # + # The below models work with OpenAIVideoTarget - either pass via environment variables + # or copy to OPENAI_VIDEO_ENDPOINT + ################################### # Note: Use the base URL without API path -AZURE_OPENAI_VIDEO_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/openai/v1" + +AZURE_OPENAI_VIDEO_ENDPOINT="" AZURE_OPENAI_VIDEO_KEY="xxxxxxx" AZURE_OPENAI_VIDEO_MODEL="sora-2" AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" @@ -254,68 +317,75 @@ OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" - ################################## + # AML TARGET SECRETS + # The below models work with AzureMLChatTarget - either pass via environment variables + # or copy to AZURE_ML_MANAGED_ENDPOINT + ################################### -AZURE_ML_PHI_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +AZURE_ML_PHI_ENDPOINT="" AZURE_ML_PHI_KEY="xxxxx" -# The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed. +# The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed + AZURE_ML_MANAGED_ENDPOINT=${AZURE_ML_PHI_ENDPOINT} AZURE_ML_KEY=${AZURE_ML_PHI_KEY} - ################################## + # MISC TARGET SECRETS -################################### +################################### -OPENAI_COMPLETION_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +OPENAI_COMPLETION_ENDPOINT="" OPENAI_COMPLETION_API_KEY="xxxxx" OPENAI_COMPLETION_MODEL="davinci-002" -OPENAI_EMBEDDING_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +OPENAI_EMBEDDING_ENDPOINT="" OPENAI_EMBEDDING_KEY="xxxxx" OPENAI_EMBEDDING_MODEL="text-embedding-3-small" -AZURE_STORAGE_ACCOUNT_CONTAINER_URL="https://xxxxxx.blob.core.windows.net/xpia" +AZURE_STORAGE_ACCOUNT_CONTAINER_URL="" AZURE_STORAGE_ACCOUNT_SAS_TOKEN="xxxxx" - AZURE_SPEECH_REGION = "eastus2" AZURE_SPEECH_KEY = "xxxxx" + # Resource ID is needed when using Entra authentication + AZURE_SPEECH_RESOURCE_ID = "xxxxx" AZURE_CONTENT_SAFETY_API_KEY="xxxxx" -AZURE_CONTENT_SAFETY_API_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/" +AZURE_CONTENT_SAFETY_API_ENDPOINT="" HUGGINGFACE_TOKEN="hf_xxxxxxx" -HUGGINGFACE_ENDPOINT="https://router.huggingface.co/v1" +HUGGINGFACE_ENDPOINT="" -GOOGLE_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/openai" +GOOGLE_GEMINI_ENDPOINT = "" GOOGLE_GEMINI_API_KEY = "xxxxx" GOOGLE_GEMINI_MODEL="gemini-2.0-flash" - ######################### + # AZURE SQL SECRETS -######################### +######################### # This connects to the test database + AZURE_SQL_DB_CONNECTION_STRING_TEST = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" -AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="https://xxxxx.blob.core.windows.net/dbdata" +AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="" # This connects to the prod database + AZURE_SQL_DB_CONNECTION_STRING_PROD = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" -AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="https://xxxxx.blob.core.windows.net/dbdata" +AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="" +# The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local -# The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local. AZURE_SQL_DB_CONNECTION_STRING = ${AZURE_SQL_DB_CONNECTION_STRING_PROD} AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL=${AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD} diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 3ed13b71b9..5dc456f25d 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -88,13 +88,12 @@ operation: op_trash_panda # - 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 -# - Local files retain standard dotenv parsing and interpolation. After source -# precedence is applied, PyRIT resolves complete-value kv:/env: references in -# winning values from any source. Overridden references are not fetched. +# - Local files retain standard dotenv parsing and ${NAME} interpolation. +# Key Vault references in local files remain literal. # - Interpolation follows load order: .env.local can reference .env, but .env # cannot see variables introduced only by the later .env.local. -# - During PyRIT initialization, selected environment sources are staged and -# committed together only after every source loads successfully. +# - Loading is non-transactional. If a later source fails, values loaded by +# earlier sources remain in the process environment. # # Example: # env_files: @@ -103,18 +102,20 @@ operation: op_trash_panda # Azure Key Vault Environment References # --------------------------------------- -# AKV secret URL whose value is the bootstrap .env document. -# Winning values may reference another merged environment key with env:NAME, -# falling back to the existing process environment, or reference a scalar -# secret in the same vault using a full URL: +# Ordered AKV secret URLs whose values are bootstrap .env documents. +# Documents load in list order before local files and use standard ${NAME} +# interpolation. Complete values in a bootstrap document may reference a +# scalar secret in that document's vault using a full URL: # kv:https://my-vault.vault.azure.net/secrets/SECRET_NAME # Include a version to pin a secret: # kv:https://my-vault.vault.azure.net/secrets/SECRET_NAME/SECRET_VERSION # Short secret names such as kv:SECRET_NAME are rejected. # Cross-vault child references are rejected. +# Only .vault.azure.net, .vault.azure.cn, and .vault.usgovcloudapi.net hosts +# are accepted. Arbitrary HTTPS hosts and malformed secret paths are rejected. # Referenced secrets are not cached; each kv: occurrence performs a vault read. # Referenced values are terminal scalars; they are not parsed for more references. -# Source precedence is AKV bootstrap -> ~/.pyrit/.env -> ~/.pyrit/.env.local. +# Source precedence is AKV bootstraps -> ~/.pyrit/.env -> ~/.pyrit/.env.local. # Explicit env_files replace the default files and load after the AKV bootstrap. # PyRIT emits a warning when these local files coexist with env_akv_ref so stale # configuration cannot silently mask or be mistaken for the Key Vault document. @@ -129,7 +130,8 @@ operation: op_trash_panda # Requires: pip install azure-keyvault-secrets # # Example: -# env_akv_ref: https://my-vault.vault.azure.net/secrets/my-pyrit-env +# env_akv_ref: +# - https://my-vault.vault.azure.net/secrets/my-pyrit-env # # Strict validation applies only to the Key Vault bootstrap and is enabled by # default. Set this to false to skip malformed or valueless bootstrap entries From 9fe85a9c5b2144b03e9889b5793c25f0d532ed40 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 14 Aug 2026 12:24:40 -0400 Subject: [PATCH 12/16] FEAT: Fixing .env_example drift --- .env_example | 52 +++++++++++-------- doc/getting_started/pyrit_conf.md | 2 +- infra/env.demo.template | 16 +++--- pyrit/setup/initializers/targets.py | 32 ++++++------ .../targets/test_targets_and_secrets.py | 46 ++++++++-------- tests/unit/setup/test_targets_initializer.py | 6 +-- 6 files changed, 81 insertions(+), 73 deletions(-) diff --git a/.env_example b/.env_example index 039544d8a4..9870c6eae8 100644 --- a/.env_example +++ b/.env_example @@ -249,23 +249,29 @@ OPENAI_REALTIME_UNDERLYING_MODEL = "" # or copy to OPENAI_IMAGE_ENDPOINT +# Entra auth should be enabled + ################################### -OPENAI_IMAGE_ENDPOINT1 = "" -OPENAI_IMAGE_API_KEY1 = "xxxxxx" -OPENAI_IMAGE_MODEL1 = "deployment-name" -OPENAI_IMAGE_UNDERLYING_MODEL1 = "dall-e-3" +AZURE_OPENAI_IMAGE_ENDPOINT1 = "" +AZURE_OPENAI_IMAGE_API_KEY1 = "xxxxxx" +AZURE_OPENAI_IMAGE_MODEL1 = "deployment-name" +AZURE_OPENAI_IMAGE_UNDERLYING_MODEL1 = "dall-e-3" -OPENAI_IMAGE_ENDPOINT2 = "" -OPENAI_IMAGE_API_KEY2 = "sk-xxxxx" -OPENAI_IMAGE_MODEL2 = "dall-e-3" -OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" +AZURE_OPENAI_IMAGE_ENDPOINT2 = "" +AZURE_OPENAI_IMAGE_API_KEY2 = "xxxxxx" +AZURE_OPENAI_IMAGE_MODEL2 = "dall-e-3" +AZURE_OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" -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_IMAGE_UNDERLYING_MODEL = "" +OPENAI_IMAGE_STRICT_FILTER_ENDPOINT = "" +OPENAI_IMAGE_STRICT_FILTER_MODEL = "gpt-image" +OPENAI_IMAGE_STRICT_FILTER_UNDERLYING_MODEL = "gpt-image" + ################################## # TTS TARGET SECRETS @@ -276,21 +282,23 @@ OPENAI_IMAGE_UNDERLYING_MODEL = "" # or copy to OPENAI_TTS_ENDPOINT +# Entra auth should be enabled + ################################### -OPENAI_TTS_ENDPOINT1 = "" -OPENAI_TTS_KEY1 = "xxxxxxx" -OPENAI_TTS_MODEL1 = "tts" -OPENAI_TTS_UNDERLYING_MODEL1 = "tts" +AZURE_OPENAI_TTS_ENDPOINT1 = "" +AZURE_OPENAI_TTS_KEY1 = "xxxxxxx" +AZURE_OPENAI_TTS_MODEL1 = "tts" +AZURE_OPENAI_TTS_UNDERLYING_MODEL1 = "tts" -OPENAI_TTS_ENDPOINT2 = "" -OPENAI_TTS_KEY2 = "xxxxxx" -OPENAI_TTS_MODEL2 = "tts-1" -OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" +AZURE_OPENAI_TTS_ENDPOINT2 = "" +AZURE_OPENAI_TTS_KEY2 = "xxxxxx" +AZURE_OPENAI_TTS_MODEL2 = "tts-1" +AZURE_OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" -OPENAI_TTS_ENDPOINT = ${OPENAI_TTS_ENDPOINT2} -OPENAI_TTS_KEY = ${OPENAI_TTS_KEY2} -OPENAI_TTS_MODEL = ${OPENAI_TTS_MODEL2} +OPENAI_TTS_ENDPOINT = ${AZURE_OPENAI_TTS_ENDPOINT2} +OPENAI_TTS_KEY = ${AZURE_OPENAI_TTS_KEY2} +OPENAI_TTS_MODEL = ${AZURE_OPENAI_TTS_MODEL2} OPENAI_TTS_UNDERLYING_MODEL = "" ################################## diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 2cdb0bdd13..89b6489be4 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -10,7 +10,7 @@ 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 diff --git a/infra/env.demo.template b/infra/env.demo.template index 4c5e7dac7e..77f2af2b56 100644 --- a/infra/env.demo.template +++ b/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 diff --git a/pyrit/setup/initializers/targets.py b/pyrit/setup/initializers/targets.py index 308366f734..790b51a87f 100644 --- a/pyrit/setup/initializers/targets.py +++ b/pyrit/setup/initializers/targets.py @@ -338,18 +338,18 @@ class TargetConfig: TargetConfig( registry_name="openai_image_azure", target_class=OpenAIImageTarget, - endpoint_var="OPENAI_IMAGE_ENDPOINT1", - key_var="OPENAI_IMAGE_API_KEY1", - model_var="OPENAI_IMAGE_MODEL1", - underlying_model_var="OPENAI_IMAGE_UNDERLYING_MODEL1", + endpoint_var="AZURE_OPENAI_IMAGE_ENDPOINT1", + key_var="AZURE_OPENAI_IMAGE_API_KEY1", + model_var="AZURE_OPENAI_IMAGE_MODEL1", + underlying_model_var="AZURE_OPENAI_IMAGE_UNDERLYING_MODEL1", ), TargetConfig( registry_name="openai_image_platform", target_class=OpenAIImageTarget, - endpoint_var="OPENAI_IMAGE_ENDPOINT2", - key_var="OPENAI_IMAGE_API_KEY2", - model_var="OPENAI_IMAGE_MODEL2", - underlying_model_var="OPENAI_IMAGE_UNDERLYING_MODEL2", + endpoint_var="AZURE_OPENAI_IMAGE_ENDPOINT2", + key_var="AZURE_OPENAI_IMAGE_API_KEY2", + model_var="AZURE_OPENAI_IMAGE_MODEL2", + underlying_model_var="AZURE_OPENAI_IMAGE_UNDERLYING_MODEL2", ), # ============================================ # TTS Targets (OpenAITTSTarget) @@ -357,18 +357,18 @@ class TargetConfig: TargetConfig( registry_name="openai_tts_azure", target_class=OpenAITTSTarget, - endpoint_var="OPENAI_TTS_ENDPOINT1", - key_var="OPENAI_TTS_KEY1", - model_var="OPENAI_TTS_MODEL1", - underlying_model_var="OPENAI_TTS_UNDERLYING_MODEL1", + endpoint_var="AZURE_OPENAI_TTS_ENDPOINT1", + key_var="AZURE_OPENAI_TTS_KEY1", + model_var="AZURE_OPENAI_TTS_MODEL1", + underlying_model_var="AZURE_OPENAI_TTS_UNDERLYING_MODEL1", ), TargetConfig( registry_name="openai_tts_platform", target_class=OpenAITTSTarget, - endpoint_var="OPENAI_TTS_ENDPOINT2", - key_var="OPENAI_TTS_KEY2", - model_var="OPENAI_TTS_MODEL2", - underlying_model_var="OPENAI_TTS_UNDERLYING_MODEL2", + endpoint_var="AZURE_OPENAI_TTS_ENDPOINT2", + key_var="AZURE_OPENAI_TTS_KEY2", + model_var="AZURE_OPENAI_TTS_MODEL2", + underlying_model_var="AZURE_OPENAI_TTS_UNDERLYING_MODEL2", ), # ============================================ # Video Targets (OpenAIVideoTarget) diff --git a/tests/integration/targets/test_targets_and_secrets.py b/tests/integration/targets/test_targets_and_secrets.py index e2ec9da733..2a15ae6397 100644 --- a/tests/integration/targets/test_targets_and_secrets.py +++ b/tests/integration/targets/test_targets_and_secrets.py @@ -561,23 +561,23 @@ async def test_connect_openai_completion(sqlite_instance: SQLiteMemory) -> None: [ ("OPENAI_IMAGE_ENDPOINT", None, "OPENAI_IMAGE_MODEL"), pytest.param( - "OPENAI_IMAGE_ENDPOINT1", + "AZURE_OPENAI_IMAGE_ENDPOINT1", None, - "OPENAI_IMAGE_MODEL1", + "AZURE_OPENAI_IMAGE_MODEL1", marks=pytest.mark.run_only_if_all_tests, ), # gpt-image-1.5 pytest.param( - "OPENAI_IMAGE_ENDPOINT1", - "OPENAI_IMAGE_API_KEY1", - "OPENAI_IMAGE_MODEL1", + "AZURE_OPENAI_IMAGE_ENDPOINT1", + "AZURE_OPENAI_IMAGE_API_KEY1", + "AZURE_OPENAI_IMAGE_MODEL1", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-image1-api-key", ), - ("OPENAI_IMAGE_ENDPOINT2", None, "OPENAI_IMAGE_MODEL2"), # gpt-image-1 + ("AZURE_OPENAI_IMAGE_ENDPOINT2", None, "AZURE_OPENAI_IMAGE_MODEL2"), # gpt-image-1 pytest.param( - "OPENAI_IMAGE_ENDPOINT2", - "OPENAI_IMAGE_API_KEY2", - "OPENAI_IMAGE_MODEL2", + "AZURE_OPENAI_IMAGE_ENDPOINT2", + "AZURE_OPENAI_IMAGE_API_KEY2", + "AZURE_OPENAI_IMAGE_MODEL2", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-image2-api-key", ), @@ -626,7 +626,7 @@ async def test_connect_image( [ pytest.param(None, id="entra"), pytest.param( - "OPENAI_IMAGE_API_KEY2", + "AZURE_OPENAI_IMAGE_API_KEY2", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="api-key", ), @@ -645,8 +645,8 @@ async def test_image_editing_single_image( 2. The edit endpoint is correctly called 3. The output image file is created """ - endpoint_value = _get_required_env_var("OPENAI_IMAGE_ENDPOINT2") - model_name_value = os.getenv("OPENAI_IMAGE_MODEL2") or "gpt-image-1" + endpoint_value = _get_required_env_var("AZURE_OPENAI_IMAGE_ENDPOINT2") + model_name_value = os.getenv("AZURE_OPENAI_IMAGE_MODEL2") or "gpt-image-1" target = OpenAIImageTarget( endpoint=endpoint_value, @@ -686,7 +686,7 @@ async def test_image_editing_single_image( [ pytest.param(None, id="entra"), pytest.param( - "OPENAI_IMAGE_API_KEY2", + "AZURE_OPENAI_IMAGE_API_KEY2", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="api-key", ), @@ -704,8 +704,8 @@ async def test_image_editing_multiple_images( 1. Multiple images can be passed to the edit endpoint 2. The model processes multiple image inputs correctly """ - endpoint_value = _get_required_env_var("OPENAI_IMAGE_ENDPOINT2") - model_name_value = os.getenv("OPENAI_IMAGE_MODEL2") or "gpt-image-1" + endpoint_value = _get_required_env_var("AZURE_OPENAI_IMAGE_ENDPOINT2") + model_name_value = os.getenv("AZURE_OPENAI_IMAGE_MODEL2") or "gpt-image-1" target = OpenAIImageTarget( endpoint=endpoint_value, @@ -749,19 +749,19 @@ async def test_image_editing_multiple_images( @pytest.mark.parametrize( ("endpoint", "api_key_env_var", "model_name"), [ - ("OPENAI_TTS_ENDPOINT1", None, "OPENAI_TTS_MODEL1"), + ("AZURE_OPENAI_TTS_ENDPOINT1", None, "AZURE_OPENAI_TTS_MODEL1"), pytest.param( - "OPENAI_TTS_ENDPOINT1", - "OPENAI_TTS_KEY1", - "OPENAI_TTS_MODEL1", + "AZURE_OPENAI_TTS_ENDPOINT1", + "AZURE_OPENAI_TTS_KEY1", + "AZURE_OPENAI_TTS_MODEL1", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-tts1-api-key", ), - ("OPENAI_TTS_ENDPOINT2", None, "OPENAI_TTS_MODEL2"), + ("AZURE_OPENAI_TTS_ENDPOINT2", None, "AZURE_OPENAI_TTS_MODEL2"), pytest.param( - "OPENAI_TTS_ENDPOINT2", - "OPENAI_TTS_KEY2", - "OPENAI_TTS_MODEL2", + "AZURE_OPENAI_TTS_ENDPOINT2", + "AZURE_OPENAI_TTS_KEY2", + "AZURE_OPENAI_TTS_MODEL2", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-tts2-api-key", ), diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index 52141904e1..06c7d4c4d1 100644 --- a/tests/unit/setup/test_targets_initializer.py +++ b/tests/unit/setup/test_targets_initializer.py @@ -104,9 +104,9 @@ async def test_registers_multiple_targets(self): os.environ["PLATFORM_OPENAI_CHAT_MODEL"] = "gpt-4o" # Set up openai_image_platform (uses ENDPOINT2/KEY2/MODEL2) - os.environ["OPENAI_IMAGE_ENDPOINT2"] = "https://api.openai.com/v1" - os.environ["OPENAI_IMAGE_API_KEY2"] = "test_image_key" - os.environ["OPENAI_IMAGE_MODEL2"] = "dall-e-3" + os.environ["AZURE_OPENAI_IMAGE_ENDPOINT2"] = "https://api.openai.com/v1" + os.environ["AZURE_OPENAI_IMAGE_API_KEY2"] = "test_image_key" + os.environ["AZURE_OPENAI_IMAGE_MODEL2"] = "dall-e-3" init = TargetInitializer() await init.initialize_async() From 2063af3fba13c4d7f78f862303f07f3daf64ae26 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 14 Aug 2026 13:45:23 -0400 Subject: [PATCH 13/16] FIX: Remove extra newlines from .env_example --- .env_example | 31 ------------------------------- 1 file changed, 31 deletions(-) diff --git a/.env_example b/.env_example index 9870c6eae8..6940f40c13 100644 --- a/.env_example +++ b/.env_example @@ -1,53 +1,30 @@ # ============================================================================ - # PyRIT Environment File Example - # ============================================================================ - # - # Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need - # - # MOST USERS ONLY NEED 3 VARIABLES to get started - # - # OPENAI_CHAT_ENDPOINT="" # or any OpenAI-compatible API - # OPENAI_CHAT_KEY="your-key-here" - # OPENAI_CHAT_MODEL="gpt-4o" - # - # These work with OpenAI, Azure OpenAI, Ollama, Groq, OpenRouter, and any - # other OpenAI-compatible endpoint. See doc/setup/populating_secrets.md - # for provider-specific examples - # - # If you are using Entra authentication for Azure resources - # keys for those resources are not needed. PyRIT auto-detects: if an API key - # is set, it uses key auth; otherwise it falls back to Entra ID automatically - # - # ============================================================================ ################################### # OPENAI TARGET SECRETS - # - # The below models work with OpenAIChatTarget - either pass via environment variables - # or copy to OPENAI_CHAT_ENDPOINT ################################### @@ -57,7 +34,6 @@ PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" PLATFORM_OPENAI_CHAT_MODEL="gpt-4o" # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately - # Example: AZURE_OPENAI_GPT4O_ENDPOINT="" @@ -65,19 +41,14 @@ AZURE_OPENAI_GPT4O_KEY="xxxxx" AZURE_OPENAI_GPT4O_MODEL="deployment-name" # Since Azure deployment name may be custom and differ from the actual underlying model - # you can specify the underlying model for identifier purposes. If not specified - # identifiers will default to the value of the standard MODEL environment variable AZURE_OPENAI_GPT4O_UNDERLYING_MODEL="gpt-4o" # Optional second GPT-4o endpoint (that can be used for round-robin distribution) - # TargetInitializer creates RoundRobinTargets that automatically group together - # targets with identical underlying model names and behavioral params, allowing - # for distribution of requests across them for rate-limit relief AZURE_OPENAI_GPT4O_ENDPOINT2="" @@ -106,7 +77,6 @@ AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" # Endpoints that host models with fewer safety mechanisms (e.g. via adversarial fine tuning - # or content filters turned off) can be defined below and used in adversarial attack testing scenarios AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="" @@ -120,7 +90,6 @@ AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" # Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) - # Default endpoint goes here; specialized ones below ADVERSARIAL_CHAT_ENDPOINT="" From e8612b2be42b91c9274b9c2fd461f1cff706d694 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 14 Aug 2026 17:28:19 -0400 Subject: [PATCH 14/16] FEAT: Addressing latest PR comments --- .env_example | 428 ++++++---- .pyrit_conf_example | 68 +- doc/code/executor/gcg/1_gcg_azure_ml.ipynb | 2 +- doc/code/executor/gcg/1_gcg_azure_ml.py | 2 +- doc/getting_started/pyrit_conf.md | 80 +- .../executor/promptgen/gcg/experiments/run.py | 2 +- pyrit/setup/akv_initialization.py | 554 +++++++++++++ pyrit/setup/configuration_loader.py | 9 + pyrit/setup/initialization.py | 500 +----------- .../promptgen/gcg/test_gcg_aml_e2e.py | 2 +- .../test_akv_initialization_integration.py | 51 ++ tests/unit/setup/test_configuration_loader.py | 20 +- tests/unit/setup/test_initialization.py | 743 +----------------- 13 files changed, 974 insertions(+), 1487 deletions(-) create mode 100644 pyrit/setup/akv_initialization.py create mode 100644 tests/integration/setup/test_akv_initialization_integration.py diff --git a/.env_example b/.env_example index 6940f40c13..ac0005e69b 100644 --- a/.env_example +++ b/.env_example @@ -1,153 +1,175 @@ # ============================================================================ + # PyRIT Environment File Example + # ============================================================================ + # + # Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need + # + # MOST USERS ONLY NEED 3 VARIABLES to get started + # + # OPENAI_CHAT_ENDPOINT="" # or any OpenAI-compatible API + # OPENAI_CHAT_KEY="your-key-here" + # OPENAI_CHAT_MODEL="gpt-4o" + # + # These work with OpenAI, Azure OpenAI, Ollama, Groq, OpenRouter, and any + # other OpenAI-compatible endpoint. See doc/setup/populating_secrets.md + # for provider-specific examples + # + # If you are using Entra authentication for Azure resources + # keys for those resources are not needed. PyRIT auto-detects: if an API key + # is set, it uses key auth; otherwise it falls back to Entra ID automatically + # + # ============================================================================ -################################### +################################## # OPENAI TARGET SECRETS + +################################## + # -# The below models work with OpenAIChatTarget - either pass via environment variables -# or copy to OPENAI_CHAT_ENDPOINT -################################### +# The below models work with OpenAIChatTarget - either pass via environment variables -PLATFORM_OPENAI_CHAT_ENDPOINT="" -PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" -PLATFORM_OPENAI_CHAT_MODEL="gpt-4o" +# or copy to OPENAI_CHAT_ENDPOINT # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately + # Example: AZURE_OPENAI_GPT4O_ENDPOINT="" -AZURE_OPENAI_GPT4O_KEY="xxxxx" AZURE_OPENAI_GPT4O_MODEL="deployment-name" # Since Azure deployment name may be custom and differ from the actual underlying model + # you can specify the underlying model for identifier purposes. If not specified + # identifiers will default to the value of the standard MODEL environment variable AZURE_OPENAI_GPT4O_UNDERLYING_MODEL="gpt-4o" # Optional second GPT-4o endpoint (that can be used for round-robin distribution) + # TargetInitializer creates RoundRobinTargets that automatically group together + # targets with identical underlying model names and behavioral params, allowing + # for distribution of requests across them for rate-limit relief AZURE_OPENAI_GPT4O_ENDPOINT2="" -AZURE_OPENAI_GPT4O_KEY2="xxxxx" AZURE_OPENAI_GPT4O_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNDERLYING_MODEL2="gpt-4o" -AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="" -AZURE_OPENAI_INTEGRATION_TEST_KEY="xxxxx" -AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" -AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" - -AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="" -AZURE_OPENAI_GPT3_5_CHAT_KEY="xxxxx" -AZURE_OPENAI_GPT3_5_CHAT_MODEL="deployment-name" -AZURE_OPENAI_GPT3_5_CHAT_UNDERLYING_MODEL="" - -AZURE_OPENAI_GPT4_CHAT_ENDPOINT="" -AZURE_OPENAI_GPT4_CHAT_KEY="xxxxx" -AZURE_OPENAI_GPT4_CHAT_MODEL="deployment-name" -AZURE_OPENAI_GPT4_CHAT_UNDERLYING_MODEL="" - -AZURE_OPENAI_GPT5_4_ENDPOINT="" -AZURE_OPENAI_GPT5_4_KEY="xxxxx" -AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" -AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" +AZURE_OPENAI_GPT4O_AAD_ENDPOINT="" +AZURE_OPENAI_GPT4O_AAD_MODEL="deployment-name" +AZURE_OPENAI_GPT4O_AAD_UNDERLYING_MODEL="gpt-4o" # Endpoints that host models with fewer safety mechanisms (e.g. via adversarial fine tuning + # or content filters turned off) can be defined below and used in adversarial attack testing scenarios AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="" -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL="" - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2="" -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" -# Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) -# Default endpoint goes here; specialized ones below +# Objective Scorer chat target (used in scorers in scenarios) -ADVERSARIAL_CHAT_ENDPOINT="" -ADVERSARIAL_CHAT_KEY="xxxxx" -ADVERSARIAL_CHAT_MODEL="deployment-name" +OBJECTIVE_SCORER_CHAT_ENDPOINT="" +OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" +OBJECTIVE_SCORER_CHAT_UNDERLYING_MODEL="" -ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="" -ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" -ADVERSARIAL_CHAT_SINGLETURN_MODEL="deployment-name" +AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="" +AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" +AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" -ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="" -ADVERSARIAL_CHAT_MULTITURN_KEY="xxxxx" -ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" +AZURE_OPENAI_GPT5_COMPLETIONS_ENDPOINT="" +AZURE_OPENAI_GPT5_COMPLETIONS_MODEL="gpt-5" +AZURE_OPENAI_GPT5_COMPLETIONS_UNDERLYING_MODEL="gpt-5" -ADVERSARIAL_CHAT_REASONING_ENDPOINT="" -ADVERSARIAL_CHAT_REASONING_KEY="xxxxx" -ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" +AZURE_OPENAI_E2E_TEST_ENDPOINT="" +AZURE_OPENAI_E2E_TEST_MODEL="deployment-name" +AZURE_OPENAI_E2E_TEST_UNDERLYING_MODEL="" -# Objective Scorer chat target (used in scorers in scenarios) +AZURE_OPENAI_GPT5_4_ENDPOINT="" +AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" +AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" -OBJECTIVE_SCORER_CHAT_ENDPOINT="" -OBJECTIVE_SCORER_CHAT_KEY="xxxxx" -OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" +AZURE_OPENAI_GPT4O_STRICT_FILTER_ENDPOINT="" +AZURE_OPENAI_GPT4O_STRICT_FILTER_MODEL="deployment-name" +AZURE_OPENAI_GPT4O_STRICT_FILTER_UNDERLYING_MODEL="gpt-4o" + +AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="" +AZURE_OPENAI_GPT3_5_CHAT_MODEL="deployment-name" +AZURE_OPENAI_GPT3_5_CHAT_UNDERLYING_MODEL="" + +AZURE_OPENAI_GPT4_CHAT_ENDPOINT="" +AZURE_OPENAI_GPT4_CHAT_MODEL="deployment-name" +AZURE_OPENAI_GPT4_CHAT_UNDERLYING_MODEL="" + +MAI_CHAT_ENDPOINT="" +MAI_CHAT_MODEL="deployment-name" +MAI_CHAT_KEY="xxxxx" +MAI_CHAT_UNDERLYING_MODEL="" + +AZURE_OPENAI_GPTV_CHAT_ENDPOINT="" +AZURE_OPENAI_GPTV_CHAT_MODEL="deployment-name" +AZURE_OPENAI_GPTV_CHAT_UNDERLYING_MODEL="" AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="" AZURE_FOUNDRY_DEEPSEEK_KEY="xxxxx" AZURE_FOUNDRY_DEEPSEEK_MODEL="" - AZURE_FOUNDRY_PHI4_ENDPOINT="" AZURE_CHAT_PHI4_KEY="xxxxx" AZURE_CHAT_PHI4_MODEL="" - AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="" -AZURE_FOUNDRY_MISTRAL_LARGE_KEY="xxxxx" AZURE_FOUNDRY_MISTRAL_LARGE_MODEL="Mistral-Large-3" +OLLAMA_CHAT_ENDPOINT="" +OLLAMA_MODEL="llama2" +AZURE_OPENAI_RESPONSES_ENDPOINT="" +AZURE_OPENAI_RESPONSES_MODEL="o4-mini" +AZURE_OPENAI_RESPONSES_UNDERLYING_MODEL="o4-mini" -AWS_ENDPOINT="" -AWS_KEY="xxxxx" -AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" -AWS_RESPONSES_MODEL="openai.gpt-oss-120b" +AZURE_OPENAI_O4MINI_STRICT_FILTER_RESPONSES_ENDPOINT="" +AZURE_OPENAI_O4MINI_STRICT_FILTER_RESPONSES_MODEL="o4-mini" +AZURE_OPENAI_O4MINI_STRICT_FILTER_RESPONSES_UNDERLYING_MODEL="o4-mini" -GROQ_ENDPOINT="" -GROQ_KEY="gsk_xxxxxxxx" -GROQ_LLAMA_MODEL="llama3-8b-8192" +AZURE_OPENAI_GPT41_RESPONSES_ENDPOINT="" +AZURE_OPENAI_GPT41_RESPONSES_MODEL="gpt-4.1" +AZURE_OPENAI_GPT41_RESPONSES_UNDERLYING_MODEL="gpt-4.1" -OPEN_ROUTER_ENDPOINT="" -OPEN_ROUTER_KEY="sk-or-v1-xxxxx" -OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" - -OLLAMA_CHAT_ENDPOINT="" -OLLAMA_MODEL="llama2" +AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="" +AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="" +AZURE_OPENAI_GPT5_MODEL="gpt-5" +AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" DEFAULT_OPENAI_FRONTEND_ENDPOINT = ${AZURE_OPENAI_GPT4O_AAD_ENDPOINT} DEFAULT_OPENAI_FRONTEND_KEY = ${AZURE_OPENAI_GPT4O_AAD_KEY} DEFAULT_OPENAI_FRONTEND_MODEL = "gpt-4o" +DEFAULT_OPENAI_FRONTEND_UNDERLYING_MODEL = "gpt-4o" OPENAI_CHAT_ENDPOINT=${PLATFORM_OPENAI_CHAT_ENDPOINT} -OPENAI_CHAT_KEY=${PLATFORM_OPENAI_CHAT_KEY} OPENAI_CHAT_MODEL=${PLATFORM_OPENAI_CHAT_MODEL} # The following line can be populated if using an Azure OpenAI deployment @@ -155,28 +177,6 @@ OPENAI_CHAT_MODEL=${PLATFORM_OPENAI_CHAT_MODEL} # where the deployment name differs from the actual underlying model OPENAI_CHAT_UNDERLYING_MODEL="" - -################################## - -# OPENAI RESPONSES TARGET SECRETS - -################################## - -AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="" -AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="" -AZURE_OPENAI_GPT5_KEY="xxxxxxx" -AZURE_OPENAI_GPT5_MODEL="gpt-5" -AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" - -PLATFORM_OPENAI_RESPONSES_ENDPOINT="" -PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" -PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" - -AZURE_OPENAI_RESPONSES_ENDPOINT="" -AZURE_OPENAI_RESPONSES_KEY="xxxxx" -AZURE_OPENAI_RESPONSES_MODEL="o4-mini" -AZURE_OPENAI_RESPONSES_UNDERLYING_MODEL="o4-mini" - OPENAI_RESPONSES_ENDPOINT=${PLATFORM_OPENAI_RESPONSES_ENDPOINT} OPENAI_RESPONSES_KEY=${PLATFORM_OPENAI_RESPONSES_KEY} OPENAI_RESPONSES_MODEL=${PLATFORM_OPENAI_RESPONSES_MODEL} @@ -186,25 +186,12 @@ OPENAI_RESPONSES_UNDERLYING_MODEL="" # OPENAI REALTIME TARGET SECRETS -# - -# The below models work with RealtimeTarget - either pass via environment variables - -# or copy to OPENAI_REALTIME_ENDPOINT - ################################## -PLATFORM_OPENAI_REALTIME_ENDPOINT="wss://api.openai.com/v1" -PLATFORM_OPENAI_REALTIME_API_KEY="sk-xxxxx" -PLATFORM_OPENAI_REALTIME_MODEL="gpt-4o-realtime-preview" - AZURE_OPENAI_REALTIME_ENDPOINT = "wss://xxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_REALTIME_API_KEY = "xxxxx" AZURE_OPENAI_REALTIME_MODEL = "gpt-4o-realtime-preview" AZURE_OPENAI_REALTIME_UNDERLYING_MODEL = "gpt-4o-realtime-preview" - OPENAI_REALTIME_ENDPOINT = ${PLATFORM_OPENAI_REALTIME_ENDPOINT} -OPENAI_REALTIME_API_KEY = ${PLATFORM_OPENAI_REALTIME_API_KEY} OPENAI_REALTIME_MODEL = ${PLATFORM_OPENAI_REALTIME_MODEL} OPENAI_REALTIME_UNDERLYING_MODEL = "" @@ -212,31 +199,15 @@ OPENAI_REALTIME_UNDERLYING_MODEL = "" # IMAGE TARGET SECRETS -# - -# The below models work with OpenAIImageTarget - either pass via environment variables - -# or copy to OPENAI_IMAGE_ENDPOINT - -# Entra auth should be enabled - -################################### - -AZURE_OPENAI_IMAGE_ENDPOINT1 = "" -AZURE_OPENAI_IMAGE_API_KEY1 = "xxxxxx" -AZURE_OPENAI_IMAGE_MODEL1 = "deployment-name" -AZURE_OPENAI_IMAGE_UNDERLYING_MODEL1 = "dall-e-3" +################################## -AZURE_OPENAI_IMAGE_ENDPOINT2 = "" -AZURE_OPENAI_IMAGE_API_KEY2 = "xxxxxx" -AZURE_OPENAI_IMAGE_MODEL2 = "dall-e-3" -AZURE_OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" +OPENAI_IMAGE_ENDPOINT2 = "" +OPENAI_IMAGE_MODEL2 = "dall-e-3" +OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" OPENAI_IMAGE_ENDPOINT = ${AZURE_OPENAI_IMAGE_ENDPOINT2} -OPENAI_IMAGE_API_KEY = ${AZURE_OPENAI_IMAGE_API_KEY2} OPENAI_IMAGE_MODEL = ${AZURE_OPENAI_IMAGE_MODEL2} OPENAI_IMAGE_UNDERLYING_MODEL = "" - OPENAI_IMAGE_STRICT_FILTER_ENDPOINT = "" OPENAI_IMAGE_STRICT_FILTER_MODEL = "gpt-image" OPENAI_IMAGE_STRICT_FILTER_UNDERLYING_MODEL = "gpt-image" @@ -245,28 +216,17 @@ OPENAI_IMAGE_STRICT_FILTER_UNDERLYING_MODEL = "gpt-image" # TTS TARGET SECRETS -# - -# The below models work with OpenAITTSTarget - either pass via environment variables - -# or copy to OPENAI_TTS_ENDPOINT - -# Entra auth should be enabled +################################## -################################### +OPENAI_TTS_ENDPOINT1 = "" +OPENAI_TTS_MODEL1 = "tts" +OPENAI_TTS_UNDERLYING_MODEL1 = "tts" -AZURE_OPENAI_TTS_ENDPOINT1 = "" -AZURE_OPENAI_TTS_KEY1 = "xxxxxxx" -AZURE_OPENAI_TTS_MODEL1 = "tts" -AZURE_OPENAI_TTS_UNDERLYING_MODEL1 = "tts" - -AZURE_OPENAI_TTS_ENDPOINT2 = "" -AZURE_OPENAI_TTS_KEY2 = "xxxxxx" -AZURE_OPENAI_TTS_MODEL2 = "tts-1" -AZURE_OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" +OPENAI_TTS_ENDPOINT2 = "" +OPENAI_TTS_MODEL2 = "tts-1" +OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" OPENAI_TTS_ENDPOINT = ${AZURE_OPENAI_TTS_ENDPOINT2} -OPENAI_TTS_KEY = ${AZURE_OPENAI_TTS_KEY2} OPENAI_TTS_MODEL = ${AZURE_OPENAI_TTS_MODEL2} OPENAI_TTS_UNDERLYING_MODEL = "" @@ -274,36 +234,55 @@ OPENAI_TTS_UNDERLYING_MODEL = "" # VIDEO TARGET SECRETS +################################## + # # The below models work with OpenAIVideoTarget - either pass via environment variables # or copy to OPENAI_VIDEO_ENDPOINT -################################### - # Note: Use the base URL without API path AZURE_OPENAI_VIDEO_ENDPOINT="" -AZURE_OPENAI_VIDEO_KEY="xxxxxxx" AZURE_OPENAI_VIDEO_MODEL="sora-2" AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" - OPENAI_VIDEO_ENDPOINT = ${AZURE_OPENAI_VIDEO_ENDPOINT} -OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" ################################## +# ADVERSARIAL MODELS + +################################## + +# Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) + +# Default endpoint goes here; specialized ones below + +ADVERSARIAL_CHAT_ENDPOINT="" +ADVERSARIAL_CHAT_MODEL="deployment-name" +ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="" +ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" +ADVERSARIAL_CHAT_SINGLETURN_MODEL="deployment-name" +ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="" +ADVERSARIAL_CHAT_MULTITURN_KEY="xxxxx" +ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" +ADVERSARIAL_CHAT_REASONING_ENDPOINT="" +ADVERSARIAL_CHAT_REASONING_KEY="xxxxx" +ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" + +################################## + # AML TARGET SECRETS +################################## + # The below models work with AzureMLChatTarget - either pass via environment variables # or copy to AZURE_ML_MANAGED_ENDPOINT -################################### - AZURE_ML_PHI_ENDPOINT="" AZURE_ML_PHI_KEY="xxxxx" @@ -316,41 +295,27 @@ AZURE_ML_KEY=${AZURE_ML_PHI_KEY} # MISC TARGET SECRETS -################################### - -OPENAI_COMPLETION_ENDPOINT="" -OPENAI_COMPLETION_API_KEY="xxxxx" -OPENAI_COMPLETION_MODEL="davinci-002" +################################## OPENAI_EMBEDDING_ENDPOINT="" -OPENAI_EMBEDDING_KEY="xxxxx" OPENAI_EMBEDDING_MODEL="text-embedding-3-small" - -AZURE_STORAGE_ACCOUNT_CONTAINER_URL="" -AZURE_STORAGE_ACCOUNT_SAS_TOKEN="xxxxx" - AZURE_SPEECH_REGION = "eastus2" -AZURE_SPEECH_KEY = "xxxxx" # Resource ID is needed when using Entra authentication AZURE_SPEECH_RESOURCE_ID = "xxxxx" - -AZURE_CONTENT_SAFETY_API_KEY="xxxxx" AZURE_CONTENT_SAFETY_API_ENDPOINT="" - HUGGINGFACE_TOKEN="hf_xxxxxxx" HUGGINGFACE_ENDPOINT="" -GOOGLE_GEMINI_ENDPOINT = "" -GOOGLE_GEMINI_API_KEY = "xxxxx" -GOOGLE_GEMINI_MODEL="gemini-2.0-flash" - -######################### +################################## # AZURE SQL SECRETS -######################### +################################## + +AZURE_STORAGE_ACCOUNT_CONTAINER_URL_PROD="" +AZURE_STORAGE_ACCOUNT_CONTAINER_URL_TEST="" # This connects to the test database @@ -361,8 +326,139 @@ AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST=" ~/.pyrit/.env -> ~/.pyrit/.env.local. -# Explicit env_files replace the default files and load after the AKV bootstrap. -# PyRIT emits a warning when these local files coexist with env_akv_ref so stale -# configuration cannot silently mask or be mistaken for the Key Vault document. -# When migrating, remove or clear ~/.pyrit/.env and ~/.pyrit/.env.local, remove -# explicit env_files if Key Vault should be authoritative, and restart PyRIT. -# Authentication uses DefaultAzureCredential (managed identity, Azure CLI, etc.). -# Key Vault operations use up to three retries with exponential backoff and -# raise KeyVaultInitializationException on bootstrap or secret-resolution failure. -# If env_akv_ref and local files are omitted, PyRIT uses existing process -# environment variables and continues initialization. -# -# 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 -# -# Strict validation applies only to the Key Vault bootstrap and is enabled by -# default. Set this to false to skip malformed or valueless bootstrap entries -# with a warning while loading valid entries. Local files retain standard -# python-dotenv parsing regardless of this setting. -# Empty assignments (NAME=) and child secrets containing an empty string are valid. +# 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 diff --git a/doc/code/executor/gcg/1_gcg_azure_ml.ipynb b/doc/code/executor/gcg/1_gcg_azure_ml.ipynb index a4c7b20040..264805db7a 100644 --- a/doc/code/executor/gcg/1_gcg_azure_ml.ipynb +++ b/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", diff --git a/doc/code/executor/gcg/1_gcg_azure_ml.py b/doc/code/executor/gcg/1_gcg_azure_ml.py index c3c559f18c..9e05233255 100644 --- a/doc/code/executor/gcg/1_gcg_azure_ml.py +++ b/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) diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 89b6489be4..82fd52dd84 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -16,46 +16,39 @@ Then edit both files for your environment. The `.pyrit_conf` tells PyRIT _how_ t 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) - -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. - -### Environment Variable Precedence +## Environment Configuration -When PyRIT initializes, environment variables are loaded in a specific order. **Later sources override earlier ones:** - -```{mermaid} -flowchart LR - A["System environment"] --> B{"env_akv_ref configured?"} - B -->|Yes| C["AKV bootstrap documents in order"] - B -->|No| D{"Explicit env_files?"} - C --> D - D -->|Yes| E["Explicit files in order"] - D -->|No| F["~/.pyrit/.env"] - F --> G["~/.pyrit/.env.local"] +```{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. ``` -System environment variables are always the baseline. If no AKV bootstrap document or environment file is available, PyRIT continues initialization using the existing process environment only. +See [Populating Secrets](./populating_secrets.md) for provider-specific variable examples. + +### Loading Order -**Default file behavior** (no `env_akv_ref` or `env_files` field in `.pyrit_conf`): +PyRIT loads environment sources in this order. Each later source overrides matching values from earlier sources: -| 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) | +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`. -**AKV behavior** (with `env_akv_ref`): The referenced secrets load in list order before local files. Unless custom `env_files` are configured, `~/.pyrit/.env` loads afterward and `~/.pyrit/.env.local` loads last. Later bootstrap documents and local files may override earlier values. +For a Key Vault-only setup, explicitly disable local dotenv loading: -PyRIT emits a warning when `env_akv_ref` is selected and default or explicit environment files coexist with it. The warning distinguishes files that are ignored from files that load afterward and override Key Vault values, making stale migration files visible at startup. When migrating to Key Vault, clear or remove `~/.pyrit/.env` and `~/.pyrit/.env.local`, remove explicit `env_files` when Key Vault should be the only source, and re-initialize PyRIT so values already present in the process environment cannot mask the Key Vault configuration. +```yaml +env_akv_ref: + - https://my-vault.vault.azure.net/secrets/my-pyrit-env +env_files: [] +``` -**Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. They override Key Vault bootstrap values when both fields are configured, and default paths are completely ignored. +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 @@ -167,7 +160,7 @@ 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 | | ----------------- | -------------------------------------------------------- | @@ -191,7 +184,7 @@ When `env_akv_ref` is not configured, an empty `env_files` list or missing defau ### `env_akv_ref` -Ordered Azure Key Vault secret URLs used to obtain bootstrap environment documents. Each secret value must contain dotenv-formatted entries. Authentication uses `DefaultAzureCredential`. +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: @@ -226,7 +219,7 @@ 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" ``` -The bootstrap documents are held in memory and never written to disk. They load before explicit `env_files` or the default `~/.pyrit/.env` and `~/.pyrit/.env.local`, allowing local values to override shared configuration. +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` @@ -242,6 +235,18 @@ Non-strict mode does not suppress Key Vault or reference failures. Missing secre 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`. @@ -363,17 +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 - -# Optional ordered Azure Key Vault bootstrap environment documents +# Recommended: ordered Azure Key Vault bootstrap environment documents # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env -# env_akv_strict: false # Optional; defaults to true +# 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 diff --git a/pyrit/executor/promptgen/gcg/experiments/run.py b/pyrit/executor/promptgen/gcg/experiments/run.py index 3c7cbf0390..c5a2c84bb0 100644 --- a/pyrit/executor/promptgen/gcg/experiments/run.py +++ b/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: diff --git a/pyrit/setup/akv_initialization.py b/pyrit/setup/akv_initialization.py new file mode 100644 index 0000000000..c0d48d7323 --- /dev/null +++ b/pyrit/setup/akv_initialization.py @@ -0,0 +1,554 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Load dotenv files and Azure Key Vault-backed environment documents.""" + +import asyncio +import io +import logging +import os +import pathlib +import urllib.parse +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +import dotenv +from dotenv.parser import parse_stream + +from pyrit.common import path +from pyrit.exceptions import KeyVaultInitializationException + +if TYPE_CHECKING: + from azure.keyvault.secrets.aio import SecretClient + +logger = logging.getLogger(__name__) + +_AKV_REFERENCE_PREFIXES = frozenset({"akv", "kv", "azure_key_vault", "env_akv_ref"}) +_AKV_VAULT_DNS_SUFFIXES = frozenset({"vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"}) +_AKV_RETRY_TOTAL = 3 +_AKV_RETRY_BACKOFF_FACTOR = 0.8 +_AKV_ENV_FILE_NAME = ".env" + + +def _load_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool = False, + include_default_base: bool = True, +) -> bool: + """ + Load environment files in the order they are provided. + Later files override values from earlier files. + + Args: + env_files: Optional sequence of environment file paths. If None, loads default + .env and .env.local from PyRIT home directory (only if they exist). + silent: If True, suppresses print statements about environment file loading. + Defaults to False. + include_default_base: If False and env_files is None, skips the default + .env file while still loading .env.local. Defaults to True. + + Returns: + True if at least one environment file was loaded, otherwise False. + + Raises: + ValueError: If any provided env_files do not exist. + """ + selected_files = _select_environment_files( + env_files=env_files, + silent=silent, + include_default_base=include_default_base, + ) + for env_file in selected_files: + dotenv.load_dotenv(dotenv_path=env_file, override=True, interpolate=True) + if not silent: + _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) + + return bool(selected_files) + + +def _select_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool, + include_default_base: bool, +) -> list[pathlib.Path]: + """ + Select and validate environment files without reading their contents. + + Returns: + list[pathlib.Path]: Environment files in load order. + + Raises: + ValueError: If an explicitly provided environment file does not exist. + """ + if env_files is not None: + if not silent: + _print_msg(f"Loading custom environment files: {[str(f) for f in env_files]}", quiet=silent, log=True) + for env_file in env_files: + if not env_file.exists(): + raise ValueError(f"Environment file not found: {env_file}") + + # By default load .env and .env.local from home directory of the package + else: + default_files = [] + base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" + local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" + + if include_default_base and base_file.exists(): + default_files.append(base_file) + if local_file.exists(): + default_files.append(local_file) + + if not silent: + if default_files: + _print_msg( + f"Found default environment files: {[str(f) for f in default_files]}", quiet=silent, log=True + ) + else: + _print_msg( + "No default environment files found. Using system environment variables only.", + quiet=silent, + log=True, + ) + + env_files = default_files + + return list(env_files) + + +def _print_msg(message: str, quiet: bool, log: bool) -> None: + """ + Print a standard initialization message unless quiet is True. + + Args: + message (str): The message to print and/or log. + quiet (bool): If True, suppresses the initialization message. + log (bool): If True, logs the message using the logger. + """ + if not quiet: + print(message) + if log: + logger.info(message) + + +def _warn_about_akv_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool = False, +) -> None: + """Warn when local environment files coexist with an AKV environment source.""" + base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" + local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" + messages: list[str] = [] + + if base_file.exists(): + if env_files is None: + messages.append(f"{base_file} will load after Key Vault and override matching values") + else: + messages.append(f"{base_file} exists but will be ignored because env_files was explicitly configured") + + if local_file.exists(): + if env_files is None: + messages.append(f"{local_file} will load after Key Vault and override matching values") + else: + messages.append(f"{local_file} exists but will be ignored because env_files was explicitly configured") + + if env_files: + messages.append(f"explicit env_files will load after Key Vault and override matching values: {list(env_files)}") + + if not messages: + return + + message = ( + "env_akv_ref is configured, but local environment files were also found:\n- " + + "\n- ".join(messages) + + "\nWhen migrating to Key Vault, clear or remove ~/.pyrit/.env and ~/.pyrit/.env.local, " + "remove explicit env_files when Key Vault should be the only source, and restart PyRIT so stale " + "process values cannot mask Key Vault configuration." + ) + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + + +def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: + """ + Parse an AKV secret URL into vault URL, secret name, and optional version. + + Args: + secret_url (str): Full AKV secret URL in the format + ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. + + Returns: + tuple[str, str, str | None]: (vault_url, secret_name, secret_version) + + Raises: + ValueError: If the URL does not match the expected format. + """ + error_message = ( + f"Invalid AKV secret URL: '{secret_url}'. Expected an HTTPS Azure Key Vault URL in the format " + "https://{vault}.{vault-dns-suffix}/secrets/{name}[/{version}]." + ) + try: + parsed_url = urllib.parse.urlsplit(secret_url) + port = parsed_url.port + except (TypeError, ValueError) as error: + raise ValueError(error_message) from error + + hostname = parsed_url.hostname + vault_name, separator, dns_suffix = hostname.partition(".") if hostname else ("", "", "") + valid_vault_name = 1 <= len(vault_name) <= 63 and all( + char.isascii() and (char.isalnum() or char == "-") for char in vault_name + ) + valid_authority = ( + parsed_url.scheme.casefold() == "https" + and parsed_url.username is None + and parsed_url.password is None + and port is None + and separator == "." + and dns_suffix in _AKV_VAULT_DNS_SUFFIXES + and valid_vault_name + ) + path_parts = parsed_url.path.split("/") + valid_path = ( + len(path_parts) in {3, 4} and path_parts[0] == "" and path_parts[1] == "secrets" and all(path_parts[2:]) + ) + if not valid_authority or not valid_path or parsed_url.query or parsed_url.fragment: + raise ValueError(error_message) + + secret_name = path_parts[2] + secret_version = path_parts[3] if len(path_parts) == 4 else None + if not _is_valid_akv_identifier(secret_name) or ( + secret_version is not None and not _is_valid_akv_identifier(secret_version) + ): + raise ValueError(error_message) + + return f"https://{hostname}", secret_name, secret_version + + +def _is_valid_akv_identifier(identifier: str) -> bool: + """ + Check whether a Key Vault secret name or version uses URL-safe characters. + + Returns: + bool: True when the identifier is valid. + """ + return 1 <= len(identifier) <= 127 and all( + char.isascii() and (char.isalnum() or char == "-") for char in identifier + ) + + +def _create_akv_secret_client(*, vault_url: str, credential: Any) -> "SecretClient": + """ + Create an asynchronous Key Vault client with an explicit retry policy. + + Returns: + SecretClient: Configured asynchronous secret client. + """ + from azure.core.pipeline.policies import AsyncRetryPolicy + from azure.keyvault.secrets.aio import SecretClient + + retry_policy = AsyncRetryPolicy( + retry_total=_AKV_RETRY_TOTAL, + retry_connect=_AKV_RETRY_TOTAL, + retry_read=_AKV_RETRY_TOTAL, + retry_status=_AKV_RETRY_TOTAL, + retry_backoff_factor=_AKV_RETRY_BACKOFF_FACTOR, + ) + return SecretClient(vault_url=vault_url, credential=credential, retry_policy=retry_policy) + + +def _key_vault_initialization_error(*, message: str, error: Exception) -> KeyVaultInitializationException: + """ + Create a contextual Key Vault exception without losing the original cause. + + Returns: + KeyVaultInitializationException: Wrapped contextual exception. + """ + status_code = getattr(error, "status_code", None) + return KeyVaultInitializationException( + status_code=status_code if isinstance(status_code, int) else 500, + message=f"{message}: {error}", + ) + + +def _validate_dotenv_document( + document: str, + *, + strict: bool = True, + silent: bool = False, +) -> str: + """ + Validate that every dotenv binding uses ``NAME=VALUE`` syntax. + + Args: + document (str): The dotenv document to validate. + strict (bool): If True, reject any invalid entry. If False, warn and + allow python-dotenv to skip invalid entries. Defaults to True. + silent (bool): If True, suppress the console warning. Defaults to False. + + Returns: + str: The original document, or a sanitized document when strict is False. + + Raises: + ValueError: If strict is True and the document contains invalid entries. + """ + bindings = list(parse_stream(io.StringIO(document))) + malformed_lines = [str(binding.original.line) for binding in bindings if binding.error] + valueless_names = [binding.key for binding in bindings if binding.key is not None and binding.value is None] + issues: list[str] = [] + if malformed_lines: + issues.append("malformed entries at lines: " + ", ".join(malformed_lines)) + if valueless_names: + issues.append("variables without values: " + ", ".join(valueless_names)) + if not issues: + return document + + details = "; ".join(issues) + if strict: + raise ValueError("AKV environment document contains " + details) + + message = "AKV environment document contains invalid entries that will be skipped: " + details + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + return "".join( + binding.original.string + for binding in bindings + if not binding.error and not (binding.key is not None and binding.value is None) + ) + + +async def _load_env_from_akv_async( + *, + secret_url: str, + strict: bool = True, + silent: bool = False, +) -> str: + """ + Load a bootstrap dotenv document and resolve its same-vault secret references. + + References are resolved once. Referenced secret values are treated as terminal + strings and are not interpreted as additional references. + + Authentication uses ``DefaultAzureCredential``, which silently tries managed + identity, Azure CLI, VS Code credentials, etc., and falls back to interactive + browser authentication when running locally. + + Args: + secret_url (str): AKV secret URL in the format + ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. + strict (bool): If True, reject malformed or valueless dotenv entries. + If False, warn and skip those entries. Defaults to True. + silent (bool): If True, suppresses print statements. Defaults to False. + + Returns: + str: The validated bootstrap dotenv document before child-secret resolution. + + Raises: + ImportError: If ``azure-keyvault-secrets`` is not installed. + KeyVaultInitializationException: If the root URL is malformed or the bootstrap environment + document cannot be fully resolved. + ValueError: Compatibility base of ``KeyVaultInitializationException``. + """ + from azure.identity.aio import DefaultAzureCredential + + try: + _print_msg(f"Loading environment from AKV secret: {secret_url}", quiet=silent, log=True) + vault_url, secret_name, secret_version = _parse_akv_secret_url(secret_url) + async with DefaultAzureCredential() as credential: + async with _create_akv_secret_client(vault_url=vault_url, credential=credential) as client: + secret = await client.get_secret(secret_name, version=secret_version) + + if not secret.value: + raise ValueError(f"AKV environment secret has no value: {secret_url}") + + validated_document = _validate_dotenv_document(secret.value, strict=strict, silent=silent) + parsed_environment = dotenv.dotenv_values(stream=io.StringIO(validated_document), interpolate=True) + if not parsed_environment: + raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") + loaded = dotenv.load_dotenv( + stream=io.StringIO(validated_document), + override=True, + interpolate=True, + ) + if not loaded: + return validated_document + + for variable_name, value in parsed_environment.items(): + if value is None: + continue + target = _parse_akv_reference(value) + if target is None: + continue + try: + referenced_name, referenced_version = _resolve_akv_secret_reference( + target=target, + variable_name=variable_name, + vault_url=vault_url, + ) + referenced_secret = await client.get_secret(referenced_name, version=referenced_version) + if referenced_secret.value is None: + raise ValueError( + f"AKV secret '{referenced_name}' referenced by environment variable " + f"'{variable_name}' has no value." + ) + os.environ[variable_name] = referenced_secret.value + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to resolve Key Vault reference for environment variable '{variable_name}'", + error=error, + ) + raise wrapped_error from error + return validated_document + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to load Key Vault bootstrap secret '{secret_url}'", + error=error, + ) + raise wrapped_error from error + + +async def _load_environment_async( + *, + env_akv_ref: Sequence[str] | None, + env_files: Sequence[pathlib.Path] | None, + env_akv_strict: bool, + env_akv_write_env: bool = False, + silent: bool, +) -> None: + """ + Load environment sources in precedence order. + + Args: + env_akv_ref (Sequence[str] | None): Optional ordered Key Vault bootstrap secret URLs. + env_files (Sequence[pathlib.Path] | None): Optional ordered local environment files. + env_akv_strict (bool): Whether bootstrap dotenv validation is strict. + env_akv_write_env (bool): Whether to save fetched bootstrap documents to + ``~/.pyrit/.env``. Defaults to False. + silent (bool): Whether initialization messages are suppressed. + + Raises: + ValueError: If a configured source or reference is invalid. + """ + if isinstance(env_akv_ref, str): + raise ValueError("env_akv_ref must be a sequence of Azure Key Vault secret URLs.") + bootstrap_documents: list[str] = [] + if env_akv_ref: + if any(not isinstance(secret_url, str) or not secret_url.strip() for secret_url in env_akv_ref): + raise ValueError("env_akv_ref must contain only non-empty Azure Key Vault secret URLs.") + await asyncio.to_thread( + _warn_about_akv_environment_files, + env_files=env_files, + silent=silent, + ) + bootstrap_documents.extend( + [ + await _load_env_from_akv_async( + secret_url=secret_url, + strict=env_akv_strict, + silent=silent, + ) + for secret_url in env_akv_ref + ] + ) + + written_env_file: pathlib.Path | None = None + if env_akv_write_env and bootstrap_documents: + written_env_file = await asyncio.to_thread( + _write_akv_env_file, + documents=bootstrap_documents, + silent=silent, + ) + + selected_env_files = env_files + if written_env_file is not None and env_files is not None: + written_path = written_env_file.resolve() + selected_env_files = [env_file for env_file in env_files if env_file.expanduser().resolve() != written_path] + + await asyncio.to_thread( + _load_environment_files, + env_files=selected_env_files, + silent=silent, + include_default_base=not (written_env_file is not None and env_files is None), + ) + + +def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Path: + """ + Write fetched bootstrap documents without resolved child-secret values. + + Returns: + pathlib.Path: Path to the written dotenv file. + """ + env_file = path.CONFIGURATION_DIRECTORY_PATH / _AKV_ENV_FILE_NAME + env_file.parent.mkdir(parents=True, exist_ok=True) + content = "\n".join(document.rstrip("\r\n") for document in documents) + "\n" + env_file.write_text(content, encoding="utf-8") + try: + env_file.chmod(0o600) + except OSError: + logger.warning("Could not restrict permissions on written AKV environment file: %s", env_file) + _print_msg(f"Saved Key Vault bootstrap environment file: {env_file}", quiet=silent, log=True) + return env_file + + +def _parse_akv_reference(value: str) -> str | None: + """ + Parse an exact whole-value Key Vault reference. + + Returns: + The referenced secret URL, or None for a literal value. + """ + prefix, separator, target = value.partition(":") + return target.strip() if separator and prefix in _AKV_REFERENCE_PREFIXES else None + + +def _validate_akv_secret_name(*, secret_name: str, variable_name: str) -> None: + if not _is_valid_akv_identifier(secret_name): + raise ValueError( + f"Invalid same-vault secret name '{secret_name}' referenced by environment variable '{variable_name}'. " + "Secret names must contain only letters, numbers, and hyphens." + ) + + +def _resolve_akv_secret_reference( + *, + target: str, + variable_name: str, + vault_url: str, +) -> tuple[str, str | None]: + """ + Resolve a full same-vault secret URI. + + Args: + target (str): Full Key Vault secret URI. + variable_name (str): The environment variable receiving the secret. + vault_url (str): The bootstrap document's vault URL. + + Returns: + tuple[str, str | None]: Secret name and optional version. + + Raises: + ValueError: If the target is not a full URI, is invalid, or references another vault. + """ + if not target.casefold().startswith("https://"): + raise ValueError( + f"AKV reference for environment variable '{variable_name}' must use a full secret URL, " + "for example kv:https://my-vault.vault.azure.net/secrets/my-secret." + ) + + referenced_vault_url, secret_name, secret_version = _parse_akv_secret_url(target) + if referenced_vault_url.rstrip("/").casefold() != vault_url.rstrip("/").casefold(): + raise ValueError( + f"Cross-vault AKV reference for environment variable '{variable_name}' is not supported. " + f"Expected vault '{vault_url}', got '{referenced_vault_url}'." + ) + + _validate_akv_secret_name(secret_name=secret_name, variable_name=variable_name) + return secret_name, secret_version diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index ecd11f0344..ef118ebb0f 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -99,6 +99,8 @@ class ConfigurationLoader(YamlLoadable): env_akv_ref: Ordered list of Key Vault bootstrap secret URLs. env_akv_strict: Whether malformed or valueless entries in a Key Vault bootstrap document should fail initialization. + env_akv_write_env: Whether to save fetched bootstrap documents to + ``~/.pyrit/.env`` for local inspection. silent: Whether to suppress initialization messages. operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. @@ -139,6 +141,7 @@ class ConfigurationLoader(YamlLoadable): env_files: list[str] | None = None env_akv_ref: list[str] | None = None env_akv_strict: bool = True + env_akv_write_env: bool = False silent: bool = False operator: str | None = None operation: str | None = None @@ -421,6 +424,7 @@ def load_with_overrides( env_files: Sequence[str] | None = None, env_akv_ref: Sequence[str] | None = None, env_akv_strict: bool | None = None, + env_akv_write_env: bool | None = None, ) -> "ConfigurationLoader": """ Load configuration with optional overrides. @@ -438,6 +442,7 @@ def load_with_overrides( env_files: Override for environment file paths. env_akv_ref: Override for the ordered Azure Key Vault bootstrap secret URLs. env_akv_strict: Override for strict Key Vault bootstrap validation. + env_akv_write_env: Override for writing the Key Vault bootstrap environment file. Returns: A merged ConfigurationLoader instance. @@ -505,6 +510,9 @@ def to_init_data(config: ConfigurationLoader) -> dict[str, Any]: if env_akv_strict is not None: config_data["env_akv_strict"] = env_akv_strict + if env_akv_write_env is not None: + config_data["env_akv_write_env"] = env_akv_write_env + return cls.from_dict(config_data) @classmethod @@ -641,6 +649,7 @@ async def initialize_pyrit_async(self) -> None: env_files=resolved_env_files, env_akv_ref=self.env_akv_ref, env_akv_strict=self.env_akv_strict, + env_akv_write_env=self.env_akv_write_env, silent=self.silent, ) diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 5d3c713ec2..78a4e4b804 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -1,25 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import asyncio -import io import logging -import os import pathlib -import urllib.parse from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Literal, get_args -import dotenv -from dotenv.parser import parse_stream - -from pyrit.common import path from pyrit.common.apply_defaults import reset_default_values -from pyrit.exceptions import KeyVaultInitializationException from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory +from pyrit.setup.akv_initialization import _load_environment_async if TYPE_CHECKING: - from azure.keyvault.secrets.aio import SecretClient - from pyrit.setup.pyrit_initializer import PyRITInitializer logger = logging.getLogger(__name__) @@ -29,490 +19,6 @@ AZURE_SQL = "AzureSQL" MemoryDatabaseType = Literal["InMemory", "SQLite", "AzureSQL"] -_AKV_REFERENCE_PREFIXES = frozenset({"akv", "kv", "azure_key_vault", "env_akv_ref"}) -_AKV_VAULT_DNS_SUFFIXES = frozenset({"vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"}) -_AKV_RETRY_TOTAL = 3 -_AKV_RETRY_BACKOFF_FACTOR = 0.8 - - -def _load_environment_files( - env_files: Sequence[pathlib.Path] | None, - *, - silent: bool = False, - include_default_base: bool = True, -) -> bool: - """ - Load environment files in the order they are provided. - Later files override values from earlier files. - - Args: - env_files: Optional sequence of environment file paths. If None, loads default - .env and .env.local from PyRIT home directory (only if they exist). - silent: If True, suppresses print statements about environment file loading. - Defaults to False. - include_default_base: If False and env_files is None, skips the default - .env file while still loading .env.local. Defaults to True. - - Returns: - True if at least one environment file was loaded, otherwise False. - - Raises: - ValueError: If any provided env_files do not exist. - """ - selected_files = _select_environment_files( - env_files=env_files, - silent=silent, - include_default_base=include_default_base, - ) - for env_file in selected_files: - dotenv.load_dotenv(dotenv_path=env_file, override=True, interpolate=True) - if not silent: - _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) - - return bool(selected_files) - - -def _select_environment_files( - env_files: Sequence[pathlib.Path] | None, - *, - silent: bool, - include_default_base: bool, -) -> list[pathlib.Path]: - """ - Select and validate environment files without reading their contents. - - Returns: - list[pathlib.Path]: Environment files in load order. - - Raises: - ValueError: If an explicitly provided environment file does not exist. - """ - if env_files is not None: - if not silent: - _print_msg(f"Loading custom environment files: {[str(f) for f in env_files]}", quiet=silent, log=True) - for env_file in env_files: - if not env_file.exists(): - raise ValueError(f"Environment file not found: {env_file}") - - # By default load .env and .env.local from home directory of the package - else: - default_files = [] - base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" - local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" - - if include_default_base and base_file.exists(): - default_files.append(base_file) - if local_file.exists(): - default_files.append(local_file) - - if not silent: - if default_files: - _print_msg( - f"Found default environment files: {[str(f) for f in default_files]}", quiet=silent, log=True - ) - else: - _print_msg( - "No default environment files found. Using system environment variables only.", - quiet=silent, - log=True, - ) - - env_files = default_files - - return list(env_files) - - -def _print_msg(message: str, quiet: bool, log: bool) -> None: - """ - Print a standard initialization message unless quiet is True. - - Args: - message (str): The message to print and/or log. - quiet (bool): If True, suppresses the initialization message. - log (bool): If True, logs the message using the logger. - """ - if not quiet: - print(message) - if log: - logger.info(message) - - -def _warn_about_akv_environment_files( - env_files: Sequence[pathlib.Path] | None, - *, - silent: bool = False, -) -> None: - """Warn when local environment files coexist with an AKV environment source.""" - base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" - local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" - messages: list[str] = [] - - if base_file.exists(): - if env_files is None: - messages.append(f"{base_file} will load after Key Vault and override matching values") - else: - messages.append(f"{base_file} exists but will be ignored because env_files was explicitly configured") - - if local_file.exists(): - if env_files is None: - messages.append(f"{local_file} will load after Key Vault and override matching values") - else: - messages.append(f"{local_file} exists but will be ignored because env_files was explicitly configured") - - if env_files: - messages.append(f"explicit env_files will load after Key Vault and override matching values: {list(env_files)}") - - if not messages: - return - - message = ( - "env_akv_ref is configured, but local environment files were also found:\n- " - + "\n- ".join(messages) - + "\nWhen migrating to Key Vault, clear or remove ~/.pyrit/.env and ~/.pyrit/.env.local, " - "remove explicit env_files when Key Vault should be the only source, and restart PyRIT so stale " - "process values cannot mask Key Vault configuration." - ) - if not silent: - print(f"WARNING: {message}") - logger.warning(message) - - -def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: - """ - Parse an AKV secret URL into vault URL, secret name, and optional version. - - Args: - secret_url (str): Full AKV secret URL in the format - ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. - - Returns: - tuple[str, str, str | None]: (vault_url, secret_name, secret_version) - - Raises: - ValueError: If the URL does not match the expected format. - """ - error_message = ( - f"Invalid AKV secret URL: '{secret_url}'. Expected an HTTPS Azure Key Vault URL in the format " - "https://{vault}.{vault-dns-suffix}/secrets/{name}[/{version}]." - ) - try: - parsed_url = urllib.parse.urlsplit(secret_url) - port = parsed_url.port - except (TypeError, ValueError) as error: - raise ValueError(error_message) from error - - hostname = parsed_url.hostname - vault_name, separator, dns_suffix = hostname.partition(".") if hostname else ("", "", "") - valid_vault_name = 1 <= len(vault_name) <= 63 and all( - char.isascii() and (char.isalnum() or char == "-") for char in vault_name - ) - valid_authority = ( - parsed_url.scheme.casefold() == "https" - and parsed_url.username is None - and parsed_url.password is None - and port is None - and separator == "." - and dns_suffix in _AKV_VAULT_DNS_SUFFIXES - and valid_vault_name - ) - path_parts = parsed_url.path.split("/") - valid_path = ( - len(path_parts) in {3, 4} and path_parts[0] == "" and path_parts[1] == "secrets" and all(path_parts[2:]) - ) - if not valid_authority or not valid_path or parsed_url.query or parsed_url.fragment: - raise ValueError(error_message) - - secret_name = path_parts[2] - secret_version = path_parts[3] if len(path_parts) == 4 else None - if not _is_valid_akv_identifier(secret_name) or ( - secret_version is not None and not _is_valid_akv_identifier(secret_version) - ): - raise ValueError(error_message) - - return f"https://{hostname}", secret_name, secret_version - - -def _is_valid_akv_identifier(identifier: str) -> bool: - """ - Check whether a Key Vault secret name or version uses URL-safe characters. - - Returns: - bool: True when the identifier is valid. - """ - return 1 <= len(identifier) <= 127 and all( - char.isascii() and (char.isalnum() or char == "-") for char in identifier - ) - - -def _create_akv_secret_client(*, vault_url: str, credential: Any) -> "SecretClient": - """ - Create an asynchronous Key Vault client with an explicit retry policy. - - Returns: - SecretClient: Configured asynchronous secret client. - """ - from azure.core.pipeline.policies import AsyncRetryPolicy - from azure.keyvault.secrets.aio import SecretClient - - retry_policy = AsyncRetryPolicy( - retry_total=_AKV_RETRY_TOTAL, - retry_connect=_AKV_RETRY_TOTAL, - retry_read=_AKV_RETRY_TOTAL, - retry_status=_AKV_RETRY_TOTAL, - retry_backoff_factor=_AKV_RETRY_BACKOFF_FACTOR, - ) - return SecretClient(vault_url=vault_url, credential=credential, retry_policy=retry_policy) - - -def _key_vault_initialization_error(*, message: str, error: Exception) -> KeyVaultInitializationException: - """ - Create a contextual Key Vault exception without losing the original cause. - - Returns: - KeyVaultInitializationException: Wrapped contextual exception. - """ - status_code = getattr(error, "status_code", None) - return KeyVaultInitializationException( - status_code=status_code if isinstance(status_code, int) else 500, - message=f"{message}: {error}", - ) - - -def _validate_dotenv_document( - document: str, - *, - strict: bool = True, - silent: bool = False, -) -> str: - """ - Validate that every dotenv binding uses ``NAME=VALUE`` syntax. - - Args: - document (str): The dotenv document to validate. - strict (bool): If True, reject any invalid entry. If False, warn and - allow python-dotenv to skip invalid entries. Defaults to True. - silent (bool): If True, suppress the console warning. Defaults to False. - - Returns: - str: The original document, or a sanitized document when strict is False. - - Raises: - ValueError: If strict is True and the document contains invalid entries. - """ - bindings = list(parse_stream(io.StringIO(document))) - malformed_lines = [str(binding.original.line) for binding in bindings if binding.error] - valueless_names = [binding.key for binding in bindings if binding.key is not None and binding.value is None] - issues: list[str] = [] - if malformed_lines: - issues.append("malformed entries at lines: " + ", ".join(malformed_lines)) - if valueless_names: - issues.append("variables without values: " + ", ".join(valueless_names)) - if not issues: - return document - - details = "; ".join(issues) - if strict: - raise ValueError("AKV environment document contains " + details) - - message = "AKV environment document contains invalid entries that will be skipped: " + details - if not silent: - print(f"WARNING: {message}") - logger.warning(message) - return "".join( - binding.original.string - for binding in bindings - if not binding.error and not (binding.key is not None and binding.value is None) - ) - - -async def _load_env_from_akv_async( - *, - secret_url: str, - strict: bool = True, - silent: bool = False, -) -> None: - """ - Load a bootstrap dotenv document and resolve its same-vault secret references. - - References are resolved once. Referenced secret values are treated as terminal - strings and are not interpreted as additional references. - - Authentication uses ``DefaultAzureCredential``, which silently tries managed - identity, Azure CLI, VS Code credentials, etc., and falls back to interactive - browser authentication when running locally. - - Args: - secret_url (str): AKV secret URL in the format - ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. - strict (bool): If True, reject malformed or valueless dotenv entries. - If False, warn and skip those entries. Defaults to True. - silent (bool): If True, suppresses print statements. Defaults to False. - - Raises: - ImportError: If ``azure-keyvault-secrets`` is not installed. - KeyVaultInitializationException: If the root URL is malformed or the bootstrap environment - document cannot be fully resolved. - ValueError: Compatibility base of ``KeyVaultInitializationException``. - """ - from azure.identity.aio import DefaultAzureCredential - - try: - _print_msg(f"Loading environment from AKV secret: {secret_url}", quiet=silent, log=True) - vault_url, secret_name, secret_version = _parse_akv_secret_url(secret_url) - async with DefaultAzureCredential() as credential: - async with _create_akv_secret_client(vault_url=vault_url, credential=credential) as client: - secret = await client.get_secret(secret_name, version=secret_version) - - if not secret.value: - raise ValueError(f"AKV environment secret has no value: {secret_url}") - - validated_document = _validate_dotenv_document(secret.value, strict=strict, silent=silent) - parsed_environment = dotenv.dotenv_values(stream=io.StringIO(validated_document), interpolate=True) - if not parsed_environment: - raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") - loaded = dotenv.load_dotenv( - stream=io.StringIO(validated_document), - override=True, - interpolate=True, - ) - if not loaded: - return - - for variable_name, value in parsed_environment.items(): - if value is None: - continue - target = _parse_akv_reference(value) - if target is None: - continue - try: - referenced_name, referenced_version = _resolve_akv_secret_reference( - target=target, - variable_name=variable_name, - vault_url=vault_url, - ) - referenced_secret = await client.get_secret(referenced_name, version=referenced_version) - if referenced_secret.value is None: - raise ValueError( - f"AKV secret '{referenced_name}' referenced by environment variable " - f"'{variable_name}' has no value." - ) - os.environ[variable_name] = referenced_secret.value - except KeyVaultInitializationException: - raise - except Exception as error: - wrapped_error = _key_vault_initialization_error( - message=f"Failed to resolve Key Vault reference for environment variable '{variable_name}'", - error=error, - ) - raise wrapped_error from error - except KeyVaultInitializationException: - raise - except Exception as error: - wrapped_error = _key_vault_initialization_error( - message=f"Failed to load Key Vault bootstrap secret '{secret_url}'", - error=error, - ) - raise wrapped_error from error - - -async def _load_environment_async( - *, - env_akv_ref: Sequence[str] | None, - env_files: Sequence[pathlib.Path] | None, - env_akv_strict: bool, - silent: bool, -) -> None: - """ - Load environment sources in precedence order. - - Args: - env_akv_ref (Sequence[str] | None): Optional ordered Key Vault bootstrap secret URLs. - env_files (Sequence[pathlib.Path] | None): Optional ordered local environment files. - env_akv_strict (bool): Whether bootstrap dotenv validation is strict. - silent (bool): Whether initialization messages are suppressed. - - Raises: - ValueError: If a configured source or reference is invalid. - """ - if isinstance(env_akv_ref, str): - raise ValueError("env_akv_ref must be a sequence of Azure Key Vault secret URLs.") - if env_akv_ref: - if any(not isinstance(secret_url, str) or not secret_url.strip() for secret_url in env_akv_ref): - raise ValueError("env_akv_ref must contain only non-empty Azure Key Vault secret URLs.") - await asyncio.to_thread( - _warn_about_akv_environment_files, - env_files=env_files, - silent=silent, - ) - for secret_url in env_akv_ref: - await _load_env_from_akv_async( - secret_url=secret_url, - strict=env_akv_strict, - silent=silent, - ) - - await asyncio.to_thread( - _load_environment_files, - env_files=env_files, - silent=silent, - ) - - -def _parse_akv_reference(value: str) -> str | None: - """ - Parse an exact whole-value Key Vault reference. - - Returns: - The referenced secret URL, or None for a literal value. - """ - prefix, separator, target = value.partition(":") - return target.strip() if separator and prefix in _AKV_REFERENCE_PREFIXES else None - - -def _validate_akv_secret_name(*, secret_name: str, variable_name: str) -> None: - if not _is_valid_akv_identifier(secret_name): - raise ValueError( - f"Invalid same-vault secret name '{secret_name}' referenced by environment variable '{variable_name}'. " - "Secret names must contain only letters, numbers, and hyphens." - ) - - -def _resolve_akv_secret_reference( - *, - target: str, - variable_name: str, - vault_url: str, -) -> tuple[str, str | None]: - """ - Resolve a full same-vault secret URI. - - Args: - target (str): Full Key Vault secret URI. - variable_name (str): The environment variable receiving the secret. - vault_url (str): The bootstrap document's vault URL. - - Returns: - tuple[str, str | None]: Secret name and optional version. - - Raises: - ValueError: If the target is not a full URI, is invalid, or references another vault. - """ - if not target.casefold().startswith("https://"): - raise ValueError( - f"AKV reference for environment variable '{variable_name}' must use a full secret URL, " - "for example kv:https://my-vault.vault.azure.net/secrets/my-secret." - ) - - referenced_vault_url, secret_name, secret_version = _parse_akv_secret_url(target) - if referenced_vault_url.rstrip("/").casefold() != vault_url.rstrip("/").casefold(): - raise ValueError( - f"Cross-vault AKV reference for environment variable '{variable_name}' is not supported. " - f"Expected vault '{vault_url}', got '{referenced_vault_url}'." - ) - - _validate_akv_secret_name(secret_name=secret_name, variable_name=variable_name) - return secret_name, secret_version - async def _execute_initializers_async(*, initializers: Sequence["PyRITInitializer"]) -> None: """ @@ -564,6 +70,7 @@ async def initialize_pyrit_async( env_files: Sequence[pathlib.Path] | None = None, env_akv_ref: Sequence[str] | None = None, env_akv_strict: bool = True, + env_akv_write_env: bool = False, silent: bool = False, **memory_instance_kwargs: Any, ) -> None: @@ -595,6 +102,8 @@ async def initialize_pyrit_async( and local files take precedence. Requires ``azure-keyvault-secrets``. env_akv_strict (bool): If True, reject malformed or valueless entries in the Key Vault bootstrap document. If False, warn and skip those entries. Defaults to True. + env_akv_write_env (bool): If True, save fetched bootstrap documents with unresolved + child references to ``~/.pyrit/.env``. Defaults to False. silent (bool): If True, suppresses print statements about environment file loading and schema migration. Defaults to False. **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. @@ -606,6 +115,7 @@ async def initialize_pyrit_async( env_akv_ref=env_akv_ref, env_files=env_files, env_akv_strict=env_akv_strict, + env_akv_write_env=env_akv_write_env, silent=silent, ) diff --git a/tests/end_to_end/executor/promptgen/gcg/test_gcg_aml_e2e.py b/tests/end_to_end/executor/promptgen/gcg/test_gcg_aml_e2e.py index 83d5078d80..8b6ed6ce40 100644 --- a/tests/end_to_end/executor/promptgen/gcg/test_gcg_aml_e2e.py +++ b/tests/end_to_end/executor/promptgen/gcg/test_gcg_aml_e2e.py @@ -47,7 +47,7 @@ pytest.importorskip("azure.identity", reason="azure-identity not installed") from pyrit.common.path import HOME_PATH # noqa: E402 -from pyrit.setup.initialization import _load_environment_files # noqa: E402 +from pyrit.setup.akv_initialization import _load_environment_files # noqa: E402 _REQUIRED_ENV_VARS = ( "AZURE_ML_SUBSCRIPTION_ID", diff --git a/tests/integration/setup/test_akv_initialization_integration.py b/tests/integration/setup/test_akv_initialization_integration.py new file mode 100644 index 0000000000..a976666c29 --- /dev/null +++ b/tests/integration/setup/test_akv_initialization_integration.py @@ -0,0 +1,51 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import os + +import pytest + +from pyrit.setup import IN_MEMORY, initialize_pyrit_async + +_SECRET_URL_ENV = "PYRIT_AKV_INTEGRATION_TEST_SECRET_URL" +_VARIABLE_NAME_ENV = "PYRIT_AKV_INTEGRATION_TEST_VARIABLE" +_EXPECTED_VALUE_ENV = "PYRIT_AKV_INTEGRATION_TEST_EXPECTED_VALUE" + + +def _get_bootstrap_secret_url() -> str: + """ + Get the explicitly configured integration bootstrap URL. + + Returns: + str: Key Vault bootstrap secret URL. + """ + configured_url = os.getenv(_SECRET_URL_ENV) + if not configured_url: + pytest.skip(f"Set {_SECRET_URL_ENV} to run this integration test.") + return configured_url + + +@pytest.mark.run_only_if_all_tests +async def test_akv_bootstrap_initialization_populates_process_environment() -> None: + variable_name = os.getenv(_VARIABLE_NAME_ENV, "TEST_KEY") + expected_value = os.getenv(_EXPECTED_VALUE_ENV, "surprise") + bootstrap_secret_url = _get_bootstrap_secret_url() + + missing = object() + original_value: object = os.environ.pop(variable_name, missing) + try: + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=[bootstrap_secret_url], + env_files=[], + load_defaults=False, + silent=True, + ) + + if os.environ.get(variable_name) != expected_value: + raise AssertionError(f"{variable_name} did not resolve to the expected integration-test sentinel.") + finally: + if original_value is missing: + os.environ.pop(variable_name, None) + else: + os.environ[variable_name] = str(original_value) diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 9682e9b2ae..93c812943a 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -43,6 +43,7 @@ def test_default_values(self): assert config.env_files is None # None means "use defaults" assert config.env_akv_ref is None assert config.env_akv_strict is True + assert config.env_akv_write_env is False assert config.silent is False def test_valid_memory_db_types_snake_case(self): @@ -149,6 +150,7 @@ def test_from_dict_with_all_fields(self): "env_files": ["/path/to/.env"], "env_akv_ref": ["https://vault.vault.azure.net/secrets/one"], "env_akv_strict": False, + "env_akv_write_env": True, "silent": True, } config = ConfigurationLoader.from_dict(data) @@ -158,6 +160,7 @@ def test_from_dict_with_all_fields(self): assert config.env_files == ["/path/to/.env"] assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] assert config.env_akv_strict is False + assert config.env_akv_write_env is True assert config.silent is True def test_from_dict_filters_none_values(self): @@ -346,6 +349,7 @@ async def test_initialize_pyrit_async_basic(self, mock_init): assert call_kwargs["env_files"] is None assert call_kwargs["env_akv_ref"] is None assert call_kwargs["env_akv_strict"] is True + assert call_kwargs["env_akv_write_env"] is False assert call_kwargs["silent"] is False @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") @@ -355,7 +359,12 @@ async def test_initialize_pyrit_async_with_env_akv_ref(self, mock_init): "https://vault.vault.azure.net/secrets/first", "https://vault.vault.azure.net/secrets/second/version", ] - config = ConfigurationLoader(memory_db_type="in_memory", env_akv_ref=refs, env_akv_strict=False) + config = ConfigurationLoader( + memory_db_type="in_memory", + env_akv_ref=refs, + env_akv_strict=False, + env_akv_write_env=True, + ) await config.initialize_pyrit_async() @@ -363,6 +372,7 @@ async def test_initialize_pyrit_async_with_env_akv_ref(self, mock_init): call_kwargs = mock_init.call_args.kwargs assert call_kwargs["env_akv_ref"] == refs assert call_kwargs["env_akv_strict"] is False + assert call_kwargs["env_akv_write_env"] is True @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") @mock.patch("pyrit.registry.InitializerRegistry") @@ -530,6 +540,14 @@ def test_load_with_overrides_env_akv_ref_override(self, mock_default_path): assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] + @mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH") + def test_load_with_overrides_env_akv_write_env_override(self, mock_default_path): + mock_default_path.exists.return_value = False + + config = ConfigurationLoader.load_with_overrides(env_akv_write_env=True) + + assert config.env_akv_write_env is True + @mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH") def test_load_with_overrides_converts_sequence_to_list(self, mock_default_path): """Test that Sequence inputs are converted to list for dataclass compatibility.""" diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 0e53da0deb..0abea4ab90 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -4,24 +4,14 @@ import os import pathlib import tempfile -import types from unittest import mock import pytest -from azure.core.exceptions import ResourceNotFoundError from pyrit.common.apply_defaults import reset_default_values from pyrit.common.singleton import Singleton -from pyrit.exceptions import KeyVaultInitializationException from pyrit.registry import InitializerRegistry from pyrit.setup import IN_MEMORY, initialize_pyrit_async -from pyrit.setup.initialization import ( - _load_env_from_akv_async, - _load_environment_files, - _parse_akv_reference, - _parse_akv_secret_url, - _warn_about_akv_environment_files, -) class TestLoadInitializersFromScripts: @@ -129,7 +119,7 @@ def setup_method(self) -> None: reset_default_values() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=False) async def test_initialize_basic(self, mock_load_env, mock_set_memory): """Test basic initialization.""" await initialize_pyrit_async(memory_db_type=IN_MEMORY, load_defaults=False) @@ -138,7 +128,7 @@ async def test_initialize_basic(self, mock_load_env, mock_set_memory): mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=False) async def test_initialize_with_script(self, mock_load_env, mock_set_memory): """Test initialization with a script.""" with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f: @@ -168,15 +158,15 @@ async def initialize_async(self) -> None: finally: os.unlink(script_path) - @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=False) async def test_invalid_memory_type_raises_error(self, mock_load_env): """Test that invalid memory type raises ValueError.""" with pytest.raises(ValueError, match="is not a supported type"): await initialize_pyrit_async(memory_db_type="InvalidType", load_defaults=False) # type: ignore[arg-type] @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) - @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) + @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, mock_set_memory): """Test that env_akv_ref loads bootstrap secrets in order.""" refs = [ @@ -196,8 +186,8 @@ async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, m mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) - @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) + @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) async def test_initialize_with_empty_env_akv_ref_does_not_load_akv( self, mock_load_akv, mock_load_env, mock_set_memory ): @@ -224,9 +214,9 @@ async def test_initialize_keeps_akv_values_when_local_file_loading_fails(self, m with mock.patch.dict(os.environ, {}, clear=True): with ( - mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch("pyrit.setup.akv_initialization._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.initialization._load_env_from_akv_async", + "pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock, side_effect=lambda **_: os.environ.update({"FROM_AKV": "resolved"}), ), @@ -252,9 +242,9 @@ async def test_initialize_loads_local_overrides_on_akv_environment(self, mock_se with ( mock.patch.dict(os.environ, {}, clear=True), - mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch("pyrit.setup.akv_initialization._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.initialization._load_env_from_akv_async", + "pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock, side_effect=lambda **_: os.environ.update({"BASE": "akv", "ONLY_AKV": "shared"}), ), @@ -282,10 +272,10 @@ async def test_initialize_default_files_override_akv_in_order(self, mock_set_mem with ( mock.patch.dict(os.environ, {}, clear=True), - mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.akv_initialization._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.initialization._load_env_from_akv_async", + "pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock, side_effect=lambda **_: os.environ.update({"VALUE": "akv"}), ), @@ -319,9 +309,9 @@ async def test_initialize_resolves_bootstrap_references_before_local_overrides(s with ( mock.patch.dict(os.environ, {"SOURCE_VALUE": "ambient-value"}, clear=True), - mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch("pyrit.setup.akv_initialization._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.initialization._load_env_from_akv_async", + "pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock, side_effect=lambda **_: os.environ.update(bootstrap_environment), ), @@ -367,7 +357,7 @@ def setup_method(self) -> None: """Clear default values before each test.""" reset_default_values() - @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=True) + @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=True) async def test_initialize_silent_produces_no_output(self, mock_load_env, capsys): """initialize_pyrit_async with silent=True must not print anything to stdout.""" await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=True, load_defaults=False) @@ -375,707 +365,10 @@ async def test_initialize_silent_produces_no_output(self, mock_load_env, capsys) captured = capsys.readouterr() assert captured.out == "" - @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=True) + @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=True) async def test_initialize_not_silent_prints_migration_message(self, mock_load_env, capsys): """Without silent, the Alembic schema-check message is printed and tagged as Alembic output.""" await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=False, load_defaults=False) captured = capsys.readouterr() assert "[pyrit:alembic] No new upgrade operations detected." in captured.out - - -class TestLoadEnvironmentFiles: - """Tests for _load_environment_files function and env_files parameter in initialize_pyrit_async.""" - - @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_loads_default_env_files_when_none_provided(self, mock_config_path): - """Test that default .env and .env.local files are loaded when env_files is None.""" - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - env_local_file = temp_path / ".env.local" - env_file.write_text("VAR1=value1") - env_local_file.write_text("VAR2=value2") - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - with mock.patch.dict(os.environ, {}, clear=True): - loaded = _load_environment_files(env_files=None) - - assert loaded is True - assert os.environ["VAR1"] == "value1" - assert os.environ["VAR2"] == "value2" - - @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_only_loads_existing_default_files(self, mock_config_path): - """Test that only existing default files are loaded.""" - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - env_file.write_text("VAR1=value1") - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - with mock.patch.dict(os.environ, {}, clear=True): - loaded = _load_environment_files(env_files=None) - - assert loaded is True - assert os.environ["VAR1"] == "value1" - - @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_excludes_default_env_when_loading_local_override(self, mock_config_path): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - env_local_file = temp_path / ".env.local" - env_file.write_text("VAR=base") - env_local_file.write_text("VAR=local") - - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - with mock.patch.dict(os.environ, {}, clear=True): - loaded = _load_environment_files(env_files=None, include_default_base=False) - - assert loaded is True - assert os.environ["VAR"] == "local" - - @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_returns_false_when_no_default_files_exist(self, mock_config_path): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - with mock.patch.dict(os.environ, {}, clear=True): - loaded = _load_environment_files(env_files=None) - - assert loaded is False - assert os.environ == {} - - @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - def test_warns_when_default_files_coexist_with_akv(self, mock_config_path, caplog, capsys): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - env_local_file = temp_path / ".env.local" - env_file.write_text("VAR=base") - env_local_file.write_text("VAR=local") - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - with caplog.at_level("WARNING", logger="pyrit.setup.initialization"): - _warn_about_akv_environment_files(env_files=None) - - output = capsys.readouterr().out - assert output.startswith("WARNING: env_akv_ref is configured") - assert f"{env_file} will load after Key Vault and override matching values" in output - assert f"{env_local_file} will load after Key Vault and override matching values" in output - assert "clear or remove ~/.pyrit/.env and ~/.pyrit/.env.local" in output - assert "remove explicit env_files when Key Vault should be the only source" in output - assert "restart PyRIT" in output - assert caplog.records[0].levelname == "WARNING" - - @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - def test_warns_when_explicit_files_replace_defaults_with_akv(self, mock_config_path, capsys): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - env_local_file = temp_path / ".env.local" - custom_file = temp_path / ".env.custom" - env_file.write_text("VAR=base") - env_local_file.write_text("VAR=local") - custom_file.write_text("VAR=custom") - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - _warn_about_akv_environment_files(env_files=[custom_file]) - - output = capsys.readouterr().out - assert f"{env_file} exists but will be ignored because env_files was explicitly configured" in output - assert f"{env_local_file} exists but will be ignored because env_files was explicitly configured" in output - assert f"explicit env_files will load after Key Vault and override matching values: {[custom_file]}" in output - - @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - def test_akv_environment_file_warning_respects_silent(self, mock_config_path, caplog, capsys): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - (temp_path / ".env").write_text("VAR=base") - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - with caplog.at_level("WARNING", logger="pyrit.setup.initialization"): - _warn_about_akv_environment_files(env_files=None, silent=True) - - assert capsys.readouterr().out == "" - assert "will load after Key Vault and override matching values" in caplog.text - assert "restart PyRIT" in caplog.text - - async def test_loads_custom_env_files_in_order(self): - """Test that custom env_files are loaded in the order provided.""" - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env1 = temp_path / ".env.test" - env2 = temp_path / ".env.prod" - env3 = temp_path / ".env.local" - - # Create files - env1.write_text("VAR=test") - env2.write_text("VAR=prod") - env3.write_text("VAR=local") - - with mock.patch.dict(os.environ, {}, clear=True): - loaded = _load_environment_files(env_files=[env1, env2, env3]) - - assert loaded is True - assert os.environ["VAR"] == "local" - - async def test_load_environment_files_interpolates_in_assignment_order(self): - with tempfile.TemporaryDirectory() as temp_dir: - env_file = pathlib.Path(temp_dir) / ".env" - env_file.write_text("A=one\nB=${A}\nA=two\nC=${A}") - - with mock.patch.dict(os.environ, {}, clear=True): - loaded = _load_environment_files(env_files=[env_file], silent=True) - - assert loaded is True - assert os.environ["A"] == "two" - assert os.environ["B"] == "one" - assert os.environ["C"] == "two" - - async def test_load_environment_files_honors_python_dotenv_disabled(self): - with tempfile.TemporaryDirectory() as temp_dir: - env_file = pathlib.Path(temp_dir) / ".env" - env_file.write_text("DISABLED_VALUE=not-loaded") - - with mock.patch.dict(os.environ, {"PYTHON_DOTENV_DISABLED": "true"}, clear=True): - loaded = _load_environment_files(env_files=[env_file], silent=True) - - assert loaded is True - assert "DISABLED_VALUE" not in os.environ - - async def test_direct_local_file_loader_keeps_pyrit_references_literal(self): - with tempfile.TemporaryDirectory() as temp_dir: - env_file = pathlib.Path(temp_dir) / ".env" - env_file.write_text( - "BASE_VALUE=base\nKV_REFERENCE=kv:api-key\nENV_REFERENCE=env:SOURCE_VALUE\nINTERPOLATED=${BASE_VALUE}" - ) - - with mock.patch.dict(os.environ, {}, clear=True): - loaded = _load_environment_files(env_files=[env_file], silent=True) - - assert loaded is True - assert os.environ["KV_REFERENCE"] == "kv:api-key" - assert os.environ["ENV_REFERENCE"] == "env:SOURCE_VALUE" - assert os.environ["INTERPOLATED"] == "base" - - @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - def test_default_local_file_can_interpolate_base_file_but_not_reverse(self, mock_config_path): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - env_local_file = temp_path / ".env.local" - env_file.write_text( - "OPENAI_CHAT_ENDPOINT=https://example.openai.azure.com/openai/v1\nFROM_LATER_LOCAL=${LOCAL_ONLY}" - ) - env_local_file.write_text("FOOBAR=${OPENAI_CHAT_ENDPOINT}\nLOCAL_ONLY=local") - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - with mock.patch.dict(os.environ, {}, clear=True): - loaded = _load_environment_files(env_files=None, silent=True) - - assert loaded is True - assert os.environ["FOOBAR"] == "https://example.openai.azure.com/openai/v1" - assert os.environ["FROM_LATER_LOCAL"] == "" - assert os.environ["LOCAL_ONLY"] == "local" - - async def test_env_akv_strict_does_not_validate_local_environment_files(self): - with tempfile.TemporaryDirectory() as temp_dir: - env_file = pathlib.Path(temp_dir) / ".env" - env_file.write_text("GOOD=resolved\n=malformed\nOTHER=also-resolved") - - with mock.patch.dict(os.environ, {}, clear=True): - await initialize_pyrit_async( - memory_db_type=IN_MEMORY, - env_files=[env_file], - env_akv_strict=True, - load_defaults=False, - silent=True, - ) - - assert os.environ["GOOD"] == "resolved" - assert os.environ["OTHER"] == "also-resolved" - - @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_keeps_local_akv_reference_literal_without_bootstrap(self, mock_set_memory): - with tempfile.TemporaryDirectory() as temp_dir: - env_file = pathlib.Path(temp_dir) / ".env" - env_file.write_text("API_KEY=kv:https://myvault.vault.azure.net/secrets/api-key") - - with mock.patch.dict(os.environ, {}, clear=True): - await initialize_pyrit_async( - memory_db_type=IN_MEMORY, - env_files=[env_file], - load_defaults=False, - silent=True, - ) - - assert os.environ["API_KEY"] == "kv:https://myvault.vault.azure.net/secrets/api-key" - - mock_set_memory.assert_called_once() - - async def test_raises_error_for_nonexistent_env_file(self): - """Test that ValueError is raised for non-existent env file.""" - nonexistent = pathlib.Path("/nonexistent/path/.env") - - with pytest.raises(ValueError, match="Environment file not found"): - _load_environment_files(env_files=[nonexistent]) - - @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_pyrit_with_custom_env_files(self, mock_set_memory): - """Test initialize_pyrit_async with custom env_files.""" - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env.custom" - env_file.write_text("CUSTOM_VAR=custom_value") - - # Should not raise an error - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[env_file], load_defaults=False) - - mock_set_memory.assert_called_once() - - @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_pyrit_raises_for_nonexistent_env_file(self, mock_set_memory): - """Test that initialize_pyrit_async raises ValueError for non-existent env file.""" - nonexistent = pathlib.Path("/nonexistent/.env") - - with pytest.raises(ValueError, match="Environment file not found"): - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[nonexistent]) - - @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_custom_env_files_override_default_behavior(self, mock_set_memory): - """Test that passing custom env_files prevents loading default files.""" - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - - # Create default files - default_env = temp_path / ".env" - default_env_local = temp_path / ".env.local" - default_env.write_text("DEFAULT=value") - default_env_local.write_text("DEFAULT_LOCAL=value") - - # Create custom file - custom_env = temp_path / ".env.custom" - custom_env.write_text("CUSTOM=value") - - with mock.patch.dict(os.environ, {}, clear=True): - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[custom_env], load_defaults=False) - - assert os.environ["CUSTOM"] == "value" - assert "DEFAULT" not in os.environ - assert "DEFAULT_LOCAL" not in os.environ - - -def _create_mock_akv_clients() -> tuple[mock.MagicMock, mock.MagicMock]: - credential = mock.MagicMock() - credential.__aenter__ = mock.AsyncMock(return_value=credential) - credential.__aexit__ = mock.AsyncMock(return_value=None) - client = mock.MagicMock() - client.__aenter__ = mock.AsyncMock(return_value=client) - client.__aexit__ = mock.AsyncMock(return_value=None) - return credential, client - - -def _assert_mock_akv_client_created( - mock_client_cls: mock.MagicMock, - *, - vault_url: str, - credential: mock.MagicMock, -) -> None: - mock_client_cls.assert_called_once() - call_kwargs = mock_client_cls.call_args.kwargs - assert call_kwargs["vault_url"] == vault_url - assert call_kwargs["credential"] is credential - retry_policy = call_kwargs["retry_policy"] - assert retry_policy.total_retries == 3 - assert retry_policy.connect_retries == 3 - assert retry_policy.read_retries == 3 - assert retry_policy.status_retries == 3 - assert retry_policy.backoff_factor == 0.8 - - -class TestAkvEnvironmentLoading: - """Tests for AKV URL parsing and env loading helpers.""" - - @pytest.mark.parametrize("prefix", ["kv", "akv", "azure_key_vault", "env_akv_ref"]) - def test_parse_akv_reference_accepts_aliases(self, prefix): - secret_url = "https://myvault.vault.azure.net/secrets/api-key" - - assert _parse_akv_reference(f"{prefix}:{secret_url}") == secret_url - - @pytest.mark.parametrize( - "value", - [ - "env:SOURCE_VALUE", - "literal:kv:https://myvault.vault.azure.net/secrets/api-key", - "@Microsoft.KeyVault(SecretUri=https://myvault.vault.azure.net/secrets/api-key)", - ], - ) - def test_parse_akv_reference_ignores_non_akv_syntax(self, value): - assert _parse_akv_reference(value) is None - - def test_parse_akv_secret_url_with_version(self): - url = "https://myvault.vault.azure.net/secrets/my-secret/abc123" - - vault_url, secret_name, secret_version = _parse_akv_secret_url(url) - - assert vault_url == "https://myvault.vault.azure.net" - assert secret_name == "my-secret" - assert secret_version == "abc123" - - def test_parse_akv_secret_url_without_version(self): - url = "https://myvault.vault.azure.net/secrets/my-secret" - - vault_url, secret_name, secret_version = _parse_akv_secret_url(url) - - assert vault_url == "https://myvault.vault.azure.net" - assert secret_name == "my-secret" - assert secret_version is None - - @pytest.mark.parametrize("dns_suffix", ["vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"]) - def test_parse_akv_secret_url_accepts_supported_clouds(self, dns_suffix): - url = f"https://myvault.{dns_suffix}/secrets/my-secret/version-1" - - vault_url, secret_name, secret_version = _parse_akv_secret_url(url) - - assert vault_url == f"https://myvault.{dns_suffix}" - assert secret_name == "my-secret" - assert secret_version == "version-1" - - @pytest.mark.parametrize( - "url", - [ - "http://myvault.vault.azure.net/secrets/my-secret", - "https://attacker.example/secrets/my-secret", - "https://myvault.vault.azure.net.attacker.example/secrets/my-secret", - "https://nested.myvault.vault.azure.net/secrets/my-secret", - "https://user@myvault.vault.azure.net/secrets/my-secret", - "https://myvault.vault.azure.net:443/secrets/my-secret", - "https://myvault.vault.azure.net/not-secrets/my-secret", - "https://myvault.vault.azure.net/secrets", - "https://myvault.vault.azure.net/secrets/my-secret/", - "https://myvault.vault.azure.net/secrets/my-secret/version/extra", - "https://myvault.vault.azure.net/secrets/my-secret?api-version=7.4", - "https://myvault.vault.azure.net/secrets/my-secret#fragment", - "https://myvault.vault.azure.net/secrets/my%2Fsecret", - ], - ) - def test_parse_akv_secret_url_invalid_raises(self, url): - with pytest.raises(ValueError, match="Invalid AKV secret URL"): - _parse_akv_secret_url(url) - - async def test_load_env_from_akv_async_rejects_non_azure_host_before_authentication(self): - with ( - mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, - mock.patch("pyrit.setup.initialization._create_akv_secret_client") as mock_create_client, - pytest.raises(KeyVaultInitializationException, match="attacker.example"), - ): - await _load_env_from_akv_async( - secret_url="https://attacker.example/secrets/bootstrap", - silent=True, - ) - - mock_credential_cls.assert_not_called() - mock_create_client.assert_not_called() - - async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_secrets(self): - credential, client = _create_mock_akv_clients() - root_document = ( - "DIRECT=from-bootstrap\n" - "FROM_ENV=${SOURCE_VALUE}\n" - "FROM_KV=kv:https://myvault.vault.azure.net/secrets/api-key\n" - "PINNED_KV=kv:https://myvault.vault.azure.net/secrets/api-key/version-2\n" - "TERMINAL=kv:https://myvault.vault.azure.net/secrets/terminal\n" - "A=one\nB=${A}\nA=two\nC=${A}" - ) - client.get_secret = mock.AsyncMock( - side_effect=[ - types.SimpleNamespace(value=root_document), - types.SimpleNamespace(value="api-key-value"), - types.SimpleNamespace(value="pinned-key-value"), - types.SimpleNamespace(value="kv:https://myvault.vault.azure.net/secrets/not-followed"), - ] - ) - secret_url = "https://myvault.vault.azure.net/secrets/bootstrap/v1" - - with ( - mock.patch.dict(os.environ, {"SOURCE_VALUE": "ambient-value"}, clear=True), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential) as mock_credential_cls, - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, - mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, - ): - await _load_env_from_akv_async(secret_url=secret_url, silent=True) - - assert os.environ["DIRECT"] == "from-bootstrap" - assert os.environ["FROM_ENV"] == "ambient-value" - assert os.environ["FROM_KV"] == "api-key-value" - assert os.environ["PINNED_KV"] == "pinned-key-value" - assert os.environ["TERMINAL"] == "kv:https://myvault.vault.azure.net/secrets/not-followed" - assert os.environ["A"] == "two" - assert os.environ["B"] == "one" - assert os.environ["C"] == "two" - - mock_credential_cls.assert_called_once_with() - _assert_mock_akv_client_created( - mock_client_cls, - vault_url="https://myvault.vault.azure.net", - credential=credential, - ) - assert client.get_secret.await_args_list == [ - mock.call("bootstrap", version="v1"), - mock.call("api-key", version=None), - mock.call("api-key", version="version-2"), - mock.call("terminal", version=None), - ] - credential.__aenter__.assert_awaited_once() - credential.__aexit__.assert_awaited_once() - client.__aenter__.assert_awaited_once() - client.__aexit__.assert_awaited_once() - mock_print_msg.assert_called_once() - - async def test_load_env_from_akv_async_rejects_short_secret_name(self): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="API_KEY=kv:api-key")) - - with ( - mock.patch.dict(os.environ, {}, clear=True), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - pytest.raises(ValueError, match="must use a full secret URL"), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - ) - - @pytest.mark.parametrize( - "reference_url", - [ - "https://other-vault.vault.azure.net/secrets/api-key", - "https://other-vault.vault.azure.net/secrets/api-key/version-1", - ], - ) - async def test_load_env_from_akv_async_rejects_cross_vault_reference(self, reference_url): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=f"API_KEY=kv:{reference_url}")) - - with ( - mock.patch.dict(os.environ, {}, clear=True), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - pytest.raises(ValueError, match="Cross-vault AKV reference"), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - ) - - async def test_load_env_from_akv_async_empty_secret_raises(self): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=None)) - - with ( - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - pytest.raises(ValueError, match="has no value"), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/my-secret", - silent=True, - ) - - credential.__aexit__.assert_awaited_once() - client.__aexit__.assert_awaited_once() - - async def test_load_env_from_akv_async_without_entries_raises(self): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="# comments only\n")) - - with ( - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - pytest.raises(ValueError, match="contains no environment entries"), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/my-secret", - silent=True, - ) - - credential.__aexit__.assert_awaited_once() - client.__aexit__.assert_awaited_once() - - @pytest.mark.parametrize( - ("document", "error"), - [ - ("GOOD=resolved\n=malformed\nOTHER=resolved", "malformed entries at lines: 2"), - ("MISSING_VALUE\n", "variables without values: MISSING_VALUE"), - ], - ) - async def test_load_env_from_akv_async_rejects_non_assignments(self, document, error): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) - - with mock.patch.dict(os.environ, {}, clear=True): - with ( - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - pytest.raises(ValueError, match=error), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - ) - - assert "GOOD" not in os.environ - assert "OTHER" not in os.environ - - async def test_load_env_from_akv_async_wraps_malformed_bootstrap(self): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="=malformed")) - - with ( - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - pytest.raises(KeyVaultInitializationException, match="malformed entries") as exc_info, - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - ) - - assert isinstance(exc_info.value.__cause__, ValueError) - - async def test_load_env_from_akv_async_wraps_missing_child_secret(self): - credential, client = _create_mock_akv_clients() - missing_error = ResourceNotFoundError(message="Secret was not found") - client.get_secret = mock.AsyncMock( - side_effect=[ - types.SimpleNamespace(value="API_KEY=kv:https://myvault.vault.azure.net/secrets/missing"), - missing_error, - ] - ) - - with ( - mock.patch.dict(os.environ, {}, clear=True), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - pytest.raises(KeyVaultInitializationException, match="Failed to resolve Key Vault reference") as exc_info, - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - ) - - assert exc_info.value.__cause__ is missing_error - - async def test_load_env_from_akv_async_allows_empty_assignment(self): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="EMPTY=")) - - with ( - mock.patch.dict(os.environ, {}, clear=True), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - ) - - assert os.environ["EMPTY"] == "" - - async def test_load_env_from_akv_async_allows_empty_child_secret(self): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock( - side_effect=[ - types.SimpleNamespace(value="EMPTY=kv:https://myvault.vault.azure.net/secrets/empty-secret"), - types.SimpleNamespace(value=""), - ] - ) - - with ( - mock.patch.dict(os.environ, {}, clear=True), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - ) - - assert os.environ["EMPTY"] == "" - assert client.get_secret.await_args_list[-1] == mock.call("empty-secret", version=None) - - async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entries(self, caplog, capsys): - credential, client = _create_mock_akv_clients() - document = "GOOD=resolved\n=malformed\nMISSING_VALUE\nOTHER=also-resolved" - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) - - with ( - mock.patch.dict(os.environ, {}, clear=True), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - caplog.at_level("WARNING", logger="pyrit.setup.initialization"), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - strict=False, - silent=False, - ) - - assert os.environ["GOOD"] == "resolved" - assert os.environ["OTHER"] == "also-resolved" - - output = capsys.readouterr().out - assert "WARNING: AKV environment document contains invalid entries that will be skipped" in output - assert "malformed entries at lines: 2" in output - assert "variables without values: MISSING_VALUE" in output - assert "GOOD" not in caplog.text - assert "resolved" not in caplog.text - - async def test_load_env_from_akv_async_non_strict_silent_logs_warning(self, caplog, capsys): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="GOOD=resolved\nMISSING_VALUE")) - - with ( - mock.patch.dict(os.environ, {}, clear=True), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - caplog.at_level("WARNING", logger="pyrit.setup.initialization"), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - strict=False, - silent=True, - ) - - assert capsys.readouterr().out == "" - assert "variables without values: MISSING_VALUE" in caplog.text - - async def test_load_env_from_akv_async_child_failure_keeps_loaded_bootstrap_values(self): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock( - side_effect=[ - types.SimpleNamespace( - value=("GOOD=resolved\nBAD=kv:https://myvault.vault.azure.net/secrets/missing-value") - ), - types.SimpleNamespace(value=None), - ] - ) - - with mock.patch.dict(os.environ, {}, clear=True): - with ( - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), - pytest.raises(ValueError, match="has no value"), - ): - await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - ) - - assert os.environ["GOOD"] == "resolved" - assert os.environ["BAD"] == "kv:https://myvault.vault.azure.net/secrets/missing-value" From 0d8c5add91cfeebbb7968528dff0b84095289327 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 14 Aug 2026 17:30:25 -0400 Subject: [PATCH 15/16] FEAT: env local integratoin test --- build_scripts/env_local_integration_test | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/build_scripts/env_local_integration_test b/build_scripts/env_local_integration_test index b61cfe7d20..ded45e2bea 100644 --- a/build_scripts/env_local_integration_test +++ b/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} From cfd24ffe0c35f28d4d92cc5c598b8bf36597d831 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 14 Aug 2026 17:42:36 -0400 Subject: [PATCH 16/16] FEAT: Update .env_example --- .env_example | 53 +- tests/unit/setup/test_akv_initialization.py | 784 ++++++++++++++++++++ 2 files changed, 785 insertions(+), 52 deletions(-) create mode 100644 tests/unit/setup/test_akv_initialization.py diff --git a/.env_example b/.env_example index ac0005e69b..94629e75cd 100644 --- a/.env_example +++ b/.env_example @@ -1,78 +1,46 @@ # ============================================================================ - # PyRIT Environment File Example - # ============================================================================ - # - # Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need - # - # MOST USERS ONLY NEED 3 VARIABLES to get started - # - # OPENAI_CHAT_ENDPOINT="" # or any OpenAI-compatible API - # OPENAI_CHAT_KEY="your-key-here" - # OPENAI_CHAT_MODEL="gpt-4o" - # - # These work with OpenAI, Azure OpenAI, Ollama, Groq, OpenRouter, and any - # other OpenAI-compatible endpoint. See doc/setup/populating_secrets.md - # for provider-specific examples - # - # If you are using Entra authentication for Azure resources - # keys for those resources are not needed. PyRIT auto-detects: if an API key - # is set, it uses key auth; otherwise it falls back to Entra ID automatically - # - # ============================================================================ ################################## - # OPENAI TARGET SECRETS - ################################## - # - # The below models work with OpenAIChatTarget - either pass via environment variables - # or copy to OPENAI_CHAT_ENDPOINT - # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately - # Example: AZURE_OPENAI_GPT4O_ENDPOINT="" AZURE_OPENAI_GPT4O_MODEL="deployment-name" # Since Azure deployment name may be custom and differ from the actual underlying model - # you can specify the underlying model for identifier purposes. If not specified - # identifiers will default to the value of the standard MODEL environment variable AZURE_OPENAI_GPT4O_UNDERLYING_MODEL="gpt-4o" # Optional second GPT-4o endpoint (that can be used for round-robin distribution) - # TargetInitializer creates RoundRobinTargets that automatically group together - # targets with identical underlying model names and behavioral params, allowing - # for distribution of requests across them for rate-limit relief AZURE_OPENAI_GPT4O_ENDPOINT2="" @@ -84,7 +52,6 @@ AZURE_OPENAI_GPT4O_AAD_MODEL="deployment-name" AZURE_OPENAI_GPT4O_AAD_UNDERLYING_MODEL="gpt-4o" # Endpoints that host models with fewer safety mechanisms (e.g. via adversarial fine tuning - # or content filters turned off) can be defined below and used in adversarial attack testing scenarios AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="" @@ -173,7 +140,6 @@ OPENAI_CHAT_ENDPOINT=${PLATFORM_OPENAI_CHAT_ENDPOINT} OPENAI_CHAT_MODEL=${PLATFORM_OPENAI_CHAT_MODEL} # The following line can be populated if using an Azure OpenAI deployment - # where the deployment name differs from the actual underlying model OPENAI_CHAT_UNDERLYING_MODEL="" @@ -237,11 +203,8 @@ OPENAI_TTS_UNDERLYING_MODEL = "" ################################## # - # The below models work with OpenAIVideoTarget - either pass via environment variables - # or copy to OPENAI_VIDEO_ENDPOINT - # Note: Use the base URL without API path AZURE_OPENAI_VIDEO_ENDPOINT="" @@ -258,7 +221,6 @@ OPENAI_VIDEO_UNDERLYING_MODEL = "" ################################## # Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) - # Default endpoint goes here; specialized ones below ADVERSARIAL_CHAT_ENDPOINT="" @@ -280,7 +242,6 @@ ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" ################################## # The below models work with AzureMLChatTarget - either pass via environment variables - # or copy to AZURE_ML_MANAGED_ENDPOINT AZURE_ML_PHI_ENDPOINT="" @@ -374,9 +335,7 @@ PLATFORM_OPENAI_EMBEDDING_KEY="sk-xxxxx" PLATFORM_OPENAI_EMBEDDING_MODEL="text-embedding-3-small" # - # The below models work with RealtimeTarget - either pass via environment variables - # or copy to OPENAI_REALTIME_ENDPOINT PLATFORM_OPENAI_REALTIME_ENDPOINT="wss://api.openai.com/v1" @@ -391,12 +350,10 @@ PROMPTINTEL_API_KEY="xxxxx" ################################## -# EXAMPLE-ONLY VARIABLES - REVIEW +# Additional entries referenced in PyRIT ################################## -# These existing dummy assignments are not present in the modern private .env key inventory - AZURE_OPENAI_GPT4O_KEY="xxxxx" AZURE_OPENAI_GPT4O_KEY2="xxxxx" AZURE_OPENAI_INTEGRATION_TEST_KEY="xxxxx" @@ -421,12 +378,8 @@ PLATFORM_OPENAI_REALTIME_API_KEY="sk-xxxxx" AZURE_OPENAI_REALTIME_API_KEY = "xxxxx" OPENAI_REALTIME_API_KEY = ${PLATFORM_OPENAI_REALTIME_API_KEY} -# - # The below models work with OpenAIImageTarget - either pass via environment variables - # or copy to OPENAI_IMAGE_ENDPOINT - # Entra auth should be enabled AZURE_OPENAI_IMAGE_ENDPOINT1 = "" @@ -439,12 +392,8 @@ AZURE_OPENAI_IMAGE_MODEL2 = "dall-e-3" AZURE_OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" OPENAI_IMAGE_API_KEY = ${AZURE_OPENAI_IMAGE_API_KEY2} -# - # The below models work with OpenAITTSTarget - either pass via environment variables - # or copy to OPENAI_TTS_ENDPOINT - # Entra auth should be enabled AZURE_OPENAI_TTS_ENDPOINT1 = "" diff --git a/tests/unit/setup/test_akv_initialization.py b/tests/unit/setup/test_akv_initialization.py new file mode 100644 index 0000000000..81f80b92ea --- /dev/null +++ b/tests/unit/setup/test_akv_initialization.py @@ -0,0 +1,784 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import os +import pathlib +import tempfile +import types +from unittest import mock + +import pytest +from azure.core.exceptions import ResourceNotFoundError + +from pyrit.exceptions import KeyVaultInitializationException +from pyrit.setup import IN_MEMORY, initialize_pyrit_async +from pyrit.setup.akv_initialization import ( + _load_env_from_akv_async, + _load_environment_async, + _load_environment_files, + _parse_akv_reference, + _parse_akv_secret_url, + _warn_about_akv_environment_files, +) + + +class TestLoadEnvironmentFiles: + """Tests for _load_environment_files function and env_files parameter in initialize_pyrit_async.""" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_loads_default_env_files_when_none_provided(self, mock_config_path): + """Test that default .env and .env.local files are loaded when env_files is None.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text("VAR1=value1") + env_local_file.write_text("VAR2=value2") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None) + + assert loaded is True + assert os.environ["VAR1"] == "value1" + assert os.environ["VAR2"] == "value2" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_only_loads_existing_default_files(self, mock_config_path): + """Test that only existing default files are loaded.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_file.write_text("VAR1=value1") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None) + + assert loaded is True + assert os.environ["VAR1"] == "value1" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_excludes_default_env_when_loading_local_override(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text("VAR=base") + env_local_file.write_text("VAR=local") + + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None, include_default_base=False) + + assert loaded is True + assert os.environ["VAR"] == "local" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_returns_false_when_no_default_files_exist(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None) + + assert loaded is False + assert os.environ == {} + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_warns_when_default_files_coexist_with_akv(self, mock_config_path, caplog, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text("VAR=base") + env_local_file.write_text("VAR=local") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"): + _warn_about_akv_environment_files(env_files=None) + + output = capsys.readouterr().out + assert output.startswith("WARNING: env_akv_ref is configured") + assert f"{env_file} will load after Key Vault and override matching values" in output + assert f"{env_local_file} will load after Key Vault and override matching values" in output + assert "clear or remove ~/.pyrit/.env and ~/.pyrit/.env.local" in output + assert "remove explicit env_files when Key Vault should be the only source" in output + assert "restart PyRIT" in output + assert caplog.records[0].levelname == "WARNING" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_warns_when_explicit_files_replace_defaults_with_akv(self, mock_config_path, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + custom_file = temp_path / ".env.custom" + env_file.write_text("VAR=base") + env_local_file.write_text("VAR=local") + custom_file.write_text("VAR=custom") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + _warn_about_akv_environment_files(env_files=[custom_file]) + + output = capsys.readouterr().out + assert f"{env_file} exists but will be ignored because env_files was explicitly configured" in output + assert f"{env_local_file} exists but will be ignored because env_files was explicitly configured" in output + assert f"explicit env_files will load after Key Vault and override matching values: {[custom_file]}" in output + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_akv_environment_file_warning_respects_silent(self, mock_config_path, caplog, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VAR=base") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"): + _warn_about_akv_environment_files(env_files=None, silent=True) + + assert capsys.readouterr().out == "" + assert "will load after Key Vault and override matching values" in caplog.text + assert "restart PyRIT" in caplog.text + + async def test_loads_custom_env_files_in_order(self): + """Test that custom env_files are loaded in the order provided.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env1 = temp_path / ".env.test" + env2 = temp_path / ".env.prod" + env3 = temp_path / ".env.local" + + # Create files + env1.write_text("VAR=test") + env2.write_text("VAR=prod") + env3.write_text("VAR=local") + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env1, env2, env3]) + + assert loaded is True + assert os.environ["VAR"] == "local" + + async def test_load_environment_files_interpolates_in_assignment_order(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("A=one\nB=${A}\nA=two\nC=${A}") + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert os.environ["A"] == "two" + assert os.environ["B"] == "one" + assert os.environ["C"] == "two" + + async def test_load_environment_files_honors_python_dotenv_disabled(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("DISABLED_VALUE=not-loaded") + + with mock.patch.dict(os.environ, {"PYTHON_DOTENV_DISABLED": "true"}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert "DISABLED_VALUE" not in os.environ + + async def test_load_environment_async_write_env_writes_unresolved_bootstrap_documents(self): + references = [ + "https://vault.vault.azure.net/secrets/first", + "https://vault.vault.azure.net/secrets/second", + ] + documents = [ + 'ENDPOINT="https://example.test"\nAPI_KEY="kv:https://vault.vault.azure.net/secrets/api-key"\n', + 'MODEL="model-name"\n', + ] + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch( + "pyrit.setup.akv_initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + side_effect=documents, + ), + mock.patch( + "pyrit.setup.akv_initialization._load_environment_files", return_value=False + ) as mock_load_environment_files, + ): + await _load_environment_async( + env_akv_ref=references, + env_files=None, + env_akv_strict=True, + env_akv_write_env=True, + silent=True, + ) + + assert (temp_path / ".env").read_text(encoding="utf-8") == "".join(documents) + assert "resolved-api-key" not in (temp_path / ".env").read_text(encoding="utf-8") + mock_load_environment_files.assert_called_once() + assert mock_load_environment_files.call_args.kwargs["include_default_base"] is False + + async def test_load_environment_async_write_env_filters_generated_explicit_file(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + generated_env = temp_path / ".env" + local_env = temp_path / ".env.local" + local_env.write_text("LOCAL=value", encoding="utf-8") + + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch( + "pyrit.setup.akv_initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + return_value="VALUE=bootstrap\n", + ), + mock.patch( + "pyrit.setup.akv_initialization._load_environment_files", return_value=True + ) as mock_load_environment_files, + ): + await _load_environment_async( + env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=[generated_env, local_env], + env_akv_strict=True, + env_akv_write_env=True, + silent=True, + ) + + assert mock_load_environment_files.call_args.kwargs["env_files"] == [local_env] + assert mock_load_environment_files.call_args.kwargs["include_default_base"] is True + + async def test_direct_local_file_loader_keeps_pyrit_references_literal(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text( + "BASE_VALUE=base\nKV_REFERENCE=kv:api-key\nENV_REFERENCE=env:SOURCE_VALUE\nINTERPOLATED=${BASE_VALUE}" + ) + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert os.environ["KV_REFERENCE"] == "kv:api-key" + assert os.environ["ENV_REFERENCE"] == "env:SOURCE_VALUE" + assert os.environ["INTERPOLATED"] == "base" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_default_local_file_can_interpolate_base_file_but_not_reverse(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text( + "OPENAI_CHAT_ENDPOINT=https://example.openai.azure.com/openai/v1\nFROM_LATER_LOCAL=${LOCAL_ONLY}" + ) + env_local_file.write_text("FOOBAR=${OPENAI_CHAT_ENDPOINT}\nLOCAL_ONLY=local") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None, silent=True) + + assert loaded is True + assert os.environ["FOOBAR"] == "https://example.openai.azure.com/openai/v1" + assert os.environ["FROM_LATER_LOCAL"] == "" + assert os.environ["LOCAL_ONLY"] == "local" + + async def test_env_akv_strict_does_not_validate_local_environment_files(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("GOOD=resolved\n=malformed\nOTHER=also-resolved") + + with mock.patch.dict(os.environ, {}, clear=True): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_files=[env_file], + env_akv_strict=True, + load_defaults=False, + silent=True, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_keeps_local_akv_reference_literal_without_bootstrap(self, mock_set_memory): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("API_KEY=kv:https://myvault.vault.azure.net/secrets/api-key") + + with mock.patch.dict(os.environ, {}, clear=True): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_files=[env_file], + load_defaults=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "kv:https://myvault.vault.azure.net/secrets/api-key" + + mock_set_memory.assert_called_once() + + async def test_raises_error_for_nonexistent_env_file(self): + """Test that ValueError is raised for non-existent env file.""" + nonexistent = pathlib.Path("/nonexistent/path/.env") + + with pytest.raises(ValueError, match="Environment file not found"): + _load_environment_files(env_files=[nonexistent]) + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_pyrit_with_custom_env_files(self, mock_set_memory): + """Test initialize_pyrit_async with custom env_files.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env.custom" + env_file.write_text("CUSTOM_VAR=custom_value") + + # Should not raise an error + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[env_file], load_defaults=False) + + mock_set_memory.assert_called_once() + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_pyrit_raises_for_nonexistent_env_file(self, mock_set_memory): + """Test that initialize_pyrit_async raises ValueError for non-existent env file.""" + nonexistent = pathlib.Path("/nonexistent/.env") + + with pytest.raises(ValueError, match="Environment file not found"): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[nonexistent]) + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_custom_env_files_override_default_behavior(self, mock_set_memory): + """Test that passing custom env_files prevents loading default files.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + + # Create default files + default_env = temp_path / ".env" + default_env_local = temp_path / ".env.local" + default_env.write_text("DEFAULT=value") + default_env_local.write_text("DEFAULT_LOCAL=value") + + # Create custom file + custom_env = temp_path / ".env.custom" + custom_env.write_text("CUSTOM=value") + + with mock.patch.dict(os.environ, {}, clear=True): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[custom_env], load_defaults=False) + + assert os.environ["CUSTOM"] == "value" + assert "DEFAULT" not in os.environ + assert "DEFAULT_LOCAL" not in os.environ + + +def _create_mock_akv_clients() -> tuple[mock.MagicMock, mock.MagicMock]: + credential = mock.MagicMock() + credential.__aenter__ = mock.AsyncMock(return_value=credential) + credential.__aexit__ = mock.AsyncMock(return_value=None) + client = mock.MagicMock() + client.__aenter__ = mock.AsyncMock(return_value=client) + client.__aexit__ = mock.AsyncMock(return_value=None) + return credential, client + + +def _assert_mock_akv_client_created( + mock_client_cls: mock.MagicMock, + *, + vault_url: str, + credential: mock.MagicMock, +) -> None: + mock_client_cls.assert_called_once() + call_kwargs = mock_client_cls.call_args.kwargs + assert call_kwargs["vault_url"] == vault_url + assert call_kwargs["credential"] is credential + retry_policy = call_kwargs["retry_policy"] + assert retry_policy.total_retries == 3 + assert retry_policy.connect_retries == 3 + assert retry_policy.read_retries == 3 + assert retry_policy.status_retries == 3 + assert retry_policy.backoff_factor == 0.8 + + +class TestAkvEnvironmentLoading: + """Tests for AKV URL parsing and env loading helpers.""" + + @pytest.mark.parametrize("prefix", ["kv", "akv", "azure_key_vault", "env_akv_ref"]) + def test_parse_akv_reference_accepts_aliases(self, prefix): + secret_url = "https://myvault.vault.azure.net/secrets/api-key" + + assert _parse_akv_reference(f"{prefix}:{secret_url}") == secret_url + + @pytest.mark.parametrize( + "value", + [ + "env:SOURCE_VALUE", + "literal:kv:https://myvault.vault.azure.net/secrets/api-key", + "@Microsoft.KeyVault(SecretUri=https://myvault.vault.azure.net/secrets/api-key)", + ], + ) + def test_parse_akv_reference_ignores_non_akv_syntax(self, value): + assert _parse_akv_reference(value) is None + + def test_parse_akv_secret_url_with_version(self): + url = "https://myvault.vault.azure.net/secrets/my-secret/abc123" + + vault_url, secret_name, secret_version = _parse_akv_secret_url(url) + + assert vault_url == "https://myvault.vault.azure.net" + assert secret_name == "my-secret" + assert secret_version == "abc123" + + def test_parse_akv_secret_url_without_version(self): + url = "https://myvault.vault.azure.net/secrets/my-secret" + + vault_url, secret_name, secret_version = _parse_akv_secret_url(url) + + assert vault_url == "https://myvault.vault.azure.net" + assert secret_name == "my-secret" + assert secret_version is None + + @pytest.mark.parametrize("dns_suffix", ["vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"]) + def test_parse_akv_secret_url_accepts_supported_clouds(self, dns_suffix): + url = f"https://myvault.{dns_suffix}/secrets/my-secret/version-1" + + vault_url, secret_name, secret_version = _parse_akv_secret_url(url) + + assert vault_url == f"https://myvault.{dns_suffix}" + assert secret_name == "my-secret" + assert secret_version == "version-1" + + @pytest.mark.parametrize( + "url", + [ + "http://myvault.vault.azure.net/secrets/my-secret", + "https://attacker.example/secrets/my-secret", + "https://myvault.vault.azure.net.attacker.example/secrets/my-secret", + "https://nested.myvault.vault.azure.net/secrets/my-secret", + "https://user@myvault.vault.azure.net/secrets/my-secret", + "https://myvault.vault.azure.net:443/secrets/my-secret", + "https://myvault.vault.azure.net/not-secrets/my-secret", + "https://myvault.vault.azure.net/secrets", + "https://myvault.vault.azure.net/secrets/my-secret/", + "https://myvault.vault.azure.net/secrets/my-secret/version/extra", + "https://myvault.vault.azure.net/secrets/my-secret?api-version=7.4", + "https://myvault.vault.azure.net/secrets/my-secret#fragment", + "https://myvault.vault.azure.net/secrets/my%2Fsecret", + ], + ) + def test_parse_akv_secret_url_invalid_raises(self, url): + with pytest.raises(ValueError, match="Invalid AKV secret URL"): + _parse_akv_secret_url(url) + + async def test_load_env_from_akv_async_rejects_non_azure_host_before_authentication(self): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + mock.patch("pyrit.setup.akv_initialization._create_akv_secret_client") as mock_create_client, + pytest.raises(KeyVaultInitializationException, match="attacker.example"), + ): + await _load_env_from_akv_async( + secret_url="https://attacker.example/secrets/bootstrap", + silent=True, + ) + + mock_credential_cls.assert_not_called() + mock_create_client.assert_not_called() + + async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_secrets(self): + credential, client = _create_mock_akv_clients() + root_document = ( + "DIRECT=from-bootstrap\n" + "FROM_ENV=${SOURCE_VALUE}\n" + "FROM_KV=kv:https://myvault.vault.azure.net/secrets/api-key\n" + "PINNED_KV=kv:https://myvault.vault.azure.net/secrets/api-key/version-2\n" + "TERMINAL=kv:https://myvault.vault.azure.net/secrets/terminal\n" + "A=one\nB=${A}\nA=two\nC=${A}" + ) + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value=root_document), + types.SimpleNamespace(value="api-key-value"), + types.SimpleNamespace(value="pinned-key-value"), + types.SimpleNamespace(value="kv:https://myvault.vault.azure.net/secrets/not-followed"), + ] + ) + secret_url = "https://myvault.vault.azure.net/secrets/bootstrap/v1" + + with ( + mock.patch.dict(os.environ, {"SOURCE_VALUE": "ambient-value"}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential) as mock_credential_cls, + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, + mock.patch("pyrit.setup.akv_initialization._print_msg") as mock_print_msg, + ): + await _load_env_from_akv_async(secret_url=secret_url, silent=True) + + assert os.environ["DIRECT"] == "from-bootstrap" + assert os.environ["FROM_ENV"] == "ambient-value" + assert os.environ["FROM_KV"] == "api-key-value" + assert os.environ["PINNED_KV"] == "pinned-key-value" + assert os.environ["TERMINAL"] == "kv:https://myvault.vault.azure.net/secrets/not-followed" + assert os.environ["A"] == "two" + assert os.environ["B"] == "one" + assert os.environ["C"] == "two" + + mock_credential_cls.assert_called_once_with() + _assert_mock_akv_client_created( + mock_client_cls, + vault_url="https://myvault.vault.azure.net", + credential=credential, + ) + assert client.get_secret.await_args_list == [ + mock.call("bootstrap", version="v1"), + mock.call("api-key", version=None), + mock.call("api-key", version="version-2"), + mock.call("terminal", version=None), + ] + credential.__aenter__.assert_awaited_once() + credential.__aexit__.assert_awaited_once() + client.__aenter__.assert_awaited_once() + client.__aexit__.assert_awaited_once() + mock_print_msg.assert_called_once() + + async def test_load_env_from_akv_async_rejects_short_secret_name(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="API_KEY=kv:api-key")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="must use a full secret URL"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + @pytest.mark.parametrize( + "reference_url", + [ + "https://other-vault.vault.azure.net/secrets/api-key", + "https://other-vault.vault.azure.net/secrets/api-key/version-1", + ], + ) + async def test_load_env_from_akv_async_rejects_cross_vault_reference(self, reference_url): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=f"API_KEY=kv:{reference_url}")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="Cross-vault AKV reference"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + async def test_load_env_from_akv_async_empty_secret_raises(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=None)) + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="has no value"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/my-secret", + silent=True, + ) + + credential.__aexit__.assert_awaited_once() + client.__aexit__.assert_awaited_once() + + async def test_load_env_from_akv_async_without_entries_raises(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="# comments only\n")) + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="contains no environment entries"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/my-secret", + silent=True, + ) + + credential.__aexit__.assert_awaited_once() + client.__aexit__.assert_awaited_once() + + @pytest.mark.parametrize( + ("document", "error"), + [ + ("GOOD=resolved\n=malformed\nOTHER=resolved", "malformed entries at lines: 2"), + ("MISSING_VALUE\n", "variables without values: MISSING_VALUE"), + ], + ) + async def test_load_env_from_akv_async_rejects_non_assignments(self, document, error): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) + + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match=error), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert "GOOD" not in os.environ + assert "OTHER" not in os.environ + + async def test_load_env_from_akv_async_wraps_malformed_bootstrap(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="=malformed")) + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(KeyVaultInitializationException, match="malformed entries") as exc_info, + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert isinstance(exc_info.value.__cause__, ValueError) + + async def test_load_env_from_akv_async_wraps_missing_child_secret(self): + credential, client = _create_mock_akv_clients() + missing_error = ResourceNotFoundError(message="Secret was not found") + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value="API_KEY=kv:https://myvault.vault.azure.net/secrets/missing"), + missing_error, + ] + ) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(KeyVaultInitializationException, match="Failed to resolve Key Vault reference") as exc_info, + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert exc_info.value.__cause__ is missing_error + + async def test_load_env_from_akv_async_allows_empty_assignment(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="EMPTY=")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert os.environ["EMPTY"] == "" + + async def test_load_env_from_akv_async_allows_empty_child_secret(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value="EMPTY=kv:https://myvault.vault.azure.net/secrets/empty-secret"), + types.SimpleNamespace(value=""), + ] + ) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert os.environ["EMPTY"] == "" + assert client.get_secret.await_args_list[-1] == mock.call("empty-secret", version=None) + + async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entries(self, caplog, capsys): + credential, client = _create_mock_akv_clients() + document = "GOOD=resolved\n=malformed\nMISSING_VALUE\nOTHER=also-resolved" + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=False, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" + + output = capsys.readouterr().out + assert "WARNING: AKV environment document contains invalid entries that will be skipped" in output + assert "malformed entries at lines: 2" in output + assert "variables without values: MISSING_VALUE" in output + assert "GOOD" not in caplog.text + assert "resolved" not in caplog.text + + async def test_load_env_from_akv_async_non_strict_silent_logs_warning(self, caplog, capsys): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="GOOD=resolved\nMISSING_VALUE")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=True, + ) + + assert capsys.readouterr().out == "" + assert "variables without values: MISSING_VALUE" in caplog.text + + async def test_load_env_from_akv_async_child_failure_keeps_loaded_bootstrap_values(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace( + value=("GOOD=resolved\nBAD=kv:https://myvault.vault.azure.net/secrets/missing-value") + ), + types.SimpleNamespace(value=None), + ] + ) + + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="has no value"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["BAD"] == "kv:https://myvault.vault.azure.net/secrets/missing-value"