diff --git a/.azuredevops/test-job-template.yml b/.azuredevops/test-job-template.yml index 65905af5cc..7ed5c8ef5e 100644 --- a/.azuredevops/test-job-template.yml +++ b/.azuredevops/test-job-template.yml @@ -26,24 +26,40 @@ jobs: versionSpec: '3.12' addToPath: true - bash: | - mkdir -p ~/.pyrit + install -d -m 700 ~/.pyrit displayName: "Create PyRIT configuration directory" name: create_pyrit_dir - task: AzureKeyVault@2 - displayName: Azure Key Vault - retrieve .env file secret + displayName: Azure Key Vault - retrieve environment secrets inputs: azureSubscription: 'integration-test-service-connection' KeyVaultName: 'pyrit-environment' SecretsFilter: 'env-global' RunAsPreJob: false - bash: | - python -c " - import os; - secret = os.environ.get('PYRIT_TEST_SECRET'); + python - <<'PY' + import os + import pathlib + import tempfile + + secret = os.environ.get("PYRIT_TEST_SECRET") if not secret: - raise ValueError('PYRIT_TEST_SECRET is not set'); - with open(os.path.expanduser('~/.pyrit/.env'), 'w') as file: - file.write(secret)" + raise ValueError("PYRIT_TEST_SECRET is not set") + + env_file = pathlib.Path.home() / ".pyrit" / ".env" + file_descriptor, temporary_name = tempfile.mkstemp(prefix=".env.", suffix=".tmp", dir=env_file.parent) + temporary_file = pathlib.Path(temporary_name) + try: + os.fchmod(file_descriptor, 0o600) + with os.fdopen(file_descriptor, "w", encoding="utf-8", newline="") as stream: + file_descriptor = -1 + stream.write(secret) + os.replace(temporary_file, env_file) + finally: + if file_descriptor >= 0: + os.close(file_descriptor) + temporary_file.unlink(missing_ok=True) + PY env: PYRIT_TEST_SECRET: $(env-global) name: create_env_file @@ -92,10 +108,14 @@ jobs: cp -r $PyRIT_DIR/doc $NEW_DIR cp -r $PyRIT_DIR/assets $NEW_DIR cp -r $PyRIT_DIR/tests/${{ parameters.testsFolder }} $NEW_DIR/tests + cp $PyRIT_DIR/.env_example $NEW_DIR/.env_example cd $NEW_DIR displayName: "Create and switch to new test directory" - task: AzureCLI@2 displayName: "Authenticate with service principal, cache Cognitive Services access token, and run tests" + env: + PYRIT_ENV_EXAMPLE_PATH: $(Build.SourcesDirectory)/../${{ parameters.newDir }}/.env_example + PYRIT_REPOSITORY_ROOT: $(Build.SourcesDirectory) inputs: azureSubscription: ${{ parameters.testAzureSubscription }} scriptType: 'bash' diff --git a/.env_example b/.env_example index 0d70a48cfd..34468eba8c 100644 --- a/.env_example +++ b/.env_example @@ -2,138 +2,134 @@ # 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="https://api.openai.com/v1" # or any OpenAI-compatible API +# OPENAI_CHAT_KEY="your-key-here" +# OPENAI_CHAT_MODEL="gpt-4o" +# +# URL placeholders intentionally use plain dotenv values. Earlier versions used +# angle brackets as documentation styling, but python-dotenv preserves them literally # # These work with OpenAI, Azure OpenAI, Ollama, Groq, OpenRouter, and any # other OpenAI-compatible endpoint. See doc/setup/populating_secrets.md -# for provider-specific examples. +# for provider-specific examples # -# If you are using Entra authentication for Azure resources, +# If you are using Entra authentication for Azure resources # keys for those resources are not needed. PyRIT auto-detects: if an API key -# is set, it uses key auth; otherwise it falls back to Entra ID automatically. +# is set, it uses key auth; otherwise it falls back to Entra ID automatically # # ============================================================================ - - -################################### +################################## # OPENAI TARGET SECRETS +################################## # # The below models work with OpenAIChatTarget - either pass via environment variables # or copy to OPENAI_CHAT_ENDPOINT -################################### - -PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" -PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" -PLATFORM_OPENAI_CHAT_MODEL="gpt-4o" - # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately # Example: https://xxxx.openai.azure.com/openai/v1 + AZURE_OPENAI_GPT4O_ENDPOINT="https://xxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT4O_KEY="xxxxx" AZURE_OPENAI_GPT4O_MODEL="deployment-name" -# Since Azure deployment name may be custom and differ from the actual underlying model, -# you can specify the underlying model for identifier purposes. If not specified, -# identifiers will default to the value of the standard MODEL environment variable. + +# Since Azure deployment name may be custom and differ from the actual underlying model +# you can specify the underlying model for identifier purposes. If not specified +# identifiers will default to the value of the standard MODEL environment variable + AZURE_OPENAI_GPT4O_UNDERLYING_MODEL="gpt-4o" -# Optional second GPT-4o endpoint (that can be used for round-robin distribution). +# Optional second GPT-4o endpoint (that can be used for round-robin distribution) # TargetInitializer creates RoundRobinTargets that automatically group together # targets with identical underlying model names and behavioral params, allowing -# for distribution of requests across them for rate-limit relief. +# for distribution of requests across them for rate-limit relief + AZURE_OPENAI_GPT4O_ENDPOINT2="https://xxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT4O_KEY2="xxxxx" AZURE_OPENAI_GPT4O_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNDERLYING_MODEL2="gpt-4o" -AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_INTEGRATION_TEST_KEY="xxxxx" -AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" -AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" - -AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT3_5_CHAT_KEY="xxxxx" -AZURE_OPENAI_GPT3_5_CHAT_MODEL="deployment-name" -AZURE_OPENAI_GPT3_5_CHAT_UNDERLYING_MODEL="" - -AZURE_OPENAI_GPT4_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT4_CHAT_KEY="xxxxx" -AZURE_OPENAI_GPT4_CHAT_MODEL="deployment-name" -AZURE_OPENAI_GPT4_CHAT_UNDERLYING_MODEL="" - -AZURE_OPENAI_GPT5_4_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT5_4_KEY="xxxxx" -AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" -AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" +AZURE_OPENAI_GPT4O_AAD_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT4O_AAD_KEY="xxxxx" +AZURE_OPENAI_GPT4O_AAD_MODEL="deployment-name" # Endpoints that host models with fewer safety mechanisms (e.g. via adversarial fine tuning -# or content filters turned off) can be defined below and used in adversarial attack testing scenarios. +# or content filters turned off) can be defined below and used in adversarial attack testing scenarios + AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL="" - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" -# Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) -# Default endpoint goes here; specialized ones below -ADVERSARIAL_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -ADVERSARIAL_CHAT_KEY="xxxxx" -ADVERSARIAL_CHAT_MODEL="deployment-name" +# Objective Scorer chat target (used in scorers in scenarios) -ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" -ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" -ADVERSARIAL_CHAT_SINGLETURN_MODEL="deployment-name" +OBJECTIVE_SCORER_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" +OBJECTIVE_SCORER_CHAT_UNDERLYING_MODEL="" -ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" -ADVERSARIAL_CHAT_MULTITURN_KEY="xxxxx" -ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" +AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" +AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" -ADVERSARIAL_CHAT_REASONING_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" -ADVERSARIAL_CHAT_REASONING_KEY="xxxxx" -ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" +AZURE_OPENAI_GPT5_COMPLETIONS_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT5_COMPLETIONS_MODEL="gpt-5" +AZURE_OPENAI_GPT5_COMPLETIONS_UNDERLYING_MODEL="gpt-5" +AZURE_OPENAI_GPT5_4_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" +AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" -# Objective Scorer chat target (used in scorers in scenarios) -OBJECTIVE_SCORER_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -OBJECTIVE_SCORER_CHAT_KEY="xxxxx" -OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" +AZURE_OPENAI_GPT4O_STRICT_FILTER_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT4O_STRICT_FILTER_MODEL="deployment-name" + +AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT3_5_CHAT_MODEL="deployment-name" +AZURE_OPENAI_GPT3_5_CHAT_UNDERLYING_MODEL="" + +AZURE_OPENAI_GPT4_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT4_CHAT_MODEL="deployment-name" +AZURE_OPENAI_GPT4_CHAT_UNDERLYING_MODEL="" + +MAI_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +MAI_CHAT_MODEL="deployment-name" +MAI_CHAT_KEY="xxxxx" +MAI_CHAT_UNDERLYING_MODEL="" + +AZURE_OPENAI_GPTV_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPTV_CHAT_MODEL="deployment-name" AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="https://xxxxx.eastus2.models.ai.azure.com" AZURE_FOUNDRY_DEEPSEEK_KEY="xxxxx" AZURE_FOUNDRY_DEEPSEEK_MODEL="" - AZURE_FOUNDRY_PHI4_ENDPOINT="https://xxxxx.models.ai.azure.com" AZURE_CHAT_PHI4_KEY="xxxxx" AZURE_CHAT_PHI4_MODEL="" - AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="https://xxxxx.services.ai.azure.com/openai/v1/" -AZURE_FOUNDRY_MISTRAL_LARGE_KEY="xxxxx" AZURE_FOUNDRY_MISTRAL_LARGE_MODEL="Mistral-Large-3" +OLLAMA_CHAT_ENDPOINT="http://127.0.0.1:11434/v1" +OLLAMA_MODEL="llama2" +AZURE_OPENAI_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_RESPONSES_MODEL="o4-mini" +AZURE_OPENAI_RESPONSES_UNDERLYING_MODEL="o4-mini" -AWS_ENDPOINT="https://bedrock-mantle.us-east-1.api.aws/v1" -AWS_KEY="xxxxx" -AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" -AWS_RESPONSES_MODEL="openai.gpt-oss-120b" +AZURE_OPENAI_O4MINI_STRICT_FILTER_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_O4MINI_STRICT_FILTER_RESPONSES_MODEL="o4-mini" -GROQ_ENDPOINT="https://api.groq.com/openai/v1" -GROQ_KEY="gsk_xxxxxxxx" -GROQ_LLAMA_MODEL="llama3-8b-8192" +AZURE_OPENAI_GPT41_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT41_RESPONSES_MODEL="gpt-4.1" -OPEN_ROUTER_ENDPOINT="https://openrouter.ai/api/v1" -OPEN_ROUTER_KEY="sk-or-v1-xxxxx" -OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" +AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" +AZURE_OPENAI_GPT5_MODEL="gpt-5" +AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" -OLLAMA_CHAT_ENDPOINT="http://127.0.0.1:11434/v1" -OLLAMA_MODEL="llama2" +PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" +PLATFORM_OPENAI_CHAT_MODEL="gpt-4o" +PLATFORM_OPENAI_RESPONSES_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" +PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" DEFAULT_OPENAI_FRONTEND_ENDPOINT = ${AZURE_OPENAI_GPT4O_AAD_ENDPOINT} DEFAULT_OPENAI_FRONTEND_KEY = ${AZURE_OPENAI_GPT4O_AAD_KEY} @@ -142,29 +138,11 @@ DEFAULT_OPENAI_FRONTEND_MODEL = "gpt-4o" OPENAI_CHAT_ENDPOINT=${PLATFORM_OPENAI_CHAT_ENDPOINT} OPENAI_CHAT_KEY=${PLATFORM_OPENAI_CHAT_KEY} OPENAI_CHAT_MODEL=${PLATFORM_OPENAI_CHAT_MODEL} + # The following line can be populated if using an Azure OpenAI deployment # where the deployment name differs from the actual underlying model -OPENAI_CHAT_UNDERLYING_MODEL="" - -################################## -# OPENAI RESPONSES TARGET SECRETS -################################## - -AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" -AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" -AZURE_OPENAI_GPT5_KEY="xxxxxxx" -AZURE_OPENAI_GPT5_MODEL="gpt-5" -AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" - -PLATFORM_OPENAI_RESPONSES_ENDPOINT="https://api.openai.com/v1" -PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" -PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" - -AZURE_OPENAI_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -AZURE_OPENAI_RESPONSES_KEY="xxxxx" -AZURE_OPENAI_RESPONSES_MODEL="o4-mini" -AZURE_OPENAI_RESPONSES_UNDERLYING_MODEL="o4-mini" +OPENAI_CHAT_UNDERLYING_MODEL="" OPENAI_RESPONSES_ENDPOINT=${PLATFORM_OPENAI_RESPONSES_ENDPOINT} OPENAI_RESPONSES_KEY=${PLATFORM_OPENAI_RESPONSES_KEY} OPENAI_RESPONSES_MODEL=${PLATFORM_OPENAI_RESPONSES_MODEL} @@ -172,20 +150,16 @@ OPENAI_RESPONSES_UNDERLYING_MODEL="" ################################## # OPENAI REALTIME TARGET SECRETS -# -# The below models work with RealtimeTarget - either pass via environment variables -# or copy to OPENAI_REALTIME_ENDPOINT ################################## -PLATFORM_OPENAI_REALTIME_ENDPOINT="wss://api.openai.com/v1" -PLATFORM_OPENAI_REALTIME_API_KEY="sk-xxxxx" -PLATFORM_OPENAI_REALTIME_MODEL="gpt-4o-realtime-preview" - AZURE_OPENAI_REALTIME_ENDPOINT = "wss://xxxx.openai.azure.com/openai/v1" AZURE_OPENAI_REALTIME_API_KEY = "xxxxx" AZURE_OPENAI_REALTIME_MODEL = "gpt-4o-realtime-preview" AZURE_OPENAI_REALTIME_UNDERLYING_MODEL = "gpt-4o-realtime-preview" - +PLATFORM_OPENAI_REALTIME_ENDPOINT="wss://api.openai.com/v1" +PLATFORM_OPENAI_REALTIME_KEY="sk-xxxxx" +PLATFORM_OPENAI_REALTIME_API_KEY="sk-xxxxx" +PLATFORM_OPENAI_REALTIME_MODEL="gpt-4o-realtime-preview" OPENAI_REALTIME_ENDPOINT = ${PLATFORM_OPENAI_REALTIME_ENDPOINT} OPENAI_REALTIME_API_KEY = ${PLATFORM_OPENAI_REALTIME_API_KEY} OPENAI_REALTIME_MODEL = ${PLATFORM_OPENAI_REALTIME_MODEL} @@ -193,129 +167,182 @@ 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_API_KEY1 = "xxxxxx" -OPENAI_IMAGE_MODEL1 = "deployment-name" -OPENAI_IMAGE_UNDERLYING_MODEL1 = "dall-e-3" - -OPENAI_IMAGE_ENDPOINT2 = "https://api.openai.com/v1" -OPENAI_IMAGE_API_KEY2 = "sk-xxxxx" -OPENAI_IMAGE_MODEL2 = "dall-e-3" -OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" - -OPENAI_IMAGE_ENDPOINT = ${OPENAI_IMAGE_ENDPOINT2} -OPENAI_IMAGE_API_KEY = ${OPENAI_IMAGE_API_KEY2} -OPENAI_IMAGE_MODEL = ${OPENAI_IMAGE_MODEL2} -OPENAI_IMAGE_UNDERLYING_MODEL = "" +################################## +AZURE_OPENAI_IMAGE_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_IMAGE_API_KEY1 = "xxxxxx" +AZURE_OPENAI_IMAGE_MODEL1 = "deployment-name" +AZURE_OPENAI_IMAGE_UNDERLYING_MODEL1 = "dall-e-3" +AZURE_OPENAI_IMAGE_ENDPOINT2 = "https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_IMAGE_API_KEY2 = "xxxxxx" +AZURE_OPENAI_IMAGE_MODEL2 = "dall-e-3" +AZURE_OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" +OPENAI_IMAGE_ENDPOINT = ${AZURE_OPENAI_IMAGE_ENDPOINT2} +OPENAI_IMAGE_API_KEY = ${AZURE_OPENAI_IMAGE_API_KEY2} +OPENAI_IMAGE_MODEL = ${AZURE_OPENAI_IMAGE_MODEL2} +OPENAI_IMAGE_UNDERLYING_MODEL = "" +OPENAI_IMAGE_STRICT_FILTER_ENDPOINT = "https://xxxxx.openai.azure.com/openai/v1" +OPENAI_IMAGE_STRICT_FILTER_MODEL = "gpt-image" ################################## # TTS TARGET SECRETS -# -# The below models work with OpenAITTSTarget - either pass via environment variables -# or copy to OPENAI_TTS_ENDPOINT -################################### - -OPENAI_TTS_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" -OPENAI_TTS_KEY1 = "xxxxxxx" -OPENAI_TTS_MODEL1 = "tts" -OPENAI_TTS_UNDERLYING_MODEL1 = "tts" - -OPENAI_TTS_ENDPOINT2 = "https://api.openai.com/v1" -OPENAI_TTS_KEY2 = "xxxxxx" -OPENAI_TTS_MODEL2 = "tts-1" -OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" - -OPENAI_TTS_ENDPOINT = ${OPENAI_TTS_ENDPOINT2} -OPENAI_TTS_KEY = ${OPENAI_TTS_KEY2} -OPENAI_TTS_MODEL = ${OPENAI_TTS_MODEL2} +################################## + +AZURE_OPENAI_TTS_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_TTS_KEY1 = "xxxxxxx" +AZURE_OPENAI_TTS_MODEL1 = "tts" +AZURE_OPENAI_TTS_UNDERLYING_MODEL1 = "tts" +AZURE_OPENAI_TTS_ENDPOINT2 = "https://xxxxx.openai.azure.com/v1" +AZURE_OPENAI_TTS_KEY2 = "xxxxxx" +AZURE_OPENAI_TTS_MODEL2 = "tts-1" +AZURE_OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" +OPENAI_TTS_ENDPOINT = ${AZURE_OPENAI_TTS_ENDPOINT2} +OPENAI_TTS_KEY = ${AZURE_OPENAI_TTS_KEY2} +OPENAI_TTS_MODEL = ${AZURE_OPENAI_TTS_MODEL2} 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_KEY="xxxxxxx" AZURE_OPENAI_VIDEO_MODEL="sora-2" AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" - OPENAI_VIDEO_ENDPOINT = ${AZURE_OPENAI_VIDEO_ENDPOINT} OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" +################################## +# ADVERSARIAL MODELS +################################## +# Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) +# Default endpoint goes here; specialized ones below + +ADVERSARIAL_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +ADVERSARIAL_CHAT_MODEL="deployment-name" +ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" +ADVERSARIAL_CHAT_SINGLETURN_MODEL="deployment-name" +ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_MULTITURN_KEY="xxxxx" +ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" +ADVERSARIAL_CHAT_REASONING_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +ADVERSARIAL_CHAT_REASONING_KEY="xxxxx" +ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" ################################## # AML TARGET SECRETS +################################## # The below models work with AzureMLChatTarget - either pass via environment variables # or copy to AZURE_ML_MANAGED_ENDPOINT -################################### AZURE_ML_PHI_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" AZURE_ML_PHI_KEY="xxxxx" -# The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed. +# The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed + AZURE_ML_MANAGED_ENDPOINT=${AZURE_ML_PHI_ENDPOINT} AZURE_ML_KEY=${AZURE_ML_PHI_KEY} - ################################## # MISC TARGET SECRETS -################################### - - -OPENAI_COMPLETION_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -OPENAI_COMPLETION_API_KEY="xxxxx" -OPENAI_COMPLETION_MODEL="davinci-002" +################################## OPENAI_EMBEDDING_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" -OPENAI_EMBEDDING_KEY="xxxxx" OPENAI_EMBEDDING_MODEL="text-embedding-3-small" - -AZURE_STORAGE_ACCOUNT_CONTAINER_URL="https://xxxxxx.blob.core.windows.net/xpia" -AZURE_STORAGE_ACCOUNT_SAS_TOKEN="xxxxx" - - AZURE_SPEECH_REGION = "eastus2" -AZURE_SPEECH_KEY = "xxxxx" + # Resource ID is needed when using Entra authentication -AZURE_SPEECH_RESOURCE_ID = "xxxxx" -AZURE_CONTENT_SAFETY_API_KEY="xxxxx" +AZURE_SPEECH_RESOURCE_ID = "xxxxx" AZURE_CONTENT_SAFETY_API_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/" - HUGGINGFACE_TOKEN="hf_xxxxxxx" HUGGINGFACE_ENDPOINT="https://router.huggingface.co/v1" -GOOGLE_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/openai" -GOOGLE_GEMINI_API_KEY = "xxxxx" -GOOGLE_GEMINI_MODEL="gemini-2.0-flash" - - -######################### +################################## # AZURE SQL SECRETS -######################### - - +################################## # This connects to the test database + AZURE_SQL_DB_CONNECTION_STRING_TEST = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="https://xxxxx.blob.core.windows.net/dbdata" # This connects to the prod database + AZURE_SQL_DB_CONNECTION_STRING_PROD = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="https://xxxxx.blob.core.windows.net/dbdata" +AZURE_STORAGE_ACCOUNT_CONTAINER_URL="https://xxxxxx.blob.core.windows.net/xpia" +# The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local -# The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local. AZURE_SQL_DB_CONNECTION_STRING = ${AZURE_SQL_DB_CONNECTION_STRING_PROD} AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL=${AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD} + +################################## +# INTEGRATION TEST ONLY SECRETS +################################## + +GOOGLE_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/openai" +GOOGLE_GEMINI_API_KEY = "xxxxx" +GOOGLE_GEMINI_MODEL="gemini-2.0-flash" + +ANTHROPIC_CHAT_ENDPOINT="https://api.anthropic.com/v1" +ANTHROPIC_CHAT_KEY="xxxxx" +ANTHROPIC_CHAT_MODEL="claude-3-7-sonnet-latest" + +AWS_KEY="xxxxx" +AWS_ENDPOINT="https://bedrock-mantle.us-east-1.api.aws/v1" +AWS_RESPONSES_MODEL="openai.gpt-oss-120b" +AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" + +PLATFORM_OPENAI_VIDEO_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_VIDEO_KEY="sk-xxxxx" +PLATFORM_OPENAI_VIDEO_MODEL="sora-2" + +PLATFORM_OPENAI_IMAGE_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_IMAGE_KEY="sk-xxxxx" +PLATFORM_OPENAI_IMAGE_MODEL="gpt-image-1" + +PLATFORM_OPENAI_EMBEDDING_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_EMBEDDING_KEY="sk-xxxxx" +PLATFORM_OPENAI_EMBEDDING_MODEL="text-embedding-3-small" + +OPENAI_COMPLETION_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +OPENAI_COMPLETION_API_KEY="xxxxx" +OPENAI_COMPLETION_MODEL="davinci-002" + +PROMPTINTEL_API_KEY="xxxxx" + +################################## +# Additional entries referenced in PyRIT +################################## + +AZURE_OPENAI_GPT4O_KEY="xxxxx" +AZURE_OPENAI_GPT4O_KEY2="xxxxx" +AZURE_OPENAI_INTEGRATION_TEST_KEY="xxxxx" +AZURE_OPENAI_GPT3_5_CHAT_KEY="xxxxx" +AZURE_OPENAI_GPT4_CHAT_KEY="xxxxx" +AZURE_OPENAI_GPT5_4_KEY="xxxxx" +AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY="xxxxx" +AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" +ADVERSARIAL_CHAT_KEY="xxxxx" +OBJECTIVE_SCORER_CHAT_KEY="xxxxx" +AZURE_FOUNDRY_MISTRAL_LARGE_KEY="xxxxx" +GROQ_ENDPOINT="https://api.groq.com/openai/v1" +GROQ_KEY="gsk_xxxxxxxx" +GROQ_LLAMA_MODEL="llama3-8b-8192" +OPEN_ROUTER_ENDPOINT="https://openrouter.ai/api/v1" +OPEN_ROUTER_KEY="sk-or-v1-xxxxx" +OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" +AZURE_OPENAI_GPT5_KEY="xxxxxxx" +AZURE_OPENAI_RESPONSES_KEY="xxxxx" +OPENAI_EMBEDDING_KEY="xxxxx" +AZURE_STORAGE_ACCOUNT_SAS_TOKEN="xxxxx" +AZURE_SPEECH_KEY = "xxxxx" +AZURE_CONTENT_SAFETY_API_KEY="xxxxx" diff --git a/.pyrit_conf_example b/.pyrit_conf_example index b41c13e060..7bcecf5ec0 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -4,7 +4,7 @@ # or specify a custom path when loading via --config-file. # # For documentation on configuration options, see: -# https://github.com/microsoft/PyRIT/blob/main/doc/setup/configuration.md +# https://github.com/microsoft/PyRIT/blob/main/doc/getting_started/pyrit_conf.md # Memory Database Type # -------------------- @@ -79,33 +79,22 @@ operation: op_trash_panda # - /path/to/my_custom_initializer.py # - ./local_initializer.py -# Environment Files -# ----------------- -# List of .env file paths to load during initialization. -# Later files override values from earlier files. -# -# Behavior: -# - Omit this field (or set to null): Load default .env and .env.local from ~/.pyrit/ if they exist -# - Set to []: Explicitly load NO environment files -# - Set to list of paths: Load only the specified files -# -# Example: -# env_files: -# - /path/to/.env -# - /path/to/.env.local - -# Azure Key Vault Environment References -# --------------------------------------- -# List of AKV secret URLs to load during initialization. -# Each secret's value must be the full contents of a .env file. -# Loaded after env_files, so AKV secrets take precedence. -# Authentication uses DefaultAzureCredential (managed identity, Azure CLI, etc.). -# -# Requires: pip install azure-keyvault-secrets -# -# Example: +# Environment Configuration +# ------------------------- +# Azure Key Vault is the canonical source for shared and deployed configuration. +# See doc/getting_started/pyrit_conf.md for loading order, references, and migration guidance. # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env +# env_akv_strict: true +# env_akv_write_env: false # Debug only: write a fully resolved, sensitive ~/.pyrit/.env. + +# Auto-discovered ~/.pyrit/.env is legacy and will be rejected in PyRIT 1.3.0. +# Use ~/.pyrit/.env.local for quick local plaintext patches or when Azure is unavailable. +# Process values remain authoritative; AKV and ordinary env_files fill gaps in load order. +# Only a file named .env.local overrides existing values. +# Explicit env_files remain supported regardless of name or location and may contain full kv: URLs. +# env_files: +# - /path/to/.env.local # Max Concurrent Scenario Runs # ---------------------------- diff --git a/build_scripts/env_local_integration_test b/build_scripts/env_local_integration_test index b61cfe7d20..ded45e2bea 100644 --- a/build_scripts/env_local_integration_test +++ b/build_scripts/env_local_integration_test @@ -7,12 +7,12 @@ OPENAI_CHAT_ENDPOINT=${AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT} OPENAI_CHAT_KEY=${AZURE_OPENAI_INTEGRATION_TEST_KEY} OPENAI_CHAT_MODEL=${AZURE_OPENAI_INTEGRATION_TEST_MODEL} -OPENAI_IMAGE_ENDPOINT=${OPENAI_IMAGE_ENDPOINT2} -OPENAI_IMAGE_API_KEY=${OPENAI_IMAGE_API_KEY2} -OPENAI_IMAGE_MODEL=${OPENAI_IMAGE_MODEL2} +OPENAI_IMAGE_ENDPOINT=${AZURE_OPENAI_IMAGE_ENDPOINT2} +OPENAI_IMAGE_API_KEY=${AZURE_OPENAI_IMAGE_API_KEY2} +OPENAI_IMAGE_MODEL=${AZURE_OPENAI_IMAGE_MODEL2} -OPENAI_TTS_ENDPOINT=${OPENAI_TTS_ENDPOINT2} -OPENAI_TTS_KEY=${OPENAI_TTS_KEY2} +OPENAI_TTS_ENDPOINT=${AZURE_OPENAI_TTS_ENDPOINT2} +OPENAI_TTS_KEY=${AZURE_OPENAI_TTS_KEY2} AZURE_SQL_DB_CONNECTION_STRING=${AZURE_SQL_DB_CONNECTION_STRING_TEST} AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL=${AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST} diff --git a/doc/code/executor/gcg/1_gcg_azure_ml.ipynb b/doc/code/executor/gcg/1_gcg_azure_ml.ipynb index a4c7b20040..264805db7a 100644 --- a/doc/code/executor/gcg/1_gcg_azure_ml.ipynb +++ b/doc/code/executor/gcg/1_gcg_azure_ml.ipynb @@ -57,7 +57,7 @@ "source": [ "import os\n", "\n", - "from pyrit.setup.initialization import _load_environment_files\n", + "from pyrit.setup.akv_initialization import _load_environment_files\n", "\n", "_load_environment_files(env_files=None)\n", "\n", diff --git a/doc/code/executor/gcg/1_gcg_azure_ml.py b/doc/code/executor/gcg/1_gcg_azure_ml.py index c3c559f18c..9e05233255 100644 --- a/doc/code/executor/gcg/1_gcg_azure_ml.py +++ b/doc/code/executor/gcg/1_gcg_azure_ml.py @@ -29,7 +29,7 @@ # %% import os -from pyrit.setup.initialization import _load_environment_files +from pyrit.setup.akv_initialization import _load_environment_files _load_environment_files(env_files=None) diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 30fd1443a9..cdb4695780 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -4,62 +4,56 @@ The recommended way to configure PyRIT. A `.pyrit_conf` file declares your datab ## Quick Setup -```bash -mkdir -p ~/.pyrit -cp .pyrit_conf_example ~/.pyrit/.pyrit_conf -cp .env_example ~/.pyrit/.env +Create `~/.pyrit/.pyrit_conf` and configure a Key Vault bootstrap document: + +```yaml +env_akv_ref: + - https://my-vault.vault.azure.net/secrets/my-pyrit-env ``` -Then edit both files for your environment. The `.pyrit_conf` tells PyRIT _how_ to initialize; the `.env` tells it _where_ your AI targets are. +The Key Vault secret value uses dotenv syntax. When Azure is unavailable or you need a quick local patch, put only those values in `~/.pyrit/.env.local`. ## File Location The default configuration file path is: -``` +```text ~/.pyrit/.pyrit_conf ``` PyRIT looks for this file automatically on startup (via the CLI, shell, or `ConfigurationLoader`). If the file does not exist, PyRIT falls back to built-in defaults. -## Setting Up Secrets (.env files) +## Environment Configuration -The `.pyrit_conf` file works hand-in-hand with `.env` files for your API credentials. See [Populating Secrets](./populating_secrets.md) for provider-specific examples of what to put in your `.env` file. +```{important} +Azure Key Vault is PyRIT's canonical environment source for shared, CI/CD, and deployed configuration. Auto-discovered `~/.pyrit/.env` is supported only as a legacy source and will be rejected in PyRIT 1.3.0. Use `~/.pyrit/.env.local` for deliberate plaintext local iteration or when Azure is unavailable. +``` -### Environment Variable Precedence +See [Populating Secrets](./populating_secrets.md) for provider-specific variable examples. -When PyRIT initializes, environment variables are loaded in a specific order. **Later sources override earlier ones:** +### Loading Order -```{mermaid} -flowchart LR - A["1. System Environment"] --> B{"env_files in .pyrit_conf?"} - B -->|No| C["2. ~/.pyrit/.env"] - C --> D["3. ~/.pyrit/.env.local"] - B -->|Yes| E["2. Your specified files (in order)"] -``` - -**Default behavior** (no `env_files` field in `.pyrit_conf`): +PyRIT loads environment sources in this order: -| Priority | Source | Description | -|----------|--------|-------------| -| Lowest | System environment variables | Always loaded as the baseline | -| Medium | `~/.pyrit/.env` | Default config file (loaded if it exists) | -| Highest | `~/.pyrit/.env.local` | Local overrides (loaded if it exists) | +1. Existing process environment variables. +2. Key Vault bootstrap documents, legacy auto-discovered `.env`, or explicit `env_files`. These sources fill only missing values. +3. Files named `.env.local`. These are the only dotenv sources that override existing values. -**Custom behavior** (with `env_files` field): Only your specified files are loaded, in order. Default paths are completely ignored. +When `env_akv_ref` is configured, PyRIT ignores an auto-discovered `~/.pyrit/.env`, emits its deprecation warning, and still loads `~/.pyrit/.env.local`. Explicit `env_files` are never blocked or deprecated based on their filename or location. ### Using .env.local for Overrides -You can use `~/.pyrit/.env.local` to override values in `~/.pyrit/.env` without modifying the base file. This is useful for: +Use `~/.pyrit/.env.local` to override process or Key Vault values deliberately. This is useful for: + - Testing different targets - Using personal credentials instead of shared ones - Switching between configurations quickly -Simply create `.env.local` in your `~/.pyrit/` directory and add any variables you want to override. +Only put the values you need to patch in this file. Because it contains plaintext secrets, do not commit it. ### Authentication Options -**API Keys (Default):** The simplest approach — set `OPENAI_CHAT_KEY` and similar variables in your `.env` file. Most targets support this method. +**API keys:** Store shared API keys as Key Vault scalar secrets and reference them from the bootstrap document with `kv:`. For local-only work, place them in `.env.local`. **Azure Entra Authentication (Optional):** For Azure resources, you can use Entra auth instead of API keys. This requires the [Azure CLI](https://learn.microsoft.com/en-us/cli/azure/install-azure-cli) and `az login`. When using Entra auth, you don't need to set API keys for Azure resources. @@ -107,7 +101,7 @@ Use `pyrit list initializers` in the CLI to see all registered initializers. See Most users should enable the following initializers. These are what the `.pyrit_conf_example` ships with and are required for features like `pyrit_scan` and automated scenarios. | Initializer | What It Registers | When You Need It | -|---|---|---| +| --- | --- | --- | | `target` | Prompt targets (OpenAI, Azure, AML, etc.) into the `TargetRegistry` | **Required for `pyrit_scan`** and any registry-based workflows | | `scorer` | Scorers (refusal, content safety, harm-category, Likert, etc.) into the `ScorerRegistry` | **Required for automated scoring** and `pyrit_scan` evaluations | | `technique` | Attack techniques into the `AttackTechniqueRegistry` | **Required for `pyrit_scan` scenarios** that select techniques | @@ -157,13 +151,13 @@ initialization_scripts: ### `env_files` -Environment file paths to load during initialization. Later files override values from earlier files. +Optional local dotenv paths. Key Vault remains the canonical shared source; explicit files support local and non-Azure workflows. -| Value | Behavior | -| ----------------- | -------------------------------------------------------------------- | -| Omitted or `null` | Load default `~/.pyrit/.env` and `~/.pyrit/.env.local` if they exist | -| `[]` (empty list) | Load **no** environment files | -| List of paths | Load **only** the specified files (defaults are skipped) | +| Value | Behavior | +| ----------------- | -------------------------------------------------------- | +| Omitted or `null` | Auto-discover legacy `~/.pyrit/.env` and supported `~/.pyrit/.env.local` | +| `[]` (empty list) | Load **no** environment files | +| List of paths | Load **only** the specified files (defaults are skipped) | ```yaml env_files: @@ -171,6 +165,81 @@ env_files: - /path/to/.env.local ``` +Local files use standard python-dotenv parsing and `${NAME}` interpolation. Ordinary files fill missing values; any file whose basename is `.env.local` overrides existing values. Explicit files load in their listed order. + +Complete-value `kv:`, `akv:`, `azure_key_vault:`, and `env_akv_ref:` references resolve in local files as well as remote bootstrap documents. Local references may use any validated supported Key Vault URL; remote child references must remain in the bootstrap document's vault. A local assignment that loses to an existing value does not fetch its secret. + +Ordinary malformed dotenv lines retain python-dotenv's permissive behavior. `env_akv_strict` controls malformed Key Vault reference syntax in all sources: strict mode raises; non-strict mode warns and skips that assignment. Authentication, authorization, transport, missing-secret, and missing-value failures always raise. + +Environment loading preserves the historical non-transactional dotenv behavior. Each bootstrap document and local file updates `os.environ` as it loads. If a later source or child-secret lookup fails, assignments made by earlier sources remain in the process environment. + +When `env_akv_ref` is not configured, an empty `env_files` list or missing default files leaves existing process environment variables unchanged and initialization continues. + +### `env_akv_ref` + +Ordered Azure Key Vault secret URLs used to obtain bootstrap environment documents. This is the canonical configuration path. Each secret value contains 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 and fill values missing from the process environment. Each document uses native dotenv interpolation against the process environment and assignments already parsed. A bootstrap document can mix literal values, `${NAME}` interpolation, and complete-value references to scalar secrets in the same vault: + +```dotenv +OPENAI_CHAT_ENDPOINT="https://example.openai.azure.com/openai/v1" +OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key" +PINNED_OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version-id" +OPENAI_CHAT_MODEL="${PYRIT_OPENAI_CHAT_MODEL}" +``` + +Resolution is limited to one child-secret lookup: + +1. PyRIT validates and loads the bootstrap dotenv document. +2. For each complete-value Key Vault reference in that document, PyRIT fetches the same-vault scalar secret and replaces the environment value. + +For example, if `OPENAI_CHAT_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key"`, the value of the `openai-chat-key` secret becomes `OPENAI_CHAT_KEY` verbatim. If that secret happens to contain `kv:another-secret`, the final environment value is the string `kv:another-secret`; PyRIT does not fetch `another-secret`. + +References must occupy the entire value. `kv:` is the canonical Key Vault prefix; `akv:`, `azure_key_vault:`, and `env_akv_ref:` are accepted aliases. + +A Key Vault reference must use a full HTTPS secret URL from the bootstrap document's vault. Supported vault DNS suffixes are `.vault.azure.net`, `.vault.azure.cn`, and `.vault.usgovcloudapi.net`. An unversioned URL reads the latest secret version at initialization. Include the version in the URL to pin it. Short names, malformed paths, arbitrary hosts, and cross-vault child references are rejected before a client is created. + +PyRIT does not cache referenced secrets. Each winning `kv:` occurrence performs a Key Vault read during initialization. References that lose to an existing process or earlier source are not fetched. Debug output is the exception: it resolves bootstrap references for the written file without changing runtime precedence. + +```dotenv +LATEST_KEY_URI="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key" +PINNED_KEY="kv:https://my-vault.vault.azure.net/secrets/openai-chat-key/version-id" +``` + +Bootstrap documents stay in memory by default. Use `.env.local` when an intentional local override is required. + +### `env_akv_strict` + +Controls Key Vault bootstrap validation and Key Vault reference syntax in local files. It defaults to `true`. + +```yaml +env_akv_strict: false +``` + +In strict mode, malformed bootstrap dotenv lines, valueless bootstrap entries, and malformed Key Vault references stop initialization. Empty assignments such as `OPTIONAL_VALUE=` remain valid. With `env_akv_strict: false`, PyRIT warns and skips malformed bootstrap entries and malformed reference assignments without logging secret values. + +Non-strict mode does not suppress operational failures. Missing secrets, authentication, authorization, transport errors, and bootstrap documents with no valid assignments still stop initialization. Loading remains non-transactional, so earlier successful assignments remain. + +Key Vault clients use an explicit Azure retry policy with up to three retries and exponential backoff. Bootstrap parsing, invalid or missing secrets, authentication, authorization, and Azure transport failures are raised as `KeyVaultInitializationException` with the original exception preserved as the cause. The exception remains `ValueError`-compatible for callers migrating from the previous contract. + +### `env_akv_write_env` + +Defaults to `false`. Set it to `true` only while debugging to write a fully resolved bootstrap document to `~/.pyrit/.env`: + +```yaml +env_akv_write_env: true +``` + +The written file contains only bootstrap assignments, comments, and fully resolved child-secret values. It excludes unrelated process and `.env.local` values, safely round-trips terminal secret text, and is always named `.env`, never `.env.new`. + +PyRIT refuses debug mode when `~/.pyrit/.env` already exists; rename or remove the existing file first. The file is created with owner-only permissions where supported and replaced atomically, but it contains plaintext secrets. Remove it when debugging is complete. `.env.local` still loads afterward and can override runtime values without changing the generated file. + ### `silent` If `true`, suppresses print statements during initialization. Useful for non-interactive environments or when embedding PyRIT in other tools. Defaults to `false`. @@ -180,7 +249,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 +285,7 @@ This means you can set sensible defaults in `~/.pyrit/.pyrit_conf` and override The 3-layer model above determines **which config values are selected**. Once resolved, the values are applied in a fixed runtime order: -1. Environment files are loaded +1. Process values are retained, AKV or ordinary local sources fill gaps, and `.env.local` applies final overrides 2. Default values are reset 3. Memory database is configured (from `memory_db_type`) 4. Initializers are executed in listed order @@ -292,11 +361,14 @@ initializers: # initialization_scripts: # - /path/to/my_custom_initializer.py -# Environment files (optional) -# Omit or set to null to use defaults (~/.pyrit/.env, ~/.pyrit/.env.local) -# Set to [] to load no env files +# Canonical: ordered Azure Key Vault bootstrap environment documents +# env_akv_ref: +# - https://my-vault.vault.azure.net/secrets/my-pyrit-env +# env_akv_strict: true +# env_akv_write_env: false # Debug only: writes fully resolved plaintext secrets + +# Optional plaintext local patch or non-Azure workflow # env_files: -# - /path/to/.env # - /path/to/.env.local # Suppress initialization messages diff --git a/infra/env.demo.template b/infra/env.demo.template index 4c5e7dac7e..77f2af2b56 100644 --- a/infra/env.demo.template +++ b/infra/env.demo.template @@ -36,16 +36,16 @@ AZURE_CONTENT_SAFETY_API_ENDPOINT=https://YOUR_CONTENT_SAFETY.cognitiveservices. AZURE_CONTENT_SAFETY_API_KEY= # ─── Image Target (optional — for image generation demos) ─── -# OPENAI_IMAGE_ENDPOINT1=https://YOUR_IMAGE_ENDPOINT.openai.azure.com/openai/v1 -# OPENAI_IMAGE_API_KEY1= -# OPENAI_IMAGE_MODEL1=dall-e-3 -# OPENAI_IMAGE_UNDERLYING_MODEL1=dall-e-3 +# AZURE_OPENAI_IMAGE_ENDPOINT1=https://YOUR_IMAGE_ENDPOINT.openai.azure.com/openai/v1 +# AZURE_OPENAI_IMAGE_API_KEY1= +# AZURE_OPENAI_IMAGE_MODEL1=dall-e-3 +# AZURE_OPENAI_IMAGE_UNDERLYING_MODEL1=dall-e-3 # ─── TTS Target (optional — for text-to-speech demos) ─── -# OPENAI_TTS_ENDPOINT1=https://YOUR_TTS_ENDPOINT.openai.azure.com/openai/v1 -# OPENAI_TTS_KEY1= -# OPENAI_TTS_MODEL1=tts-1 -# OPENAI_TTS_UNDERLYING_MODEL1=tts-1 +# AZURE_OPENAI_TTS_ENDPOINT1=https://YOUR_TTS_ENDPOINT.openai.azure.com/openai/v1 +# AZURE_OPENAI_TTS_KEY1= +# AZURE_OPENAI_TTS_MODEL1=tts-1 +# AZURE_OPENAI_TTS_UNDERLYING_MODEL1=tts-1 # ─── Video Target (optional — for video generation demos) ─── # AZURE_OPENAI_VIDEO_ENDPOINT=https://YOUR_VIDEO_ENDPOINT.openai.azure.com/openai/v1 diff --git a/pyrit/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/executor/promptgen/gcg/experiments/run.py b/pyrit/executor/promptgen/gcg/experiments/run.py index 3c7cbf0390..c5a2c84bb0 100644 --- a/pyrit/executor/promptgen/gcg/experiments/run.py +++ b/pyrit/executor/promptgen/gcg/experiments/run.py @@ -27,7 +27,7 @@ from pyrit.executor.promptgen.gcg.config import GCGConfig, GCGDataConfig, GCGOutputConfig from pyrit.executor.promptgen.gcg.data import load_goals_and_targets from pyrit.executor.promptgen.gcg.generator import GCGGenerator -from pyrit.setup.initialization import _load_environment_files +from pyrit.setup.akv_initialization import _load_environment_files def _parse_arguments() -> argparse.Namespace: diff --git a/pyrit/setup/akv_initialization.py b/pyrit/setup/akv_initialization.py new file mode 100644 index 0000000000..3751d66205 --- /dev/null +++ b/pyrit/setup/akv_initialization.py @@ -0,0 +1,851 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Load dotenv files and Azure Key Vault-backed environment documents.""" + +import asyncio +import contextlib +import io +import logging +import os +import pathlib +import tempfile +import urllib.parse +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any + +import dotenv +from dotenv.parser import parse_stream + +from pyrit.common import path, print_deprecation_message +from pyrit.exceptions import KeyVaultInitializationException + +if TYPE_CHECKING: + from azure.keyvault.secrets.aio import SecretClient + +logger = logging.getLogger(__name__) + +_AKV_REFERENCE_PREFIXES = frozenset({"akv", "kv", "azure_key_vault", "env_akv_ref"}) +_AKV_VAULT_DNS_SUFFIXES = frozenset({"vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"}) +_AKV_RETRY_TOTAL = 3 +_AKV_RETRY_BACKOFF_FACTOR = 0.8 +_AKV_ENV_FILE_NAME = ".env" +_LEGACY_ENV_REMOVED_IN = "1.3.0" + + +def _load_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool = False, + include_default_base: bool = True, + assignment_fallbacks: dict[str, str | None] | None = None, +) -> bool: + """ + Load environment files in the order they are provided. + + Files fill values missing from the process environment. A file named + ``.env.local`` is the only local source that overrides existing values. + + Args: + env_files: Optional sequence of environment file paths. If None, loads default + .env and .env.local from PyRIT home directory (only if they exist). + silent: If True, suppresses print statements about environment file loading. + Defaults to False. + include_default_base: If False and env_files is None, skips the default + .env file while still loading .env.local. Defaults to True. + assignment_fallbacks: Optional output mapping from assignments that win + precedence to the value they replaced, if any. + + Returns: + True if at least one environment file was loaded, otherwise False. + + Raises: + ValueError: If any provided env_files do not exist. + """ + selected_files = _select_environment_files( + env_files=env_files, + silent=silent, + include_default_base=include_default_base, + ) + for env_file in selected_files: + override = env_file.name == ".env.local" + if assignment_fallbacks is not None: + assignment_names = dotenv.dotenv_values(dotenv_path=env_file, interpolate=False) + for variable_name in assignment_names: + if override or variable_name not in os.environ: + assignment_fallbacks[variable_name] = os.environ.get(variable_name) + dotenv.load_dotenv( + dotenv_path=env_file, + override=override, + interpolate=True, + ) + if not silent: + _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) + + return bool(selected_files) + + +def _select_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool, + include_default_base: bool, +) -> list[pathlib.Path]: + """ + Select and validate environment files without reading their contents. + + Returns: + list[pathlib.Path]: Environment files in load order. + + Raises: + ValueError: If an explicitly provided environment file does not exist. + """ + if env_files is not None: + if not silent: + _print_msg(f"Loading custom environment files: {[str(f) for f in env_files]}", quiet=silent, log=True) + for env_file in env_files: + if not env_file.exists(): + raise ValueError(f"Environment file not found: {env_file}") + + # By default load .env and .env.local from home directory of the package + else: + default_files = [] + base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" + local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" + + if include_default_base and base_file.exists(): + _warn_about_legacy_env(env_file=base_file, ignored_for_akv=False, silent=silent) + default_files.append(base_file) + if local_file.exists(): + default_files.append(local_file) + + if not silent: + if default_files: + _print_msg( + f"Found default environment files: {[str(f) for f in default_files]}", quiet=silent, log=True + ) + else: + _print_msg( + "No default environment files found. Using system environment variables only.", + quiet=silent, + log=True, + ) + + env_files = default_files + + return list(env_files) + + +def _print_msg(message: str, quiet: bool, log: bool) -> None: + """ + Print a standard initialization message unless quiet is True. + + Args: + message (str): The message to print and/or log. + quiet (bool): If True, suppresses the initialization message. + log (bool): If True, logs the message using the logger. + """ + if not quiet: + print(message) + if log: + logger.info(message) + + +def _warn_about_akv_environment_files( + env_files: Sequence[pathlib.Path] | None, + *, + silent: bool = False, +) -> None: + """Warn when an auto-discovered legacy environment file coexists with AKV.""" + if env_files is not None: + return + + base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" + if base_file.exists(): + _warn_about_legacy_env(env_file=base_file, ignored_for_akv=True, silent=silent) + + +def _warn_about_legacy_env(*, env_file: pathlib.Path, ignored_for_akv: bool, silent: bool) -> None: + """Emit the standard and visible warnings for auto-discovered legacy ``.env`` loading.""" + print_deprecation_message( + old_item=f"Auto-discovered {env_file}", + new_item="env_akv_ref or ~/.pyrit/.env.local", + removed_in=_LEGACY_ENV_REMOVED_IN, + ) + behavior = "will be ignored because env_akv_ref is configured" if ignored_for_akv else "will still be loaded" + message = ( + f"Auto-discovered {env_file} is deprecated and {behavior}. " + f"Support will be removed in {_LEGACY_ENV_REMOVED_IN}. Use env_akv_ref or ~/.pyrit/.env.local instead." + ) + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + + +def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: + """ + Parse an AKV secret URL into vault URL, secret name, and optional version. + + Args: + secret_url (str): Full AKV secret URL in the format + ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. + + Returns: + tuple[str, str, str | None]: (vault_url, secret_name, secret_version) + + Raises: + ValueError: If the URL does not match the expected format. + """ + error_message = ( + f"Invalid AKV secret URL: '{secret_url}'. Expected an HTTPS Azure Key Vault URL in the format " + "https://{vault}.{vault-dns-suffix}/secrets/{name}[/{version}]." + ) + try: + parsed_url = urllib.parse.urlsplit(secret_url) + port = parsed_url.port + except (TypeError, ValueError) as error: + raise ValueError(error_message) from error + + hostname = parsed_url.hostname + vault_name, separator, dns_suffix = hostname.partition(".") if hostname else ("", "", "") + valid_vault_name = 1 <= len(vault_name) <= 63 and all( + char.isascii() and (char.isalnum() or char == "-") for char in vault_name + ) + valid_authority = ( + parsed_url.scheme.casefold() == "https" + and parsed_url.username is None + and parsed_url.password is None + and port is None + and separator == "." + and dns_suffix in _AKV_VAULT_DNS_SUFFIXES + and valid_vault_name + ) + path_parts = parsed_url.path.split("/") + valid_path = ( + len(path_parts) in {3, 4} and path_parts[0] == "" and path_parts[1] == "secrets" and all(path_parts[2:]) + ) + if not valid_authority or not valid_path or parsed_url.query or parsed_url.fragment: + raise ValueError(error_message) + + secret_name = path_parts[2] + secret_version = path_parts[3] if len(path_parts) == 4 else None + if not _is_valid_akv_identifier(secret_name) or ( + secret_version is not None and not _is_valid_akv_identifier(secret_version) + ): + raise ValueError(error_message) + + return f"https://{hostname}", secret_name, secret_version + + +def _is_valid_akv_identifier(identifier: str) -> bool: + """ + Check whether a Key Vault secret name or version uses URL-safe characters. + + Returns: + bool: True when the identifier is valid. + """ + return 1 <= len(identifier) <= 127 and all( + char.isascii() and (char.isalnum() or char == "-") for char in identifier + ) + + +def _create_akv_secret_client(*, vault_url: str, credential: Any) -> "SecretClient": + """ + Create an asynchronous Key Vault client with an explicit retry policy. + + Returns: + SecretClient: Configured asynchronous secret client. + """ + from azure.core.pipeline.policies import AsyncRetryPolicy + from azure.keyvault.secrets.aio import SecretClient + + retry_policy = AsyncRetryPolicy( + retry_total=_AKV_RETRY_TOTAL, + retry_connect=_AKV_RETRY_TOTAL, + retry_read=_AKV_RETRY_TOTAL, + retry_status=_AKV_RETRY_TOTAL, + retry_backoff_factor=_AKV_RETRY_BACKOFF_FACTOR, + ) + return SecretClient(vault_url=vault_url, credential=credential, retry_policy=retry_policy) + + +async def _fetch_akv_secret_value_async( + *, + client: Any, + secret_name: str, + secret_version: str | None, + variable_name: str, +) -> str: + """ + Fetch a referenced Key Vault secret value. + + Returns: + str: The secret value, including an empty string. + + Raises: + ValueError: If the referenced secret has no value. + """ + referenced_secret = await client.get_secret(secret_name, version=secret_version) + if referenced_secret.value is None: + raise ValueError( + f"AKV secret '{secret_name}' referenced by environment variable '{variable_name}' has no value." + ) + return referenced_secret.value + + +def _key_vault_initialization_error(*, message: str, error: Exception) -> KeyVaultInitializationException: + """ + Create a contextual Key Vault exception without losing the original cause. + + 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, + resolve_references_for_output: bool = False, +) -> str: + """ + Load a bootstrap dotenv document and resolve its same-vault secret references. + + References are resolved once. Referenced secret values are treated as terminal + strings and are not interpreted as additional references. + + Authentication uses ``DefaultAzureCredential``, which silently tries managed + identity, Azure CLI, VS Code credentials, etc., and falls back to interactive + browser authentication when running locally. + + Args: + secret_url (str): AKV secret URL in the format + ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. + strict (bool): If True, reject malformed or valueless dotenv entries. + If False, warn and skip those entries. Defaults to True. + silent (bool): If True, suppresses print statements. Defaults to False. + resolve_references_for_output (bool): If True, resolve child references even + when their runtime assignment loses to an existing process value, and + return a native dotenv document containing those resolved values. + + Returns: + str: The validated bootstrap dotenv document, with child-secret references + replaced when ``resolve_references_for_output`` is True. + + Raises: + ImportError: If ``azure-keyvault-secrets`` is not installed. + KeyVaultInitializationException: If the root URL is malformed or the bootstrap environment + document cannot be fully resolved. + ValueError: Compatibility base of ``KeyVaultInitializationException``. + """ + from azure.identity.aio import DefaultAzureCredential + + try: + _print_msg(f"Loading environment from AKV secret: {secret_url}", quiet=silent, log=True) + vault_url, secret_name, secret_version = _parse_akv_secret_url(secret_url) + async with DefaultAzureCredential() as credential: + async with _create_akv_secret_client(vault_url=vault_url, credential=credential) as client: + secret = await client.get_secret(secret_name, version=secret_version) + + if not secret.value: + raise ValueError(f"AKV environment secret has no value: {secret_url}") + + validated_document = _validate_dotenv_document(secret.value, strict=strict, silent=silent) + parsed_environment = dotenv.dotenv_values(stream=io.StringIO(validated_document), interpolate=True) + if not parsed_environment: + raise ValueError(f"AKV environment secret contains no environment entries: {secret_url}") + existing_environment_names = set(os.environ) + loaded = dotenv.load_dotenv( + stream=io.StringIO(validated_document), + override=False, + interpolate=True, + ) + if not loaded: + return validated_document + + resolved_reference_values: dict[str, str] = {} + skipped_reference_names: set[str] = set() + for variable_name, value in parsed_environment.items(): + if value is None: + continue + target = _parse_akv_reference(value) + if target is None: + continue + assignment_wins = variable_name not in existing_environment_names + if not assignment_wins and not resolve_references_for_output: + continue + try: + _, referenced_name, referenced_version = _parse_akv_reference_url( + target=target, + variable_name=variable_name, + expected_vault_url=vault_url, + ) + except ValueError as error: + if strict: + wrapped_error = _key_vault_initialization_error( + message=f"Invalid AKV reference for environment variable '{variable_name}'", + error=error, + ) + raise wrapped_error from error + if assignment_wins: + os.environ.pop(variable_name, None) + skipped_reference_names.add(variable_name) + _warn_about_invalid_akv_reference( + variable_name=variable_name, + error=error, + silent=silent, + ) + continue + try: + resolved_value = await _fetch_akv_secret_value_async( + client=client, + secret_name=referenced_name, + secret_version=referenced_version, + variable_name=variable_name, + ) + resolved_reference_values[variable_name] = resolved_value + if assignment_wins: + os.environ[variable_name] = resolved_value + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to resolve Key Vault reference for environment variable '{variable_name}'", + error=error, + ) + raise wrapped_error from error + if resolve_references_for_output: + return _render_resolved_akv_document( + document=validated_document, + resolved_reference_values=resolved_reference_values, + skipped_reference_names=skipped_reference_names, + ) + return validated_document + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to load Key Vault bootstrap secret '{secret_url}'", + error=error, + ) + raise wrapped_error from error + + +async def _load_environment_async( + *, + env_akv_ref: Sequence[str] | None, + env_files: Sequence[pathlib.Path] | None, + env_akv_strict: bool, + env_akv_write_env: bool = False, + silent: bool, +) -> None: + """ + Load environment sources in precedence order. + + Args: + env_akv_ref (Sequence[str] | None): Optional ordered Key Vault bootstrap secret URLs. + env_files (Sequence[pathlib.Path] | None): Optional ordered local environment files. + env_akv_strict (bool): Whether bootstrap dotenv validation is strict. + env_akv_write_env (bool): Whether to save fetched bootstrap documents to + ``~/.pyrit/.env``. Defaults to False. + silent (bool): Whether initialization messages are suppressed. + + Raises: + ValueError: If a configured source or reference is invalid. + """ + if isinstance(env_akv_ref, str): + raise ValueError("env_akv_ref must be a sequence of Azure Key Vault secret URLs.") + bootstrap_documents: list[str] = [] + if env_akv_ref: + if any(not isinstance(secret_url, str) or not secret_url.strip() for secret_url in env_akv_ref): + raise ValueError("env_akv_ref must contain only non-empty Azure Key Vault secret URLs.") + env_file = path.CONFIGURATION_DIRECTORY_PATH / _AKV_ENV_FILE_NAME + if env_akv_write_env and (env_file.exists() or env_file.is_symlink()): + raise ValueError( + f"Cannot write the resolved Key Vault environment because {env_file} already exists; " + "rename or remove it before enabling env_akv_write_env." + ) + await asyncio.to_thread( + _warn_about_akv_environment_files, + env_files=env_files, + silent=silent, + ) + bootstrap_documents.extend( + [ + await _load_env_from_akv_async( + secret_url=secret_url, + strict=env_akv_strict, + silent=silent, + resolve_references_for_output=env_akv_write_env, + ) + for secret_url in env_akv_ref + ] + ) + + written_env_file: pathlib.Path | None = None + if env_akv_write_env and bootstrap_documents: + written_env_file = await asyncio.to_thread( + _write_akv_env_file, + documents=bootstrap_documents, + silent=silent, + ) + + selected_env_files = env_files + if written_env_file is not None and env_files is not None: + written_path = written_env_file.resolve() + selected_env_files = [env_file for env_file in env_files if env_file.expanduser().resolve() != written_path] + + assignment_fallbacks: dict[str, str | None] = {} + await asyncio.to_thread( + _load_environment_files, + env_files=selected_env_files, + silent=silent, + include_default_base=not (env_akv_ref and env_files is None), + assignment_fallbacks=assignment_fallbacks, + ) + await _resolve_local_akv_references_async( + assignment_fallbacks=assignment_fallbacks, + strict=env_akv_strict, + silent=silent, + ) + + +def _write_akv_env_file(*, documents: Sequence[str], silent: bool) -> pathlib.Path: + """ + Write fetched bootstrap documents with resolved child-secret values. + + Returns: + pathlib.Path: Path to the written dotenv file. + + Raises: + ValueError: If the destination is a symbolic link. + """ + env_file = path.CONFIGURATION_DIRECTORY_PATH / _AKV_ENV_FILE_NAME + env_file.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + if env_file.is_symlink(): + raise ValueError(f"Refusing to write the AKV environment through a symbolic link: {env_file}") + + content = _merge_akv_documents_for_debug(documents=documents) + file_descriptor: int | None = None + temporary_file: pathlib.Path | None = None + try: + file_descriptor, temporary_name = tempfile.mkstemp( + prefix=f"{env_file.name}.", + suffix=".tmp", + dir=env_file.parent, + ) + temporary_file = pathlib.Path(temporary_name) + file_chmod = getattr(os, "fchmod", None) + if file_chmod is not None: + file_chmod(file_descriptor, 0o600) + else: + os.chmod(temporary_file, 0o600) + stream = os.fdopen(file_descriptor, "w", encoding="utf-8", newline="") + file_descriptor = None + with stream: + stream.write(content) + if env_file.is_symlink(): + raise ValueError(f"Refusing to replace a symbolic link with the AKV environment: {env_file}") + os.replace(temporary_file, env_file) + temporary_file = None + finally: + if file_descriptor is not None: + os.close(file_descriptor) + if temporary_file is not None: + with contextlib.suppress(FileNotFoundError): + temporary_file.unlink() + + _print_msg(f"Saved Key Vault bootstrap environment file: {env_file}", quiet=silent, log=True) + return env_file + + +def _merge_akv_documents_for_debug(*, documents: Sequence[str]) -> str: + """ + Merge resolved bootstrap documents using runtime first-document precedence. + + Duplicate assignments within one document are retained because interpolation + depends on assignment order. Assignments established by an earlier document + are omitted from later documents. + + Returns: + str: A native dotenv document with equivalent bootstrap precedence. + """ + established_names: set[str] = set() + merged_bindings: list[str] = [] + for document in documents: + document_names: set[str] = set() + for binding in parse_stream(io.StringIO(document)): + if binding.key is None or binding.key not in established_names: + merged_bindings.append(binding.original.string) + if binding.key is not None: + document_names.add(binding.key) + established_names.update(document_names) + + return "".join(merged_bindings).rstrip("\r\n") + "\n" + + +def _render_resolved_akv_document( + *, + document: str, + resolved_reference_values: Mapping[str, str], + skipped_reference_names: set[str] | None = None, +) -> str: + """ + Replace resolved Key Vault reference assignments with native dotenv values. + + Returns: + str: Dotenv text that preserves non-reference bindings and comments. + """ + skipped_reference_names = skipped_reference_names or set() + rendered_bindings: list[str] = [] + for binding in parse_stream(io.StringIO(document)): + variable_name = binding.key + is_reference = binding.value is not None and _parse_akv_reference(binding.value) is not None + if variable_name in skipped_reference_names and is_reference: + continue + if variable_name is not None and variable_name in resolved_reference_values and is_reference: + original = binding.original.string + export_prefix = "export " if original.lstrip().startswith("export ") else "" + if original.endswith("\r\n"): + newline = "\r\n" + elif original.endswith("\n"): + newline = "\n" + else: + newline = "" + rendered_bindings.append( + f"{export_prefix}{variable_name}=" + f"{_serialize_terminal_dotenv_value(resolved_reference_values[variable_name])}{newline}" + ) + else: + rendered_bindings.append(binding.original.string) + return "".join(rendered_bindings) + + +def _serialize_terminal_dotenv_value(value: str) -> str: + """ + Quote a terminal secret value for a native python-dotenv round trip. + + The empty-name default expression produces a literal dollar sign during + interpolation, preventing terminal ``${NAME}`` text from being reinterpreted. + + Returns: + str: A single-quoted dotenv value. + """ + escaped_value = value.replace("'", "\\'").replace("${", "${:-$}{") + return f"'{escaped_value}'" + + +def _warn_about_invalid_akv_reference(*, variable_name: str, error: ValueError, silent: bool) -> None: + """Warn that a malformed Key Vault reference assignment is being skipped.""" + message = f"Invalid AKV reference for environment variable '{variable_name}' will be skipped: {error}" + if not silent: + print(f"WARNING: {message}") + logger.warning(message) + + +def _parse_akv_reference(value: str) -> str | None: + """ + Parse an exact whole-value Key Vault reference. + + 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 _parse_akv_reference_url( + *, + target: str, + variable_name: str, + expected_vault_url: str | None = None, +) -> tuple[str, str, str | None]: + """ + Parse and optionally constrain a complete Key Vault secret reference. + + Returns: + tuple[str, str, str | None]: Vault URL, secret name, and optional secret version. + + Raises: + ValueError: If the reference is malformed or violates the expected vault constraint. + """ + if not target.casefold().startswith("https://"): + raise ValueError( + f"AKV reference for environment variable '{variable_name}' must use a full secret URL, " + "for example kv:https://my-vault.vault.azure.net/secrets/my-secret." + ) + + referenced_vault_url, secret_name, secret_version = _parse_akv_secret_url(target) + if expected_vault_url and referenced_vault_url.rstrip("/").casefold() != expected_vault_url.rstrip("/").casefold(): + raise ValueError( + f"Cross-vault AKV reference for environment variable '{variable_name}' is not supported. " + f"Expected vault '{expected_vault_url}', got '{referenced_vault_url}'." + ) + + _validate_akv_secret_name(secret_name=secret_name, variable_name=variable_name) + return referenced_vault_url, secret_name, secret_version + + +async def _resolve_local_akv_references_async( + *, + assignment_fallbacks: Mapping[str, str | None], + strict: bool, + silent: bool, +) -> None: + """ + Resolve complete Key Vault references from winning local assignments. + + Raises: + KeyVaultInitializationException: If strict validation or secret retrieval fails. + """ + parsed_references: list[tuple[str, str, str, str | None]] = [] + for variable_name, fallback_value in assignment_fallbacks.items(): + value = os.environ.get(variable_name) + if value is None: + continue + target = _parse_akv_reference(value) + if target is None: + continue + try: + vault_url, secret_name, secret_version = _parse_akv_reference_url( + target=target, + variable_name=variable_name, + ) + except ValueError as error: + if strict: + wrapped_error = _key_vault_initialization_error( + message=f"Invalid AKV reference for environment variable '{variable_name}'", + error=error, + ) + raise wrapped_error from error + if fallback_value is None: + os.environ.pop(variable_name, None) + else: + os.environ[variable_name] = fallback_value + _warn_about_invalid_akv_reference( + variable_name=variable_name, + error=error, + silent=silent, + ) + continue + parsed_references.append((variable_name, vault_url, secret_name, secret_version)) + + if not parsed_references: + return + + from azure.identity.aio import DefaultAzureCredential + + async with DefaultAzureCredential() as credential: + async with contextlib.AsyncExitStack() as client_stack: + clients: dict[str, Any] = {} + for variable_name, vault_url, secret_name, secret_version in parsed_references: + try: + client = clients.get(vault_url) + if client is None: + client = await client_stack.enter_async_context( + _create_akv_secret_client(vault_url=vault_url, credential=credential) + ) + clients[vault_url] = client + os.environ[variable_name] = await _fetch_akv_secret_value_async( + client=client, + secret_name=secret_name, + secret_version=secret_version, + variable_name=variable_name, + ) + except KeyVaultInitializationException: + raise + except Exception as error: + wrapped_error = _key_vault_initialization_error( + message=f"Failed to resolve Key Vault reference for environment variable '{variable_name}'", + error=error, + ) + raise wrapped_error from error + + +def _resolve_akv_secret_reference( + *, + target: str, + variable_name: str, + vault_url: str, +) -> tuple[str, str | None]: + """ + Resolve a full same-vault secret URI. + + Args: + target (str): Full Key Vault secret URI. + variable_name (str): The environment variable receiving the secret. + vault_url (str): The bootstrap document's vault URL. + + Returns: + tuple[str, str | None]: Secret name and optional version. + + Raises: + ValueError: If the target is not a full URI, is invalid, or references another vault. + """ + _, secret_name, secret_version = _parse_akv_reference_url( + target=target, + variable_name=variable_name, + expected_vault_url=vault_url, + ) + return secret_name, secret_version diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 26a29c45b9..6caeafb431 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -95,7 +95,13 @@ class ConfigurationLoader(YamlLoadable): initialization_scripts: List of paths to custom initialization scripts. None means "use defaults", [] means "load nothing". env_files: List of environment file paths to load. - None means "use defaults (.env, .env.local)", [] means "load nothing". + None means auto-discover legacy ``.env`` and supported ``.env.local``; + [] means "load nothing". + env_akv_ref: Ordered list of Key Vault bootstrap secret URLs. + env_akv_strict: Whether malformed or valueless entries in a Key Vault + bootstrap document should fail initialization. + env_akv_write_env: Whether to save fully resolved bootstrap documents with + plaintext child-secret values to ``~/.pyrit/.env`` for debugging. silent: Whether to suppress initialization messages. operator: Name for the current operator, e.g. a team or username. operation: Name for the current operation. @@ -135,6 +141,8 @@ 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 + env_akv_write_env: bool = False silent: bool = False operator: str | None = None operation: str | None = None @@ -147,8 +155,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 +424,8 @@ 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, + env_akv_write_env: bool | None = None, ) -> "ConfigurationLoader": """ Load configuration with optional overrides. @@ -416,7 +441,9 @@ 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. + env_akv_write_env: Override for writing the Key Vault bootstrap environment file. Returns: A merged ConfigurationLoader instance. @@ -477,8 +504,16 @@ 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 + + if env_akv_write_env is not None: + config_data["env_akv_write_env"] = env_akv_write_env + return cls.from_dict(config_data) @classmethod @@ -582,10 +617,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 +649,8 @@ 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, + env_akv_write_env=self.env_akv_write_env, silent=self.silent, ) diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index eb0cf04ff8..91ce1e6112 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -1,16 +1,13 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import io import logging import pathlib from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Literal, get_args -import dotenv - -from pyrit.common import path from pyrit.common.apply_defaults import reset_default_values from pyrit.memory import AzureSQLMemory, CentralMemory, MemoryInterface, SQLiteMemory +from pyrit.setup.akv_initialization import _load_environment_async if TYPE_CHECKING: from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -23,137 +20,6 @@ MemoryDatabaseType = Literal["InMemory", "SQLite", "AzureSQL"] -def _load_environment_files(env_files: Sequence[pathlib.Path] | None, *, silent: bool = False) -> None: - """ - Load environment files in the order they are provided. - Later files override values from earlier files. - - Args: - env_files: Optional sequence of environment file paths. If None, loads default - .env and .env.local from PyRIT home directory (only if they exist). - silent: If True, suppresses print statements about environment file loading. - Defaults to False. - - Raises: - ValueError: If any provided env_files do not exist. - """ - # Validate env_files exist if they were provided - if env_files is not None: - if not silent: - _print_msg(f"Loading custom environment files: {[str(f) for f in env_files]}", quiet=silent, log=True) - for env_file in env_files: - if not env_file.exists(): - raise ValueError(f"Environment file not found: {env_file}") - - # By default load .env and .env.local from home directory of the package - else: - default_files = [] - base_file = path.CONFIGURATION_DIRECTORY_PATH / ".env" - local_file = path.CONFIGURATION_DIRECTORY_PATH / ".env.local" - - if base_file.exists(): - default_files.append(base_file) - if local_file.exists(): - default_files.append(local_file) - - if not silent: - if default_files: - _print_msg( - f"Found default environment files: {[str(f) for f in default_files]}", quiet=silent, log=True - ) - else: - _print_msg( - "No default environment files found. Using system environment variables only.", - quiet=silent, - log=True, - ) - - env_files = default_files - - for env_file in env_files: - dotenv.load_dotenv(env_file, override=True, interpolate=True) - if not silent: - _print_msg(f"Loaded environment file: {env_file}", quiet=silent, log=True) - - -def _print_msg(message: str, quiet: bool, log: bool) -> None: - """ - Print a standard initialization message unless quiet is True. - - Args: - message (str): The message to print and/or log. - quiet (bool): If True, suppresses the initialization message. - log (bool): If True, logs the message using the logger. - """ - if not quiet: - print(message) - if log: - logger.info(message) - - -def _parse_akv_secret_url(secret_url: str) -> tuple[str, str, str | None]: - """ - Parse an AKV secret URL into vault URL, secret name, and optional version. - - Args: - secret_url (str): Full AKV secret URL in the format - ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. - - Returns: - tuple[str, str, str | None]: (vault_url, secret_name, secret_version) - - Raises: - ValueError: If the URL does not match the expected format. - """ - parts = secret_url.split("/secrets/") - if len(parts) != 2: - raise ValueError( - f"Invalid AKV secret URL: '{secret_url}'. " - "Expected format: https://{{vault}}.vault.azure.net/secrets/{{name}}[/{{version}}]" - ) - vault_url = parts[0] - name_parts = parts[1].rstrip("/").split("/") - secret_name = name_parts[0] - secret_version = name_parts[1] if len(name_parts) > 1 else None - return vault_url, secret_name, secret_version - - -async def _load_env_from_akv_async(*, secret_urls: Sequence[str], silent: bool = False) -> None: - """ - Load environment variables from Azure Key Vault secrets. - - Each secret's value is treated as the full contents of a ``.env`` file and - parsed accordingly. Later secrets override values from earlier ones. - - Authentication uses ``DefaultAzureCredential``, which silently tries managed - identity, Azure CLI, VS Code credentials, etc., and falls back to interactive - browser authentication when running locally. - - Args: - secret_urls (Sequence[str]): Sequence of AKV secret URLs to load, each in - the format ``https://{vault}.vault.azure.net/secrets/{name}[/{version}]``. - silent (bool): If True, suppresses print statements. Defaults to False. - - Raises: - ImportError: If ``azure-keyvault-secrets`` is not installed. - ValueError: If a secret URL is malformed. - """ - if not secret_urls: - return - from azure.identity.aio import DefaultAzureCredential - from azure.keyvault.secrets.aio import SecretClient - - credential = DefaultAzureCredential() - for secret_url in secret_urls: - _print_msg(f"Loading environment from AKV secret: {secret_url}", quiet=silent, log=True) - vault_url, secret_name, secret_version = _parse_akv_secret_url(secret_url) - client = SecretClient(vault_url=vault_url, credential=credential) - secret = await client.get_secret(secret_name, version=secret_version) - if secret.value: - dotenv.load_dotenv(stream=io.StringIO(secret.value), override=True) - _print_msg(f"Loaded environment from AKV secret: {secret_url}", quiet=silent, log=True) - - async def _execute_initializers_async(*, initializers: Sequence["PyRITInitializer"]) -> None: """ Execute PyRITInitializer instances in the order provided. @@ -203,6 +69,8 @@ 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, + env_akv_write_env: bool = False, silent: bool = False, **memory_instance_kwargs: Any, ) -> None: @@ -227,22 +95,29 @@ async def initialize_pyrit_async( ``core`` techniques and ``default`` targets are loaded — ``extra`` / per-source technique groups and ``scorer`` target variants remain opt-in. env_files (Sequence[pathlib.Path] | None): Optional sequence of environment file paths to load - in order. If not provided, will load default .env and .env.local files from PyRIT home if they exist. - All paths must be valid pathlib.Path objects. - env_akv_ref (Sequence[str] | None): Optional sequence of Azure Key Vault secret URLs to load. - Each secret's value must be the full contents of a .env file. Loaded before ``env_files`` - so local files take precedence over AKV. Requires ``azure-keyvault-secrets``. + in order. Ordinary files fill missing process values; files named ``.env.local`` override. + If omitted, PyRIT auto-discovers legacy ``.env`` and supported ``.env.local`` files. + env_akv_ref (Sequence[str] | None): Optional ordered Azure Key Vault URLs whose secret values + contain bootstrap dotenv documents. Documents fill missing process values and support + complete-value references to scalar secrets. Requires ``azure-keyvault-secrets``. + env_akv_strict (bool): If True, reject malformed bootstrap entries and Key Vault reference + syntax. If False, warn and skip those entries. Operational Key Vault failures always raise. + env_akv_write_env (bool): If True, write fully resolved bootstrap documents with plaintext + child-secret values to ``~/.pyrit/.env`` for debugging. Defaults to False. silent (bool): If True, suppresses print statements about environment file loading and schema migration. Defaults to False. **memory_instance_kwargs (Any | None): Additional keyword arguments to pass to the memory instance. Raises: - 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, + env_akv_write_env=env_akv_write_env, + silent=silent, + ) # Reset all default values before executing initialization scripts # This ensures a clean state for each initialization diff --git a/pyrit/setup/initializers/targets.py b/pyrit/setup/initializers/targets.py index 308366f734..790b51a87f 100644 --- a/pyrit/setup/initializers/targets.py +++ b/pyrit/setup/initializers/targets.py @@ -338,18 +338,18 @@ class TargetConfig: TargetConfig( registry_name="openai_image_azure", target_class=OpenAIImageTarget, - endpoint_var="OPENAI_IMAGE_ENDPOINT1", - key_var="OPENAI_IMAGE_API_KEY1", - model_var="OPENAI_IMAGE_MODEL1", - underlying_model_var="OPENAI_IMAGE_UNDERLYING_MODEL1", + endpoint_var="AZURE_OPENAI_IMAGE_ENDPOINT1", + key_var="AZURE_OPENAI_IMAGE_API_KEY1", + model_var="AZURE_OPENAI_IMAGE_MODEL1", + underlying_model_var="AZURE_OPENAI_IMAGE_UNDERLYING_MODEL1", ), TargetConfig( registry_name="openai_image_platform", target_class=OpenAIImageTarget, - endpoint_var="OPENAI_IMAGE_ENDPOINT2", - key_var="OPENAI_IMAGE_API_KEY2", - model_var="OPENAI_IMAGE_MODEL2", - underlying_model_var="OPENAI_IMAGE_UNDERLYING_MODEL2", + endpoint_var="AZURE_OPENAI_IMAGE_ENDPOINT2", + key_var="AZURE_OPENAI_IMAGE_API_KEY2", + model_var="AZURE_OPENAI_IMAGE_MODEL2", + underlying_model_var="AZURE_OPENAI_IMAGE_UNDERLYING_MODEL2", ), # ============================================ # TTS Targets (OpenAITTSTarget) @@ -357,18 +357,18 @@ class TargetConfig: TargetConfig( registry_name="openai_tts_azure", target_class=OpenAITTSTarget, - endpoint_var="OPENAI_TTS_ENDPOINT1", - key_var="OPENAI_TTS_KEY1", - model_var="OPENAI_TTS_MODEL1", - underlying_model_var="OPENAI_TTS_UNDERLYING_MODEL1", + endpoint_var="AZURE_OPENAI_TTS_ENDPOINT1", + key_var="AZURE_OPENAI_TTS_KEY1", + model_var="AZURE_OPENAI_TTS_MODEL1", + underlying_model_var="AZURE_OPENAI_TTS_UNDERLYING_MODEL1", ), TargetConfig( registry_name="openai_tts_platform", target_class=OpenAITTSTarget, - endpoint_var="OPENAI_TTS_ENDPOINT2", - key_var="OPENAI_TTS_KEY2", - model_var="OPENAI_TTS_MODEL2", - underlying_model_var="OPENAI_TTS_UNDERLYING_MODEL2", + endpoint_var="AZURE_OPENAI_TTS_ENDPOINT2", + key_var="AZURE_OPENAI_TTS_KEY2", + model_var="AZURE_OPENAI_TTS_MODEL2", + underlying_model_var="AZURE_OPENAI_TTS_UNDERLYING_MODEL2", ), # ============================================ # Video Targets (OpenAIVideoTarget) diff --git a/tests/end_to_end/executor/promptgen/gcg/test_gcg_aml_e2e.py b/tests/end_to_end/executor/promptgen/gcg/test_gcg_aml_e2e.py index 83d5078d80..8b6ed6ce40 100644 --- a/tests/end_to_end/executor/promptgen/gcg/test_gcg_aml_e2e.py +++ b/tests/end_to_end/executor/promptgen/gcg/test_gcg_aml_e2e.py @@ -47,7 +47,7 @@ pytest.importorskip("azure.identity", reason="azure-identity not installed") from pyrit.common.path import HOME_PATH # noqa: E402 -from pyrit.setup.initialization import _load_environment_files # noqa: E402 +from pyrit.setup.akv_initialization import _load_environment_files # noqa: E402 _REQUIRED_ENV_VARS = ( "AZURE_ML_SUBSCRIPTION_ID", diff --git a/tests/integration/setup/test_env_example_drift.py b/tests/integration/setup/test_env_example_drift.py new file mode 100644 index 0000000000..9daa065d23 --- /dev/null +++ b/tests/integration/setup/test_env_example_drift.py @@ -0,0 +1,145 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import os +import pathlib +import re +import subprocess +from unittest import mock + +from dotenv import dotenv_values + +_ENV_EXAMPLE_PATH_ENV = "PYRIT_ENV_EXAMPLE_PATH" +_REPOSITORY_ROOT_ENV = "PYRIT_REPOSITORY_ROOT" +_ENVIRONMENT_NAME_PATTERN = re.compile(r"(? pathlib.Path: + configured_root = os.getenv(_REPOSITORY_ROOT_ENV) + if configured_root: + root = pathlib.Path(configured_root) + if root.is_dir(): + return root + raise AssertionError(f"{_REPOSITORY_ROOT_ENV} does not identify a directory: {root}") + + result = subprocess.run( + ["git", "rev-parse", "--show-toplevel"], + capture_output=True, + check=False, + text=True, + ) + if result.returncode != 0: + raise AssertionError("Could not locate the repository root with git.") + return pathlib.Path(result.stdout.strip()) + + +def _get_env_example_path(*, repository_root: pathlib.Path) -> pathlib.Path: + configured_path = os.getenv(_ENV_EXAMPLE_PATH_ENV) + path = pathlib.Path(configured_path) if configured_path else repository_root / ".env_example" + if not path.is_file(): + raise AssertionError(f"Could not locate .env_example at {path}.") + return path + + +def _grep_repository_for_environment_names(*, environment_names: set[str], repository_root: pathlib.Path) -> set[str]: + grep_pattern = "(" + "|".join(sorted(environment_names)) + ")" + result = subprocess.run( + ["git", "grep", "-I", "-h", "-E", grep_pattern, "--", ".", ":(exclude).env_example"], + capture_output=True, + check=False, + cwd=repository_root, + text=True, + ) + if result.returncode not in {0, 1}: + raise AssertionError(f"Could not search tracked repository files with git: {result.stderr.strip()}") + return environment_names & set(_ENVIRONMENT_NAME_PATTERN.findall(result.stdout)) + + +def _find_referenced_environment_names( + *, + environment_names: set[str], + repository_root: pathlib.Path, + env_example_path: pathlib.Path, +) -> set[str]: + example_contents = env_example_path.read_text(encoding="utf-8") + example_without_assignment_names = _DOTENV_ASSIGNMENT_NAME_PATTERN.sub("=", example_contents) + referenced_names = environment_names & set(_ENVIRONMENT_NAME_PATTERN.findall(example_without_assignment_names)) + referenced_names.update( + _grep_repository_for_environment_names( + environment_names=environment_names, + repository_root=repository_root, + ) + ) + return referenced_names + + +def test_env_example_names_are_referenced_in_repository() -> None: + """Catch example entries with no weak textual reference in tracked repository files.""" + repository_root = _get_repository_root() + env_example_path = _get_env_example_path(repository_root=repository_root) + environment_names = set(dotenv_values(dotenv_path=env_example_path, interpolate=False)) + assert environment_names, ".env_example contains no dotenv assignments." + + referenced_names = _find_referenced_environment_names( + environment_names=environment_names, + repository_root=repository_root, + env_example_path=env_example_path, + ) + unreferenced_names = environment_names - referenced_names + assert not unreferenced_names, ".env_example contains names with no tracked repository reference: " + ", ".join( + sorted(unreferenced_names) + ) + + +def test_env_example_url_values_are_not_wrapped_in_angle_brackets() -> None: + """Ensure URL placeholder styling does not become part of parsed dotenv values.""" + repository_root = _get_repository_root() + env_example_path = _get_env_example_path(repository_root=repository_root) + values = dotenv_values(dotenv_path=env_example_path, interpolate=False) + + wrapped_names = {name for name, value in values.items() if value and ("<" in value or ">" in value)} + assert not wrapped_names, ".env_example contains values wrapped in angle brackets: " + ", ".join( + sorted(wrapped_names) + ) + + +def test_env_example_comment_blocks_do_not_contain_blank_lines() -> None: + """Keep consecutive comment lines together so the example remains compact.""" + repository_root = _get_repository_root() + env_example_path = _get_env_example_path(repository_root=repository_root) + contents = env_example_path.read_text(encoding="utf-8") + + assert not _BLANK_LINE_BETWEEN_COMMENTS_PATTERN.search(contents), ( + ".env_example contains a blank line between consecutive comment lines." + ) + + +def test_env_example_aliases_resolve_in_assignment_order() -> None: + """Ensure complete-value aliases resolve to their sources without ambient environment values.""" + repository_root = _get_repository_root() + env_example_path = _get_env_example_path(repository_root=repository_root) + raw_values = dotenv_values(dotenv_path=env_example_path, interpolate=False) + aliases = { + name: match.group(1) + for name, value in raw_values.items() + if value and (match := _DOTENV_COMPLETE_REFERENCE_PATTERN.fullmatch(value)) + } + assert aliases, ".env_example contains no complete-value aliases." + + with mock.patch.dict(os.environ, {}, clear=True): + resolved_values = dotenv_values(dotenv_path=env_example_path, interpolate=True) + + unresolved_names = {name for name in aliases if not resolved_values.get(name)} + assert not unresolved_names, ".env_example contains aliases that resolve to empty values: " + ", ".join( + sorted(unresolved_names) + ) + + mismatched_names = { + name for name, source_name in aliases.items() if resolved_values[name] != resolved_values.get(source_name) + } + assert not mismatched_names, ".env_example contains aliases that differ from their sources: " + ", ".join( + sorted(mismatched_names) + ) diff --git a/tests/integration/targets/test_targets_and_secrets.py b/tests/integration/targets/test_targets_and_secrets.py index e2ec9da733..2a15ae6397 100644 --- a/tests/integration/targets/test_targets_and_secrets.py +++ b/tests/integration/targets/test_targets_and_secrets.py @@ -561,23 +561,23 @@ async def test_connect_openai_completion(sqlite_instance: SQLiteMemory) -> None: [ ("OPENAI_IMAGE_ENDPOINT", None, "OPENAI_IMAGE_MODEL"), pytest.param( - "OPENAI_IMAGE_ENDPOINT1", + "AZURE_OPENAI_IMAGE_ENDPOINT1", None, - "OPENAI_IMAGE_MODEL1", + "AZURE_OPENAI_IMAGE_MODEL1", marks=pytest.mark.run_only_if_all_tests, ), # gpt-image-1.5 pytest.param( - "OPENAI_IMAGE_ENDPOINT1", - "OPENAI_IMAGE_API_KEY1", - "OPENAI_IMAGE_MODEL1", + "AZURE_OPENAI_IMAGE_ENDPOINT1", + "AZURE_OPENAI_IMAGE_API_KEY1", + "AZURE_OPENAI_IMAGE_MODEL1", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-image1-api-key", ), - ("OPENAI_IMAGE_ENDPOINT2", None, "OPENAI_IMAGE_MODEL2"), # gpt-image-1 + ("AZURE_OPENAI_IMAGE_ENDPOINT2", None, "AZURE_OPENAI_IMAGE_MODEL2"), # gpt-image-1 pytest.param( - "OPENAI_IMAGE_ENDPOINT2", - "OPENAI_IMAGE_API_KEY2", - "OPENAI_IMAGE_MODEL2", + "AZURE_OPENAI_IMAGE_ENDPOINT2", + "AZURE_OPENAI_IMAGE_API_KEY2", + "AZURE_OPENAI_IMAGE_MODEL2", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-image2-api-key", ), @@ -626,7 +626,7 @@ async def test_connect_image( [ pytest.param(None, id="entra"), pytest.param( - "OPENAI_IMAGE_API_KEY2", + "AZURE_OPENAI_IMAGE_API_KEY2", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="api-key", ), @@ -645,8 +645,8 @@ async def test_image_editing_single_image( 2. The edit endpoint is correctly called 3. The output image file is created """ - endpoint_value = _get_required_env_var("OPENAI_IMAGE_ENDPOINT2") - model_name_value = os.getenv("OPENAI_IMAGE_MODEL2") or "gpt-image-1" + endpoint_value = _get_required_env_var("AZURE_OPENAI_IMAGE_ENDPOINT2") + model_name_value = os.getenv("AZURE_OPENAI_IMAGE_MODEL2") or "gpt-image-1" target = OpenAIImageTarget( endpoint=endpoint_value, @@ -686,7 +686,7 @@ async def test_image_editing_single_image( [ pytest.param(None, id="entra"), pytest.param( - "OPENAI_IMAGE_API_KEY2", + "AZURE_OPENAI_IMAGE_API_KEY2", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="api-key", ), @@ -704,8 +704,8 @@ async def test_image_editing_multiple_images( 1. Multiple images can be passed to the edit endpoint 2. The model processes multiple image inputs correctly """ - endpoint_value = _get_required_env_var("OPENAI_IMAGE_ENDPOINT2") - model_name_value = os.getenv("OPENAI_IMAGE_MODEL2") or "gpt-image-1" + endpoint_value = _get_required_env_var("AZURE_OPENAI_IMAGE_ENDPOINT2") + model_name_value = os.getenv("AZURE_OPENAI_IMAGE_MODEL2") or "gpt-image-1" target = OpenAIImageTarget( endpoint=endpoint_value, @@ -749,19 +749,19 @@ async def test_image_editing_multiple_images( @pytest.mark.parametrize( ("endpoint", "api_key_env_var", "model_name"), [ - ("OPENAI_TTS_ENDPOINT1", None, "OPENAI_TTS_MODEL1"), + ("AZURE_OPENAI_TTS_ENDPOINT1", None, "AZURE_OPENAI_TTS_MODEL1"), pytest.param( - "OPENAI_TTS_ENDPOINT1", - "OPENAI_TTS_KEY1", - "OPENAI_TTS_MODEL1", + "AZURE_OPENAI_TTS_ENDPOINT1", + "AZURE_OPENAI_TTS_KEY1", + "AZURE_OPENAI_TTS_MODEL1", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-tts1-api-key", ), - ("OPENAI_TTS_ENDPOINT2", None, "OPENAI_TTS_MODEL2"), + ("AZURE_OPENAI_TTS_ENDPOINT2", None, "AZURE_OPENAI_TTS_MODEL2"), pytest.param( - "OPENAI_TTS_ENDPOINT2", - "OPENAI_TTS_KEY2", - "OPENAI_TTS_MODEL2", + "AZURE_OPENAI_TTS_ENDPOINT2", + "AZURE_OPENAI_TTS_KEY2", + "AZURE_OPENAI_TTS_MODEL2", marks=pytest.mark.skip(reason=_AZURE_KEY_AUTH_DISABLED_REASON), id="openai-tts2-api-key", ), diff --git a/tests/unit/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_akv_initialization.py b/tests/unit/setup/test_akv_initialization.py new file mode 100644 index 0000000000..c0f3d465e4 --- /dev/null +++ b/tests/unit/setup/test_akv_initialization.py @@ -0,0 +1,1216 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import os +import pathlib +import tempfile +import types +import warnings +from unittest import mock + +import pytest +from azure.core.exceptions import ResourceNotFoundError +from dotenv import dotenv_values + +from pyrit.exceptions import KeyVaultInitializationException +from pyrit.setup import IN_MEMORY, initialize_pyrit_async +from pyrit.setup.akv_initialization import ( + _load_env_from_akv_async, + _load_environment_async, + _load_environment_files, + _parse_akv_reference, + _parse_akv_secret_url, + _warn_about_akv_environment_files, + _write_akv_env_file, +) + + +class TestLoadEnvironmentFiles: + """Tests for _load_environment_files function and env_files parameter in initialize_pyrit_async.""" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_loads_default_env_files_when_none_provided(self, mock_config_path): + """Test that default .env and .env.local files are loaded when env_files is None.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text("VAR1=value1") + env_local_file.write_text("VAR2=value2") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch.dict(os.environ, {}, clear=True), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): + loaded = _load_environment_files(env_files=None) + + assert loaded is True + assert os.environ["VAR1"] == "value1" + assert os.environ["VAR2"] == "value2" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_only_loads_existing_default_files(self, mock_config_path): + """Test that only existing default files are loaded.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_file.write_text("VAR1=value1") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch.dict(os.environ, {}, clear=True), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): + loaded = _load_environment_files(env_files=None) + + assert loaded is True + assert os.environ["VAR1"] == "value1" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_default_env_preserves_process_environment(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VAR=legacy\nLEGACY_ONLY=legacy") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch.dict(os.environ, {"VAR": "process"}, clear=True), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): + loaded = _load_environment_files(env_files=None, silent=True) + + assert loaded is True + assert os.environ["VAR"] == "process" + assert os.environ["LEGACY_ONLY"] == "legacy" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_default_env_local_overrides_process_environment_and_env(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VAR=legacy") + (temp_path / ".env.local").write_text("VAR=local") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch.dict(os.environ, {"VAR": "process"}, clear=True), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): + loaded = _load_environment_files(env_files=None, silent=True) + + assert loaded is True + assert os.environ["VAR"] == "local" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_excludes_default_env_when_loading_local_override(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text("VAR=base") + env_local_file.write_text("VAR=local") + + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None, include_default_base=False) + + assert loaded is True + assert os.environ["VAR"] == "local" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_returns_false_when_no_default_files_exist(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=None) + + assert loaded is False + assert os.environ == {} + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_auto_discovered_env_warns_with_removal_version(self, mock_config_path, caplog, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_file.write_text("VAR=base") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + pytest.warns(DeprecationWarning, match=r"\.env.*removed in 1\.3\.0.*\.env\.local"), + ): + _load_environment_files(env_files=None) + + output = capsys.readouterr().out + assert f"WARNING: Auto-discovered {env_file} is deprecated" in output + assert "Use env_akv_ref or ~/.pyrit/.env.local instead" in output + assert f"Auto-discovered {env_file} is deprecated" in caplog.text + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_explicit_env_file_does_not_emit_legacy_deprecation(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + explicit_env = temp_path / ".env" + explicit_env.write_text("VAR=explicit") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + loaded = _load_environment_files(env_files=[explicit_env], silent=True) + + assert loaded is True + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_akv_legacy_env_warning_respects_silent(self, mock_config_path, caplog, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VAR=base") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): + _warn_about_akv_environment_files(env_files=None, silent=True) + + assert capsys.readouterr().out == "" + assert "will be ignored because env_akv_ref is configured" in caplog.text + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_akv_ignores_auto_discovered_env_and_loads_env_local(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env").write_text("VALUE=legacy") + (temp_path / ".env.local").write_text("VALUE=local") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch( + "pyrit.setup.akv_initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + return_value="VALUE=akv\n", + ), + mock.patch("pyrit.setup.akv_initialization._load_environment_files") as mock_load_files, + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): + await _load_environment_async( + env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=None, + env_akv_strict=True, + env_akv_write_env=False, + silent=True, + ) + + assert mock_load_files.call_args.kwargs["env_files"] is None + assert mock_load_files.call_args.kwargs["include_default_base"] is False + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + async def test_akv_debug_mode_rejects_existing_env_before_fetch(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_file.write_text("VALUE=legacy") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch( + "pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock + ) as mock_load_akv, + pytest.raises(ValueError, match=r"already exists.*rename or remove"), + ): + await _load_environment_async( + env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=None, + env_akv_strict=True, + env_akv_write_env=True, + silent=True, + ) + + mock_load_akv.assert_not_awaited() + + async def test_loads_custom_env_files_in_order(self): + """Test that custom env_files are loaded in the order provided.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env1 = temp_path / ".env.test" + env2 = temp_path / ".env.prod" + env3 = temp_path / ".env.local" + + # Create files + env1.write_text("VAR=test") + env2.write_text("VAR=prod") + env3.write_text("VAR=local") + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env1, env2, env3]) + + assert loaded is True + assert os.environ["VAR"] == "local" + + async def test_explicit_files_only_override_when_named_env_local(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + first_file = temp_path / "first.env" + second_file = temp_path / "second.env" + local_file = temp_path / "nested" / ".env.local" + local_file.parent.mkdir() + first_file.write_text("PROCESS_VALUE=first\nFILE_VALUE=first") + second_file.write_text("PROCESS_VALUE=second\nFILE_VALUE=second\nSECOND_ONLY=second") + local_file.write_text("PROCESS_VALUE=local\nFILE_VALUE=local") + + with mock.patch.dict(os.environ, {"PROCESS_VALUE": "process"}, clear=True): + loaded = _load_environment_files(env_files=[first_file, second_file, local_file], silent=True) + + assert loaded is True + assert os.environ["PROCESS_VALUE"] == "local" + assert os.environ["FILE_VALUE"] == "local" + assert os.environ["SECOND_ONLY"] == "second" + + with mock.patch.dict(os.environ, {"PROCESS_VALUE": "process"}, clear=True): + _load_environment_files(env_files=[first_file, second_file], silent=True) + + assert os.environ["PROCESS_VALUE"] == "process" + assert os.environ["FILE_VALUE"] == "first" + + async def test_load_environment_files_interpolates_in_assignment_order(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("A=one\nB=${A}\nA=two\nC=${A}") + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert os.environ["A"] == "two" + assert os.environ["B"] == "one" + assert os.environ["C"] == "two" + + async def test_load_environment_files_honors_python_dotenv_disabled(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text("DISABLED_VALUE=not-loaded") + + with mock.patch.dict(os.environ, {"PYTHON_DOTENV_DISABLED": "true"}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert "DISABLED_VALUE" not in os.environ + + async def test_load_environment_async_write_env_writes_resolved_native_bootstrap(self): + credential, client = _create_mock_akv_clients() + document = ( + "# Bootstrap values\n" + "BASE=bootstrap\n" + "DERIVED=${BASE}\n" + "API_KEY=kv:https://vault.vault.azure.net/secrets/api-key\n" + ) + resolved_api_key = "line one\nquote' and literal ${UNRELATED}" + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value=document), + types.SimpleNamespace(value=resolved_api_key), + ] + ) + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + (temp_path / ".env.local").write_text("API_KEY=local-key\nLOCAL_ONLY=local", encoding="utf-8") + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch.dict( + os.environ, + { + "BASE": "process", + "API_KEY": "process-key", + "PROCESS_ONLY": "not-written", + "UNRELATED": "changed", + }, + clear=True, + ), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + await _load_environment_async( + env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=None, + env_akv_strict=True, + env_akv_write_env=True, + silent=True, + ) + + assert os.environ["BASE"] == "process" + assert os.environ["API_KEY"] == "local-key" + assert os.environ["LOCAL_ONLY"] == "local" + + written_env = temp_path / ".env" + assert written_env.is_file() + assert not (temp_path / ".env.new").exists() + content = written_env.read_text(encoding="utf-8") + assert "# Bootstrap values" in content + assert "kv:" not in content + assert "PROCESS_ONLY" not in content + assert "LOCAL_ONLY" not in content + + with mock.patch.dict(os.environ, {}, clear=True): + written_values = dotenv_values(dotenv_path=written_env, interpolate=True) + + assert written_values == { + "BASE": "bootstrap", + "DERIVED": "bootstrap", + "API_KEY": resolved_api_key, + } + assert client.get_secret.await_args_list == [ + mock.call("bootstrap", version=None), + mock.call("api-key", version=None), + ] + + async def test_load_environment_async_write_env_filters_generated_explicit_file(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + generated_env = temp_path / ".env" + local_env = temp_path / ".env.local" + local_env.write_text("LOCAL=value", encoding="utf-8") + + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch( + "pyrit.setup.akv_initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + return_value="VALUE=bootstrap\n", + ), + mock.patch( + "pyrit.setup.akv_initialization._load_environment_files", return_value=True + ) as mock_load_environment_files, + ): + await _load_environment_async( + env_akv_ref=["https://vault.vault.azure.net/secrets/bootstrap"], + env_files=[generated_env, local_env], + env_akv_strict=True, + env_akv_write_env=True, + silent=True, + ) + + assert mock_load_environment_files.call_args.kwargs["env_files"] == [local_env] + assert mock_load_environment_files.call_args.kwargs["include_default_base"] is True + + async def test_load_environment_async_write_env_preserves_first_bootstrap_value(self): + documents = [ + "SHARED=first\nFIRST_ONLY=first\n", + "SHARED=second\nSECOND_ONLY=second\n", + ] + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch( + "pyrit.setup.akv_initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + side_effect=lambda **kwargs: documents.pop(0), + ), + ): + await _load_environment_async( + env_akv_ref=[ + "https://vault.vault.azure.net/secrets/first", + "https://vault.vault.azure.net/secrets/second", + ], + env_files=[], + env_akv_strict=True, + env_akv_write_env=True, + silent=True, + ) + + with mock.patch.dict(os.environ, {}, clear=True): + written_values = dotenv_values(dotenv_path=temp_path / ".env", interpolate=True) + + assert written_values == { + "SHARED": "first", + "FIRST_ONLY": "first", + "SECOND_ONLY": "second", + } + + def test_write_akv_env_file_secures_descriptor_before_writing(self): + events: list[str] = [] + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + temporary_file = temp_path / ".env.test.tmp" + stream = mock.MagicMock() + stream.write.side_effect = lambda content: events.append(f"write:{content}") + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch( + "pyrit.setup.akv_initialization.tempfile.mkstemp", + side_effect=lambda **kwargs: events.append("create") or (7, str(temporary_file)), + ), + mock.patch( + "pyrit.setup.akv_initialization.os.fchmod", + side_effect=lambda *args: events.append("fchmod"), + create=True, + ), + mock.patch( + "pyrit.setup.akv_initialization.os.fdopen", + side_effect=lambda *args, **kwargs: events.append("fdopen") or stream, + ), + mock.patch( + "pyrit.setup.akv_initialization.os.replace", + side_effect=lambda *args: events.append("replace"), + ), + ): + _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) + + assert events == ["create", "fchmod", "fdopen", "write:VALUE=bootstrap\n", "replace"] + + def test_write_akv_env_file_preserves_existing_file_when_replace_fails(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_file.write_text("ORIGINAL=value\n", encoding="utf-8") + + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.akv_initialization.os.replace", side_effect=OSError("replace failed")), + pytest.raises(OSError, match="replace failed"), + ): + _write_akv_env_file(documents=["NEW=value\n"], silent=True) + + assert env_file.read_text(encoding="utf-8") == "ORIGINAL=value\n" + assert list(temp_path.glob(".env.*.tmp")) == [] + + @pytest.mark.skipif(os.name != "posix", reason="POSIX permission bits are not enforced on this platform.") + def test_write_akv_env_file_uses_owner_only_permissions(self): + with tempfile.TemporaryDirectory() as temp_dir: + configuration_directory = pathlib.Path(temp_dir) / ".pyrit" + with mock.patch( + "pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", configuration_directory + ): + env_file = _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) + + assert configuration_directory.stat().st_mode & 0o777 == 0o700 + assert env_file.stat().st_mode & 0o777 == 0o600 + + def test_write_akv_env_file_rejects_symbolic_link(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + target = temp_path / "target" + target.write_text("unchanged", encoding="utf-8") + env_file = temp_path / ".env" + try: + env_file.symlink_to(target) + except OSError: + pytest.skip("Symbolic links are unavailable on this platform.") + + with ( + mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + pytest.raises(ValueError, match="symbolic link"), + ): + _write_akv_env_file(documents=["VALUE=bootstrap\n"], silent=True) + + assert target.read_text(encoding="utf-8") == "unchanged" + + async def test_direct_local_file_loader_keeps_pyrit_references_literal(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / ".env" + env_file.write_text( + "BASE_VALUE=base\nKV_REFERENCE=kv:api-key\nENV_REFERENCE=env:SOURCE_VALUE\nINTERPOLATED=${BASE_VALUE}" + ) + + with mock.patch.dict(os.environ, {}, clear=True): + loaded = _load_environment_files(env_files=[env_file], silent=True) + + assert loaded is True + assert os.environ["KV_REFERENCE"] == "kv:api-key" + assert os.environ["ENV_REFERENCE"] == "env:SOURCE_VALUE" + assert os.environ["INTERPOLATED"] == "base" + + @mock.patch("pyrit.setup.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH") + def test_default_local_file_can_interpolate_base_file_but_not_reverse(self, mock_config_path): + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env" + env_local_file = temp_path / ".env.local" + env_file.write_text( + "OPENAI_CHAT_ENDPOINT=https://example.openai.azure.com/openai/v1\nFROM_LATER_LOCAL=${LOCAL_ONLY}" + ) + env_local_file.write_text("FOOBAR=${OPENAI_CHAT_ENDPOINT}\nLOCAL_ONLY=local") + mock_config_path.__truediv__ = lambda self, other: temp_path / other + + with ( + mock.patch.dict(os.environ, {}, clear=True), + pytest.warns(DeprecationWarning, match=r"removed in 1\.3\.0"), + ): + loaded = _load_environment_files(env_files=None, silent=True) + + assert loaded is True + assert os.environ["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" + + @pytest.mark.parametrize( + ("file_name", "initial_environment"), + [ + ("custom.env", {}), + (".env.local", {"API_KEY": "process-key"}), + ], + ) + async def test_load_environment_async_resolves_local_akv_reference(self, file_name, initial_environment): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="resolved-key")) + + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / file_name + env_file.write_text("API_KEY=kv:https://local-vault.vault.azure.net/secrets/api-key/version-1") + + with ( + mock.patch.dict(os.environ, initial_environment, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, + ): + await _load_environment_async( + env_akv_ref=None, + env_files=[env_file], + env_akv_strict=True, + env_akv_write_env=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "resolved-key" + + _assert_mock_akv_client_created( + mock_client_cls, + vault_url="https://local-vault.vault.azure.net", + credential=credential, + ) + client.get_secret.assert_awaited_once_with("api-key", version="version-1") + + async def test_load_environment_async_does_not_fetch_local_reference_that_loses_to_process_value(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / "custom.env" + env_file.write_text("API_KEY=kv:https://local-vault.vault.azure.net/secrets/api-key") + + with ( + mock.patch.dict(os.environ, {"API_KEY": "process-key"}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + ): + await _load_environment_async( + env_akv_ref=None, + env_files=[env_file], + env_akv_strict=True, + env_akv_write_env=False, + silent=True, + ) + + assert os.environ["API_KEY"] == "process-key" + + mock_credential_cls.assert_not_called() + + async def test_load_environment_async_strict_rejects_malformed_local_akv_reference_before_authentication(self): + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / "custom.env" + env_file.write_text("API_KEY=kv:api-key") + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + pytest.raises(KeyVaultInitializationException, match="must use a full secret URL"), + ): + await _load_environment_async( + env_akv_ref=None, + env_files=[env_file], + env_akv_strict=True, + env_akv_write_env=False, + silent=True, + ) + + mock_credential_cls.assert_not_called() + + async def test_load_environment_async_non_strict_skips_malformed_local_akv_reference(self, caplog, capsys): + with tempfile.TemporaryDirectory() as temp_dir: + env_local_file = pathlib.Path(temp_dir) / ".env.local" + env_local_file.write_text("API_KEY=kv:api-key") + + with ( + mock.patch.dict(os.environ, {"API_KEY": "process-key"}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + ): + await _load_environment_async( + env_akv_ref=None, + env_files=[env_local_file], + env_akv_strict=False, + env_akv_write_env=False, + silent=False, + ) + + assert os.environ["API_KEY"] == "process-key" + + mock_credential_cls.assert_not_called() + assert ( + "WARNING: Invalid AKV reference for environment variable 'API_KEY' will be skipped" + in capsys.readouterr().out + ) + assert "API_KEY" in caplog.text + + async def test_load_environment_async_non_strict_still_raises_for_missing_local_secret(self): + credential, client = _create_mock_akv_clients() + missing_error = ResourceNotFoundError(message="Secret was not found") + client.get_secret = mock.AsyncMock(side_effect=missing_error) + + with tempfile.TemporaryDirectory() as temp_dir: + env_file = pathlib.Path(temp_dir) / "custom.env" + env_file.write_text("API_KEY=kv:https://local-vault.vault.azure.net/secrets/missing") + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises( + KeyVaultInitializationException, match="Failed to resolve Key Vault reference" + ) as exc_info, + ): + await _load_environment_async( + env_akv_ref=None, + env_files=[env_file], + env_akv_strict=False, + env_akv_write_env=False, + silent=True, + ) + + assert exc_info.value.__cause__ is missing_error + + async def test_raises_error_for_nonexistent_env_file(self): + """Test that ValueError is raised for non-existent env file.""" + nonexistent = pathlib.Path("/nonexistent/path/.env") + + with pytest.raises(ValueError, match="Environment file not found"): + _load_environment_files(env_files=[nonexistent]) + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_pyrit_with_custom_env_files(self, mock_set_memory): + """Test initialize_pyrit_async with custom env_files.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + env_file = temp_path / ".env.custom" + env_file.write_text("CUSTOM_VAR=custom_value") + + # Should not raise an error + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[env_file], load_defaults=False) + + mock_set_memory.assert_called_once() + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_pyrit_raises_for_nonexistent_env_file(self, mock_set_memory): + """Test that initialize_pyrit_async raises ValueError for non-existent env file.""" + nonexistent = pathlib.Path("/nonexistent/.env") + + with pytest.raises(ValueError, match="Environment file not found"): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[nonexistent]) + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_custom_env_files_override_default_behavior(self, mock_set_memory): + """Test that passing custom env_files prevents loading default files.""" + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = pathlib.Path(temp_dir) + + # Create default files + default_env = temp_path / ".env" + default_env_local = temp_path / ".env.local" + default_env.write_text("DEFAULT=value") + default_env_local.write_text("DEFAULT_LOCAL=value") + + # Create custom file + custom_env = temp_path / ".env.custom" + custom_env.write_text("CUSTOM=value") + + with mock.patch.dict(os.environ, {}, clear=True): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[custom_env], load_defaults=False) + + assert os.environ["CUSTOM"] == "value" + assert "DEFAULT" not in os.environ + assert "DEFAULT_LOCAL" not in os.environ + + +def _create_mock_akv_clients() -> tuple[mock.MagicMock, mock.MagicMock]: + credential = mock.MagicMock() + credential.__aenter__ = mock.AsyncMock(return_value=credential) + credential.__aexit__ = mock.AsyncMock(return_value=None) + client = mock.MagicMock() + client.__aenter__ = mock.AsyncMock(return_value=client) + client.__aexit__ = mock.AsyncMock(return_value=None) + return credential, client + + +def _assert_mock_akv_client_created( + mock_client_cls: mock.MagicMock, + *, + vault_url: str, + credential: mock.MagicMock, +) -> None: + mock_client_cls.assert_called_once() + call_kwargs = mock_client_cls.call_args.kwargs + assert call_kwargs["vault_url"] == vault_url + assert call_kwargs["credential"] is credential + retry_policy = call_kwargs["retry_policy"] + assert retry_policy.total_retries == 3 + assert retry_policy.connect_retries == 3 + assert retry_policy.read_retries == 3 + assert retry_policy.status_retries == 3 + assert retry_policy.backoff_factor == 0.8 + + +class TestAkvEnvironmentLoading: + """Tests for AKV URL parsing and env loading helpers.""" + + @pytest.mark.parametrize("prefix", ["kv", "akv", "azure_key_vault", "env_akv_ref"]) + def test_parse_akv_reference_accepts_aliases(self, prefix): + secret_url = "https://myvault.vault.azure.net/secrets/api-key" + + assert _parse_akv_reference(f"{prefix}:{secret_url}") == secret_url + + @pytest.mark.parametrize( + "value", + [ + "env:SOURCE_VALUE", + "literal:kv:https://myvault.vault.azure.net/secrets/api-key", + "@Microsoft.KeyVault(SecretUri=https://myvault.vault.azure.net/secrets/api-key)", + ], + ) + def test_parse_akv_reference_ignores_non_akv_syntax(self, value): + assert _parse_akv_reference(value) is None + + def test_parse_akv_secret_url_with_version(self): + url = "https://myvault.vault.azure.net/secrets/my-secret/abc123" + + vault_url, secret_name, secret_version = _parse_akv_secret_url(url) + + assert vault_url == "https://myvault.vault.azure.net" + assert secret_name == "my-secret" + assert secret_version == "abc123" + + def test_parse_akv_secret_url_without_version(self): + url = "https://myvault.vault.azure.net/secrets/my-secret" + + vault_url, secret_name, secret_version = _parse_akv_secret_url(url) + + assert vault_url == "https://myvault.vault.azure.net" + assert secret_name == "my-secret" + assert secret_version is None + + @pytest.mark.parametrize("dns_suffix", ["vault.azure.net", "vault.azure.cn", "vault.usgovcloudapi.net"]) + def test_parse_akv_secret_url_accepts_supported_clouds(self, dns_suffix): + url = f"https://myvault.{dns_suffix}/secrets/my-secret/version-1" + + vault_url, secret_name, secret_version = _parse_akv_secret_url(url) + + assert vault_url == f"https://myvault.{dns_suffix}" + assert secret_name == "my-secret" + assert secret_version == "version-1" + + @pytest.mark.parametrize( + "url", + [ + "http://myvault.vault.azure.net/secrets/my-secret", + "https://attacker.example/secrets/my-secret", + "https://myvault.vault.azure.net.attacker.example/secrets/my-secret", + "https://nested.myvault.vault.azure.net/secrets/my-secret", + "https://user@myvault.vault.azure.net/secrets/my-secret", + "https://myvault.vault.azure.net:443/secrets/my-secret", + "https://myvault.vault.azure.net/not-secrets/my-secret", + "https://myvault.vault.azure.net/secrets", + "https://myvault.vault.azure.net/secrets/my-secret/", + "https://myvault.vault.azure.net/secrets/my-secret/version/extra", + "https://myvault.vault.azure.net/secrets/my-secret?api-version=7.4", + "https://myvault.vault.azure.net/secrets/my-secret#fragment", + "https://myvault.vault.azure.net/secrets/my%2Fsecret", + ], + ) + def test_parse_akv_secret_url_invalid_raises(self, url): + with pytest.raises(ValueError, match="Invalid AKV secret URL"): + _parse_akv_secret_url(url) + + async def test_load_env_from_akv_async_rejects_non_azure_host_before_authentication(self): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential") as mock_credential_cls, + mock.patch("pyrit.setup.akv_initialization._create_akv_secret_client") as mock_create_client, + pytest.raises(KeyVaultInitializationException, match="attacker.example"), + ): + await _load_env_from_akv_async( + secret_url="https://attacker.example/secrets/bootstrap", + silent=True, + ) + + mock_credential_cls.assert_not_called() + mock_create_client.assert_not_called() + + async def test_load_env_from_akv_async_loads_bootstrap_and_resolves_child_secrets(self): + credential, client = _create_mock_akv_clients() + root_document = ( + "DIRECT=from-bootstrap\n" + "FROM_ENV=${SOURCE_VALUE}\n" + "FROM_KV=kv:https://myvault.vault.azure.net/secrets/api-key\n" + "PINNED_KV=kv:https://myvault.vault.azure.net/secrets/api-key/version-2\n" + "TERMINAL=kv:https://myvault.vault.azure.net/secrets/terminal\n" + "A=one\nB=${A}\nA=two\nC=${A}" + ) + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value=root_document), + types.SimpleNamespace(value="api-key-value"), + types.SimpleNamespace(value="pinned-key-value"), + types.SimpleNamespace(value="kv:https://myvault.vault.azure.net/secrets/not-followed"), + ] + ) + secret_url = "https://myvault.vault.azure.net/secrets/bootstrap/v1" + + with ( + mock.patch.dict(os.environ, {"SOURCE_VALUE": "ambient-value"}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential) as mock_credential_cls, + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client) as mock_client_cls, + mock.patch("pyrit.setup.akv_initialization._print_msg") as mock_print_msg, + ): + await _load_env_from_akv_async(secret_url=secret_url, silent=True) + + assert os.environ["DIRECT"] == "from-bootstrap" + assert os.environ["FROM_ENV"] == "ambient-value" + assert os.environ["FROM_KV"] == "api-key-value" + assert os.environ["PINNED_KV"] == "pinned-key-value" + assert os.environ["TERMINAL"] == "kv:https://myvault.vault.azure.net/secrets/not-followed" + assert os.environ["A"] == "two" + assert os.environ["B"] == "one" + assert os.environ["C"] == "two" + + mock_credential_cls.assert_called_once_with() + _assert_mock_akv_client_created( + mock_client_cls, + vault_url="https://myvault.vault.azure.net", + credential=credential, + ) + assert client.get_secret.await_args_list == [ + mock.call("bootstrap", version="v1"), + mock.call("api-key", version=None), + mock.call("api-key", version="version-2"), + mock.call("terminal", version=None), + ] + credential.__aenter__.assert_awaited_once() + credential.__aexit__.assert_awaited_once() + client.__aenter__.assert_awaited_once() + client.__aexit__.assert_awaited_once() + mock_print_msg.assert_called_once() + + async def test_load_env_from_akv_async_preserves_process_values_without_fetching_overridden_child(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + return_value=types.SimpleNamespace( + value=("DIRECT=from-bootstrap\nFROM_KV=kv:https://myvault.vault.azure.net/secrets/api-key") + ) + ) + secret_url = "https://myvault.vault.azure.net/secrets/bootstrap" + + with ( + mock.patch.dict( + os.environ, + {"DIRECT": "from-process", "FROM_KV": "process-key"}, + clear=True, + ), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + await _load_env_from_akv_async(secret_url=secret_url, silent=True) + + assert os.environ["DIRECT"] == "from-process" + assert os.environ["FROM_KV"] == "process-key" + + client.get_secret.assert_awaited_once_with("bootstrap", version=None) + + async def test_load_env_from_akv_async_rejects_short_secret_name(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="API_KEY=kv:api-key")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="must use a full secret URL"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + @pytest.mark.parametrize( + "reference_url", + [ + "https://other-vault.vault.azure.net/secrets/api-key", + "https://other-vault.vault.azure.net/secrets/api-key/version-1", + ], + ) + async def test_load_env_from_akv_async_rejects_cross_vault_reference(self, reference_url): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=f"API_KEY=kv:{reference_url}")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="Cross-vault AKV reference"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + async def test_load_env_from_akv_async_empty_secret_raises(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=None)) + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="has no value"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/my-secret", + silent=True, + ) + + credential.__aexit__.assert_awaited_once() + client.__aexit__.assert_awaited_once() + + async def test_load_env_from_akv_async_without_entries_raises(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="# comments only\n")) + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="contains no environment entries"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/my-secret", + silent=True, + ) + + credential.__aexit__.assert_awaited_once() + client.__aexit__.assert_awaited_once() + + @pytest.mark.parametrize( + ("document", "error"), + [ + ("GOOD=resolved\n=malformed\nOTHER=resolved", "malformed entries at lines: 2"), + ("MISSING_VALUE\n", "variables without values: MISSING_VALUE"), + ], + ) + async def test_load_env_from_akv_async_rejects_non_assignments(self, document, error): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) + + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match=error), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert "GOOD" not in os.environ + assert "OTHER" not in os.environ + + async def test_load_env_from_akv_async_wraps_malformed_bootstrap(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="=malformed")) + + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(KeyVaultInitializationException, match="malformed entries") as exc_info, + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert isinstance(exc_info.value.__cause__, ValueError) + + async def test_load_env_from_akv_async_wraps_missing_child_secret(self): + credential, client = _create_mock_akv_clients() + missing_error = ResourceNotFoundError(message="Secret was not found") + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value="API_KEY=kv:https://myvault.vault.azure.net/secrets/missing"), + missing_error, + ] + ) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(KeyVaultInitializationException, match="Failed to resolve Key Vault reference") as exc_info, + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert exc_info.value.__cause__ is missing_error + + async def test_load_env_from_akv_async_allows_empty_assignment(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="EMPTY=")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert os.environ["EMPTY"] == "" + + async def test_load_env_from_akv_async_allows_empty_child_secret(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace(value="EMPTY=kv:https://myvault.vault.azure.net/secrets/empty-secret"), + types.SimpleNamespace(value=""), + ] + ) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert os.environ["EMPTY"] == "" + assert client.get_secret.await_args_list[-1] == mock.call("empty-secret", version=None) + + async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_entries(self, caplog, capsys): + credential, client = _create_mock_akv_clients() + document = "GOOD=resolved\n=malformed\nMISSING_VALUE\nOTHER=also-resolved" + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=False, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" + + output = capsys.readouterr().out + assert "WARNING: AKV environment document contains invalid entries that will be skipped" in output + assert "malformed entries at lines: 2" in output + assert "variables without values: MISSING_VALUE" in output + assert "GOOD" not in caplog.text + assert "resolved" not in caplog.text + + async def test_load_env_from_akv_async_non_strict_warns_and_skips_invalid_reference(self, caplog, capsys): + credential, client = _create_mock_akv_clients() + document = "GOOD=resolved\nBAD=kv:short-name\nOTHER=also-resolved" + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value=document)) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + ): + resolved_document = await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=False, + resolve_references_for_output=True, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["OTHER"] == "also-resolved" + assert "BAD" not in os.environ + + assert "BAD=" not in resolved_document + assert ( + "WARNING: Invalid AKV reference for environment variable 'BAD' will be skipped" in capsys.readouterr().out + ) + assert "BAD" in caplog.text + client.get_secret.assert_awaited_once_with("bootstrap", version=None) + + async def test_load_env_from_akv_async_non_strict_silent_logs_warning(self, caplog, capsys): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock(return_value=types.SimpleNamespace(value="GOOD=resolved\nMISSING_VALUE")) + + with ( + mock.patch.dict(os.environ, {}, clear=True), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + caplog.at_level("WARNING", logger="pyrit.setup.akv_initialization"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + strict=False, + silent=True, + ) + + assert capsys.readouterr().out == "" + assert "variables without values: MISSING_VALUE" in caplog.text + + async def test_load_env_from_akv_async_child_failure_keeps_loaded_bootstrap_values(self): + credential, client = _create_mock_akv_clients() + client.get_secret = mock.AsyncMock( + side_effect=[ + types.SimpleNamespace( + value=("GOOD=resolved\nBAD=kv:https://myvault.vault.azure.net/secrets/missing-value") + ), + types.SimpleNamespace(value=None), + ] + ) + + with mock.patch.dict(os.environ, {}, clear=True): + with ( + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("azure.keyvault.secrets.aio.SecretClient", return_value=client), + pytest.raises(ValueError, match="has no value"), + ): + await _load_env_from_akv_async( + secret_url="https://myvault.vault.azure.net/secrets/bootstrap", + silent=True, + ) + + assert os.environ["GOOD"] == "resolved" + assert os.environ["BAD"] == "kv:https://myvault.vault.azure.net/secrets/missing-value" diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index 99bd2c5fbc..93c812943a 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -42,6 +42,8 @@ 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.env_akv_write_env is False assert config.silent is False def test_valid_memory_db_types_snake_case(self): @@ -147,6 +149,8 @@ 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, + "env_akv_write_env": True, "silent": True, } config = ConfigurationLoader.from_dict(data) @@ -155,6 +159,8 @@ 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.env_akv_write_env is True assert config.silent is True def test_from_dict_filters_none_values(self): @@ -307,7 +313,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 +321,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 +348,8 @@ 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["env_akv_write_env"] is False assert call_kwargs["silent"] is False @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") @@ -343,13 +359,20 @@ 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, + env_akv_write_env=True, + ) 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 + assert call_kwargs["env_akv_write_env"] is True @mock.patch("pyrit.setup.configuration_loader.initialize_pyrit_async") @mock.patch("pyrit.registry.InitializerRegistry") @@ -517,6 +540,14 @@ def test_load_with_overrides_env_akv_ref_override(self, mock_default_path): assert config.env_akv_ref == ["https://vault.vault.azure.net/secrets/one"] + @mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH") + def test_load_with_overrides_env_akv_write_env_override(self, mock_default_path): + mock_default_path.exists.return_value = False + + config = ConfigurationLoader.load_with_overrides(env_akv_write_env=True) + + assert config.env_akv_write_env is True + @mock.patch("pyrit.setup.configuration_loader.DEFAULT_CONFIG_PATH") def test_load_with_overrides_converts_sequence_to_list(self, mock_default_path): """Test that Sequence inputs are converted to list for dataclass compatibility.""" diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index b919df4338..8e445c9d9c 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -3,9 +3,7 @@ import os import pathlib -import sys import tempfile -import types from unittest import mock import pytest @@ -14,7 +12,6 @@ from pyrit.common.singleton import Singleton from pyrit.registry import InitializerRegistry from pyrit.setup import IN_MEMORY, initialize_pyrit_async -from pyrit.setup.initialization import _load_env_from_akv_async, _load_environment_files, _parse_akv_secret_url class TestLoadInitializersFromScripts: @@ -122,16 +119,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.akv_initialization._load_environment_files", return_value=False) async def test_initialize_basic(self, mock_load_env, mock_set_memory): """Test basic initialization.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY) + 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.akv_initialization._load_environment_files", return_value=False) async def test_initialize_with_script(self, mock_load_env, mock_set_memory): """Test initialization with a script.""" with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f: @@ -161,59 +158,195 @@ async def initialize_async(self) -> None: finally: os.unlink(script_path) - async def test_invalid_memory_type_raises_error(self): + @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=False) + async def test_invalid_memory_type_raises_error(self, mock_load_env): """Test that invalid memory type raises ValueError.""" with pytest.raises(ValueError, match="is not a supported type"): - await initialize_pyrit_async(memory_db_type="InvalidType") # type: ignore[arg-type] + await initialize_pyrit_async(memory_db_type="InvalidType", load_defaults=False) # type: ignore[arg-type] @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files") - @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) + @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) async def test_initialize_with_env_akv_ref(self, mock_load_akv, mock_load_env, mock_set_memory): - """Test that env_akv_ref triggers AKV env loading.""" - refs = ["https://vault.vault.azure.net/secrets/test-secret"] - - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs) - - mock_load_akv.assert_awaited_once() - assert mock_load_akv.await_args.kwargs["secret_urls"] == refs - assert mock_load_akv.await_args.kwargs["silent"] is False + """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 + + with mock.patch("pyrit.setup.akv_initialization._warn_about_akv_environment_files") as mock_warn: + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_akv_ref=refs, load_defaults=False) + + assert mock_load_akv.await_args_list == [ + mock.call(secret_url=refs[0], strict=True, silent=False, resolve_references_for_output=False), + mock.call(secret_url=refs[1], strict=True, silent=False, resolve_references_for_output=False), + ] + mock_warn.assert_called_once() mock_load_env.assert_called_once() mock_set_memory.assert_called_once() @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - @mock.patch("pyrit.setup.initialization._load_environment_files") - @mock.patch("pyrit.setup.initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) + @mock.patch("pyrit.setup.akv_initialization._load_environment_files", return_value=False) + @mock.patch("pyrit.setup.akv_initialization._load_env_from_akv_async", new_callable=mock.AsyncMock) async def test_initialize_with_empty_env_akv_ref_does_not_load_akv( self, mock_load_akv, mock_load_env, mock_set_memory ): - """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.akv_initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.akv_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.akv_initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.akv_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.akv_initialization.path.CONFIGURATION_DIRECTORY_PATH", temp_path), + mock.patch("pyrit.setup.akv_initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.akv_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"] + credential = mock.MagicMock() + credential.__aenter__ = mock.AsyncMock(return_value=credential) + credential.__aexit__ = mock.AsyncMock(return_value=None) + client = mock.MagicMock() + client.__aenter__ = mock.AsyncMock(return_value=client) + client.__aexit__ = mock.AsyncMock(return_value=None) + client.get_secret = mock.AsyncMock(return_value=mock.MagicMock(value="local-secret-value")) + with tempfile.TemporaryDirectory() as temp_dir: + local_file = pathlib.Path(temp_dir) / ".env.local" + local_file.write_text( + "OVERRIDDEN=local\n" + "LOCAL_SECRET=kv:https://vault.vault.azure.net/secrets/local-secret\n" + "LOCAL_ENV=env:BOOTSTRAP_SOURCE" + ) + bootstrap_environment = { + "OVERRIDDEN": "unused-secret-value", + "BOOTSTRAP_SECRET": "bootstrap-secret-value", + "BOOTSTRAP_SOURCE": "bootstrap-value", + } + + with ( + mock.patch.dict(os.environ, {"SOURCE_VALUE": "ambient-value"}, clear=True), + mock.patch("pyrit.setup.akv_initialization._warn_about_akv_environment_files"), + mock.patch( + "pyrit.setup.akv_initialization._load_env_from_akv_async", + new_callable=mock.AsyncMock, + side_effect=lambda **_: os.environ.update(bootstrap_environment), + ), + mock.patch("azure.identity.aio.DefaultAzureCredential", return_value=credential), + mock.patch("pyrit.setup.akv_initialization._create_akv_secret_client", return_value=client), + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, + env_akv_ref=refs, + env_files=[local_file], + load_defaults=False, + ) + + assert os.environ["OVERRIDDEN"] == "local" + assert os.environ["BOOTSTRAP_SECRET"] == "bootstrap-secret-value" + assert os.environ["LOCAL_SECRET"] == "local-secret-value" + assert os.environ["LOCAL_ENV"] == "env:BOOTSTRAP_SOURCE" + + client.get_secret.assert_awaited_once_with("local-secret", version=None) + + mock_set_memory.assert_called_once() + + @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") + async def test_initialize_without_environment_file_uses_system_environment(self, mock_set_memory): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[], load_defaults=False) - assert call_order == ["akv", "env_files"] mock_set_memory.assert_called_once() @@ -237,227 +370,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.akv_initialization._load_environment_files", return_value=True) + async def test_initialize_silent_produces_no_output(self, mock_load_env, capsys): """initialize_pyrit_async with silent=True must not print anything to stdout.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=True) + 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.akv_initialization._load_environment_files", return_value=True) + async def test_initialize_not_silent_prints_migration_message(self, mock_load_env, capsys): """Without silent, the Alembic schema-check message is printed and tagged as Alembic output.""" - await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=False) + await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=False, load_defaults=False) captured = capsys.readouterr() assert "[pyrit:alembic] No new upgrade operations detected." in captured.out - - -class TestLoadEnvironmentFiles: - """Tests for _load_environment_files function and env_files parameter in initialize_pyrit_async.""" - - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") - @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_loads_default_env_files_when_none_provided(self, mock_config_path, mock_load_dotenv): - """Test that default .env and .env.local files are loaded when env_files is None.""" - # Create temporary directory and files - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - env_local_file = temp_path / ".env.local" - - # Create the files - env_file.write_text("VAR1=value1") - env_local_file.write_text("VAR2=value2") - - # Mock CONFIGURATION_DIRECTORY_PATH to point to our temp directory - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - # Call the function with None (default behavior) - _load_environment_files(env_files=None) - - # Verify both files were loaded - assert mock_load_dotenv.call_count == 2 - calls = [call[0][0] for call in mock_load_dotenv.call_args_list] - assert env_file in calls - assert env_local_file in calls - - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") - @mock.patch("pyrit.setup.initialization.path.CONFIGURATION_DIRECTORY_PATH") - async def test_only_loads_existing_default_files(self, mock_config_path, mock_load_dotenv): - """Test that only existing default files are loaded.""" - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env" - - # Only create .env, not .env.local - env_file.write_text("VAR1=value1") - - mock_config_path.__truediv__ = lambda self, other: temp_path / other - - _load_environment_files(env_files=None) - - # Verify only one file was loaded - assert mock_load_dotenv.call_count == 1 - assert mock_load_dotenv.call_args[0][0] == env_file - - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") - async def test_loads_custom_env_files_in_order(self, mock_load_dotenv): - """Test that custom env_files are loaded in the order provided.""" - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env1 = temp_path / ".env.test" - env2 = temp_path / ".env.prod" - env3 = temp_path / ".env.local" - - # Create files - env1.write_text("VAR=test") - env2.write_text("VAR=prod") - env3.write_text("VAR=local") - - # Pass custom files - _load_environment_files(env_files=[env1, env2, env3]) - - # Verify all three files were loaded in order - assert mock_load_dotenv.call_count == 3 - call_args = [call[0][0] for call in mock_load_dotenv.call_args_list] - assert call_args == [env1, env2, env3] - - async def test_raises_error_for_nonexistent_env_file(self): - """Test that ValueError is raised for non-existent env file.""" - nonexistent = pathlib.Path("/nonexistent/path/.env") - - with pytest.raises(ValueError, match="Environment file not found"): - _load_environment_files(env_files=[nonexistent]) - - @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_pyrit_with_custom_env_files(self, mock_set_memory): - """Test initialize_pyrit_async with custom env_files.""" - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - env_file = temp_path / ".env.custom" - env_file.write_text("CUSTOM_VAR=custom_value") - - # Should not raise an error - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[env_file]) - - mock_set_memory.assert_called_once() - - @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_initialize_pyrit_raises_for_nonexistent_env_file(self, mock_set_memory): - """Test that initialize_pyrit_async raises ValueError for non-existent env file.""" - nonexistent = pathlib.Path("/nonexistent/.env") - - with pytest.raises(ValueError, match="Environment file not found"): - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[nonexistent]) - - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") - @mock.patch("pyrit.setup.initialization.path.HOME_PATH") - @mock.patch("pyrit.memory.central_memory.CentralMemory.set_memory_instance") - async def test_custom_env_files_override_default_behavior(self, mock_set_memory, mock_home_path, mock_load_dotenv): - """Test that passing custom env_files prevents loading default files.""" - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = pathlib.Path(temp_dir) - - # Create default files - default_env = temp_path / ".env" - default_env_local = temp_path / ".env.local" - default_env.write_text("DEFAULT=value") - default_env_local.write_text("DEFAULT_LOCAL=value") - - # Create custom file - custom_env = temp_path / ".env.custom" - custom_env.write_text("CUSTOM=value") - - mock_home_path.__truediv__ = lambda self, other: temp_path / other - - # Pass custom env_files - should NOT load defaults - await initialize_pyrit_async(memory_db_type=IN_MEMORY, env_files=[custom_env]) - - # Verify only custom file was loaded, not the default ones - assert mock_load_dotenv.call_count == 1 - assert mock_load_dotenv.call_args[0][0] == custom_env - - -class TestAkvEnvironmentLoading: - """Tests for AKV URL parsing and env loading helpers.""" - - def test_parse_akv_secret_url_with_version(self): - url = "https://myvault.vault.azure.net/secrets/my-secret/abc123" - - vault_url, secret_name, secret_version = _parse_akv_secret_url(url) - - assert vault_url == "https://myvault.vault.azure.net" - assert secret_name == "my-secret" - assert secret_version == "abc123" - - def test_parse_akv_secret_url_without_version(self): - url = "https://myvault.vault.azure.net/secrets/my-secret" - - vault_url, secret_name, secret_version = _parse_akv_secret_url(url) - - assert vault_url == "https://myvault.vault.azure.net" - assert secret_name == "my-secret" - assert secret_version is None - - def test_parse_akv_secret_url_invalid_raises(self): - with pytest.raises(ValueError, match="Invalid AKV secret URL"): - _parse_akv_secret_url("https://myvault.vault.azure.net/not-secrets/my-secret") - - @mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") - async def test_load_env_from_akv_async_empty_urls_noop(self, mock_load_dotenv): - await _load_env_from_akv_async(secret_urls=[]) - mock_load_dotenv.assert_not_called() - - async def test_load_env_from_akv_async_loads_secret_content(self): - class FakeCredential: - pass - - client_calls: list[tuple[str, object, object]] = [] - - class FakeSecretClient: - def __init__(self, *, vault_url, credential): - client_calls.append(("init", vault_url, credential)) - - async def get_secret(self, name, version=None): - client_calls.append(("get_secret", name, version)) - return types.SimpleNamespace(value="AKV_VAR=from_secret\n") - - azure_module = types.ModuleType("azure") - identity_module = types.ModuleType("azure.identity") - identity_aio_module = types.ModuleType("azure.identity.aio") - keyvault_module = types.ModuleType("azure.keyvault") - keyvault_secrets_module = types.ModuleType("azure.keyvault.secrets") - keyvault_secrets_aio_module = types.ModuleType("azure.keyvault.secrets.aio") - - identity_aio_module.DefaultAzureCredential = FakeCredential - keyvault_secrets_aio_module.SecretClient = FakeSecretClient - - with ( - mock.patch.dict( - sys.modules, - { - "azure": azure_module, - "azure.identity": identity_module, - "azure.identity.aio": identity_aio_module, - "azure.keyvault": keyvault_module, - "azure.keyvault.secrets": keyvault_secrets_module, - "azure.keyvault.secrets.aio": keyvault_secrets_aio_module, - }, - ), - mock.patch("pyrit.setup.initialization.dotenv.load_dotenv") as mock_load_dotenv, - mock.patch("pyrit.setup.initialization._print_msg") as mock_print_msg, - ): - await _load_env_from_akv_async( - secret_urls=["https://myvault.vault.azure.net/secrets/my-secret/v1"], - silent=True, - ) - - assert client_calls[0][0] == "init" - assert client_calls[0][1] == "https://myvault.vault.azure.net" - assert isinstance(client_calls[0][2], FakeCredential) - assert client_calls[1] == ("get_secret", "my-secret", "v1") - - stream = mock_load_dotenv.call_args.kwargs["stream"] - assert stream.getvalue() == "AKV_VAR=from_secret\n" - assert mock_load_dotenv.call_args.kwargs["override"] is True - assert mock_print_msg.call_count == 2 diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index 52141904e1..06c7d4c4d1 100644 --- a/tests/unit/setup/test_targets_initializer.py +++ b/tests/unit/setup/test_targets_initializer.py @@ -104,9 +104,9 @@ async def test_registers_multiple_targets(self): os.environ["PLATFORM_OPENAI_CHAT_MODEL"] = "gpt-4o" # Set up openai_image_platform (uses ENDPOINT2/KEY2/MODEL2) - os.environ["OPENAI_IMAGE_ENDPOINT2"] = "https://api.openai.com/v1" - os.environ["OPENAI_IMAGE_API_KEY2"] = "test_image_key" - os.environ["OPENAI_IMAGE_MODEL2"] = "dall-e-3" + os.environ["AZURE_OPENAI_IMAGE_ENDPOINT2"] = "https://api.openai.com/v1" + os.environ["AZURE_OPENAI_IMAGE_API_KEY2"] = "test_image_key" + os.environ["AZURE_OPENAI_IMAGE_MODEL2"] = "dall-e-3" init = TargetInitializer() await init.initialize_async()