diff --git a/.env_example b/.env_example index 0d70a48cfd..039544d8a4 100644 --- a/.env_example +++ b/.env_example @@ -1,138 +1,176 @@ # ============================================================================ + # PyRIT Environment File Example + # ============================================================================ + # -# Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need. + +# Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need + # -# MOST USERS ONLY NEED 3 VARIABLES to get started: + +# MOST USERS ONLY NEED 3 VARIABLES to get started + # -# OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" # or any OpenAI-compatible API -# OPENAI_CHAT_KEY="your-key-here" -# OPENAI_CHAT_MODEL="gpt-4o" + +# OPENAI_CHAT_ENDPOINT="" # or any OpenAI-compatible API + +# OPENAI_CHAT_KEY="your-key-here" + +# OPENAI_CHAT_MODEL="gpt-4o" + # + # These work with OpenAI, Azure OpenAI, Ollama, Groq, OpenRouter, and any + # other OpenAI-compatible endpoint. See doc/setup/populating_secrets.md -# for provider-specific examples. + +# for provider-specific examples + # -# If you are using Entra authentication for Azure resources, + +# If you are using Entra authentication for Azure resources + # keys for those resources are not needed. PyRIT auto-detects: if an API key -# is set, it uses key auth; otherwise it falls back to Entra ID automatically. + +# is set, it uses key auth; otherwise it falls back to Entra ID automatically + # -# ============================================================================ +# ============================================================================ ################################### + # OPENAI TARGET SECRETS + # + # The below models work with OpenAIChatTarget - either pass via environment variables + # or copy to OPENAI_CHAT_ENDPOINT + ################################### -PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_CHAT_ENDPOINT="" PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" PLATFORM_OPENAI_CHAT_MODEL="gpt-4o" # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately -# Example: https://xxxx.openai.azure.com/openai/v1 -AZURE_OPENAI_GPT4O_ENDPOINT="https://xxxx.openai.azure.com/openai/v1" + +# Example: + +AZURE_OPENAI_GPT4O_ENDPOINT="" AZURE_OPENAI_GPT4O_KEY="xxxxx" AZURE_OPENAI_GPT4O_MODEL="deployment-name" -# Since Azure deployment name may be custom and differ from the actual underlying model, -# you can specify the underlying model for identifier purposes. If not specified, -# identifiers will default to the value of the standard MODEL environment variable. + +# Since Azure deployment name may be custom and differ from the actual underlying model + +# you can specify the underlying model for identifier purposes. If not specified + +# identifiers will default to the value of the standard MODEL environment variable + AZURE_OPENAI_GPT4O_UNDERLYING_MODEL="gpt-4o" -# Optional second GPT-4o endpoint (that can be used for round-robin distribution). +# Optional second GPT-4o endpoint (that can be used for round-robin distribution) + # TargetInitializer creates RoundRobinTargets that automatically group together + # targets with identical underlying model names and behavioral params, allowing -# for distribution of requests across them for rate-limit relief. -AZURE_OPENAI_GPT4O_ENDPOINT2="https://xxxx.openai.azure.com/openai/v1" + +# for distribution of requests across them for rate-limit relief + +AZURE_OPENAI_GPT4O_ENDPOINT2="" AZURE_OPENAI_GPT4O_KEY2="xxxxx" AZURE_OPENAI_GPT4O_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNDERLYING_MODEL2="gpt-4o" -AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="" AZURE_OPENAI_INTEGRATION_TEST_KEY="xxxxx" AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="" AZURE_OPENAI_GPT3_5_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT3_5_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT3_5_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT4_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT4_CHAT_ENDPOINT="" AZURE_OPENAI_GPT4_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT5_4_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT5_4_ENDPOINT="" AZURE_OPENAI_GPT5_4_KEY="xxxxx" AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" # Endpoints that host models with fewer safety mechanisms (e.g. via adversarial fine tuning -# or content filters turned off) can be defined below and used in adversarial attack testing scenarios. -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" + +# or content filters turned off) can be defined below and used in adversarial attack testing scenarios + +AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2="" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" # Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) + # Default endpoint goes here; specialized ones below -ADVERSARIAL_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" + +ADVERSARIAL_CHAT_ENDPOINT="" ADVERSARIAL_CHAT_KEY="xxxxx" ADVERSARIAL_CHAT_MODEL="deployment-name" -ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="" ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" ADVERSARIAL_CHAT_SINGLETURN_MODEL="deployment-name" -ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="" ADVERSARIAL_CHAT_MULTITURN_KEY="xxxxx" ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" -ADVERSARIAL_CHAT_REASONING_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_REASONING_ENDPOINT="" ADVERSARIAL_CHAT_REASONING_KEY="xxxxx" ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" - # Objective Scorer chat target (used in scorers in scenarios) -OBJECTIVE_SCORER_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" + +OBJECTIVE_SCORER_CHAT_ENDPOINT="" OBJECTIVE_SCORER_CHAT_KEY="xxxxx" OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" -AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="https://xxxxx.eastus2.models.ai.azure.com" +AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="" AZURE_FOUNDRY_DEEPSEEK_KEY="xxxxx" AZURE_FOUNDRY_DEEPSEEK_MODEL="" -AZURE_FOUNDRY_PHI4_ENDPOINT="https://xxxxx.models.ai.azure.com" +AZURE_FOUNDRY_PHI4_ENDPOINT="" AZURE_CHAT_PHI4_KEY="xxxxx" AZURE_CHAT_PHI4_MODEL="" -AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="https://xxxxx.services.ai.azure.com/openai/v1/" +AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="" AZURE_FOUNDRY_MISTRAL_LARGE_KEY="xxxxx" AZURE_FOUNDRY_MISTRAL_LARGE_MODEL="Mistral-Large-3" -AWS_ENDPOINT="https://bedrock-mantle.us-east-1.api.aws/v1" +AWS_ENDPOINT="" AWS_KEY="xxxxx" AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" AWS_RESPONSES_MODEL="openai.gpt-oss-120b" -GROQ_ENDPOINT="https://api.groq.com/openai/v1" +GROQ_ENDPOINT="" GROQ_KEY="gsk_xxxxxxxx" GROQ_LLAMA_MODEL="llama3-8b-8192" -OPEN_ROUTER_ENDPOINT="https://openrouter.ai/api/v1" +OPEN_ROUTER_ENDPOINT="" OPEN_ROUTER_KEY="sk-or-v1-xxxxx" OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" -OLLAMA_CHAT_ENDPOINT="http://127.0.0.1:11434/v1" +OLLAMA_CHAT_ENDPOINT="" OLLAMA_MODEL="llama2" DEFAULT_OPENAI_FRONTEND_ENDPOINT = ${AZURE_OPENAI_GPT4O_AAD_ENDPOINT} @@ -142,25 +180,30 @@ DEFAULT_OPENAI_FRONTEND_MODEL = "gpt-4o" OPENAI_CHAT_ENDPOINT=${PLATFORM_OPENAI_CHAT_ENDPOINT} OPENAI_CHAT_KEY=${PLATFORM_OPENAI_CHAT_KEY} OPENAI_CHAT_MODEL=${PLATFORM_OPENAI_CHAT_MODEL} + # The following line can be populated if using an Azure OpenAI deployment + # where the deployment name differs from the actual underlying model + OPENAI_CHAT_UNDERLYING_MODEL="" ################################## + # OPENAI RESPONSES TARGET SECRETS + ################################## -AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" -AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" +AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="" +AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="" AZURE_OPENAI_GPT5_KEY="xxxxxxx" AZURE_OPENAI_GPT5_MODEL="gpt-5" AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" -PLATFORM_OPENAI_RESPONSES_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_RESPONSES_ENDPOINT="" PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" -AZURE_OPENAI_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_RESPONSES_ENDPOINT="" AZURE_OPENAI_RESPONSES_KEY="xxxxx" AZURE_OPENAI_RESPONSES_MODEL="o4-mini" AZURE_OPENAI_RESPONSES_UNDERLYING_MODEL="o4-mini" @@ -171,10 +214,15 @@ OPENAI_RESPONSES_MODEL=${PLATFORM_OPENAI_RESPONSES_MODEL} OPENAI_RESPONSES_UNDERLYING_MODEL="" ################################## + # OPENAI REALTIME TARGET SECRETS + # + # The below models work with RealtimeTarget - either pass via environment variables + # or copy to OPENAI_REALTIME_ENDPOINT + ################################## PLATFORM_OPENAI_REALTIME_ENDPOINT="wss://api.openai.com/v1" @@ -192,18 +240,23 @@ OPENAI_REALTIME_MODEL = ${PLATFORM_OPENAI_REALTIME_MODEL} OPENAI_REALTIME_UNDERLYING_MODEL = "" ################################## + # IMAGE TARGET SECRETS + # + # The below models work with OpenAIImageTarget - either pass via environment variables + # or copy to OPENAI_IMAGE_ENDPOINT + ################################### -OPENAI_IMAGE_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" +OPENAI_IMAGE_ENDPOINT1 = "" OPENAI_IMAGE_API_KEY1 = "xxxxxx" OPENAI_IMAGE_MODEL1 = "deployment-name" OPENAI_IMAGE_UNDERLYING_MODEL1 = "dall-e-3" -OPENAI_IMAGE_ENDPOINT2 = "https://api.openai.com/v1" +OPENAI_IMAGE_ENDPOINT2 = "" OPENAI_IMAGE_API_KEY2 = "sk-xxxxx" OPENAI_IMAGE_MODEL2 = "dall-e-3" OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" @@ -213,20 +266,24 @@ OPENAI_IMAGE_API_KEY = ${OPENAI_IMAGE_API_KEY2} OPENAI_IMAGE_MODEL = ${OPENAI_IMAGE_MODEL2} OPENAI_IMAGE_UNDERLYING_MODEL = "" - ################################## + # TTS TARGET SECRETS + # + # The below models work with OpenAITTSTarget - either pass via environment variables + # or copy to OPENAI_TTS_ENDPOINT + ################################### -OPENAI_TTS_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" +OPENAI_TTS_ENDPOINT1 = "" OPENAI_TTS_KEY1 = "xxxxxxx" OPENAI_TTS_MODEL1 = "tts" OPENAI_TTS_UNDERLYING_MODEL1 = "tts" -OPENAI_TTS_ENDPOINT2 = "https://api.openai.com/v1" +OPENAI_TTS_ENDPOINT2 = "" OPENAI_TTS_KEY2 = "xxxxxx" OPENAI_TTS_MODEL2 = "tts-1" OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" @@ -237,14 +294,20 @@ OPENAI_TTS_MODEL = ${OPENAI_TTS_MODEL2} OPENAI_TTS_UNDERLYING_MODEL = "" ################################## + # VIDEO TARGET SECRETS + # + # The below models work with OpenAIVideoTarget - either pass via environment variables + # or copy to OPENAI_VIDEO_ENDPOINT + ################################### # Note: Use the base URL without API path -AZURE_OPENAI_VIDEO_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/openai/v1" + +AZURE_OPENAI_VIDEO_ENDPOINT="" AZURE_OPENAI_VIDEO_KEY="xxxxxxx" AZURE_OPENAI_VIDEO_MODEL="sora-2" AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" @@ -254,68 +317,75 @@ OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" - ################################## + # AML TARGET SECRETS + # The below models work with AzureMLChatTarget - either pass via environment variables + # or copy to AZURE_ML_MANAGED_ENDPOINT + ################################### -AZURE_ML_PHI_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +AZURE_ML_PHI_ENDPOINT="" AZURE_ML_PHI_KEY="xxxxx" -# The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed. +# The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed + AZURE_ML_MANAGED_ENDPOINT=${AZURE_ML_PHI_ENDPOINT} AZURE_ML_KEY=${AZURE_ML_PHI_KEY} - ################################## + # MISC TARGET SECRETS -################################### +################################### -OPENAI_COMPLETION_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +OPENAI_COMPLETION_ENDPOINT="" OPENAI_COMPLETION_API_KEY="xxxxx" OPENAI_COMPLETION_MODEL="davinci-002" -OPENAI_EMBEDDING_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +OPENAI_EMBEDDING_ENDPOINT="" OPENAI_EMBEDDING_KEY="xxxxx" OPENAI_EMBEDDING_MODEL="text-embedding-3-small" -AZURE_STORAGE_ACCOUNT_CONTAINER_URL="https://xxxxxx.blob.core.windows.net/xpia" +AZURE_STORAGE_ACCOUNT_CONTAINER_URL="" AZURE_STORAGE_ACCOUNT_SAS_TOKEN="xxxxx" - AZURE_SPEECH_REGION = "eastus2" AZURE_SPEECH_KEY = "xxxxx" + # Resource ID is needed when using Entra authentication + AZURE_SPEECH_RESOURCE_ID = "xxxxx" AZURE_CONTENT_SAFETY_API_KEY="xxxxx" -AZURE_CONTENT_SAFETY_API_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/" +AZURE_CONTENT_SAFETY_API_ENDPOINT="" HUGGINGFACE_TOKEN="hf_xxxxxxx" -HUGGINGFACE_ENDPOINT="https://router.huggingface.co/v1" +HUGGINGFACE_ENDPOINT="" -GOOGLE_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/openai" +GOOGLE_GEMINI_ENDPOINT = "" GOOGLE_GEMINI_API_KEY = "xxxxx" GOOGLE_GEMINI_MODEL="gemini-2.0-flash" - ######################### + # AZURE SQL SECRETS -######################### +######################### # This connects to the test database + AZURE_SQL_DB_CONNECTION_STRING_TEST = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" -AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="https://xxxxx.blob.core.windows.net/dbdata" +AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="" # This connects to the prod database + AZURE_SQL_DB_CONNECTION_STRING_PROD = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" -AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="https://xxxxx.blob.core.windows.net/dbdata" +AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="" +# The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local -# The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local. AZURE_SQL_DB_CONNECTION_STRING = ${AZURE_SQL_DB_CONNECTION_STRING_PROD} AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL=${AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD} diff --git a/.pyrit_conf_example b/.pyrit_conf_example index b1b1ecf932..5dc456f25d 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -88,6 +88,12 @@ operation: op_trash_panda # - Omit this field (or set to null): Load default .env and .env.local from ~/.pyrit/ if they exist # - Set to []: Explicitly load NO environment files # - Set to list of paths: Load only the specified files +# - Local files retain standard dotenv parsing and ${NAME} interpolation. +# Key Vault references in local files remain literal. +# - Interpolation follows load order: .env.local can reference .env, but .env +# cannot see variables introduced only by the later .env.local. +# - Loading is non-transactional. If a later source fails, values loaded by +# earlier sources remain in the process environment. # # Example: # env_files: @@ -96,16 +102,43 @@ operation: op_trash_panda # Azure Key Vault Environment References # --------------------------------------- -# List of AKV secret URLs to load during initialization. -# Each secret's value must be the full contents of a .env file. -# Loaded after env_files, so AKV secrets take precedence. +# Ordered AKV secret URLs whose values are bootstrap .env documents. +# Documents load in list order before local files and use standard ${NAME} +# interpolation. Complete values in a bootstrap document may reference a +# scalar secret in that document's vault using a full URL: +# kv:https://my-vault.vault.azure.net/secrets/SECRET_NAME +# Include a version to pin a secret: +# kv:https://my-vault.vault.azure.net/secrets/SECRET_NAME/SECRET_VERSION +# Short secret names such as kv:SECRET_NAME are rejected. +# Cross-vault child references are rejected. +# Only .vault.azure.net, .vault.azure.cn, and .vault.usgovcloudapi.net hosts +# are accepted. Arbitrary HTTPS hosts and malformed secret paths are rejected. +# Referenced secrets are not cached; each kv: occurrence performs a vault read. +# Referenced values are terminal scalars; they are not parsed for more references. +# Source precedence is AKV bootstraps -> ~/.pyrit/.env -> ~/.pyrit/.env.local. +# Explicit env_files replace the default files and load after the AKV bootstrap. +# PyRIT emits a warning when these local files coexist with env_akv_ref so stale +# configuration cannot silently mask or be mistaken for the Key Vault document. +# When migrating, remove or clear ~/.pyrit/.env and ~/.pyrit/.env.local, remove +# explicit env_files if Key Vault should be authoritative, and restart PyRIT. # Authentication uses DefaultAzureCredential (managed identity, Azure CLI, etc.). +# Key Vault operations use up to three retries with exponential backoff and +# raise KeyVaultInitializationException on bootstrap or secret-resolution failure. +# If env_akv_ref and local files are omitted, PyRIT uses existing process +# environment variables and continues initialization. # # Requires: pip install azure-keyvault-secrets # # Example: # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env +# +# Strict validation applies only to the Key Vault bootstrap and is enabled by +# default. Set this to false to skip malformed or valueless bootstrap entries +# with a warning while loading valid entries. Local files retain standard +# python-dotenv parsing regardless of this setting. +# Empty assignments (NAME=) and child secrets containing an empty string are valid. +# env_akv_strict: false # Max Concurrent Scenario Runs # ---------------------------- diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 878246c140..2cdb0bdd13 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -32,25 +32,35 @@ When PyRIT initializes, environment variables are loaded in a specific order. ** ```{mermaid} flowchart LR - A["1. System Environment"] --> B{"env_files in .pyrit_conf?"} - B -->|No| C["2. ~/.pyrit/.env"] - C --> D["3. ~/.pyrit/.env.local"] - B -->|Yes| E["2. Your specified files (in order)"] + A["System environment"] --> B{"env_akv_ref configured?"} + B -->|Yes| C["AKV bootstrap documents in order"] + B -->|No| D{"Explicit env_files?"} + C --> D + D -->|Yes| E["Explicit files in order"] + D -->|No| F["~/.pyrit/.env"] + F --> G["~/.pyrit/.env.local"] ``` -**Default behavior** (no `env_files` field in `.pyrit_conf`): +System environment variables are always the baseline. If no AKV bootstrap document or environment file is available, PyRIT continues initialization using the existing process environment only. + +**Default file behavior** (no `env_akv_ref` or `env_files` field in `.pyrit_conf`): | Priority | Source | Description | -|----------|--------|-------------| +| ---------- | -------- | ------------- | | Lowest | System environment variables | Always loaded as the baseline | | Medium | `~/.pyrit/.env` | Default config file (loaded if it exists) | | Highest | `~/.pyrit/.env.local` | Local overrides (loaded if it exists) | -**Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. Default paths are completely ignored. +**AKV behavior** (with `env_akv_ref`): The referenced secrets load in list order before local files. Unless custom `env_files` are configured, `~/.pyrit/.env` loads afterward and `~/.pyrit/.env.local` loads last. Later bootstrap documents and local files may override earlier values. + +PyRIT emits a warning when `env_akv_ref` is selected and default or explicit environment files coexist with it. The warning distinguishes files that are ignored from files that load afterward and override Key Vault values, making stale migration files visible at startup. When migrating to Key Vault, clear or remove `~/.pyrit/.env` and `~/.pyrit/.env.local`, remove explicit `env_files` when Key Vault should be the only source, and re-initialize PyRIT so values already present in the process environment cannot mask the Key Vault configuration. + +**Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. They override Key Vault bootstrap values when both fields are configured, and default paths are completely ignored. ### Using .env.local for Overrides You can use `~/.pyrit/.env.local` to override values in `~/.pyrit/.env` without modifying the base file. This is useful for: + - Testing different targets - Using personal credentials instead of shared ones - Switching between configurations quickly @@ -107,7 +117,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 | @@ -159,11 +169,11 @@ initialization_scripts: Environment file paths to load during initialization. Later files override values from earlier files. -| Value | Behavior | -| ----------------- | -------------------------------------------------------------------- | -| Omitted or `null` | Load default `~/.pyrit/.env` and `~/.pyrit/.env.local` if they exist | -| `[]` (empty list) | Load **no** environment files | -| List of paths | Load **only** the specified files (defaults are skipped) | +| Value | Behavior | +| ----------------- | -------------------------------------------------------- | +| Omitted or `null` | Load default `~/.pyrit/.env` and `~/.pyrit/.env.local` | +| `[]` (empty list) | Load **no** environment files | +| List of paths | Load **only** the specified files (defaults are skipped) | ```yaml env_files: @@ -171,6 +181,67 @@ env_files: - /path/to/.env.local ``` +Local environment files use standard python-dotenv parsing and `${NAME}` interpolation. Interpolation follows assignment and file load order. The default `.env.local` can reference a value loaded earlier from `.env`, for example `FOOBAR=${OPENAI_CHAT_ENDPOINT}`. A `.env` value cannot reference a variable introduced only by the later `.env.local`; values are not resolved retroactively. Explicit `env_files` follow the order in which they are listed. + +`env_akv_strict` does not apply to local files: malformed local lines retain python-dotenv's existing permissive skip-and-warn behavior. Local `kv:`, `akv:`, `azure_key_vault:`, and `env_akv_ref:` values remain literal; child-secret resolution is limited to Key Vault bootstrap documents. PyRIT does not define `env:` or `literal:` interpolation syntax. Use standard `${NAME}` interpolation instead. + +Environment loading preserves the historical non-transactional dotenv behavior. Each bootstrap document and local file updates `os.environ` as it loads. If a later source or child-secret lookup fails, assignments made by earlier sources remain in the process environment. + +When `env_akv_ref` is not configured, an empty `env_files` list or missing default files leaves existing process environment variables unchanged and initialization continues. + +### `env_akv_ref` + +Ordered Azure Key Vault secret URLs used to obtain bootstrap environment documents. Each secret value must contain dotenv-formatted entries. Authentication uses `DefaultAzureCredential`. + +```yaml +env_akv_ref: + - https://my-vault.vault.azure.net/secrets/shared-pyrit-env + - https://my-vault.vault.azure.net/secrets/team-pyrit-env +``` + +Bootstrap documents load in list order with `override=True`; local environment files load afterward. Each document uses native dotenv interpolation against the process environment and assignments already parsed. A bootstrap document can mix literal values, `${NAME}` interpolation, and complete-value references to scalar secrets in the same vault: + +```dotenv +OPENAI_CHAT_ENDPOINT="https://example.openai.azure.com/openai/v1" +OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key" +PINNED_OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version-id" +OPENAI_CHAT_MODEL="${PYRIT_OPENAI_CHAT_MODEL}" +``` + +Resolution is limited to one child-secret lookup: + +1. PyRIT validates and loads the bootstrap dotenv document. +2. For each complete-value Key Vault reference in that document, PyRIT fetches the same-vault scalar secret and replaces the environment value. + +For example, if `OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key"`, the value of the `openai-chat-key` secret becomes `OPENAI_CHAT_KEY` verbatim. If that secret happens to contain `kv:another-secret`, the final environment value is the string `kv:another-secret`; PyRIT does not fetch `another-secret`. + +References must occupy the entire value. `kv:` is the canonical Key Vault prefix; `akv:`, `azure_key_vault:`, and `env_akv_ref:` are accepted aliases. + +A Key Vault reference must use a full HTTPS secret URL from the bootstrap document's vault. Supported vault DNS suffixes are `.vault.azure.net`, `.vault.azure.cn`, and `.vault.usgovcloudapi.net`. An unversioned URL reads the latest secret version at initialization. Include the version in the URL to pin it. Short names, malformed paths, arbitrary hosts, and cross-vault child references are rejected before a client is created. + +PyRIT does not cache referenced secrets. Each `kv:` occurrence in a bootstrap document performs a Key Vault read during initialization, including repeated references to the same URI. A later bootstrap or local file may override a reference after it has already been fetched. + +```dotenv +LATEST_KEY_URI="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key" +PINNED_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version-id" +``` + +The bootstrap documents are held in memory and never written to disk. They load before explicit `env_files` or the default `~/.pyrit/.env` and `~/.pyrit/.env.local`, allowing local values to override shared configuration. + +### `env_akv_strict` + +Controls validation only of the Key Vault bootstrap document and defaults to `true`. It does not change parsing of `.env`, `.env.local`, or explicit `env_files`. + +```yaml +env_akv_strict: false +``` + +In strict mode, any malformed dotenv line or variable without an equals sign stops that bootstrap document before it mutates the environment. Empty assignments such as `OPTIONAL_VALUE=` remain valid and set the variable to an empty string. A referenced Key Vault secret whose value is an empty string is also valid. A missing value represented by `None` is treated as an error. With `env_akv_strict: false`, PyRIT emits a warning containing only malformed line numbers and valueless variable names, skips those entries, and loads the valid assignments. Secret values are never included in the warning. + +Non-strict mode does not suppress Key Vault or reference failures. Missing secrets, invalid `kv:` URLs, and bootstrap documents with no valid assignments still stop initialization. Because loading is non-transactional, values from earlier bootstrap documents remain if a later document fails, and raw values from the current document may remain if a child-secret lookup fails. + +Key Vault clients use an explicit Azure retry policy with up to three retries and exponential backoff. Bootstrap parsing, invalid or missing secrets, authentication, authorization, and Azure transport failures are raised as `KeyVaultInitializationException` with the original exception preserved as the cause. The exception remains `ValueError`-compatible for callers migrating from the previous contract. + ### `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. Configured AKV bootstrap documents load in order, followed by selected environment files 2. Default values are reset 3. Memory database is configured (from `memory_db_type`) 4. Initializers are executed in listed order @@ -299,6 +370,11 @@ initializers: # - /path/to/.env # - /path/to/.env.local +# Optional ordered Azure Key Vault bootstrap environment documents +# env_akv_ref: +# - https://my-vault.vault.azure.net/secrets/my-pyrit-env +# env_akv_strict: false # Optional; defaults to true + # Suppress initialization messages silent: false 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 26a29c45b9..ecd11f0344 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -96,6 +96,9 @@ class ConfigurationLoader(YamlLoadable): None means "use defaults", [] means "load nothing". env_files: List of environment file paths to load. None means "use defaults (.env, .env.local)", [] means "load nothing". + env_akv_ref: Ordered list of Key Vault bootstrap secret URLs. + env_akv_strict: Whether malformed or valueless entries in a Key Vault + bootstrap document should fail initialization. silent: Whether to suppress initialization messages. operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. @@ -135,6 +138,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 @@ -147,8 +151,23 @@ def __post_init__(self) -> None: """Validate and normalize the configuration after loading.""" self._normalize_memory_db_type() self._normalize_initializers() + self._validate_env_akv_ref() self._normalize_server() + def _validate_env_akv_ref(self) -> None: + """ + Validate the Key Vault bootstrap secret reference. + + Raises: + ValueError: If env_akv_ref is not 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 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 +420,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 +436,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 for the ordered Azure Key Vault bootstrap secret URLs. + env_akv_strict: Override for strict Key Vault bootstrap validation. Returns: A merged ConfigurationLoader instance. @@ -477,8 +498,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 +608,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 +640,7 @@ async def initialize_pyrit_async(self) -> None: initializers=resolved_initializers if resolved_initializers else None, env_files=resolved_env_files, env_akv_ref=self.env_akv_ref, + env_akv_strict=self.env_akv_strict, silent=self.silent, ) diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index eb0cf04ff8..5d3c713ec2 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -1,18 +1,25 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio import io import logging +import os import pathlib +import urllib.parse from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Literal, get_args import dotenv +from dotenv.parser import parse_stream from pyrit.common import path from pyrit.common.apply_defaults import reset_default_values +from pyrit.exceptions import KeyVaultInitializationException from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory if TYPE_CHECKING: + from azure.keyvault.secrets.aio import SecretClient + from pyrit.setup.pyrit_initializer import PyRITInitializer logger = logging.getLogger(__name__) @@ -22,8 +29,18 @@ AZURE_SQL = "AzureSQL" MemoryDatabaseType = Literal["InMemory", "SQLite", "AzureSQL"] +_AKV_REFERENCE_PREFIXES = frozenset({"akv", "kv", "azure_key_vault", "env_akv_ref"}) +_AKV_VAULT_DNS_SUFFIXES = frozenset({"vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"}) +_AKV_RETRY_TOTAL = 3 +_AKV_RETRY_BACKOFF_FACTOR = 0.8 -def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: bool = False) -> None: + +def _load_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool = False, + include_default_base: bool = True, +) -> bool: """ Load environment files in the order they are provided. Later files override values from earlier files. @@ -33,11 +50,43 @@ def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: .env and .env.local from PyRIT home directory (only if they exist). silent: If True, suppresses print statements about environment file loading. Defaults to False. + include_default_base: If False and env_files is None, skips the default + .env file while still loading .env.local. Defaults to True. + + Returns: + True if at least one environment file was loaded, otherwise False. Raises: ValueError: If any provided env_files do not exist. """ - # Validate env_files exist if they were provided + selected_files = _select_environment_files( + env_files=env_files, + silent=silent, + include_default_base=include_default_base, + ) + for env_file in selected_files: + dotenv.load_dotenv(dotenv_path=env_file, override=True, interpolate=True) + if not silent: + _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) + + return bool(selected_files) + + +def _select_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool, + include_default_base: bool, +) -> list[pathlib.Path]: + """ + Select and validate environment files without reading their contents. + + Returns: + list[pathlib.Path]: Environment files in load order. + + Raises: + ValueError: If an explicitly provided environment file does not exist. + """ if env_files is not None: if not silent: _print_msg(f"Loading custom environment files: {[str(f) for f in env_files]}", quiet=silent, log=True) @@ -51,7 +100,7 @@ def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" - if base_file.exists(): + if include_default_base and base_file.exists(): default_files.append(base_file) if local_file.exists(): default_files.append(local_file) @@ -70,10 +119,7 @@ def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: env_files = default_files - for env_file in env_files: - dotenv.load_dotenv(env_file, override=True, interpolate=True) - if not silent: - _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) + return list(env_files) def _print_msg(message: str, quiet: bool, log: bool) -> None: @@ -91,6 +137,46 @@ def _print_msg(message: str, quiet: bool, log: bool) -> None: logger.info(message) +def _warn_about_akv_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool = False, +) -> None: + """Warn when local environment files coexist with an AKV environment source.""" + base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" + local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" + messages: list[str] = [] + + if base_file.exists(): + if env_files is None: + messages.append(f"{base_file} will load after Key Vault and override matching values") + else: + messages.append(f"{base_file} exists but will be ignored because env_files was explicitly configured") + + if local_file.exists(): + if env_files is None: + messages.append(f"{local_file} will load after Key Vault and override matching values") + else: + messages.append(f"{local_file} exists but will be ignored because env_files was explicitly configured") + + if env_files: + messages.append(f"explicit env_files will load after Key Vault and override matching values: {list(env_files)}") + + if not messages: + return + + message = ( + "env_akv_ref is configured, but local environment files were also found:\n- " + + "\n- ".join(messages) + + "\nWhen migrating to Key Vault, clear or remove ~/.pyrit/.env and ~/.pyrit/.env.local, " + "remove explicit env_files when Key Vault should be the only source, and restart PyRIT so stale " + "process values cannot mask Key Vault configuration." + ) + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + + def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: """ Parse an AKV secret URL into vault URL, secret name, and optional version. @@ -105,53 +191,327 @@ def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: Raises: ValueError: If the URL does not match the expected format. """ - parts = secret_url.split("/secrets/") - if len(parts) != 2: - raise ValueError( - f"Invalid AKV secret URL: '{secret_url}'. " - "Expected format: https://{{vault}}.vault.azure.net/secrets/{{name}}[/{{version}}]" - ) - vault_url = parts[0] - name_parts = parts[1].rstrip("/").split("/") - secret_name = name_parts[0] - secret_version = name_parts[1] if len(name_parts) > 1 else None - return vault_url, secret_name, secret_version + error_message = ( + f"Invalid AKV secret URL: '{secret_url}'. Expected an HTTPS Azure Key Vault URL in the format " + "https://{vault}.{vault-dns-suffix}/secrets/{name}[/{version}]." + ) + try: + parsed_url = urllib.parse.urlsplit(secret_url) + port = parsed_url.port + except (TypeError, ValueError) as error: + raise ValueError(error_message) from error + + hostname = parsed_url.hostname + vault_name, separator, dns_suffix = hostname.partition(".") if hostname else ("", "", "") + valid_vault_name = 1 <= len(vault_name) <= 63 and all( + char.isascii() and (char.isalnum() or char == "-") for char in vault_name + ) + valid_authority = ( + parsed_url.scheme.casefold() == "https" + and parsed_url.username is None + and parsed_url.password is None + and port is None + and separator == "." + and dns_suffix in _AKV_VAULT_DNS_SUFFIXES + and valid_vault_name + ) + path_parts = parsed_url.path.split("/") + valid_path = ( + len(path_parts) in {3, 4} and path_parts[0] == "" and path_parts[1] == "secrets" and all(path_parts[2:]) + ) + if not valid_authority or not valid_path or parsed_url.query or parsed_url.fragment: + raise ValueError(error_message) + + secret_name = path_parts[2] + secret_version = path_parts[3] if len(path_parts) == 4 else None + if not _is_valid_akv_identifier(secret_name) or ( + secret_version is not None and not _is_valid_akv_identifier(secret_version) + ): + raise ValueError(error_message) + + return f"https://{hostname}", secret_name, secret_version + + +def _is_valid_akv_identifier(identifier: str) -> bool: + """ + Check whether a Key Vault secret name or version uses URL-safe characters. + + Returns: + bool: True when the identifier is valid. + """ + return 1 <= len(identifier) <= 127 and all( + char.isascii() and (char.isalnum() or char == "-") for char in identifier + ) + +def _create_akv_secret_client(*, vault_url: str, credential: Any) -> "SecretClient": + """ + Create an asynchronous Key Vault client with an explicit retry policy. -async def _load_env_from_akv_async(*, secret_urls: Sequence[str], silent: bool = False) -> None: + Returns: + SecretClient: Configured asynchronous secret client. """ - Load environment variables from Azure Key Vault secrets. + 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) - Each secret's value is treated as the full contents of a ``.env`` file and - parsed accordingly. Later secrets override values from earlier ones. + +def _key_vault_initialization_error(*, message: str, error: Exception) -> KeyVaultInitializationException: + """ + Create a contextual Key Vault exception without losing the original cause. + + Returns: + KeyVaultInitializationException: Wrapped contextual exception. + """ + status_code = getattr(error, "status_code", None) + return KeyVaultInitializationException( + status_code=status_code if isinstance(status_code, int) else 500, + message=f"{message}: {error}", + ) + + +def _validate_dotenv_document( + document: str, + *, + strict: bool = True, + silent: bool = False, +) -> str: + """ + Validate that every dotenv binding uses ``NAME=VALUE`` syntax. + + Args: + document (str): The dotenv document to validate. + strict (bool): If True, reject any invalid entry. If False, warn and + allow python-dotenv to skip invalid entries. Defaults to True. + silent (bool): If True, suppress the console warning. Defaults to False. + + Returns: + str: The original document, or a sanitized document when strict is False. + + Raises: + ValueError: If strict is True and the document contains invalid entries. + """ + bindings = list(parse_stream(io.StringIO(document))) + malformed_lines = [str(binding.original.line) for binding in bindings if binding.error] + valueless_names = [binding.key for binding in bindings if binding.key is not None and binding.value is None] + issues: list[str] = [] + if malformed_lines: + issues.append("malformed entries at lines: " + ", ".join(malformed_lines)) + if valueless_names: + issues.append("variables without values: " + ", ".join(valueless_names)) + if not issues: + return document + + details = "; ".join(issues) + if strict: + raise ValueError("AKV environment document contains " + details) + + message = "AKV environment document contains invalid entries that will be skipped: " + details + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + return "".join( + binding.original.string + for binding in bindings + if not binding.error and not (binding.key is not None and binding.value is None) + ) + + +async def _load_env_from_akv_async( + *, + secret_url: str, + strict: bool = True, + silent: bool = False, +) -> None: + """ + Load a bootstrap dotenv document and resolve its same-vault secret references. + + References are resolved once. Referenced secret values are treated as terminal + strings and are not interpreted as additional references. Authentication uses ``DefaultAzureCredential``, which silently tries managed identity, Azure CLI, VS Code credentials, etc., and falls back to interactive browser authentication when running locally. Args: - secret_urls (Sequence[str]): Sequence of AKV secret URLs to load, each in - the format ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. + secret_url (str): AKV secret URL in the format + ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. + strict (bool): If True, reject malformed or valueless dotenv entries. + If False, warn and skip those entries. Defaults to True. silent (bool): If True, suppresses print statements. Defaults to False. Raises: ImportError: If ``azure-keyvault-secrets`` is not installed. - ValueError: If a secret URL is malformed. + KeyVaultInitializationException: If the root URL is malformed or the bootstrap environment + document cannot be fully resolved. + ValueError: Compatibility base of ``KeyVaultInitializationException``. """ - 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: + 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) - 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) + async with DefaultAzureCredential() as credential: + async with _create_akv_secret_client(vault_url=vault_url, credential=credential) as client: + secret = await client.get_secret(secret_name, version=secret_version) + + if not secret.value: + raise ValueError(f"AKV environment secret has no value: {secret_url}") + + validated_document = _validate_dotenv_document(secret.value, strict=strict, silent=silent) + parsed_environment = dotenv.dotenv_values(stream=io.StringIO(validated_document), interpolate=True) + if not parsed_environment: + raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") + loaded = dotenv.load_dotenv( + stream=io.StringIO(validated_document), + override=True, + interpolate=True, + ) + if not loaded: + return + + for variable_name, value in parsed_environment.items(): + if value is None: + continue + target = _parse_akv_reference(value) + if target is None: + continue + try: + referenced_name, referenced_version = _resolve_akv_secret_reference( + target=target, + variable_name=variable_name, + vault_url=vault_url, + ) + referenced_secret = await client.get_secret(referenced_name, version=referenced_version) + if referenced_secret.value is None: + raise ValueError( + f"AKV secret '{referenced_name}' referenced by environment variable " + f"'{variable_name}' has no value." + ) + os.environ[variable_name] = referenced_secret.value + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to resolve Key Vault reference for environment variable '{variable_name}'", + error=error, + ) + raise wrapped_error from error + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to load Key Vault bootstrap secret '{secret_url}'", + error=error, + ) + raise wrapped_error from error + + +async def _load_environment_async( + *, + env_akv_ref: Sequence[str] | None, + env_files: Sequence[pathlib.Path] | None, + env_akv_strict: bool, + silent: bool, +) -> None: + """ + Load environment sources in precedence order. + + Args: + env_akv_ref (Sequence[str] | None): Optional ordered Key Vault bootstrap secret URLs. + env_files (Sequence[pathlib.Path] | None): Optional ordered local environment files. + env_akv_strict (bool): Whether bootstrap dotenv validation is strict. + silent (bool): Whether initialization messages are suppressed. + + Raises: + ValueError: If a configured source or reference is invalid. + """ + if isinstance(env_akv_ref, str): + raise ValueError("env_akv_ref must be a sequence of Azure Key Vault secret URLs.") + if env_akv_ref: + if any(not isinstance(secret_url, str) or not secret_url.strip() for secret_url in env_akv_ref): + raise ValueError("env_akv_ref must contain only non-empty Azure Key Vault secret URLs.") + await asyncio.to_thread( + _warn_about_akv_environment_files, + env_files=env_files, + silent=silent, + ) + for secret_url in env_akv_ref: + await _load_env_from_akv_async( + secret_url=secret_url, + strict=env_akv_strict, + silent=silent, + ) + + await asyncio.to_thread( + _load_environment_files, + env_files=env_files, + silent=silent, + ) + + +def _parse_akv_reference(value: str) -> str | None: + """ + Parse an exact whole-value Key Vault reference. + + Returns: + The referenced secret URL, or None for a literal value. + """ + prefix, separator, target = value.partition(":") + return target.strip() if separator and prefix in _AKV_REFERENCE_PREFIXES else None + + +def _validate_akv_secret_name(*, secret_name: str, variable_name: str) -> None: + if not _is_valid_akv_identifier(secret_name): + raise ValueError( + f"Invalid same-vault secret name '{secret_name}' referenced by environment variable '{variable_name}'. " + "Secret names must contain only letters, numbers, and hyphens." + ) + + +def _resolve_akv_secret_reference( + *, + target: str, + variable_name: str, + vault_url: str, +) -> tuple[str, str | None]: + """ + Resolve a full same-vault secret URI. + + Args: + target (str): Full Key Vault secret URI. + variable_name (str): The environment variable receiving the secret. + vault_url (str): The bootstrap document's vault URL. + + Returns: + tuple[str, str | None]: Secret name and optional version. + + Raises: + ValueError: If the target is not a full URI, is invalid, or references another vault. + """ + if not target.casefold().startswith("https://"): + raise ValueError( + f"AKV reference for environment variable '{variable_name}' must use a full secret URL, " + "for example kv:https://my-vault.vault.azure.net/secrets/my-secret." + ) + + referenced_vault_url, secret_name, secret_version = _parse_akv_secret_url(target) + if referenced_vault_url.rstrip("/").casefold() != vault_url.rstrip("/").casefold(): + raise ValueError( + f"Cross-vault AKV reference for environment variable '{variable_name}' is not supported. " + f"Expected vault '{vault_url}', got '{referenced_vault_url}'." + ) + + _validate_akv_secret_name(secret_name=secret_name, variable_name=variable_name) + return secret_name, secret_version async def _execute_initializers_async(*, initializers: Sequence["PyRITInitializer"]) -> None: @@ -203,6 +563,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: @@ -229,20 +590,24 @@ async def initialize_pyrit_async( env_files (Sequence[pathlib.Path] | None): Optional sequence of environment file paths to load in order. If not provided, will load default .env and .env.local files from PyRIT home if they exist. All paths must be valid pathlib.Path objects. - env_akv_ref (Sequence[str] | None): Optional sequence of Azure Key Vault secret URLs to load. - 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``. + env_akv_ref (Sequence[str] | None): Optional ordered Azure Key Vault URLs whose secret values + contain bootstrap .env documents. Loaded before ``env_files`` so later bootstrap documents + and local files take precedence. Requires ``azure-keyvault-secrets``. + env_akv_strict (bool): If True, reject malformed or valueless entries in the Key Vault + bootstrap document. If False, warn and skip those entries. Defaults to True. silent (bool): If True, suppresses print statements about environment file loading and 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. + 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) + 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/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..9682e9b2ae 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -42,6 +42,7 @@ def test_default_values(self): assert config.initialization_scripts is None # None means "use defaults" assert config.env_files is None # None means "use defaults" assert config.env_akv_ref is None + assert config.env_akv_strict is True assert config.silent is False def test_valid_memory_db_types_snake_case(self): @@ -147,6 +148,7 @@ def test_from_dict_with_all_fields(self): "initialization_scripts": ["/path/to/script.py"], "env_files": ["/path/to/.env"], "env_akv_ref": ["https://vault.vault.azure.net/secrets/one"], + "env_akv_strict": False, "silent": True, } config = ConfigurationLoader.from_dict(data) @@ -155,6 +157,7 @@ def test_from_dict_with_all_fields(self): assert config.initialization_scripts == ["/path/to/script.py"] assert config.env_files == ["/path/to/.env"] assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] + assert config.env_akv_strict is False assert config.silent is True def test_from_dict_filters_none_values(self): @@ -307,7 +310,7 @@ 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.""" + """Test that the configured AKV references are returned unchanged.""" refs = [ "https://vault.vault.azure.net/secrets/first", "https://vault.vault.azure.net/secrets/second/version", @@ -315,6 +318,14 @@ def testresolve_env_akv_ref_returns_configured_values(self): 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] + @pytest.mark.usefixtures("patch_central_database") class TestConfigurationLoaderInitialization: @@ -334,6 +345,7 @@ async def test_initialize_pyrit_async_basic(self, mock_init): assert call_kwargs["initializers"] is None assert call_kwargs["env_files"] is None assert call_kwargs["env_akv_ref"] is None + assert call_kwargs["env_akv_strict"] is True assert call_kwargs["silent"] is False @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") @@ -343,13 +355,14 @@ async def test_initialize_pyrit_async_with_env_akv_ref(self, mock_init): "https://vault.vault.azure.net/secrets/first", "https://vault.vault.azure.net/secrets/second/version", ] - config = ConfigurationLoader(memory_db_type="in_memory", env_akv_ref=refs) + config = ConfigurationLoader(memory_db_type="in_memory", env_akv_ref=refs, env_akv_strict=False) await config.initialize_pyrit_async() mock_init.assert_called_once() call_kwargs = mock_init.call_args.kwargs assert call_kwargs["env_akv_ref"] == refs + assert call_kwargs["env_akv_strict"] is False @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") @mock.patch("pyrit.registry.InitializerRegistry") diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index b919df4338..0e53da0deb 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -3,18 +3,25 @@ import os import pathlib -import sys import tempfile import types from unittest import mock import pytest +from azure.core.exceptions import ResourceNotFoundError from pyrit.common.apply_defaults import reset_default_values from pyrit.common.singleton import Singleton +from pyrit.exceptions import KeyVaultInitializationException from pyrit.registry import InitializerRegistry from pyrit.setup import IN_MEMORY, initialize_pyrit_async -from pyrit.setup.initialization import _load_env_from_akv_async, _load_environment_files, _parse_akv_secret_url +from pyrit.setup.initialization import ( + _load_env_from_akv_async, + _load_environment_files, + _parse_akv_reference, + _parse_akv_secret_url, + _warn_about_akv_environment_files, +) class TestLoadInitializersFromScripts: @@ -122,16 +129,16 @@ 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") + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) async def test_initialize_basic(self, mock_load_env, mock_set_memory): """Test basic initialization.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, load_defaults=False) mock_load_env.assert_called_once() mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files") + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) async def test_initialize_with_script(self, mock_load_env, mock_set_memory): """Test initialization with a script.""" with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f: @@ -161,59 +168,182 @@ async def initialize_async(self) -> None: finally: os.unlink(script_path) - async def test_invalid_memory_type_raises_error(self): + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) + async def test_invalid_memory_type_raises_error(self, mock_load_env): """Test that invalid memory type raises ValueError.""" with pytest.raises(ValueError, match="is not a supported type"): - await initialize_pyrit_async(memory_db_type="InvalidType") # 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_environment_files", return_value=False) @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, mock_set_memory): - """Test that env_akv_ref triggers AKV env loading.""" - refs = ["https://vault.vault.azure.net/secrets/test-secret"] + """Test that env_akv_ref loads bootstrap secrets in order.""" + refs = [ + "https://vault.vault.azure.net/secrets/first", + "https://vault.vault.azure.net/secrets/second/version", + ] + + mock_load_akv.return_value = None - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) - mock_load_akv.assert_awaited_once() - assert mock_load_akv.await_args.kwargs["secret_urls"] == refs - assert mock_load_akv.await_args.kwargs["silent"] is False + assert mock_load_akv.await_args_list == [ + mock.call(secret_url=refs[0], strict=True, silent=False), + mock.call(secret_url=refs[1], strict=True, silent=False), + ] mock_load_env.assert_called_once() mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files") + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=False) @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) async def test_initialize_with_empty_env_akv_ref_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=[]) + """Test that an empty env_akv_ref list skips AKV loading.""" + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=[], load_defaults=False) mock_load_akv.assert_not_called() mock_load_env.assert_called_once() mock_set_memory.assert_called_once() + @pytest.mark.parametrize("env_akv_ref", ["https://vault.vault.azure.net/secrets/one", [""], [None]]) + async def test_initialize_rejects_invalid_env_akv_ref(self, env_akv_ref): + with pytest.raises(ValueError, match="env_akv_ref must"): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=env_akv_ref, # type: ignore[arg-type] + load_defaults=False, + ) + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_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] = [] + async def test_initialize_keeps_akv_values_when_local_file_loading_fails(self, mock_set_memory): + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + nonexistent = pathlib.Path("/nonexistent/.env") - async def _record_akv_call(*, secret_urls, silent=False): - call_order.append("akv") + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + side_effect=lambda **_: os.environ.update({"FROM_AKV": "resolved"}), + ), + pytest.raises(ValueError, match="Environment file not found"), + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=refs, + env_files=[nonexistent], + load_defaults=False, + ) + + assert os.environ["FROM_AKV"] == "resolved" + + mock_set_memory.assert_not_called() - def _record_env_file_call(*, env_files, silent=False): - call_order.append("env_files") + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_loads_local_overrides_on_akv_environment(self, mock_set_memory): + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + with tempfile.TemporaryDirectory() as temp_dir: + local_file = pathlib.Path(temp_dir) / ".env.local" + local_file.write_text("DERIVED=${BASE}\nBASE=local") + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + side_effect=lambda **_: os.environ.update({"BASE": "akv", "ONLY_AKV": "shared"}), + ), + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=refs, + env_files=[local_file], + load_defaults=False, + ) + + assert os.environ["BASE"] == "local" + assert os.environ["DERIVED"] == "akv" + assert os.environ["ONLY_AKV"] == "shared" - refs = ["https://vault.vault.azure.net/secrets/test-secret"] + mock_set_memory.assert_called_once() - 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) + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_default_files_override_akv_in_order(self, mock_set_memory): + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VALUE=env") + (temp_path / ".env.local").write_text("VALUE=local") + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + side_effect=lambda **_: os.environ.update({"VALUE": "akv"}), + ), + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=refs, + load_defaults=False, + silent=True, + ) + + assert os.environ["VALUE"] == "local" + + mock_set_memory.assert_called_once() + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_resolves_bootstrap_references_before_local_overrides(self, mock_set_memory): + refs = ["https://vault.vault.azure.net/secrets/bootstrap"] + with tempfile.TemporaryDirectory() as temp_dir: + local_file = pathlib.Path(temp_dir) / ".env.local" + local_file.write_text( + "OVERRIDDEN=local\n" + "LOCAL_SECRET=kv:https://vault.vault.azure.net/secrets/local-secret\n" + "LOCAL_ENV=env:BOOTSTRAP_SOURCE" + ) + bootstrap_environment = { + "OVERRIDDEN": "unused-secret-value", + "BOOTSTRAP_SECRET": "bootstrap-secret-value", + "BOOTSTRAP_SOURCE": "bootstrap-value", + } + + with ( + mock.patch.dict(os.environ, {"SOURCE_VALUE": "ambient-value"}, clear=True), + mock.patch("pyrit.setup.initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + side_effect=lambda **_: os.environ.update(bootstrap_environment), + ), + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=refs, + env_files=[local_file], + load_defaults=False, + ) + + assert os.environ["OVERRIDDEN"] == "local" + assert os.environ["BOOTSTRAP_SECRET"] == "bootstrap-secret-value" + assert os.environ["LOCAL_SECRET"] == "kv:https://vault.vault.azure.net/secrets/local-secret" + assert os.environ["LOCAL_ENV"] == "env:BOOTSTRAP_SOURCE" + + mock_set_memory.assert_called_once() + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_without_environment_file_uses_system_environment(self, mock_set_memory): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[], load_defaults=False) - assert call_order == ["akv", "env_files"] mock_set_memory.assert_called_once() @@ -237,16 +367,18 @@ def setup_method(self) -> None: """Clear default values before each test.""" reset_default_values() - async def test_initialize_silent_produces_no_output(self, capsys): + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=True) + async def test_initialize_silent_produces_no_output(self, mock_load_env, capsys): """initialize_pyrit_async with silent=True must not print anything to stdout.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=True) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=True, load_defaults=False) captured = capsys.readouterr() assert captured.out == "" - async def test_initialize_not_silent_prints_migration_message(self, capsys): + @mock.patch("pyrit.setup.initialization._load_environment_files", return_value=True) + async def test_initialize_not_silent_prints_migration_message(self, mock_load_env, capsys): """Without silent, the Alembic schema-check message is printed and tagged as Alembic output.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=False) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=False, load_defaults=False) captured = capsys.readouterr() assert "[pyrit:alembic] No new upgrade operations detected." in captured.out @@ -255,53 +387,124 @@ async def test_initialize_not_silent_prints_migration_message(self, capsys): class TestLoadEnvironmentFiles: """Tests for _load_environment_files function and env_files parameter in initialize_pyrit_async.""" - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_loads_default_env_files_when_none_provided(self, mock_config_path, mock_load_dotenv): + async def test_loads_default_env_files_when_none_provided(self, mock_config_path): """Test that default .env and .env.local files are loaded when env_files is None.""" - # Create temporary directory and files with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) env_file = temp_path / ".env" env_local_file = temp_path / ".env.local" - - # Create the files env_file.write_text("VAR1=value1") env_local_file.write_text("VAR2=value2") - - # Mock CONFIGURATION_DIRECTORY_PATH to point to our temp directory mock_config_path.__truediv__ = lambda self, other: temp_path / other - # Call the function with None (default behavior) - _load_environment_files(env_files=None) + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None) - # Verify both files were loaded - assert mock_load_dotenv.call_count == 2 - calls = [call[0][0] for call in mock_load_dotenv.call_args_list] - assert env_file in calls - assert env_local_file in calls + assert loaded is True + assert os.environ["VAR1"] == "value1" + assert os.environ["VAR2"] == "value2" - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_only_loads_existing_default_files(self, mock_config_path, mock_load_dotenv): + async def test_only_loads_existing_default_files(self, mock_config_path): """Test that only existing default files are loaded.""" with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) env_file = temp_path / ".env" - - # Only create .env, not .env.local env_file.write_text("VAR1=value1") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None) + + assert loaded is True + assert os.environ["VAR1"] == "value1" + + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_excludes_default_env_when_loading_local_override(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text("VAR=base") + env_local_file.write_text("VAR=local") + + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None, include_default_base=False) + + assert loaded is True + assert os.environ["VAR"] == "local" + + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_returns_false_when_no_default_files_exist(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None) + + assert loaded is False + assert os.environ == {} + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_warns_when_default_files_coexist_with_akv(self, mock_config_path, caplog, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text("VAR=base") + env_local_file.write_text("VAR=local") mock_config_path.__truediv__ = lambda self, other: temp_path / other - _load_environment_files(env_files=None) + with caplog.at_level("WARNING", logger="pyrit.setup.initialization"): + _warn_about_akv_environment_files(env_files=None) + + output = capsys.readouterr().out + assert output.startswith("WARNING: env_akv_ref is configured") + assert f"{env_file} will load after Key Vault and override matching values" in output + assert f"{env_local_file} will load after Key Vault and override matching values" in output + assert "clear or remove ~/.pyrit/.env and ~/.pyrit/.env.local" in output + assert "remove explicit env_files when Key Vault should be the only source" in output + assert "restart PyRIT" in output + assert caplog.records[0].levelname == "WARNING" - # 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.path.CONFIGURATION_DIRECTORY_PATH") + def test_warns_when_explicit_files_replace_defaults_with_akv(self, mock_config_path, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + custom_file = temp_path / ".env.custom" + env_file.write_text("VAR=base") + env_local_file.write_text("VAR=local") + custom_file.write_text("VAR=custom") + mock_config_path.__truediv__ = lambda self, other: temp_path / other - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") - async def test_loads_custom_env_files_in_order(self, mock_load_dotenv): + _warn_about_akv_environment_files(env_files=[custom_file]) + + output = capsys.readouterr().out + assert f"{env_file} exists but will be ignored because env_files was explicitly configured" in output + assert f"{env_local_file} exists but will be ignored because env_files was explicitly configured" in output + assert f"explicit env_files will load after Key Vault and override matching values: {[custom_file]}" in output + + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_akv_environment_file_warning_respects_silent(self, mock_config_path, caplog, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VAR=base") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with caplog.at_level("WARNING", logger="pyrit.setup.initialization"): + _warn_about_akv_environment_files(env_files=None, silent=True) + + assert capsys.readouterr().out == "" + assert "will load after Key Vault and override matching values" in caplog.text + assert "restart PyRIT" in caplog.text + + async def test_loads_custom_env_files_in_order(self): """Test that custom env_files are loaded in the order provided.""" with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -314,13 +517,105 @@ async def test_loads_custom_env_files_in_order(self, mock_load_dotenv): env2.write_text("VAR=prod") env3.write_text("VAR=local") - # Pass custom files - _load_environment_files(env_files=[env1, env2, env3]) + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env1, env2, env3]) + + assert loaded is True + assert os.environ["VAR"] == "local" + + async def test_load_environment_files_interpolates_in_assignment_order(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("A=one\nB=${A}\nA=two\nC=${A}") + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert os.environ["A"] == "two" + assert os.environ["B"] == "one" + assert os.environ["C"] == "two" + + async def test_load_environment_files_honors_python_dotenv_disabled(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("DISABLED_VALUE=not-loaded") + + with mock.patch.dict(os.environ, {"PYTHON_DOTENV_DISABLED": "true"}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert "DISABLED_VALUE" not in os.environ + + async def test_direct_local_file_loader_keeps_pyrit_references_literal(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text( + "BASE_VALUE=base\nKV_REFERENCE=kv:api-key\nENV_REFERENCE=env:SOURCE_VALUE\nINTERPOLATED=${BASE_VALUE}" + ) + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert os.environ["KV_REFERENCE"] == "kv:api-key" + assert os.environ["ENV_REFERENCE"] == "env:SOURCE_VALUE" + assert os.environ["INTERPOLATED"] == "base" - # 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] + @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_default_local_file_can_interpolate_base_file_but_not_reverse(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text( + "OPENAI_CHAT_ENDPOINT=https://example.openai.azure.com/openai/v1\nFROM_LATER_LOCAL=${LOCAL_ONLY}" + ) + env_local_file.write_text("FOOBAR=${OPENAI_CHAT_ENDPOINT}\nLOCAL_ONLY=local") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None, silent=True) + + assert loaded is True + assert os.environ["FOOBAR"] == "https://example.openai.azure.com/openai/v1" + assert os.environ["FROM_LATER_LOCAL"] == "" + assert os.environ["LOCAL_ONLY"] == "local" + + async def test_env_akv_strict_does_not_validate_local_environment_files(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("GOOD=resolved\n=malformed\nOTHER=also-resolved") + + with mock.patch.dict(os.environ, {}, clear=True): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_files=[env_file], + env_akv_strict=True, + load_defaults=False, + silent=True, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_keeps_local_akv_reference_literal_without_bootstrap(self, mock_set_memory): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("API_KEY=kv:https://myvault.vault.azure.net/secrets/api-key") + + with mock.patch.dict(os.environ, {}, clear=True): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_files=[env_file], + load_defaults=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "kv:https://myvault.vault.azure.net/secrets/api-key" + + mock_set_memory.assert_called_once() async def test_raises_error_for_nonexistent_env_file(self): """Test that ValueError is raised for non-existent env file.""" @@ -338,7 +633,7 @@ async def test_initialize_pyrit_with_custom_env_files(self, mock_set_memory): env_file.write_text("CUSTOM_VAR=custom_value") # Should not raise an error - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[env_file]) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[env_file], load_defaults=False) mock_set_memory.assert_called_once() @@ -350,10 +645,8 @@ async def test_initialize_pyrit_raises_for_nonexistent_env_file(self, mock_set_m with pytest.raises(ValueError, match="Environment file not found"): await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[nonexistent]) - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") - @mock.patch("pyrit.setup.initialization.path.HOME_PATH") @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_custom_env_files_override_default_behavior(self, mock_set_memory, mock_home_path, mock_load_dotenv): + async def test_custom_env_files_override_default_behavior(self, mock_set_memory): """Test that passing custom env_files prevents loading default files.""" with tempfile.TemporaryDirectory() as temp_dir: temp_path = pathlib.Path(temp_dir) @@ -368,19 +661,62 @@ async def test_custom_env_files_override_default_behavior(self, mock_set_memory, custom_env = temp_path / ".env.custom" custom_env.write_text("CUSTOM=value") - mock_home_path.__truediv__ = lambda self, other: temp_path / other - - # 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 + with mock.patch.dict(os.environ, {}, clear=True): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[custom_env], load_defaults=False) + + assert os.environ["CUSTOM"] == "value" + assert "DEFAULT" not in os.environ + assert "DEFAULT_LOCAL" not in os.environ + + +def _create_mock_akv_clients() -> tuple[mock.MagicMock, mock.MagicMock]: + credential = mock.MagicMock() + credential.__aenter__ = mock.AsyncMock(return_value=credential) + credential.__aexit__ = mock.AsyncMock(return_value=None) + client = mock.MagicMock() + client.__aenter__ = mock.AsyncMock(return_value=client) + client.__aexit__ = mock.AsyncMock(return_value=None) + return credential, client + + +def _assert_mock_akv_client_created( + mock_client_cls: mock.MagicMock, + *, + vault_url: str, + credential: mock.MagicMock, +) -> None: + mock_client_cls.assert_called_once() + call_kwargs = mock_client_cls.call_args.kwargs + assert call_kwargs["vault_url"] == vault_url + assert call_kwargs["credential"] is credential + retry_policy = call_kwargs["retry_policy"] + assert retry_policy.total_retries == 3 + assert retry_policy.connect_retries == 3 + assert retry_policy.read_retries == 3 + assert retry_policy.status_retries == 3 + assert retry_policy.backoff_factor == 0.8 class TestAkvEnvironmentLoading: """Tests for AKV URL parsing and env loading helpers.""" + @pytest.mark.parametrize("prefix", ["kv", "akv", "azure_key_vault", "env_akv_ref"]) + def test_parse_akv_reference_accepts_aliases(self, prefix): + secret_url = "https://myvault.vault.azure.net/secrets/api-key" + + assert _parse_akv_reference(f"{prefix}:{secret_url}") == secret_url + + @pytest.mark.parametrize( + "value", + [ + "env:SOURCE_VALUE", + "literal:kv:https://myvault.vault.azure.net/secrets/api-key", + "@Microsoft.KeyVault(SecretUri=https://myvault.vault.azure.net/secrets/api-key)", + ], + ) + def test_parse_akv_reference_ignores_non_akv_syntax(self, value): + assert _parse_akv_reference(value) is None + def test_parse_akv_secret_url_with_version(self): url = "https://myvault.vault.azure.net/secrets/my-secret/abc123" @@ -399,65 +735,347 @@ def test_parse_akv_secret_url_without_version(self): assert secret_name == "my-secret" assert secret_version is None - def test_parse_akv_secret_url_invalid_raises(self): + @pytest.mark.parametrize("dns_suffix", ["vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"]) + def test_parse_akv_secret_url_accepts_supported_clouds(self, dns_suffix): + url = f"https://myvault.{dns_suffix}/secrets/my-secret/version-1" + + vault_url, secret_name, secret_version = _parse_akv_secret_url(url) + + assert vault_url == f"https://myvault.{dns_suffix}" + assert secret_name == "my-secret" + assert secret_version == "version-1" + + @pytest.mark.parametrize( + "url", + [ + "http://myvault.vault.azure.net/secrets/my-secret", + "https://attacker.example/secrets/my-secret", + "https://myvault.vault.azure.net.attacker.example/secrets/my-secret", + "https://nested.myvault.vault.azure.net/secrets/my-secret", + "https://user@myvault.vault.azure.net/secrets/my-secret", + "https://myvault.vault.azure.net:443/secrets/my-secret", + "https://myvault.vault.azure.net/not-secrets/my-secret", + "https://myvault.vault.azure.net/secrets", + "https://myvault.vault.azure.net/secrets/my-secret/", + "https://myvault.vault.azure.net/secrets/my-secret/version/extra", + "https://myvault.vault.azure.net/secrets/my-secret?api-version=7.4", + "https://myvault.vault.azure.net/secrets/my-secret#fragment", + "https://myvault.vault.azure.net/secrets/my%2Fsecret", + ], + ) + def test_parse_akv_secret_url_invalid_raises(self, url): with pytest.raises(ValueError, match="Invalid AKV secret URL"): - _parse_akv_secret_url("https://myvault.vault.azure.net/not-secrets/my-secret") + _parse_akv_secret_url(url) + + async def test_load_env_from_akv_async_rejects_non_azure_host_before_authentication(self): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + mock.patch("pyrit.setup.initialization._create_akv_secret_client") as mock_create_client, + pytest.raises(KeyVaultInitializationException, match="attacker.example"), + ): + await _load_env_from_akv_async( + secret_url="https://attacker.example/secrets/bootstrap", + silent=True, + ) - @mock.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() + mock_credential_cls.assert_not_called() + mock_create_client.assert_not_called() + + async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_secrets(self): + credential, client = _create_mock_akv_clients() + root_document = ( + "DIRECT=from-bootstrap\n" + "FROM_ENV=${SOURCE_VALUE}\n" + "FROM_KV=kv:https://myvault.vault.azure.net/secrets/api-key\n" + "PINNED_KV=kv:https://myvault.vault.azure.net/secrets/api-key/version-2\n" + "TERMINAL=kv:https://myvault.vault.azure.net/secrets/terminal\n" + "A=one\nB=${A}\nA=two\nC=${A}" + ) + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value=root_document), + types.SimpleNamespace(value="api-key-value"), + types.SimpleNamespace(value="pinned-key-value"), + types.SimpleNamespace(value="kv:https://myvault.vault.azure.net/secrets/not-followed"), + ] + ) + secret_url = "https://myvault.vault.azure.net/secrets/bootstrap/v1" - async def test_load_env_from_akv_async_loads_secret_content(self): - class FakeCredential: - pass + with ( + mock.patch.dict(os.environ, {"SOURCE_VALUE": "ambient-value"}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential) as mock_credential_cls, + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, + mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, + ): + await _load_env_from_akv_async(secret_url=secret_url, silent=True) + + assert os.environ["DIRECT"] == "from-bootstrap" + assert os.environ["FROM_ENV"] == "ambient-value" + assert os.environ["FROM_KV"] == "api-key-value" + assert os.environ["PINNED_KV"] == "pinned-key-value" + assert os.environ["TERMINAL"] == "kv:https://myvault.vault.azure.net/secrets/not-followed" + assert os.environ["A"] == "two" + assert os.environ["B"] == "one" + assert os.environ["C"] == "two" + + mock_credential_cls.assert_called_once_with() + _assert_mock_akv_client_created( + mock_client_cls, + vault_url="https://myvault.vault.azure.net", + credential=credential, + ) + assert client.get_secret.await_args_list == [ + mock.call("bootstrap", version="v1"), + mock.call("api-key", version=None), + mock.call("api-key", version="version-2"), + mock.call("terminal", version=None), + ] + credential.__aenter__.assert_awaited_once() + credential.__aexit__.assert_awaited_once() + client.__aenter__.assert_awaited_once() + client.__aexit__.assert_awaited_once() + mock_print_msg.assert_called_once() + + async def test_load_env_from_akv_async_rejects_short_secret_name(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="API_KEY=kv:api-key")) - client_calls: list[tuple[str, object, object]] = [] + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="must use a full secret URL"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) - class FakeSecretClient: - def __init__(self, *, vault_url, credential): - client_calls.append(("init", vault_url, credential)) + @pytest.mark.parametrize( + "reference_url", + [ + "https://other-vault.vault.azure.net/secrets/api-key", + "https://other-vault.vault.azure.net/secrets/api-key/version-1", + ], + ) + async def test_load_env_from_akv_async_rejects_cross_vault_reference(self, reference_url): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=f"API_KEY=kv:{reference_url}")) - async def get_secret(self, name, version=None): - client_calls.append(("get_secret", name, version)) - return types.SimpleNamespace(value="AKV_VAR=from_secret\n") + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="Cross-vault AKV reference"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) - 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") + async def test_load_env_from_akv_async_empty_secret_raises(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=None)) - identity_aio_module.DefaultAzureCredential = FakeCredential - keyvault_secrets_aio_module.SecretClient = FakeSecretClient + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="has no value"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/my-secret", + silent=True, + ) + + credential.__aexit__.assert_awaited_once() + client.__aexit__.assert_awaited_once() + + async def test_load_env_from_akv_async_without_entries_raises(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="# comments only\n")) with ( - mock.patch.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, + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="contains no environment entries"), ): await _load_env_from_akv_async( - secret_urls=["https://myvault.vault.azure.net/secrets/my-secret/v1"], + secret_url="https://myvault.vault.azure.net/secrets/my-secret", silent=True, ) - assert client_calls[0][0] == "init" - assert client_calls[0][1] == "https://myvault.vault.azure.net" - assert isinstance(client_calls[0][2], FakeCredential) - assert client_calls[1] == ("get_secret", "my-secret", "v1") + credential.__aexit__.assert_awaited_once() + client.__aexit__.assert_awaited_once() + + @pytest.mark.parametrize( + ("document", "error"), + [ + ("GOOD=resolved\n=malformed\nOTHER=resolved", "malformed entries at lines: 2"), + ("MISSING_VALUE\n", "variables without values: MISSING_VALUE"), + ], + ) + async def test_load_env_from_akv_async_rejects_non_assignments(self, document, error): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) + + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match=error), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert "GOOD" not in os.environ + assert "OTHER" not in os.environ + + async def test_load_env_from_akv_async_wraps_malformed_bootstrap(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="=malformed")) + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(KeyVaultInitializationException, match="malformed entries") as exc_info, + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert isinstance(exc_info.value.__cause__, ValueError) + + async def test_load_env_from_akv_async_wraps_missing_child_secret(self): + credential, client = _create_mock_akv_clients() + missing_error = ResourceNotFoundError(message="Secret was not found") + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value="API_KEY=kv:https://myvault.vault.azure.net/secrets/missing"), + missing_error, + ] + ) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(KeyVaultInitializationException, match="Failed to resolve Key Vault reference") as exc_info, + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert exc_info.value.__cause__ is missing_error + + async def test_load_env_from_akv_async_allows_empty_assignment(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="EMPTY=")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert os.environ["EMPTY"] == "" + + async def test_load_env_from_akv_async_allows_empty_child_secret(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value="EMPTY=kv:https://myvault.vault.azure.net/secrets/empty-secret"), + types.SimpleNamespace(value=""), + ] + ) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert os.environ["EMPTY"] == "" + assert client.get_secret.await_args_list[-1] == mock.call("empty-secret", version=None) + + async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entries(self, caplog, capsys): + credential, client = _create_mock_akv_clients() + document = "GOOD=resolved\n=malformed\nMISSING_VALUE\nOTHER=also-resolved" + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + caplog.at_level("WARNING", logger="pyrit.setup.initialization"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=False, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" + + output = capsys.readouterr().out + assert "WARNING: AKV environment document contains invalid entries that will be skipped" in output + assert "malformed entries at lines: 2" in output + assert "variables without values: MISSING_VALUE" in output + assert "GOOD" not in caplog.text + assert "resolved" not in caplog.text + + async def test_load_env_from_akv_async_non_strict_silent_logs_warning(self, caplog, capsys): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="GOOD=resolved\nMISSING_VALUE")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + caplog.at_level("WARNING", logger="pyrit.setup.initialization"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=True, + ) - 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 + assert capsys.readouterr().out == "" + assert "variables without values: MISSING_VALUE" in caplog.text + + async def test_load_env_from_akv_async_child_failure_keeps_loaded_bootstrap_values(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace( + value=("GOOD=resolved\nBAD=kv:https://myvault.vault.azure.net/secrets/missing-value") + ), + types.SimpleNamespace(value=None), + ] + ) + + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="has no value"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["BAD"] == "kv:https://myvault.vault.azure.net/secrets/missing-value"