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/.pyrit_conf_example b/.pyrit_conf_example index b41c13e060..8a129f06e6 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 # -------------------- @@ -79,33 +79,23 @@ operation: op_trash_panda # - /path/to/my_custom_initializer.py # - ./local_initializer.py -# Environment Files -# ----------------- -# List of .env file paths to load during initialization. -# Later files override values from earlier files. -# -# Behavior: -# - Omit this field (or set to null): Load default .env and .env.local from ~/.pyrit/ if they exist -# - Set to []: Explicitly load NO environment files -# - Set to list of paths: Load only the specified files -# -# Example: -# env_files: -# - /path/to/.env -# - /path/to/.env.local - -# Azure Key Vault Environment References -# --------------------------------------- -# List of AKV secret URLs to load during initialization. -# Each secret's value must be the full contents of a .env file. -# Loaded after env_files, so AKV secrets take precedence. -# Authentication uses DefaultAzureCredential (managed identity, Azure CLI, etc.). -# -# Requires: pip install azure-keyvault-secrets -# -# Example: +# Environment Configuration +# ------------------------- +# Azure Key Vault is 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 + +# 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. +# Explicit env_files remain supported regardless of name or location and may contain full kv: URLs. +# env_files: +# - /path/to/.env.local # Max Concurrent Scenario Runs # ---------------------------- diff --git a/build_scripts/export_akv_environment.py b/build_scripts/export_akv_environment.py new file mode 100644 index 0000000000..7a0df8b3d3 --- /dev/null +++ b/build_scripts/export_akv_environment.py @@ -0,0 +1,246 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Export a resolved Azure Key Vault bootstrap document to ``~/.pyrit/.env_akv``.""" + +import argparse +import contextlib +import logging +import os +import pathlib +import tempfile +from collections.abc import Mapping, Sequence +from io import StringIO +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 + +logger = logging.getLogger(__name__) + +DEFAULT_OUTPUT_FILE = pathlib.Path.home() / ".pyrit" / ".env_akv" + + +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 _fetch_document( + *, + secret_url: str, + credential: "TokenCredential", + clients: dict[str, "SecretClient"], + strict: bool, + silent: bool, +) -> 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(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: + escaped = value.replace("\\", "\\\\").replace("'", "\\'").replace("${", "${:-$}{") + return f"'{escaped}'" + + +def _render( + *, + document: tuple[str, str], + credential: "TokenCredential", + clients: dict[str, "SecretClient"], + strict: bool, + silent: bool, +) -> str: + 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, + ) + 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: + 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") + 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 + + owned_credential = None + if credential is None: + owned_credential = DefaultAzureCredential() + active_credential = owned_credential + else: + active_credential = credential + clients: dict[str, SecretClient] = {} + try: + document = _fetch_document( + secret_url=secret_urls[0], + credential=active_credential, + clients=clients, + strict=strict, + silent=silent, + ) + document = _render( + document=document, + 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", 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_url], + 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 30fd1443a9..0692bf286f 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -4,62 +4,56 @@ 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 AI 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 The default configuration file path is: -``` +```text ~/.pyrit/.pyrit_conf ``` PyRIT looks for this file automatically on startup (via the CLI, shell, or `ConfigurationLoader`). If the file does not exist, PyRIT falls back to built-in defaults. -## Setting Up Secrets (.env files) +## Environment Configuration -The `.pyrit_conf` file works hand-in-hand with `.env` files for your API credentials. See [Populating Secrets](./populating_secrets.md) for provider-specific examples of what to put in your `.env` file. +```{important} +Azure Key Vault is 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. +``` -### Environment Variable Precedence +See [Populating Secrets](./populating_secrets.md) for provider-specific variable examples. -When PyRIT initializes, environment variables are loaded in a specific order. **Later sources override earlier ones:** +### Loading Order -```{mermaid} -flowchart LR - A["1. System Environment"] --> B{"env_files in .pyrit_conf?"} - B -->|No| C["2. ~/.pyrit/.env"] - C --> D["3. ~/.pyrit/.env.local"] - B -->|Yes| E["2. Your specified files (in order)"] -``` - -**Default behavior** (no `env_files` field in `.pyrit_conf`): +PyRIT loads environment sources in this order: -| 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. 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. -**Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. Default paths are completely ignored. +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 -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. @@ -107,7 +101,7 @@ Use `pyrit list initializers` in the CLI to see all registered initializers. See Most users should enable the following initializers. These are what the `.pyrit_conf_example` ships with and are required for features like `pyrit_scan` and automated scenarios. | Initializer | What It Registers | When You Need It | -|---|---|---| +| --- | --- | --- | | `target` | Prompt targets (OpenAI, Azure, AML, etc.) into the `TargetRegistry` | **Required for `pyrit_scan`** and any registry-based workflows | | `scorer` | Scorers (refusal, content safety, harm-category, Likert, etc.) into the `ScorerRegistry` | **Required for automated scoring** and `pyrit_scan` evaluations | | `technique` | Attack techniques into the `AttackTechniqueRegistry` | **Required for `pyrit_scan` scenarios** that select techniques | @@ -157,13 +151,13 @@ initialization_scripts: ### `env_files` -Environment file paths to load during initialization. Later files override values from earlier files. +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` if they exist | -| `[]` (empty list) | Load **no** environment files | -| List of paths | Load **only** the specified files (defaults are skipped) | +| Value | Behavior | +| ----------------- | -------------------------------------------------------- | +| Omitted or `null` | 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) | ```yaml env_files: @@ -171,6 +165,83 @@ env_files: - /path/to/.env.local ``` +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 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. 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. + +`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` + +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 +``` + +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" +OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key" +PINNED_OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version-id" +OPENAI_CHAT_MODEL="${PYRIT_OPENAI_CHAT_MODEL}" +``` + +Resolution is limited to one child-secret lookup: + +1. PyRIT validates and loads the bootstrap dotenv document. +2. For each complete-value Key Vault reference in that document, PyRIT fetches the same-vault scalar secret and replaces the environment value. + +For example, if `OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key"`, the value of the `openai-chat-key` secret becomes `OPENAI_CHAT_KEY` verbatim. If that secret happens to contain `kv:another-secret`, the final environment value is the string `kv:another-secret`; PyRIT does not fetch `another-secret`. + +References must occupy the entire value. `kv:` is the canonical Key Vault prefix; `akv:`, `azure_key_vault:`, and `env_akv_ref:` are accepted aliases. + +A Key Vault reference must use a full HTTPS secret URL from the bootstrap document's vault. Supported vault DNS suffixes are `.vault.azure.net`, `.vault.azure.cn`, and `.vault.usgovcloudapi.net`. An unversioned URL reads the latest secret version at initialization. Include the version in the URL to pin it. Short names, malformed paths, arbitrary hosts, and cross-vault child references are rejected before a client is created. + +PyRIT does not cache referenced secrets. Each 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" +PINNED_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version-id" +``` + +The bootstrap document stays in memory. Use `.env.local` when an intentional local override is required. + +### `env_akv_strict` + +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, 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 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. + +### Exporting AKV Configuration for Debugging + +Environment initialization never writes secrets to disk. To inspect an AKV-only configuration explicitly from a source checkout, run the standalone helper: + +```powershell +python -m build_scripts.export_akv_environment ` + --secret-url https://my-vault.vault.azure.net/secrets/my-pyrit-env +``` + +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. + ### `silent` If `true`, suppresses print statements during initialization. Useful for non-interactive environments or when embedding PyRIT in other tools. Defaults to `false`. @@ -180,7 +251,7 @@ If `true`, suppresses print statements during initialization. Useful for non-int Client settings for connecting to or launching a PyRIT backend. | Field | Description | Default | -|---|---|---| +| --- | --- | --- | | `url` | Backend URL used when `--server-url` is omitted | `http://localhost:8000` | | `startup_timeout` | Seconds `pyrit_scan start-server` waits for a healthy backend before terminating the spawned process | `120` | @@ -216,7 +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. Environment files are loaded +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 @@ -292,11 +363,13 @@ 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 +# 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 + +# Optional plaintext local patch or non-Azure workflow # env_files: -# - /path/to/.env # - /path/to/.env.local # 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 86f6b402e2..d78ae8cca4 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.environment_loading import validate_env_akv_strict from pyrit.setup.initialization import ( AZURE_SQL, IN_MEMORY, @@ -95,7 +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 "use defaults (.env, .env.local)", [] means "load nothing". + None means auto-discover supported ``.env`` and ``.env.local``; + [] means "load nothing". + 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. operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. @@ -135,6 +140,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 @@ -145,10 +151,28 @@ class ConfigurationLoader(YamlLoadable): def __post_init__(self) -> None: """Validate and normalize the configuration after loading.""" + validate_env_akv_strict(env_akv_strict=self.env_akv_strict) 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 a list of non-empty strings. + """ + if self.env_akv_ref is 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.") + def _normalize_memory_db_type(self) -> None: """ Normalize and validate memory_db_type. @@ -401,6 +425,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. @@ -416,7 +441,8 @@ 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 containing at most one Azure Key Vault bootstrap secret URL. + env_akv_strict: Override for strict Key Vault bootstrap validation. Returns: A merged ConfigurationLoader instance. @@ -477,8 +503,13 @@ def to_init_data(config: ConfigurationLoader) -> dict[str, Any]: config_data["env_files"] = list(env_files) if env_akv_ref is not None: + 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 + return cls.from_dict(config_data) @classmethod @@ -582,10 +613,10 @@ def resolve_env_files(self) -> Sequence[pathlib.Path] | None: def resolve_env_akv_ref(self) -> list[str] | None: """ - Return the list of AKV secret URLs, or ``None`` when not configured. + Return the AKV bootstrap secret URLs, or ``None`` when not configured. Returns: - list[str] | None: The configured AKV secret URLs, or ``None``. + list[str] | None: The configured AKV bootstrap secret URLs, or ``None``. """ return self.env_akv_ref @@ -614,6 +645,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/environment_loading.py b/pyrit/setup/environment_loading.py new file mode 100644 index 0000000000..834f66e7d7 --- /dev/null +++ b/pyrit/setup/environment_loading.py @@ -0,0 +1,581 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Load dotenv files and Azure Key Vault-backed environment documents.""" + +import asyncio +import contextlib +import logging +import os +import pathlib +import urllib.parse +from collections.abc import Mapping, Sequence +from io import StringIO +from typing import TYPE_CHECKING + +import dotenv +from dotenv.main import DotEnv +from dotenv.parser import parse_stream + +from pyrit.common import path +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__) + +__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 +_DOTENV_DISABLED_VALUES = frozenset({"1", "true", "t", "yes", "y"}) + + +def validate_env_akv_strict(*, env_akv_strict: object) -> None: + """ + Require a real boolean for Key Vault strict-mode behavior. + + Raises: + TypeError: If env_akv_strict is not a bool. + """ + 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( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool = False, + include_default_base: bool = True, + _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. + + 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 + .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. + _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. + + Raises: + ValueError: If any provided env_files do not exist. + """ + 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: + _load_dotenv_source( + dotenv_path=env_file, + override=env_file.name == ".env.local", + ordinary_candidates=_ordinary_candidates, + override_candidates=_override_candidates, + ) + 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, + 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, +) -> None: + """ + Load one dotenv source and record values that participate in precedence. + + 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 + + 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() + 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)) + + +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_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 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}") + 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, 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 _create_akv_secret_client(*, vault_url: str, credential: "AsyncTokenCredential") -> "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. + + 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. + """ + 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(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 _fetch_akv_document_async( + *, + secret_url: str, + strict: bool = True, + silent: bool = False, +) -> tuple[str, str]: + """ + 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 + 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: + tuple[str, str]: Validated document text and source vault URL. + + Raises: + ImportError: If ``azure-keyvault-secrets`` is not installed. + 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 + + 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=StringIO(validated_document), interpolate=False) + if not parsed_environment: + raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") + return validated_document, 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 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 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.") + 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.") + 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, + ) + 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), + _ordinary_candidates=ordinary_candidates, + _override_candidates=override_candidates, + ) + await _resolve_environment_candidates_async( + process_environment=process_environment, + ordinary_candidates=ordinary_candidates, + override_candidates=override_candidates, + strict=env_akv_strict, + silent=silent, + ) + + +def _parse_akv_reference( + *, + value: str, + variable_name: str, + expected_vault_url: str | None = None, +) -> tuple[str, str, str | None] | None: + """ + Parse and validate an exact whole-value Key Vault reference. + + Returns: + 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, " + "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 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 '{expected_vault_url}', got '{referenced_vault_url}'." + ) + + return referenced_vault_url, secret_name, secret_version + + +async def _resolve_environment_candidates_async( + *, + 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: + """ + Resolve complete Key Vault references from winning environment assignments. + + 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, ())) + ] + 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: + reference = _parse_akv_reference( + value=value, + variable_name=variable_name, + expected_vault_url=expected_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 + 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 + 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: + os.environ.pop(variable_name, None) + + 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, SecretClient] = {} + 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 + 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: + 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 diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index eb0cf04ff8..7d99e2783e 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -1,16 +1,17 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import io import logging import pathlib from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Literal, get_args -import dotenv - -from pyrit.common import path 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, + load_environment_files, + validate_env_akv_strict, +) if TYPE_CHECKING: from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -22,136 +23,7 @@ AZURE_SQL = "AzureSQL" MemoryDatabaseType = Literal["InMemory", "SQLite", "AzureSQL"] - -def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: bool = False) -> None: - """ - 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. - - Raises: - ValueError: If any provided env_files do not exist. - """ - # Validate env_files exist if they were provided - 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 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 - - 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) - - -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 _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. - """ - 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 - - -async def _load_env_from_akv_async(*, secret_urls: Sequence[str], silent: bool = False) -> None: - """ - Load environment variables from Azure Key Vault secrets. - - Each secret's value is treated as the full contents of a ``.env`` file and - parsed accordingly. Later secrets override values from earlier ones. - - 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}]``. - 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. - """ - if not secret_urls: - return - 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) +_load_environment_files = load_environment_files async def _execute_initializers_async(*, initializers: Sequence["PyRITInitializer"]) -> None: @@ -203,6 +75,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: @@ -227,22 +100,28 @@ 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. - 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``. + 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 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 schema migration. Defaults to False. **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. + 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. """ - if env_akv_ref: - await _load_env_from_akv_async(secret_urls=env_akv_ref, silent=silent) - - _load_environment_files(env_files=env_files, silent=silent) + 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, + silent=silent, + ) # Reset all default values before executing initialization scripts # This ensures a clean state for each initialization 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..de9f988a20 --- /dev/null +++ b/tests/unit/build_scripts/test_export_akv_environment.py @@ -0,0 +1,287 @@ +# 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, + _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 = ( + ("# 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() + client.get_secret.return_value = SimpleNamespace(value="resolved-key") + + rendered = _render( + document=document, + credential=mock.MagicMock(), + 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"} + 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 = (content, "https://vault.vault.azure.net") + client = mock.MagicMock() + client.get_secret.return_value = SimpleNamespace(value="resolved-key") + + rendered = _render( + document=document, + credential=mock.MagicMock(), + clients={document[1]: 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_non_strict_warns_and_skips_invalid_reference(caplog: pytest.LogCaptureFixture, capsys) -> None: + 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( + document=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_document") 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_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") + + 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/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 99bd2c5fbc..fa0c81be60 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -42,8 +42,14 @@ 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 + @pytest.mark.parametrize("invalid_value", ["false", "true", 0, 1, None, [], {}]) + 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.""" for db_type in ["in_memory", "sqlite", "azure_sql"]: @@ -147,6 +153,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 +162,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): @@ -231,6 +239,21 @@ def test_from_yaml_file(self): finally: pathlib.Path(yaml_path).unlink() + def test_from_yaml_rejects_quoted_env_akv_strict(self, tmp_path): + yaml_path = tmp_path / "quoted-boolean.yaml" + yaml_path.write_text('env_akv_strict: "false"\n', encoding="utf-8") + + 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_env_akv_strict(self, tmp_path): + yaml_path = tmp_path / "native-booleans.yaml" + 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 + 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" @@ -307,14 +330,28 @@ 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", - ] + """Test that the configured AKV references are returned unchanged.""" + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] 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] + + 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: @@ -334,22 +371,25 @@ 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") 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) + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + 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_environment_loading.py b/tests/unit/setup/test_environment_loading.py new file mode 100644 index 0000000000..c8a29080e6 --- /dev/null +++ b/tests/unit/setup/test_environment_loading.py @@ -0,0 +1,1135 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import os +import pathlib +import tempfile +import types +import warnings +from unittest import mock + +import pytest +from azure.core.exceptions import ResourceNotFoundError + +from pyrit.exceptions import KeyVaultInitializationException +from pyrit.setup.environment_loading import ( + _fetch_akv_document_async, + _parse_akv_reference, + _parse_akv_secret_url, + _warn_about_dotenv_file, + load_environment_async, + load_environment_files, +) + + +class TestLoadEnvironmentFiles: + """Tests for load_environment_files and the env_files initialization parameter.""" + + @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) + + 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 + + @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) + + 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): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_file.write_text("VAR=base") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + caplog.at_level("WARNING", logger="pyrit.setup.environment_loading"), + warnings.catch_warnings(), + ): + warnings.simplefilter("error", DeprecationWarning) + load_environment_files(env_files=None) + + output = capsys.readouterr().out + 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_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 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_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"): + _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): + 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.environment_loading._fetch_akv_document_async", + 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, + ): + await load_environment_async( + env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=None, + env_akv_strict=True, + 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 + + 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_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" + 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 + + @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_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() + + @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" + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=bootstrap_document)) + + with tempfile.TemporaryDirectory() as temp_dir: + local_file = pathlib.Path(temp_dir) / ".env.local" + local_file.write_text("API_KEY=local-key\n", encoding="utf-8") + + 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=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=[local_file], + env_akv_strict=True, + silent=True, + ) + + assert os.environ["API_KEY"] == "local-key" + assert os.environ["AKV_ONLY"] == "akv" + + client.get_secret.assert_awaited_once_with("bootstrap", 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" + + 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) + 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") + + 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=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=[ordinary_file, local_file], + env_akv_strict=True, + silent=True, + ) + + assert os.environ["A"] == "literal" + assert os.environ["B"] == "resolved-for-B" + + assert client.get_secret.await_args_list == [ + mock.call("bootstrap", version=None), + mock.call("key", version=None), + ] + + async def test_non_strict_runtime_uses_file_candidate_after_malformed_akv_winner(self): + credential, client = _create_mock_akv_clients() + + with tempfile.TemporaryDirectory() as temp_dir: + 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, + 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=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=[ordinary_file], + env_akv_strict=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "resolved-fallback" + + 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: + 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.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) + 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 load_environment_async( + env_akv_ref=None, + env_files=[env_file], + env_akv_strict=True, + silent=True, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" + + @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) / 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, 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], + env_akv_strict=True, + 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, + 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, + silent=True, + ) + + 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") + process_value = "kv:https://process-vault.vault.azure.net/secrets/do-not-resolve" + + with ( + 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.environment_loading"), + ): + await load_environment_async( + env_akv_ref=None, + env_files=[env_local_file], + env_akv_strict=False, + silent=False, + ) + + assert os.environ["API_KEY"] == process_value + + 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 + + @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, + 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, + 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, + 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, + 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, + 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") + 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, + 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.""" + nonexistent = pathlib.Path("/nonexistent/path/.env") + + 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) + 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(value=f"{prefix}:{secret_url}", variable_name="API_KEY") == ( + "https://myvault.vault.azure.net", + "api-key", + 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=value, variable_name="API_KEY") is None + + @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", + [ + "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_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 _fetch_akv_document_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_fetch_akv_document_async_returns_validated_document(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(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.environment_loading._print_msg") as mock_print_msg, + ): + document = await _fetch_akv_document_async(secret_url=secret_url, silent=True) + + assert document == (root_document, "https://myvault.vault.azure.net") + assert os.environ == {"SOURCE_VALUE": "ambient-value"} + + mock_credential_cls.assert_called_once_with() + _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_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( + 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_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_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")) + + 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_environment_async( + env_akv_ref=["https://myvault.vault.azure.net/secrets/bootstrap"], + env_files=[], + env_akv_strict=True, + 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_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}")) + + 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_environment_async( + env_akv_ref=["https://myvault.vault.azure.net/secrets/bootstrap"], + env_files=[], + env_akv_strict=True, + silent=True, + ) + + @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_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), + 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 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() + 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_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_runtime_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_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_runtime_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_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_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)) + + 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"), + ): + fetched_document = await _fetch_akv_document_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=False, + ) + + assert fetched_document == ( + "GOOD=resolved\nOTHER=also-resolved", + "https://myvault.vault.azure.net", + ) + assert 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_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")) + + 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 _fetch_akv_document_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_runtime_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_environment_async( + env_akv_ref=["https://myvault.vault.azure.net/secrets/bootstrap"], + env_files=[], + env_akv_strict=True, + silent=True, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["BAD"] == "kv:https://myvault.vault.azure.net/secrets/missing-value" diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index b919df4338..b7b8518104 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -3,9 +3,7 @@ import os import pathlib -import sys import tempfile -import types from unittest import mock import pytest @@ -14,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.initialization import _load_env_from_akv_async, _load_environment_files, _parse_akv_secret_url class TestLoadInitializersFromScripts: @@ -122,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.initialization._load_environment_files") - 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) + 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.initialization._load_environment_files") - 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( @@ -156,65 +153,55 @@ 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) - async def test_invalid_memory_type_raises_error(self): + @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") # 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") - @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"] - - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs) - - mock_load_akv.assert_awaited_once() - assert mock_load_akv.await_args.kwargs["secret_urls"] == refs - assert mock_load_akv.await_args.kwargs["silent"] is False - mock_load_env.assert_called_once() - mock_set_memory.assert_called_once() + mock_load_environment.assert_awaited_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @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_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_akv.assert_not_called() - mock_load_env.assert_called_once() + @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_akv_strict=False, + silent=True, + load_defaults=False, + ) + + 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() - @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] = [] + @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( + "pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock + ) as mock_load_environment: + 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, + ) - async def _record_akv_call(*, secret_urls, silent=False): - call_order.append("akv") - - def _record_env_file_call(*, env_files, silent=False): - call_order.append("env_files") - - refs = ["https://vault.vault.azure.net/secrets/test-secret"] - - with ( - 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) - - assert call_order == ["akv", "env_files"] - mock_set_memory.assert_called_once() + mock_load_environment.assert_not_awaited() @pytest.fixture @@ -237,227 +224,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_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) + 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_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) + 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.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): - """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) - _load_environment_files(env_files=None) - - # Verify both files were loaded - 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 - - @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): - """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 - - _load_environment_files(env_files=None) - - # Verify only one file was loaded - 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") - async def test_loads_custom_env_files_in_order(self, mock_load_dotenv): - """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") - - # Pass custom files - _load_environment_files(env_files=[env1, env2, env3]) - - # 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] - - 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]) - - 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.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): - """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") - - 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]) - - # 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 - - -class TestAkvEnvironmentLoading: - """Tests for AKV URL parsing and env loading helpers.""" - - 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 - - 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]] = [] - - class FakeSecretClient: - def __init__(self, *, vault_url, credential): - client_calls.append(("init", vault_url, credential)) - - async def get_secret(self, name, version=None): - client_calls.append(("get_secret", name, version)) - return types.SimpleNamespace(value="AKV_VAR=from_secret\n") - - 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") - - identity_aio_module.DefaultAzureCredential = FakeCredential - keyvault_secrets_aio_module.SecretClient = FakeSecretClient - - 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("pyrit.setup.initialization.dotenv.load_dotenv") as mock_load_dotenv, - mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, - ): - await _load_env_from_akv_async( - secret_urls=["https://myvault.vault.azure.net/secrets/my-secret/v1"], - 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") - - 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