From 6941013c37fad8bdd672ebc5156a9cc81b6b501e Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 11 Aug 2026 11:03:39 -0400 Subject: [PATCH 01/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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" From 294fa0e0f5d766d29f7fad448c6ff04981c4b248 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 17 Aug 2026 13:59:10 -0400 Subject: [PATCH 17/28] FEAT: Integration tests for env drift --- .azuredevops/test-job-template.yml | 10 +- .env_example | 29 ------ .../test_akv_initialization_integration.py | 62 +++++++++++-- .../setup/test_env_example_drift.py | 91 +++++++++++++++++++ 4 files changed, 155 insertions(+), 37 deletions(-) create mode 100644 tests/integration/setup/test_env_example_drift.py diff --git a/.azuredevops/test-job-template.yml b/.azuredevops/test-job-template.yml index 65905af5cc..d469c52602 100644 --- a/.azuredevops/test-job-template.yml +++ b/.azuredevops/test-job-template.yml @@ -30,11 +30,11 @@ jobs: displayName: "Create PyRIT configuration directory" name: create_pyrit_dir - task: AzureKeyVault@2 - displayName: Azure Key Vault - retrieve .env file secret + displayName: Azure Key Vault - retrieve environment secrets inputs: azureSubscription: 'integration-test-service-connection' KeyVaultName: 'pyrit-environment' - SecretsFilter: 'env-global' + SecretsFilter: 'env-global,env-new' RunAsPreJob: false - bash: | python -c " @@ -92,10 +92,16 @@ jobs: cp -r $PyRIT_DIR/doc $NEW_DIR cp -r $PyRIT_DIR/assets $NEW_DIR cp -r $PyRIT_DIR/tests/${{ parameters.testsFolder }} $NEW_DIR/tests + cp $PyRIT_DIR/.env_example $NEW_DIR/.env_example cd $NEW_DIR displayName: "Create and switch to new test directory" - task: AzureCLI@2 displayName: "Authenticate with service principal, cache Cognitive Services access token, and run tests" + env: + PYRIT_AKV_INTEGRATION_TEST_ENV: $(env-new) + PYRIT_AKV_INTEGRATION_TEST_REQUIRED: 'true' + PYRIT_ENV_EXAMPLE_PATH: $(Build.SourcesDirectory)/../${{ parameters.newDir }}/.env_example + PYRIT_REPOSITORY_ROOT: $(Build.SourcesDirectory) inputs: azureSubscription: ${{ parameters.testAzureSubscription }} scriptType: 'bash' diff --git a/.env_example b/.env_example index 94629e75cd..c1097d99c4 100644 --- a/.env_example +++ b/.env_example @@ -49,7 +49,6 @@ AZURE_OPENAI_GPT4O_UNDERLYING_MODEL2="gpt-4o" 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 @@ -75,17 +74,12 @@ AZURE_OPENAI_GPT5_COMPLETIONS_ENDPOINT=" str: - """ - Get the explicitly configured integration bootstrap URL. +def _get_env_example_path() -> pathlib.Path: + configured_path = os.getenv(_ENV_EXAMPLE_PATH_ENV) + if configured_path: + path = pathlib.Path(configured_path) + if path.is_file(): + return path + raise AssertionError(f"{_ENV_EXAMPLE_PATH_ENV} does not identify a file: {path}") + + candidates = [pathlib.Path.cwd() / ".env_example"] + candidates.extend(parent / ".env_example" for parent in pathlib.Path(__file__).resolve().parents) + for path in candidates: + if path.is_file(): + return path + + raise AssertionError("Could not locate .env_example.") + + +def _get_akv_environment_keys() -> set[str]: + document = os.getenv(_AKV_ENVIRONMENT_ENV) + if not document: + if os.getenv(_AKV_ENVIRONMENT_REQUIRED_ENV, "").casefold() == "true": + raise AssertionError(f"{_AKV_ENVIRONMENT_ENV} is required but was not populated.") + pytest.skip(f"Set {_AKV_ENVIRONMENT_ENV} to run the AKV schema integration test.") + + keys = set(dotenv_values(stream=io.StringIO(document), interpolate=False)) + if not keys: + raise AssertionError("The env-new Key Vault secret contains no dotenv assignments.") + return keys + - Returns: - str: Key Vault bootstrap secret URL. - """ +def _get_env_example_keys() -> set[str]: + keys = set(dotenv_values(dotenv_path=_get_env_example_path(), interpolate=False)) + if not keys: + raise AssertionError(".env_example contains no dotenv assignments.") + return keys + + +def _get_bootstrap_secret_url() -> str: + """Get the explicitly configured integration bootstrap URL.""" configured_url = os.getenv(_SECRET_URL_ENV) if not configured_url: pytest.skip(f"Set {_SECRET_URL_ENV} to run this integration test.") @@ -49,3 +87,15 @@ async def test_akv_bootstrap_initialization_populates_process_environment() -> N os.environ.pop(variable_name, None) else: os.environ[variable_name] = str(original_value) + + +def test_akv_environment_keys_are_represented_in_env_example() -> None: + """Ensure the public example covers every environment name in the new AKV bootstrap document.""" + env_example_keys = _get_env_example_keys() + akv_environment_keys = _get_akv_environment_keys() + + missing_from_example = akv_environment_keys - env_example_keys + assert not missing_from_example, ( + "The env-new Key Vault secret contains names absent from .env_example: " + + ", ".join(sorted(missing_from_example)) + ) diff --git a/tests/integration/setup/test_env_example_drift.py b/tests/integration/setup/test_env_example_drift.py new file mode 100644 index 0000000000..763a78b410 --- /dev/null +++ b/tests/integration/setup/test_env_example_drift.py @@ -0,0 +1,91 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import os +import pathlib +import re +import subprocess + +from dotenv import dotenv_values + +_ENV_EXAMPLE_PATH_ENV = "PYRIT_ENV_EXAMPLE_PATH" +_REPOSITORY_ROOT_ENV = "PYRIT_REPOSITORY_ROOT" +_ENVIRONMENT_NAME_PATTERN = re.compile(r"(? pathlib.Path: + configured_root = os.getenv(_REPOSITORY_ROOT_ENV) + if configured_root: + root = pathlib.Path(configured_root) + if root.is_dir(): + return root + raise AssertionError(f"{_REPOSITORY_ROOT_ENV} does not identify a directory: {root}") + + result = subprocess.run( + ["git", "rev-parse", "--show-toplevel"], + capture_output=True, + check=False, + text=True, + ) + if result.returncode != 0: + raise AssertionError("Could not locate the repository root with git.") + return pathlib.Path(result.stdout.strip()) + + +def _get_env_example_path(*, repository_root: pathlib.Path) -> pathlib.Path: + configured_path = os.getenv(_ENV_EXAMPLE_PATH_ENV) + path = pathlib.Path(configured_path) if configured_path else repository_root / ".env_example" + if not path.is_file(): + raise AssertionError(f"Could not locate .env_example at {path}.") + return path + + +def _grep_repository_for_environment_names(*, environment_names: set[str], repository_root: pathlib.Path) -> set[str]: + grep_pattern = "(" + "|".join(sorted(environment_names)) + ")" + result = subprocess.run( + ["git", "grep", "-I", "-h", "-E", grep_pattern, "--", ".", ":(exclude).env_example"], + capture_output=True, + check=False, + cwd=repository_root, + text=True, + ) + if result.returncode not in {0, 1}: + raise AssertionError(f"Could not search tracked repository files with git: {result.stderr.strip()}") + return environment_names & set(_ENVIRONMENT_NAME_PATTERN.findall(result.stdout)) + + +def _find_referenced_environment_names( + *, + environment_names: set[str], + repository_root: pathlib.Path, + env_example_path: pathlib.Path, +) -> set[str]: + example_contents = env_example_path.read_text(encoding="utf-8") + example_without_assignment_names = _DOTENV_ASSIGNMENT_NAME_PATTERN.sub("=", example_contents) + referenced_names = environment_names & set(_ENVIRONMENT_NAME_PATTERN.findall(example_without_assignment_names)) + referenced_names.update( + _grep_repository_for_environment_names( + environment_names=environment_names, + repository_root=repository_root, + ) + ) + return referenced_names + + +def test_env_example_names_are_referenced_in_repository() -> None: + """Catch example entries with no weak textual reference in tracked repository files.""" + repository_root = _get_repository_root() + env_example_path = _get_env_example_path(repository_root=repository_root) + environment_names = set(dotenv_values(dotenv_path=env_example_path, interpolate=False)) + assert environment_names, ".env_example contains no dotenv assignments." + + referenced_names = _find_referenced_environment_names( + environment_names=environment_names, + repository_root=repository_root, + env_example_path=env_example_path, + ) + unreferenced_names = environment_names - referenced_names + assert not unreferenced_names, ".env_example contains names with no tracked repository reference: " + ", ".join( + sorted(unreferenced_names) + ) From 74f4bef6c889386aadcf15a6d74e5abd4aa8c642 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 17 Aug 2026 15:57:20 -0400 Subject: [PATCH 18/28] FIX: Outdated integration test --- .../test_akv_initialization_integration.py | 50 ++----------------- 1 file changed, 5 insertions(+), 45 deletions(-) diff --git a/tests/integration/setup/test_akv_initialization_integration.py b/tests/integration/setup/test_akv_initialization_integration.py index 6428776823..0bab039e6b 100644 --- a/tests/integration/setup/test_akv_initialization_integration.py +++ b/tests/integration/setup/test_akv_initialization_integration.py @@ -8,14 +8,9 @@ import pytest from dotenv import dotenv_values -from pyrit.setup import IN_MEMORY, initialize_pyrit_async - _AKV_ENVIRONMENT_ENV = "PYRIT_AKV_INTEGRATION_TEST_ENV" _AKV_ENVIRONMENT_REQUIRED_ENV = "PYRIT_AKV_INTEGRATION_TEST_REQUIRED" _ENV_EXAMPLE_PATH_ENV = "PYRIT_ENV_EXAMPLE_PATH" -_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_env_example_path() -> pathlib.Path: @@ -55,47 +50,12 @@ def _get_env_example_keys() -> set[str]: return keys -def _get_bootstrap_secret_url() -> str: - """Get the explicitly configured integration bootstrap 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) - - -def test_akv_environment_keys_are_represented_in_env_example() -> None: - """Ensure the public example covers every environment name in the new AKV bootstrap document.""" +def test_env_example_keys_are_represented_in_akv_environment() -> None: + """Ensure the new AKV bootstrap document covers every environment name in the public example.""" env_example_keys = _get_env_example_keys() akv_environment_keys = _get_akv_environment_keys() - missing_from_example = akv_environment_keys - env_example_keys - assert not missing_from_example, ( - "The env-new Key Vault secret contains names absent from .env_example: " - + ", ".join(sorted(missing_from_example)) + missing_from_akv = env_example_keys - akv_environment_keys + assert not missing_from_akv, "The env-new Key Vault secret is missing names defined in .env_example: " + ", ".join( + sorted(missing_from_akv) ) From b635045346fd86401b574c58c711d43b60b5d3f8 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 18 Aug 2026 14:32:27 -0400 Subject: [PATCH 19/28] FIX: Updates to tests and atomic file writing --- .env_example | 129 +++++++++++------- pyrit/setup/akv_initialization.py | 41 +++++- .../setup/test_env_example_drift.py | 42 ++++++ tests/unit/setup/test_akv_initialization.py | 80 +++++++++++ 4 files changed, 238 insertions(+), 54 deletions(-) diff --git a/.env_example b/.env_example index c1097d99c4..c604652c88 100644 --- a/.env_example +++ b/.env_example @@ -1,46 +1,84 @@ # ============================================================================ + # 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" + +# + +# URL placeholders intentionally use plain dotenv values. Earlier versions used + +# angle brackets as documentation styling, but python-dotenv preserves them literally + # + # 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="" @@ -48,9 +86,11 @@ AZURE_OPENAI_GPT4O_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNDERLYING_MODEL2="gpt-4o" AZURE_OPENAI_GPT4O_AAD_ENDPOINT="" +AZURE_OPENAI_GPT4O_AAD_KEY="xxxxx" AZURE_OPENAI_GPT4O_AAD_MODEL="deployment-name" # 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="" @@ -121,14 +161,23 @@ AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="" AZURE_OPENAI_GPT5_MODEL="gpt-5" AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" +PLATFORM_OPENAI_CHAT_ENDPOINT="" +PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" +PLATFORM_OPENAI_CHAT_MODEL="gpt-4o" +PLATFORM_OPENAI_RESPONSES_ENDPOINT="" +PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" +PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" + DEFAULT_OPENAI_FRONTEND_ENDPOINT = ${AZURE_OPENAI_GPT4O_AAD_ENDPOINT} DEFAULT_OPENAI_FRONTEND_KEY = ${AZURE_OPENAI_GPT4O_AAD_KEY} 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="" @@ -144,9 +193,15 @@ OPENAI_RESPONSES_UNDERLYING_MODEL="" ################################## 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" +PLATFORM_OPENAI_REALTIME_ENDPOINT="wss://api.openai.com/v1" +PLATFORM_OPENAI_REALTIME_KEY="sk-xxxxx" +PLATFORM_OPENAI_REALTIME_API_KEY="sk-xxxxx" +PLATFORM_OPENAI_REALTIME_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 = "" @@ -156,7 +211,16 @@ OPENAI_REALTIME_UNDERLYING_MODEL = "" ################################## +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_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 = "" @@ -168,7 +232,16 @@ OPENAI_IMAGE_STRICT_FILTER_MODEL = "gpt-image" ################################## +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_ENDPOINT = ${AZURE_OPENAI_TTS_ENDPOINT2} +OPENAI_TTS_KEY = ${AZURE_OPENAI_TTS_KEY2} OPENAI_TTS_MODEL = ${AZURE_OPENAI_TTS_MODEL2} OPENAI_TTS_UNDERLYING_MODEL = "" @@ -179,14 +252,19 @@ 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="" +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 = "" @@ -197,6 +275,7 @@ 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="" @@ -218,6 +297,7 @@ 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="" @@ -286,13 +366,6 @@ AWS_ENDPOINT="" AWS_RESPONSES_MODEL="openai.gpt-oss-120b" AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" -PLATFORM_OPENAI_CHAT_ENDPOINT="" -PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" -PLATFORM_OPENAI_CHAT_MODEL="gpt-4o" -PLATFORM_OPENAI_RESPONSES_ENDPOINT="" -PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" -PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" - PLATFORM_OPENAI_VIDEO_ENDPOINT="" PLATFORM_OPENAI_VIDEO_KEY="sk-xxxxx" PLATFORM_OPENAI_VIDEO_MODEL="sora-2" @@ -305,14 +378,6 @@ PLATFORM_OPENAI_EMBEDDING_ENDPOINT="" 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" -PLATFORM_OPENAI_REALTIME_KEY="sk-xxxxx" -PLATFORM_OPENAI_REALTIME_MODEL="gpt-4o-realtime-preview" - OPENAI_COMPLETION_ENDPOINT="" OPENAI_COMPLETION_API_KEY="xxxxx" OPENAI_COMPLETION_MODEL="davinci-002" @@ -342,42 +407,8 @@ GROQ_LLAMA_MODEL="llama3-8b-8192" OPEN_ROUTER_ENDPOINT="" OPEN_ROUTER_KEY="sk-or-v1-xxxxx" OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" -OPENAI_CHAT_KEY=${PLATFORM_OPENAI_CHAT_KEY} AZURE_OPENAI_GPT5_KEY="xxxxxxx" AZURE_OPENAI_RESPONSES_KEY="xxxxx" -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 = "" -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_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 = "" -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_KEY = ${AZURE_OPENAI_TTS_KEY2} -AZURE_OPENAI_VIDEO_KEY="xxxxxxx" -OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_EMBEDDING_KEY="xxxxx" AZURE_STORAGE_ACCOUNT_SAS_TOKEN="xxxxx" AZURE_SPEECH_KEY = "xxxxx" diff --git a/pyrit/setup/akv_initialization.py b/pyrit/setup/akv_initialization.py index c0d48d7323..9941375ca2 100644 --- a/pyrit/setup/akv_initialization.py +++ b/pyrit/setup/akv_initialization.py @@ -4,10 +4,12 @@ """Load dotenv files and Azure Key Vault-backed environment documents.""" import asyncio +import contextlib import io import logging import os import pathlib +import tempfile import urllib.parse from collections.abc import Sequence from typing import TYPE_CHECKING, Any @@ -485,15 +487,44 @@ def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Pa Returns: pathlib.Path: Path to the written dotenv file. + + Raises: + ValueError: If the destination is a symbolic link. """ env_file = path.CONFIGURATION_DIRECTORY_PATH / _AKV_ENV_FILE_NAME - env_file.parent.mkdir(parents=True, exist_ok=True) + env_file.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + if env_file.is_symlink(): + raise ValueError(f"Refusing to write the AKV environment through a symbolic link: {env_file}") + content = "\n".join(document.rstrip("\r\n") for document in documents) + "\n" - env_file.write_text(content, encoding="utf-8") + file_descriptor: int | None = None + temporary_file: pathlib.Path | None = None try: - env_file.chmod(0o600) - except OSError: - logger.warning("Could not restrict permissions on written AKV environment file: %s", env_file) + file_descriptor, temporary_name = tempfile.mkstemp( + prefix=f"{env_file.name}.", + suffix=".tmp", + dir=env_file.parent, + ) + temporary_file = pathlib.Path(temporary_name) + if hasattr(os, "fchmod"): + os.fchmod(file_descriptor, 0o600) + else: + os.chmod(temporary_file, 0o600) + stream = os.fdopen(file_descriptor, "w", encoding="utf-8", newline="") + file_descriptor = None + with stream: + stream.write(content) + if env_file.is_symlink(): + raise ValueError(f"Refusing to replace a symbolic link with the AKV environment: {env_file}") + os.replace(temporary_file, env_file) + temporary_file = None + finally: + if file_descriptor is not None: + os.close(file_descriptor) + if temporary_file is not None: + with contextlib.suppress(FileNotFoundError): + temporary_file.unlink() + _print_msg(f"Saved Key Vault bootstrap environment file: {env_file}", quiet=silent, log=True) return env_file diff --git a/tests/integration/setup/test_env_example_drift.py b/tests/integration/setup/test_env_example_drift.py index 763a78b410..81913948db 100644 --- a/tests/integration/setup/test_env_example_drift.py +++ b/tests/integration/setup/test_env_example_drift.py @@ -5,6 +5,7 @@ import pathlib import re import subprocess +from unittest import mock from dotenv import dotenv_values @@ -12,6 +13,7 @@ _REPOSITORY_ROOT_ENV = "PYRIT_REPOSITORY_ROOT" _ENVIRONMENT_NAME_PATTERN = re.compile(r"(? pathlib.Path: @@ -89,3 +91,43 @@ def test_env_example_names_are_referenced_in_repository() -> None: assert not unreferenced_names, ".env_example contains names with no tracked repository reference: " + ", ".join( sorted(unreferenced_names) ) + + +def test_env_example_url_values_are_not_wrapped_in_angle_brackets() -> None: + """Ensure URL placeholder styling does not become part of parsed dotenv values.""" + repository_root = _get_repository_root() + env_example_path = _get_env_example_path(repository_root=repository_root) + values = dotenv_values(dotenv_path=env_example_path, interpolate=False) + + wrapped_names = {name for name, value in values.items() if value and ("<" in value or ">" in value)} + assert not wrapped_names, ".env_example contains values wrapped in angle brackets: " + ", ".join( + sorted(wrapped_names) + ) + + +def test_env_example_aliases_resolve_in_assignment_order() -> None: + """Ensure complete-value aliases resolve to their sources without ambient environment values.""" + repository_root = _get_repository_root() + env_example_path = _get_env_example_path(repository_root=repository_root) + raw_values = dotenv_values(dotenv_path=env_example_path, interpolate=False) + aliases = { + name: match.group(1) + for name, value in raw_values.items() + if value and (match := _DOTENV_COMPLETE_REFERENCE_PATTERN.fullmatch(value)) + } + assert aliases, ".env_example contains no complete-value aliases." + + with mock.patch.dict(os.environ, {}, clear=True): + resolved_values = dotenv_values(dotenv_path=env_example_path, interpolate=True) + + unresolved_names = {name for name in aliases if not resolved_values.get(name)} + assert not unresolved_names, ".env_example contains aliases that resolve to empty values: " + ", ".join( + sorted(unresolved_names) + ) + + mismatched_names = { + name for name, source_name in aliases.items() if resolved_values[name] != resolved_values.get(source_name) + } + assert not mismatched_names, ".env_example contains aliases that differ from their sources: " + ", ".join( + sorted(mismatched_names) + ) diff --git a/tests/unit/setup/test_akv_initialization.py b/tests/unit/setup/test_akv_initialization.py index 81f80b92ea..0cd51db544 100644 --- a/tests/unit/setup/test_akv_initialization.py +++ b/tests/unit/setup/test_akv_initialization.py @@ -19,6 +19,7 @@ _parse_akv_reference, _parse_akv_secret_url, _warn_about_akv_environment_files, + _write_akv_env_file, ) @@ -250,6 +251,85 @@ async def test_load_environment_async_write_env_filters_generated_explicit_file( 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 + def test_write_akv_env_file_secures_descriptor_before_writing(self): + events: list[str] = [] + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + temporary_file = temp_path / ".env.test.tmp" + stream = mock.MagicMock() + stream.write.side_effect = lambda content: events.append(f"write:{content}") + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch( + "pyrit.setup.akv_initialization.tempfile.mkstemp", + side_effect=lambda **kwargs: events.append("create") or (7, str(temporary_file)), + ), + mock.patch( + "pyrit.setup.akv_initialization.os.fchmod", + side_effect=lambda *args: events.append("fchmod"), + create=True, + ), + mock.patch( + "pyrit.setup.akv_initialization.os.fdopen", + side_effect=lambda *args, **kwargs: events.append("fdopen") or stream, + ), + mock.patch( + "pyrit.setup.akv_initialization.os.replace", + side_effect=lambda *args: events.append("replace"), + ), + ): + _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) + + assert events == ["create", "fchmod", "fdopen", "write:VALUE=bootstrap\n", "replace"] + + def test_write_akv_env_file_preserves_existing_file_when_replace_fails(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_file.write_text("ORIGINAL=value\n", encoding="utf-8") + + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.akv_initialization.os.replace", side_effect=OSError("replace failed")), + pytest.raises(OSError, match="replace failed"), + ): + _write_akv_env_file(documents=["NEW=value\n"], silent=True) + + assert env_file.read_text(encoding="utf-8") == "ORIGINAL=value\n" + assert list(temp_path.glob(".env.*.tmp")) == [] + + @pytest.mark.skipif(os.name != "posix", reason="POSIX permission bits are not enforced on this platform.") + def test_write_akv_env_file_uses_owner_only_permissions(self): + with tempfile.TemporaryDirectory() as temp_dir: + configuration_directory = pathlib.Path(temp_dir) / ".pyrit" + with mock.patch( + "pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", configuration_directory + ): + env_file = _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) + + assert configuration_directory.stat().st_mode & 0o777 == 0o700 + assert env_file.stat().st_mode & 0o777 == 0o600 + + def test_write_akv_env_file_rejects_symbolic_link(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + target = temp_path / "target" + target.write_text("unchanged", encoding="utf-8") + env_file = temp_path / ".env" + try: + env_file.symlink_to(target) + except OSError: + pytest.skip("Symbolic links are unavailable on this platform.") + + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + pytest.raises(ValueError, match="symbolic link"), + ): + _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) + + assert target.read_text(encoding="utf-8") == "unchanged" + 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" From 3df3f625c3fc9b9478af042a328a202900cd142c Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 19 Aug 2026 11:58:39 -0400 Subject: [PATCH 20/28] FIX: Incorporating latest feedback --- .azuredevops/test-job-template.yml | 34 +- .env_example | 171 ++----- .pyrit_conf_example | 14 +- doc/getting_started/pyrit_conf.md | 78 ++- pyrit/setup/akv_initialization.py | 378 +++++++++++--- pyrit/setup/configuration_loader.py | 7 +- pyrit/setup/initialization.py | 16 +- .../test_akv_initialization_integration.py | 61 --- .../setup/test_env_example_drift.py | 12 + tests/unit/setup/test_akv_initialization.py | 476 +++++++++++++++--- tests/unit/setup/test_initialization.py | 21 +- 11 files changed, 897 insertions(+), 371 deletions(-) delete mode 100644 tests/integration/setup/test_akv_initialization_integration.py diff --git a/.azuredevops/test-job-template.yml b/.azuredevops/test-job-template.yml index d469c52602..7ed5c8ef5e 100644 --- a/.azuredevops/test-job-template.yml +++ b/.azuredevops/test-job-template.yml @@ -26,7 +26,7 @@ jobs: versionSpec: '3.12' addToPath: true - bash: | - mkdir -p ~/.pyrit + install -d -m 700 ~/.pyrit displayName: "Create PyRIT configuration directory" name: create_pyrit_dir - task: AzureKeyVault@2 @@ -34,16 +34,32 @@ jobs: inputs: azureSubscription: 'integration-test-service-connection' KeyVaultName: 'pyrit-environment' - SecretsFilter: 'env-global,env-new' + SecretsFilter: 'env-global' RunAsPreJob: false - bash: | - python -c " - import os; - secret = os.environ.get('PYRIT_TEST_SECRET'); + python - <<'PY' + import os + import pathlib + import tempfile + + secret = os.environ.get("PYRIT_TEST_SECRET") if not secret: - raise ValueError('PYRIT_TEST_SECRET is not set'); - with open(os.path.expanduser('~/.pyrit/.env'), 'w') as file: - file.write(secret)" + raise ValueError("PYRIT_TEST_SECRET is not set") + + env_file = pathlib.Path.home() / ".pyrit" / ".env" + file_descriptor, temporary_name = tempfile.mkstemp(prefix=".env.", suffix=".tmp", dir=env_file.parent) + temporary_file = pathlib.Path(temporary_name) + try: + os.fchmod(file_descriptor, 0o600) + with os.fdopen(file_descriptor, "w", encoding="utf-8", newline="") as stream: + file_descriptor = -1 + stream.write(secret) + os.replace(temporary_file, env_file) + finally: + if file_descriptor >= 0: + os.close(file_descriptor) + temporary_file.unlink(missing_ok=True) + PY env: PYRIT_TEST_SECRET: $(env-global) name: create_env_file @@ -98,8 +114,6 @@ jobs: - task: AzureCLI@2 displayName: "Authenticate with service principal, cache Cognitive Services access token, and run tests" env: - PYRIT_AKV_INTEGRATION_TEST_ENV: $(env-new) - PYRIT_AKV_INTEGRATION_TEST_REQUIRED: 'true' PYRIT_ENV_EXAMPLE_PATH: $(Build.SourcesDirectory)/../${{ parameters.newDir }}/.env_example PYRIT_REPOSITORY_ROOT: $(Build.SourcesDirectory) inputs: diff --git a/.env_example b/.env_example index c604652c88..34468eba8c 100644 --- a/.env_example +++ b/.env_example @@ -1,170 +1,133 @@ # ============================================================================ - # 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_ENDPOINT="https://api.openai.com/v1" # or any OpenAI-compatible API # OPENAI_CHAT_KEY="your-key-here" - # OPENAI_CHAT_MODEL="gpt-4o" - # - # URL placeholders intentionally use plain dotenv values. Earlier versions used - # angle brackets as documentation styling, but python-dotenv preserves them literally - # - # 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: https://xxxx.openai.azure.com/openai/v1 -# Example: - -AZURE_OPENAI_GPT4O_ENDPOINT="" +AZURE_OPENAI_GPT4O_ENDPOINT="https://xxxx.openai.azure.com/openai/v1" 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_ENDPOINT2="https://xxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4O_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNDERLYING_MODEL2="gpt-4o" -AZURE_OPENAI_GPT4O_AAD_ENDPOINT="" +AZURE_OPENAI_GPT4O_AAD_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4O_AAD_KEY="xxxxx" AZURE_OPENAI_GPT4O_AAD_MODEL="deployment-name" # 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_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" 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_ENDPOINT2="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" # Objective Scorer chat target (used in scorers in scenarios) -OBJECTIVE_SCORER_CHAT_ENDPOINT="" +OBJECTIVE_SCORER_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" OBJECTIVE_SCORER_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="" +AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT5_COMPLETIONS_ENDPOINT="" +AZURE_OPENAI_GPT5_COMPLETIONS_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT5_COMPLETIONS_MODEL="gpt-5" AZURE_OPENAI_GPT5_COMPLETIONS_UNDERLYING_MODEL="gpt-5" -AZURE_OPENAI_GPT5_4_ENDPOINT="" +AZURE_OPENAI_GPT5_4_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" -AZURE_OPENAI_GPT4O_STRICT_FILTER_ENDPOINT="" +AZURE_OPENAI_GPT4O_STRICT_FILTER_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4O_STRICT_FILTER_MODEL="deployment-name" -AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="" +AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT3_5_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT3_5_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT4_CHAT_ENDPOINT="" +AZURE_OPENAI_GPT4_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4_CHAT_UNDERLYING_MODEL="" -MAI_CHAT_ENDPOINT="" +MAI_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" MAI_CHAT_MODEL="deployment-name" MAI_CHAT_KEY="xxxxx" MAI_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPTV_CHAT_ENDPOINT="" +AZURE_OPENAI_GPTV_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPTV_CHAT_MODEL="deployment-name" -AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="" +AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="https://xxxxx.eastus2.models.ai.azure.com" AZURE_FOUNDRY_DEEPSEEK_KEY="xxxxx" AZURE_FOUNDRY_DEEPSEEK_MODEL="" -AZURE_FOUNDRY_PHI4_ENDPOINT="" +AZURE_FOUNDRY_PHI4_ENDPOINT="https://xxxxx.models.ai.azure.com" AZURE_CHAT_PHI4_KEY="xxxxx" AZURE_CHAT_PHI4_MODEL="" -AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="" +AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="https://xxxxx.services.ai.azure.com/openai/v1/" AZURE_FOUNDRY_MISTRAL_LARGE_MODEL="Mistral-Large-3" -OLLAMA_CHAT_ENDPOINT="" +OLLAMA_CHAT_ENDPOINT="http://127.0.0.1:11434/v1" OLLAMA_MODEL="llama2" -AZURE_OPENAI_RESPONSES_ENDPOINT="" +AZURE_OPENAI_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_RESPONSES_MODEL="o4-mini" AZURE_OPENAI_RESPONSES_UNDERLYING_MODEL="o4-mini" -AZURE_OPENAI_O4MINI_STRICT_FILTER_RESPONSES_ENDPOINT="" +AZURE_OPENAI_O4MINI_STRICT_FILTER_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_O4MINI_STRICT_FILTER_RESPONSES_MODEL="o4-mini" -AZURE_OPENAI_GPT41_RESPONSES_ENDPOINT="" +AZURE_OPENAI_GPT41_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT41_RESPONSES_MODEL="gpt-4.1" -AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="" +AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" AZURE_OPENAI_GPT5_MODEL="gpt-5" AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" -PLATFORM_OPENAI_CHAT_ENDPOINT="" +PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" PLATFORM_OPENAI_CHAT_MODEL="gpt-4o" -PLATFORM_OPENAI_RESPONSES_ENDPOINT="" +PLATFORM_OPENAI_RESPONSES_ENDPOINT="https://api.openai.com/v1" PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" @@ -177,7 +140,6 @@ 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="" @@ -187,9 +149,7 @@ OPENAI_RESPONSES_MODEL=${PLATFORM_OPENAI_RESPONSES_MODEL} OPENAI_RESPONSES_UNDERLYING_MODEL="" ################################## - # OPENAI REALTIME TARGET SECRETS - ################################## AZURE_OPENAI_REALTIME_ENDPOINT = "wss://xxxx.openai.azure.com/openai/v1" @@ -206,16 +166,14 @@ OPENAI_REALTIME_MODEL = ${PLATFORM_OPENAI_REALTIME_MODEL} OPENAI_REALTIME_UNDERLYING_MODEL = "" ################################## - # IMAGE TARGET SECRETS - ################################## -AZURE_OPENAI_IMAGE_ENDPOINT1 = "" +AZURE_OPENAI_IMAGE_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" 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_ENDPOINT2 = "https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_IMAGE_API_KEY2 = "xxxxxx" AZURE_OPENAI_IMAGE_MODEL2 = "dall-e-3" AZURE_OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" @@ -223,20 +181,18 @@ 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_ENDPOINT = "https://xxxxx.openai.azure.com/openai/v1" OPENAI_IMAGE_STRICT_FILTER_MODEL = "gpt-image" ################################## - # TTS TARGET SECRETS - ################################## -AZURE_OPENAI_TTS_ENDPOINT1 = "" +AZURE_OPENAI_TTS_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_TTS_KEY1 = "xxxxxxx" AZURE_OPENAI_TTS_MODEL1 = "tts" AZURE_OPENAI_TTS_UNDERLYING_MODEL1 = "tts" -AZURE_OPENAI_TTS_ENDPOINT2 = "" +AZURE_OPENAI_TTS_ENDPOINT2 = "https://xxxxx.openai.azure.com/v1" AZURE_OPENAI_TTS_KEY2 = "xxxxxx" AZURE_OPENAI_TTS_MODEL2 = "tts-1" AZURE_OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" @@ -246,20 +202,14 @@ OPENAI_TTS_MODEL = ${AZURE_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="" +AZURE_OPENAI_VIDEO_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/openai/v1" AZURE_OPENAI_VIDEO_KEY="xxxxxxx" AZURE_OPENAI_VIDEO_MODEL="sora-2" AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" @@ -269,38 +219,30 @@ 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_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" ADVERSARIAL_CHAT_MODEL="deployment-name" -ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="" +ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" ADVERSARIAL_CHAT_SINGLETURN_MODEL="deployment-name" -ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="" +ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" ADVERSARIAL_CHAT_MULTITURN_KEY="xxxxx" ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" -ADVERSARIAL_CHAT_REASONING_ENDPOINT="" +ADVERSARIAL_CHAT_REASONING_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" 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_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" AZURE_ML_PHI_KEY="xxxxx" # The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed @@ -309,38 +251,33 @@ AZURE_ML_MANAGED_ENDPOINT=${AZURE_ML_PHI_ENDPOINT} AZURE_ML_KEY=${AZURE_ML_PHI_KEY} ################################## - # MISC TARGET SECRETS - ################################## -OPENAI_EMBEDDING_ENDPOINT="" +OPENAI_EMBEDDING_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" OPENAI_EMBEDDING_MODEL="text-embedding-3-small" AZURE_SPEECH_REGION = "eastus2" # Resource ID is needed when using Entra authentication AZURE_SPEECH_RESOURCE_ID = "xxxxx" -AZURE_CONTENT_SAFETY_API_ENDPOINT="" +AZURE_CONTENT_SAFETY_API_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/" HUGGINGFACE_TOKEN="hf_xxxxxxx" -HUGGINGFACE_ENDPOINT="" +HUGGINGFACE_ENDPOINT="https://router.huggingface.co/v1" ################################## - # 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="" +AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="https://xxxxx.blob.core.windows.net/dbdata" # 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="" -AZURE_STORAGE_ACCOUNT_CONTAINER_URL="" +AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="https://xxxxx.blob.core.windows.net/dbdata" +AZURE_STORAGE_ACCOUNT_CONTAINER_URL="https://xxxxxx.blob.core.windows.net/xpia" # The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local @@ -348,46 +285,42 @@ 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} ################################## - # INTEGRATION TEST ONLY SECRETS - ################################## -GOOGLE_GEMINI_ENDPOINT = "" +GOOGLE_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/openai" GOOGLE_GEMINI_API_KEY = "xxxxx" GOOGLE_GEMINI_MODEL="gemini-2.0-flash" -ANTHROPIC_CHAT_ENDPOINT="" +ANTHROPIC_CHAT_ENDPOINT="https://api.anthropic.com/v1" ANTHROPIC_CHAT_KEY="xxxxx" ANTHROPIC_CHAT_MODEL="claude-3-7-sonnet-latest" AWS_KEY="xxxxx" -AWS_ENDPOINT="" +AWS_ENDPOINT="https://bedrock-mantle.us-east-1.api.aws/v1" AWS_RESPONSES_MODEL="openai.gpt-oss-120b" AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" -PLATFORM_OPENAI_VIDEO_ENDPOINT="" +PLATFORM_OPENAI_VIDEO_ENDPOINT="https://api.openai.com/v1" PLATFORM_OPENAI_VIDEO_KEY="sk-xxxxx" PLATFORM_OPENAI_VIDEO_MODEL="sora-2" -PLATFORM_OPENAI_IMAGE_ENDPOINT="" +PLATFORM_OPENAI_IMAGE_ENDPOINT="https://api.openai.com/v1" PLATFORM_OPENAI_IMAGE_KEY="sk-xxxxx" PLATFORM_OPENAI_IMAGE_MODEL="gpt-image-1" -PLATFORM_OPENAI_EMBEDDING_ENDPOINT="" +PLATFORM_OPENAI_EMBEDDING_ENDPOINT="https://api.openai.com/v1" PLATFORM_OPENAI_EMBEDDING_KEY="sk-xxxxx" PLATFORM_OPENAI_EMBEDDING_MODEL="text-embedding-3-small" -OPENAI_COMPLETION_ENDPOINT="" +OPENAI_COMPLETION_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" OPENAI_COMPLETION_API_KEY="xxxxx" OPENAI_COMPLETION_MODEL="davinci-002" PROMPTINTEL_API_KEY="xxxxx" ################################## - # Additional entries referenced in PyRIT - ################################## AZURE_OPENAI_GPT4O_KEY="xxxxx" @@ -401,10 +334,10 @@ AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" ADVERSARIAL_CHAT_KEY="xxxxx" OBJECTIVE_SCORER_CHAT_KEY="xxxxx" AZURE_FOUNDRY_MISTRAL_LARGE_KEY="xxxxx" -GROQ_ENDPOINT="" +GROQ_ENDPOINT="https://api.groq.com/openai/v1" GROQ_KEY="gsk_xxxxxxxx" GROQ_LLAMA_MODEL="llama3-8b-8192" -OPEN_ROUTER_ENDPOINT="" +OPEN_ROUTER_ENDPOINT="https://openrouter.ai/api/v1" OPEN_ROUTER_KEY="sk-or-v1-xxxxx" OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" AZURE_OPENAI_GPT5_KEY="xxxxxxx" diff --git a/.pyrit_conf_example b/.pyrit_conf_example index a9c1ae8482..389a9b4a5c 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -4,7 +4,7 @@ # or specify a custom path when loading via --config-file. # # For documentation on configuration options, see: -# https://github.com/microsoft/PyRIT/blob/main/doc/setup/configuration.md +# https://github.com/microsoft/PyRIT/blob/main/doc/getting_started/pyrit_conf.md # Memory Database Type # -------------------- @@ -81,18 +81,20 @@ operation: op_trash_panda # Environment Configuration # ------------------------- -# Azure Key Vault is recommended for shared and deployed configurations. +# Azure Key Vault is the canonical source for shared and deployed configuration. # See doc/getting_started/pyrit_conf.md for loading order, references, and migration guidance. # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env # env_akv_strict: true -# env_akv_write_env: false # Opt in to writing ~/.pyrit/.env for inspection. +# env_akv_write_env: false # Debug only: write a fully resolved, sensitive ~/.pyrit/.env. -# 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. +# Auto-discovered ~/.pyrit/.env is legacy and will be rejected in PyRIT 1.3.0. +# Use ~/.pyrit/.env.local for quick local plaintext patches or when Azure is unavailable. +# Process values remain authoritative; AKV and ordinary env_files fill gaps in load order. +# Only a file named .env.local overrides existing values. +# Explicit env_files remain supported regardless of name or location and may contain full kv: URLs. # env_files: # - /path/to/.env.local -# 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 82fd52dd84..fab8cacf7b 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -4,13 +4,14 @@ The recommended way to configure PyRIT. A `.pyrit_conf` file declares your datab ## Quick Setup -```bash -mkdir -p ~/.pyrit -cp .pyrit_conf_example ~/.pyrit/.pyrit_conf -cp .env_example ~/.pyrit/.env +Create `~/.pyrit/.pyrit_conf` and configure a Key Vault bootstrap document: + +```yaml +env_akv_ref: + - https://my-vault.vault.azure.net/secrets/my-pyrit-env ``` -Then edit both files for your environment. The `.pyrit_conf` tells PyRIT _how_ to initialize; the `.env` tells it _where_ your targets are. +The Key Vault secret value uses dotenv syntax. When Azure is unavailable or you need a quick local patch, put only those values in `~/.pyrit/.env.local`. ## File Location @@ -25,44 +26,34 @@ PyRIT looks for this file automatically on startup (via the CLI, shell, or `Conf ## Environment Configuration ```{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. +Azure Key Vault is PyRIT's canonical environment source for shared, CI/CD, and deployed configuration. Auto-discovered `~/.pyrit/.env` is supported only as a legacy source and will be rejected in PyRIT 1.3.0. Use `~/.pyrit/.env.local` for deliberate plaintext local iteration or when Azure is unavailable. ``` See [Populating Secrets](./populating_secrets.md) for provider-specific variable examples. ### Loading Order -PyRIT loads environment sources in this order. Each later source overrides matching values from earlier sources: +PyRIT loads environment sources in this order: 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`. - -For a Key Vault-only setup, explicitly disable local dotenv loading: +2. Key Vault bootstrap documents, legacy auto-discovered `.env`, or explicit `env_files`. These sources fill only missing values. +3. Files named `.env.local`. These are the only dotenv sources that override existing values. -```yaml -env_akv_ref: - - https://my-vault.vault.azure.net/secrets/my-pyrit-env -env_files: [] -``` - -If Key Vault and local files are both configured, PyRIT warns that local values may override the fetched configuration. Remove stale local files when Key Vault should be authoritative. +When `env_akv_ref` is configured, PyRIT ignores an auto-discovered `~/.pyrit/.env`, emits its deprecation warning, and still loads `~/.pyrit/.env.local`. Explicit `env_files` are never blocked or deprecated based on their filename or location. ### 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: +Use `~/.pyrit/.env.local` to override process or Key Vault values deliberately. This is useful for: - Testing different targets - Using personal credentials instead of shared ones - Switching between configurations quickly -Simply create `.env.local` in your `~/.pyrit/` directory and add any variables you want to override. +Only put the values you need to patch in this file. Because it contains plaintext secrets, do not commit it. ### Authentication Options -**API Keys (Default):** The simplest approach — set `OPENAI_CHAT_KEY` and similar variables in your `.env` file. Most targets support this method. +**API keys:** Store shared API keys as Key Vault scalar secrets and reference them from the bootstrap document with `kv:`. For local-only work, place them in `.env.local`. **Azure Entra Authentication (Optional):** For Azure resources, you can use Entra auth instead of API keys. This requires the [Azure CLI](https://learn.microsoft.com/en-us/cli/azure/install-azure-cli) and `az login`. When using Entra auth, you don't need to set API keys for Azure resources. @@ -160,11 +151,11 @@ initialization_scripts: ### `env_files` -Optional local dotenv paths. Key Vault is recommended for shared or deployed configuration; use local files for backward compatibility and deliberate local overrides. +Optional local dotenv paths. Key Vault remains the canonical shared source; explicit files support local and non-Azure workflows. | Value | Behavior | | ----------------- | -------------------------------------------------------- | -| Omitted or `null` | Load default `~/.pyrit/.env` and `~/.pyrit/.env.local` | +| Omitted or `null` | Auto-discover legacy `~/.pyrit/.env` and supported `~/.pyrit/.env.local` | | `[]` (empty list) | Load **no** environment files | | List of paths | Load **only** the specified files (defaults are skipped) | @@ -174,9 +165,11 @@ env_files: - /path/to/.env.local ``` -Local environment files use standard python-dotenv parsing and `${NAME}` interpolation. Interpolation follows assignment and file load order. The default `.env.local` can reference a value loaded earlier from `.env`, for example `FOOBAR=${OPENAI_CHAT_ENDPOINT}`. A `.env` value cannot reference a variable introduced only by the later `.env.local`; values are not resolved retroactively. Explicit `env_files` follow the order in which they are listed. +Local files use standard python-dotenv parsing and `${NAME}` interpolation. Ordinary files fill missing values; any file whose basename is `.env.local` overrides existing values. Explicit files load in their listed order. + +Complete-value `kv:`, `akv:`, `azure_key_vault:`, and `env_akv_ref:` references resolve in local files as well as remote bootstrap documents. Local references may use any validated supported Key Vault URL; remote child references must remain in the bootstrap document's vault. A local assignment that loses to an existing value does not fetch its secret. -`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. +Ordinary malformed dotenv lines retain python-dotenv's permissive behavior. `env_akv_strict` controls malformed Key Vault reference syntax in all sources: strict mode raises; non-strict mode warns and skips that assignment. Authentication, authorization, transport, missing-secret, and missing-value failures always raise. 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. @@ -184,7 +177,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. This is the recommended configuration path. 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 canonical configuration path. Each secret value contains dotenv-formatted entries. Authentication uses `DefaultAzureCredential`. ```yaml env_akv_ref: @@ -192,7 +185,7 @@ env_akv_ref: - https://my-vault.vault.azure.net/secrets/team-pyrit-env ``` -Bootstrap documents load in list order with `override=True`; local environment files load afterward. Each document uses native dotenv interpolation against the process environment and assignments already parsed. A bootstrap document can mix literal values, `${NAME}` interpolation, and complete-value references to scalar secrets in the same vault: +Bootstrap documents load in list order and fill values missing from the process environment. 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" @@ -212,40 +205,40 @@ References must occupy the entire value. `kv:` is the canonical Key Vault prefix A Key Vault reference must use a full HTTPS secret URL from the bootstrap document's vault. Supported vault DNS suffixes are `.vault.azure.net`, `.vault.azure.cn`, and `.vault.usgovcloudapi.net`. An unversioned URL reads the latest secret version at initialization. Include the version in the URL to pin it. Short names, malformed paths, arbitrary hosts, and cross-vault child references are rejected before a client is created. -PyRIT does not cache referenced secrets. Each `kv:` occurrence in a bootstrap document performs a Key Vault read during initialization, including repeated references to the same URI. A later bootstrap or local file may override a reference after it has already been fetched. +PyRIT does not cache referenced secrets. Each winning `kv:` occurrence performs a Key Vault read during initialization. References that lose to an existing process or earlier source are not fetched. Debug output is the exception: it resolves bootstrap references for the written file without changing runtime precedence. ```dotenv LATEST_KEY_URI="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key" PINNED_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version-id" ``` -Bootstrap documents stay in memory by default. They load before explicit `env_files` or the default `~/.pyrit/.env` and `~/.pyrit/.env.local`, allowing intentional local overrides. +Bootstrap documents stay in memory by default. Use `.env.local` when an intentional local override is required. ### `env_akv_strict` -Controls validation only of the Key Vault bootstrap document and defaults to `true`. It does not change parsing of `.env`, `.env.local`, or explicit `env_files`. +Controls Key Vault bootstrap validation and Key Vault reference syntax in local files. It defaults to `true`. ```yaml env_akv_strict: false ``` -In strict mode, any malformed dotenv line or variable without an equals sign stops that bootstrap document before it mutates the environment. Empty assignments such as `OPTIONAL_VALUE=` remain valid and set the variable to an empty string. A referenced Key Vault secret whose value is an empty string is also valid. A missing value represented by `None` is treated as an error. With `env_akv_strict: false`, PyRIT emits a warning containing only malformed line numbers and valueless variable names, skips those entries, and loads the valid assignments. Secret values are never included in the warning. +In strict mode, malformed bootstrap dotenv lines, valueless bootstrap entries, and malformed Key Vault references stop initialization. Empty assignments such as `OPTIONAL_VALUE=` remain valid. With `env_akv_strict: false`, PyRIT warns and skips malformed bootstrap entries and malformed reference assignments without logging secret values. -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. +Non-strict mode does not suppress operational failures. Missing secrets, authentication, authorization, transport errors, and bootstrap documents with no valid assignments still stop initialization. Loading remains non-transactional, so earlier successful assignments remain. 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: +Defaults to `false`. Set it to `true` only while debugging to write a fully resolved bootstrap document to `~/.pyrit/.env`: ```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. +The written file contains only bootstrap assignments, comments, and fully resolved child-secret values. It excludes unrelated process and `.env.local` values, safely round-trips terminal secret text, and is always named `.env`, never `.env.new`. -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. +PyRIT refuses debug mode when `~/.pyrit/.env` already exists; rename or remove the existing file first. The file is created with owner-only permissions where supported and replaced atomically, but it contains plaintext secrets. Remove it when debugging is complete. `.env.local` still loads afterward and can override runtime values without changing the generated file. ### `silent` @@ -292,7 +285,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. Configured AKV bootstrap documents load in order, followed by selected environment files +1. Process values are retained, AKV or ordinary local sources fill gaps, and `.env.local` applies final overrides 2. Default values are reset 3. Memory database is configured (from `memory_db_type`) 4. Initializers are executed in listed order @@ -368,14 +361,15 @@ initializers: # initialization_scripts: # - /path/to/my_custom_initializer.py -# Recommended: ordered Azure Key Vault bootstrap environment documents +# Canonical: ordered Azure Key Vault bootstrap environment documents # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env # env_akv_strict: true -# env_akv_write_env: false # Opt in to writing ~/.pyrit/.env for inspection +# env_akv_write_env: false # Debug only: writes fully resolved plaintext secrets -# Recommended with Key Vault: disable local dotenv overrides -# env_files: [] +# Optional plaintext local patch or non-Azure workflow +# env_files: +# - /path/to/.env.local # Suppress initialization messages silent: false diff --git a/pyrit/setup/akv_initialization.py b/pyrit/setup/akv_initialization.py index 9941375ca2..3751d66205 100644 --- a/pyrit/setup/akv_initialization.py +++ b/pyrit/setup/akv_initialization.py @@ -11,13 +11,13 @@ import pathlib import tempfile import urllib.parse -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any import dotenv from dotenv.parser import parse_stream -from pyrit.common import path +from pyrit.common import path, print_deprecation_message from pyrit.exceptions import KeyVaultInitializationException if TYPE_CHECKING: @@ -30,6 +30,7 @@ _AKV_RETRY_TOTAL = 3 _AKV_RETRY_BACKOFF_FACTOR = 0.8 _AKV_ENV_FILE_NAME = ".env" +_LEGACY_ENV_REMOVED_IN = "1.3.0" def _load_environment_files( @@ -37,10 +38,13 @@ def _load_environment_files( *, silent: bool = False, include_default_base: bool = True, + assignment_fallbacks: dict[str, str | None] | None = None, ) -> bool: """ Load environment files in the order they are provided. - Later files override values from earlier files. + + Files fill values missing from the process environment. A file named + ``.env.local`` is the only local source that overrides existing values. Args: env_files: Optional sequence of environment file paths. If None, loads default @@ -49,6 +53,8 @@ def _load_environment_files( 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. + assignment_fallbacks: Optional output mapping from assignments that win + precedence to the value they replaced, if any. Returns: True if at least one environment file was loaded, otherwise False. @@ -62,7 +68,17 @@ def _load_environment_files( include_default_base=include_default_base, ) for env_file in selected_files: - dotenv.load_dotenv(dotenv_path=env_file, override=True, interpolate=True) + override = env_file.name == ".env.local" + if assignment_fallbacks is not None: + assignment_names = dotenv.dotenv_values(dotenv_path=env_file, interpolate=False) + for variable_name in assignment_names: + if override or variable_name not in os.environ: + assignment_fallbacks[variable_name] = os.environ.get(variable_name) + dotenv.load_dotenv( + dotenv_path=env_file, + override=override, + interpolate=True, + ) if not silent: _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) @@ -98,6 +114,7 @@ def _select_environment_files( local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" if include_default_base and base_file.exists(): + _warn_about_legacy_env(env_file=base_file, ignored_for_akv=False, silent=silent) default_files.append(base_file) if local_file.exists(): default_files.append(local_file) @@ -139,35 +156,26 @@ def _warn_about_akv_environment_files( *, 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] = [] + """Warn when an auto-discovered legacy environment file coexists with AKV.""" + if env_files is not None: + return + base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" 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") + _warn_about_legacy_env(env_file=base_file, ignored_for_akv=True, silent=silent) - 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 +def _warn_about_legacy_env(*, env_file: pathlib.Path, ignored_for_akv: bool, silent: bool) -> None: + """Emit the standard and visible warnings for auto-discovered legacy ``.env`` loading.""" + print_deprecation_message( + old_item=f"Auto-discovered {env_file}", + new_item="env_akv_ref or ~/.pyrit/.env.local", + removed_in=_LEGACY_ENV_REMOVED_IN, + ) + behavior = "will be ignored because env_akv_ref is configured" if ignored_for_akv else "will still be loaded" 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." + f"Auto-discovered {env_file} is deprecated and {behavior}. " + f"Support will be removed in {_LEGACY_ENV_REMOVED_IN}. Use env_akv_ref or ~/.pyrit/.env.local instead." ) if not silent: print(f"WARNING: {message}") @@ -261,6 +269,30 @@ def _create_akv_secret_client(*, vault_url: str, credential: Any) -> "SecretClie return SecretClient(vault_url=vault_url, credential=credential, retry_policy=retry_policy) +async def _fetch_akv_secret_value_async( + *, + client: Any, + secret_name: str, + secret_version: str | None, + variable_name: str, +) -> str: + """ + Fetch a referenced Key Vault secret value. + + Returns: + str: The secret value, including an empty string. + + Raises: + ValueError: If the referenced secret has no value. + """ + referenced_secret = await client.get_secret(secret_name, version=secret_version) + if referenced_secret.value is None: + raise ValueError( + f"AKV secret '{secret_name}' referenced by environment variable '{variable_name}' has no value." + ) + return referenced_secret.value + + def _key_vault_initialization_error(*, message: str, error: Exception) -> KeyVaultInitializationException: """ Create a contextual Key Vault exception without losing the original cause. @@ -327,6 +359,7 @@ async def _load_env_from_akv_async( secret_url: str, strict: bool = True, silent: bool = False, + resolve_references_for_output: bool = False, ) -> str: """ Load a bootstrap dotenv document and resolve its same-vault secret references. @@ -344,9 +377,13 @@ async def _load_env_from_akv_async( 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. + resolve_references_for_output (bool): If True, resolve child references even + when their runtime assignment loses to an existing process value, and + return a native dotenv document containing those resolved values. Returns: - str: The validated bootstrap dotenv document before child-secret resolution. + str: The validated bootstrap dotenv document, with child-secret references + replaced when ``resolve_references_for_output`` is True. Raises: ImportError: If ``azure-keyvault-secrets`` is not installed. @@ -370,33 +407,58 @@ 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}") + existing_environment_names = set(os.environ) loaded = dotenv.load_dotenv( stream=io.StringIO(validated_document), - override=True, + override=False, interpolate=True, ) if not loaded: return validated_document + resolved_reference_values: dict[str, str] = {} + skipped_reference_names: set[str] = set() for variable_name, value in parsed_environment.items(): if value is None: continue target = _parse_akv_reference(value) if target is None: continue + assignment_wins = variable_name not in existing_environment_names + if not assignment_wins and not resolve_references_for_output: + continue try: - referenced_name, referenced_version = _resolve_akv_secret_reference( + _, referenced_name, referenced_version = _parse_akv_reference_url( target=target, variable_name=variable_name, - vault_url=vault_url, + expected_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." + except ValueError as error: + if strict: + wrapped_error = _key_vault_initialization_error( + message=f"Invalid AKV reference for environment variable '{variable_name}'", + error=error, ) - os.environ[variable_name] = referenced_secret.value + raise wrapped_error from error + if assignment_wins: + os.environ.pop(variable_name, None) + skipped_reference_names.add(variable_name) + _warn_about_invalid_akv_reference( + variable_name=variable_name, + error=error, + silent=silent, + ) + continue + try: + resolved_value = await _fetch_akv_secret_value_async( + client=client, + secret_name=referenced_name, + secret_version=referenced_version, + variable_name=variable_name, + ) + resolved_reference_values[variable_name] = resolved_value + if assignment_wins: + os.environ[variable_name] = resolved_value except KeyVaultInitializationException: raise except Exception as error: @@ -405,6 +467,12 @@ async def _load_env_from_akv_async( error=error, ) raise wrapped_error from error + if resolve_references_for_output: + return _render_resolved_akv_document( + document=validated_document, + resolved_reference_values=resolved_reference_values, + skipped_reference_names=skipped_reference_names, + ) return validated_document except KeyVaultInitializationException: raise @@ -444,6 +512,12 @@ async def _load_environment_async( 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.") + env_file = path.CONFIGURATION_DIRECTORY_PATH / _AKV_ENV_FILE_NAME + if env_akv_write_env and (env_file.exists() or env_file.is_symlink()): + raise ValueError( + f"Cannot write the resolved Key Vault environment because {env_file} already exists; " + "rename or remove it before enabling env_akv_write_env." + ) await asyncio.to_thread( _warn_about_akv_environment_files, env_files=env_files, @@ -455,6 +529,7 @@ async def _load_environment_async( secret_url=secret_url, strict=env_akv_strict, silent=silent, + resolve_references_for_output=env_akv_write_env, ) for secret_url in env_akv_ref ] @@ -473,17 +548,24 @@ async def _load_environment_async( written_path = written_env_file.resolve() selected_env_files = [env_file for env_file in env_files if env_file.expanduser().resolve() != written_path] + assignment_fallbacks: dict[str, str | None] = {} 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), + include_default_base=not (env_akv_ref and env_files is None), + assignment_fallbacks=assignment_fallbacks, + ) + await _resolve_local_akv_references_async( + assignment_fallbacks=assignment_fallbacks, + strict=env_akv_strict, + silent=silent, ) def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Path: """ - Write fetched bootstrap documents without resolved child-secret values. + Write fetched bootstrap documents with resolved child-secret values. Returns: pathlib.Path: Path to the written dotenv file. @@ -496,7 +578,7 @@ def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Pa if env_file.is_symlink(): raise ValueError(f"Refusing to write the AKV environment through a symbolic link: {env_file}") - content = "\n".join(document.rstrip("\r\n") for document in documents) + "\n" + content = _merge_akv_documents_for_debug(documents=documents) file_descriptor: int | None = None temporary_file: pathlib.Path | None = None try: @@ -506,8 +588,9 @@ def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Pa dir=env_file.parent, ) temporary_file = pathlib.Path(temporary_name) - if hasattr(os, "fchmod"): - os.fchmod(file_descriptor, 0o600) + file_chmod = getattr(os, "fchmod", None) + if file_chmod is not None: + file_chmod(file_descriptor, 0o600) else: os.chmod(temporary_file, 0o600) stream = os.fdopen(file_descriptor, "w", encoding="utf-8", newline="") @@ -529,6 +612,90 @@ def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Pa return env_file +def _merge_akv_documents_for_debug(*, documents: Sequence[str]) -> str: + """ + Merge resolved bootstrap documents using runtime first-document precedence. + + Duplicate assignments within one document are retained because interpolation + depends on assignment order. Assignments established by an earlier document + are omitted from later documents. + + Returns: + str: A native dotenv document with equivalent bootstrap precedence. + """ + established_names: set[str] = set() + merged_bindings: list[str] = [] + for document in documents: + document_names: set[str] = set() + for binding in parse_stream(io.StringIO(document)): + if binding.key is None or binding.key not in established_names: + merged_bindings.append(binding.original.string) + if binding.key is not None: + document_names.add(binding.key) + established_names.update(document_names) + + return "".join(merged_bindings).rstrip("\r\n") + "\n" + + +def _render_resolved_akv_document( + *, + document: str, + resolved_reference_values: Mapping[str, str], + skipped_reference_names: set[str] | None = None, +) -> str: + """ + Replace resolved Key Vault reference assignments with native dotenv values. + + Returns: + str: Dotenv text that preserves non-reference bindings and comments. + """ + skipped_reference_names = skipped_reference_names or set() + rendered_bindings: list[str] = [] + for binding in parse_stream(io.StringIO(document)): + variable_name = binding.key + is_reference = binding.value is not None and _parse_akv_reference(binding.value) is not None + if variable_name in skipped_reference_names and is_reference: + continue + if variable_name is not None and variable_name in resolved_reference_values and is_reference: + original = binding.original.string + export_prefix = "export " if original.lstrip().startswith("export ") else "" + if original.endswith("\r\n"): + newline = "\r\n" + elif original.endswith("\n"): + newline = "\n" + else: + newline = "" + rendered_bindings.append( + f"{export_prefix}{variable_name}=" + f"{_serialize_terminal_dotenv_value(resolved_reference_values[variable_name])}{newline}" + ) + else: + rendered_bindings.append(binding.original.string) + return "".join(rendered_bindings) + + +def _serialize_terminal_dotenv_value(value: str) -> str: + """ + Quote a terminal secret value for a native python-dotenv round trip. + + The empty-name default expression produces a literal dollar sign during + interpolation, preventing terminal ``${NAME}`` text from being reinterpreted. + + Returns: + str: A single-quoted dotenv value. + """ + escaped_value = value.replace("'", "\\'").replace("${", "${:-$}{") + return f"'{escaped_value}'" + + +def _warn_about_invalid_akv_reference(*, variable_name: str, error: ValueError, silent: bool) -> None: + """Warn that a malformed Key Vault reference assignment is being skipped.""" + message = f"Invalid AKV reference for environment variable '{variable_name}' will be skipped: {error}" + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + + def _parse_akv_reference(value: str) -> str | None: """ Parse an exact whole-value Key Vault reference. @@ -548,25 +715,20 @@ def _validate_akv_secret_name(*, secret_name: str, variable_name: str) -> None: ) -def _resolve_akv_secret_reference( +def _parse_akv_reference_url( *, target: str, variable_name: str, - vault_url: str, -) -> tuple[str, str | None]: + expected_vault_url: str | None = None, +) -> tuple[str, 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. + Parse and optionally constrain a complete Key Vault secret reference. Returns: - tuple[str, str | None]: Secret name and optional version. + tuple[str, str, str | None]: Vault URL, secret name, and optional secret version. Raises: - ValueError: If the target is not a full URI, is invalid, or references another vault. + ValueError: If the reference is malformed or violates the expected vault constraint. """ if not target.casefold().startswith("https://"): raise ValueError( @@ -575,11 +737,115 @@ def _resolve_akv_secret_reference( ) referenced_vault_url, secret_name, secret_version = _parse_akv_secret_url(target) - if referenced_vault_url.rstrip("/").casefold() != vault_url.rstrip("/").casefold(): + if expected_vault_url and referenced_vault_url.rstrip("/").casefold() != expected_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}'." + f"Expected vault '{expected_vault_url}', got '{referenced_vault_url}'." ) _validate_akv_secret_name(secret_name=secret_name, variable_name=variable_name) + return referenced_vault_url, secret_name, secret_version + + +async def _resolve_local_akv_references_async( + *, + assignment_fallbacks: Mapping[str, str | None], + strict: bool, + silent: bool, +) -> None: + """ + Resolve complete Key Vault references from winning local assignments. + + Raises: + KeyVaultInitializationException: If strict validation or secret retrieval fails. + """ + parsed_references: list[tuple[str, str, str, str | None]] = [] + for variable_name, fallback_value in assignment_fallbacks.items(): + value = os.environ.get(variable_name) + if value is None: + continue + target = _parse_akv_reference(value) + if target is None: + continue + try: + vault_url, secret_name, secret_version = _parse_akv_reference_url( + target=target, + variable_name=variable_name, + ) + except ValueError as error: + if strict: + wrapped_error = _key_vault_initialization_error( + message=f"Invalid AKV reference for environment variable '{variable_name}'", + error=error, + ) + raise wrapped_error from error + if fallback_value is None: + os.environ.pop(variable_name, None) + else: + os.environ[variable_name] = fallback_value + _warn_about_invalid_akv_reference( + variable_name=variable_name, + error=error, + silent=silent, + ) + continue + parsed_references.append((variable_name, vault_url, secret_name, secret_version)) + + if not parsed_references: + return + + from azure.identity.aio import DefaultAzureCredential + + async with DefaultAzureCredential() as credential: + async with contextlib.AsyncExitStack() as client_stack: + clients: dict[str, Any] = {} + for variable_name, vault_url, secret_name, secret_version in parsed_references: + try: + client = clients.get(vault_url) + if client is None: + client = await client_stack.enter_async_context( + _create_akv_secret_client(vault_url=vault_url, credential=credential) + ) + clients[vault_url] = client + os.environ[variable_name] = await _fetch_akv_secret_value_async( + client=client, + secret_name=secret_name, + secret_version=secret_version, + variable_name=variable_name, + ) + 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 + + +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. + """ + _, secret_name, secret_version = _parse_akv_reference_url( + target=target, + variable_name=variable_name, + expected_vault_url=vault_url, + ) return secret_name, secret_version diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index ef118ebb0f..6caeafb431 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -95,12 +95,13 @@ class ConfigurationLoader(YamlLoadable): initialization_scripts: List of paths to custom initialization scripts. 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". + None means auto-discover legacy ``.env`` and supported ``.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. - env_akv_write_env: Whether to save fetched bootstrap documents to - ``~/.pyrit/.env`` for local inspection. + env_akv_write_env: Whether to save fully resolved bootstrap documents with + plaintext child-secret values to ``~/.pyrit/.env`` for debugging. silent: Whether to suppress initialization messages. operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 78a4e4b804..91ce1e6112 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -95,15 +95,15 @@ async def initialize_pyrit_async( ``core`` techniques and ``default`` targets are loaded — ``extra`` / per-source technique groups and ``scorer`` target variants remain opt-in. 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. + in order. Ordinary files fill missing process values; files named ``.env.local`` override. + If omitted, PyRIT auto-discovers legacy ``.env`` and supported ``.env.local`` files. 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. - env_akv_write_env (bool): If True, save fetched bootstrap documents with unresolved - child references to ``~/.pyrit/.env``. Defaults to False. + contain bootstrap dotenv documents. Documents fill missing process values and support + complete-value references to scalar secrets. Requires ``azure-keyvault-secrets``. + env_akv_strict (bool): If True, reject malformed bootstrap entries and Key Vault reference + syntax. If False, warn and skip those entries. Operational Key Vault failures always raise. + env_akv_write_env (bool): If True, write fully resolved bootstrap documents with plaintext + child-secret values to ``~/.pyrit/.env`` for debugging. 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. diff --git a/tests/integration/setup/test_akv_initialization_integration.py b/tests/integration/setup/test_akv_initialization_integration.py deleted file mode 100644 index 0bab039e6b..0000000000 --- a/tests/integration/setup/test_akv_initialization_integration.py +++ /dev/null @@ -1,61 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import io -import os -import pathlib - -import pytest -from dotenv import dotenv_values - -_AKV_ENVIRONMENT_ENV = "PYRIT_AKV_INTEGRATION_TEST_ENV" -_AKV_ENVIRONMENT_REQUIRED_ENV = "PYRIT_AKV_INTEGRATION_TEST_REQUIRED" -_ENV_EXAMPLE_PATH_ENV = "PYRIT_ENV_EXAMPLE_PATH" - - -def _get_env_example_path() -> pathlib.Path: - configured_path = os.getenv(_ENV_EXAMPLE_PATH_ENV) - if configured_path: - path = pathlib.Path(configured_path) - if path.is_file(): - return path - raise AssertionError(f"{_ENV_EXAMPLE_PATH_ENV} does not identify a file: {path}") - - candidates = [pathlib.Path.cwd() / ".env_example"] - candidates.extend(parent / ".env_example" for parent in pathlib.Path(__file__).resolve().parents) - for path in candidates: - if path.is_file(): - return path - - raise AssertionError("Could not locate .env_example.") - - -def _get_akv_environment_keys() -> set[str]: - document = os.getenv(_AKV_ENVIRONMENT_ENV) - if not document: - if os.getenv(_AKV_ENVIRONMENT_REQUIRED_ENV, "").casefold() == "true": - raise AssertionError(f"{_AKV_ENVIRONMENT_ENV} is required but was not populated.") - pytest.skip(f"Set {_AKV_ENVIRONMENT_ENV} to run the AKV schema integration test.") - - keys = set(dotenv_values(stream=io.StringIO(document), interpolate=False)) - if not keys: - raise AssertionError("The env-new Key Vault secret contains no dotenv assignments.") - return keys - - -def _get_env_example_keys() -> set[str]: - keys = set(dotenv_values(dotenv_path=_get_env_example_path(), interpolate=False)) - if not keys: - raise AssertionError(".env_example contains no dotenv assignments.") - return keys - - -def test_env_example_keys_are_represented_in_akv_environment() -> None: - """Ensure the new AKV bootstrap document covers every environment name in the public example.""" - env_example_keys = _get_env_example_keys() - akv_environment_keys = _get_akv_environment_keys() - - missing_from_akv = env_example_keys - akv_environment_keys - assert not missing_from_akv, "The env-new Key Vault secret is missing names defined in .env_example: " + ", ".join( - sorted(missing_from_akv) - ) diff --git a/tests/integration/setup/test_env_example_drift.py b/tests/integration/setup/test_env_example_drift.py index 81913948db..9daa065d23 100644 --- a/tests/integration/setup/test_env_example_drift.py +++ b/tests/integration/setup/test_env_example_drift.py @@ -14,6 +14,7 @@ _ENVIRONMENT_NAME_PATTERN = re.compile(r"(? pathlib.Path: @@ -105,6 +106,17 @@ def test_env_example_url_values_are_not_wrapped_in_angle_brackets() -> None: ) +def test_env_example_comment_blocks_do_not_contain_blank_lines() -> None: + """Keep consecutive comment lines together so the example remains compact.""" + repository_root = _get_repository_root() + env_example_path = _get_env_example_path(repository_root=repository_root) + contents = env_example_path.read_text(encoding="utf-8") + + assert not _BLANK_LINE_BETWEEN_COMMENTS_PATTERN.search(contents), ( + ".env_example contains a blank line between consecutive comment lines." + ) + + def test_env_example_aliases_resolve_in_assignment_order() -> None: """Ensure complete-value aliases resolve to their sources without ambient environment values.""" repository_root = _get_repository_root() diff --git a/tests/unit/setup/test_akv_initialization.py b/tests/unit/setup/test_akv_initialization.py index 0cd51db544..c0f3d465e4 100644 --- a/tests/unit/setup/test_akv_initialization.py +++ b/tests/unit/setup/test_akv_initialization.py @@ -5,10 +5,12 @@ import pathlib import tempfile import types +import warnings from unittest import mock import pytest from azure.core.exceptions import ResourceNotFoundError +from dotenv import dotenv_values from pyrit.exceptions import KeyVaultInitializationException from pyrit.setup import IN_MEMORY, initialize_pyrit_async @@ -37,7 +39,10 @@ async def test_loads_default_env_files_when_none_provided(self, mock_config_path 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): + with ( + mock.patch.dict(os.environ, {}, clear=True), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): loaded = _load_environment_files(env_files=None) assert loaded is True @@ -53,12 +58,49 @@ async def test_only_loads_existing_default_files(self, mock_config_path): env_file.write_text("VAR1=value1") mock_config_path.__truediv__ = lambda self, other: temp_path / other - with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch.dict(os.environ, {}, clear=True), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): 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_default_env_preserves_process_environment(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VAR=legacy\nLEGACY_ONLY=legacy") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch.dict(os.environ, {"VAR": "process"}, clear=True), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): + loaded = _load_environment_files(env_files=None, silent=True) + + assert loaded is True + assert os.environ["VAR"] == "process" + assert os.environ["LEGACY_ONLY"] == "legacy" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_default_env_local_overrides_process_environment_and_env(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VAR=legacy") + (temp_path / ".env.local").write_text("VAR=local") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch.dict(os.environ, {"VAR": "process"}, clear=True), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): + loaded = _load_environment_files(env_files=None, silent=True) + + assert loaded is True + assert os.environ["VAR"] == "local" + @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: @@ -89,59 +131,105 @@ async def test_returns_false_when_no_default_files_exist(self, mock_config_path) 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): + def test_auto_discovered_env_warns_with_removal_version(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) + with ( + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + pytest.warns(DeprecationWarning, match=r"\.env.*removed in 1\.3\.0.*\.env\.local"), + ): + _load_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" + assert f"WARNING: Auto-discovered {env_file} is deprecated" in output + assert "Use env_akv_ref or ~/.pyrit/.env.local instead" in output + assert f"Auto-discovered {env_file} is deprecated" in caplog.text @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): + def test_explicit_env_file_does_not_emit_legacy_deprecation(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" - 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") + explicit_env = temp_path / ".env" + explicit_env.write_text("VAR=explicit") mock_config_path.__truediv__ = lambda self, other: temp_path / other - _warn_about_akv_environment_files(env_files=[custom_file]) + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + loaded = _load_environment_files(env_files=[explicit_env], silent=True) - 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 + assert loaded is True @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") - def test_akv_environment_file_warning_respects_silent(self, mock_config_path, caplog, capsys): + def test_akv_legacy_env_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"): + with ( + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): _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 + assert "will be ignored because env_akv_ref is configured" in caplog.text + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_akv_ignores_auto_discovered_env_and_loads_env_local(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VALUE=legacy") + (temp_path / ".env.local").write_text("VALUE=local") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch( + "pyrit.setup.akv_initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + return_value="VALUE=akv\n", + ), + mock.patch("pyrit.setup.akv_initialization._load_environment_files") as mock_load_files, + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): + await _load_environment_async( + env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=None, + env_akv_strict=True, + env_akv_write_env=False, + silent=True, + ) + + assert mock_load_files.call_args.kwargs["env_files"] is None + assert mock_load_files.call_args.kwargs["include_default_base"] is False + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_akv_debug_mode_rejects_existing_env_before_fetch(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_file.write_text("VALUE=legacy") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch( + "pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock + ) as mock_load_akv, + pytest.raises(ValueError, match=r"already exists.*rename or remove"), + ): + await _load_environment_async( + env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=None, + env_akv_strict=True, + env_akv_write_env=True, + silent=True, + ) + + mock_load_akv.assert_not_awaited() async def test_loads_custom_env_files_in_order(self): """Test that custom env_files are loaded in the order provided.""" @@ -162,6 +250,31 @@ async def test_loads_custom_env_files_in_order(self): assert loaded is True assert os.environ["VAR"] == "local" + async def test_explicit_files_only_override_when_named_env_local(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + first_file = temp_path / "first.env" + second_file = temp_path / "second.env" + local_file = temp_path / "nested" / ".env.local" + local_file.parent.mkdir() + first_file.write_text("PROCESS_VALUE=first\nFILE_VALUE=first") + second_file.write_text("PROCESS_VALUE=second\nFILE_VALUE=second\nSECOND_ONLY=second") + local_file.write_text("PROCESS_VALUE=local\nFILE_VALUE=local") + + with mock.patch.dict(os.environ, {"PROCESS_VALUE": "process"}, clear=True): + loaded = _load_environment_files(env_files=[first_file, second_file, local_file], silent=True) + + assert loaded is True + assert os.environ["PROCESS_VALUE"] == "local" + assert os.environ["FILE_VALUE"] == "local" + assert os.environ["SECOND_ONLY"] == "second" + + with mock.patch.dict(os.environ, {"PROCESS_VALUE": "process"}, clear=True): + _load_environment_files(env_files=[first_file, second_file], silent=True) + + assert os.environ["PROCESS_VALUE"] == "process" + assert os.environ["FILE_VALUE"] == "first" + async def test_load_environment_files_interpolates_in_assignment_order(self): with tempfile.TemporaryDirectory() as temp_dir: env_file = pathlib.Path(temp_dir) / ".env" @@ -186,41 +299,73 @@ async def test_load_environment_files_honors_python_dotenv_disabled(self): 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', - ] + async def test_load_environment_async_write_env_writes_resolved_native_bootstrap(self): + credential, client = _create_mock_akv_clients() + document = ( + "# Bootstrap values\n" + "BASE=bootstrap\n" + "DERIVED=${BASE}\n" + "API_KEY=kv:https://vault.vault.azure.net/secrets/api-key\n" + ) + resolved_api_key = "line one\nquote' and literal ${UNRELATED}" + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value=document), + types.SimpleNamespace(value=resolved_api_key), + ] + ) with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) + (temp_path / ".env.local").write_text("API_KEY=local-key\nLOCAL_ONLY=local", 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, - side_effect=documents, + mock.patch.dict( + os.environ, + { + "BASE": "process", + "API_KEY": "process-key", + "PROCESS_ONLY": "not-written", + "UNRELATED": "changed", + }, + clear=True, ), - mock.patch( - "pyrit.setup.akv_initialization._load_environment_files", return_value=False - ) as mock_load_environment_files, + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), ): await _load_environment_async( - env_akv_ref=references, + env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], 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 + assert os.environ["BASE"] == "process" + assert os.environ["API_KEY"] == "local-key" + assert os.environ["LOCAL_ONLY"] == "local" + + written_env = temp_path / ".env" + assert written_env.is_file() + assert not (temp_path / ".env.new").exists() + content = written_env.read_text(encoding="utf-8") + assert "# Bootstrap values" in content + assert "kv:" not in content + assert "PROCESS_ONLY" not in content + assert "LOCAL_ONLY" not in content + + with mock.patch.dict(os.environ, {}, clear=True): + written_values = dotenv_values(dotenv_path=written_env, interpolate=True) + + assert written_values == { + "BASE": "bootstrap", + "DERIVED": "bootstrap", + "API_KEY": resolved_api_key, + } + assert client.get_secret.await_args_list == [ + mock.call("bootstrap", version=None), + mock.call("api-key", version=None), + ] async def test_load_environment_async_write_env_filters_generated_explicit_file(self): with tempfile.TemporaryDirectory() as temp_dir: @@ -251,6 +396,42 @@ async def test_load_environment_async_write_env_filters_generated_explicit_file( 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_load_environment_async_write_env_preserves_first_bootstrap_value(self): + documents = [ + "SHARED=first\nFIRST_ONLY=first\n", + "SHARED=second\nSECOND_ONLY=second\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=lambda **kwargs: documents.pop(0), + ), + ): + await _load_environment_async( + env_akv_ref=[ + "https://vault.vault.azure.net/secrets/first", + "https://vault.vault.azure.net/secrets/second", + ], + env_files=[], + env_akv_strict=True, + env_akv_write_env=True, + silent=True, + ) + + with mock.patch.dict(os.environ, {}, clear=True): + written_values = dotenv_values(dotenv_path=temp_path / ".env", interpolate=True) + + assert written_values == { + "SHARED": "first", + "FIRST_ONLY": "first", + "SECOND_ONLY": "second", + } + def test_write_akv_env_file_secures_descriptor_before_writing(self): events: list[str] = [] @@ -357,7 +538,10 @@ 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 - with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch.dict(os.environ, {}, clear=True), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): loaded = _load_environment_files(env_files=None, silent=True) assert loaded is True @@ -382,23 +566,137 @@ 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_keeps_local_akv_reference_literal_without_bootstrap(self, mock_set_memory): + @pytest.mark.parametrize( + ("file_name", "initial_environment"), + [ + ("custom.env", {}), + (".env.local", {"API_KEY": "process-key"}), + ], + ) + async def test_load_environment_async_resolves_local_akv_reference(self, file_name, initial_environment): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="resolved-key")) + 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") + env_file = pathlib.Path(temp_dir) / file_name + env_file.write_text("API_KEY=kv:https://local-vault.vault.azure.net/secrets/api-key/version-1") - with mock.patch.dict(os.environ, {}, clear=True): - await initialize_pyrit_async( - memory_db_type=IN_MEMORY, + with ( + mock.patch.dict(os.environ, initial_environment, 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 _load_environment_async( + env_akv_ref=None, env_files=[env_file], - load_defaults=False, + env_akv_strict=True, + env_akv_write_env=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "resolved-key" + + _assert_mock_akv_client_created( + mock_client_cls, + vault_url="https://local-vault.vault.azure.net", + credential=credential, + ) + client.get_secret.assert_awaited_once_with("api-key", version="version-1") + + async def test_load_environment_async_does_not_fetch_local_reference_that_loses_to_process_value(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / "custom.env" + env_file.write_text("API_KEY=kv:https://local-vault.vault.azure.net/secrets/api-key") + + with ( + mock.patch.dict(os.environ, {"API_KEY": "process-key"}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + ): + await _load_environment_async( + env_akv_ref=None, + env_files=[env_file], + env_akv_strict=True, + env_akv_write_env=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "process-key" + + mock_credential_cls.assert_not_called() + + async def test_load_environment_async_strict_rejects_malformed_local_akv_reference_before_authentication(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / "custom.env" + env_file.write_text("API_KEY=kv:api-key") + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + pytest.raises(KeyVaultInitializationException, match="must use a full secret URL"), + ): + await _load_environment_async( + env_akv_ref=None, + env_files=[env_file], + env_akv_strict=True, + env_akv_write_env=False, silent=True, ) - assert os.environ["API_KEY"] == "kv:https://myvault.vault.azure.net/secrets/api-key" + mock_credential_cls.assert_not_called() + + async def test_load_environment_async_non_strict_skips_malformed_local_akv_reference(self, caplog, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + env_local_file = pathlib.Path(temp_dir) / ".env.local" + env_local_file.write_text("API_KEY=kv:api-key") + + with ( + mock.patch.dict(os.environ, {"API_KEY": "process-key"}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + ): + await _load_environment_async( + env_akv_ref=None, + env_files=[env_local_file], + env_akv_strict=False, + env_akv_write_env=False, + silent=False, + ) - mock_set_memory.assert_called_once() + assert os.environ["API_KEY"] == "process-key" + + mock_credential_cls.assert_not_called() + assert ( + "WARNING: Invalid AKV reference for environment variable 'API_KEY' will be skipped" + in capsys.readouterr().out + ) + assert "API_KEY" in caplog.text + + async def test_load_environment_async_non_strict_still_raises_for_missing_local_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 tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / "custom.env" + env_file.write_text("API_KEY=kv:https://local-vault.vault.azure.net/secrets/missing") + + 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_environment_async( + env_akv_ref=None, + env_files=[env_file], + env_akv_strict=False, + env_akv_write_env=False, + silent=True, + ) + + assert exc_info.value.__cause__ is missing_error async def test_raises_error_for_nonexistent_env_file(self): """Test that ValueError is raised for non-existent env file.""" @@ -619,6 +917,31 @@ async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_secret client.__aexit__.assert_awaited_once() mock_print_msg.assert_called_once() + async def test_load_env_from_akv_async_preserves_process_values_without_fetching_overridden_child(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + return_value=types.SimpleNamespace( + value=("DIRECT=from-bootstrap\nFROM_KV=kv:https://myvault.vault.azure.net/secrets/api-key") + ) + ) + secret_url = "https://myvault.vault.azure.net/secrets/bootstrap" + + with ( + mock.patch.dict( + os.environ, + {"DIRECT": "from-process", "FROM_KV": "process-key"}, + 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=secret_url, silent=True) + + assert os.environ["DIRECT"] == "from-process" + assert os.environ["FROM_KV"] == "process-key" + + client.get_secret.assert_awaited_once_with("bootstrap", version=None) + 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")) @@ -819,6 +1142,35 @@ async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entrie assert "GOOD" not in caplog.text assert "resolved" not in caplog.text + async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_reference(self, caplog, capsys): + credential, client = _create_mock_akv_clients() + document = "GOOD=resolved\nBAD=kv:short-name\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"), + ): + resolved_document = await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=False, + resolve_references_for_output=True, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" + assert "BAD" not in os.environ + + assert "BAD=" not in resolved_document + assert ( + "WARNING: Invalid AKV reference for environment variable 'BAD' will be skipped" in capsys.readouterr().out + ) + assert "BAD" in caplog.text + client.get_secret.assert_awaited_once_with("bootstrap", version=None) + 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")) diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 0abea4ab90..8e445c9d9c 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -176,12 +176,14 @@ async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, m mock_load_akv.return_value = None - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) + with mock.patch("pyrit.setup.akv_initialization._warn_about_akv_environment_files") as mock_warn: + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) 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.call(secret_url=refs[0], strict=True, silent=False, resolve_references_for_output=False), + mock.call(secret_url=refs[1], strict=True, silent=False, resolve_references_for_output=False), ] + mock_warn.assert_called_once() mock_load_env.assert_called_once() mock_set_memory.assert_called_once() @@ -294,6 +296,13 @@ async def test_initialize_default_files_override_akv_in_order(self, mock_set_mem @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") async def test_initialize_resolves_bootstrap_references_before_local_overrides(self, mock_set_memory): refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + 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=mock.MagicMock(value="local-secret-value")) with tempfile.TemporaryDirectory() as temp_dir: local_file = pathlib.Path(temp_dir) / ".env.local" local_file.write_text( @@ -315,6 +324,8 @@ async def test_initialize_resolves_bootstrap_references_before_local_overrides(s new_callable=mock.AsyncMock, side_effect=lambda **_: os.environ.update(bootstrap_environment), ), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("pyrit.setup.akv_initialization._create_akv_secret_client", return_value=client), ): await initialize_pyrit_async( memory_db_type=IN_MEMORY, @@ -325,9 +336,11 @@ async def test_initialize_resolves_bootstrap_references_before_local_overrides(s assert os.environ["OVERRIDDEN"] == "local" assert os.environ["BOOTSTRAP_SECRET"] == "bootstrap-secret-value" - assert os.environ["LOCAL_SECRET"] == "kv:https://vault.vault.azure.net/secrets/local-secret" + assert os.environ["LOCAL_SECRET"] == "local-secret-value" assert os.environ["LOCAL_ENV"] == "env:BOOTSTRAP_SOURCE" + client.get_secret.assert_awaited_once_with("local-secret", version=None) + mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") From 933f3bc2201162163ffade0961d5367749c09ebf Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 20 Aug 2026 11:43:19 -0400 Subject: [PATCH 21/28] FIX: Strict boolean checks for env_akv_strict and env_akv_write_env --- pyrit/setup/akv_initialization.py | 15 +++++++ pyrit/setup/configuration_loader.py | 5 +++ pyrit/setup/initialization.py | 7 ++- tests/unit/setup/test_configuration_loader.py | 26 +++++++++++ tests/unit/setup/test_initialization.py | 43 +++++++++++++++++++ 5 files changed, 95 insertions(+), 1 deletion(-) diff --git a/pyrit/setup/akv_initialization.py b/pyrit/setup/akv_initialization.py index 3751d66205..6f005a5d51 100644 --- a/pyrit/setup/akv_initialization.py +++ b/pyrit/setup/akv_initialization.py @@ -33,6 +33,21 @@ _LEGACY_ENV_REMOVED_IN = "1.3.0" +def _validate_akv_boolean_options(*, env_akv_strict: object, env_akv_write_env: object) -> None: + """ + Require real booleans for Key Vault behavior flags. + + Raises: + TypeError: If either option is not a bool. + """ + for option_name, option_value in ( + ("env_akv_strict", env_akv_strict), + ("env_akv_write_env", env_akv_write_env), + ): + if not isinstance(option_value, bool): + raise TypeError(f"{option_name} must be a bool, got {type(option_value).__name__}.") + + def _load_environment_files( env_files: Sequence[pathlib.Path] | None, *, diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 6caeafb431..31e4f3e4cf 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -21,6 +21,7 @@ from pyrit.common.utils import verify_and_resolve_path from pyrit.common.yaml_loadable import YamlLoadable from pyrit.models import class_name_to_snake_case +from pyrit.setup.akv_initialization import _validate_akv_boolean_options from pyrit.setup.initialization import ( AZURE_SQL, IN_MEMORY, @@ -153,6 +154,10 @@ class ConfigurationLoader(YamlLoadable): def __post_init__(self) -> None: """Validate and normalize the configuration after loading.""" + _validate_akv_boolean_options( + env_akv_strict=self.env_akv_strict, + env_akv_write_env=self.env_akv_write_env, + ) self._normalize_memory_db_type() self._normalize_initializers() self._validate_env_akv_ref() diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 91ce1e6112..fba743e2fc 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -7,7 +7,7 @@ from pyrit.common.apply_defaults import reset_default_values from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory -from pyrit.setup.akv_initialization import _load_environment_async +from pyrit.setup.akv_initialization import _load_environment_async, _validate_akv_boolean_options if TYPE_CHECKING: from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -109,8 +109,13 @@ async def initialize_pyrit_async( **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. Raises: + TypeError: If ``env_akv_strict`` or ``env_akv_write_env`` is not a bool. ValueError: If an unsupported memory_db_type is provided or env_files contains non-existent files. """ + _validate_akv_boolean_options( + env_akv_strict=env_akv_strict, + env_akv_write_env=env_akv_write_env, + ) await _load_environment_async( env_akv_ref=env_akv_ref, env_files=env_files, diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 93c812943a..65de4b684e 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -46,6 +46,15 @@ def test_default_values(self): assert config.env_akv_write_env is False assert config.silent is False + @pytest.mark.parametrize("option_name", ["env_akv_strict", "env_akv_write_env"]) + @pytest.mark.parametrize("invalid_value", ["false", "true", 0, 1, None, [], {}]) + def test_rejects_non_boolean_akv_options(self, option_name, invalid_value): + with pytest.raises(TypeError, match=rf"{option_name} must be a bool"): + if option_name == "env_akv_strict": + ConfigurationLoader(env_akv_strict=invalid_value) # type: ignore[arg-type] + else: + ConfigurationLoader(env_akv_write_env=invalid_value) # type: ignore[arg-type] + def test_valid_memory_db_types_snake_case(self): """Test all valid memory database types in snake_case.""" for db_type in ["in_memory", "sqlite", "azure_sql"]: @@ -237,6 +246,23 @@ def test_from_yaml_file(self): finally: pathlib.Path(yaml_path).unlink() + @pytest.mark.parametrize("option_name", ["env_akv_strict", "env_akv_write_env"]) + def test_from_yaml_rejects_quoted_boolean_akv_options(self, tmp_path, option_name): + yaml_path = tmp_path / "quoted-boolean.yaml" + yaml_path.write_text(f'{option_name}: "false"\n', encoding="utf-8") + + with pytest.raises(TypeError, match=rf"{option_name} must be a bool"): + ConfigurationLoader.from_yaml_file(yaml_path) + + def test_from_yaml_accepts_native_boolean_akv_options(self, tmp_path): + yaml_path = tmp_path / "native-booleans.yaml" + yaml_path.write_text("env_akv_strict: false\nenv_akv_write_env: true\n", encoding="utf-8") + + config = ConfigurationLoader.from_yaml_file(yaml_path) + + assert config.env_akv_strict is False + assert config.env_akv_write_env is True + def test_from_empty_yaml_file_raises_value_error(self, tmp_path): """Test that an empty YAML file raises a clear ValueError.""" yaml_path = tmp_path / "empty.yaml" diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 8e445c9d9c..56b2fa479d 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -209,6 +209,49 @@ async def test_initialize_rejects_invalid_env_akv_ref(self, env_akv_ref): load_defaults=False, ) + @pytest.mark.parametrize("option_name", ["env_akv_strict", "env_akv_write_env"]) + @pytest.mark.parametrize("invalid_value", ["false", "true", 0, 1, None, [], {}]) + async def test_initialize_rejects_non_boolean_akv_options_before_loading(self, option_name, invalid_value): + with mock.patch( + "pyrit.setup.initialization._load_environment_async", new_callable=mock.AsyncMock + ) as mock_load_environment: + with pytest.raises(TypeError, match=rf"{option_name} must be a bool"): + if option_name == "env_akv_strict": + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_strict=invalid_value, # type: ignore[arg-type] + load_defaults=False, + ) + else: + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_write_env=invalid_value, # type: ignore[arg-type] + load_defaults=False, + ) + + mock_load_environment.assert_not_awaited() + + @pytest.mark.parametrize( + ("env_akv_strict", "env_akv_write_env"), + [(True, False), (False, True)], + ) + async def test_initialize_forwards_boolean_akv_options(self, env_akv_strict, env_akv_write_env): + with ( + mock.patch( + "pyrit.setup.initialization._load_environment_async", new_callable=mock.AsyncMock + ) as mock_load_environment, + mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance"), + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_strict=env_akv_strict, + env_akv_write_env=env_akv_write_env, + load_defaults=False, + ) + + assert mock_load_environment.await_args.kwargs["env_akv_strict"] is env_akv_strict + assert mock_load_environment.await_args.kwargs["env_akv_write_env"] is env_akv_write_env + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") async def test_initialize_keeps_akv_values_when_local_file_loading_fails(self, mock_set_memory): refs = ["https://vault.vault.azure.net/secrets/bootstrap"] From 9be37dc6163053016e55b97436e446a7cebe8de7 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 20 Aug 2026 12:30:05 -0400 Subject: [PATCH 22/28] FIX: Backslash parsing, local interpolation fix, no-clobber fix --- pyrit/setup/akv_initialization.py | 73 +++++++++---- tests/unit/setup/test_akv_initialization.py | 114 ++++++++++++++++++-- tests/unit/setup/test_initialization.py | 6 +- 3 files changed, 162 insertions(+), 31 deletions(-) diff --git a/pyrit/setup/akv_initialization.py b/pyrit/setup/akv_initialization.py index 6f005a5d51..0d3c564ab0 100644 --- a/pyrit/setup/akv_initialization.py +++ b/pyrit/setup/akv_initialization.py @@ -431,8 +431,9 @@ async def _load_env_from_akv_async( if not loaded: return validated_document - resolved_reference_values: dict[str, str] = {} - skipped_reference_names: set[str] = set() + final_assignment_indexes = _get_final_assignment_indexes(document=validated_document) + resolved_reference_values: dict[int, str] = {} + skipped_reference_indexes: set[int] = set() for variable_name, value in parsed_environment.items(): if value is None: continue @@ -457,12 +458,12 @@ async def _load_env_from_akv_async( raise wrapped_error from error if assignment_wins: os.environ.pop(variable_name, None) - skipped_reference_names.add(variable_name) _warn_about_invalid_akv_reference( variable_name=variable_name, error=error, silent=silent, ) + skipped_reference_indexes.add(final_assignment_indexes[variable_name]) continue try: resolved_value = await _fetch_akv_secret_value_async( @@ -471,7 +472,7 @@ async def _load_env_from_akv_async( secret_version=referenced_version, variable_name=variable_name, ) - resolved_reference_values[variable_name] = resolved_value + resolved_reference_values[final_assignment_indexes[variable_name]] = resolved_value if assignment_wins: os.environ[variable_name] = resolved_value except KeyVaultInitializationException: @@ -486,7 +487,7 @@ async def _load_env_from_akv_async( return _render_resolved_akv_document( document=validated_document, resolved_reference_values=resolved_reference_values, - skipped_reference_names=skipped_reference_names, + skipped_reference_indexes=skipped_reference_indexes, ) return validated_document except KeyVaultInitializationException: @@ -529,10 +530,7 @@ async def _load_environment_async( raise ValueError("env_akv_ref must contain only non-empty Azure Key Vault secret URLs.") env_file = path.CONFIGURATION_DIRECTORY_PATH / _AKV_ENV_FILE_NAME if env_akv_write_env and (env_file.exists() or env_file.is_symlink()): - raise ValueError( - f"Cannot write the resolved Key Vault environment because {env_file} already exists; " - "rename or remove it before enabling env_akv_write_env." - ) + raise ValueError(_get_akv_env_file_exists_message(env_file=env_file)) await asyncio.to_thread( _warn_about_akv_environment_files, env_files=env_files, @@ -586,12 +584,15 @@ def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Pa pathlib.Path: Path to the written dotenv file. Raises: - ValueError: If the destination is a symbolic link. + ValueError: If the destination already exists or is a symbolic link. + OSError: If the filesystem cannot atomically publish the completed file. """ env_file = path.CONFIGURATION_DIRECTORY_PATH / _AKV_ENV_FILE_NAME env_file.parent.mkdir(mode=0o700, parents=True, exist_ok=True) if env_file.is_symlink(): raise ValueError(f"Refusing to write the AKV environment through a symbolic link: {env_file}") + if env_file.exists(): + raise ValueError(_get_akv_env_file_exists_message(env_file=env_file)) content = _merge_akv_documents_for_debug(documents=documents) file_descriptor: int | None = None @@ -612,10 +613,10 @@ def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Pa file_descriptor = None with stream: stream.write(content) - if env_file.is_symlink(): - raise ValueError(f"Refusing to replace a symbolic link with the AKV environment: {env_file}") - os.replace(temporary_file, env_file) - temporary_file = None + try: + os.link(temporary_file, env_file) + except FileExistsError as error: + raise ValueError(_get_akv_env_file_exists_message(env_file=env_file)) from error finally: if file_descriptor is not None: os.close(file_descriptor) @@ -627,6 +628,19 @@ def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Pa return env_file +def _get_akv_env_file_exists_message(*, env_file: pathlib.Path) -> str: + """ + Create the message used when debug output would clobber an existing path. + + Returns: + str: Error message containing recovery guidance for the user. + """ + return ( + f"Cannot write the resolved Key Vault environment because {env_file} already exists; " + "rename or remove it before enabling env_akv_write_env." + ) + + def _merge_akv_documents_for_debug(*, documents: Sequence[str]) -> str: """ Merge resolved bootstrap documents using runtime first-document precedence. @@ -655,8 +669,8 @@ def _merge_akv_documents_for_debug(*, documents: Sequence[str]) -> str: def _render_resolved_akv_document( *, document: str, - resolved_reference_values: Mapping[str, str], - skipped_reference_names: set[str] | None = None, + resolved_reference_values: Mapping[int, str], + skipped_reference_indexes: set[int] | None = None, ) -> str: """ Replace resolved Key Vault reference assignments with native dotenv values. @@ -664,14 +678,13 @@ def _render_resolved_akv_document( Returns: str: Dotenv text that preserves non-reference bindings and comments. """ - skipped_reference_names = skipped_reference_names or set() + skipped_reference_indexes = skipped_reference_indexes or set() rendered_bindings: list[str] = [] - for binding in parse_stream(io.StringIO(document)): + for binding_index, binding in enumerate(parse_stream(io.StringIO(document))): variable_name = binding.key - is_reference = binding.value is not None and _parse_akv_reference(binding.value) is not None - if variable_name in skipped_reference_names and is_reference: + if binding_index in skipped_reference_indexes: continue - if variable_name is not None and variable_name in resolved_reference_values and is_reference: + if variable_name is not None and binding_index in resolved_reference_values: original = binding.original.string export_prefix = "export " if original.lstrip().startswith("export ") else "" if original.endswith("\r\n"): @@ -682,13 +695,27 @@ def _render_resolved_akv_document( newline = "" rendered_bindings.append( f"{export_prefix}{variable_name}=" - f"{_serialize_terminal_dotenv_value(resolved_reference_values[variable_name])}{newline}" + f"{_serialize_terminal_dotenv_value(resolved_reference_values[binding_index])}{newline}" ) else: rendered_bindings.append(binding.original.string) return "".join(rendered_bindings) +def _get_final_assignment_indexes(*, document: str) -> dict[str, int]: + """ + Map each variable name to its final assignment occurrence in a dotenv document. + + Returns: + dict[str, int]: Final parsed binding index for each assigned variable. + """ + return { + binding.key: binding_index + for binding_index, binding in enumerate(parse_stream(io.StringIO(document))) + if binding.key is not None + } + + def _serialize_terminal_dotenv_value(value: str) -> str: """ Quote a terminal secret value for a native python-dotenv round trip. @@ -699,7 +726,7 @@ def _serialize_terminal_dotenv_value(value: str) -> str: Returns: str: A single-quoted dotenv value. """ - escaped_value = value.replace("'", "\\'").replace("${", "${:-$}{") + escaped_value = value.replace("\\", "\\\\").replace("'", "\\'").replace("${", "${:-$}{") return f"'{escaped_value}'" diff --git a/tests/unit/setup/test_akv_initialization.py b/tests/unit/setup/test_akv_initialization.py index c0f3d465e4..b981c7d1e3 100644 --- a/tests/unit/setup/test_akv_initialization.py +++ b/tests/unit/setup/test_akv_initialization.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import io import os import pathlib import tempfile @@ -20,6 +21,7 @@ _load_environment_files, _parse_akv_reference, _parse_akv_secret_url, + _serialize_terminal_dotenv_value, _warn_about_akv_environment_files, _write_akv_env_file, ) @@ -299,6 +301,25 @@ async def test_load_environment_files_honors_python_dotenv_disabled(self): assert loaded is True assert "DISABLED_VALUE" not in os.environ + @pytest.mark.parametrize( + "value", + [ + "single\\backslash", + "double\\\\backslash", + "four\\\\\\\\backslashes", + "\\leading-and-trailing\\", + r"C:\Users\name\secret.txt", + "quote'\\${LITERAL}\nline\\two", + ], + ) + def test_serialize_terminal_dotenv_value_preserves_backslashes(self, value): + document = f"VALUE={_serialize_terminal_dotenv_value(value)}\n" + + with mock.patch.dict(os.environ, {}, clear=True): + reloaded_value = dotenv_values(stream=io.StringIO(document), interpolate=True)["VALUE"] + + assert reloaded_value == value + async def test_load_environment_async_write_env_writes_resolved_native_bootstrap(self): credential, client = _create_mock_akv_clients() document = ( @@ -456,15 +477,15 @@ def test_write_akv_env_file_secures_descriptor_before_writing(self): side_effect=lambda *args, **kwargs: events.append("fdopen") or stream, ), mock.patch( - "pyrit.setup.akv_initialization.os.replace", - side_effect=lambda *args: events.append("replace"), + "pyrit.setup.akv_initialization.os.link", + side_effect=lambda *args: events.append("link"), ), ): _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) - assert events == ["create", "fchmod", "fdopen", "write:VALUE=bootstrap\n", "replace"] + assert events == ["create", "fchmod", "fdopen", "write:VALUE=bootstrap\n", "link"] - def test_write_akv_env_file_preserves_existing_file_when_replace_fails(self): + def test_write_akv_env_file_rejects_existing_file(self): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) env_file = temp_path / ".env" @@ -472,14 +493,44 @@ def test_write_akv_env_file_preserves_existing_file_when_replace_fails(self): with ( mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - mock.patch("pyrit.setup.akv_initialization.os.replace", side_effect=OSError("replace failed")), - pytest.raises(OSError, match="replace failed"), + pytest.raises(ValueError, match="already exists.*rename or remove"), ): _write_akv_env_file(documents=["NEW=value\n"], silent=True) assert env_file.read_text(encoding="utf-8") == "ORIGINAL=value\n" assert list(temp_path.glob(".env.*.tmp")) == [] + def test_write_akv_env_file_does_not_clobber_file_created_before_publish(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + real_fdopen = os.fdopen + + class CompetingFileStream: + def __init__(self, *args, **kwargs): + self._stream = real_fdopen(*args, **kwargs) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + result = self._stream.__exit__(exc_type, exc_value, traceback) + env_file.write_text("CREATED_BY_OTHER_PROCESS=value\n", encoding="utf-8") + return result + + def write(self, content): + return self._stream.write(content) + + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.akv_initialization.os.fdopen", side_effect=CompetingFileStream), + pytest.raises(ValueError, match="already exists.*rename or remove"), + ): + _write_akv_env_file(documents=["NEW=value\n"], silent=True) + + assert env_file.read_text(encoding="utf-8") == "CREATED_BY_OTHER_PROCESS=value\n" + assert list(temp_path.glob(".env.*.tmp")) == [] + @pytest.mark.skipif(os.name != "posix", reason="POSIX permission bits are not enforced on this platform.") def test_write_akv_env_file_uses_owner_only_permissions(self): with tempfile.TemporaryDirectory() as temp_dir: @@ -917,6 +968,57 @@ async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_secret client.__aexit__.assert_awaited_once() mock_print_msg.assert_called_once() + @pytest.mark.parametrize( + ("document", "expected_values", "expected_child_fetches"), + [ + ( + "A=kv:https://myvault.vault.azure.net/secrets/key\nB=${A}\nA=literal\n", + {"A": "literal", "B": "resolved-key"}, + 1, + ), + ( + "A=literal\nB=${A}\nA=kv:https://myvault.vault.azure.net/secrets/key\n", + {"A": "resolved-key", "B": "literal"}, + 1, + ), + ( + "A=kv:https://myvault.vault.azure.net/secrets/key\nB=${A}\nC=${B}\nB=literal\n", + {"A": "resolved-key", "B": "literal", "C": "resolved-key"}, + 2, + ), + ], + ) + async def test_debug_document_matches_runtime_for_interpolated_reference_assignments( + self, document, expected_values, expected_child_fetches + ): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + side_effect=[types.SimpleNamespace(value=document)] + + [types.SimpleNamespace(value="resolved-key")] * expected_child_fetches + ) + + 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), + ): + rendered_document = await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + resolve_references_for_output=True, + ) + runtime_values = {name: os.environ[name] for name in expected_values} + + with mock.patch.dict(os.environ, {}, clear=True): + reloaded_values = dict(dotenv_values(stream=io.StringIO(rendered_document), interpolate=True)) + + assert runtime_values == expected_values + assert reloaded_values == expected_values + assert ( + client.get_secret.await_args_list + == [mock.call("bootstrap", version=None)] + [mock.call("key", version=None)] * expected_child_fetches + ) + async def test_load_env_from_akv_async_preserves_process_values_without_fetching_overridden_child(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock( diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 56b2fa479d..5299a38441 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -249,8 +249,10 @@ async def test_initialize_forwards_boolean_akv_options(self, env_akv_strict, env load_defaults=False, ) - assert mock_load_environment.await_args.kwargs["env_akv_strict"] is env_akv_strict - assert mock_load_environment.await_args.kwargs["env_akv_write_env"] is env_akv_write_env + await_args = mock_load_environment.await_args + assert await_args is not None + assert await_args.kwargs["env_akv_strict"] is env_akv_strict + assert await_args.kwargs["env_akv_write_env"] is env_akv_write_env @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") async def test_initialize_keeps_akv_values_when_local_file_loading_fails(self, mock_set_memory): From b7b80d0aecd742de6c43d1d0c442c8b5b44e7467 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 20 Aug 2026 16:11:06 -0400 Subject: [PATCH 23/28] FIX: Removed .env_example drift features, renamed files, rescoped imports --- .azuredevops/test-job-template.yml | 4 - .env_example | 381 ++++++++---------- build_scripts/env_local_integration_test | 10 +- doc/code/executor/gcg/1_gcg_azure_ml.ipynb | 4 +- doc/code/executor/gcg/1_gcg_azure_ml.py | 4 +- .../executor/promptgen/gcg/experiments/run.py | 4 +- pyrit/setup/configuration_loader.py | 4 +- ...itialization.py => environment_loading.py} | 139 ++++--- pyrit/setup/initialization.py | 6 +- pyrit/setup/initializers/targets.py | 32 +- .../promptgen/gcg/test_gcg_aml_e2e.py | 4 +- .../setup/test_env_example_drift.py | 145 ------- .../targets/test_targets_and_secrets.py | 46 +-- .../promptgen/gcg/test_data_and_config.py | 4 +- ...ization.py => test_environment_loading.py} | 294 ++++++++++---- tests/unit/setup/test_initialization.py | 44 +- tests/unit/setup/test_targets_initializer.py | 54 ++- 17 files changed, 615 insertions(+), 564 deletions(-) rename pyrit/setup/{akv_initialization.py => environment_loading.py} (88%) delete mode 100644 tests/integration/setup/test_env_example_drift.py rename tests/unit/setup/{test_akv_initialization.py => test_environment_loading.py} (81%) diff --git a/.azuredevops/test-job-template.yml b/.azuredevops/test-job-template.yml index 7ed5c8ef5e..bd07c9370e 100644 --- a/.azuredevops/test-job-template.yml +++ b/.azuredevops/test-job-template.yml @@ -108,14 +108,10 @@ jobs: cp -r $PyRIT_DIR/doc $NEW_DIR cp -r $PyRIT_DIR/assets $NEW_DIR cp -r $PyRIT_DIR/tests/${{ parameters.testsFolder }} $NEW_DIR/tests - cp $PyRIT_DIR/.env_example $NEW_DIR/.env_example cd $NEW_DIR displayName: "Create and switch to new test directory" - task: AzureCLI@2 displayName: "Authenticate with service principal, cache Cognitive Services access token, and run tests" - env: - PYRIT_ENV_EXAMPLE_PATH: $(Build.SourcesDirectory)/../${{ parameters.newDir }}/.env_example - PYRIT_REPOSITORY_ROOT: $(Build.SourcesDirectory) inputs: azureSubscription: ${{ parameters.testAzureSubscription }} scriptType: 'bash' diff --git a/.env_example b/.env_example index 34468eba8c..0d70a48cfd 100644 --- a/.env_example +++ b/.env_example @@ -2,134 +2,138 @@ # 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" -# -# URL placeholders intentionally use plain dotenv values. Earlier versions used -# angle brackets as documentation styling, but python-dotenv preserves them literally +# OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" # 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_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" +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 - +# for distribution of requests across them for rate-limit relief. AZURE_OPENAI_GPT4O_ENDPOINT2="https://xxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT4O_KEY2="xxxxx" AZURE_OPENAI_GPT4O_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNDERLYING_MODEL2="gpt-4o" -AZURE_OPENAI_GPT4O_AAD_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT4O_AAD_KEY="xxxxx" -AZURE_OPENAI_GPT4O_AAD_MODEL="deployment-name" +AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_INTEGRATION_TEST_KEY="xxxxx" +AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" +AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" -# 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_GPT3_5_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +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_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_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" +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_KEY2="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" -# Objective Scorer chat target (used in scorers in scenarios) - -OBJECTIVE_SCORER_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" -OBJECTIVE_SCORER_CHAT_UNDERLYING_MODEL="" - -AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" -AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" - -AZURE_OPENAI_GPT5_COMPLETIONS_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT5_COMPLETIONS_MODEL="gpt-5" -AZURE_OPENAI_GPT5_COMPLETIONS_UNDERLYING_MODEL="gpt-5" - -AZURE_OPENAI_GPT5_4_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" -AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" +# 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_KEY="xxxxx" +ADVERSARIAL_CHAT_MODEL="deployment-name" -AZURE_OPENAI_GPT4O_STRICT_FILTER_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT4O_STRICT_FILTER_MODEL="deployment-name" +ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" +ADVERSARIAL_CHAT_SINGLETURN_MODEL="deployment-name" -AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT3_5_CHAT_MODEL="deployment-name" -AZURE_OPENAI_GPT3_5_CHAT_UNDERLYING_MODEL="" +ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_MULTITURN_KEY="xxxxx" +ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" -AZURE_OPENAI_GPT4_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT4_CHAT_MODEL="deployment-name" -AZURE_OPENAI_GPT4_CHAT_UNDERLYING_MODEL="" +ADVERSARIAL_CHAT_REASONING_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_REASONING_KEY="xxxxx" +ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" -MAI_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -MAI_CHAT_MODEL="deployment-name" -MAI_CHAT_KEY="xxxxx" -MAI_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPTV_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPTV_CHAT_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_KEY="xxxxx" +OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="https://xxxxx.eastus2.models.ai.azure.com" AZURE_FOUNDRY_DEEPSEEK_KEY="xxxxx" AZURE_FOUNDRY_DEEPSEEK_MODEL="" + AZURE_FOUNDRY_PHI4_ENDPOINT="https://xxxxx.models.ai.azure.com" 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_KEY="xxxxx" AZURE_FOUNDRY_MISTRAL_LARGE_MODEL="Mistral-Large-3" -OLLAMA_CHAT_ENDPOINT="http://127.0.0.1:11434/v1" -OLLAMA_MODEL="llama2" -AZURE_OPENAI_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_RESPONSES_MODEL="o4-mini" -AZURE_OPENAI_RESPONSES_UNDERLYING_MODEL="o4-mini" -AZURE_OPENAI_O4MINI_STRICT_FILTER_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_O4MINI_STRICT_FILTER_RESPONSES_MODEL="o4-mini" +AWS_ENDPOINT="https://bedrock-mantle.us-east-1.api.aws/v1" +AWS_KEY="xxxxx" +AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" +AWS_RESPONSES_MODEL="openai.gpt-oss-120b" -AZURE_OPENAI_GPT41_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT41_RESPONSES_MODEL="gpt-4.1" +GROQ_ENDPOINT="https://api.groq.com/openai/v1" +GROQ_KEY="gsk_xxxxxxxx" +GROQ_LLAMA_MODEL="llama3-8b-8192" -AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" -AZURE_OPENAI_GPT5_MODEL="gpt-5" -AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" +OPEN_ROUTER_ENDPOINT="https://openrouter.ai/api/v1" +OPEN_ROUTER_KEY="sk-or-v1-xxxxx" +OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" -PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" -PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" -PLATFORM_OPENAI_CHAT_MODEL="gpt-4o" -PLATFORM_OPENAI_RESPONSES_ENDPOINT="https://api.openai.com/v1" -PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" -PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" +OLLAMA_CHAT_ENDPOINT="http://127.0.0.1:11434/v1" +OLLAMA_MODEL="llama2" DEFAULT_OPENAI_FRONTEND_ENDPOINT = ${AZURE_OPENAI_GPT4O_AAD_ENDPOINT} DEFAULT_OPENAI_FRONTEND_KEY = ${AZURE_OPENAI_GPT4O_AAD_KEY} @@ -138,11 +142,29 @@ 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_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_KEY="sk-xxxxx" +PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" + +AZURE_OPENAI_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +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} @@ -150,16 +172,20 @@ 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" -PLATFORM_OPENAI_REALTIME_ENDPOINT="wss://api.openai.com/v1" -PLATFORM_OPENAI_REALTIME_KEY="sk-xxxxx" -PLATFORM_OPENAI_REALTIME_API_KEY="sk-xxxxx" -PLATFORM_OPENAI_REALTIME_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} @@ -167,182 +193,129 @@ OPENAI_REALTIME_UNDERLYING_MODEL = "" ################################## # IMAGE TARGET SECRETS -################################## - -AZURE_OPENAI_IMAGE_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_IMAGE_API_KEY1 = "xxxxxx" -AZURE_OPENAI_IMAGE_MODEL1 = "deployment-name" -AZURE_OPENAI_IMAGE_UNDERLYING_MODEL1 = "dall-e-3" -AZURE_OPENAI_IMAGE_ENDPOINT2 = "https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_IMAGE_API_KEY2 = "xxxxxx" -AZURE_OPENAI_IMAGE_MODEL2 = "dall-e-3" -AZURE_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} +# +# 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_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_API_KEY2 = "sk-xxxxx" +OPENAI_IMAGE_MODEL2 = "dall-e-3" +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_UNDERLYING_MODEL = "" -OPENAI_IMAGE_STRICT_FILTER_ENDPOINT = "https://xxxxx.openai.azure.com/openai/v1" -OPENAI_IMAGE_STRICT_FILTER_MODEL = "gpt-image" + ################################## # TTS TARGET SECRETS -################################## - -AZURE_OPENAI_TTS_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_TTS_KEY1 = "xxxxxxx" -AZURE_OPENAI_TTS_MODEL1 = "tts" -AZURE_OPENAI_TTS_UNDERLYING_MODEL1 = "tts" -AZURE_OPENAI_TTS_ENDPOINT2 = "https://xxxxx.openai.azure.com/v1" -AZURE_OPENAI_TTS_KEY2 = "xxxxxx" -AZURE_OPENAI_TTS_MODEL2 = "tts-1" -AZURE_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} +# +# 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_KEY1 = "xxxxxxx" +OPENAI_TTS_MODEL1 = "tts" +OPENAI_TTS_UNDERLYING_MODEL1 = "tts" + +OPENAI_TTS_ENDPOINT2 = "https://api.openai.com/v1" +OPENAI_TTS_KEY2 = "xxxxxx" +OPENAI_TTS_MODEL2 = "tts-1" +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_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 +################################### +# Note: Use the base URL without API path AZURE_OPENAI_VIDEO_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/openai/v1" 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="https://xxxxx.openai.azure.com/openai/v1" -ADVERSARIAL_CHAT_MODEL="deployment-name" -ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" -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_KEY="xxxxx" -ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" -ADVERSARIAL_CHAT_REASONING_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" -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="https://xxxxxx.westus3.inference.ml.azure.com/score" 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_API_KEY="xxxxx" +OPENAI_COMPLETION_MODEL="davinci-002" OPENAI_EMBEDDING_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +OPENAI_EMBEDDING_KEY="xxxxx" OPENAI_EMBEDDING_MODEL="text-embedding-3-small" -AZURE_SPEECH_REGION = "eastus2" -# Resource ID is needed when using Entra authentication +AZURE_STORAGE_ACCOUNT_CONTAINER_URL="https://xxxxxx.blob.core.windows.net/xpia" +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/" + HUGGINGFACE_TOKEN="hf_xxxxxxx" HUGGINGFACE_ENDPOINT="https://router.huggingface.co/v1" -################################## +GOOGLE_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/openai" +GOOGLE_GEMINI_API_KEY = "xxxxx" +GOOGLE_GEMINI_MODEL="gemini-2.0-flash" + + +######################### # AZURE SQL SECRETS -################################## -# This connects to the test database +######################### + +# 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" # 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_CONTAINER_URL="https://xxxxxx.blob.core.windows.net/xpia" -# 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} - -################################## -# INTEGRATION TEST ONLY SECRETS -################################## - -GOOGLE_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/openai" -GOOGLE_GEMINI_API_KEY = "xxxxx" -GOOGLE_GEMINI_MODEL="gemini-2.0-flash" - -ANTHROPIC_CHAT_ENDPOINT="https://api.anthropic.com/v1" -ANTHROPIC_CHAT_KEY="xxxxx" -ANTHROPIC_CHAT_MODEL="claude-3-7-sonnet-latest" - -AWS_KEY="xxxxx" -AWS_ENDPOINT="https://bedrock-mantle.us-east-1.api.aws/v1" -AWS_RESPONSES_MODEL="openai.gpt-oss-120b" -AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" - -PLATFORM_OPENAI_VIDEO_ENDPOINT="https://api.openai.com/v1" -PLATFORM_OPENAI_VIDEO_KEY="sk-xxxxx" -PLATFORM_OPENAI_VIDEO_MODEL="sora-2" - -PLATFORM_OPENAI_IMAGE_ENDPOINT="https://api.openai.com/v1" -PLATFORM_OPENAI_IMAGE_KEY="sk-xxxxx" -PLATFORM_OPENAI_IMAGE_MODEL="gpt-image-1" - -PLATFORM_OPENAI_EMBEDDING_ENDPOINT="https://api.openai.com/v1" -PLATFORM_OPENAI_EMBEDDING_KEY="sk-xxxxx" -PLATFORM_OPENAI_EMBEDDING_MODEL="text-embedding-3-small" - -OPENAI_COMPLETION_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -OPENAI_COMPLETION_API_KEY="xxxxx" -OPENAI_COMPLETION_MODEL="davinci-002" - -PROMPTINTEL_API_KEY="xxxxx" - -################################## -# Additional entries referenced in PyRIT -################################## - -AZURE_OPENAI_GPT4O_KEY="xxxxx" -AZURE_OPENAI_GPT4O_KEY2="xxxxx" -AZURE_OPENAI_INTEGRATION_TEST_KEY="xxxxx" -AZURE_OPENAI_GPT3_5_CHAT_KEY="xxxxx" -AZURE_OPENAI_GPT4_CHAT_KEY="xxxxx" -AZURE_OPENAI_GPT5_4_KEY="xxxxx" -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY="xxxxx" -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" -ADVERSARIAL_CHAT_KEY="xxxxx" -OBJECTIVE_SCORER_CHAT_KEY="xxxxx" -AZURE_FOUNDRY_MISTRAL_LARGE_KEY="xxxxx" -GROQ_ENDPOINT="https://api.groq.com/openai/v1" -GROQ_KEY="gsk_xxxxxxxx" -GROQ_LLAMA_MODEL="llama3-8b-8192" -OPEN_ROUTER_ENDPOINT="https://openrouter.ai/api/v1" -OPEN_ROUTER_KEY="sk-or-v1-xxxxx" -OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" -AZURE_OPENAI_GPT5_KEY="xxxxxxx" -AZURE_OPENAI_RESPONSES_KEY="xxxxx" -OPENAI_EMBEDDING_KEY="xxxxx" -AZURE_STORAGE_ACCOUNT_SAS_TOKEN="xxxxx" -AZURE_SPEECH_KEY = "xxxxx" -AZURE_CONTENT_SAFETY_API_KEY="xxxxx" diff --git a/build_scripts/env_local_integration_test b/build_scripts/env_local_integration_test index ded45e2bea..b61cfe7d20 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=${AZURE_OPENAI_IMAGE_ENDPOINT2} -OPENAI_IMAGE_API_KEY=${AZURE_OPENAI_IMAGE_API_KEY2} -OPENAI_IMAGE_MODEL=${AZURE_OPENAI_IMAGE_MODEL2} +OPENAI_IMAGE_ENDPOINT=${OPENAI_IMAGE_ENDPOINT2} +OPENAI_IMAGE_API_KEY=${OPENAI_IMAGE_API_KEY2} +OPENAI_IMAGE_MODEL=${OPENAI_IMAGE_MODEL2} -OPENAI_TTS_ENDPOINT=${AZURE_OPENAI_TTS_ENDPOINT2} -OPENAI_TTS_KEY=${AZURE_OPENAI_TTS_KEY2} +OPENAI_TTS_ENDPOINT=${OPENAI_TTS_ENDPOINT2} +OPENAI_TTS_KEY=${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} diff --git a/doc/code/executor/gcg/1_gcg_azure_ml.ipynb b/doc/code/executor/gcg/1_gcg_azure_ml.ipynb index 264805db7a..f2053f9fb2 100644 --- a/doc/code/executor/gcg/1_gcg_azure_ml.ipynb +++ b/doc/code/executor/gcg/1_gcg_azure_ml.ipynb @@ -57,9 +57,9 @@ "source": [ "import os\n", "\n", - "from pyrit.setup.akv_initialization import _load_environment_files\n", + "from pyrit.setup.environment_loading import load_environment_files\n", "\n", - "_load_environment_files(env_files=None)\n", + "load_environment_files(env_files=None)\n", "\n", "subscription_id = os.environ.get(\"AZURE_ML_SUBSCRIPTION_ID\")\n", "resource_group = os.environ.get(\"AZURE_ML_RESOURCE_GROUP\")\n", diff --git a/doc/code/executor/gcg/1_gcg_azure_ml.py b/doc/code/executor/gcg/1_gcg_azure_ml.py index 9e05233255..0e4b5fe40c 100644 --- a/doc/code/executor/gcg/1_gcg_azure_ml.py +++ b/doc/code/executor/gcg/1_gcg_azure_ml.py @@ -29,9 +29,9 @@ # %% import os -from pyrit.setup.akv_initialization import _load_environment_files +from pyrit.setup.environment_loading import load_environment_files -_load_environment_files(env_files=None) +load_environment_files(env_files=None) subscription_id = os.environ.get("AZURE_ML_SUBSCRIPTION_ID") resource_group = os.environ.get("AZURE_ML_RESOURCE_GROUP") diff --git a/pyrit/executor/promptgen/gcg/experiments/run.py b/pyrit/executor/promptgen/gcg/experiments/run.py index c5a2c84bb0..3bad3dc2f9 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.akv_initialization import _load_environment_files +from pyrit.setup.environment_loading import load_environment_files def _parse_arguments() -> argparse.Namespace: @@ -85,7 +85,7 @@ def _resolve_output(*, output: GCGOutputConfig, output_dir: str | None) -> GCGOu async def _main_async(config_path: str, data_path: str, output_dir: str | None = None) -> None: - _load_environment_files(env_files=None) + load_environment_files(env_files=None) config = GCGConfig.from_json_file(config_path) data = GCGDataConfig.from_json_file(data_path) if config.hf_token is None: diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 31e4f3e4cf..178142e0dc 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -21,7 +21,7 @@ from pyrit.common.utils import verify_and_resolve_path from pyrit.common.yaml_loadable import YamlLoadable from pyrit.models import class_name_to_snake_case -from pyrit.setup.akv_initialization import _validate_akv_boolean_options +from pyrit.setup.environment_loading import validate_akv_boolean_options from pyrit.setup.initialization import ( AZURE_SQL, IN_MEMORY, @@ -154,7 +154,7 @@ class ConfigurationLoader(YamlLoadable): def __post_init__(self) -> None: """Validate and normalize the configuration after loading.""" - _validate_akv_boolean_options( + validate_akv_boolean_options( env_akv_strict=self.env_akv_strict, env_akv_write_env=self.env_akv_write_env, ) diff --git a/pyrit/setup/akv_initialization.py b/pyrit/setup/environment_loading.py similarity index 88% rename from pyrit/setup/akv_initialization.py rename to pyrit/setup/environment_loading.py index 0d3c564ab0..eba981b678 100644 --- a/pyrit/setup/akv_initialization.py +++ b/pyrit/setup/environment_loading.py @@ -5,14 +5,15 @@ import asyncio import contextlib -import io import logging import os import pathlib import tempfile import urllib.parse from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any +from dataclasses import dataclass +from io import StringIO +from typing import TYPE_CHECKING import dotenv from dotenv.parser import parse_stream @@ -21,6 +22,7 @@ from pyrit.exceptions import KeyVaultInitializationException if TYPE_CHECKING: + from azure.core.credentials_async import AsyncTokenCredential from azure.keyvault.secrets.aio import SecretClient logger = logging.getLogger(__name__) @@ -33,7 +35,15 @@ _LEGACY_ENV_REMOVED_IN = "1.3.0" -def _validate_akv_boolean_options(*, env_akv_strict: object, env_akv_write_env: object) -> None: +@dataclass(frozen=True) +class _EnvironmentValueCandidate: + """An ordered environment value and whether local AKV resolution applies.""" + + value: str + resolve_akv_reference: bool + + +def validate_akv_boolean_options(*, env_akv_strict: object, env_akv_write_env: object) -> None: """ Require real booleans for Key Vault behavior flags. @@ -48,12 +58,12 @@ def _validate_akv_boolean_options(*, env_akv_strict: object, env_akv_write_env: raise TypeError(f"{option_name} must be a bool, got {type(option_value).__name__}.") -def _load_environment_files( +def load_environment_files( env_files: Sequence[pathlib.Path] | None, *, silent: bool = False, include_default_base: bool = True, - assignment_fallbacks: dict[str, str | None] | None = None, + assignment_candidates: dict[str, list[_EnvironmentValueCandidate]] | None = None, ) -> bool: """ Load environment files in the order they are provided. @@ -68,8 +78,9 @@ def _load_environment_files( 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. - assignment_fallbacks: Optional output mapping from assignments that win - precedence to the value they replaced, if any. + assignment_candidates: Optional output mapping containing each applicable + local value in ascending precedence order. Existing process or AKV values + are retained as non-resolvable baseline candidates. Returns: True if at least one environment file was loaded, otherwise False. @@ -84,16 +95,30 @@ def _load_environment_files( ) for env_file in selected_files: override = env_file.name == ".env.local" - if assignment_fallbacks is not None: - assignment_names = dotenv.dotenv_values(dotenv_path=env_file, interpolate=False) - for variable_name in assignment_names: - if override or variable_name not in os.environ: - assignment_fallbacks[variable_name] = os.environ.get(variable_name) - dotenv.load_dotenv( + applicable_names: list[str] = [] + previous_values: dict[str, str | None] = {} + if assignment_candidates is not None: + assignment_values = dotenv.dotenv_values(dotenv_path=env_file, interpolate=False) + applicable_names = [ + variable_name + for variable_name, value in assignment_values.items() + if value is not None and (override or variable_name not in os.environ) + ] + previous_values = {variable_name: os.environ.get(variable_name) for variable_name in applicable_names} + loaded = dotenv.load_dotenv( dotenv_path=env_file, override=override, interpolate=True, ) + if assignment_candidates is not None and loaded: + for variable_name in applicable_names: + candidates = assignment_candidates.setdefault(variable_name, []) + previous_value = previous_values[variable_name] + if not candidates and previous_value is not None: + candidates.append(_EnvironmentValueCandidate(value=previous_value, resolve_akv_reference=False)) + loaded_value = os.environ.get(variable_name) + if loaded_value is not None: + candidates.append(_EnvironmentValueCandidate(value=loaded_value, resolve_akv_reference=True)) if not silent: _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) @@ -264,7 +289,7 @@ def _is_valid_akv_identifier(identifier: str) -> bool: ) -def _create_akv_secret_client(*, vault_url: str, credential: Any) -> "SecretClient": +def _create_akv_secret_client(*, vault_url: str, credential: "AsyncTokenCredential") -> "SecretClient": """ Create an asynchronous Key Vault client with an explicit retry policy. @@ -286,7 +311,7 @@ def _create_akv_secret_client(*, vault_url: str, credential: Any) -> "SecretClie async def _fetch_akv_secret_value_async( *, - client: Any, + client: "SecretClient", secret_name: str, secret_version: str | None, variable_name: str, @@ -343,7 +368,7 @@ def _validate_dotenv_document( Raises: ValueError: If strict is True and the document contains invalid entries. """ - bindings = list(parse_stream(io.StringIO(document))) + bindings = list(parse_stream(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] = [] @@ -419,12 +444,12 @@ async def _load_env_from_akv_async( 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) + parsed_environment = dotenv.dotenv_values(stream=StringIO(validated_document), interpolate=True) if not parsed_environment: raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") existing_environment_names = set(os.environ) loaded = dotenv.load_dotenv( - stream=io.StringIO(validated_document), + stream=StringIO(validated_document), override=False, interpolate=True, ) @@ -500,7 +525,7 @@ async def _load_env_from_akv_async( raise wrapped_error from error -async def _load_environment_async( +async def load_environment_async( *, env_akv_ref: Sequence[str] | None, env_files: Sequence[pathlib.Path] | None, @@ -561,16 +586,16 @@ async def _load_environment_async( written_path = written_env_file.resolve() selected_env_files = [env_file for env_file in env_files if env_file.expanduser().resolve() != written_path] - assignment_fallbacks: dict[str, str | None] = {} + assignment_candidates: dict[str, list[_EnvironmentValueCandidate]] = {} await asyncio.to_thread( - _load_environment_files, + load_environment_files, env_files=selected_env_files, silent=silent, include_default_base=not (env_akv_ref and env_files is None), - assignment_fallbacks=assignment_fallbacks, + assignment_candidates=assignment_candidates, ) await _resolve_local_akv_references_async( - assignment_fallbacks=assignment_fallbacks, + assignment_candidates=assignment_candidates, strict=env_akv_strict, silent=silent, ) @@ -656,7 +681,7 @@ def _merge_akv_documents_for_debug(*, documents: Sequence[str]) -> str: merged_bindings: list[str] = [] for document in documents: document_names: set[str] = set() - for binding in parse_stream(io.StringIO(document)): + for binding in parse_stream(StringIO(document)): if binding.key is None or binding.key not in established_names: merged_bindings.append(binding.original.string) if binding.key is not None: @@ -680,7 +705,7 @@ def _render_resolved_akv_document( """ skipped_reference_indexes = skipped_reference_indexes or set() rendered_bindings: list[str] = [] - for binding_index, binding in enumerate(parse_stream(io.StringIO(document))): + for binding_index, binding in enumerate(parse_stream(StringIO(document))): variable_name = binding.key if binding_index in skipped_reference_indexes: continue @@ -711,7 +736,7 @@ def _get_final_assignment_indexes(*, document: str) -> dict[str, int]: """ return { binding.key: binding_index - for binding_index, binding in enumerate(parse_stream(io.StringIO(document))) + for binding_index, binding in enumerate(parse_stream(StringIO(document))) if binding.key is not None } @@ -791,7 +816,7 @@ def _parse_akv_reference_url( async def _resolve_local_akv_references_async( *, - assignment_fallbacks: Mapping[str, str | None], + assignment_candidates: Mapping[str, Sequence[_EnvironmentValueCandidate]], strict: bool, silent: bool, ) -> None: @@ -802,36 +827,38 @@ async def _resolve_local_akv_references_async( KeyVaultInitializationException: If strict validation or secret retrieval fails. """ parsed_references: list[tuple[str, str, str, str | None]] = [] - for variable_name, fallback_value in assignment_fallbacks.items(): - value = os.environ.get(variable_name) - if value is None: - continue - target = _parse_akv_reference(value) - if target is None: - continue - try: - vault_url, secret_name, secret_version = _parse_akv_reference_url( - target=target, - variable_name=variable_name, - ) - except ValueError as error: - if strict: - wrapped_error = _key_vault_initialization_error( - message=f"Invalid AKV reference for environment variable '{variable_name}'", + for variable_name, candidates in assignment_candidates.items(): + for candidate in reversed(candidates): + if not candidate.resolve_akv_reference: + os.environ[variable_name] = candidate.value + break + target = _parse_akv_reference(candidate.value) + if target is None: + os.environ[variable_name] = candidate.value + break + try: + vault_url, secret_name, secret_version = _parse_akv_reference_url( + target=target, + variable_name=variable_name, + ) + except ValueError as error: + if strict: + wrapped_error = _key_vault_initialization_error( + message=f"Invalid AKV reference for environment variable '{variable_name}'", + error=error, + ) + raise wrapped_error from error + _warn_about_invalid_akv_reference( + variable_name=variable_name, error=error, + silent=silent, ) - raise wrapped_error from error - if fallback_value is None: - os.environ.pop(variable_name, None) - else: - os.environ[variable_name] = fallback_value - _warn_about_invalid_akv_reference( - variable_name=variable_name, - error=error, - silent=silent, - ) - continue - parsed_references.append((variable_name, vault_url, secret_name, secret_version)) + continue + os.environ[variable_name] = candidate.value + parsed_references.append((variable_name, vault_url, secret_name, secret_version)) + break + else: + os.environ.pop(variable_name, None) if not parsed_references: return @@ -840,7 +867,7 @@ async def _resolve_local_akv_references_async( async with DefaultAzureCredential() as credential: async with contextlib.AsyncExitStack() as client_stack: - clients: dict[str, Any] = {} + clients: dict[str, SecretClient] = {} for variable_name, vault_url, secret_name, secret_version in parsed_references: try: client = clients.get(vault_url) diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index fba743e2fc..df921d5630 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -7,7 +7,7 @@ from pyrit.common.apply_defaults import reset_default_values from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory -from pyrit.setup.akv_initialization import _load_environment_async, _validate_akv_boolean_options +from pyrit.setup.environment_loading import load_environment_async, validate_akv_boolean_options if TYPE_CHECKING: from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -112,11 +112,11 @@ async def initialize_pyrit_async( TypeError: If ``env_akv_strict`` or ``env_akv_write_env`` is not a bool. ValueError: If an unsupported memory_db_type is provided or env_files contains non-existent files. """ - _validate_akv_boolean_options( + validate_akv_boolean_options( env_akv_strict=env_akv_strict, env_akv_write_env=env_akv_write_env, ) - await _load_environment_async( + await load_environment_async( env_akv_ref=env_akv_ref, env_files=env_files, env_akv_strict=env_akv_strict, diff --git a/pyrit/setup/initializers/targets.py b/pyrit/setup/initializers/targets.py index 790b51a87f..308366f734 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="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", + endpoint_var="OPENAI_IMAGE_ENDPOINT1", + key_var="OPENAI_IMAGE_API_KEY1", + model_var="OPENAI_IMAGE_MODEL1", + underlying_model_var="OPENAI_IMAGE_UNDERLYING_MODEL1", ), TargetConfig( registry_name="openai_image_platform", target_class=OpenAIImageTarget, - 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", + endpoint_var="OPENAI_IMAGE_ENDPOINT2", + key_var="OPENAI_IMAGE_API_KEY2", + model_var="OPENAI_IMAGE_MODEL2", + underlying_model_var="OPENAI_IMAGE_UNDERLYING_MODEL2", ), # ============================================ # TTS Targets (OpenAITTSTarget) @@ -357,18 +357,18 @@ class TargetConfig: TargetConfig( registry_name="openai_tts_azure", target_class=OpenAITTSTarget, - 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", + endpoint_var="OPENAI_TTS_ENDPOINT1", + key_var="OPENAI_TTS_KEY1", + model_var="OPENAI_TTS_MODEL1", + underlying_model_var="OPENAI_TTS_UNDERLYING_MODEL1", ), TargetConfig( registry_name="openai_tts_platform", target_class=OpenAITTSTarget, - 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", + endpoint_var="OPENAI_TTS_ENDPOINT2", + key_var="OPENAI_TTS_KEY2", + model_var="OPENAI_TTS_MODEL2", + underlying_model_var="OPENAI_TTS_UNDERLYING_MODEL2", ), # ============================================ # Video Targets (OpenAIVideoTarget) 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 8b6ed6ce40..34b67f0cac 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.akv_initialization import _load_environment_files # noqa: E402 +from pyrit.setup.environment_loading import load_environment_files # noqa: E402 _REQUIRED_ENV_VARS = ( "AZURE_ML_SUBSCRIPTION_ID", @@ -70,7 +70,7 @@ def test_gcg_aml_notebook_runs_to_completion() -> None: MLClient from its namespace, then polls until the job reaches a terminal state and asserts ``Completed``. """ - _load_environment_files(env_files=None, silent=True) + load_environment_files(env_files=None, silent=True) missing = [name for name in _REQUIRED_ENV_VARS if not os.environ.get(name)] if missing: pytest.skip(f"Missing required env vars for GCG AML e2e test: {', '.join(missing)}") diff --git a/tests/integration/setup/test_env_example_drift.py b/tests/integration/setup/test_env_example_drift.py deleted file mode 100644 index 9daa065d23..0000000000 --- a/tests/integration/setup/test_env_example_drift.py +++ /dev/null @@ -1,145 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import os -import pathlib -import re -import subprocess -from unittest import mock - -from dotenv import dotenv_values - -_ENV_EXAMPLE_PATH_ENV = "PYRIT_ENV_EXAMPLE_PATH" -_REPOSITORY_ROOT_ENV = "PYRIT_REPOSITORY_ROOT" -_ENVIRONMENT_NAME_PATTERN = re.compile(r"(? pathlib.Path: - configured_root = os.getenv(_REPOSITORY_ROOT_ENV) - if configured_root: - root = pathlib.Path(configured_root) - if root.is_dir(): - return root - raise AssertionError(f"{_REPOSITORY_ROOT_ENV} does not identify a directory: {root}") - - result = subprocess.run( - ["git", "rev-parse", "--show-toplevel"], - capture_output=True, - check=False, - text=True, - ) - if result.returncode != 0: - raise AssertionError("Could not locate the repository root with git.") - return pathlib.Path(result.stdout.strip()) - - -def _get_env_example_path(*, repository_root: pathlib.Path) -> pathlib.Path: - configured_path = os.getenv(_ENV_EXAMPLE_PATH_ENV) - path = pathlib.Path(configured_path) if configured_path else repository_root / ".env_example" - if not path.is_file(): - raise AssertionError(f"Could not locate .env_example at {path}.") - return path - - -def _grep_repository_for_environment_names(*, environment_names: set[str], repository_root: pathlib.Path) -> set[str]: - grep_pattern = "(" + "|".join(sorted(environment_names)) + ")" - result = subprocess.run( - ["git", "grep", "-I", "-h", "-E", grep_pattern, "--", ".", ":(exclude).env_example"], - capture_output=True, - check=False, - cwd=repository_root, - text=True, - ) - if result.returncode not in {0, 1}: - raise AssertionError(f"Could not search tracked repository files with git: {result.stderr.strip()}") - return environment_names & set(_ENVIRONMENT_NAME_PATTERN.findall(result.stdout)) - - -def _find_referenced_environment_names( - *, - environment_names: set[str], - repository_root: pathlib.Path, - env_example_path: pathlib.Path, -) -> set[str]: - example_contents = env_example_path.read_text(encoding="utf-8") - example_without_assignment_names = _DOTENV_ASSIGNMENT_NAME_PATTERN.sub("=", example_contents) - referenced_names = environment_names & set(_ENVIRONMENT_NAME_PATTERN.findall(example_without_assignment_names)) - referenced_names.update( - _grep_repository_for_environment_names( - environment_names=environment_names, - repository_root=repository_root, - ) - ) - return referenced_names - - -def test_env_example_names_are_referenced_in_repository() -> None: - """Catch example entries with no weak textual reference in tracked repository files.""" - repository_root = _get_repository_root() - env_example_path = _get_env_example_path(repository_root=repository_root) - environment_names = set(dotenv_values(dotenv_path=env_example_path, interpolate=False)) - assert environment_names, ".env_example contains no dotenv assignments." - - referenced_names = _find_referenced_environment_names( - environment_names=environment_names, - repository_root=repository_root, - env_example_path=env_example_path, - ) - unreferenced_names = environment_names - referenced_names - assert not unreferenced_names, ".env_example contains names with no tracked repository reference: " + ", ".join( - sorted(unreferenced_names) - ) - - -def test_env_example_url_values_are_not_wrapped_in_angle_brackets() -> None: - """Ensure URL placeholder styling does not become part of parsed dotenv values.""" - repository_root = _get_repository_root() - env_example_path = _get_env_example_path(repository_root=repository_root) - values = dotenv_values(dotenv_path=env_example_path, interpolate=False) - - wrapped_names = {name for name, value in values.items() if value and ("<" in value or ">" in value)} - assert not wrapped_names, ".env_example contains values wrapped in angle brackets: " + ", ".join( - sorted(wrapped_names) - ) - - -def test_env_example_comment_blocks_do_not_contain_blank_lines() -> None: - """Keep consecutive comment lines together so the example remains compact.""" - repository_root = _get_repository_root() - env_example_path = _get_env_example_path(repository_root=repository_root) - contents = env_example_path.read_text(encoding="utf-8") - - assert not _BLANK_LINE_BETWEEN_COMMENTS_PATTERN.search(contents), ( - ".env_example contains a blank line between consecutive comment lines." - ) - - -def test_env_example_aliases_resolve_in_assignment_order() -> None: - """Ensure complete-value aliases resolve to their sources without ambient environment values.""" - repository_root = _get_repository_root() - env_example_path = _get_env_example_path(repository_root=repository_root) - raw_values = dotenv_values(dotenv_path=env_example_path, interpolate=False) - aliases = { - name: match.group(1) - for name, value in raw_values.items() - if value and (match := _DOTENV_COMPLETE_REFERENCE_PATTERN.fullmatch(value)) - } - assert aliases, ".env_example contains no complete-value aliases." - - with mock.patch.dict(os.environ, {}, clear=True): - resolved_values = dotenv_values(dotenv_path=env_example_path, interpolate=True) - - unresolved_names = {name for name in aliases if not resolved_values.get(name)} - assert not unresolved_names, ".env_example contains aliases that resolve to empty values: " + ", ".join( - sorted(unresolved_names) - ) - - mismatched_names = { - name for name, source_name in aliases.items() if resolved_values[name] != resolved_values.get(source_name) - } - assert not mismatched_names, ".env_example contains aliases that differ from their sources: " + ", ".join( - sorted(mismatched_names) - ) diff --git a/tests/integration/targets/test_targets_and_secrets.py b/tests/integration/targets/test_targets_and_secrets.py index 2a15ae6397..e2ec9da733 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( - "AZURE_OPENAI_IMAGE_ENDPOINT1", + "OPENAI_IMAGE_ENDPOINT1", None, - "AZURE_OPENAI_IMAGE_MODEL1", + "OPENAI_IMAGE_MODEL1", marks=pytest.mark.run_only_if_all_tests, ), # gpt-image-1.5 pytest.param( - "AZURE_OPENAI_IMAGE_ENDPOINT1", - "AZURE_OPENAI_IMAGE_API_KEY1", - "AZURE_OPENAI_IMAGE_MODEL1", + "OPENAI_IMAGE_ENDPOINT1", + "OPENAI_IMAGE_API_KEY1", + "OPENAI_IMAGE_MODEL1", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-image1-api-key", ), - ("AZURE_OPENAI_IMAGE_ENDPOINT2", None, "AZURE_OPENAI_IMAGE_MODEL2"), # gpt-image-1 + ("OPENAI_IMAGE_ENDPOINT2", None, "OPENAI_IMAGE_MODEL2"), # gpt-image-1 pytest.param( - "AZURE_OPENAI_IMAGE_ENDPOINT2", - "AZURE_OPENAI_IMAGE_API_KEY2", - "AZURE_OPENAI_IMAGE_MODEL2", + "OPENAI_IMAGE_ENDPOINT2", + "OPENAI_IMAGE_API_KEY2", + "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( - "AZURE_OPENAI_IMAGE_API_KEY2", + "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("AZURE_OPENAI_IMAGE_ENDPOINT2") - model_name_value = os.getenv("AZURE_OPENAI_IMAGE_MODEL2") or "gpt-image-1" + endpoint_value = _get_required_env_var("OPENAI_IMAGE_ENDPOINT2") + model_name_value = os.getenv("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( - "AZURE_OPENAI_IMAGE_API_KEY2", + "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("AZURE_OPENAI_IMAGE_ENDPOINT2") - model_name_value = os.getenv("AZURE_OPENAI_IMAGE_MODEL2") or "gpt-image-1" + endpoint_value = _get_required_env_var("OPENAI_IMAGE_ENDPOINT2") + model_name_value = os.getenv("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"), [ - ("AZURE_OPENAI_TTS_ENDPOINT1", None, "AZURE_OPENAI_TTS_MODEL1"), + ("OPENAI_TTS_ENDPOINT1", None, "OPENAI_TTS_MODEL1"), pytest.param( - "AZURE_OPENAI_TTS_ENDPOINT1", - "AZURE_OPENAI_TTS_KEY1", - "AZURE_OPENAI_TTS_MODEL1", + "OPENAI_TTS_ENDPOINT1", + "OPENAI_TTS_KEY1", + "OPENAI_TTS_MODEL1", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-tts1-api-key", ), - ("AZURE_OPENAI_TTS_ENDPOINT2", None, "AZURE_OPENAI_TTS_MODEL2"), + ("OPENAI_TTS_ENDPOINT2", None, "OPENAI_TTS_MODEL2"), pytest.param( - "AZURE_OPENAI_TTS_ENDPOINT2", - "AZURE_OPENAI_TTS_KEY2", - "AZURE_OPENAI_TTS_MODEL2", + "OPENAI_TTS_ENDPOINT2", + "OPENAI_TTS_KEY2", + "OPENAI_TTS_MODEL2", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-tts2-api-key", ), diff --git a/tests/unit/executor/promptgen/gcg/test_data_and_config.py b/tests/unit/executor/promptgen/gcg/test_data_and_config.py index d41ffef2db..adc0eaa9ac 100644 --- a/tests/unit/executor/promptgen/gcg/test_data_and_config.py +++ b/tests/unit/executor/promptgen/gcg/test_data_and_config.py @@ -153,7 +153,7 @@ def test_override_falls_back_to_default_basename(self, tmp_path: Path) -> None: class TestMainAsyncCli: """Tests for ``run.py``'s ``--config`` + ``--data`` CLI wrapper around GCGGenerator.execute_async.""" - @patch("pyrit.executor.promptgen.gcg.experiments.run._load_environment_files") + @patch("pyrit.executor.promptgen.gcg.experiments.run.load_environment_files") async def test_raises_when_no_token_anywhere(self, mock_load_env: MagicMock, tmp_path: Path) -> None: config = GCGConfig(models=[GCGModelConfig(name="org/model")]) config_path = tmp_path / "config.json" @@ -166,7 +166,7 @@ async def test_raises_when_no_token_anywhere(self, mock_load_env: MagicMock, tmp with pytest.raises(ValueError, match="No HuggingFace token available"): await _main_async(str(config_path), str(data_path)) - @patch("pyrit.executor.promptgen.gcg.experiments.run._load_environment_files") + @patch("pyrit.executor.promptgen.gcg.experiments.run.load_environment_files") @patch("pyrit.executor.promptgen.gcg.experiments.run.load_goals_and_targets") async def test_passes_loaded_goals_to_generator_and_uses_env_token( self, diff --git a/tests/unit/setup/test_akv_initialization.py b/tests/unit/setup/test_environment_loading.py similarity index 81% rename from tests/unit/setup/test_akv_initialization.py rename to tests/unit/setup/test_environment_loading.py index b981c7d1e3..896da5d0d8 100644 --- a/tests/unit/setup/test_akv_initialization.py +++ b/tests/unit/setup/test_environment_loading.py @@ -15,22 +15,22 @@ from pyrit.exceptions import KeyVaultInitializationException from pyrit.setup import IN_MEMORY, initialize_pyrit_async -from pyrit.setup.akv_initialization import ( +from pyrit.setup.environment_loading import ( _load_env_from_akv_async, - _load_environment_async, - _load_environment_files, _parse_akv_reference, _parse_akv_secret_url, _serialize_terminal_dotenv_value, _warn_about_akv_environment_files, _write_akv_env_file, + load_environment_async, + load_environment_files, ) class TestLoadEnvironmentFiles: - """Tests for _load_environment_files function and env_files parameter in initialize_pyrit_async.""" + """Tests for load_environment_files and the env_files initialization parameter.""" - @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + @mock.patch("pyrit.setup.environment_loading.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: @@ -45,13 +45,13 @@ async def test_loads_default_env_files_when_none_provided(self, mock_config_path mock.patch.dict(os.environ, {}, clear=True), pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), ): - loaded = _load_environment_files(env_files=None) + 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") + @mock.patch("pyrit.setup.environment_loading.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: @@ -64,12 +64,12 @@ async def test_only_loads_existing_default_files(self, mock_config_path): mock.patch.dict(os.environ, {}, clear=True), pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), ): - loaded = _load_environment_files(env_files=None) + 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") + @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") async def test_default_env_preserves_process_environment(self, mock_config_path): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -80,13 +80,13 @@ async def test_default_env_preserves_process_environment(self, mock_config_path) mock.patch.dict(os.environ, {"VAR": "process"}, clear=True), pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), ): - loaded = _load_environment_files(env_files=None, silent=True) + loaded = load_environment_files(env_files=None, silent=True) assert loaded is True assert os.environ["VAR"] == "process" assert os.environ["LEGACY_ONLY"] == "legacy" - @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") async def test_default_env_local_overrides_process_environment_and_env(self, mock_config_path): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -98,12 +98,12 @@ async def test_default_env_local_overrides_process_environment_and_env(self, moc mock.patch.dict(os.environ, {"VAR": "process"}, clear=True), pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), ): - loaded = _load_environment_files(env_files=None, silent=True) + loaded = load_environment_files(env_files=None, silent=True) assert loaded is True assert os.environ["VAR"] == "local" - @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + @mock.patch("pyrit.setup.environment_loading.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) @@ -115,24 +115,24 @@ async def test_excludes_default_env_when_loading_local_override(self, mock_confi 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) + 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") + @mock.patch("pyrit.setup.environment_loading.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) + loaded = load_environment_files(env_files=None) assert loaded is False assert os.environ == {} - @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") def test_auto_discovered_env_warns_with_removal_version(self, mock_config_path, caplog, capsys): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -141,17 +141,17 @@ def test_auto_discovered_env_warns_with_removal_version(self, mock_config_path, mock_config_path.__truediv__ = lambda self, other: temp_path / other with ( - caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"), pytest.warns(DeprecationWarning, match=r"\.env.*removed in 1\.3\.0.*\.env\.local"), ): - _load_environment_files(env_files=None) + load_environment_files(env_files=None) output = capsys.readouterr().out assert f"WARNING: Auto-discovered {env_file} is deprecated" in output assert "Use env_akv_ref or ~/.pyrit/.env.local instead" in output assert f"Auto-discovered {env_file} is deprecated" in caplog.text - @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") def test_explicit_env_file_does_not_emit_legacy_deprecation(self, mock_config_path): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -161,11 +161,11 @@ def test_explicit_env_file_does_not_emit_legacy_deprecation(self, mock_config_pa with warnings.catch_warnings(): warnings.simplefilter("error", DeprecationWarning) - loaded = _load_environment_files(env_files=[explicit_env], silent=True) + loaded = load_environment_files(env_files=[explicit_env], silent=True) assert loaded is True - @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") def test_akv_legacy_env_warning_respects_silent(self, mock_config_path, caplog, capsys): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -173,7 +173,7 @@ def test_akv_legacy_env_warning_respects_silent(self, mock_config_path, caplog, mock_config_path.__truediv__ = lambda self, other: temp_path / other with ( - caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"), pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), ): _warn_about_akv_environment_files(env_files=None, silent=True) @@ -181,7 +181,7 @@ def test_akv_legacy_env_warning_respects_silent(self, mock_config_path, caplog, assert capsys.readouterr().out == "" assert "will be ignored because env_akv_ref is configured" in caplog.text - @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") async def test_akv_ignores_auto_discovered_env_and_loads_env_local(self, mock_config_path): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -191,14 +191,14 @@ async def test_akv_ignores_auto_discovered_env_and_loads_env_local(self, mock_co with ( mock.patch( - "pyrit.setup.akv_initialization._load_env_from_akv_async", + "pyrit.setup.environment_loading._load_env_from_akv_async", new_callable=mock.AsyncMock, return_value="VALUE=akv\n", ), - mock.patch("pyrit.setup.akv_initialization._load_environment_files") as mock_load_files, + mock.patch("pyrit.setup.environment_loading.load_environment_files") as mock_load_files, pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), ): - await _load_environment_async( + await load_environment_async( env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], env_files=None, env_akv_strict=True, @@ -209,7 +209,7 @@ async def test_akv_ignores_auto_discovered_env_and_loads_env_local(self, mock_co assert mock_load_files.call_args.kwargs["env_files"] is None assert mock_load_files.call_args.kwargs["include_default_base"] is False - @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") async def test_akv_debug_mode_rejects_existing_env_before_fetch(self, mock_config_path): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -219,11 +219,11 @@ async def test_akv_debug_mode_rejects_existing_env_before_fetch(self, mock_confi with ( mock.patch( - "pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock + "pyrit.setup.environment_loading._load_env_from_akv_async", new_callable=mock.AsyncMock ) as mock_load_akv, pytest.raises(ValueError, match=r"already exists.*rename or remove"), ): - await _load_environment_async( + await load_environment_async( env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], env_files=None, env_akv_strict=True, @@ -247,7 +247,7 @@ async def test_loads_custom_env_files_in_order(self): env3.write_text("VAR=local") with mock.patch.dict(os.environ, {}, clear=True): - loaded = _load_environment_files(env_files=[env1, env2, env3]) + loaded = load_environment_files(env_files=[env1, env2, env3]) assert loaded is True assert os.environ["VAR"] == "local" @@ -264,7 +264,7 @@ async def test_explicit_files_only_override_when_named_env_local(self): local_file.write_text("PROCESS_VALUE=local\nFILE_VALUE=local") with mock.patch.dict(os.environ, {"PROCESS_VALUE": "process"}, clear=True): - loaded = _load_environment_files(env_files=[first_file, second_file, local_file], silent=True) + loaded = load_environment_files(env_files=[first_file, second_file, local_file], silent=True) assert loaded is True assert os.environ["PROCESS_VALUE"] == "local" @@ -272,7 +272,7 @@ async def test_explicit_files_only_override_when_named_env_local(self): assert os.environ["SECOND_ONLY"] == "second" with mock.patch.dict(os.environ, {"PROCESS_VALUE": "process"}, clear=True): - _load_environment_files(env_files=[first_file, second_file], silent=True) + load_environment_files(env_files=[first_file, second_file], silent=True) assert os.environ["PROCESS_VALUE"] == "process" assert os.environ["FILE_VALUE"] == "first" @@ -283,7 +283,7 @@ async def test_load_environment_files_interpolates_in_assignment_order(self): 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) + loaded = load_environment_files(env_files=[env_file], silent=True) assert loaded is True assert os.environ["A"] == "two" @@ -296,7 +296,7 @@ async def test_load_environment_files_honors_python_dotenv_disabled(self): 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) + loaded = load_environment_files(env_files=[env_file], silent=True) assert loaded is True assert "DISABLED_VALUE" not in os.environ @@ -340,7 +340,7 @@ async def test_load_environment_async_write_env_writes_resolved_native_bootstrap temp_path = pathlib.Path(temp_dir) (temp_path / ".env.local").write_text("API_KEY=local-key\nLOCAL_ONLY=local", encoding="utf-8") with ( - mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), mock.patch.dict( os.environ, { @@ -354,7 +354,7 @@ async def test_load_environment_async_write_env_writes_resolved_native_bootstrap mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), ): - await _load_environment_async( + await load_environment_async( env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], env_files=None, env_akv_strict=True, @@ -396,17 +396,17 @@ async def test_load_environment_async_write_env_filters_generated_explicit_file( 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.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), mock.patch( - "pyrit.setup.akv_initialization._load_env_from_akv_async", + "pyrit.setup.environment_loading._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 + "pyrit.setup.environment_loading.load_environment_files", return_value=True ) as mock_load_environment_files, ): - await _load_environment_async( + await load_environment_async( env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], env_files=[generated_env, local_env], env_akv_strict=True, @@ -426,14 +426,14 @@ async def test_load_environment_async_write_env_preserves_first_bootstrap_value( 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.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), mock.patch( - "pyrit.setup.akv_initialization._load_env_from_akv_async", + "pyrit.setup.environment_loading._load_env_from_akv_async", new_callable=mock.AsyncMock, side_effect=lambda **kwargs: documents.pop(0), ), ): - await _load_environment_async( + await load_environment_async( env_akv_ref=[ "https://vault.vault.azure.net/secrets/first", "https://vault.vault.azure.net/secrets/second", @@ -462,22 +462,22 @@ def test_write_akv_env_file_secures_descriptor_before_writing(self): stream = mock.MagicMock() stream.write.side_effect = lambda content: events.append(f"write:{content}") with ( - mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), mock.patch( - "pyrit.setup.akv_initialization.tempfile.mkstemp", + "pyrit.setup.environment_loading.tempfile.mkstemp", side_effect=lambda **kwargs: events.append("create") or (7, str(temporary_file)), ), mock.patch( - "pyrit.setup.akv_initialization.os.fchmod", + "pyrit.setup.environment_loading.os.fchmod", side_effect=lambda *args: events.append("fchmod"), create=True, ), mock.patch( - "pyrit.setup.akv_initialization.os.fdopen", + "pyrit.setup.environment_loading.os.fdopen", side_effect=lambda *args, **kwargs: events.append("fdopen") or stream, ), mock.patch( - "pyrit.setup.akv_initialization.os.link", + "pyrit.setup.environment_loading.os.link", side_effect=lambda *args: events.append("link"), ), ): @@ -492,7 +492,7 @@ def test_write_akv_env_file_rejects_existing_file(self): env_file.write_text("ORIGINAL=value\n", encoding="utf-8") with ( - mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), pytest.raises(ValueError, match="already exists.*rename or remove"), ): _write_akv_env_file(documents=["NEW=value\n"], silent=True) @@ -522,8 +522,8 @@ def write(self, content): return self._stream.write(content) with ( - mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - mock.patch("pyrit.setup.akv_initialization.os.fdopen", side_effect=CompetingFileStream), + mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.environment_loading.os.fdopen", side_effect=CompetingFileStream), pytest.raises(ValueError, match="already exists.*rename or remove"), ): _write_akv_env_file(documents=["NEW=value\n"], silent=True) @@ -536,7 +536,7 @@ def test_write_akv_env_file_uses_owner_only_permissions(self): with tempfile.TemporaryDirectory() as temp_dir: configuration_directory = pathlib.Path(temp_dir) / ".pyrit" with mock.patch( - "pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", configuration_directory + "pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", configuration_directory ): env_file = _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) @@ -555,7 +555,7 @@ def test_write_akv_env_file_rejects_symbolic_link(self): pytest.skip("Symbolic links are unavailable on this platform.") with ( - mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), pytest.raises(ValueError, match="symbolic link"), ): _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) @@ -570,14 +570,14 @@ async def test_direct_local_file_loader_keeps_pyrit_references_literal(self): ) with mock.patch.dict(os.environ, {}, clear=True): - loaded = _load_environment_files(env_files=[env_file], silent=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") + @mock.patch("pyrit.setup.environment_loading.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) @@ -593,7 +593,7 @@ def test_default_local_file_can_interpolate_base_file_but_not_reverse(self, mock mock.patch.dict(os.environ, {}, clear=True), pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), ): - loaded = _load_environment_files(env_files=None, silent=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" @@ -637,7 +637,7 @@ async def test_load_environment_async_resolves_local_akv_reference(self, file_na mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, ): - await _load_environment_async( + await load_environment_async( env_akv_ref=None, env_files=[env_file], env_akv_strict=True, @@ -663,7 +663,7 @@ async def test_load_environment_async_does_not_fetch_local_reference_that_loses_ mock.patch.dict(os.environ, {"API_KEY": "process-key"}, clear=True), mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, ): - await _load_environment_async( + await load_environment_async( env_akv_ref=None, env_files=[env_file], env_akv_strict=True, @@ -685,7 +685,7 @@ async def test_load_environment_async_strict_rejects_malformed_local_akv_referen mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, pytest.raises(KeyVaultInitializationException, match="must use a full secret URL"), ): - await _load_environment_async( + await load_environment_async( env_akv_ref=None, env_files=[env_file], env_akv_strict=True, @@ -699,13 +699,14 @@ async def test_load_environment_async_non_strict_skips_malformed_local_akv_refer with tempfile.TemporaryDirectory() as temp_dir: env_local_file = pathlib.Path(temp_dir) / ".env.local" env_local_file.write_text("API_KEY=kv:api-key") + process_value = "kv:https://process-vault.vault.azure.net/secrets/do-not-resolve" with ( - mock.patch.dict(os.environ, {"API_KEY": "process-key"}, clear=True), + mock.patch.dict(os.environ, {"API_KEY": process_value}, clear=True), mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, - caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"), ): - await _load_environment_async( + await load_environment_async( env_akv_ref=None, env_files=[env_local_file], env_akv_strict=False, @@ -713,7 +714,7 @@ async def test_load_environment_async_non_strict_skips_malformed_local_akv_refer silent=False, ) - assert os.environ["API_KEY"] == "process-key" + assert os.environ["API_KEY"] == process_value mock_credential_cls.assert_not_called() assert ( @@ -722,6 +723,157 @@ async def test_load_environment_async_non_strict_skips_malformed_local_akv_refer ) assert "API_KEY" in caplog.text + @pytest.mark.parametrize("malformed_override_count", [1, 2]) + async def test_non_strict_falls_back_to_valid_reference_after_malformed_overrides( + self, malformed_override_count, caplog + ): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="resolved-key")) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + ordinary_file = temp_path / "ordinary.env" + ordinary_file.write_text("API_KEY=kv:https://local-vault.vault.azure.net/secrets/api-key") + env_files = [ordinary_file] + for index in range(malformed_override_count): + local_file = temp_path / str(index) / ".env.local" + local_file.parent.mkdir() + local_file.write_text(f"API_KEY=kv:short-{index}") + env_files.append(local_file) + + 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.environment_loading"), + ): + await load_environment_async( + env_akv_ref=None, + env_files=env_files, + env_akv_strict=False, + env_akv_write_env=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "resolved-key" + + client.get_secret.assert_awaited_once_with("api-key", version=None) + warnings = [record for record in caplog.records if "Invalid AKV reference" in record.message] + assert len(warnings) == malformed_override_count + + @pytest.mark.parametrize( + ("fallback_value", "expected_value"), + [("plain-value", "plain-value"), (None, None)], + ) + async def test_non_strict_malformed_override_uses_literal_or_no_fallback(self, fallback_value, expected_value): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_files: list[pathlib.Path] = [] + if fallback_value is not None: + ordinary_file = temp_path / "ordinary.env" + ordinary_file.write_text(f"API_KEY={fallback_value}") + env_files.append(ordinary_file) + local_file = temp_path / ".env.local" + local_file.write_text("API_KEY=kv:short") + env_files.append(local_file) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + ): + await load_environment_async( + env_akv_ref=None, + env_files=env_files, + env_akv_strict=False, + env_akv_write_env=False, + silent=True, + ) + + assert os.environ.get("API_KEY") == expected_value + + mock_credential_cls.assert_not_called() + + async def test_strict_malformed_override_does_not_resolve_valid_fallback(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + ordinary_file = temp_path / "ordinary.env" + ordinary_file.write_text("API_KEY=kv:https://local-vault.vault.azure.net/secrets/api-key") + local_file = temp_path / ".env.local" + local_file.write_text("API_KEY=kv:short") + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + pytest.raises(KeyVaultInitializationException, match="must use a full secret URL"), + ): + await load_environment_async( + env_akv_ref=None, + env_files=[ordinary_file, local_file], + env_akv_strict=True, + env_akv_write_env=False, + silent=True, + ) + + mock_credential_cls.assert_not_called() + + async def test_highest_valid_local_reference_wins(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="local-key")) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + ordinary_file = temp_path / "ordinary.env" + ordinary_file.write_text("API_KEY=kv:https://local-vault.vault.azure.net/secrets/ordinary-key") + local_file = temp_path / ".env.local" + local_file.write_text("API_KEY=kv:https://local-vault.vault.azure.net/secrets/local-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), + ): + await load_environment_async( + env_akv_ref=None, + env_files=[ordinary_file, local_file], + env_akv_strict=True, + env_akv_write_env=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "local-key" + + client.get_secret.assert_awaited_once_with("local-key", version=None) + + async def test_non_strict_resolves_interpolated_fallback_candidate(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="resolved-key")) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + ordinary_file = temp_path / "ordinary.env" + ordinary_file.write_text( + "SECRET_URL=https://local-vault.vault.azure.net/secrets/api-key\nAPI_KEY=kv:${SECRET_URL}" + ) + local_file = temp_path / ".env.local" + local_file.write_text("API_KEY=kv:short") + + 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_environment_async( + env_akv_ref=None, + env_files=[ordinary_file, local_file], + env_akv_strict=False, + env_akv_write_env=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "resolved-key" + + client.get_secret.assert_awaited_once_with("api-key", version=None) + async def test_load_environment_async_non_strict_still_raises_for_missing_local_secret(self): credential, client = _create_mock_akv_clients() missing_error = ResourceNotFoundError(message="Secret was not found") @@ -739,7 +891,7 @@ async def test_load_environment_async_non_strict_still_raises_for_missing_local_ KeyVaultInitializationException, match="Failed to resolve Key Vault reference" ) as exc_info, ): - await _load_environment_async( + await load_environment_async( env_akv_ref=None, env_files=[env_file], env_akv_strict=False, @@ -754,7 +906,7 @@ async def test_raises_error_for_nonexistent_env_file(self): nonexistent = pathlib.Path("/nonexistent/path/.env") with pytest.raises(ValueError, match="Environment file not found"): - _load_environment_files(env_files=[nonexistent]) + 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): @@ -902,7 +1054,7 @@ def test_parse_akv_secret_url_invalid_raises(self, 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, + mock.patch("pyrit.setup.environment_loading._create_akv_secret_client") as mock_create_client, pytest.raises(KeyVaultInitializationException, match="attacker.example"), ): await _load_env_from_akv_async( @@ -937,7 +1089,7 @@ async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_secret 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, + mock.patch("pyrit.setup.environment_loading._print_msg") as mock_print_msg, ): await _load_env_from_akv_async(secret_url=secret_url, silent=True) @@ -1226,7 +1378,7 @@ async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entrie 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"), + caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"), ): await _load_env_from_akv_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", @@ -1253,7 +1405,7 @@ async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_refere 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"), + caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"), ): resolved_document = await _load_env_from_akv_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", @@ -1281,7 +1433,7 @@ async def test_load_env_from_akv_async_non_strict_silent_logs_warning(self, capl 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"), + caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"), ): await _load_env_from_akv_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 5299a38441..884383ee9f 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -119,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.akv_initialization._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.environment_loading.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) @@ -128,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.akv_initialization._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.environment_loading.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: @@ -158,15 +158,15 @@ async def initialize_async(self) -> None: finally: os.unlink(script_path) - @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.environment_loading.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.akv_initialization._load_environment_files", return_value=False) - @mock.patch("pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) + @mock.patch("pyrit.setup.environment_loading.load_environment_files", return_value=False) + @mock.patch("pyrit.setup.environment_loading._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 = [ @@ -176,7 +176,7 @@ async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, m mock_load_akv.return_value = None - with mock.patch("pyrit.setup.akv_initialization._warn_about_akv_environment_files") as mock_warn: + with mock.patch("pyrit.setup.environment_loading._warn_about_akv_environment_files") as mock_warn: await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) assert mock_load_akv.await_args_list == [ @@ -188,8 +188,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.akv_initialization._load_environment_files", return_value=False) - @mock.patch("pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) + @mock.patch("pyrit.setup.environment_loading.load_environment_files", return_value=False) + @mock.patch("pyrit.setup.environment_loading._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 ): @@ -213,7 +213,7 @@ async def test_initialize_rejects_invalid_env_akv_ref(self, env_akv_ref): @pytest.mark.parametrize("invalid_value", ["false", "true", 0, 1, None, [], {}]) async def test_initialize_rejects_non_boolean_akv_options_before_loading(self, option_name, invalid_value): with mock.patch( - "pyrit.setup.initialization._load_environment_async", new_callable=mock.AsyncMock + "pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock ) as mock_load_environment: with pytest.raises(TypeError, match=rf"{option_name} must be a bool"): if option_name == "env_akv_strict": @@ -238,7 +238,7 @@ async def test_initialize_rejects_non_boolean_akv_options_before_loading(self, o async def test_initialize_forwards_boolean_akv_options(self, env_akv_strict, env_akv_write_env): with ( mock.patch( - "pyrit.setup.initialization._load_environment_async", new_callable=mock.AsyncMock + "pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock ) as mock_load_environment, mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance"), ): @@ -261,9 +261,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.akv_initialization._warn_about_akv_environment_files"), + mock.patch("pyrit.setup.environment_loading._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.akv_initialization._load_env_from_akv_async", + "pyrit.setup.environment_loading._load_env_from_akv_async", new_callable=mock.AsyncMock, side_effect=lambda **_: os.environ.update({"FROM_AKV": "resolved"}), ), @@ -289,9 +289,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.akv_initialization._warn_about_akv_environment_files"), + mock.patch("pyrit.setup.environment_loading._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.akv_initialization._load_env_from_akv_async", + "pyrit.setup.environment_loading._load_env_from_akv_async", new_callable=mock.AsyncMock, side_effect=lambda **_: os.environ.update({"BASE": "akv", "ONLY_AKV": "shared"}), ), @@ -319,10 +319,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.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - mock.patch("pyrit.setup.akv_initialization._warn_about_akv_environment_files"), + mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.environment_loading._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.akv_initialization._load_env_from_akv_async", + "pyrit.setup.environment_loading._load_env_from_akv_async", new_callable=mock.AsyncMock, side_effect=lambda **_: os.environ.update({"VALUE": "akv"}), ), @@ -363,14 +363,14 @@ 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.akv_initialization._warn_about_akv_environment_files"), + mock.patch("pyrit.setup.environment_loading._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.akv_initialization._load_env_from_akv_async", + "pyrit.setup.environment_loading._load_env_from_akv_async", new_callable=mock.AsyncMock, side_effect=lambda **_: os.environ.update(bootstrap_environment), ), mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("pyrit.setup.akv_initialization._create_akv_secret_client", return_value=client), + mock.patch("pyrit.setup.environment_loading._create_akv_secret_client", return_value=client), ): await initialize_pyrit_async( memory_db_type=IN_MEMORY, @@ -415,7 +415,7 @@ def setup_method(self) -> None: """Clear default values before each test.""" reset_default_values() - @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=True) + @mock.patch("pyrit.setup.environment_loading.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) @@ -423,7 +423,7 @@ async def test_initialize_silent_produces_no_output(self, mock_load_env, capsys) captured = capsys.readouterr() assert captured.out == "" - @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=True) + @mock.patch("pyrit.setup.environment_loading.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) diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index 06c7d4c4d1..d0793f90b2 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["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" + 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" init = TargetInitializer() await init.initialize_async() @@ -116,6 +116,54 @@ async def test_registers_multiple_targets(self): assert "platform_openai_chat" in registry.instances assert "openai_image_platform" in registry.instances + @pytest.mark.parametrize( + ("registry_name", "endpoint_var", "key_var", "model_var", "endpoint"), + [ + ( + "openai_image_azure", + "OPENAI_IMAGE_ENDPOINT1", + "OPENAI_IMAGE_API_KEY1", + "OPENAI_IMAGE_MODEL1", + "https://image.openai.azure.com/openai/v1", + ), + ( + "openai_image_platform", + "OPENAI_IMAGE_ENDPOINT2", + "OPENAI_IMAGE_API_KEY2", + "OPENAI_IMAGE_MODEL2", + "https://api.openai.com/v1", + ), + ( + "openai_tts_azure", + "OPENAI_TTS_ENDPOINT1", + "OPENAI_TTS_KEY1", + "OPENAI_TTS_MODEL1", + "https://tts.openai.azure.com/openai/v1", + ), + ( + "openai_tts_platform", + "OPENAI_TTS_ENDPOINT2", + "OPENAI_TTS_KEY2", + "OPENAI_TTS_MODEL2", + "https://api.openai.com/v1", + ), + ], + ) + async def test_media_targets_use_main_environment_contract( + self, registry_name, endpoint_var, key_var, model_var, endpoint + ): + with patch.dict( + os.environ, + {endpoint_var: endpoint, key_var: "test-key", model_var: "test-model"}, + clear=True, + ): + await TargetInitializer().initialize_async() + + target = TargetRegistry.get_registry_singleton().instances.get(registry_name) + assert target is not None + assert target._endpoint == endpoint + assert target._model_name == "test-model" + async def test_registers_azure_content_safety_without_model(self): """Test that PromptShieldTarget is registered without model_name (it doesn't use one).""" os.environ["AZURE_CONTENT_SAFETY_API_ENDPOINT"] = "https://test.cognitiveservices.azure.com" From 485b1bf804ecd0675f9fa501deb60fed8197c0fc Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 21 Aug 2026 10:45:02 -0400 Subject: [PATCH 24/28] FIX: Addressed PR comments, removed writing AKV to disk, removed deprecation schedule, removed changes to .env_example --- .pyrit_conf_example | 4 +- build_scripts/export_akv_environment.py | 366 +++++++++++ doc/getting_started/pyrit_conf.md | 24 +- pyrit/setup/configuration_loader.py | 18 +- pyrit/setup/environment_loading.py | 530 +++++---------- pyrit/setup/initialization.py | 15 +- .../test_export_akv_environment.py | 303 +++++++++ tests/unit/setup/test_configuration_loader.py | 36 +- tests/unit/setup/test_environment_loading.py | 613 +++++------------- tests/unit/setup/test_initialization.py | 101 +-- 10 files changed, 1081 insertions(+), 929 deletions(-) create mode 100644 build_scripts/export_akv_environment.py create mode 100644 tests/unit/build_scripts/test_export_akv_environment.py diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 7bcecf5ec0..033b90c274 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -86,9 +86,9 @@ operation: op_trash_panda # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env # env_akv_strict: true -# env_akv_write_env: false # Debug only: write a fully resolved, sensitive ~/.pyrit/.env. -# Auto-discovered ~/.pyrit/.env is legacy and will be rejected in PyRIT 1.3.0. +# Auto-discovered ~/.pyrit/.env remains supported but emits a security warning. +# Prefer env_akv_ref for shared or deployed secrets. # Use ~/.pyrit/.env.local for quick local plaintext patches or when Azure is unavailable. # Process values remain authoritative; AKV and ordinary env_files fill gaps in load order. # Only a file named .env.local overrides existing values. diff --git a/build_scripts/export_akv_environment.py b/build_scripts/export_akv_environment.py new file mode 100644 index 0000000000..8ea6c0cec1 --- /dev/null +++ b/build_scripts/export_akv_environment.py @@ -0,0 +1,366 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Export resolved Azure Key Vault bootstrap documents to ``~/.pyrit/.env_akv``.""" + +import argparse +import contextlib +import logging +import os +import pathlib +import tempfile +import urllib.parse +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from io import StringIO +from typing import TYPE_CHECKING, Any + +import dotenv +from dotenv.parser import parse_stream +from dotenv.variables import parse_variables + +if TYPE_CHECKING: + from azure.core.credentials import TokenCredential + from azure.keyvault.secrets import SecretClient + +logger = logging.getLogger(__name__) + +DEFAULT_OUTPUT_FILE = pathlib.Path.home() / ".pyrit" / ".env_akv" +_REFERENCE_PREFIXES = frozenset({"akv", "kv", "azure_key_vault", "env_akv_ref"}) +_VAULT_DNS_SUFFIXES = frozenset({"vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"}) + + +@dataclass(frozen=True) +class _Document: + content: str + vault_url: str + + +@dataclass(frozen=True) +class _Candidate: + document_index: int + binding_index: int + name: str + value: str + vault_url: str + + +def _parse_secret_url(url: str) -> tuple[str, str, str | None]: + """Return vault URL, secret name, and optional version from a full AKV URL.""" + error_message = f"Invalid Azure Key Vault secret URL: {url}" + try: + parsed = urllib.parse.urlsplit(url) + port = parsed.port + except (TypeError, ValueError) as error: + raise ValueError(error_message) from error + hostname = parsed.hostname + vault_name, separator, suffix = hostname.partition(".") if hostname else ("", "", "") + valid_vault = 1 <= len(vault_name) <= 63 and all( + char.isascii() and (char.isalnum() or char == "-") for char in vault_name + ) + parts = parsed.path.split("/") + valid_path = len(parts) in {3, 4} and parts[:2] == ["", "secrets"] and all(parts[2:]) + if ( + parsed.scheme.casefold() != "https" + or parsed.username is not None + or parsed.password is not None + or port is not None + or separator != "." + or suffix not in _VAULT_DNS_SUFFIXES + or not valid_vault + or not valid_path + or parsed.query + or parsed.fragment + ): + raise ValueError(error_message) + secret_name = parts[2] + secret_version = parts[3] if len(parts) == 4 else None + identifiers = [secret_name] + ([secret_version] if secret_version else []) + if any(not (1 <= len(item) <= 127 and all(char.isalnum() or char == "-" for char in item)) for item in identifiers): + raise ValueError(error_message) + return f"https://{hostname}", secret_name, secret_version + + +def _create_client(*, vault_url: str, credential: "TokenCredential") -> "SecretClient": + """Create a Key Vault client with explicit retry settings.""" + from azure.core.pipeline.policies import RetryPolicy + from azure.keyvault.secrets import SecretClient + + return SecretClient( + vault_url=vault_url, + credential=credential, + retry_policy=RetryPolicy( + retry_total=3, + retry_connect=3, + retry_read=3, + retry_status=3, + retry_backoff_factor=0.8, + ), + ) + + +def _client_for(*, vault_url: str, credential: "TokenCredential", clients: dict[str, "SecretClient"]) -> "SecretClient": + client = clients.get(vault_url) + if client is None: + client = _create_client(vault_url=vault_url, credential=credential) + clients[vault_url] = client + return client + + +def _validate_document(*, document: str, strict: bool, silent: bool) -> str: + bindings = list(parse_stream(StringIO(document))) + malformed = [str(binding.original.line) for binding in bindings if binding.error] + valueless = [binding.key for binding in bindings if binding.key is not None and binding.value is None] + issues: list[str] = [] + if malformed: + issues.append("malformed entries at lines: " + ", ".join(malformed)) + if valueless: + issues.append("variables without values: " + ", ".join(valueless)) + 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) + ) + + +def _fetch_documents( + *, + secret_urls: Sequence[str], + credential: "TokenCredential", + clients: dict[str, "SecretClient"], + strict: bool, + silent: bool, +) -> list[_Document]: + documents: list[_Document] = [] + for url in secret_urls: + vault_url, name, version = _parse_secret_url(url) + secret = _client_for(vault_url=vault_url, credential=credential, clients=clients).get_secret( + name, version=version + ) + if not secret.value: + raise ValueError(f"AKV environment secret has no value: {url}") + content = _validate_document(document=secret.value, strict=strict, silent=silent) + if not dotenv.dotenv_values(stream=StringIO(content), interpolate=False): + raise ValueError(f"AKV environment secret contains no assignments: {url}") + documents.append(_Document(content=content, vault_url=vault_url)) + return documents + + +def _resolve_interpolation(*, value: str, environment: Mapping[str, str | None]) -> str: + return "".join(atom.resolve(environment) for atom in parse_variables(value)) + + +def _build_candidates(documents: Sequence[_Document]) -> tuple[list[list[Any]], dict[str, list[_Candidate]]]: + effective: dict[str, str | None] = {} + chains: dict[str, list[_Candidate]] = {} + all_bindings: list[list[Any]] = [] + for document_index, document in enumerate(documents): + bindings = list(parse_stream(StringIO(document.content))) + all_bindings.append(bindings) + current: dict[str, str | None] = {} + final_indexes: dict[str, int] = {} + for binding_index, binding in enumerate(bindings): + if binding.key is None or binding.value is None: + continue + environment = dict(current) + environment.update(effective) + current[binding.key] = _resolve_interpolation(value=binding.value, environment=environment) + final_indexes[binding.key] = binding_index + for name, value in current.items(): + if value is None: + continue + chains.setdefault(name, []).append( + _Candidate(document_index, final_indexes[name], name, value, document.vault_url) + ) + effective.setdefault(name, value) + return all_bindings, chains + + +def _reference_target(value: str) -> str | None: + prefix, separator, target = value.partition(":") + return target.strip() if separator and prefix in _REFERENCE_PREFIXES else None + + +def _serialize(value: str) -> str: + escaped = value.replace("\\", "\\\\").replace("'", "\\'").replace("${", "${:-$}{") + return f"'{escaped}'" + + +def _render( + *, + documents: Sequence[_Document], + credential: "TokenCredential", + clients: dict[str, "SecretClient"], + strict: bool, + silent: bool, +) -> str: + all_bindings, chains = _build_candidates(documents) + selected: dict[str, _Candidate] = {} + resolved: dict[tuple[int, int], str] = {} + for name, candidates in chains.items(): + for candidate in candidates: + target = _reference_target(candidate.value) + if target is None: + selected[name] = candidate + break + try: + vault_url, secret_name, version = _parse_secret_url(target) + if vault_url.casefold() != candidate.vault_url.casefold(): + raise ValueError(f"Cross-vault AKV reference for '{name}' is not supported") + except ValueError as error: + if strict: + raise + message = f"Invalid AKV reference for '{name}' will be skipped: {error}" + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + continue + secret = _client_for(vault_url=vault_url, credential=credential, clients=clients).get_secret( + secret_name, version=version + ) + if secret.value is None: + raise ValueError(f"AKV secret '{secret_name}' referenced by '{name}' has no value") + selected[name] = candidate + resolved[(candidate.document_index, candidate.binding_index)] = secret.value + break + + output: list[str] = [] + for document_index, bindings in enumerate(all_bindings): + for binding_index, binding in enumerate(bindings): + name = binding.key + if name is None: + output.append(binding.original.string) + continue + winner = selected.get(name) + if winner is None or winner.document_index != document_index: + continue + value = resolved.get((document_index, binding_index)) + if value is None: + output.append(binding.original.string) + continue + original = binding.original.string + export = "export " if original.lstrip().startswith("export ") else "" + newline = "\r\n" if original.endswith("\r\n") else "\n" if original.endswith("\n") else "" + output.append(f"{export}{name}={_serialize(value)}{newline}") + return "".join(output).rstrip("\r\n") + "\n" + + +def _ensure_output_available(output_file: pathlib.Path) -> pathlib.Path: + output_file = output_file.expanduser() + if output_file.is_symlink(): + raise ValueError(f"Output path is a symbolic link: {output_file}") + if output_file.exists(): + raise ValueError(f"Output already exists: {output_file}. Rename or remove it before exporting") + return output_file + + +def _write_output(*, output_file: pathlib.Path, document: str) -> pathlib.Path: + output_file = _ensure_output_available(output_file) + output_file.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + descriptor: int | None = None + temporary: pathlib.Path | None = None + try: + descriptor, name = tempfile.mkstemp(prefix=f"{output_file.name}.", suffix=".tmp", dir=output_file.parent) + temporary = pathlib.Path(name) + file_chmod = getattr(os, "fchmod", None) + if file_chmod is not None: + file_chmod(descriptor, 0o600) + else: + os.chmod(temporary, 0o600) + stream = os.fdopen(descriptor, "w", encoding="utf-8", newline="") + descriptor = None + with stream: + stream.write(document) + try: + os.link(temporary, output_file) + except FileExistsError as error: + raise ValueError(f"Output already exists: {output_file}. Rename or remove it before exporting") from error + finally: + if descriptor is not None: + os.close(descriptor) + if temporary is not None: + with contextlib.suppress(FileNotFoundError): + temporary.unlink() + return output_file + + +def export_akv_environment( + *, + secret_urls: Sequence[str], + output_file: pathlib.Path = DEFAULT_OUTPUT_FILE, + strict: bool = True, + silent: bool = False, + credential: "TokenCredential | None" = None, +) -> pathlib.Path: + """Fetch, resolve, and securely export AKV-only configuration. + + A caller-provided credential remains caller-owned and is not closed. + """ + if not secret_urls: + raise ValueError("At least one secret URL is required") + output_file = _ensure_output_available(output_file) + from azure.identity import DefaultAzureCredential + + owned_credential = None + if credential is None: + owned_credential = DefaultAzureCredential() + active_credential = owned_credential + else: + active_credential = credential + clients: dict[str, SecretClient] = {} + try: + documents = _fetch_documents( + secret_urls=secret_urls, + credential=active_credential, + clients=clients, + strict=strict, + silent=silent, + ) + document = _render( + documents=documents, + credential=active_credential, + clients=clients, + strict=strict, + silent=silent, + ) + output = _write_output(output_file=output_file, document=document) + finally: + for client in clients.values(): + client.close() + if owned_credential is not None: + owned_credential.close() + if not silent: + print(f"Exported resolved AKV environment to {output}") + return output + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--secret-url", dest="secret_urls", action="append", required=True) + parser.add_argument("--output", type=pathlib.Path, default=DEFAULT_OUTPUT_FILE) + parser.add_argument("--non-strict", action="store_true") + parser.add_argument("--silent", action="store_true") + args = parser.parse_args() + try: + export_akv_environment( + secret_urls=args.secret_urls, + output_file=args.output, + strict=not args.non_strict, + silent=args.silent, + ) + except Exception as error: + parser.exit(1, f"Export failed: {error}\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index cdb4695780..0f98f7a6b5 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -26,7 +26,7 @@ PyRIT looks for this file automatically on startup (via the CLI, shell, or `Conf ## Environment Configuration ```{important} -Azure Key Vault is PyRIT's canonical environment source for shared, CI/CD, and deployed configuration. Auto-discovered `~/.pyrit/.env` is supported only as a legacy source and will be rejected in PyRIT 1.3.0. Use `~/.pyrit/.env.local` for deliberate plaintext local iteration or when Azure is unavailable. +Azure Key Vault is PyRIT's canonical environment source for shared, CI/CD, and deployed configuration. Auto-discovered `~/.pyrit/.env` remains supported, but PyRIT warns because plaintext files are less secure. Use `~/.pyrit/.env.local` for deliberate plaintext local iteration or when Azure is unavailable. ``` See [Populating Secrets](./populating_secrets.md) for provider-specific variable examples. @@ -36,10 +36,10 @@ See [Populating Secrets](./populating_secrets.md) for provider-specific variable PyRIT loads environment sources in this order: 1. Existing process environment variables. -2. Key Vault bootstrap documents, legacy auto-discovered `.env`, or explicit `env_files`. These sources fill only missing values. +2. Key Vault bootstrap documents, auto-discovered `.env`, or explicit `env_files`. These sources fill only missing values. 3. Files named `.env.local`. These are the only dotenv sources that override existing values. -When `env_akv_ref` is configured, PyRIT ignores an auto-discovered `~/.pyrit/.env`, emits its deprecation warning, and still loads `~/.pyrit/.env.local`. Explicit `env_files` are never blocked or deprecated based on their filename or location. +When `env_akv_ref` is configured, PyRIT ignores an auto-discovered `~/.pyrit/.env`, emits a security warning, and still loads `~/.pyrit/.env.local`. Explicit `env_files` are never blocked based on their filename or location. ### Using .env.local for Overrides @@ -155,7 +155,7 @@ Optional local dotenv paths. Key Vault remains the canonical shared source; expl | Value | Behavior | | ----------------- | -------------------------------------------------------- | -| Omitted or `null` | Auto-discover legacy `~/.pyrit/.env` and supported `~/.pyrit/.env.local` | +| Omitted or `null` | Auto-discover `~/.pyrit/.env` and `~/.pyrit/.env.local` | | `[]` (empty list) | Load **no** environment files | | List of paths | Load **only** the specified files (defaults are skipped) | @@ -205,7 +205,7 @@ References must occupy the entire value. `kv:` is the canonical Key Vault prefix 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 winning `kv:` occurrence performs a Key Vault read during initialization. References that lose to an existing process or earlier source are not fetched. Debug output is the exception: it resolves bootstrap references for the written file without changing runtime precedence. +PyRIT does not cache referenced secrets. Each winning `kv:` occurrence performs a Key Vault read during initialization. References that lose to an existing process or earlier source are not fetched. The standalone exporter described below resolves bootstrap references independently and does not change runtime precedence. ```dotenv LATEST_KEY_URI="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key" @@ -228,17 +228,18 @@ Non-strict mode does not suppress operational failures. Missing secrets, authent 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` +### Exporting AKV Configuration for Debugging -Defaults to `false`. Set it to `true` only while debugging to write a fully resolved bootstrap document to `~/.pyrit/.env`: +Environment initialization never writes secrets to disk. To inspect an AKV-only configuration explicitly from a source checkout, run the standalone helper: -```yaml -env_akv_write_env: true +```powershell +python -m build_scripts.export_akv_environment ` + --secret-url https://my-vault.vault.azure.net/secrets/my-pyrit-env ``` -The written file contains only bootstrap assignments, comments, and fully resolved child-secret values. It excludes unrelated process and `.env.local` values, safely round-trips terminal secret text, and is always named `.env`, never `.env.new`. +Repeat `--secret-url` to preserve multiple bootstrap documents in load order. The helper writes `~/.pyrit/.env_akv` by default. This file is not auto-loaded by PyRIT and excludes process, `.env`, explicit file, and `.env.local` values. -PyRIT refuses debug mode when `~/.pyrit/.env` already exists; rename or remove the existing file first. The file is created with owner-only permissions where supported and replaced atomically, but it contains plaintext secrets. Remove it when debugging is complete. `.env.local` still loads afterward and can override runtime values without changing the generated file. +The helper resolves child-secret references and writes plaintext secrets with owner-only permissions where supported. It refuses to overwrite an existing path; remove the file when debugging is complete. Use `--output` to select a different path and `--non-strict` to skip malformed entries or references with warnings. ### `silent` @@ -365,7 +366,6 @@ initializers: # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env # env_akv_strict: true -# env_akv_write_env: false # Debug only: writes fully resolved plaintext secrets # Optional plaintext local patch or non-Azure workflow # env_files: diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 178142e0dc..d2d494de3e 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -21,7 +21,7 @@ from pyrit.common.utils import verify_and_resolve_path from pyrit.common.yaml_loadable import YamlLoadable from pyrit.models import class_name_to_snake_case -from pyrit.setup.environment_loading import validate_akv_boolean_options +from pyrit.setup.environment_loading import validate_env_akv_strict from pyrit.setup.initialization import ( AZURE_SQL, IN_MEMORY, @@ -96,13 +96,11 @@ class ConfigurationLoader(YamlLoadable): initialization_scripts: List of paths to custom initialization scripts. None means "use defaults", [] means "load nothing". env_files: List of environment file paths to load. - None means auto-discover legacy ``.env`` and supported ``.env.local``; + None means auto-discover supported ``.env`` and ``.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. - env_akv_write_env: Whether to save fully resolved bootstrap documents with - plaintext child-secret values to ``~/.pyrit/.env`` for debugging. silent: Whether to suppress initialization messages. operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. @@ -143,7 +141,6 @@ 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 @@ -154,10 +151,7 @@ class ConfigurationLoader(YamlLoadable): def __post_init__(self) -> None: """Validate and normalize the configuration after loading.""" - validate_akv_boolean_options( - env_akv_strict=self.env_akv_strict, - env_akv_write_env=self.env_akv_write_env, - ) + validate_env_akv_strict(env_akv_strict=self.env_akv_strict) self._normalize_memory_db_type() self._normalize_initializers() self._validate_env_akv_ref() @@ -430,7 +424,6 @@ 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. @@ -448,7 +441,6 @@ 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. @@ -516,9 +508,6 @@ 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 @@ -655,7 +644,6 @@ 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/environment_loading.py b/pyrit/setup/environment_loading.py index eba981b678..7bc74d22d8 100644 --- a/pyrit/setup/environment_loading.py +++ b/pyrit/setup/environment_loading.py @@ -8,7 +8,6 @@ import logging import os import pathlib -import tempfile import urllib.parse from collections.abc import Mapping, Sequence from dataclasses import dataclass @@ -16,9 +15,10 @@ from typing import TYPE_CHECKING import dotenv +from dotenv.main import DotEnv from dotenv.parser import parse_stream -from pyrit.common import path, print_deprecation_message +from pyrit.common import path from pyrit.exceptions import KeyVaultInitializationException if TYPE_CHECKING: @@ -27,35 +27,44 @@ logger = logging.getLogger(__name__) +__all__ = [ + "load_environment_async", + "load_environment_files", + "validate_env_akv_strict", +] + _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" -_LEGACY_ENV_REMOVED_IN = "1.3.0" @dataclass(frozen=True) class _EnvironmentValueCandidate: - """An ordered environment value and whether local AKV resolution applies.""" + """An ordered environment value and its Key Vault resolution policy.""" value: str resolve_akv_reference: bool + expected_vault_url: str | None = None + + +@dataclass(frozen=True) +class _AkvEnvironmentDocument: + """A validated Key Vault bootstrap document and its source vault.""" + + content: str + vault_url: str -def validate_akv_boolean_options(*, env_akv_strict: object, env_akv_write_env: object) -> None: +def validate_env_akv_strict(*, env_akv_strict: object) -> None: """ - Require real booleans for Key Vault behavior flags. + Require a real boolean for Key Vault strict-mode behavior. Raises: - TypeError: If either option is not a bool. + TypeError: If env_akv_strict is not a bool. """ - for option_name, option_value in ( - ("env_akv_strict", env_akv_strict), - ("env_akv_write_env", env_akv_write_env), - ): - if not isinstance(option_value, bool): - raise TypeError(f"{option_name} must be a bool, got {type(option_value).__name__}.") + if not isinstance(env_akv_strict, bool): + raise TypeError(f"env_akv_strict must be a bool, got {type(env_akv_strict).__name__}.") def load_environment_files( @@ -63,6 +72,26 @@ def load_environment_files( *, silent: bool = False, include_default_base: bool = True, +) -> bool: + """ + Load local environment files using PyRIT's standard precedence. + + Returns: + bool: Whether at least one environment file was selected. + """ + return _load_environment_files( + env_files=env_files, + silent=silent, + include_default_base=include_default_base, + assignment_candidates=None, + ) + + +def _load_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool, + include_default_base: bool, assignment_candidates: dict[str, list[_EnvironmentValueCandidate]] | None = None, ) -> bool: """ @@ -94,37 +123,82 @@ def load_environment_files( include_default_base=include_default_base, ) for env_file in selected_files: - override = env_file.name == ".env.local" - applicable_names: list[str] = [] - previous_values: dict[str, str | None] = {} - if assignment_candidates is not None: - assignment_values = dotenv.dotenv_values(dotenv_path=env_file, interpolate=False) - applicable_names = [ - variable_name - for variable_name, value in assignment_values.items() - if value is not None and (override or variable_name not in os.environ) - ] - previous_values = {variable_name: os.environ.get(variable_name) for variable_name in applicable_names} - loaded = dotenv.load_dotenv( + loaded = _load_dotenv_source( dotenv_path=env_file, - override=override, - interpolate=True, + override=env_file.name == ".env.local", + assignment_candidates=assignment_candidates, ) - if assignment_candidates is not None and loaded: - for variable_name in applicable_names: - candidates = assignment_candidates.setdefault(variable_name, []) - previous_value = previous_values[variable_name] - if not candidates and previous_value is not None: - candidates.append(_EnvironmentValueCandidate(value=previous_value, resolve_akv_reference=False)) - loaded_value = os.environ.get(variable_name) - if loaded_value is not None: - candidates.append(_EnvironmentValueCandidate(value=loaded_value, resolve_akv_reference=True)) if not silent: _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) return bool(selected_files) +def _load_dotenv_source( + *, + override: bool, + assignment_candidates: dict[str, list[_EnvironmentValueCandidate]] | None, + dotenv_path: pathlib.Path | None = None, + document: str | None = None, + expected_vault_url: str | None = None, +) -> bool: + """ + Load one dotenv source and record values that participate in precedence. + + Returns: + bool: Whether python-dotenv loaded at least one assignment. + + Raises: + ValueError: If both or neither source representations are provided. + """ + if (dotenv_path is None) == (document is None): + raise ValueError("Exactly one dotenv_path or document must be provided.") + + if dotenv_path is not None: + assignment_values = DotEnv( + dotenv_path=dotenv_path, + override=override, + interpolate=True, + ).dict() + else: + assignment_values = DotEnv( + dotenv_path=None, + stream=StringIO(document or ""), + override=override, + interpolate=True, + ).dict() + previous_values = {variable_name: os.environ.get(variable_name) for variable_name in assignment_values} + + if dotenv_path is not None: + loaded = dotenv.load_dotenv(dotenv_path=dotenv_path, override=override, interpolate=True) + else: + loaded = dotenv.load_dotenv(stream=StringIO(document or ""), override=override, interpolate=True) + if assignment_candidates is None or not loaded: + return loaded + + for variable_name, loaded_value in assignment_values.items(): + if loaded_value is None: + continue + candidates = assignment_candidates.setdefault(variable_name, []) + previous_value = previous_values[variable_name] + if not candidates and previous_value is not None: + candidates.append(_EnvironmentValueCandidate(value=previous_value, resolve_akv_reference=False)) + candidate = _EnvironmentValueCandidate( + value=loaded_value, + resolve_akv_reference=True, + expected_vault_url=expected_vault_url, + ) + if override: + candidates.append(candidate) + elif candidates and not candidates[-1].resolve_akv_reference: + continue + elif candidates: + candidates.insert(0, candidate) + else: + candidates.append(candidate) + return loaded + + def _select_environment_files( env_files: Sequence[pathlib.Path] | None, *, @@ -154,7 +228,7 @@ def _select_environment_files( local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" if include_default_base and base_file.exists(): - _warn_about_legacy_env(env_file=base_file, ignored_for_akv=False, silent=silent) + _warn_about_dotenv_file(env_file=base_file, ignored_for_akv=False, silent=silent) default_files.append(base_file) if local_file.exists(): default_files.append(local_file) @@ -191,31 +265,14 @@ 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 an auto-discovered legacy environment file coexists with AKV.""" - if env_files is not None: - return - - base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" - if base_file.exists(): - _warn_about_legacy_env(env_file=base_file, ignored_for_akv=True, silent=silent) - - -def _warn_about_legacy_env(*, env_file: pathlib.Path, ignored_for_akv: bool, silent: bool) -> None: - """Emit the standard and visible warnings for auto-discovered legacy ``.env`` loading.""" - print_deprecation_message( - old_item=f"Auto-discovered {env_file}", - new_item="env_akv_ref or ~/.pyrit/.env.local", - removed_in=_LEGACY_ENV_REMOVED_IN, - ) - behavior = "will be ignored because env_akv_ref is configured" if ignored_for_akv else "will still be loaded" +def _warn_about_dotenv_file(*, env_file: pathlib.Path, ignored_for_akv: bool, silent: bool) -> None: + """Warn that Azure Key Vault is safer than an auto-discovered plaintext ``.env`` file.""" + behavior = "will be ignored because env_akv_ref is configured" if ignored_for_akv else "will be loaded" message = ( - f"Auto-discovered {env_file} is deprecated and {behavior}. " - f"Support will be removed in {_LEGACY_ENV_REMOVED_IN}. Use env_akv_ref or ~/.pyrit/.env.local instead." + f"Auto-discovered plaintext environment file {env_file} {behavior}. Azure Key Vault through env_akv_ref " + "is more secure for shared or deployed secrets; use .env.local only for deliberate local overrides. " + "To inspect a resolved AKV-only configuration from a source checkout, run " + "`python -m build_scripts.export_akv_environment`; it writes ~/.pyrit/.env_akv." ) if not silent: print(f"WARNING: {message}") @@ -394,18 +451,14 @@ def _validate_dotenv_document( ) -async def _load_env_from_akv_async( +async def _fetch_akv_document_async( *, secret_url: str, strict: bool = True, silent: bool = False, - resolve_references_for_output: bool = False, -) -> str: +) -> _AkvEnvironmentDocument: """ - 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. + Fetch and validate one Key Vault bootstrap dotenv document. Authentication uses ``DefaultAzureCredential``, which silently tries managed identity, Azure CLI, VS Code credentials, etc., and falls back to interactive @@ -417,18 +470,14 @@ async def _load_env_from_akv_async( 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. - resolve_references_for_output (bool): If True, resolve child references even - when their runtime assignment loses to an existing process value, and - return a native dotenv document containing those resolved values. Returns: - str: The validated bootstrap dotenv document, with child-secret references - replaced when ``resolve_references_for_output`` is True. + _AkvEnvironmentDocument: Validated document text and source vault metadata. 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. + KeyVaultInitializationException: If the root URL is malformed or the bootstrap + document cannot be fetched and validated. ValueError: Compatibility base of ``KeyVaultInitializationException``. """ from azure.identity.aio import DefaultAzureCredential @@ -444,77 +493,10 @@ async def _load_env_from_akv_async( 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=StringIO(validated_document), interpolate=True) + parsed_environment = dotenv.dotenv_values(stream=StringIO(validated_document), interpolate=False) if not parsed_environment: raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") - existing_environment_names = set(os.environ) - loaded = dotenv.load_dotenv( - stream=StringIO(validated_document), - override=False, - interpolate=True, - ) - if not loaded: - return validated_document - - final_assignment_indexes = _get_final_assignment_indexes(document=validated_document) - resolved_reference_values: dict[int, str] = {} - skipped_reference_indexes: set[int] = set() - for variable_name, value in parsed_environment.items(): - if value is None: - continue - target = _parse_akv_reference(value) - if target is None: - continue - assignment_wins = variable_name not in existing_environment_names - if not assignment_wins and not resolve_references_for_output: - continue - try: - _, referenced_name, referenced_version = _parse_akv_reference_url( - target=target, - variable_name=variable_name, - expected_vault_url=vault_url, - ) - except ValueError as error: - if strict: - wrapped_error = _key_vault_initialization_error( - message=f"Invalid AKV reference for environment variable '{variable_name}'", - error=error, - ) - raise wrapped_error from error - if assignment_wins: - os.environ.pop(variable_name, None) - _warn_about_invalid_akv_reference( - variable_name=variable_name, - error=error, - silent=silent, - ) - skipped_reference_indexes.add(final_assignment_indexes[variable_name]) - continue - try: - resolved_value = await _fetch_akv_secret_value_async( - client=client, - secret_name=referenced_name, - secret_version=referenced_version, - variable_name=variable_name, - ) - resolved_reference_values[final_assignment_indexes[variable_name]] = resolved_value - if assignment_wins: - os.environ[variable_name] = resolved_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 - if resolve_references_for_output: - return _render_resolved_akv_document( - document=validated_document, - resolved_reference_values=resolved_reference_values, - skipped_reference_indexes=skipped_reference_indexes, - ) - return validated_document + return _AkvEnvironmentDocument(content=validated_document, vault_url=vault_url) except KeyVaultInitializationException: raise except Exception as error: @@ -530,7 +512,6 @@ 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: """ @@ -540,8 +521,6 @@ async def load_environment_async( 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: @@ -549,212 +528,47 @@ async def load_environment_async( """ 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] = [] + assignment_candidates: dict[str, list[_EnvironmentValueCandidate]] = {} 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.") - env_file = path.CONFIGURATION_DIRECTORY_PATH / _AKV_ENV_FILE_NAME - if env_akv_write_env and (env_file.exists() or env_file.is_symlink()): - raise ValueError(_get_akv_env_file_exists_message(env_file=env_file)) - 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, + if env_files is None: + dotenv_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" + if dotenv_file.exists(): + await asyncio.to_thread( + _warn_about_dotenv_file, + env_file=dotenv_file, + ignored_for_akv=True, silent=silent, - resolve_references_for_output=env_akv_write_env, ) - 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] + for secret_url in env_akv_ref: + document = await _fetch_akv_document_async( + secret_url=secret_url, + strict=env_akv_strict, + silent=silent, + ) + await asyncio.to_thread( + _load_dotenv_source, + document=document.content, + override=False, + assignment_candidates=assignment_candidates, + expected_vault_url=document.vault_url, + ) - assignment_candidates: dict[str, list[_EnvironmentValueCandidate]] = {} await asyncio.to_thread( - load_environment_files, - env_files=selected_env_files, + _load_environment_files, + env_files=env_files, silent=silent, include_default_base=not (env_akv_ref and env_files is None), assignment_candidates=assignment_candidates, ) - await _resolve_local_akv_references_async( + await _resolve_environment_candidates_async( assignment_candidates=assignment_candidates, strict=env_akv_strict, silent=silent, ) -def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Path: - """ - Write fetched bootstrap documents with resolved child-secret values. - - Returns: - pathlib.Path: Path to the written dotenv file. - - Raises: - ValueError: If the destination already exists or is a symbolic link. - OSError: If the filesystem cannot atomically publish the completed file. - """ - env_file = path.CONFIGURATION_DIRECTORY_PATH / _AKV_ENV_FILE_NAME - env_file.parent.mkdir(mode=0o700, parents=True, exist_ok=True) - if env_file.is_symlink(): - raise ValueError(f"Refusing to write the AKV environment through a symbolic link: {env_file}") - if env_file.exists(): - raise ValueError(_get_akv_env_file_exists_message(env_file=env_file)) - - content = _merge_akv_documents_for_debug(documents=documents) - file_descriptor: int | None = None - temporary_file: pathlib.Path | None = None - try: - file_descriptor, temporary_name = tempfile.mkstemp( - prefix=f"{env_file.name}.", - suffix=".tmp", - dir=env_file.parent, - ) - temporary_file = pathlib.Path(temporary_name) - file_chmod = getattr(os, "fchmod", None) - if file_chmod is not None: - file_chmod(file_descriptor, 0o600) - else: - os.chmod(temporary_file, 0o600) - stream = os.fdopen(file_descriptor, "w", encoding="utf-8", newline="") - file_descriptor = None - with stream: - stream.write(content) - try: - os.link(temporary_file, env_file) - except FileExistsError as error: - raise ValueError(_get_akv_env_file_exists_message(env_file=env_file)) from error - finally: - if file_descriptor is not None: - os.close(file_descriptor) - if temporary_file is not None: - with contextlib.suppress(FileNotFoundError): - temporary_file.unlink() - - _print_msg(f"Saved Key Vault bootstrap environment file: {env_file}", quiet=silent, log=True) - return env_file - - -def _get_akv_env_file_exists_message(*, env_file: pathlib.Path) -> str: - """ - Create the message used when debug output would clobber an existing path. - - Returns: - str: Error message containing recovery guidance for the user. - """ - return ( - f"Cannot write the resolved Key Vault environment because {env_file} already exists; " - "rename or remove it before enabling env_akv_write_env." - ) - - -def _merge_akv_documents_for_debug(*, documents: Sequence[str]) -> str: - """ - Merge resolved bootstrap documents using runtime first-document precedence. - - Duplicate assignments within one document are retained because interpolation - depends on assignment order. Assignments established by an earlier document - are omitted from later documents. - - Returns: - str: A native dotenv document with equivalent bootstrap precedence. - """ - established_names: set[str] = set() - merged_bindings: list[str] = [] - for document in documents: - document_names: set[str] = set() - for binding in parse_stream(StringIO(document)): - if binding.key is None or binding.key not in established_names: - merged_bindings.append(binding.original.string) - if binding.key is not None: - document_names.add(binding.key) - established_names.update(document_names) - - return "".join(merged_bindings).rstrip("\r\n") + "\n" - - -def _render_resolved_akv_document( - *, - document: str, - resolved_reference_values: Mapping[int, str], - skipped_reference_indexes: set[int] | None = None, -) -> str: - """ - Replace resolved Key Vault reference assignments with native dotenv values. - - Returns: - str: Dotenv text that preserves non-reference bindings and comments. - """ - skipped_reference_indexes = skipped_reference_indexes or set() - rendered_bindings: list[str] = [] - for binding_index, binding in enumerate(parse_stream(StringIO(document))): - variable_name = binding.key - if binding_index in skipped_reference_indexes: - continue - if variable_name is not None and binding_index in resolved_reference_values: - original = binding.original.string - export_prefix = "export " if original.lstrip().startswith("export ") else "" - if original.endswith("\r\n"): - newline = "\r\n" - elif original.endswith("\n"): - newline = "\n" - else: - newline = "" - rendered_bindings.append( - f"{export_prefix}{variable_name}=" - f"{_serialize_terminal_dotenv_value(resolved_reference_values[binding_index])}{newline}" - ) - else: - rendered_bindings.append(binding.original.string) - return "".join(rendered_bindings) - - -def _get_final_assignment_indexes(*, document: str) -> dict[str, int]: - """ - Map each variable name to its final assignment occurrence in a dotenv document. - - Returns: - dict[str, int]: Final parsed binding index for each assigned variable. - """ - return { - binding.key: binding_index - for binding_index, binding in enumerate(parse_stream(StringIO(document))) - if binding.key is not None - } - - -def _serialize_terminal_dotenv_value(value: str) -> str: - """ - Quote a terminal secret value for a native python-dotenv round trip. - - The empty-name default expression produces a literal dollar sign during - interpolation, preventing terminal ``${NAME}`` text from being reinterpreted. - - Returns: - str: A single-quoted dotenv value. - """ - escaped_value = value.replace("\\", "\\\\").replace("'", "\\'").replace("${", "${:-$}{") - return f"'{escaped_value}'" - - def _warn_about_invalid_akv_reference(*, variable_name: str, error: ValueError, silent: bool) -> None: """Warn that a malformed Key Vault reference assignment is being skipped.""" message = f"Invalid AKV reference for environment variable '{variable_name}' will be skipped: {error}" @@ -774,14 +588,6 @@ def _parse_akv_reference(value: str) -> str | None: 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 _parse_akv_reference_url( *, target: str, @@ -810,18 +616,17 @@ def _parse_akv_reference_url( f"Expected vault '{expected_vault_url}', got '{referenced_vault_url}'." ) - _validate_akv_secret_name(secret_name=secret_name, variable_name=variable_name) return referenced_vault_url, secret_name, secret_version -async def _resolve_local_akv_references_async( +async def _resolve_environment_candidates_async( *, assignment_candidates: Mapping[str, Sequence[_EnvironmentValueCandidate]], strict: bool, silent: bool, ) -> None: """ - Resolve complete Key Vault references from winning local assignments. + Resolve complete Key Vault references from winning environment assignments. Raises: KeyVaultInitializationException: If strict validation or secret retrieval fails. @@ -840,6 +645,7 @@ async def _resolve_local_akv_references_async( vault_url, secret_name, secret_version = _parse_akv_reference_url( target=target, variable_name=variable_name, + expected_vault_url=candidate.expected_vault_url, ) except ValueError as error: if strict: @@ -890,31 +696,3 @@ async def _resolve_local_akv_references_async( error=error, ) raise wrapped_error from error - - -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. - """ - _, secret_name, secret_version = _parse_akv_reference_url( - target=target, - variable_name=variable_name, - expected_vault_url=vault_url, - ) - return secret_name, secret_version diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index df921d5630..0c345aa595 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -7,7 +7,7 @@ from pyrit.common.apply_defaults import reset_default_values from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory -from pyrit.setup.environment_loading import load_environment_async, validate_akv_boolean_options +from pyrit.setup.environment_loading import load_environment_async, validate_env_akv_strict if TYPE_CHECKING: from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -70,7 +70,6 @@ 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: @@ -96,31 +95,25 @@ async def initialize_pyrit_async( and ``scorer`` target variants remain opt-in. env_files (Sequence[pathlib.Path] | None): Optional sequence of environment file paths to load in order. Ordinary files fill missing process values; files named ``.env.local`` override. - If omitted, PyRIT auto-discovers legacy ``.env`` and supported ``.env.local`` files. + If omitted, PyRIT auto-discovers supported ``.env`` and ``.env.local`` files. env_akv_ref (Sequence[str] | None): Optional ordered Azure Key Vault URLs whose secret values contain bootstrap dotenv documents. Documents fill missing process values and support complete-value references to scalar secrets. Requires ``azure-keyvault-secrets``. env_akv_strict (bool): If True, reject malformed bootstrap entries and Key Vault reference syntax. If False, warn and skip those entries. Operational Key Vault failures always raise. - env_akv_write_env (bool): If True, write fully resolved bootstrap documents with plaintext - child-secret values to ``~/.pyrit/.env`` for debugging. 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. Raises: - TypeError: If ``env_akv_strict`` or ``env_akv_write_env`` is not a bool. + TypeError: If ``env_akv_strict`` is not a bool. ValueError: If an unsupported memory_db_type is provided or env_files contains non-existent files. """ - validate_akv_boolean_options( - env_akv_strict=env_akv_strict, - env_akv_write_env=env_akv_write_env, - ) + validate_env_akv_strict(env_akv_strict=env_akv_strict) await load_environment_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/unit/build_scripts/test_export_akv_environment.py b/tests/unit/build_scripts/test_export_akv_environment.py new file mode 100644 index 0000000000..9b5531f753 --- /dev/null +++ b/tests/unit/build_scripts/test_export_akv_environment.py @@ -0,0 +1,303 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import io +import os +import pathlib +from types import SimpleNamespace +from unittest import mock + +import dotenv +import pytest + +from build_scripts.export_akv_environment import ( + DEFAULT_OUTPUT_FILE, + _Document, + _render, + _serialize, + _write_output, + export_akv_environment, +) + + +@pytest.mark.parametrize( + "value", + [ + "single\\backslash", + "double\\\\backslash", + "four\\\\\\\\backslashes", + "\\leading-and-trailing\\", + r"C:\Users\name\secret.txt", + "quote'\\${LITERAL}\nline\\two", + ], +) +def test_serialize_round_trips_terminal_values(value: str) -> None: + document = f"VALUE={_serialize(value)}\n" + assert dotenv.dotenv_values(stream=io.StringIO(document), interpolate=True)["VALUE"] == value + + +def test_render_resolves_akv_only_values() -> None: + document = _Document( + content=( + "# AKV config\nBASE=bootstrap\nDERIVED=${BASE}\nAPI_KEY=kv:https://vault.vault.azure.net/secrets/api-key\n" + ), + vault_url="https://vault.vault.azure.net", + ) + client = mock.MagicMock() + client.get_secret.return_value = SimpleNamespace(value="resolved-key") + + rendered = _render( + documents=[document], + credential=mock.MagicMock(), + clients={document.vault_url: client}, + strict=True, + silent=True, + ) + + values = dotenv.dotenv_values(stream=io.StringIO(rendered), interpolate=True) + assert values == {"BASE": "bootstrap", "DERIVED": "bootstrap", "API_KEY": "resolved-key"} + assert "# AKV config" in rendered + client.get_secret.assert_called_once_with("api-key", version=None) + + +@pytest.mark.parametrize( + ("content", "expected_values", "expected_child_fetches"), + [ + ( + "A=kv:https://vault.vault.azure.net/secrets/key\nB=${A}\nA=literal\n", + {"A": "literal", "B": "resolved-key"}, + 1, + ), + ( + "A=literal\nB=${A}\nA=kv:https://vault.vault.azure.net/secrets/key\n", + {"A": "resolved-key", "B": "literal"}, + 1, + ), + ( + "A=kv:https://vault.vault.azure.net/secrets/key\nB=${A}\nC=${B}\nB=literal\n", + {"A": "resolved-key", "B": "literal", "C": "resolved-key"}, + 2, + ), + ], +) +def test_render_resolves_interpolated_reference_assignments( + content: str, expected_values: dict[str, str], expected_child_fetches: int +) -> None: + document = _Document(content=content, vault_url="https://vault.vault.azure.net") + client = mock.MagicMock() + client.get_secret.return_value = SimpleNamespace(value="resolved-key") + + rendered = _render( + documents=[document], + credential=mock.MagicMock(), + clients={document.vault_url: client}, + strict=True, + silent=True, + ) + + values = dict(dotenv.dotenv_values(stream=io.StringIO(rendered), interpolate=True)) + assert values == expected_values + assert client.get_secret.call_args_list == [mock.call("key", version=None)] * expected_child_fetches + + +def test_render_preserves_first_document_values() -> None: + documents = [ + _Document( + content="SHARED=first\nFIRST_ONLY=first\n", + vault_url="https://vault.vault.azure.net", + ), + _Document( + content="SHARED=second\nSECOND_ONLY=second\n", + vault_url="https://vault.vault.azure.net", + ), + ] + + rendered = _render( + documents=documents, + credential=mock.MagicMock(), + clients={}, + strict=True, + silent=True, + ) + + assert dotenv.dotenv_values(stream=io.StringIO(rendered), interpolate=True) == { + "SHARED": "first", + "FIRST_ONLY": "first", + "SECOND_ONLY": "second", + } + + +def test_render_non_strict_warns_and_skips_invalid_reference(caplog: pytest.LogCaptureFixture, capsys) -> None: + document = _Document( + content="GOOD=resolved\nBAD=kv:short-name\nOTHER=also-resolved", + vault_url="https://vault.vault.azure.net", + ) + + with caplog.at_level("WARNING", logger="build_scripts.export_akv_environment"): + rendered = _render( + documents=[document], + credential=mock.MagicMock(), + clients={}, + strict=False, + silent=False, + ) + + assert "BAD=" not in rendered + assert dotenv.dotenv_values(stream=io.StringIO(rendered)) == { + "GOOD": "resolved", + "OTHER": "also-resolved", + } + assert "WARNING: Invalid AKV reference for 'BAD' will be skipped" in capsys.readouterr().out + assert "BAD" in caplog.text + + +def test_export_writes_env_akv_without_process_values(tmp_path: pathlib.Path) -> None: + credential = mock.MagicMock() + client = mock.MagicMock() + client.get_secret.side_effect = [ + SimpleNamespace(value="VALUE=bootstrap\nKEY=kv:https://vault.vault.azure.net/secrets/key\n"), + SimpleNamespace(value="resolved-key"), + ] + output_file = tmp_path / ".env_akv" + + with ( + mock.patch.dict(os.environ, {"PROCESS_ONLY": "not-written"}, clear=True), + mock.patch("build_scripts.export_akv_environment._create_client", return_value=client), + ): + output = export_akv_environment( + secret_urls=["https://vault.vault.azure.net/secrets/bootstrap"], + output_file=output_file, + credential=credential, + silent=True, + ) + + assert output == output_file + assert output.name == ".env_akv" + assert DEFAULT_OUTPUT_FILE.name == ".env_akv" + assert not (tmp_path / ".env").exists() + assert dotenv.dotenv_values(dotenv_path=output_file) == {"VALUE": "bootstrap", "KEY": "resolved-key"} + assert "PROCESS_ONLY" not in output_file.read_text(encoding="utf-8") + client.close.assert_called_once() + credential.close.assert_not_called() + + +def test_export_closes_client_but_not_provided_credential_on_failure(tmp_path: pathlib.Path) -> None: + credential = mock.MagicMock() + client = mock.MagicMock() + client.get_secret.side_effect = RuntimeError("fetch failed") + + with ( + mock.patch("build_scripts.export_akv_environment._create_client", return_value=client), + pytest.raises(RuntimeError, match="fetch failed"), + ): + export_akv_environment( + secret_urls=["https://vault.vault.azure.net/secrets/bootstrap"], + output_file=tmp_path / ".env_akv", + credential=credential, + silent=True, + ) + + client.close.assert_called_once() + credential.close.assert_not_called() + + +def test_export_rejects_existing_output_before_fetch(tmp_path: pathlib.Path) -> None: + output_file = tmp_path / ".env_akv" + output_file.write_text("ORIGINAL=value\n", encoding="utf-8") + + with mock.patch("build_scripts.export_akv_environment._fetch_documents") as mock_fetch: + with pytest.raises(ValueError, match="already exists"): + export_akv_environment( + secret_urls=["https://vault.vault.azure.net/secrets/bootstrap"], + output_file=output_file, + credential=mock.MagicMock(), + silent=True, + ) + + mock_fetch.assert_not_called() + + +def test_write_output_does_not_clobber_existing_file(tmp_path: pathlib.Path) -> None: + output_file = tmp_path / ".env_akv" + output_file.write_text("ORIGINAL=value\n", encoding="utf-8") + + with pytest.raises(ValueError, match="already exists"): + _write_output(output_file=output_file, document="NEW=value\n") + + assert output_file.read_text(encoding="utf-8") == "ORIGINAL=value\n" + assert list(tmp_path.glob(".env_akv.*.tmp")) == [] + + +def test_write_output_secures_descriptor_before_writing(tmp_path: pathlib.Path) -> None: + events: list[str] = [] + temporary_file = tmp_path / ".env_akv.test.tmp" + stream = mock.MagicMock() + stream.write.side_effect = lambda content: events.append(f"write:{content}") + + with ( + mock.patch( + "build_scripts.export_akv_environment.tempfile.mkstemp", + side_effect=lambda **kwargs: events.append("create") or (7, str(temporary_file)), + ), + mock.patch( + "build_scripts.export_akv_environment.os.fchmod", + side_effect=lambda *args: events.append("fchmod"), + create=True, + ), + mock.patch( + "build_scripts.export_akv_environment.os.fdopen", + side_effect=lambda *args, **kwargs: events.append("fdopen") or stream, + ), + mock.patch( + "build_scripts.export_akv_environment.os.link", + side_effect=lambda *args: events.append("link"), + ), + ): + _write_output(output_file=tmp_path / ".env_akv", document="VALUE=bootstrap\n") + + assert events == ["create", "fchmod", "fdopen", "write:VALUE=bootstrap\n", "link"] + + +def test_write_output_does_not_clobber_file_created_before_publish(tmp_path: pathlib.Path) -> None: + output_file = tmp_path / ".env_akv" + real_link = os.link + + def competing_link(source: str | os.PathLike[str], destination: str | os.PathLike[str]) -> None: + output_file.write_text("CREATED_BY_OTHER_PROCESS=value\n", encoding="utf-8") + real_link(source, destination) + + with ( + mock.patch("build_scripts.export_akv_environment.os.link", side_effect=competing_link), + pytest.raises(ValueError, match="already exists"), + ): + _write_output(output_file=output_file, document="NEW=value\n") + + assert output_file.read_text(encoding="utf-8") == "CREATED_BY_OTHER_PROCESS=value\n" + assert list(tmp_path.glob(".env_akv.*.tmp")) == [] + + +@pytest.mark.skipif(os.name != "posix", reason="POSIX permission bits are not enforced on this platform.") +def test_write_output_uses_owner_only_permissions(tmp_path: pathlib.Path) -> None: + configuration_directory = tmp_path / ".pyrit" + output_file = _write_output( + output_file=configuration_directory / ".env_akv", + document="VALUE=bootstrap\n", + ) + + assert configuration_directory.stat().st_mode & 0o777 == 0o700 + assert output_file.stat().st_mode & 0o777 == 0o600 + + +def test_write_output_rejects_symbolic_link(tmp_path: pathlib.Path) -> None: + target = tmp_path / "target" + target.write_text("unchanged", encoding="utf-8") + output_file = tmp_path / ".env_akv" + try: + output_file.symlink_to(target) + except OSError: + pytest.skip("Symbolic links are unavailable on this platform.") + + with pytest.raises(ValueError, match="symbolic link"): + _write_output(output_file=output_file, document="VALUE=bootstrap\n") + + assert target.read_text(encoding="utf-8") == "unchanged" diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 65de4b684e..8d74a209c9 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -43,17 +43,12 @@ 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 - @pytest.mark.parametrize("option_name", ["env_akv_strict", "env_akv_write_env"]) @pytest.mark.parametrize("invalid_value", ["false", "true", 0, 1, None, [], {}]) - def test_rejects_non_boolean_akv_options(self, option_name, invalid_value): - with pytest.raises(TypeError, match=rf"{option_name} must be a bool"): - if option_name == "env_akv_strict": - ConfigurationLoader(env_akv_strict=invalid_value) # type: ignore[arg-type] - else: - ConfigurationLoader(env_akv_write_env=invalid_value) # type: ignore[arg-type] + def test_rejects_non_boolean_env_akv_strict(self, invalid_value): + with pytest.raises(TypeError, match=r"env_akv_strict must be a bool"): + ConfigurationLoader(env_akv_strict=invalid_value) # type: ignore[arg-type] def test_valid_memory_db_types_snake_case(self): """Test all valid memory database types in snake_case.""" @@ -159,7 +154,6 @@ 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) @@ -169,7 +163,6 @@ 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): @@ -246,22 +239,20 @@ def test_from_yaml_file(self): finally: pathlib.Path(yaml_path).unlink() - @pytest.mark.parametrize("option_name", ["env_akv_strict", "env_akv_write_env"]) - def test_from_yaml_rejects_quoted_boolean_akv_options(self, tmp_path, option_name): + def test_from_yaml_rejects_quoted_env_akv_strict(self, tmp_path): yaml_path = tmp_path / "quoted-boolean.yaml" - yaml_path.write_text(f'{option_name}: "false"\n', encoding="utf-8") + yaml_path.write_text('env_akv_strict: "false"\n', encoding="utf-8") - with pytest.raises(TypeError, match=rf"{option_name} must be a bool"): + with pytest.raises(TypeError, match=r"env_akv_strict must be a bool"): ConfigurationLoader.from_yaml_file(yaml_path) - def test_from_yaml_accepts_native_boolean_akv_options(self, tmp_path): + def test_from_yaml_accepts_native_env_akv_strict(self, tmp_path): yaml_path = tmp_path / "native-booleans.yaml" - yaml_path.write_text("env_akv_strict: false\nenv_akv_write_env: true\n", encoding="utf-8") + yaml_path.write_text("env_akv_strict: false\n", encoding="utf-8") config = ConfigurationLoader.from_yaml_file(yaml_path) assert config.env_akv_strict is False - assert config.env_akv_write_env is True def test_from_empty_yaml_file_raises_value_error(self, tmp_path): """Test that an empty YAML file raises a clear ValueError.""" @@ -375,7 +366,6 @@ 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") @@ -389,7 +379,6 @@ async def test_initialize_pyrit_async_with_env_akv_ref(self, mock_init): memory_db_type="in_memory", env_akv_ref=refs, env_akv_strict=False, - env_akv_write_env=True, ) await config.initialize_pyrit_async() @@ -398,7 +387,6 @@ 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") @@ -566,14 +554,6 @@ 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_environment_loading.py b/tests/unit/setup/test_environment_loading.py index 896da5d0d8..016c980b2b 100644 --- a/tests/unit/setup/test_environment_loading.py +++ b/tests/unit/setup/test_environment_loading.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import io import os import pathlib import tempfile @@ -11,17 +10,15 @@ import pytest from azure.core.exceptions import ResourceNotFoundError -from dotenv import dotenv_values from pyrit.exceptions import KeyVaultInitializationException from pyrit.setup import IN_MEMORY, initialize_pyrit_async from pyrit.setup.environment_loading import ( - _load_env_from_akv_async, + _AkvEnvironmentDocument, + _fetch_akv_document_async, _parse_akv_reference, _parse_akv_secret_url, - _serialize_terminal_dotenv_value, - _warn_about_akv_environment_files, - _write_akv_env_file, + _warn_about_dotenv_file, load_environment_async, load_environment_files, ) @@ -41,10 +38,7 @@ async def test_loads_default_env_files_when_none_provided(self, mock_config_path 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), - pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), - ): + with mock.patch.dict(os.environ, {}, clear=True): loaded = load_environment_files(env_files=None) assert loaded is True @@ -60,10 +54,7 @@ async def test_only_loads_existing_default_files(self, mock_config_path): env_file.write_text("VAR1=value1") mock_config_path.__truediv__ = lambda self, other: temp_path / other - with ( - mock.patch.dict(os.environ, {}, clear=True), - pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), - ): + with mock.patch.dict(os.environ, {}, clear=True): loaded = load_environment_files(env_files=None) assert loaded is True @@ -76,10 +67,7 @@ async def test_default_env_preserves_process_environment(self, mock_config_path) (temp_path / ".env").write_text("VAR=legacy\nLEGACY_ONLY=legacy") mock_config_path.__truediv__ = lambda self, other: temp_path / other - with ( - mock.patch.dict(os.environ, {"VAR": "process"}, clear=True), - pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), - ): + with mock.patch.dict(os.environ, {"VAR": "process"}, clear=True): loaded = load_environment_files(env_files=None, silent=True) assert loaded is True @@ -94,10 +82,7 @@ async def test_default_env_local_overrides_process_environment_and_env(self, moc (temp_path / ".env.local").write_text("VAR=local") mock_config_path.__truediv__ = lambda self, other: temp_path / other - with ( - mock.patch.dict(os.environ, {"VAR": "process"}, clear=True), - pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), - ): + with mock.patch.dict(os.environ, {"VAR": "process"}, clear=True): loaded = load_environment_files(env_files=None, silent=True) assert loaded is True @@ -133,7 +118,7 @@ async def test_returns_false_when_no_default_files_exist(self, mock_config_path) assert os.environ == {} @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") - def test_auto_discovered_env_warns_with_removal_version(self, mock_config_path, caplog, capsys): + def test_auto_discovered_env_warns_about_plaintext(self, mock_config_path, caplog, capsys): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) env_file = temp_path / ".env" @@ -142,44 +127,61 @@ def test_auto_discovered_env_warns_with_removal_version(self, mock_config_path, with ( caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"), - pytest.warns(DeprecationWarning, match=r"\.env.*removed in 1\.3\.0.*\.env\.local"), + warnings.catch_warnings(), ): + warnings.simplefilter("error", DeprecationWarning) load_environment_files(env_files=None) output = capsys.readouterr().out - assert f"WARNING: Auto-discovered {env_file} is deprecated" in output - assert "Use env_akv_ref or ~/.pyrit/.env.local instead" in output - assert f"Auto-discovered {env_file} is deprecated" in caplog.text + assert f"Auto-discovered plaintext environment file {env_file} will be loaded" in output + assert "Azure Key Vault through env_akv_ref is more secure" in output + assert "build_scripts.export_akv_environment" in output + assert "~/.pyrit/.env_akv" in caplog.text @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") - def test_explicit_env_file_does_not_emit_legacy_deprecation(self, mock_config_path): + def test_default_discovery_does_not_load_env_akv(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env_akv").write_text("EXPORTED_SECRET=not-loaded", encoding="utf-8") + 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 False + assert "EXPORTED_SECRET" not in os.environ + + @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") + def test_explicit_env_file_does_not_emit_auto_discovery_warning(self, mock_config_path, caplog): with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) explicit_env = temp_path / ".env" explicit_env.write_text("VAR=explicit") mock_config_path.__truediv__ = lambda self, other: temp_path / other - with warnings.catch_warnings(): - warnings.simplefilter("error", DeprecationWarning) + with caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"): loaded = load_environment_files(env_files=[explicit_env], silent=True) assert loaded is True + assert "Azure Key Vault through env_akv_ref is more secure" not in caplog.text @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") - def test_akv_legacy_env_warning_respects_silent(self, mock_config_path, caplog, capsys): + def test_akv_dotenv_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.environment_loading"), - pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), - ): - _warn_about_akv_environment_files(env_files=None, silent=True) + with caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"): + _warn_about_dotenv_file( + env_file=temp_path / ".env", + ignored_for_akv=True, + silent=True, + ) assert capsys.readouterr().out == "" assert "will be ignored because env_akv_ref is configured" in caplog.text + assert "build_scripts.export_akv_environment" in caplog.text @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") async def test_akv_ignores_auto_discovered_env_and_loads_env_local(self, mock_config_path): @@ -191,48 +193,25 @@ async def test_akv_ignores_auto_discovered_env_and_loads_env_local(self, mock_co with ( mock.patch( - "pyrit.setup.environment_loading._load_env_from_akv_async", + "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock, - return_value="VALUE=akv\n", + return_value=_AkvEnvironmentDocument( + content="VALUE=akv\n", + vault_url="https://vault.vault.azure.net", + ), ), - mock.patch("pyrit.setup.environment_loading.load_environment_files") as mock_load_files, - pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + mock.patch("pyrit.setup.environment_loading._load_environment_files") as mock_load_files, ): await load_environment_async( env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], env_files=None, env_akv_strict=True, - env_akv_write_env=False, silent=True, ) assert mock_load_files.call_args.kwargs["env_files"] is None assert mock_load_files.call_args.kwargs["include_default_base"] is False - @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") - async def test_akv_debug_mode_rejects_existing_env_before_fetch(self, mock_config_path): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - env_file.write_text("VALUE=legacy") - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - with ( - mock.patch( - "pyrit.setup.environment_loading._load_env_from_akv_async", new_callable=mock.AsyncMock - ) as mock_load_akv, - pytest.raises(ValueError, match=r"already exists.*rename or remove"), - ): - await load_environment_async( - env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], - env_files=None, - env_akv_strict=True, - env_akv_write_env=True, - silent=True, - ) - - mock_load_akv.assert_not_awaited() - 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: @@ -301,266 +280,121 @@ async def test_load_environment_files_honors_python_dotenv_disabled(self): assert loaded is True assert "DISABLED_VALUE" not in os.environ - @pytest.mark.parametrize( - "value", - [ - "single\\backslash", - "double\\\\backslash", - "four\\\\\\\\backslashes", - "\\leading-and-trailing\\", - r"C:\Users\name\secret.txt", - "quote'\\${LITERAL}\nline\\two", - ], - ) - def test_serialize_terminal_dotenv_value_preserves_backslashes(self, value): - document = f"VALUE={_serialize_terminal_dotenv_value(value)}\n" - - with mock.patch.dict(os.environ, {}, clear=True): - reloaded_value = dotenv_values(stream=io.StringIO(document), interpolate=True)["VALUE"] - - assert reloaded_value == value - - async def test_load_environment_async_write_env_writes_resolved_native_bootstrap(self): + async def test_runtime_does_not_fetch_akv_reference_overridden_by_env_local(self): credential, client = _create_mock_akv_clients() - document = ( - "# Bootstrap values\n" - "BASE=bootstrap\n" - "DERIVED=${BASE}\n" - "API_KEY=kv:https://vault.vault.azure.net/secrets/api-key\n" - ) - resolved_api_key = "line one\nquote' and literal ${UNRELATED}" - client.get_secret = mock.AsyncMock( - side_effect=[ - types.SimpleNamespace(value=document), - types.SimpleNamespace(value=resolved_api_key), - ] - ) + bootstrap_document = "API_KEY=kv:https://vault.vault.azure.net/secrets/api-key\nAKV_ONLY=akv\n" + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=bootstrap_document)) with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - (temp_path / ".env.local").write_text("API_KEY=local-key\nLOCAL_ONLY=local", encoding="utf-8") + local_file = pathlib.Path(temp_dir) / ".env.local" + local_file.write_text("API_KEY=local-key\n", encoding="utf-8") + with ( - mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - mock.patch.dict( - os.environ, - { - "BASE": "process", - "API_KEY": "process-key", - "PROCESS_ONLY": "not-written", - "UNRELATED": "changed", - }, - clear=True, - ), + 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_environment_async( env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], - env_files=None, + env_files=[local_file], env_akv_strict=True, - env_akv_write_env=True, silent=True, ) - assert os.environ["BASE"] == "process" assert os.environ["API_KEY"] == "local-key" - assert os.environ["LOCAL_ONLY"] == "local" + assert os.environ["AKV_ONLY"] == "akv" - written_env = temp_path / ".env" - assert written_env.is_file() - assert not (temp_path / ".env.new").exists() - content = written_env.read_text(encoding="utf-8") - assert "# Bootstrap values" in content - assert "kv:" not in content - assert "PROCESS_ONLY" not in content - assert "LOCAL_ONLY" not in content + client.get_secret.assert_awaited_once_with("bootstrap", version=None) - with mock.patch.dict(os.environ, {}, clear=True): - written_values = dotenv_values(dotenv_path=written_env, interpolate=True) - - assert written_values == { - "BASE": "bootstrap", - "DERIVED": "bootstrap", - "API_KEY": resolved_api_key, - } - assert client.get_secret.await_args_list == [ - mock.call("bootstrap", version=None), - mock.call("api-key", version=None), - ] + async def test_runtime_resolves_local_alias_after_all_sources_are_loaded(self): + credential, client = _create_mock_akv_clients() + bootstrap_document = "A=kv:https://vault.vault.azure.net/secrets/key\n" + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=bootstrap_document)) - 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") + ordinary_file = temp_path / "ordinary.env" + ordinary_file.write_text("B=${A}\n", encoding="utf-8") + local_file = temp_path / ".env.local" + local_file.write_text("A=literal\n", encoding="utf-8") + + async def resolve_by_variable_name(**kwargs): + return f"resolved-for-{kwargs['variable_name']}" with ( - mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + 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), mock.patch( - "pyrit.setup.environment_loading._load_env_from_akv_async", + "pyrit.setup.environment_loading._fetch_akv_secret_value_async", new_callable=mock.AsyncMock, - return_value="VALUE=bootstrap\n", - ), - mock.patch( - "pyrit.setup.environment_loading.load_environment_files", return_value=True - ) as mock_load_environment_files, + side_effect=resolve_by_variable_name, + ) as mock_fetch_child, ): await load_environment_async( env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], - env_files=[generated_env, local_env], + env_files=[ordinary_file, local_file], 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 + assert os.environ["A"] == "literal" + assert os.environ["B"] == "resolved-for-B" + + mock_fetch_child.assert_awaited_once() + assert mock_fetch_child.await_args is not None + assert mock_fetch_child.await_args.kwargs["variable_name"] == "B" - async def test_load_environment_async_write_env_preserves_first_bootstrap_value(self): - documents = [ - "SHARED=first\nFIRST_ONLY=first\n", - "SHARED=second\nSECOND_ONLY=second\n", - ] + @pytest.mark.parametrize("fallback_source", ["akv", "file"]) + async def test_non_strict_runtime_uses_blocked_lower_candidate_after_malformed_akv_winner(self, fallback_source): + credential, client = _create_mock_akv_clients() + documents = { + "https://vault.vault.azure.net/secrets/first": "API_KEY=kv:short\n", + "https://vault.vault.azure.net/secrets/second": ( + "API_KEY=kv:https://vault.vault.azure.net/secrets/fallback-key\n" + ), + } + + async def fetch_document(secret_url, **kwargs): + return _AkvEnvironmentDocument( + content=documents[secret_url], + vault_url="https://vault.vault.azure.net", + ) with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) + env_files: list[pathlib.Path] = [] + env_akv_ref = ["https://vault.vault.azure.net/secrets/first"] + if fallback_source == "akv": + env_akv_ref.append("https://vault.vault.azure.net/secrets/second") + else: + ordinary_file = pathlib.Path(temp_dir) / "ordinary.env" + ordinary_file.write_text( + "API_KEY=kv:https://local-vault.vault.azure.net/secrets/fallback-key\n", + encoding="utf-8", + ) + env_files.append(ordinary_file) + with ( - mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch.dict(os.environ, {}, clear=True), mock.patch( - "pyrit.setup.environment_loading._load_env_from_akv_async", + "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock, - side_effect=lambda **kwargs: documents.pop(0), + side_effect=fetch_document, ), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), ): + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="resolved-fallback")) await load_environment_async( - env_akv_ref=[ - "https://vault.vault.azure.net/secrets/first", - "https://vault.vault.azure.net/secrets/second", - ], - env_files=[], - env_akv_strict=True, - env_akv_write_env=True, + env_akv_ref=env_akv_ref, + env_files=env_files, + env_akv_strict=False, silent=True, ) - with mock.patch.dict(os.environ, {}, clear=True): - written_values = dotenv_values(dotenv_path=temp_path / ".env", interpolate=True) - - assert written_values == { - "SHARED": "first", - "FIRST_ONLY": "first", - "SECOND_ONLY": "second", - } - - def test_write_akv_env_file_secures_descriptor_before_writing(self): - events: list[str] = [] - - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - temporary_file = temp_path / ".env.test.tmp" - stream = mock.MagicMock() - stream.write.side_effect = lambda content: events.append(f"write:{content}") - with ( - mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - mock.patch( - "pyrit.setup.environment_loading.tempfile.mkstemp", - side_effect=lambda **kwargs: events.append("create") or (7, str(temporary_file)), - ), - mock.patch( - "pyrit.setup.environment_loading.os.fchmod", - side_effect=lambda *args: events.append("fchmod"), - create=True, - ), - mock.patch( - "pyrit.setup.environment_loading.os.fdopen", - side_effect=lambda *args, **kwargs: events.append("fdopen") or stream, - ), - mock.patch( - "pyrit.setup.environment_loading.os.link", - side_effect=lambda *args: events.append("link"), - ), - ): - _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) - - assert events == ["create", "fchmod", "fdopen", "write:VALUE=bootstrap\n", "link"] - - def test_write_akv_env_file_rejects_existing_file(self): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - env_file.write_text("ORIGINAL=value\n", encoding="utf-8") - - with ( - mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - pytest.raises(ValueError, match="already exists.*rename or remove"), - ): - _write_akv_env_file(documents=["NEW=value\n"], silent=True) - - assert env_file.read_text(encoding="utf-8") == "ORIGINAL=value\n" - assert list(temp_path.glob(".env.*.tmp")) == [] - - def test_write_akv_env_file_does_not_clobber_file_created_before_publish(self): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - real_fdopen = os.fdopen - - class CompetingFileStream: - def __init__(self, *args, **kwargs): - self._stream = real_fdopen(*args, **kwargs) - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - result = self._stream.__exit__(exc_type, exc_value, traceback) - env_file.write_text("CREATED_BY_OTHER_PROCESS=value\n", encoding="utf-8") - return result - - def write(self, content): - return self._stream.write(content) - - with ( - mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - mock.patch("pyrit.setup.environment_loading.os.fdopen", side_effect=CompetingFileStream), - pytest.raises(ValueError, match="already exists.*rename or remove"), - ): - _write_akv_env_file(documents=["NEW=value\n"], silent=True) - - assert env_file.read_text(encoding="utf-8") == "CREATED_BY_OTHER_PROCESS=value\n" - assert list(temp_path.glob(".env.*.tmp")) == [] - - @pytest.mark.skipif(os.name != "posix", reason="POSIX permission bits are not enforced on this platform.") - def test_write_akv_env_file_uses_owner_only_permissions(self): - with tempfile.TemporaryDirectory() as temp_dir: - configuration_directory = pathlib.Path(temp_dir) / ".pyrit" - with mock.patch( - "pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", configuration_directory - ): - env_file = _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) - - assert configuration_directory.stat().st_mode & 0o777 == 0o700 - assert env_file.stat().st_mode & 0o777 == 0o600 - - def test_write_akv_env_file_rejects_symbolic_link(self): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - target = temp_path / "target" - target.write_text("unchanged", encoding="utf-8") - env_file = temp_path / ".env" - try: - env_file.symlink_to(target) - except OSError: - pytest.skip("Symbolic links are unavailable on this platform.") - - with ( - mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - pytest.raises(ValueError, match="symbolic link"), - ): - _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) + assert os.environ["API_KEY"] == "resolved-fallback" - assert target.read_text(encoding="utf-8") == "unchanged" + client.get_secret.assert_awaited_once_with("fallback-key", version=None) async def test_direct_local_file_loader_keeps_pyrit_references_literal(self): with tempfile.TemporaryDirectory() as temp_dir: @@ -589,10 +423,7 @@ 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 - with ( - mock.patch.dict(os.environ, {}, clear=True), - pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), - ): + with mock.patch.dict(os.environ, {}, clear=True): loaded = load_environment_files(env_files=None, silent=True) assert loaded is True @@ -641,7 +472,6 @@ async def test_load_environment_async_resolves_local_akv_reference(self, file_na env_akv_ref=None, env_files=[env_file], env_akv_strict=True, - env_akv_write_env=False, silent=True, ) @@ -667,7 +497,6 @@ async def test_load_environment_async_does_not_fetch_local_reference_that_loses_ env_akv_ref=None, env_files=[env_file], env_akv_strict=True, - env_akv_write_env=False, silent=True, ) @@ -689,7 +518,6 @@ async def test_load_environment_async_strict_rejects_malformed_local_akv_referen env_akv_ref=None, env_files=[env_file], env_akv_strict=True, - env_akv_write_env=False, silent=True, ) @@ -710,7 +538,6 @@ async def test_load_environment_async_non_strict_skips_malformed_local_akv_refer env_akv_ref=None, env_files=[env_local_file], env_akv_strict=False, - env_akv_write_env=False, silent=False, ) @@ -751,7 +578,6 @@ async def test_non_strict_falls_back_to_valid_reference_after_malformed_override env_akv_ref=None, env_files=env_files, env_akv_strict=False, - env_akv_write_env=False, silent=True, ) @@ -785,7 +611,6 @@ async def test_non_strict_malformed_override_uses_literal_or_no_fallback(self, f env_akv_ref=None, env_files=env_files, env_akv_strict=False, - env_akv_write_env=False, silent=True, ) @@ -810,7 +635,6 @@ async def test_strict_malformed_override_does_not_resolve_valid_fallback(self): env_akv_ref=None, env_files=[ordinary_file, local_file], env_akv_strict=True, - env_akv_write_env=False, silent=True, ) @@ -836,7 +660,6 @@ async def test_highest_valid_local_reference_wins(self): env_akv_ref=None, env_files=[ordinary_file, local_file], env_akv_strict=True, - env_akv_write_env=False, silent=True, ) @@ -866,7 +689,6 @@ async def test_non_strict_resolves_interpolated_fallback_candidate(self): env_akv_ref=None, env_files=[ordinary_file, local_file], env_akv_strict=False, - env_akv_write_env=False, silent=True, ) @@ -895,7 +717,6 @@ async def test_load_environment_async_non_strict_still_raises_for_missing_local_ env_akv_ref=None, env_files=[env_file], env_akv_strict=False, - env_akv_write_env=False, silent=True, ) @@ -1051,13 +872,13 @@ 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): + async def test_fetch_akv_document_async_rejects_non_azure_host_before_authentication(self): with ( mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, mock.patch("pyrit.setup.environment_loading._create_akv_secret_client") as mock_create_client, pytest.raises(KeyVaultInitializationException, match="attacker.example"), ): - await _load_env_from_akv_async( + await _fetch_akv_document_async( secret_url="https://attacker.example/secrets/bootstrap", silent=True, ) @@ -1065,7 +886,7 @@ async def test_load_env_from_akv_async_rejects_non_azure_host_before_authenticat 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): + async def test_fetch_akv_document_async_returns_validated_document(self): credential, client = _create_mock_akv_clients() root_document = ( "DIRECT=from-bootstrap\n" @@ -1075,14 +896,7 @@ async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_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 ( @@ -1091,16 +905,11 @@ async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_secret mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, mock.patch("pyrit.setup.environment_loading._print_msg") as mock_print_msg, ): - await _load_env_from_akv_async(secret_url=secret_url, silent=True) + document = await _fetch_akv_document_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" + assert document.content == root_document + assert document.vault_url == "https://myvault.vault.azure.net" + assert os.environ == {"SOURCE_VALUE": "ambient-value"} mock_credential_cls.assert_called_once_with() _assert_mock_akv_client_created( @@ -1108,70 +917,14 @@ async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_secret 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), - ] + 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() - @pytest.mark.parametrize( - ("document", "expected_values", "expected_child_fetches"), - [ - ( - "A=kv:https://myvault.vault.azure.net/secrets/key\nB=${A}\nA=literal\n", - {"A": "literal", "B": "resolved-key"}, - 1, - ), - ( - "A=literal\nB=${A}\nA=kv:https://myvault.vault.azure.net/secrets/key\n", - {"A": "resolved-key", "B": "literal"}, - 1, - ), - ( - "A=kv:https://myvault.vault.azure.net/secrets/key\nB=${A}\nC=${B}\nB=literal\n", - {"A": "resolved-key", "B": "literal", "C": "resolved-key"}, - 2, - ), - ], - ) - async def test_debug_document_matches_runtime_for_interpolated_reference_assignments( - self, document, expected_values, expected_child_fetches - ): - credential, client = _create_mock_akv_clients() - client.get_secret = mock.AsyncMock( - side_effect=[types.SimpleNamespace(value=document)] - + [types.SimpleNamespace(value="resolved-key")] * expected_child_fetches - ) - - 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), - ): - rendered_document = await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - resolve_references_for_output=True, - ) - runtime_values = {name: os.environ[name] for name in expected_values} - - with mock.patch.dict(os.environ, {}, clear=True): - reloaded_values = dict(dotenv_values(stream=io.StringIO(rendered_document), interpolate=True)) - - assert runtime_values == expected_values - assert reloaded_values == expected_values - assert ( - client.get_secret.await_args_list - == [mock.call("bootstrap", version=None)] + [mock.call("key", version=None)] * expected_child_fetches - ) - - async def test_load_env_from_akv_async_preserves_process_values_without_fetching_overridden_child(self): + async def test_runtime_preserves_process_values_without_fetching_overridden_akv_child(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock( return_value=types.SimpleNamespace( @@ -1189,14 +942,19 @@ async def test_load_env_from_akv_async_preserves_process_values_without_fetching 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=secret_url, silent=True) + await load_environment_async( + env_akv_ref=[secret_url], + env_files=[], + env_akv_strict=True, + silent=True, + ) assert os.environ["DIRECT"] == "from-process" assert os.environ["FROM_KV"] == "process-key" client.get_secret.assert_awaited_once_with("bootstrap", version=None) - async def test_load_env_from_akv_async_rejects_short_secret_name(self): + async def test_runtime_rejects_short_akv_secret_name(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="API_KEY=kv:api-key")) @@ -1206,8 +964,10 @@ async def test_load_env_from_akv_async_rejects_short_secret_name(self): 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", + await load_environment_async( + env_akv_ref=["https://myvault.vault.azure.net/secrets/bootstrap"], + env_files=[], + env_akv_strict=True, silent=True, ) @@ -1218,7 +978,7 @@ async def test_load_env_from_akv_async_rejects_short_secret_name(self): "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): + async def test_runtime_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}")) @@ -1228,12 +988,14 @@ async def test_load_env_from_akv_async_rejects_cross_vault_reference(self, refer 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", + await load_environment_async( + env_akv_ref=["https://myvault.vault.azure.net/secrets/bootstrap"], + env_files=[], + env_akv_strict=True, silent=True, ) - async def test_load_env_from_akv_async_empty_secret_raises(self): + async def test_fetch_akv_document_async_empty_secret_raises(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=None)) @@ -1242,7 +1004,7 @@ async def test_load_env_from_akv_async_empty_secret_raises(self): mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), pytest.raises(ValueError, match="has no value"), ): - await _load_env_from_akv_async( + await _fetch_akv_document_async( secret_url="https://myvault.vault.azure.net/secrets/my-secret", silent=True, ) @@ -1250,7 +1012,7 @@ async def test_load_env_from_akv_async_empty_secret_raises(self): credential.__aexit__.assert_awaited_once() client.__aexit__.assert_awaited_once() - async def test_load_env_from_akv_async_without_entries_raises(self): + async def test_fetch_akv_document_async_without_entries_raises(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="# comments only\n")) @@ -1259,7 +1021,7 @@ async def test_load_env_from_akv_async_without_entries_raises(self): mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), pytest.raises(ValueError, match="contains no environment entries"), ): - await _load_env_from_akv_async( + await _fetch_akv_document_async( secret_url="https://myvault.vault.azure.net/secrets/my-secret", silent=True, ) @@ -1274,7 +1036,7 @@ async def test_load_env_from_akv_async_without_entries_raises(self): ("MISSING_VALUE\n", "variables without values: MISSING_VALUE"), ], ) - async def test_load_env_from_akv_async_rejects_non_assignments(self, document, error): + async def test_fetch_akv_document_async_rejects_non_assignments(self, document, error): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) @@ -1284,7 +1046,7 @@ async def test_load_env_from_akv_async_rejects_non_assignments(self, document, e mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), pytest.raises(ValueError, match=error), ): - await _load_env_from_akv_async( + await _fetch_akv_document_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", silent=True, ) @@ -1292,7 +1054,7 @@ 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): + async def test_fetch_akv_document_async_wraps_malformed_bootstrap(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="=malformed")) @@ -1301,14 +1063,14 @@ async def test_load_env_from_akv_async_wraps_malformed_bootstrap(self): 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( + await _fetch_akv_document_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): + async def test_runtime_wraps_missing_child_secret(self): credential, client = _create_mock_akv_clients() missing_error = ResourceNotFoundError(message="Secret was not found") client.get_secret = mock.AsyncMock( @@ -1324,14 +1086,16 @@ async def test_load_env_from_akv_async_wraps_missing_child_secret(self): 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", + await load_environment_async( + env_akv_ref=["https://myvault.vault.azure.net/secrets/bootstrap"], + env_files=[], + env_akv_strict=True, silent=True, ) assert exc_info.value.__cause__ is missing_error - async def test_load_env_from_akv_async_allows_empty_assignment(self): + async def test_runtime_allows_empty_assignment(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="EMPTY=")) @@ -1340,14 +1104,16 @@ 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( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + await load_environment_async( + env_akv_ref=["https://myvault.vault.azure.net/secrets/bootstrap"], + env_files=[], + env_akv_strict=True, silent=True, ) assert os.environ["EMPTY"] == "" - async def test_load_env_from_akv_async_allows_empty_child_secret(self): + async def test_runtime_allows_empty_child_secret(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock( side_effect=[ @@ -1361,15 +1127,17 @@ async def test_load_env_from_akv_async_allows_empty_child_secret(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( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + await load_environment_async( + env_akv_ref=["https://myvault.vault.azure.net/secrets/bootstrap"], + env_files=[], + env_akv_strict=True, 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): + async def test_fetch_akv_document_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)) @@ -1380,14 +1148,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.environment_loading"), ): - await _load_env_from_akv_async( + fetched_document = await _fetch_akv_document_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 fetched_document.content == "GOOD=resolved\nOTHER=also-resolved" + assert os.environ == {} output = capsys.readouterr().out assert "WARNING: AKV environment document contains invalid entries that will be skipped" in output @@ -1396,36 +1164,7 @@ async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entrie assert "GOOD" not in caplog.text assert "resolved" not in caplog.text - async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_reference(self, caplog, capsys): - credential, client = _create_mock_akv_clients() - document = "GOOD=resolved\nBAD=kv:short-name\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.environment_loading"), - ): - resolved_document = await _load_env_from_akv_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - strict=False, - silent=False, - resolve_references_for_output=True, - ) - - assert os.environ["GOOD"] == "resolved" - assert os.environ["OTHER"] == "also-resolved" - assert "BAD" not in os.environ - - assert "BAD=" not in resolved_document - assert ( - "WARNING: Invalid AKV reference for environment variable 'BAD' will be skipped" in capsys.readouterr().out - ) - assert "BAD" in caplog.text - client.get_secret.assert_awaited_once_with("bootstrap", version=None) - - async def test_load_env_from_akv_async_non_strict_silent_logs_warning(self, caplog, capsys): + async def test_fetch_akv_document_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")) @@ -1435,7 +1174,7 @@ async def test_load_env_from_akv_async_non_strict_silent_logs_warning(self, capl mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"), ): - await _load_env_from_akv_async( + await _fetch_akv_document_async( secret_url="https://myvault.vault.azure.net/secrets/bootstrap", strict=False, silent=True, @@ -1444,7 +1183,7 @@ 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_child_failure_keeps_loaded_bootstrap_values(self): + async def test_runtime_child_failure_keeps_loaded_bootstrap_values(self): credential, client = _create_mock_akv_clients() client.get_secret = mock.AsyncMock( side_effect=[ @@ -1461,8 +1200,10 @@ async def test_load_env_from_akv_async_child_failure_keeps_loaded_bootstrap_valu 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", + await load_environment_async( + env_akv_ref=["https://myvault.vault.azure.net/secrets/bootstrap"], + env_files=[], + env_akv_strict=True, silent=True, ) diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 884383ee9f..57c6eb349e 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -12,6 +12,7 @@ from pyrit.common.singleton import Singleton from pyrit.registry import InitializerRegistry from pyrit.setup import IN_MEMORY, initialize_pyrit_async +from pyrit.setup.environment_loading import _AkvEnvironmentDocument class TestLoadInitializersFromScripts: @@ -119,7 +120,7 @@ def setup_method(self) -> None: reset_default_values() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.environment_loading.load_environment_files", return_value=False) + @mock.patch("pyrit.setup.environment_loading._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) @@ -128,7 +129,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.environment_loading.load_environment_files", return_value=False) + @mock.patch("pyrit.setup.environment_loading._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: @@ -158,15 +159,15 @@ async def initialize_async(self) -> None: finally: os.unlink(script_path) - @mock.patch("pyrit.setup.environment_loading.load_environment_files", return_value=False) + @mock.patch("pyrit.setup.environment_loading._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.environment_loading.load_environment_files", return_value=False) - @mock.patch("pyrit.setup.environment_loading._load_env_from_akv_async", new_callable=mock.AsyncMock) + @mock.patch("pyrit.setup.environment_loading._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.environment_loading._fetch_akv_document_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 = [ @@ -174,22 +175,28 @@ async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, m "https://vault.vault.azure.net/secrets/second/version", ] - mock_load_akv.return_value = None + mock_load_akv.side_effect = [ + _AkvEnvironmentDocument(content="FIRST=one\n", vault_url="https://vault.vault.azure.net"), + _AkvEnvironmentDocument(content="SECOND=two\n", vault_url="https://vault.vault.azure.net"), + ] - with mock.patch("pyrit.setup.environment_loading._warn_about_akv_environment_files") as mock_warn: - 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=refs, + env_files=[], + load_defaults=False, + ) assert mock_load_akv.await_args_list == [ - mock.call(secret_url=refs[0], strict=True, silent=False, resolve_references_for_output=False), - mock.call(secret_url=refs[1], strict=True, silent=False, resolve_references_for_output=False), + mock.call(secret_url=refs[0], strict=True, silent=False), + mock.call(secret_url=refs[1], strict=True, silent=False), ] - mock_warn.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.environment_loading.load_environment_files", return_value=False) - @mock.patch("pyrit.setup.environment_loading._load_env_from_akv_async", new_callable=mock.AsyncMock) + @mock.patch("pyrit.setup.environment_loading._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.environment_loading._fetch_akv_document_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 ): @@ -209,33 +216,22 @@ async def test_initialize_rejects_invalid_env_akv_ref(self, env_akv_ref): load_defaults=False, ) - @pytest.mark.parametrize("option_name", ["env_akv_strict", "env_akv_write_env"]) @pytest.mark.parametrize("invalid_value", ["false", "true", 0, 1, None, [], {}]) - async def test_initialize_rejects_non_boolean_akv_options_before_loading(self, option_name, invalid_value): + async def test_initialize_rejects_non_boolean_env_akv_strict_before_loading(self, invalid_value): with mock.patch( "pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock ) as mock_load_environment: - with pytest.raises(TypeError, match=rf"{option_name} must be a bool"): - if option_name == "env_akv_strict": - await initialize_pyrit_async( - memory_db_type=IN_MEMORY, - env_akv_strict=invalid_value, # type: ignore[arg-type] - load_defaults=False, - ) - else: - await initialize_pyrit_async( - memory_db_type=IN_MEMORY, - env_akv_write_env=invalid_value, # type: ignore[arg-type] - load_defaults=False, - ) + with pytest.raises(TypeError, match=r"env_akv_strict must be a bool"): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_strict=invalid_value, # type: ignore[arg-type] + load_defaults=False, + ) mock_load_environment.assert_not_awaited() - @pytest.mark.parametrize( - ("env_akv_strict", "env_akv_write_env"), - [(True, False), (False, True)], - ) - async def test_initialize_forwards_boolean_akv_options(self, env_akv_strict, env_akv_write_env): + @pytest.mark.parametrize("env_akv_strict", [True, False]) + async def test_initialize_forwards_env_akv_strict(self, env_akv_strict): with ( mock.patch( "pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock @@ -245,14 +241,12 @@ async def test_initialize_forwards_boolean_akv_options(self, env_akv_strict, env await initialize_pyrit_async( memory_db_type=IN_MEMORY, env_akv_strict=env_akv_strict, - env_akv_write_env=env_akv_write_env, load_defaults=False, ) await_args = mock_load_environment.await_args assert await_args is not None assert await_args.kwargs["env_akv_strict"] is env_akv_strict - assert await_args.kwargs["env_akv_write_env"] is env_akv_write_env @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") async def test_initialize_keeps_akv_values_when_local_file_loading_fails(self, mock_set_memory): @@ -261,11 +255,13 @@ 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.environment_loading._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.environment_loading._load_env_from_akv_async", + "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock, - side_effect=lambda **_: os.environ.update({"FROM_AKV": "resolved"}), + return_value=_AkvEnvironmentDocument( + content="FROM_AKV=resolved\n", + vault_url="https://vault.vault.azure.net", + ), ), pytest.raises(ValueError, match="Environment file not found"), ): @@ -289,11 +285,13 @@ 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.environment_loading._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.environment_loading._load_env_from_akv_async", + "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock, - side_effect=lambda **_: os.environ.update({"BASE": "akv", "ONLY_AKV": "shared"}), + return_value=_AkvEnvironmentDocument( + content="BASE=akv\nONLY_AKV=shared\n", + vault_url="https://vault.vault.azure.net", + ), ), ): await initialize_pyrit_async( @@ -320,11 +318,14 @@ 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.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - mock.patch("pyrit.setup.environment_loading._warn_about_akv_environment_files"), + mock.patch("pyrit.setup.environment_loading._warn_about_dotenv_file"), mock.patch( - "pyrit.setup.environment_loading._load_env_from_akv_async", + "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock, - side_effect=lambda **_: os.environ.update({"VALUE": "akv"}), + return_value=_AkvEnvironmentDocument( + content="VALUE=akv\n", + vault_url="https://vault.vault.azure.net", + ), ), ): await initialize_pyrit_async( @@ -363,11 +364,13 @@ 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.environment_loading._warn_about_akv_environment_files"), mock.patch( - "pyrit.setup.environment_loading._load_env_from_akv_async", + "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock, - side_effect=lambda **_: os.environ.update(bootstrap_environment), + return_value=_AkvEnvironmentDocument( + content="".join(f"{name}={value}\n" for name, value in bootstrap_environment.items()), + vault_url="https://vault.vault.azure.net", + ), ), mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("pyrit.setup.environment_loading._create_akv_secret_client", return_value=client), @@ -415,7 +418,7 @@ def setup_method(self) -> None: """Clear default values before each test.""" reset_default_values() - @mock.patch("pyrit.setup.environment_loading.load_environment_files", return_value=True) + @mock.patch("pyrit.setup.environment_loading._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) @@ -423,7 +426,7 @@ async def test_initialize_silent_produces_no_output(self, mock_load_env, capsys) captured = capsys.readouterr() assert captured.out == "" - @mock.patch("pyrit.setup.environment_loading.load_environment_files", return_value=True) + @mock.patch("pyrit.setup.environment_loading._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) From 8f7d5e6c62222785c6670a6ffce662fd37dff775 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 21 Aug 2026 11:09:35 -0400 Subject: [PATCH 25/28] FEAT: PYTHON_DOTENV_DISABLED implementation --- .gitignore | 1 + doc/getting_started/pyrit_conf.md | 2 + pyrit/setup/environment_loading.py | 6 +++ tests/unit/setup/test_environment_loading.py | 51 ++++++++++++++++++++ 4 files changed, 60 insertions(+) diff --git a/.gitignore b/.gitignore index f92df639f0..6b6badb88a 100644 --- a/.gitignore +++ b/.gitignore @@ -168,6 +168,7 @@ cython_debug/ # PyRIT secrets file .env +.env_akv .pyrit_cache/ # Cache for generating docs diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 0f98f7a6b5..bf3b39eb39 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -175,6 +175,8 @@ Environment loading preserves the historical non-transactional dotenv behavior. 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. +`PYTHON_DOTENV_DISABLED` disables the complete PyRIT environment-loading step using python-dotenv's accepted true values (`1`, `true`, `t`, `yes`, and `y`, case-insensitive). When enabled, default discovery and configured `env_akv_ref` or `env_files` sources are skipped. Existing process environment variables remain unchanged. + ### `env_akv_ref` Ordered Azure Key Vault secret URLs used to obtain bootstrap environment documents. This is the canonical configuration path. Each secret value contains dotenv-formatted entries. Authentication uses `DefaultAzureCredential`. diff --git a/pyrit/setup/environment_loading.py b/pyrit/setup/environment_loading.py index 7bc74d22d8..0c68102249 100644 --- a/pyrit/setup/environment_loading.py +++ b/pyrit/setup/environment_loading.py @@ -394,6 +394,9 @@ def _key_vault_initialization_error(*, message: str, error: Exception) -> KeyVau """ Create a contextual Key Vault exception without losing the original cause. + An upstream HTTP status is preserved when available. Failures without an + HTTP response use PyRIT's generic 500 status for opaque internal failures. + Returns: KeyVaultInitializationException: Wrapped contextual exception. """ @@ -526,6 +529,9 @@ async def load_environment_async( Raises: ValueError: If a configured source or reference is invalid. """ + if os.environ.get("PYTHON_DOTENV_DISABLED", "").casefold() in {"1", "true", "t", "yes", "y"}: + return + if isinstance(env_akv_ref, str): raise ValueError("env_akv_ref must be a sequence of Azure Key Vault secret URLs.") assignment_candidates: dict[str, list[_EnvironmentValueCandidate]] = {} diff --git a/tests/unit/setup/test_environment_loading.py b/tests/unit/setup/test_environment_loading.py index 016c980b2b..031b180bc9 100644 --- a/tests/unit/setup/test_environment_loading.py +++ b/tests/unit/setup/test_environment_loading.py @@ -280,6 +280,57 @@ async def test_load_environment_files_honors_python_dotenv_disabled(self): assert loaded is True assert "DISABLED_VALUE" not in os.environ + async def test_load_environment_async_skips_sources_when_python_dotenv_disabled(self): + with ( + mock.patch.dict(os.environ, {"PYTHON_DOTENV_DISABLED": "true"}, clear=True), + mock.patch( + "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock + ) as mock_fetch_akv, + mock.patch("pyrit.setup.environment_loading._load_environment_files") as mock_load_files, + ): + await load_environment_async( + env_akv_ref=None, + env_files=None, + env_akv_strict=True, + silent=True, + ) + + mock_fetch_akv.assert_not_awaited() + mock_load_files.assert_not_called() + + @pytest.mark.parametrize( + ("env_akv_ref", "env_files"), + [ + (["https://vault.vault.azure.net/secrets/bootstrap"], None), + (None, [pathlib.Path("configured.env")]), + ], + ) + async def test_load_environment_async_skips_configured_sources_when_python_dotenv_disabled( + self, env_akv_ref, env_files + ): + with ( + mock.patch.dict( + os.environ, + {"PYTHON_DOTENV_DISABLED": "true", "AMBIENT_VALUE": "preserved"}, + clear=True, + ), + mock.patch( + "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock + ) as mock_fetch_akv, + mock.patch("pyrit.setup.environment_loading._load_environment_files") as mock_load_files, + ): + await load_environment_async( + env_akv_ref=env_akv_ref, + env_files=env_files, + env_akv_strict=True, + silent=True, + ) + + assert os.environ["AMBIENT_VALUE"] == "preserved" + + mock_fetch_akv.assert_not_awaited() + mock_load_files.assert_not_called() + async def test_runtime_does_not_fetch_akv_reference_overridden_by_env_local(self): credential, client = _create_mock_akv_clients() bootstrap_document = "API_KEY=kv:https://vault.vault.azure.net/secrets/api-key\nAKV_ONLY=akv\n" From ddc59838d77d249d1234f44038dd2e5d2624049a Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 21 Aug 2026 15:51:50 -0400 Subject: [PATCH 26/28] FIX: Simplifications --- .azuredevops/test-job-template.yml | 32 +- .pyrit_conf_example | 1 + doc/code/executor/gcg/1_gcg_azure_ml.ipynb | 4 +- doc/code/executor/gcg/1_gcg_azure_ml.py | 4 +- doc/getting_started/pyrit_conf.md | 19 +- infra/env.demo.template | 16 +- pyrit/setup/configuration_loader.py | 6 +- pyrit/setup/environment_loading.py | 358 +++++--------- pyrit/setup/initialization.py | 12 +- .../promptgen/gcg/test_gcg_aml_e2e.py | 4 +- .../test_export_akv_environment.py | 72 ++- .../promptgen/gcg/test_data_and_config.py | 4 +- tests/unit/setup/test_configuration_loader.py | 19 +- tests/unit/setup/test_environment_loading.py | 450 +++++++----------- tests/unit/setup/test_initialization.py | 248 ++-------- tests/unit/setup/test_targets_initializer.py | 48 -- 16 files changed, 407 insertions(+), 890 deletions(-) diff --git a/.azuredevops/test-job-template.yml b/.azuredevops/test-job-template.yml index bd07c9370e..65905af5cc 100644 --- a/.azuredevops/test-job-template.yml +++ b/.azuredevops/test-job-template.yml @@ -26,40 +26,24 @@ jobs: versionSpec: '3.12' addToPath: true - bash: | - install -d -m 700 ~/.pyrit + mkdir -p ~/.pyrit displayName: "Create PyRIT configuration directory" name: create_pyrit_dir - task: AzureKeyVault@2 - displayName: Azure Key Vault - retrieve environment secrets + displayName: Azure Key Vault - retrieve .env file secret inputs: azureSubscription: 'integration-test-service-connection' KeyVaultName: 'pyrit-environment' SecretsFilter: 'env-global' RunAsPreJob: false - bash: | - python - <<'PY' - import os - import pathlib - import tempfile - - secret = os.environ.get("PYRIT_TEST_SECRET") + python -c " + import os; + secret = os.environ.get('PYRIT_TEST_SECRET'); if not secret: - raise ValueError("PYRIT_TEST_SECRET is not set") - - env_file = pathlib.Path.home() / ".pyrit" / ".env" - file_descriptor, temporary_name = tempfile.mkstemp(prefix=".env.", suffix=".tmp", dir=env_file.parent) - temporary_file = pathlib.Path(temporary_name) - try: - os.fchmod(file_descriptor, 0o600) - with os.fdopen(file_descriptor, "w", encoding="utf-8", newline="") as stream: - file_descriptor = -1 - stream.write(secret) - os.replace(temporary_file, env_file) - finally: - if file_descriptor >= 0: - os.close(file_descriptor) - temporary_file.unlink(missing_ok=True) - PY + raise ValueError('PYRIT_TEST_SECRET is not set'); + with open(os.path.expanduser('~/.pyrit/.env'), 'w') as file: + file.write(secret)" env: PYRIT_TEST_SECRET: $(env-global) name: create_env_file diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 033b90c274..8a129f06e6 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -83,6 +83,7 @@ operation: op_trash_panda # ------------------------- # Azure Key Vault is the canonical source for shared and deployed configuration. # See doc/getting_started/pyrit_conf.md for loading order, references, and migration guidance. +# The list may contain at most one bootstrap secret URL. # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env # env_akv_strict: true diff --git a/doc/code/executor/gcg/1_gcg_azure_ml.ipynb b/doc/code/executor/gcg/1_gcg_azure_ml.ipynb index f2053f9fb2..a4c7b20040 100644 --- a/doc/code/executor/gcg/1_gcg_azure_ml.ipynb +++ b/doc/code/executor/gcg/1_gcg_azure_ml.ipynb @@ -57,9 +57,9 @@ "source": [ "import os\n", "\n", - "from pyrit.setup.environment_loading import load_environment_files\n", + "from pyrit.setup.initialization import _load_environment_files\n", "\n", - "load_environment_files(env_files=None)\n", + "_load_environment_files(env_files=None)\n", "\n", "subscription_id = os.environ.get(\"AZURE_ML_SUBSCRIPTION_ID\")\n", "resource_group = os.environ.get(\"AZURE_ML_RESOURCE_GROUP\")\n", diff --git a/doc/code/executor/gcg/1_gcg_azure_ml.py b/doc/code/executor/gcg/1_gcg_azure_ml.py index 0e4b5fe40c..c3c559f18c 100644 --- a/doc/code/executor/gcg/1_gcg_azure_ml.py +++ b/doc/code/executor/gcg/1_gcg_azure_ml.py @@ -29,9 +29,9 @@ # %% import os -from pyrit.setup.environment_loading import load_environment_files +from pyrit.setup.initialization import _load_environment_files -load_environment_files(env_files=None) +_load_environment_files(env_files=None) subscription_id = os.environ.get("AZURE_ML_SUBSCRIPTION_ID") resource_group = os.environ.get("AZURE_ML_RESOURCE_GROUP") diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index bf3b39eb39..0692bf286f 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -36,7 +36,7 @@ See [Populating Secrets](./populating_secrets.md) for provider-specific variable PyRIT loads environment sources in this order: 1. Existing process environment variables. -2. Key Vault bootstrap documents, auto-discovered `.env`, or explicit `env_files`. These sources fill only missing values. +2. A Key Vault bootstrap document, auto-discovered `.env`, or explicit `env_files`. These sources fill only missing values. 3. Files named `.env.local`. These are the only dotenv sources that override existing values. When `env_akv_ref` is configured, PyRIT ignores an auto-discovered `~/.pyrit/.env`, emits a security warning, and still loads `~/.pyrit/.env.local`. Explicit `env_files` are never blocked based on their filename or location. @@ -167,11 +167,11 @@ env_files: Local files use standard python-dotenv parsing and `${NAME}` interpolation. Ordinary files fill missing values; any file whose basename is `.env.local` overrides existing values. Explicit files load in their listed order. -Complete-value `kv:`, `akv:`, `azure_key_vault:`, and `env_akv_ref:` references resolve in local files as well as remote bootstrap documents. Local references may use any validated supported Key Vault URL; remote child references must remain in the bootstrap document's vault. A local assignment that loses to an existing value does not fetch its secret. +Complete-value `kv:`, `akv:`, `azure_key_vault:`, and `env_akv_ref:` references resolve in local files as well as the remote bootstrap document. Local references may use any validated supported Key Vault URL; remote child references must remain in the bootstrap document's vault. A local assignment that loses to an existing value does not fetch its secret. Ordinary malformed dotenv lines retain python-dotenv's permissive behavior. `env_akv_strict` controls malformed Key Vault reference syntax in all sources: strict mode raises; non-strict mode warns and skips that assignment. Authentication, authorization, transport, missing-secret, and missing-value failures always raise. -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. +Environment loading preserves the historical non-transactional dotenv behavior. The bootstrap document and each local file update `os.environ` as they load. 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. @@ -179,15 +179,14 @@ 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. This is the canonical configuration path. Each secret value contains dotenv-formatted entries. Authentication uses `DefaultAzureCredential`. +List-shaped Azure Key Vault bootstrap configuration. It may be omitted, empty, or contain one secret URL; multiple bootstrap URLs are rejected. The secret value contains dotenv-formatted entries, and authentication uses `DefaultAzureCredential`. ```yaml env_akv_ref: - https://my-vault.vault.azure.net/secrets/shared-pyrit-env - - https://my-vault.vault.azure.net/secrets/team-pyrit-env ``` -Bootstrap documents load in list order and fill values missing from the process environment. 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: +The bootstrap document fills values missing from the process environment and uses native dotenv interpolation against the process environment and assignments already parsed. It 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" @@ -214,7 +213,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" ``` -Bootstrap documents stay in memory by default. Use `.env.local` when an intentional local override is required. +The bootstrap document stays in memory. Use `.env.local` when an intentional local override is required. ### `env_akv_strict` @@ -226,7 +225,7 @@ env_akv_strict: false In strict mode, malformed bootstrap dotenv lines, valueless bootstrap entries, and malformed Key Vault references stop initialization. Empty assignments such as `OPTIONAL_VALUE=` remain valid. With `env_akv_strict: false`, PyRIT warns and skips malformed bootstrap entries and malformed reference assignments without logging secret values. -Non-strict mode does not suppress operational failures. Missing secrets, authentication, authorization, transport errors, and bootstrap documents with no valid assignments still stop initialization. Loading remains non-transactional, so earlier successful assignments remain. +Non-strict mode does not suppress operational failures. Missing secrets, authentication, authorization, transport errors, and a bootstrap document with no valid assignments still stop initialization. Loading remains non-transactional, so earlier successful assignments remain. 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. @@ -239,7 +238,7 @@ python -m build_scripts.export_akv_environment ` --secret-url https://my-vault.vault.azure.net/secrets/my-pyrit-env ``` -Repeat `--secret-url` to preserve multiple bootstrap documents in load order. The helper writes `~/.pyrit/.env_akv` by default. This file is not auto-loaded by PyRIT and excludes process, `.env`, explicit file, and `.env.local` values. +The helper accepts exactly one `--secret-url` and writes `~/.pyrit/.env_akv` by default. This file is not auto-loaded by PyRIT and excludes process, `.env`, explicit file, and `.env.local` values. The helper resolves child-secret references and writes plaintext secrets with owner-only permissions where supported. It refuses to overwrite an existing path; remove the file when debugging is complete. Use `--output` to select a different path and `--non-strict` to skip malformed entries or references with warnings. @@ -364,7 +363,7 @@ initializers: # initialization_scripts: # - /path/to/my_custom_initializer.py -# Canonical: ordered Azure Key Vault bootstrap environment documents +# Canonical: zero or one Azure Key Vault bootstrap environment document # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env # env_akv_strict: true diff --git a/infra/env.demo.template b/infra/env.demo.template index 77f2af2b56..4c5e7dac7e 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) ─── -# 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 +# 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 # ─── TTS Target (optional — for text-to-speech demos) ─── -# 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 +# 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 # ─── 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/configuration_loader.py b/pyrit/setup/configuration_loader.py index e61ae50db9..d78ae8cca4 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -98,7 +98,7 @@ class ConfigurationLoader(YamlLoadable): env_files: List of environment file paths to load. None means auto-discover supported ``.env`` and ``.env.local``; [] means "load nothing". - env_akv_ref: Ordered list of Key Vault bootstrap secret URLs. + env_akv_ref: List containing at most one Key Vault bootstrap secret URL. env_akv_strict: Whether malformed or valueless entries in a Key Vault bootstrap document should fail initialization. silent: Whether to suppress initialization messages. @@ -168,6 +168,8 @@ def _validate_env_akv_ref(self) -> None: return if not isinstance(self.env_akv_ref, list): raise ValueError("env_akv_ref must be a list of Azure Key Vault secret URLs.") + if len(self.env_akv_ref) > 1: + raise ValueError("env_akv_ref supports at most one Azure Key Vault bootstrap secret URL.") 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.") @@ -439,7 +441,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 ordered Azure Key Vault bootstrap secret URLs. + env_akv_ref: Override containing at most one Azure Key Vault bootstrap secret URL. env_akv_strict: Override for strict Key Vault bootstrap validation. Returns: diff --git a/pyrit/setup/environment_loading.py b/pyrit/setup/environment_loading.py index 0c68102249..e3122adc9c 100644 --- a/pyrit/setup/environment_loading.py +++ b/pyrit/setup/environment_loading.py @@ -10,7 +10,6 @@ import pathlib import urllib.parse from collections.abc import Mapping, Sequence -from dataclasses import dataclass from io import StringIO from typing import TYPE_CHECKING @@ -37,23 +36,7 @@ _AKV_VAULT_DNS_SUFFIXES = frozenset({"vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"}) _AKV_RETRY_TOTAL = 3 _AKV_RETRY_BACKOFF_FACTOR = 0.8 - - -@dataclass(frozen=True) -class _EnvironmentValueCandidate: - """An ordered environment value and its Key Vault resolution policy.""" - - value: str - resolve_akv_reference: bool - expected_vault_url: str | None = None - - -@dataclass(frozen=True) -class _AkvEnvironmentDocument: - """A validated Key Vault bootstrap document and its source vault.""" - - content: str - vault_url: str +_DOTENV_DISABLED_VALUES = frozenset({"1", "true", "t", "yes", "y"}) def validate_env_akv_strict(*, env_akv_strict: object) -> None: @@ -83,7 +66,8 @@ def load_environment_files( env_files=env_files, silent=silent, include_default_base=include_default_base, - assignment_candidates=None, + ordinary_candidates=None, + override_candidates=None, ) @@ -92,7 +76,8 @@ def _load_environment_files( *, silent: bool, include_default_base: bool, - assignment_candidates: dict[str, list[_EnvironmentValueCandidate]] | None = None, + ordinary_candidates: dict[str, list[tuple[str, str | None]]] | None = None, + override_candidates: dict[str, list[tuple[str, str | None]]] | None = None, ) -> bool: """ Load environment files in the order they are provided. @@ -107,9 +92,8 @@ def _load_environment_files( 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. - assignment_candidates: Optional output mapping containing each applicable - local value in ascending precedence order. Existing process or AKV values - are retained as non-resolvable baseline candidates. + ordinary_candidates: Optional output mapping for non-overriding assignments. + override_candidates: Optional output mapping for ``.env.local`` assignments. Returns: True if at least one environment file was loaded, otherwise False. @@ -117,16 +101,40 @@ def _load_environment_files( 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, - ) + if env_files is None: + selected_files = [] + base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" + local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" + if include_default_base and base_file.exists(): + _warn_about_dotenv_file(env_file=base_file, ignored_for_akv=False, silent=silent) + selected_files.append(base_file) + if local_file.exists(): + selected_files.append(local_file) + if not silent: + message = ( + f"Found default environment files: {[str(file) for file in selected_files]}" + if selected_files + else "No default environment files found. Using system environment variables only." + ) + _print_msg(message, quiet=False, log=True) + else: + selected_files = list(env_files) + if not silent: + _print_msg( + f"Loading custom environment files: {[str(file) for file in selected_files]}", + quiet=False, + log=True, + ) + for env_file in selected_files: + if not env_file.exists(): + raise ValueError(f"Environment file not found: {env_file}") + for env_file in selected_files: loaded = _load_dotenv_source( dotenv_path=env_file, override=env_file.name == ".env.local", - assignment_candidates=assignment_candidates, + ordinary_candidates=ordinary_candidates, + override_candidates=override_candidates, ) if not silent: _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) @@ -137,7 +145,8 @@ def _load_environment_files( def _load_dotenv_source( *, override: bool, - assignment_candidates: dict[str, list[_EnvironmentValueCandidate]] | None, + ordinary_candidates: dict[str, list[tuple[str, str | None]]] | None, + override_candidates: dict[str, list[tuple[str, str | None]]] | None, dotenv_path: pathlib.Path | None = None, document: str | None = None, expected_vault_url: str | None = None, @@ -153,103 +162,28 @@ def _load_dotenv_source( """ if (dotenv_path is None) == (document is None): raise ValueError("Exactly one dotenv_path or document must be provided.") - - if dotenv_path is not None: - assignment_values = DotEnv( - dotenv_path=dotenv_path, - override=override, - interpolate=True, - ).dict() - else: - assignment_values = DotEnv( - dotenv_path=None, - stream=StringIO(document or ""), - override=override, - interpolate=True, - ).dict() - previous_values = {variable_name: os.environ.get(variable_name) for variable_name in assignment_values} - - if dotenv_path is not None: - loaded = dotenv.load_dotenv(dotenv_path=dotenv_path, override=override, interpolate=True) - else: - loaded = dotenv.load_dotenv(stream=StringIO(document or ""), override=override, interpolate=True) - if assignment_candidates is None or not loaded: + if os.environ.get("PYTHON_DOTENV_DISABLED", "").casefold() in _DOTENV_DISABLED_VALUES: + return False + + source = DotEnv( + dotenv_path=dotenv_path, + stream=StringIO(document or "") if document is not None else None, + override=override, + interpolate=True, + ) + assignment_values = source.dict() + loaded = source.set_as_environment_variables() + if ordinary_candidates is None or override_candidates is None or not loaded: return loaded + candidates = override_candidates if override else ordinary_candidates for variable_name, loaded_value in assignment_values.items(): if loaded_value is None: continue - candidates = assignment_candidates.setdefault(variable_name, []) - previous_value = previous_values[variable_name] - if not candidates and previous_value is not None: - candidates.append(_EnvironmentValueCandidate(value=previous_value, resolve_akv_reference=False)) - candidate = _EnvironmentValueCandidate( - value=loaded_value, - resolve_akv_reference=True, - expected_vault_url=expected_vault_url, - ) - if override: - candidates.append(candidate) - elif candidates and not candidates[-1].resolve_akv_reference: - continue - elif candidates: - candidates.insert(0, candidate) - else: - candidates.append(candidate) + candidates.setdefault(variable_name, []).append((loaded_value, expected_vault_url)) return loaded -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(): - _warn_about_dotenv_file(env_file=base_file, ignored_for_akv=False, silent=silent) - 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. @@ -324,28 +258,18 @@ def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: 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) + secret_name, secret_version = path_parts[2], path_parts[3] if len(path_parts) == 4 else None + identifiers = [secret_name] + ([secret_version] if secret_version else []) + if any( + not 1 <= len(identifier) <= 127 + or not all(char.isascii() and (char.isalnum() or char == "-") for char in identifier) + for identifier in identifiers ): 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: "AsyncTokenCredential") -> "SecretClient": """ Create an asynchronous Key Vault client with an explicit retry policy. @@ -366,30 +290,6 @@ def _create_akv_secret_client(*, vault_url: str, credential: "AsyncTokenCredenti return SecretClient(vault_url=vault_url, credential=credential, retry_policy=retry_policy) -async def _fetch_akv_secret_value_async( - *, - client: "SecretClient", - secret_name: str, - secret_version: str | None, - variable_name: str, -) -> str: - """ - Fetch a referenced Key Vault secret value. - - Returns: - str: The secret value, including an empty string. - - Raises: - ValueError: If the referenced secret has no value. - """ - referenced_secret = await client.get_secret(secret_name, version=secret_version) - if referenced_secret.value is None: - raise ValueError( - f"AKV secret '{secret_name}' referenced by environment variable '{variable_name}' has no value." - ) - return referenced_secret.value - - def _key_vault_initialization_error(*, message: str, error: Exception) -> KeyVaultInitializationException: """ Create a contextual Key Vault exception without losing the original cause. @@ -459,7 +359,7 @@ async def _fetch_akv_document_async( secret_url: str, strict: bool = True, silent: bool = False, -) -> _AkvEnvironmentDocument: +) -> tuple[str, str]: """ Fetch and validate one Key Vault bootstrap dotenv document. @@ -475,7 +375,7 @@ async def _fetch_akv_document_async( silent (bool): If True, suppresses print statements. Defaults to False. Returns: - _AkvEnvironmentDocument: Validated document text and source vault metadata. + tuple[str, str]: Validated document text and source vault URL. Raises: ImportError: If ``azure-keyvault-secrets`` is not installed. @@ -499,7 +399,7 @@ async def _fetch_akv_document_async( parsed_environment = dotenv.dotenv_values(stream=StringIO(validated_document), interpolate=False) if not parsed_environment: raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") - return _AkvEnvironmentDocument(content=validated_document, vault_url=vault_url) + return validated_document, vault_url except KeyVaultInitializationException: raise except Exception as error: @@ -529,12 +429,16 @@ async def load_environment_async( Raises: ValueError: If a configured source or reference is invalid. """ - if os.environ.get("PYTHON_DOTENV_DISABLED", "").casefold() in {"1", "true", "t", "yes", "y"}: + if os.environ.get("PYTHON_DOTENV_DISABLED", "").casefold() in _DOTENV_DISABLED_VALUES: return if isinstance(env_akv_ref, str): raise ValueError("env_akv_ref must be a sequence of Azure Key Vault secret URLs.") - assignment_candidates: dict[str, list[_EnvironmentValueCandidate]] = {} + if env_akv_ref is not None and len(env_akv_ref) > 1: + raise ValueError("env_akv_ref supports at most one Azure Key Vault bootstrap secret URL.") + process_environment = dict(os.environ) + ordinary_candidates: dict[str, list[tuple[str, str | None]]] = {} + override_candidates: dict[str, list[tuple[str, str | None]]] = {} 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.") @@ -547,68 +451,56 @@ async def load_environment_async( ignored_for_akv=True, silent=silent, ) - for secret_url in env_akv_ref: - document = await _fetch_akv_document_async( - secret_url=secret_url, - strict=env_akv_strict, - silent=silent, - ) - await asyncio.to_thread( - _load_dotenv_source, - document=document.content, - override=False, - assignment_candidates=assignment_candidates, - expected_vault_url=document.vault_url, - ) + document, vault_url = await _fetch_akv_document_async( + secret_url=env_akv_ref[0], + strict=env_akv_strict, + silent=silent, + ) + await asyncio.to_thread( + _load_dotenv_source, + document=document, + override=False, + ordinary_candidates=ordinary_candidates, + override_candidates=override_candidates, + expected_vault_url=vault_url, + ) await asyncio.to_thread( _load_environment_files, env_files=env_files, silent=silent, include_default_base=not (env_akv_ref and env_files is None), - assignment_candidates=assignment_candidates, + ordinary_candidates=ordinary_candidates, + override_candidates=override_candidates, ) await _resolve_environment_candidates_async( - assignment_candidates=assignment_candidates, + process_environment=process_environment, + ordinary_candidates=ordinary_candidates, + override_candidates=override_candidates, strict=env_akv_strict, silent=silent, ) -def _warn_about_invalid_akv_reference(*, variable_name: str, error: ValueError, silent: bool) -> None: - """Warn that a malformed Key Vault reference assignment is being skipped.""" - message = f"Invalid AKV reference for environment variable '{variable_name}' will be skipped: {error}" - if not silent: - print(f"WARNING: {message}") - logger.warning(message) - - -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 _parse_akv_reference_url( +def _parse_akv_reference( *, - target: str, + value: str, variable_name: str, expected_vault_url: str | None = None, -) -> tuple[str, str, str | None]: +) -> tuple[str, str, str | None] | None: """ - Parse and optionally constrain a complete Key Vault secret reference. + Parse and validate an exact whole-value Key Vault reference. Returns: - tuple[str, str, str | None]: Vault URL, secret name, and optional secret version. + tuple[str, str, str | None] | None: Parsed reference, or None for a literal value. Raises: ValueError: If the reference is malformed or violates the expected vault constraint. """ + prefix, separator, target = value.partition(":") + if not separator or prefix not in _AKV_REFERENCE_PREFIXES: + return None + target = target.strip() if not target.casefold().startswith("https://"): raise ValueError( f"AKV reference for environment variable '{variable_name}' must use a full secret URL, " @@ -627,7 +519,9 @@ def _parse_akv_reference_url( async def _resolve_environment_candidates_async( *, - assignment_candidates: Mapping[str, Sequence[_EnvironmentValueCandidate]], + process_environment: Mapping[str, str], + ordinary_candidates: Mapping[str, Sequence[tuple[str, str | None]]], + override_candidates: Mapping[str, Sequence[tuple[str, str | None]]], strict: bool, silent: bool, ) -> None: @@ -638,20 +532,26 @@ async def _resolve_environment_candidates_async( KeyVaultInitializationException: If strict validation or secret retrieval fails. """ parsed_references: list[tuple[str, str, str, str | None]] = [] - for variable_name, candidates in assignment_candidates.items(): - for candidate in reversed(candidates): - if not candidate.resolve_akv_reference: - os.environ[variable_name] = candidate.value - break - target = _parse_akv_reference(candidate.value) - if target is None: - os.environ[variable_name] = candidate.value + variable_names = ordinary_candidates.keys() | override_candidates.keys() + for variable_name in variable_names: + candidates = [ + (value, vault_url, True) + for value, vault_url in reversed(override_candidates.get(variable_name, ())) + ] + if variable_name in process_environment: + candidates.append((process_environment[variable_name], None, False)) + candidates.extend( + (value, vault_url, True) for value, vault_url in ordinary_candidates.get(variable_name, ()) + ) + for value, expected_vault_url, resolve_reference in candidates: + if not resolve_reference: + os.environ[variable_name] = value break try: - vault_url, secret_name, secret_version = _parse_akv_reference_url( - target=target, + reference = _parse_akv_reference( + value=value, variable_name=variable_name, - expected_vault_url=candidate.expected_vault_url, + expected_vault_url=expected_vault_url, ) except ValueError as error: if strict: @@ -660,13 +560,18 @@ async def _resolve_environment_candidates_async( error=error, ) raise wrapped_error from error - _warn_about_invalid_akv_reference( - variable_name=variable_name, - error=error, - silent=silent, + message = ( + f"Invalid AKV reference for environment variable '{variable_name}' will be skipped: {error}" ) + if not silent: + print(f"WARNING: {message}") + logger.warning(message) continue - os.environ[variable_name] = candidate.value + if reference is None: + os.environ[variable_name] = value + break + vault_url, secret_name, secret_version = reference + os.environ[variable_name] = value parsed_references.append((variable_name, vault_url, secret_name, secret_version)) break else: @@ -688,12 +593,13 @@ async def _resolve_environment_candidates_async( _create_akv_secret_client(vault_url=vault_url, credential=credential) ) clients[vault_url] = client - os.environ[variable_name] = await _fetch_akv_secret_value_async( - client=client, - secret_name=secret_name, - secret_version=secret_version, - variable_name=variable_name, - ) + secret = await client.get_secret(secret_name, version=secret_version) + if secret.value is None: + raise ValueError( + f"AKV secret '{secret_name}' referenced by environment variable " + f"'{variable_name}' has no value." + ) + os.environ[variable_name] = secret.value except KeyVaultInitializationException: raise except Exception as error: diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 0c345aa595..5aef5cfc61 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -7,7 +7,11 @@ from pyrit.common.apply_defaults import reset_default_values from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory -from pyrit.setup.environment_loading import load_environment_async, validate_env_akv_strict +from pyrit.setup.environment_loading import ( + load_environment_async, + load_environment_files as _load_environment_files, + validate_env_akv_strict, +) if TYPE_CHECKING: from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -96,9 +100,9 @@ async def initialize_pyrit_async( env_files (Sequence[pathlib.Path] | None): Optional sequence of environment file paths to load in order. Ordinary files fill missing process values; files named ``.env.local`` override. If omitted, PyRIT auto-discovers supported ``.env`` and ``.env.local`` files. - env_akv_ref (Sequence[str] | None): Optional ordered Azure Key Vault URLs whose secret values - contain bootstrap dotenv documents. Documents fill missing process values and support - complete-value references to scalar secrets. Requires ``azure-keyvault-secrets``. + env_akv_ref (Sequence[str] | None): Optional zero-or-one-item sequence containing an Azure Key Vault + URL whose secret value is a bootstrap dotenv document. The document fills missing process values + and supports complete-value references to scalar secrets. Requires ``azure-keyvault-secrets``. env_akv_strict (bool): If True, reject malformed bootstrap entries and Key Vault reference syntax. If False, warn and skip those entries. Operational Key Vault failures always raise. silent (bool): If True, suppresses print statements about environment file loading and 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 34b67f0cac..83d5078d80 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.environment_loading import load_environment_files # noqa: E402 +from pyrit.setup.initialization import _load_environment_files # noqa: E402 _REQUIRED_ENV_VARS = ( "AZURE_ML_SUBSCRIPTION_ID", @@ -70,7 +70,7 @@ def test_gcg_aml_notebook_runs_to_completion() -> None: MLClient from its namespace, then polls until the job reaches a terminal state and asserts ``Completed``. """ - load_environment_files(env_files=None, silent=True) + _load_environment_files(env_files=None, silent=True) missing = [name for name in _REQUIRED_ENV_VARS if not os.environ.get(name)] if missing: pytest.skip(f"Missing required env vars for GCG AML e2e test: {', '.join(missing)}") diff --git a/tests/unit/build_scripts/test_export_akv_environment.py b/tests/unit/build_scripts/test_export_akv_environment.py index 9b5531f753..24888ceb8b 100644 --- a/tests/unit/build_scripts/test_export_akv_environment.py +++ b/tests/unit/build_scripts/test_export_akv_environment.py @@ -12,7 +12,6 @@ from build_scripts.export_akv_environment import ( DEFAULT_OUTPUT_FILE, - _Document, _render, _serialize, _write_output, @@ -37,26 +36,25 @@ def test_serialize_round_trips_terminal_values(value: str) -> None: def test_render_resolves_akv_only_values() -> None: - document = _Document( - content=( + document = ( + ( "# AKV config\nBASE=bootstrap\nDERIVED=${BASE}\nAPI_KEY=kv:https://vault.vault.azure.net/secrets/api-key\n" ), - vault_url="https://vault.vault.azure.net", + "https://vault.vault.azure.net", ) client = mock.MagicMock() client.get_secret.return_value = SimpleNamespace(value="resolved-key") rendered = _render( - documents=[document], + document=document, credential=mock.MagicMock(), - clients={document.vault_url: client}, + clients={document[1]: client}, strict=True, silent=True, ) values = dotenv.dotenv_values(stream=io.StringIO(rendered), interpolate=True) assert values == {"BASE": "bootstrap", "DERIVED": "bootstrap", "API_KEY": "resolved-key"} - assert "# AKV config" in rendered client.get_secret.assert_called_once_with("api-key", version=None) @@ -83,14 +81,14 @@ def test_render_resolves_akv_only_values() -> None: def test_render_resolves_interpolated_reference_assignments( content: str, expected_values: dict[str, str], expected_child_fetches: int ) -> None: - document = _Document(content=content, vault_url="https://vault.vault.azure.net") + document = (content, "https://vault.vault.azure.net") client = mock.MagicMock() client.get_secret.return_value = SimpleNamespace(value="resolved-key") rendered = _render( - documents=[document], + document=document, credential=mock.MagicMock(), - clients={document.vault_url: client}, + clients={document[1]: client}, strict=True, silent=True, ) @@ -100,42 +98,12 @@ def test_render_resolves_interpolated_reference_assignments( assert client.get_secret.call_args_list == [mock.call("key", version=None)] * expected_child_fetches -def test_render_preserves_first_document_values() -> None: - documents = [ - _Document( - content="SHARED=first\nFIRST_ONLY=first\n", - vault_url="https://vault.vault.azure.net", - ), - _Document( - content="SHARED=second\nSECOND_ONLY=second\n", - vault_url="https://vault.vault.azure.net", - ), - ] - - rendered = _render( - documents=documents, - credential=mock.MagicMock(), - clients={}, - strict=True, - silent=True, - ) - - assert dotenv.dotenv_values(stream=io.StringIO(rendered), interpolate=True) == { - "SHARED": "first", - "FIRST_ONLY": "first", - "SECOND_ONLY": "second", - } - - def test_render_non_strict_warns_and_skips_invalid_reference(caplog: pytest.LogCaptureFixture, capsys) -> None: - document = _Document( - content="GOOD=resolved\nBAD=kv:short-name\nOTHER=also-resolved", - vault_url="https://vault.vault.azure.net", - ) + document = ("GOOD=resolved\nBAD=kv:short-name\nOTHER=also-resolved", "https://vault.vault.azure.net") with caplog.at_level("WARNING", logger="build_scripts.export_akv_environment"): rendered = _render( - documents=[document], + document=document, credential=mock.MagicMock(), clients={}, strict=False, @@ -205,7 +173,7 @@ def test_export_rejects_existing_output_before_fetch(tmp_path: pathlib.Path) -> output_file = tmp_path / ".env_akv" output_file.write_text("ORIGINAL=value\n", encoding="utf-8") - with mock.patch("build_scripts.export_akv_environment._fetch_documents") as mock_fetch: + with mock.patch("build_scripts.export_akv_environment._fetch_document") as mock_fetch: with pytest.raises(ValueError, match="already exists"): export_akv_environment( secret_urls=["https://vault.vault.azure.net/secrets/bootstrap"], @@ -217,6 +185,24 @@ def test_export_rejects_existing_output_before_fetch(tmp_path: pathlib.Path) -> mock_fetch.assert_not_called() +def test_export_rejects_multiple_bootstrap_urls_before_fetch(tmp_path: pathlib.Path) -> None: + with ( + mock.patch("build_scripts.export_akv_environment._fetch_document") as mock_fetch, + pytest.raises(ValueError, match="Only one"), + ): + export_akv_environment( + secret_urls=[ + "https://vault.vault.azure.net/secrets/first", + "https://vault.vault.azure.net/secrets/second", + ], + output_file=tmp_path / ".env_akv", + credential=mock.MagicMock(), + silent=True, + ) + + mock_fetch.assert_not_called() + + def test_write_output_does_not_clobber_existing_file(tmp_path: pathlib.Path) -> None: output_file = tmp_path / ".env_akv" output_file.write_text("ORIGINAL=value\n", encoding="utf-8") diff --git a/tests/unit/executor/promptgen/gcg/test_data_and_config.py b/tests/unit/executor/promptgen/gcg/test_data_and_config.py index adc0eaa9ac..d41ffef2db 100644 --- a/tests/unit/executor/promptgen/gcg/test_data_and_config.py +++ b/tests/unit/executor/promptgen/gcg/test_data_and_config.py @@ -153,7 +153,7 @@ def test_override_falls_back_to_default_basename(self, tmp_path: Path) -> None: class TestMainAsyncCli: """Tests for ``run.py``'s ``--config`` + ``--data`` CLI wrapper around GCGGenerator.execute_async.""" - @patch("pyrit.executor.promptgen.gcg.experiments.run.load_environment_files") + @patch("pyrit.executor.promptgen.gcg.experiments.run._load_environment_files") async def test_raises_when_no_token_anywhere(self, mock_load_env: MagicMock, tmp_path: Path) -> None: config = GCGConfig(models=[GCGModelConfig(name="org/model")]) config_path = tmp_path / "config.json" @@ -166,7 +166,7 @@ async def test_raises_when_no_token_anywhere(self, mock_load_env: MagicMock, tmp with pytest.raises(ValueError, match="No HuggingFace token available"): await _main_async(str(config_path), str(data_path)) - @patch("pyrit.executor.promptgen.gcg.experiments.run.load_environment_files") + @patch("pyrit.executor.promptgen.gcg.experiments.run._load_environment_files") @patch("pyrit.executor.promptgen.gcg.experiments.run.load_goals_and_targets") async def test_passes_loaded_goals_to_generator_and_uses_env_token( self, diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 8d74a209c9..fa0c81be60 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -331,10 +331,7 @@ def testresolve_env_akv_ref_none_returns_none(self): def testresolve_env_akv_ref_returns_configured_values(self): """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", - ] + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] config = ConfigurationLoader(env_akv_ref=refs) assert config.resolve_env_akv_ref() == refs @@ -346,6 +343,15 @@ 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] + def test_env_akv_ref_rejects_multiple_bootstrap_urls(self): + with pytest.raises(ValueError, match="at most one"): + ConfigurationLoader( + env_akv_ref=[ + "https://vault.vault.azure.net/secrets/first", + "https://vault.vault.azure.net/secrets/second", + ] + ) + @pytest.mark.usefixtures("patch_central_database") class TestConfigurationLoaderInitialization: @@ -371,10 +377,7 @@ 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", - ] + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] config = ConfigurationLoader( memory_db_type="in_memory", env_akv_ref=refs, diff --git a/tests/unit/setup/test_environment_loading.py b/tests/unit/setup/test_environment_loading.py index 031b180bc9..d4c776e5b1 100644 --- a/tests/unit/setup/test_environment_loading.py +++ b/tests/unit/setup/test_environment_loading.py @@ -12,9 +12,7 @@ from azure.core.exceptions import ResourceNotFoundError from pyrit.exceptions import KeyVaultInitializationException -from pyrit.setup import IN_MEMORY, initialize_pyrit_async from pyrit.setup.environment_loading import ( - _AkvEnvironmentDocument, _fetch_akv_document_async, _parse_akv_reference, _parse_akv_secret_url, @@ -27,95 +25,70 @@ class TestLoadEnvironmentFiles: """Tests for load_environment_files and the env_files initialization parameter.""" - @mock.patch("pyrit.setup.environment_loading.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.environment_loading.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.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") - async def test_default_env_preserves_process_environment(self, mock_config_path): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - (temp_path / ".env").write_text("VAR=legacy\nLEGACY_ONLY=legacy") - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - with mock.patch.dict(os.environ, {"VAR": "process"}, clear=True): - loaded = load_environment_files(env_files=None, silent=True) - - assert loaded is True - assert os.environ["VAR"] == "process" - assert os.environ["LEGACY_ONLY"] == "legacy" - - @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") - async def test_default_env_local_overrides_process_environment_and_env(self, mock_config_path): - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - (temp_path / ".env").write_text("VAR=legacy") - (temp_path / ".env.local").write_text("VAR=local") - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - with mock.patch.dict(os.environ, {"VAR": "process"}, clear=True): - loaded = load_environment_files(env_files=None, silent=True) - - assert loaded is True - assert os.environ["VAR"] == "local" - - @mock.patch("pyrit.setup.environment_loading.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" + @pytest.mark.parametrize( + ("files", "expected_environment", "expected_loaded"), + [ + ({".env": "VAR1=value1", ".env.local": "VAR2=value2"}, {"VAR1": "value1", "VAR2": "value2"}, True), + ({".env": "VAR1=value1"}, {"VAR1": "value1"}, True), + ({}, {}, False), + ], + ids=["base-and-local", "base-only", "none"], + ) + def test_default_file_selection( + self, + tmp_path: pathlib.Path, + files: dict[str, str], + expected_environment: dict[str, str], + expected_loaded: bool, + ) -> None: + for file_name, content in files.items(): + (tmp_path / file_name).write_text(content) - @mock.patch("pyrit.setup.environment_loading.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("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", tmp_path), + mock.patch.dict(os.environ, {}, clear=True), + ): + assert load_environment_files(env_files=None, silent=True) is expected_loaded + assert os.environ == expected_environment - with mock.patch.dict(os.environ, {}, clear=True): - loaded = load_environment_files(env_files=None) + @pytest.mark.parametrize( + ("base", "local", "initial_environment", "include_default_base", "expected_environment"), + [ + ( + "VAR=legacy\nLEGACY_ONLY=legacy", + None, + {"VAR": "process"}, + True, + {"VAR": "process", "LEGACY_ONLY": "legacy"}, + ), + ("VAR=legacy", "VAR=local", {"VAR": "process"}, True, {"VAR": "local"}), + ("VAR=base", "VAR=local", {}, False, {"VAR": "local"}), + ], + ids=["process-over-base", "local-over-process", "excluded-base"], + ) + def test_default_file_precedence( + self, + tmp_path: pathlib.Path, + base: str, + local: str | None, + initial_environment: dict[str, str], + include_default_base: bool, + expected_environment: dict[str, str], + ) -> None: + (tmp_path / ".env").write_text(base) + if local is not None: + (tmp_path / ".env.local").write_text(local) - assert loaded is False - assert os.environ == {} + with ( + mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", tmp_path), + mock.patch.dict(os.environ, initial_environment, clear=True), + ): + assert load_environment_files( + env_files=None, + silent=True, + include_default_base=include_default_base, + ) + assert os.environ == expected_environment @mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH") def test_auto_discovered_env_warns_about_plaintext(self, mock_config_path, caplog, capsys): @@ -195,10 +168,7 @@ async def test_akv_ignores_auto_discovered_env_and_loads_env_local(self, mock_co mock.patch( "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock, - return_value=_AkvEnvironmentDocument( - content="VALUE=akv\n", - vault_url="https://vault.vault.azure.net", - ), + return_value=("VALUE=akv\n", "https://vault.vault.azure.net"), ), mock.patch("pyrit.setup.environment_loading._load_environment_files") as mock_load_files, ): @@ -280,32 +250,16 @@ async def test_load_environment_files_honors_python_dotenv_disabled(self): assert loaded is True assert "DISABLED_VALUE" not in os.environ - async def test_load_environment_async_skips_sources_when_python_dotenv_disabled(self): - with ( - mock.patch.dict(os.environ, {"PYTHON_DOTENV_DISABLED": "true"}, clear=True), - mock.patch( - "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock - ) as mock_fetch_akv, - mock.patch("pyrit.setup.environment_loading._load_environment_files") as mock_load_files, - ): - await load_environment_async( - env_akv_ref=None, - env_files=None, - env_akv_strict=True, - silent=True, - ) - - mock_fetch_akv.assert_not_awaited() - mock_load_files.assert_not_called() - @pytest.mark.parametrize( ("env_akv_ref", "env_files"), [ + (None, None), (["https://vault.vault.azure.net/secrets/bootstrap"], None), (None, [pathlib.Path("configured.env")]), ], + ids=["defaults", "akv", "file"], ) - async def test_load_environment_async_skips_configured_sources_when_python_dotenv_disabled( + async def test_load_environment_async_skips_sources_when_python_dotenv_disabled( self, env_akv_ref, env_files ): with ( @@ -331,6 +285,35 @@ async def test_load_environment_async_skips_configured_sources_when_python_doten mock_fetch_akv.assert_not_awaited() mock_load_files.assert_not_called() + @pytest.mark.parametrize("env_akv_ref", ["https://vault.vault.azure.net/secrets/one", [""], [None]]) + async def test_load_environment_async_rejects_invalid_env_akv_ref(self, env_akv_ref): + with pytest.raises(ValueError, match="env_akv_ref must"): + await load_environment_async( + env_akv_ref=env_akv_ref, # type: ignore[arg-type] + env_files=[], + env_akv_strict=True, + silent=True, + ) + + async def test_load_environment_async_rejects_multiple_bootstrap_urls_before_loading(self): + with ( + mock.patch("pyrit.setup.environment_loading._fetch_akv_document_async") as mock_fetch, + mock.patch("pyrit.setup.environment_loading._load_environment_files") as mock_load_files, + pytest.raises(ValueError, match="at most one"), + ): + await load_environment_async( + env_akv_ref=[ + "https://vault.vault.azure.net/secrets/first", + "https://vault.vault.azure.net/secrets/second", + ], + env_files=[], + env_akv_strict=True, + silent=True, + ) + + mock_fetch.assert_not_called() + mock_load_files.assert_not_called() + async def test_runtime_does_not_fetch_akv_reference_overridden_by_env_local(self): credential, client = _create_mock_akv_clients() bootstrap_document = "API_KEY=kv:https://vault.vault.azure.net/secrets/api-key\nAKV_ONLY=akv\n" @@ -360,7 +343,12 @@ async def test_runtime_does_not_fetch_akv_reference_overridden_by_env_local(self async def test_runtime_resolves_local_alias_after_all_sources_are_loaded(self): credential, client = _create_mock_akv_clients() bootstrap_document = "A=kv:https://vault.vault.azure.net/secrets/key\n" - client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=bootstrap_document)) + + async def get_secret(secret_name, **kwargs): + value = bootstrap_document if secret_name == "bootstrap" else "resolved-for-B" + return types.SimpleNamespace(value=value) + + client.get_secret = mock.AsyncMock(side_effect=get_secret) with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -369,18 +357,10 @@ async def test_runtime_resolves_local_alias_after_all_sources_are_loaded(self): local_file = temp_path / ".env.local" local_file.write_text("A=literal\n", encoding="utf-8") - async def resolve_by_variable_name(**kwargs): - return f"resolved-for-{kwargs['variable_name']}" - 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), - mock.patch( - "pyrit.setup.environment_loading._fetch_akv_secret_value_async", - new_callable=mock.AsyncMock, - side_effect=resolve_by_variable_name, - ) as mock_fetch_child, ): await load_environment_async( env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], @@ -392,53 +372,35 @@ async def resolve_by_variable_name(**kwargs): assert os.environ["A"] == "literal" assert os.environ["B"] == "resolved-for-B" - mock_fetch_child.assert_awaited_once() - assert mock_fetch_child.await_args is not None - assert mock_fetch_child.await_args.kwargs["variable_name"] == "B" + assert client.get_secret.await_args_list == [ + mock.call("bootstrap", version=None), + mock.call("key", version=None), + ] - @pytest.mark.parametrize("fallback_source", ["akv", "file"]) - async def test_non_strict_runtime_uses_blocked_lower_candidate_after_malformed_akv_winner(self, fallback_source): + async def test_non_strict_runtime_uses_file_candidate_after_malformed_akv_winner(self): credential, client = _create_mock_akv_clients() - documents = { - "https://vault.vault.azure.net/secrets/first": "API_KEY=kv:short\n", - "https://vault.vault.azure.net/secrets/second": ( - "API_KEY=kv:https://vault.vault.azure.net/secrets/fallback-key\n" - ), - } - - async def fetch_document(secret_url, **kwargs): - return _AkvEnvironmentDocument( - content=documents[secret_url], - vault_url="https://vault.vault.azure.net", - ) with tempfile.TemporaryDirectory() as temp_dir: - env_files: list[pathlib.Path] = [] - env_akv_ref = ["https://vault.vault.azure.net/secrets/first"] - if fallback_source == "akv": - env_akv_ref.append("https://vault.vault.azure.net/secrets/second") - else: - ordinary_file = pathlib.Path(temp_dir) / "ordinary.env" - ordinary_file.write_text( - "API_KEY=kv:https://local-vault.vault.azure.net/secrets/fallback-key\n", - encoding="utf-8", - ) - env_files.append(ordinary_file) + ordinary_file = pathlib.Path(temp_dir) / "ordinary.env" + ordinary_file.write_text( + "API_KEY=kv:https://local-vault.vault.azure.net/secrets/fallback-key\n", + encoding="utf-8", + ) with ( mock.patch.dict(os.environ, {}, clear=True), mock.patch( "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock, - side_effect=fetch_document, + return_value=("API_KEY=kv:short\n", "https://vault.vault.azure.net"), ), mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), ): client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="resolved-fallback")) await load_environment_async( - env_akv_ref=env_akv_ref, - env_files=env_files, + env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=[ordinary_file], env_akv_strict=False, silent=True, ) @@ -488,11 +450,10 @@ async def test_env_akv_strict_does_not_validate_local_environment_files(self): 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, + await load_environment_async( + env_akv_ref=None, env_files=[env_file], env_akv_strict=True, - load_defaults=False, silent=True, ) @@ -780,51 +741,6 @@ async def test_raises_error_for_nonexistent_env_file(self): 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) @@ -860,7 +776,11 @@ class TestAkvEnvironmentLoading: 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 + assert _parse_akv_reference(value=f"{prefix}:{secret_url}", variable_name="API_KEY") == ( + "https://myvault.vault.azure.net", + "api-key", + None, + ) @pytest.mark.parametrize( "value", @@ -871,35 +791,31 @@ def test_parse_akv_reference_accepts_aliases(self, prefix): ], ) 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" + assert _parse_akv_reference(value=value, variable_name="API_KEY") is None - 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", "expected"), + [ + ( + "https://myvault.vault.azure.net/secrets/my-secret", + ("https://myvault.vault.azure.net", "my-secret", None), + ), + ( + "https://myvault.vault.azure.net/secrets/my-secret/abc123", + ("https://myvault.vault.azure.net", "my-secret", "abc123"), + ), + *[ + ( + f"https://myvault.{suffix}/secrets/my-secret/version-1", + (f"https://myvault.{suffix}", "my-secret", "version-1"), + ) + for suffix in ("vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net") + ], + ], + ids=["unversioned", "versioned", "public", "china", "us-government"], + ) + def test_parse_akv_secret_url_valid(self, url, expected): + assert _parse_akv_secret_url(url) == expected @pytest.mark.parametrize( "url", @@ -958,8 +874,7 @@ async def test_fetch_akv_document_async_returns_validated_document(self): ): document = await _fetch_akv_document_async(secret_url=secret_url, silent=True) - assert document.content == root_document - assert document.vault_url == "https://myvault.vault.azure.net" + assert document == (root_document, "https://myvault.vault.azure.net") assert os.environ == {"SOURCE_VALUE": "ambient-value"} mock_credential_cls.assert_called_once_with() @@ -1046,80 +961,36 @@ async def test_runtime_rejects_cross_vault_reference(self, reference_url): silent=True, ) - async def test_fetch_akv_document_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 _fetch_akv_document_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_fetch_akv_document_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 _fetch_akv_document_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"), [ + (None, "has no value"), + ("# comments only\n", "contains no environment entries"), ("GOOD=resolved\n=malformed\nOTHER=resolved", "malformed entries at lines: 2"), ("MISSING_VALUE\n", "variables without values: MISSING_VALUE"), ], + ids=["missing-value", "no-entries", "malformed", "valueless"], ) - async def test_fetch_akv_document_async_rejects_non_assignments(self, document, error): + async def test_fetch_akv_document_async_rejects_invalid_document(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), - ): + 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), + ): + with pytest.raises(KeyVaultInitializationException, match=error) as exc_info: await _fetch_akv_document_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_fetch_akv_document_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 _fetch_akv_document_async( - secret_url="https://myvault.vault.azure.net/secrets/bootstrap", - silent=True, - ) + assert os.environ == {} assert isinstance(exc_info.value.__cause__, ValueError) + credential.__aexit__.assert_awaited_once() + client.__aexit__.assert_awaited_once() async def test_runtime_wraps_missing_child_secret(self): credential, client = _create_mock_akv_clients() @@ -1205,7 +1076,10 @@ async def test_fetch_akv_document_async_non_strict_warns_and_skips_invalid_entri silent=False, ) - assert fetched_document.content == "GOOD=resolved\nOTHER=also-resolved" + assert fetched_document == ( + "GOOD=resolved\nOTHER=also-resolved", + "https://myvault.vault.azure.net", + ) assert os.environ == {} output = capsys.readouterr().out diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 57c6eb349e..b7b8518104 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -12,7 +12,6 @@ from pyrit.common.singleton import Singleton from pyrit.registry import InitializerRegistry from pyrit.setup import IN_MEMORY, initialize_pyrit_async -from pyrit.setup.environment_loading import _AkvEnvironmentDocument class TestLoadInitializersFromScripts: @@ -120,17 +119,17 @@ def setup_method(self) -> None: reset_default_values() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.environment_loading._load_environment_files", return_value=False) - async def test_initialize_basic(self, mock_load_env, mock_set_memory): + @mock.patch("pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock) + async def test_initialize_basic(self, mock_load_environment, mock_set_memory): """Test basic initialization.""" await initialize_pyrit_async(memory_db_type=IN_MEMORY, load_defaults=False) - mock_load_env.assert_called_once() + mock_load_environment.assert_awaited_once() mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.environment_loading._load_environment_files", return_value=False) - async def test_initialize_with_script(self, mock_load_env, mock_set_memory): + @mock.patch("pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock) + async def test_initialize_with_script(self, mock_load_environment, mock_set_memory): """Test initialization with a script.""" with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f: f.write( @@ -154,68 +153,42 @@ 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_load_environment.assert_awaited_once() mock_set_memory.assert_called_once() finally: os.unlink(script_path) - @mock.patch("pyrit.setup.environment_loading._load_environment_files", return_value=False) - async def test_invalid_memory_type_raises_error(self, mock_load_env): + @mock.patch("pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock) + async def test_invalid_memory_type_raises_error(self, mock_load_environment): """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_load_environment.assert_awaited_once() + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.environment_loading._load_environment_files", return_value=False) - @mock.patch("pyrit.setup.environment_loading._fetch_akv_document_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 = [ - "https://vault.vault.azure.net/secrets/first", - "https://vault.vault.azure.net/secrets/second/version", - ] - - mock_load_akv.side_effect = [ - _AkvEnvironmentDocument(content="FIRST=one\n", vault_url="https://vault.vault.azure.net"), - _AkvEnvironmentDocument(content="SECOND=two\n", vault_url="https://vault.vault.azure.net"), - ] + @mock.patch("pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock) + async def test_initialize_forwards_environment_options(self, mock_load_environment, mock_set_memory): + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + env_files = [pathlib.Path("custom.env")] await initialize_pyrit_async( memory_db_type=IN_MEMORY, env_akv_ref=refs, - env_files=[], + env_files=env_files, + env_akv_strict=False, + silent=True, load_defaults=False, ) - 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.environment_loading._load_environment_files", return_value=False) - @mock.patch("pyrit.setup.environment_loading._fetch_akv_document_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 - ): - """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_load_env.assert_called_once() + mock_load_environment.assert_awaited_once_with( + env_akv_ref=refs, + env_files=env_files, + env_akv_strict=False, + silent=True, + ) 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, - ) - @pytest.mark.parametrize("invalid_value", ["false", "true", 0, 1, None, [], {}]) async def test_initialize_rejects_non_boolean_env_akv_strict_before_loading(self, invalid_value): with mock.patch( @@ -230,173 +203,6 @@ async def test_initialize_rejects_non_boolean_env_akv_strict_before_loading(self mock_load_environment.assert_not_awaited() - @pytest.mark.parametrize("env_akv_strict", [True, False]) - async def test_initialize_forwards_env_akv_strict(self, env_akv_strict): - with ( - mock.patch( - "pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock - ) as mock_load_environment, - mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance"), - ): - await initialize_pyrit_async( - memory_db_type=IN_MEMORY, - env_akv_strict=env_akv_strict, - load_defaults=False, - ) - - await_args = mock_load_environment.await_args - assert await_args is not None - assert await_args.kwargs["env_akv_strict"] is env_akv_strict - - @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - 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): - with ( - mock.patch( - "pyrit.setup.environment_loading._fetch_akv_document_async", - new_callable=mock.AsyncMock, - return_value=_AkvEnvironmentDocument( - content="FROM_AKV=resolved\n", - vault_url="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_files=[nonexistent], - load_defaults=False, - ) - - 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_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") - - with ( - mock.patch.dict(os.environ, {}, clear=True), - mock.patch( - "pyrit.setup.environment_loading._fetch_akv_document_async", - new_callable=mock.AsyncMock, - return_value=_AkvEnvironmentDocument( - content="BASE=akv\nONLY_AKV=shared\n", - vault_url="https://vault.vault.azure.net", - ), - ), - ): - 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_default_files_override_akv_in_order(self, mock_set_memory): - 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") - (temp_path / ".env.local").write_text("VALUE=local") - - with ( - mock.patch.dict(os.environ, {}, clear=True), - mock.patch("pyrit.setup.environment_loading.path.CONFIGURATION_DIRECTORY_PATH", temp_path), - mock.patch("pyrit.setup.environment_loading._warn_about_dotenv_file"), - mock.patch( - "pyrit.setup.environment_loading._fetch_akv_document_async", - new_callable=mock.AsyncMock, - return_value=_AkvEnvironmentDocument( - content="VALUE=akv\n", - vault_url="https://vault.vault.azure.net", - ), - ), - ): - await initialize_pyrit_async( - memory_db_type=IN_MEMORY, - env_akv_ref=refs, - 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_bootstrap_references_before_local_overrides(self, mock_set_memory): - refs = ["https://vault.vault.azure.net/secrets/bootstrap"] - 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=mock.MagicMock(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": "unused-secret-value", - "BOOTSTRAP_SECRET": "bootstrap-secret-value", - "BOOTSTRAP_SOURCE": "bootstrap-value", - } - - with ( - mock.patch.dict(os.environ, {"SOURCE_VALUE": "ambient-value"}, clear=True), - mock.patch( - "pyrit.setup.environment_loading._fetch_akv_document_async", - new_callable=mock.AsyncMock, - return_value=_AkvEnvironmentDocument( - content="".join(f"{name}={value}\n" for name, value in bootstrap_environment.items()), - vault_url="https://vault.vault.azure.net", - ), - ), - mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), - mock.patch("pyrit.setup.environment_loading._create_akv_secret_client", return_value=client), - ): - await initialize_pyrit_async( - memory_db_type=IN_MEMORY, - 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"] == "env:BOOTSTRAP_SOURCE" - - client.get_secret.assert_awaited_once_with("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) - - mock_set_memory.assert_called_once() - @pytest.fixture def reset_memory_singletons(): @@ -418,16 +224,16 @@ def setup_method(self) -> None: """Clear default values before each test.""" reset_default_values() - @mock.patch("pyrit.setup.environment_loading._load_environment_files", return_value=True) - async def test_initialize_silent_produces_no_output(self, mock_load_env, capsys): + @mock.patch("pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock) + async def test_initialize_silent_produces_no_output(self, mock_load_environment, 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) captured = capsys.readouterr() assert captured.out == "" - @mock.patch("pyrit.setup.environment_loading._load_environment_files", return_value=True) - async def test_initialize_not_silent_prints_migration_message(self, mock_load_env, capsys): + @mock.patch("pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock) + async def test_initialize_not_silent_prints_migration_message(self, mock_load_environment, 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) diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index d0793f90b2..52141904e1 100644 --- a/tests/unit/setup/test_targets_initializer.py +++ b/tests/unit/setup/test_targets_initializer.py @@ -116,54 +116,6 @@ async def test_registers_multiple_targets(self): assert "platform_openai_chat" in registry.instances assert "openai_image_platform" in registry.instances - @pytest.mark.parametrize( - ("registry_name", "endpoint_var", "key_var", "model_var", "endpoint"), - [ - ( - "openai_image_azure", - "OPENAI_IMAGE_ENDPOINT1", - "OPENAI_IMAGE_API_KEY1", - "OPENAI_IMAGE_MODEL1", - "https://image.openai.azure.com/openai/v1", - ), - ( - "openai_image_platform", - "OPENAI_IMAGE_ENDPOINT2", - "OPENAI_IMAGE_API_KEY2", - "OPENAI_IMAGE_MODEL2", - "https://api.openai.com/v1", - ), - ( - "openai_tts_azure", - "OPENAI_TTS_ENDPOINT1", - "OPENAI_TTS_KEY1", - "OPENAI_TTS_MODEL1", - "https://tts.openai.azure.com/openai/v1", - ), - ( - "openai_tts_platform", - "OPENAI_TTS_ENDPOINT2", - "OPENAI_TTS_KEY2", - "OPENAI_TTS_MODEL2", - "https://api.openai.com/v1", - ), - ], - ) - async def test_media_targets_use_main_environment_contract( - self, registry_name, endpoint_var, key_var, model_var, endpoint - ): - with patch.dict( - os.environ, - {endpoint_var: endpoint, key_var: "test-key", model_var: "test-model"}, - clear=True, - ): - await TargetInitializer().initialize_async() - - target = TargetRegistry.get_registry_singleton().instances.get(registry_name) - assert target is not None - assert target._endpoint == endpoint - assert target._model_name == "test-model" - async def test_registers_azure_content_safety_without_model(self): """Test that PromptShieldTarget is registered without model_name (it doesn't use one).""" os.environ["AZURE_CONTENT_SAFETY_API_ENDPOINT"] = "https://test.cognitiveservices.azure.com" From 3689d91aed4894ad3bf174d2136ed8c3f9b6f7d0 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 21 Aug 2026 15:57:18 -0400 Subject: [PATCH 27/28] FIX: Further simplifications --- build_scripts/export_akv_environment.py | 242 +++++------------- .../executor/promptgen/gcg/experiments/run.py | 4 +- pyrit/setup/environment_loading.py | 51 +--- pyrit/setup/initialization.py | 4 +- .../test_export_akv_environment.py | 4 +- tests/unit/setup/test_environment_loading.py | 11 +- 6 files changed, 85 insertions(+), 231 deletions(-) diff --git a/build_scripts/export_akv_environment.py b/build_scripts/export_akv_environment.py index 8ea6c0cec1..7a0df8b3d3 100644 --- a/build_scripts/export_akv_environment.py +++ b/build_scripts/export_akv_environment.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -"""Export resolved Azure Key Vault bootstrap documents to ``~/.pyrit/.env_akv``.""" +"""Export a resolved Azure Key Vault bootstrap document to ``~/.pyrit/.env_akv``.""" import argparse import contextlib @@ -9,16 +9,20 @@ import os import pathlib import tempfile -import urllib.parse from collections.abc import Mapping, Sequence -from dataclasses import dataclass from io import StringIO -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import dotenv from dotenv.parser import parse_stream from dotenv.variables import parse_variables +from pyrit.setup.environment_loading import ( + _parse_akv_reference, + _parse_akv_secret_url, + _validate_dotenv_document, +) + if TYPE_CHECKING: from azure.core.credentials import TokenCredential from azure.keyvault.secrets import SecretClient @@ -26,59 +30,6 @@ logger = logging.getLogger(__name__) DEFAULT_OUTPUT_FILE = pathlib.Path.home() / ".pyrit" / ".env_akv" -_REFERENCE_PREFIXES = frozenset({"akv", "kv", "azure_key_vault", "env_akv_ref"}) -_VAULT_DNS_SUFFIXES = frozenset({"vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"}) - - -@dataclass(frozen=True) -class _Document: - content: str - vault_url: str - - -@dataclass(frozen=True) -class _Candidate: - document_index: int - binding_index: int - name: str - value: str - vault_url: str - - -def _parse_secret_url(url: str) -> tuple[str, str, str | None]: - """Return vault URL, secret name, and optional version from a full AKV URL.""" - error_message = f"Invalid Azure Key Vault secret URL: {url}" - try: - parsed = urllib.parse.urlsplit(url) - port = parsed.port - except (TypeError, ValueError) as error: - raise ValueError(error_message) from error - hostname = parsed.hostname - vault_name, separator, suffix = hostname.partition(".") if hostname else ("", "", "") - valid_vault = 1 <= len(vault_name) <= 63 and all( - char.isascii() and (char.isalnum() or char == "-") for char in vault_name - ) - parts = parsed.path.split("/") - valid_path = len(parts) in {3, 4} and parts[:2] == ["", "secrets"] and all(parts[2:]) - if ( - parsed.scheme.casefold() != "https" - or parsed.username is not None - or parsed.password is not None - or port is not None - or separator != "." - or suffix not in _VAULT_DNS_SUFFIXES - or not valid_vault - or not valid_path - or parsed.query - or parsed.fragment - ): - raise ValueError(error_message) - secret_name = parts[2] - secret_version = parts[3] if len(parts) == 4 else None - identifiers = [secret_name] + ([secret_version] if secret_version else []) - if any(not (1 <= len(item) <= 127 and all(char.isalnum() or char == "-" for char in item)) for item in identifiers): - raise ValueError(error_message) - return f"https://{hostname}", secret_name, secret_version def _create_client(*, vault_url: str, credential: "TokenCredential") -> "SecretClient": @@ -107,87 +58,35 @@ def _client_for(*, vault_url: str, credential: "TokenCredential", clients: dict[ return client -def _validate_document(*, document: str, strict: bool, silent: bool) -> str: - bindings = list(parse_stream(StringIO(document))) - malformed = [str(binding.original.line) for binding in bindings if binding.error] - valueless = [binding.key for binding in bindings if binding.key is not None and binding.value is None] - issues: list[str] = [] - if malformed: - issues.append("malformed entries at lines: " + ", ".join(malformed)) - if valueless: - issues.append("variables without values: " + ", ".join(valueless)) - 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) - ) - - -def _fetch_documents( +def _fetch_document( *, - secret_urls: Sequence[str], + secret_url: str, credential: "TokenCredential", clients: dict[str, "SecretClient"], strict: bool, silent: bool, -) -> list[_Document]: - documents: list[_Document] = [] - for url in secret_urls: - vault_url, name, version = _parse_secret_url(url) - secret = _client_for(vault_url=vault_url, credential=credential, clients=clients).get_secret( - name, version=version - ) - if not secret.value: - raise ValueError(f"AKV environment secret has no value: {url}") - content = _validate_document(document=secret.value, strict=strict, silent=silent) - if not dotenv.dotenv_values(stream=StringIO(content), interpolate=False): - raise ValueError(f"AKV environment secret contains no assignments: {url}") - documents.append(_Document(content=content, vault_url=vault_url)) - return documents +) -> tuple[str, str]: + vault_url, name, version = _parse_akv_secret_url(secret_url) + secret = _client_for(vault_url=vault_url, credential=credential, clients=clients).get_secret(name, version=version) + if not secret.value: + raise ValueError(f"AKV environment secret has no value: {secret_url}") + content = _validate_dotenv_document(secret.value, strict=strict, silent=silent) + if not dotenv.dotenv_values(stream=StringIO(content), interpolate=False): + raise ValueError(f"AKV environment secret contains no assignments: {secret_url}") + return content, vault_url def _resolve_interpolation(*, value: str, environment: Mapping[str, str | None]) -> str: return "".join(atom.resolve(environment) for atom in parse_variables(value)) -def _build_candidates(documents: Sequence[_Document]) -> tuple[list[list[Any]], dict[str, list[_Candidate]]]: - effective: dict[str, str | None] = {} - chains: dict[str, list[_Candidate]] = {} - all_bindings: list[list[Any]] = [] - for document_index, document in enumerate(documents): - bindings = list(parse_stream(StringIO(document.content))) - all_bindings.append(bindings) - current: dict[str, str | None] = {} - final_indexes: dict[str, int] = {} - for binding_index, binding in enumerate(bindings): - if binding.key is None or binding.value is None: - continue - environment = dict(current) - environment.update(effective) - current[binding.key] = _resolve_interpolation(value=binding.value, environment=environment) - final_indexes[binding.key] = binding_index - for name, value in current.items(): - if value is None: - continue - chains.setdefault(name, []).append( - _Candidate(document_index, final_indexes[name], name, value, document.vault_url) - ) - effective.setdefault(name, value) - return all_bindings, chains - - -def _reference_target(value: str) -> str | None: - prefix, separator, target = value.partition(":") - return target.strip() if separator and prefix in _REFERENCE_PREFIXES else None +def _build_candidates(document: tuple[str, str]) -> dict[str, tuple[str, str]]: + content, vault_url = document + values: dict[str, str] = {} + for binding in parse_stream(StringIO(content)): + if binding.key is not None and binding.value is not None: + values[binding.key] = _resolve_interpolation(value=binding.value, environment=values) + return {name: (value, vault_url) for name, value in values.items()} def _serialize(value: str) -> str: @@ -197,61 +96,40 @@ def _serialize(value: str) -> str: def _render( *, - documents: Sequence[_Document], + document: tuple[str, str], credential: "TokenCredential", clients: dict[str, "SecretClient"], strict: bool, silent: bool, ) -> str: - all_bindings, chains = _build_candidates(documents) - selected: dict[str, _Candidate] = {} - resolved: dict[tuple[int, int], str] = {} - for name, candidates in chains.items(): - for candidate in candidates: - target = _reference_target(candidate.value) - if target is None: - selected[name] = candidate - break - try: - vault_url, secret_name, version = _parse_secret_url(target) - if vault_url.casefold() != candidate.vault_url.casefold(): - raise ValueError(f"Cross-vault AKV reference for '{name}' is not supported") - except ValueError as error: - if strict: - raise - message = f"Invalid AKV reference for '{name}' will be skipped: {error}" - if not silent: - print(f"WARNING: {message}") - logger.warning(message) - continue - secret = _client_for(vault_url=vault_url, credential=credential, clients=clients).get_secret( - secret_name, version=version + resolved: dict[str, str] = {} + for name, (value, source_vault_url) in _build_candidates(document).items(): + try: + reference = _parse_akv_reference( + value=value, + variable_name=name, + expected_vault_url=source_vault_url, ) - if secret.value is None: - raise ValueError(f"AKV secret '{secret_name}' referenced by '{name}' has no value") - selected[name] = candidate - resolved[(candidate.document_index, candidate.binding_index)] = secret.value - break - - output: list[str] = [] - for document_index, bindings in enumerate(all_bindings): - for binding_index, binding in enumerate(bindings): - name = binding.key - if name is None: - output.append(binding.original.string) - continue - winner = selected.get(name) - if winner is None or winner.document_index != document_index: - continue - value = resolved.get((document_index, binding_index)) - if value is None: - output.append(binding.original.string) - continue - original = binding.original.string - export = "export " if original.lstrip().startswith("export ") else "" - newline = "\r\n" if original.endswith("\r\n") else "\n" if original.endswith("\n") else "" - output.append(f"{export}{name}={_serialize(value)}{newline}") - return "".join(output).rstrip("\r\n") + "\n" + except ValueError as error: + if strict: + raise + message = f"Invalid AKV reference for '{name}' will be skipped: {error}" + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + continue + if reference is None: + resolved[name] = value + continue + vault_url, secret_name, version = reference + secret = _client_for(vault_url=vault_url, credential=credential, clients=clients).get_secret( + secret_name, version=version + ) + if secret.value is None: + raise ValueError(f"AKV secret '{secret_name}' referenced by '{name}' has no value") + resolved[name] = secret.value + + return "".join(f"{name}={_serialize(value)}\n" for name, value in resolved.items()) def _ensure_output_available(output_file: pathlib.Path) -> pathlib.Path: @@ -307,6 +185,8 @@ def export_akv_environment( """ if not secret_urls: raise ValueError("At least one secret URL is required") + if len(secret_urls) > 1: + raise ValueError("Only one Azure Key Vault bootstrap secret URL is supported") output_file = _ensure_output_available(output_file) from azure.identity import DefaultAzureCredential @@ -318,15 +198,15 @@ def export_akv_environment( active_credential = credential clients: dict[str, SecretClient] = {} try: - documents = _fetch_documents( - secret_urls=secret_urls, + document = _fetch_document( + secret_url=secret_urls[0], credential=active_credential, clients=clients, strict=strict, silent=silent, ) document = _render( - documents=documents, + document=document, credential=active_credential, clients=clients, strict=strict, @@ -345,14 +225,14 @@ def export_akv_environment( def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--secret-url", dest="secret_urls", action="append", required=True) + parser.add_argument("--secret-url", required=True) parser.add_argument("--output", type=pathlib.Path, default=DEFAULT_OUTPUT_FILE) parser.add_argument("--non-strict", action="store_true") parser.add_argument("--silent", action="store_true") args = parser.parse_args() try: export_akv_environment( - secret_urls=args.secret_urls, + secret_urls=[args.secret_url], output_file=args.output, strict=not args.non_strict, silent=args.silent, diff --git a/pyrit/executor/promptgen/gcg/experiments/run.py b/pyrit/executor/promptgen/gcg/experiments/run.py index 3bad3dc2f9..3c7cbf0390 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.environment_loading import load_environment_files +from pyrit.setup.initialization import _load_environment_files def _parse_arguments() -> argparse.Namespace: @@ -85,7 +85,7 @@ def _resolve_output(*, output: GCGOutputConfig, output_dir: str | None) -> GCGOu async def _main_async(config_path: str, data_path: str, output_dir: str | None = None) -> None: - load_environment_files(env_files=None) + _load_environment_files(env_files=None) config = GCGConfig.from_json_file(config_path) data = GCGDataConfig.from_json_file(data_path) if config.hf_token is None: diff --git a/pyrit/setup/environment_loading.py b/pyrit/setup/environment_loading.py index e3122adc9c..3e12b91f9c 100644 --- a/pyrit/setup/environment_loading.py +++ b/pyrit/setup/environment_loading.py @@ -55,29 +55,8 @@ def load_environment_files( *, silent: bool = False, include_default_base: bool = True, -) -> bool: - """ - Load local environment files using PyRIT's standard precedence. - - Returns: - bool: Whether at least one environment file was selected. - """ - return _load_environment_files( - env_files=env_files, - silent=silent, - include_default_base=include_default_base, - ordinary_candidates=None, - override_candidates=None, - ) - - -def _load_environment_files( - env_files: Sequence[pathlib.Path] | None, - *, - silent: bool, - include_default_base: bool, - ordinary_candidates: dict[str, list[tuple[str, str | None]]] | None = None, - override_candidates: dict[str, list[tuple[str, str | None]]] | None = None, + _ordinary_candidates: dict[str, list[tuple[str, str | None]]] | None = None, + _override_candidates: dict[str, list[tuple[str, str | None]]] | None = None, ) -> bool: """ Load environment files in the order they are provided. @@ -92,8 +71,8 @@ def _load_environment_files( 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. - ordinary_candidates: Optional output mapping for non-overriding assignments. - override_candidates: Optional output mapping for ``.env.local`` assignments. + _ordinary_candidates: Internal output mapping for non-overriding assignments. + _override_candidates: Internal output mapping for ``.env.local`` assignments. Returns: True if at least one environment file was loaded, otherwise False. @@ -133,8 +112,8 @@ def _load_environment_files( loaded = _load_dotenv_source( dotenv_path=env_file, override=env_file.name == ".env.local", - ordinary_candidates=ordinary_candidates, - override_candidates=override_candidates, + ordinary_candidates=_ordinary_candidates, + override_candidates=_override_candidates, ) if not silent: _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) @@ -466,12 +445,12 @@ async def load_environment_async( ) await asyncio.to_thread( - _load_environment_files, + load_environment_files, env_files=env_files, silent=silent, include_default_base=not (env_akv_ref and env_files is None), - ordinary_candidates=ordinary_candidates, - override_candidates=override_candidates, + _ordinary_candidates=ordinary_candidates, + _override_candidates=override_candidates, ) await _resolve_environment_candidates_async( process_environment=process_environment, @@ -530,19 +509,17 @@ async def _resolve_environment_candidates_async( Raises: KeyVaultInitializationException: If strict validation or secret retrieval fails. + ValueError: If a referenced secret has no value. """ parsed_references: list[tuple[str, str, str, str | None]] = [] variable_names = ordinary_candidates.keys() | override_candidates.keys() for variable_name in variable_names: candidates = [ - (value, vault_url, True) - for value, vault_url in reversed(override_candidates.get(variable_name, ())) + (value, vault_url, True) for value, vault_url in reversed(override_candidates.get(variable_name, ())) ] if variable_name in process_environment: candidates.append((process_environment[variable_name], None, False)) - candidates.extend( - (value, vault_url, True) for value, vault_url in ordinary_candidates.get(variable_name, ()) - ) + candidates.extend((value, vault_url, True) for value, vault_url in ordinary_candidates.get(variable_name, ())) for value, expected_vault_url, resolve_reference in candidates: if not resolve_reference: os.environ[variable_name] = value @@ -560,9 +537,7 @@ async def _resolve_environment_candidates_async( error=error, ) raise wrapped_error from error - message = ( - f"Invalid AKV reference for environment variable '{variable_name}' will be skipped: {error}" - ) + message = f"Invalid AKV reference for environment variable '{variable_name}' will be skipped: {error}" if not silent: print(f"WARNING: {message}") logger.warning(message) diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index 5aef5cfc61..7d99e2783e 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -9,7 +9,7 @@ from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory from pyrit.setup.environment_loading import ( load_environment_async, - load_environment_files as _load_environment_files, + load_environment_files, validate_env_akv_strict, ) @@ -23,6 +23,8 @@ AZURE_SQL = "AzureSQL" MemoryDatabaseType = Literal["InMemory", "SQLite", "AzureSQL"] +_load_environment_files = load_environment_files + async def _execute_initializers_async(*, initializers: Sequence["PyRITInitializer"]) -> None: """ diff --git a/tests/unit/build_scripts/test_export_akv_environment.py b/tests/unit/build_scripts/test_export_akv_environment.py index 24888ceb8b..de9f988a20 100644 --- a/tests/unit/build_scripts/test_export_akv_environment.py +++ b/tests/unit/build_scripts/test_export_akv_environment.py @@ -37,9 +37,7 @@ def test_serialize_round_trips_terminal_values(value: str) -> None: def test_render_resolves_akv_only_values() -> None: document = ( - ( - "# AKV config\nBASE=bootstrap\nDERIVED=${BASE}\nAPI_KEY=kv:https://vault.vault.azure.net/secrets/api-key\n" - ), + ("# AKV config\nBASE=bootstrap\nDERIVED=${BASE}\nAPI_KEY=kv:https://vault.vault.azure.net/secrets/api-key\n"), "https://vault.vault.azure.net", ) client = mock.MagicMock() diff --git a/tests/unit/setup/test_environment_loading.py b/tests/unit/setup/test_environment_loading.py index d4c776e5b1..c8a29080e6 100644 --- a/tests/unit/setup/test_environment_loading.py +++ b/tests/unit/setup/test_environment_loading.py @@ -170,7 +170,7 @@ async def test_akv_ignores_auto_discovered_env_and_loads_env_local(self, mock_co new_callable=mock.AsyncMock, return_value=("VALUE=akv\n", "https://vault.vault.azure.net"), ), - mock.patch("pyrit.setup.environment_loading._load_environment_files") as mock_load_files, + mock.patch("pyrit.setup.environment_loading.load_environment_files") as mock_load_files, ): await load_environment_async( env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], @@ -259,9 +259,7 @@ async def test_load_environment_files_honors_python_dotenv_disabled(self): ], ids=["defaults", "akv", "file"], ) - async def test_load_environment_async_skips_sources_when_python_dotenv_disabled( - self, env_akv_ref, env_files - ): + async def test_load_environment_async_skips_sources_when_python_dotenv_disabled(self, env_akv_ref, env_files): with ( mock.patch.dict( os.environ, @@ -271,7 +269,7 @@ async def test_load_environment_async_skips_sources_when_python_dotenv_disabled( mock.patch( "pyrit.setup.environment_loading._fetch_akv_document_async", new_callable=mock.AsyncMock ) as mock_fetch_akv, - mock.patch("pyrit.setup.environment_loading._load_environment_files") as mock_load_files, + mock.patch("pyrit.setup.environment_loading.load_environment_files") as mock_load_files, ): await load_environment_async( env_akv_ref=env_akv_ref, @@ -298,7 +296,7 @@ async def test_load_environment_async_rejects_invalid_env_akv_ref(self, env_akv_ async def test_load_environment_async_rejects_multiple_bootstrap_urls_before_loading(self): with ( mock.patch("pyrit.setup.environment_loading._fetch_akv_document_async") as mock_fetch, - mock.patch("pyrit.setup.environment_loading._load_environment_files") as mock_load_files, + mock.patch("pyrit.setup.environment_loading.load_environment_files") as mock_load_files, pytest.raises(ValueError, match="at most one"), ): await load_environment_async( @@ -741,6 +739,7 @@ async def test_raises_error_for_nonexistent_env_file(self): with pytest.raises(ValueError, match="Environment file not found"): load_environment_files(env_files=[nonexistent]) + def _create_mock_akv_clients() -> tuple[mock.MagicMock, mock.MagicMock]: credential = mock.MagicMock() credential.__aenter__ = mock.AsyncMock(return_value=credential) From f561eeadd2f97c4012f6934ea7c8c3f9d99e6f73 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 21 Aug 2026 16:42:46 -0400 Subject: [PATCH 28/28] FIX: Unused variable not caught by linting --- pyrit/setup/environment_loading.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/pyrit/setup/environment_loading.py b/pyrit/setup/environment_loading.py index 3e12b91f9c..834f66e7d7 100644 --- a/pyrit/setup/environment_loading.py +++ b/pyrit/setup/environment_loading.py @@ -109,7 +109,7 @@ def load_environment_files( raise ValueError(f"Environment file not found: {env_file}") for env_file in selected_files: - loaded = _load_dotenv_source( + _load_dotenv_source( dotenv_path=env_file, override=env_file.name == ".env.local", ordinary_candidates=_ordinary_candidates, @@ -129,20 +129,17 @@ def _load_dotenv_source( dotenv_path: pathlib.Path | None = None, document: str | None = None, expected_vault_url: str | None = None, -) -> bool: +) -> None: """ Load one dotenv source and record values that participate in precedence. - Returns: - bool: Whether python-dotenv loaded at least one assignment. - Raises: ValueError: If both or neither source representations are provided. """ if (dotenv_path is None) == (document is None): raise ValueError("Exactly one dotenv_path or document must be provided.") if os.environ.get("PYTHON_DOTENV_DISABLED", "").casefold() in _DOTENV_DISABLED_VALUES: - return False + return source = DotEnv( dotenv_path=dotenv_path, @@ -151,16 +148,15 @@ def _load_dotenv_source( interpolate=True, ) assignment_values = source.dict() - loaded = source.set_as_environment_variables() - if ordinary_candidates is None or override_candidates is None or not loaded: - return loaded + source.set_as_environment_variables() + if ordinary_candidates is None or override_candidates is None: + return candidates = override_candidates if override else ordinary_candidates for variable_name, loaded_value in assignment_values.items(): if loaded_value is None: continue candidates.setdefault(variable_name, []).append((loaded_value, expected_vault_url)) - return loaded def _print_msg(message: str, quiet: bool, log: bool) -> None: