[FEAT] Resolve Key Vault-Backend Environment References

.gitignore

@@ -168,6 +168,7 @@ cython_debug/
 
 # PyRIT secrets file
 .env
+.env_akv
 .pyrit_cache/
 
 # Cache for generating docs

.pyrit_conf_example

@@ -4,7 +4,7 @@
 # or specify a custom path when loading via --config-file.
 #
 # For documentation on configuration options, see:
-# https://github.com/microsoft/PyRIT/blob/main/doc/setup/configuration.md
+# https://github.com/microsoft/PyRIT/blob/main/doc/getting_started/pyrit_conf.md
 
 # Memory Database Type
 # --------------------
@@ -79,33 +79,23 @@ operation: op_trash_panda
 #   - /path/to/my_custom_initializer.py
 #   - ./local_initializer.py
 
-# Environment Files
-# -----------------
-# List of .env file paths to load during initialization.
-# Later files override values from earlier files.
-#
-# Behavior:
-# - Omit this field (or set to null): Load default .env and .env.local from ~/.pyrit/ if they exist
-# - Set to []: Explicitly load NO environment files
-# - Set to list of paths: Load only the specified files
-#
-# Example:
-# env_files:
-#   - /path/to/.env
-#   - /path/to/.env.local
-
-# Azure Key Vault Environment References
-# ---------------------------------------
-# List of AKV secret URLs to load during initialization.
-# Each secret's value must be the full contents of a .env file.
-# Loaded after env_files, so AKV secrets take precedence.
-# Authentication uses DefaultAzureCredential (managed identity, Azure CLI, etc.).
-#
-# Requires: pip install azure-keyvault-secrets
-#
-# Example:
+# Environment Configuration
+# -------------------------
+# Azure Key Vault is the canonical source for shared and deployed configuration.
+# See doc/getting_started/pyrit_conf.md for loading order, references, and migration guidance.
+# The list may contain at most one bootstrap secret URL.
 # env_akv_ref:
 #   - https://my-vault.vault.azure.net/secrets/my-pyrit-env
+# env_akv_strict: true
+
+# Auto-discovered ~/.pyrit/.env remains supported but emits a security warning.
+# Prefer env_akv_ref for shared or deployed secrets.
+# Use ~/.pyrit/.env.local for quick local plaintext patches or when Azure is unavailable.
+# Process values remain authoritative; AKV and ordinary env_files fill gaps in load order.
+# Only a file named .env.local overrides existing values.
+# Explicit env_files remain supported regardless of name or location and may contain full kv: URLs.
+# env_files:
+#   - /path/to/.env.local
 
 # Max Concurrent Scenario Runs
 # ----------------------------

build_scripts/export_akv_environment.py

@@ -0,0 +1,246 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""Export a resolved Azure Key Vault bootstrap document to ``~/.pyrit/.env_akv``."""
+
+import argparse
+import contextlib
+import logging
+import os
+import pathlib
+import tempfile
+from collections.abc import Mapping, Sequence
+from io import StringIO
+from typing import TYPE_CHECKING
+
+import dotenv
+from dotenv.parser import parse_stream
+from dotenv.variables import parse_variables
+
+from pyrit.setup.environment_loading import (
+    _parse_akv_reference,
+    _parse_akv_secret_url,
+    _validate_dotenv_document,
+)
+
+if TYPE_CHECKING:
+    from azure.core.credentials import TokenCredential
+    from azure.keyvault.secrets import SecretClient
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_OUTPUT_FILE = pathlib.Path.home() / ".pyrit" / ".env_akv"
+
+
+def _create_client(*, vault_url: str, credential: "TokenCredential") -> "SecretClient":
+    """Create a Key Vault client with explicit retry settings."""
+    from azure.core.pipeline.policies import RetryPolicy
+    from azure.keyvault.secrets import SecretClient
+
+    return SecretClient(
+        vault_url=vault_url,
+        credential=credential,
+        retry_policy=RetryPolicy(
+            retry_total=3,
+            retry_connect=3,
+            retry_read=3,
+            retry_status=3,
+            retry_backoff_factor=0.8,
+        ),
+    )
+
+
+def _client_for(*, vault_url: str, credential: "TokenCredential", clients: dict[str, "SecretClient"]) -> "SecretClient":
+    client = clients.get(vault_url)
+    if client is None:
+        client = _create_client(vault_url=vault_url, credential=credential)
+        clients[vault_url] = client
+    return client
+
+
+def _fetch_document(
+    *,
+    secret_url: str,
+    credential: "TokenCredential",
+    clients: dict[str, "SecretClient"],
+    strict: bool,
+    silent: bool,
+) -> tuple[str, str]:
+    vault_url, name, version = _parse_akv_secret_url(secret_url)
+    secret = _client_for(vault_url=vault_url, credential=credential, clients=clients).get_secret(name, version=version)
+    if not secret.value:
+        raise ValueError(f"AKV environment secret has no value: {secret_url}")
+    content = _validate_dotenv_document(secret.value, strict=strict, silent=silent)
+    if not dotenv.dotenv_values(stream=StringIO(content), interpolate=False):
+        raise ValueError(f"AKV environment secret contains no assignments: {secret_url}")
+    return content, vault_url
+
+
+def _resolve_interpolation(*, value: str, environment: Mapping[str, str | None]) -> str:
+    return "".join(atom.resolve(environment) for atom in parse_variables(value))
+
+
+def _build_candidates(document: tuple[str, str]) -> dict[str, tuple[str, str]]:
+    content, vault_url = document
+    values: dict[str, str] = {}
+    for binding in parse_stream(StringIO(content)):
+        if binding.key is not None and binding.value is not None:
+            values[binding.key] = _resolve_interpolation(value=binding.value, environment=values)
+    return {name: (value, vault_url) for name, value in values.items()}
+
+
+def _serialize(value: str) -> str:
+    escaped = value.replace("\\", "\\\\").replace("'", "\\'").replace("${", "${:-$}{")
+    return f"'{escaped}'"
+
+
+def _render(
+    *,
+    document: tuple[str, str],
+    credential: "TokenCredential",
+    clients: dict[str, "SecretClient"],
+    strict: bool,
+    silent: bool,
+) -> str:
+    resolved: dict[str, str] = {}
+    for name, (value, source_vault_url) in _build_candidates(document).items():
+        try:
+            reference = _parse_akv_reference(
+                value=value,
+                variable_name=name,
+                expected_vault_url=source_vault_url,
+            )
+        except ValueError as error:
+            if strict:
+                raise
+            message = f"Invalid AKV reference for '{name}' will be skipped: {error}"
+            if not silent:
+                print(f"WARNING: {message}")
+            logger.warning(message)
+            continue
+        if reference is None:
+            resolved[name] = value
+            continue
+        vault_url, secret_name, version = reference
+        secret = _client_for(vault_url=vault_url, credential=credential, clients=clients).get_secret(
+            secret_name, version=version
+        )
+        if secret.value is None:
+            raise ValueError(f"AKV secret '{secret_name}' referenced by '{name}' has no value")
+        resolved[name] = secret.value
+
+    return "".join(f"{name}={_serialize(value)}\n" for name, value in resolved.items())
+
+
+def _ensure_output_available(output_file: pathlib.Path) -> pathlib.Path:
+    output_file = output_file.expanduser()
+    if output_file.is_symlink():
+        raise ValueError(f"Output path is a symbolic link: {output_file}")
+    if output_file.exists():
+        raise ValueError(f"Output already exists: {output_file}. Rename or remove it before exporting")
+    return output_file
+
+
+def _write_output(*, output_file: pathlib.Path, document: str) -> pathlib.Path:
+    output_file = _ensure_output_available(output_file)
+    output_file.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+    descriptor: int | None = None
+    temporary: pathlib.Path | None = None
+    try:
+        descriptor, name = tempfile.mkstemp(prefix=f"{output_file.name}.", suffix=".tmp", dir=output_file.parent)
+        temporary = pathlib.Path(name)
+        file_chmod = getattr(os, "fchmod", None)
+        if file_chmod is not None:
+            file_chmod(descriptor, 0o600)
+        else:
+            os.chmod(temporary, 0o600)
+        stream = os.fdopen(descriptor, "w", encoding="utf-8", newline="")
+        descriptor = None
+        with stream:
+            stream.write(document)
+        try:
+            os.link(temporary, output_file)
+        except FileExistsError as error:
+            raise ValueError(f"Output already exists: {output_file}. Rename or remove it before exporting") from error
+    finally:
+        if descriptor is not None:
+            os.close(descriptor)
+        if temporary is not None:
+            with contextlib.suppress(FileNotFoundError):
+                temporary.unlink()
+    return output_file
+
+
+def export_akv_environment(
+    *,
+    secret_urls: Sequence[str],
+    output_file: pathlib.Path = DEFAULT_OUTPUT_FILE,
+    strict: bool = True,
+    silent: bool = False,
+    credential: "TokenCredential | None" = None,
+) -> pathlib.Path:
+    """Fetch, resolve, and securely export AKV-only configuration.
+
+    A caller-provided credential remains caller-owned and is not closed.
+    """
+    if not secret_urls:
+        raise ValueError("At least one secret URL is required")
+    if len(secret_urls) > 1:
+        raise ValueError("Only one Azure Key Vault bootstrap secret URL is supported")
+    output_file = _ensure_output_available(output_file)
+    from azure.identity import DefaultAzureCredential
+
+    owned_credential = None
+    if credential is None:
+        owned_credential = DefaultAzureCredential()
+        active_credential = owned_credential
+    else:
+        active_credential = credential
+    clients: dict[str, SecretClient] = {}
+    try:
+        document = _fetch_document(
+            secret_url=secret_urls[0],
+            credential=active_credential,
+            clients=clients,
+            strict=strict,
+            silent=silent,
+        )
+        document = _render(
+            document=document,
+            credential=active_credential,
+            clients=clients,
+            strict=strict,
+            silent=silent,
+        )
+        output = _write_output(output_file=output_file, document=document)
+    finally:
+        for client in clients.values():
+            client.close()
+        if owned_credential is not None:
+            owned_credential.close()
+    if not silent:
+        print(f"Exported resolved AKV environment to {output}")
+    return output
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--secret-url", required=True)
+    parser.add_argument("--output", type=pathlib.Path, default=DEFAULT_OUTPUT_FILE)
+    parser.add_argument("--non-strict", action="store_true")
+    parser.add_argument("--silent", action="store_true")
+    args = parser.parse_args()
+    try:
+        export_akv_environment(
+            secret_urls=[args.secret_url],
+            output_file=args.output,
+            strict=not args.non_strict,
+            silent=args.silent,
+        )
+    except Exception as error:
+        parser.exit(1, f"Export failed: {error}\n")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())