From d5ccd0a1bf982208cd8577562369fef1d1c6f251 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 26 Aug 2026 16:29:43 -0700 Subject: [PATCH 1/3] FEAT: Adding remote auth support for CLI Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6f936dc9-8e9b-4295-9c3b-aefc3d7554fc --- doc/getting_started/pyrit_conf.md | 30 +++ doc/scanner/0_scanner.md | 25 ++ doc/scanner/2_pyrit_shell.md | 9 + infra/README.md | 27 +- infra/deploy_instance.py | 11 +- pyrit/backend/routes/auth.py | 19 +- pyrit/cli/_auth.py | 303 +++++++++++++++++++++++ pyrit/cli/_config_reader.py | 16 +- pyrit/cli/api_client.py | 96 ++++++- pyrit/cli/pyrit_scan.py | 30 ++- pyrit/cli/pyrit_shell.py | 78 +++++- tests/unit/backend/test_auth_route.py | 60 +++++ tests/unit/cli/test_api_client.py | 84 +++++++ tests/unit/cli/test_auth.py | 167 +++++++++++++ tests/unit/cli/test_config_reader.py | 25 ++ tests/unit/cli/test_pyrit_scan.py | 23 ++ tests/unit/cli/test_pyrit_shell.py | 78 ++++++ tests/unit/infra/test_deploy_instance.py | 24 ++ 18 files changed, 1070 insertions(+), 35 deletions(-) create mode 100644 pyrit/cli/_auth.py create mode 100644 tests/unit/backend/test_auth_route.py create mode 100644 tests/unit/cli/test_auth.py create mode 100644 tests/unit/infra/test_deploy_instance.py diff --git a/doc/getting_started/pyrit_conf.md b/doc/getting_started/pyrit_conf.md index 64235420ea..10032e1590 100644 --- a/doc/getting_started/pyrit_conf.md +++ b/doc/getting_started/pyrit_conf.md @@ -264,6 +264,7 @@ Client settings for connecting to or launching a PyRIT backend. | --- | --- | --- | | `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` | +| `auth_mode` | Backend authentication mode: `auto`, `azure_cli`, `device_code`, or `none` | `auto` | `startup_timeout` must be a finite number greater than zero. The `--startup-timeout` CLI option overrides the configured value for an individual scanner invocation. @@ -273,8 +274,37 @@ Set `server: null` to reset all server settings, including values inherited from server: url: http://localhost:8000 startup_timeout: 120 + auth_mode: auto ``` +In `auto` mode, the CLI reads the backend's public `/api/auth/config` endpoint. It sends no +token when authentication is disabled. For an authenticated backend, it uses Entra device-code +login with the exact Microsoft Graph `User.Read` scope. The encrypted persistent token cache +normally prevents a new prompt on each run. A non-interactive process fails instead of waiting +for a prompt. + +Use an explicit mode when needed: + +```yaml +server: + url: https://copyrit.example.com/ + auth_mode: azure_cli +``` + +`azure_cli` is an explicit compatibility mode. It can send a Microsoft Graph token with +permissions beyond `User.Read` because the Azure CLI application controls the token's granted +permissions. Prefer `auto` or `device_code`. If you accept this behavior, sign in to the +backend's tenant before using `azure_cli`: + +```bash +az login --tenant +pyrit_scan --config-file ./.pyrit_conf list-scenarios +``` + +Use `device_code` to require the same exact-scope interactive flow as `auto`. Use `none` only +when you intentionally need to suppress authentication discovery. Access tokens are not stored +in `.pyrit_conf`. + ## Configuration Precedence PyRIT uses a 3-layer configuration precedence model. **Later layers override earlier ones:** diff --git a/doc/scanner/0_scanner.md b/doc/scanner/0_scanner.md index 35f4ba4b5c..aaf0f976c3 100644 --- a/doc/scanner/0_scanner.md +++ b/doc/scanner/0_scanner.md @@ -26,6 +26,31 @@ PyRIT provides two command-line interfaces: pyrit_scan run foundry.red_team_agent --target openai_chat --initializers target --techniques base64 ``` +### Connecting to CoPyRIT + +Point a local configuration file at the remote backend: + +```yaml +server: + url: https://copyrit.example.com/ + auth_mode: auto +``` + +Then use the file without changing the default configuration in `~/.pyrit`: + +```bash +pyrit_scan --config-file ./.pyrit_conf list-scenarios +``` + +The CLI reads the server's public authentication configuration. Automatic mode uses an +interactive Entra device code with the exact Microsoft Graph `User.Read` scope and an encrypted +persistent token cache. Use `--auth-mode device_code` to require this flow or `--auth-mode none` +to disable authentication discovery. + +`--auth-mode azure_cli` is an explicit compatibility mode. The Azure CLI application can issue +a Graph token with permissions beyond `User.Read`, and the CLI sends that token to the backend. +Prefer automatic device-code authentication. + ## Built-in Scenarios PyRIT ships with scenarios organized into the following families: diff --git a/doc/scanner/2_pyrit_shell.md b/doc/scanner/2_pyrit_shell.md index c74416f0e5..2710e66b99 100644 --- a/doc/scanner/2_pyrit_shell.md +++ b/doc/scanner/2_pyrit_shell.md @@ -23,8 +23,17 @@ pyrit_shell --config-file ./.pyrit_conf # Set default log level pyrit_shell --log-level DEBUG + +# Connect to an authenticated remote backend +pyrit_shell --config-file ./.pyrit_conf --auth-mode auto ``` +Authentication defaults to `auto`. The shell uses exact-scope device-code login and stores the +result in an encrypted persistent token cache. The configuration file can set +`server.auth_mode` to `device_code` or `none` when automatic selection is not appropriate. +`azure_cli` remains an explicit compatibility mode, but its Graph token can contain permissions +beyond `User.Read`. + ## Available Commands Once starting the shell, you will see the list of commands you have access to. Some of them are shown below: diff --git a/infra/README.md b/infra/README.md index f4af84d426..7b425ed951 100644 --- a/infra/README.md +++ b/infra/README.md @@ -7,7 +7,7 @@ managed identity, security response headers, and no embedded secrets. ## Architecture ``` -Users ──→ MSAL PKCE auth ──→ Container App +Users ──→ MSAL PKCE or CLI device-code auth ──→ Container App ↓ Graph-backed authentication ↓ @@ -57,10 +57,11 @@ Production is opt-in via `deployToProd: true`. ## Security - **Authentication**: [MSAL](https://learn.microsoft.com/en-us/entra/msal/) - [PKCE](https://oauth.net/2/pkce/) on the frontend (`@azure/msal-browser`) + + [PKCE](https://oauth.net/2/pkce/) on the frontend (`@azure/msal-browser`) and + public-client device-code authentication for the PyRIT CLI + Microsoft Graph-backed middleware on the backend. The frontend sends a delegated - Graph token, and the backend authenticates it through Graph `/me`. PKCE (public - client) requires no client secrets or certificates. + Graph token, and the backend authenticates it through Graph `/me`. These public-client + flows require no client secrets or certificates. - **Authorization**: Entra group check via `allowedGroupObjectIds` param. Requires delegated Graph `User.Read`; the backend calls `/me/checkMemberGroups` and compares the returned transitive memberships with the configured group IDs. Each security @@ -165,7 +166,7 @@ az account show --query tenantId -o tsv > --spa-redirect-uris "https://$FQDN" > ``` -**Configure delegated Microsoft Graph access** (required): +**Configure delegated Microsoft Graph access and public-client login** (required): In Azure Portal → App registrations → your app → **API permissions**: @@ -186,9 +187,13 @@ Or via CLI: APP_OBJ_ID=$(az ad app show --id $APP_ID --query id -o tsv) az rest --method PATCH \ --url "https://graph.microsoft.com/v1.0/applications/$APP_OBJ_ID" \ - --body '{"requiredResourceAccess":[{"resourceAppId":"00000003-0000-0000-c000-000000000000","resourceAccess":[{"id":"e1fe6dd8-ba31-4d61-89e7-88639da4683d","type":"Scope"}]}]}' + --body '{"isFallbackPublicClient":true,"requiredResourceAccess":[{"resourceAppId":"00000003-0000-0000-c000-000000000000","resourceAccess":[{"id":"e1fe6dd8-ba31-4d61-89e7-88639da4683d","type":"Scope"}]}]}' ``` +`isFallbackPublicClient` enables device-code login for `pyrit_scan` and `pyrit_shell`. In the +Azure Portal, the equivalent setting is **Authentication → Advanced settings → Allow public +client flows → Yes**. + ### 3. Entra security groups (required for group-based authorization) Create one or more security groups for authorized users. Multiple groups can be @@ -350,14 +355,22 @@ az deployment group create \ ## Post-Deployment -1. **Set SPA redirect URI** on the app registration (requires the FQDN from deploy output): +1. **Configure browser and CLI public-client authentication** on the app registration: ```bash FQDN=$(az deployment group show -g -n main \ --query properties.outputs.appFqdn.value -o tsv) az ad app update --id \ --spa-redirect-uris "https://$FQDN" + + APP_OBJ_ID=$(az ad app show --id --query id -o tsv) + az rest --method PATCH \ + --url "https://graph.microsoft.com/v1.0/applications/$APP_OBJ_ID" \ + --body '{"isFallbackPublicClient":true}' ``` + For an existing deployment, run only the `APP_OBJ_ID` and `az rest` commands once. Do not + rerun `infra/deploy_instance.py`; its resource creation steps are not idempotent. + 2. **Grant managed identity RBAC** (required — the Bicep template does **not** create role assignments; the app will fail to start without AcrPull): ```bash diff --git a/infra/deploy_instance.py b/infra/deploy_instance.py index e038d518ab..2220d6285b 100644 --- a/infra/deploy_instance.py +++ b/infra/deploy_instance.py @@ -15,7 +15,7 @@ 7. Managed identity + RBAC role assignments (AcrPull, Storage Blob Data Contributor) 7b. AOAI RBAC (optional — Cognitive Services OpenAI User on specified resources) 8. Bicep deployment (Container App, networking, logging) - 9. Post-deploy: SPA redirect URI + 9. Post-deploy: SPA redirect URI + public-client device-code flow Usage: python infra/deploy_instance.py \\ @@ -853,15 +853,18 @@ def post_deploy( fqdn: str, ) -> None: """ - Run post-deployment steps: SPA redirect URI. + Run post-deployment steps for browser and CLI authentication. Args: app_object_id (str): The Entra app registration object ID (for Graph API). fqdn (str): The deployed app FQDN. """ - # Set SPA redirect URI via Graph REST API (more portable than --spa-redirect-uris flag) + # Keep browser PKCE and device-code clients on the same public app registration. logger.info("Setting SPA redirect URI: https://%s", fqdn) - spa_body = {"spa": {"redirectUris": [f"https://{fqdn}"]}} + spa_body = { + "spa": {"redirectUris": [f"https://{fqdn}"]}, + "isFallbackPublicClient": True, + } run_az( args=[ "rest", diff --git a/pyrit/backend/routes/auth.py b/pyrit/backend/routes/auth.py index 9510618985..107d957ca3 100644 --- a/pyrit/backend/routes/auth.py +++ b/pyrit/backend/routes/auth.py @@ -13,10 +13,11 @@ from fastapi import APIRouter router = APIRouter() +_GRAPH_SCOPES = ["https://graph.microsoft.com/User.Read"] @router.get("/auth/config") -async def get_auth_config_async() -> dict[str, str]: +async def get_auth_config_async() -> dict[str, str | bool | list[str]]: """ Return Entra ID configuration for the frontend MSAL client. @@ -25,10 +26,18 @@ async def get_auth_config_async() -> dict[str, str]: are included so the frontend can show appropriate error messages. Returns: - dict: Auth configuration with clientId, tenantId, allowedGroupIds. + dict: Auth configuration with enabled state, clientId, tenantId, + allowedGroupIds, and delegated Microsoft Graph scopes. """ + client_id = os.getenv("ENTRA_CLIENT_ID", "").strip() + tenant_id = os.getenv("ENTRA_TENANT_ID", "").strip() + allowed_group_ids = os.getenv("ENTRA_ALLOWED_GROUP_IDS", "").strip() + enabled = bool(client_id and tenant_id and allowed_group_ids) + return { - "clientId": os.getenv("ENTRA_CLIENT_ID", ""), - "tenantId": os.getenv("ENTRA_TENANT_ID", ""), - "allowedGroupIds": os.getenv("ENTRA_ALLOWED_GROUP_IDS", ""), + "enabled": enabled, + "clientId": client_id, + "tenantId": tenant_id, + "allowedGroupIds": allowed_group_ids, + "scopes": list(_GRAPH_SCOPES) if enabled else [], } diff --git a/pyrit/cli/_auth.py b/pyrit/cli/_auth.py new file mode 100644 index 0000000000..0cdf67e281 --- /dev/null +++ b/pyrit/cli/_auth.py @@ -0,0 +1,303 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Authentication helpers for the thin PyRIT REST clients.""" + +from __future__ import annotations + +import asyncio +import sys +import time +from dataclasses import dataclass +from typing import Any, Literal, Protocol + +AuthMode = Literal["auto", "azure_cli", "device_code", "none"] +AUTH_MODES: tuple[AuthMode, ...] = ("auto", "azure_cli", "device_code", "none") +_GRAPH_USER_READ_SCOPE = "https://graph.microsoft.com/User.Read" +_GRAPH_DEFAULT_SCOPE = "https://graph.microsoft.com/.default" +_TOKEN_REFRESH_BUFFER_SECONDS = 300 +_AZURE_CLI_WARNING = ( + "Warning: azure_cli mode sends the Azure CLI application's Microsoft Graph token to the backend. " + "That token can contain permissions beyond User.Read. Prefer auto or device_code." +) + + +class CliAuthenticationError(Exception): + """Raised when the CLI cannot authenticate to a protected backend.""" + + +@dataclass(frozen=True) +class BackendAuthConfig: + """Public Entra configuration advertised by the PyRIT backend.""" + + enabled: bool + tenant_id: str + client_id: str + scopes: tuple[str, ...] + + @classmethod + def from_payload(cls, payload: Any) -> BackendAuthConfig: + """ + Validate an ``/api/auth/config`` response. + + Args: + payload: Parsed response payload. + + Returns: + Validated backend authentication configuration. + + Raises: + CliAuthenticationError: If the server returns an unsupported contract. + """ + if not isinstance(payload, dict): + raise CliAuthenticationError( + "The server returned an unsupported authentication contract. " + "Upgrade the PyRIT backend so /api/auth/config includes 'enabled' and 'scopes'." + ) + + tenant_id = payload.get("tenantId", "") + client_id = payload.get("clientId", "") + if not isinstance(tenant_id, str) or not isinstance(client_id, str): + raise CliAuthenticationError("The server returned invalid Entra tenant or client configuration.") + if "enabled" not in payload and not tenant_id.strip() and not client_id.strip(): + return cls(enabled=False, tenant_id="", client_id="", scopes=()) + if not isinstance(payload.get("enabled"), bool): + raise CliAuthenticationError( + "The server returned an unsupported authentication contract. " + "Upgrade the PyRIT backend so /api/auth/config includes 'enabled' and 'scopes'." + ) + + enabled = payload["enabled"] + raw_scopes = payload.get("scopes", []) + if not isinstance(raw_scopes, list) or not all( + isinstance(scope, str) and scope.strip() for scope in raw_scopes + ): + raise CliAuthenticationError("The server returned invalid delegated authentication scopes.") + + scopes = tuple(scope.strip() for scope in raw_scopes) + if enabled and scopes != (_GRAPH_USER_READ_SCOPE,): + raise CliAuthenticationError( + "The server requested an unsupported authentication scope. " + f"Only {_GRAPH_USER_READ_SCOPE} is allowed." + ) + tenant_id = tenant_id.strip() + client_id = client_id.strip() + if enabled and (not tenant_id or not client_id or not scopes): + raise CliAuthenticationError( + "The server reports authentication enabled but its Entra configuration is incomplete." + ) + + return cls( + enabled=enabled, + tenant_id=tenant_id, + client_id=client_id, + scopes=scopes, + ) + + +class TokenProvider(Protocol): + """Supply and refresh access tokens for backend requests.""" + + async def get_token_async(self) -> str: + """Return a current bearer token.""" + + async def close_async(self) -> None: + """Release credential resources.""" + + +class _AzureIdentityTokenProvider: + """Adapt an asynchronous Azure Identity credential to the CLI token protocol.""" + + def __init__(self, *, credential: Any, auth_config: BackendAuthConfig, mode: AuthMode) -> None: + self._credential = credential + self._auth_config = auth_config + self._mode = mode + self._access_token: Any = None + + async def get_token_async(self) -> str: + """ + Acquire a delegated Microsoft Graph access token. + + Returns: + Access token text. + + Raises: + CliAuthenticationError: If Azure Identity cannot authenticate the user. + """ + from azure.core.exceptions import ClientAuthenticationError + from azure.identity import CredentialUnavailableError + + cached_token = self._get_current_token() + if cached_token is not None: + return cached_token + + try: + access_token = await self._credential.get_token(_GRAPH_DEFAULT_SCOPE) + except (ClientAuthenticationError, CredentialUnavailableError) as exc: + if self._mode == "azure_cli": + hint = f"Run 'az login --tenant {self._auth_config.tenant_id}' and try again." + else: + hint = "Confirm that device-code authentication is enabled for the CoPyRIT Entra application." + raise CliAuthenticationError(f"Entra authentication failed. {hint}") from exc + + token = getattr(access_token, "token", "") + if not isinstance(token, str) or not token: + raise CliAuthenticationError("Entra authentication returned an empty access token.") + self._access_token = access_token + return token + + def _get_current_token(self) -> str | None: + """Return a cached token that is valid beyond the refresh buffer.""" + if self._access_token is None: + return None + expires_on = getattr(self._access_token, "expires_on", 0) + token = getattr(self._access_token, "token", "") + if ( + isinstance(expires_on, int | float) + and expires_on > time.time() + _TOKEN_REFRESH_BUFFER_SECONDS + and isinstance(token, str) + and token + ): + return token + return None + + async def close_async(self) -> None: + """Close the underlying Azure Identity credential.""" + await self._credential.close() + + +class _DeviceCodeTokenProvider: + """Adapt the synchronous device-code credential without blocking the event loop.""" + + def __init__(self, *, credential: Any, auth_config: BackendAuthConfig) -> None: + self._credential = credential + self._auth_config = auth_config + self._access_token: Any = None + + async def get_token_async(self) -> str: + """ + Acquire a delegated Microsoft Graph token through device-code login. + + Returns: + Access token text. + + Raises: + CliAuthenticationError: If Entra rejects device-code authentication. + """ + from azure.core.exceptions import ClientAuthenticationError + from azure.identity import CredentialUnavailableError + + cached_token = self._get_current_token() + if cached_token is not None: + return cached_token + + try: + access_token = await asyncio.to_thread(self._credential.get_token, *self._auth_config.scopes) + except (ClientAuthenticationError, CredentialUnavailableError) as exc: + raise CliAuthenticationError( + "Entra authentication failed. Confirm that device-code authentication " + "is enabled for the CoPyRIT Entra application." + ) from exc + + token = getattr(access_token, "token", "") + if not isinstance(token, str) or not token: + raise CliAuthenticationError("Entra authentication returned an empty access token.") + self._access_token = access_token + return token + + def _get_current_token(self) -> str | None: + """Return a cached token that is valid beyond the refresh buffer.""" + if self._access_token is None: + return None + expires_on = getattr(self._access_token, "expires_on", 0) + token = getattr(self._access_token, "token", "") + if ( + isinstance(expires_on, int | float) + and expires_on > time.time() + _TOKEN_REFRESH_BUFFER_SECONDS + and isinstance(token, str) + and token + ): + return token + return None + + async def close_async(self) -> None: + """Close the underlying synchronous Azure Identity credential.""" + await asyncio.to_thread(self._credential.close) + + +def _is_interactive() -> bool: + """Return whether authentication may safely prompt this process.""" + return bool(sys.stdin.isatty() and sys.stderr.isatty()) + + +def _create_azure_cli_provider(*, auth_config: BackendAuthConfig) -> TokenProvider: + from azure.identity.aio import AzureCliCredential + + credential = AzureCliCredential(tenant_id=auth_config.tenant_id) + return _AzureIdentityTokenProvider(credential=credential, auth_config=auth_config, mode="azure_cli") + + +def _create_device_code_provider(*, auth_config: BackendAuthConfig) -> TokenProvider: + from azure.identity import DeviceCodeCredential, TokenCachePersistenceOptions + + cache_options = TokenCachePersistenceOptions(name=f"pyrit-copyrit-{auth_config.client_id}") + credential = DeviceCodeCredential( + tenant_id=auth_config.tenant_id, + client_id=auth_config.client_id, + cache_persistence_options=cache_options, + ) + return _DeviceCodeTokenProvider(credential=credential, auth_config=auth_config) + + +async def _verify_provider_async(*, provider: TokenProvider) -> TokenProvider: + """ + Acquire an initial token so selection fails before the first API operation. + + Returns: + The verified provider. + + Raises: + CliAuthenticationError: If the provider cannot acquire a token. + """ + try: + await provider.get_token_async() + except CliAuthenticationError: + await provider.close_async() + raise + return provider + + +async def create_token_provider_async( + *, + auth_config: BackendAuthConfig, + auth_mode: AuthMode, + interactive: bool | None = None, +) -> TokenProvider | None: + """ + Select and verify a token provider for the backend. + + Args: + auth_config: Authentication requirements discovered from the backend. + auth_mode: Requested credential selection behavior. + interactive: Optional terminal-interactivity override for tests. + + Returns: + A verified token provider, or ``None`` when authentication is disabled. + + Raises: + CliAuthenticationError: If the requested authentication flow cannot run. + """ + if auth_mode == "none" or not auth_config.enabled: + return None + + if auth_mode == "azure_cli": + print(_AZURE_CLI_WARNING, file=sys.stderr) + return await _verify_provider_async(provider=_create_azure_cli_provider(auth_config=auth_config)) + + can_prompt = _is_interactive() if interactive is None else interactive + if not can_prompt: + raise CliAuthenticationError( + "Device-code authentication requires an interactive terminal. " + "Use azure_cli only when you accept sending the Azure CLI Graph token to the backend." + ) + return await _verify_provider_async(provider=_create_device_code_provider(auth_config=auth_config)) diff --git a/pyrit/cli/_config_reader.py b/pyrit/cli/_config_reader.py index e55217bbca..d97b2a183f 100644 --- a/pyrit/cli/_config_reader.py +++ b/pyrit/cli/_config_reader.py @@ -15,12 +15,15 @@ from pathlib import Path from typing import Any +from pyrit.cli._auth import AUTH_MODES, AuthMode + # Mirror the default path from pyrit.common.path without importing it. _DEFAULT_CONFIG_DIR = Path.home() / ".pyrit" _DEFAULT_CONFIG_FILE = _DEFAULT_CONFIG_DIR / ".pyrit_conf" DEFAULT_SERVER_URL = "http://localhost:8000" DEFAULT_SERVER_STARTUP_TIMEOUT = 120.0 +DEFAULT_AUTH_MODE: AuthMode = "auto" @dataclass(frozen=True) @@ -29,6 +32,7 @@ class ServerSettings: url: str | None = None startup_timeout: float = DEFAULT_SERVER_STARTUP_TIMEOUT + auth_mode: AuthMode = DEFAULT_AUTH_MODE class ConfigError(Exception): @@ -106,7 +110,7 @@ def read_server_settings(*, config_file: Path | None = None) -> ServerSettings: config_file: Optional explicit config path. Returns: - ServerSettings: The resolved URL and startup timeout. + ServerSettings: The resolved URL, startup timeout, and authentication mode. Raises: ConfigError: If a config file exists but is malformed. @@ -198,4 +202,12 @@ def _merge_server_settings(*, settings: ServerSettings, data: dict[str, Any], pa raise ConfigError(f"Config file {path}: 'server.startup_timeout' must be a finite number greater than 0.") startup_timeout = float(raw_timeout) - return ServerSettings(url=url, startup_timeout=startup_timeout) + auth_mode = settings.auth_mode + if "auth_mode" in server_block: + raw_auth_mode = server_block["auth_mode"] + if not isinstance(raw_auth_mode, str) or raw_auth_mode not in AUTH_MODES: + supported_modes = ", ".join(AUTH_MODES) + raise ConfigError(f"Config file {path}: 'server.auth_mode' must be one of: {supported_modes}.") + auth_mode = raw_auth_mode + + return ServerSettings(url=url, startup_timeout=startup_timeout, auth_mode=auth_mode) diff --git a/pyrit/cli/api_client.py b/pyrit/cli/api_client.py index 927cf1d0ff..642403749f 100644 --- a/pyrit/cli/api_client.py +++ b/pyrit/cli/api_client.py @@ -14,8 +14,10 @@ import logging from typing import TYPE_CHECKING, Any +from urllib.parse import urlparse if TYPE_CHECKING: + from pyrit.cli._auth import AuthMode, TokenProvider from pyrit.models import ScenarioResult from pyrit.models.catalog import ( RegisteredInitializer, @@ -44,7 +46,14 @@ class PyRITApiClient: scenarios = await client.list_scenarios_async() """ - def __init__(self, *, base_url: str, request_timeout: float | None = None) -> None: + def __init__( + self, + *, + base_url: str, + request_timeout: float | None = None, + auth_mode: AuthMode = "none", + interactive: bool | None = None, + ) -> None: """ Initialize the API client. @@ -55,10 +64,15 @@ def __init__(self, *, base_url: str, request_timeout: float | None = None) -> No the live scenario-run endpoint always uses ``read=None`` regardless of this value, because the server may legitimately take many seconds to respond while a scenario is executing. Defaults to ``60.0``. + auth_mode: Authentication behavior for protected remote backends. + interactive: Optional terminal-interactivity override. """ self._base_url = base_url.rstrip("/") self._request_timeout = request_timeout if request_timeout is not None else 60.0 + self._auth_mode = auth_mode + self._interactive = interactive self._client: Any = None # httpx.AsyncClient (typed Any to avoid top-level import) + self._token_provider: TokenProvider | None = None async def __aenter__(self) -> PyRITApiClient: """ @@ -66,10 +80,28 @@ async def __aenter__(self) -> PyRITApiClient: Returns: PyRITApiClient: ``self``, with the HTTP client opened. + + Raises: + CliAuthenticationError: If authentication discovery or login fails. + httpx.HTTPError: If the authentication discovery request fails. """ import httpx - self._client = httpx.AsyncClient(base_url=self._base_url, timeout=self._request_timeout) + client_kwargs: dict[str, Any] = { + "base_url": self._base_url, + "timeout": self._request_timeout, + } + if self._auth_mode != "none": + client_kwargs["event_hooks"] = {"request": [self._add_authorization_header_async]} + self._client = httpx.AsyncClient(**client_kwargs) + if self._auth_mode != "none": + from pyrit.cli._auth import CliAuthenticationError + + try: + await self._configure_authentication_async() + except (CliAuthenticationError, httpx.HTTPError): + await self.close_async() + raise return self async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: @@ -359,14 +391,68 @@ async def get_conversation_messages_async( async def close_async(self) -> None: """Close the underlying HTTP client.""" - if self._client is not None: - await self._client.aclose() - self._client = None + client = self._client + token_provider = self._token_provider + self._client = None + self._token_provider = None + + try: + if client is not None: + await client.aclose() + finally: + if token_provider is not None: + await token_provider.close_async() # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ + async def _configure_authentication_async(self) -> None: + """ + Discover backend authentication requirements and select a credential. + + Raises: + CliAuthenticationError: If the server contract or selected credential is invalid. + httpx.HTTPError: If the discovery request fails. + """ + from pyrit.cli._auth import BackendAuthConfig, CliAuthenticationError, create_token_provider_async + + client = self._get_client() + response = await client.get("/api/auth/config") + if response.status_code == 404: + return + self._raise_for_status(response) + try: + payload = response.json() + except ValueError as exc: + raise CliAuthenticationError("The server returned invalid JSON from /api/auth/config.") from exc + + auth_config = BackendAuthConfig.from_payload(payload) + if auth_config.enabled and not self._uses_secure_auth_transport(): + raise CliAuthenticationError( + "Refusing to send an Entra access token over a non-HTTPS connection. " + "Use HTTPS for remote backends." + ) + self._token_provider = await create_token_provider_async( + auth_config=auth_config, + auth_mode=self._auth_mode, + interactive=self._interactive, + ) + + def _uses_secure_auth_transport(self) -> bool: + """Return whether the server URL protects bearer tokens in transit.""" + parsed = urlparse(self._base_url) + if parsed.scheme == "https": + return True + return parsed.scheme == "http" and parsed.hostname in {"localhost", "127.0.0.1", "::1"} + + async def _add_authorization_header_async(self, request: Any) -> None: + """Attach a current bearer token to protected backend requests.""" + if self._token_provider is None or request.url.path in {"/api/auth/config", "/api/health"}: + return + token = await self._token_provider.get_token_async() + request.headers["Authorization"] = f"Bearer {token}" + def _get_client(self) -> Any: """ Return the ``httpx.AsyncClient``, raising if not opened. diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index e14e0e9bce..bfcf6bb296 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -18,10 +18,11 @@ import sys from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter from pathlib import Path -from typing import TYPE_CHECKING, Any, get_args, get_origin +from typing import TYPE_CHECKING, Any, cast, get_args, get_origin import aiofiles +from pyrit.cli._auth import AUTH_MODES, AuthMode from pyrit.cli._cli_args import ( ARG_HELP, _parse_initializer_arg, @@ -160,6 +161,10 @@ def _positive_finite_float(value: str) -> float: "(catalog/results/cancel/etc). Defaults to 60. Polling a live " "scenario run always waits indefinitely regardless of this value." ) +_AUTH_MODE_HELP = ( + "Backend authentication mode (default: server.auth_mode or auto). " + "Auto uses exact-scope device-code authentication in an interactive terminal." +) def _add_common_options(*, parser: ArgumentParser, suppress_defaults: bool) -> None: @@ -181,6 +186,7 @@ def _add_common_options(*, parser: ArgumentParser, suppress_defaults: bool) -> N group = parser.add_argument_group("global options") group.add_argument("--server-url", type=str, default=default, help=_SERVER_URL_HELP) group.add_argument("--config-file", type=Path, default=default, help=_CONFIG_FILE_HELP) + group.add_argument("--auth-mode", choices=AUTH_MODES, default=default, help=_AUTH_MODE_HELP) group.add_argument("--log-level", type=validate_log_level_argparse, default=log_default, help=_LOG_LEVEL_HELP) @@ -700,6 +706,27 @@ def _resolve_configured_server_url(*, parsed_args: Namespace) -> str: return server_url +def _resolve_auth_mode(*, parsed_args: Namespace) -> AuthMode: + """ + Resolve the authentication mode from CLI and layered configuration. + + Returns: + The selected authentication mode. + + Raises: + ValueError: If an unsupported mode is supplied programmatically. + """ + from pyrit.cli._config_reader import read_server_settings + + configured_mode = read_server_settings(config_file=parsed_args.config_file).auth_mode + raw_auth_mode = getattr(parsed_args, "auth_mode", None) + if raw_auth_mode is None: + return configured_mode + if raw_auth_mode not in AUTH_MODES: + raise ValueError(f"Unsupported authentication mode: {raw_auth_mode}") + return cast("AuthMode", raw_auth_mode) + + async def _handle_stop_server_async(*, parsed_args: Namespace) -> int: """ Handle ``stop-server``: probe, then terminate the listening process. @@ -1121,6 +1148,7 @@ async def _run_async(*, parsed_args: Namespace) -> int: async with PyRITApiClient( base_url=base_url_result, request_timeout=getattr(parsed_args, "request_timeout", None), + auth_mode=_resolve_auth_mode(parsed_args=parsed_args), ) as client: return await _dispatch_with_client_async(client=client, parsed_args=parsed_args) except ServerNotAvailableError as exc: diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py index be4d634cef..80ea4419b8 100644 --- a/pyrit/cli/pyrit_shell.py +++ b/pyrit/cli/pyrit_shell.py @@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any, TypeVar from pyrit.cli import _banner as banner +from pyrit.cli._auth import AUTH_MODES, AuthMode if TYPE_CHECKING: from collections.abc import Coroutine @@ -92,6 +93,7 @@ def __init__( server_url: str | None = None, config_file: Path | None = None, start_server: bool = False, + auth_mode: AuthMode | None = None, ) -> None: """ Initialize the PyRIT shell. @@ -101,12 +103,14 @@ def __init__( server_url: Optional explicit server URL. config_file: Optional config file path. start_server: If True, auto-start a local backend. + auth_mode: Optional backend authentication mode override. """ super().__init__() self._no_animation = no_animation self._server_url = server_url self._config_file = config_file self._start_server = start_server + self._auth_mode = auth_mode self._api_client: Any = None # PyRITApiClient (lazy) self._base_url: str | None = None self._launcher: Any = None # ServerLauncher (lazy) @@ -165,6 +169,57 @@ def _resolve_base_url(self) -> str: return self._server_url return read_server_url(config_file=self._config_file) or DEFAULT_SERVER_URL + def _resolve_auth_mode(self) -> AuthMode: + """ + Determine the backend authentication mode. + + Returns: + The selected authentication mode. + """ + from pyrit.cli._config_reader import read_server_settings + + return self._auth_mode or read_server_settings(config_file=self._config_file).auth_mode + + def _open_client(self, *, base_url: str) -> bool: + """ + Open an API client while keeping connection failures inside the REPL. + + Returns: + ``True`` when the client is ready, otherwise ``False``. + """ + import httpx + + from pyrit.cli._auth import CliAuthenticationError + from pyrit.cli._config_reader import ConfigError + from pyrit.cli._output import print_error_with_hint + from pyrit.cli.api_client import PyRITApiClient + + self._base_url = base_url + try: + client = PyRITApiClient(base_url=base_url, auth_mode=self._resolve_auth_mode()) + self._run_async(client.__aenter__(), timeout=None) + except CliAuthenticationError as exc: + self._api_client = None + print_error_with_hint( + message=str(exc), + hint="Use --auth-mode to select auto, device_code, azure_cli, or none.", + ) + return False + except ConfigError as exc: + self._api_client = None + print(f"Error: {exc}") + return False + except httpx.HTTPError as exc: + self._api_client = None + print_error_with_hint( + message=f"Could not initialize the client for {base_url}: {exc}", + hint="Check the server's /api/auth/config endpoint and your network connection.", + ) + return False + + self._api_client = client + return True + def _ensure_client(self) -> bool: """ Ensure the API client is connected. @@ -212,11 +267,8 @@ def _ensure_client(self) -> bool: ) return False - from pyrit.cli.api_client import PyRITApiClient - - self._base_url = base_url - self._api_client = PyRITApiClient(base_url=base_url) - self._run_async(self._api_client.__aenter__()) + if not self._open_client(base_url=base_url): + return False self._start_server = False # only auto-start once return True @@ -667,7 +719,6 @@ def do_start_server(self, arg: str) -> None: print(f"Error: start-server does not accept arguments, got: {arg.strip()}") return from pyrit.cli._server_launcher import ServerLauncher - from pyrit.cli.api_client import PyRITApiClient base_url = self._resolve_base_url() @@ -675,9 +726,7 @@ def do_start_server(self, arg: str) -> None: if self._run_async(ServerLauncher.probe_health_async(base_url=base_url)): print(f"Server already running at {base_url}") if self._api_client is None: - self._base_url = base_url - self._api_client = PyRITApiClient(base_url=base_url) - self._run_async(self._api_client.__aenter__()) + self._open_client(base_url=base_url) return self._launcher = ServerLauncher() @@ -690,8 +739,8 @@ def do_start_server(self, arg: str) -> None: # Create new client for the started server if self._api_client is not None: self._run_async(self._api_client.close_async()) - self._api_client = PyRITApiClient(base_url=new_url) - self._run_async(self._api_client.__aenter__()) + self._api_client = None + self._open_client(base_url=new_url) except RuntimeError as exc: print(f"Error: {exc}") @@ -829,6 +878,12 @@ def main() -> int: help=ARG_HELP["config_file"], ) + parser.add_argument( + "--auth-mode", + choices=AUTH_MODES, + help="Backend authentication mode (default: server.auth_mode or auto)", + ) + parser.add_argument( "--log-level", type=str, @@ -870,6 +925,7 @@ def main() -> int: server_url=args.server_url, config_file=args.config_file, start_server=args.start_server, + auth_mode=args.auth_mode, ) shell.cmdloop(intro=intro) return 0 diff --git a/tests/unit/backend/test_auth_route.py b/tests/unit/backend/test_auth_route.py new file mode 100644 index 0000000000..b599205003 --- /dev/null +++ b/tests/unit/backend/test_auth_route.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for the public authentication configuration route.""" + +from unittest.mock import patch + +from pyrit.backend.routes.auth import get_auth_config_async + + +async def test_get_auth_config_returns_enabled_graph_contract() -> None: + environment = { + "ENTRA_TENANT_ID": " tenant-id ", + "ENTRA_CLIENT_ID": " client-id ", + "ENTRA_ALLOWED_GROUP_IDS": " group-1,group-2 ", + } + + with patch.dict("os.environ", environment, clear=False): + result = await get_auth_config_async() + + assert result == { + "enabled": True, + "clientId": "client-id", + "tenantId": "tenant-id", + "allowedGroupIds": "group-1,group-2", + "scopes": ["https://graph.microsoft.com/User.Read"], + } + + +async def test_get_auth_config_returns_disabled_contract_when_configuration_is_absent() -> None: + environment = { + "ENTRA_TENANT_ID": "", + "ENTRA_CLIENT_ID": "", + "ENTRA_ALLOWED_GROUP_IDS": "", + } + + with patch.dict("os.environ", environment, clear=False): + result = await get_auth_config_async() + + assert result == { + "enabled": False, + "clientId": "", + "tenantId": "", + "allowedGroupIds": "", + "scopes": [], + } + + +async def test_get_auth_config_does_not_enable_incomplete_configuration() -> None: + environment = { + "ENTRA_TENANT_ID": "tenant-id", + "ENTRA_CLIENT_ID": "", + "ENTRA_ALLOWED_GROUP_IDS": "group-1", + } + + with patch.dict("os.environ", environment, clear=False): + result = await get_auth_config_async() + + assert result["enabled"] is False + assert result["scopes"] == [] diff --git a/tests/unit/cli/test_api_client.py b/tests/unit/cli/test_api_client.py index b4e6109d0a..28988ba6d9 100644 --- a/tests/unit/cli/test_api_client.py +++ b/tests/unit/cli/test_api_client.py @@ -11,6 +11,7 @@ import httpx import pytest +from pyrit.cli._auth import CliAuthenticationError from pyrit.cli.api_client import PyRITApiClient, ServerNotAvailableError from pyrit.models import ScenarioRunState, TargetCapabilities from pyrit.models.catalog import ( @@ -156,6 +157,89 @@ async def test_close_async_is_noop_when_already_closed(): await c.close_async() # Should not raise. +async def test_context_manager_discovers_auth_and_attaches_bearer_token(mock_httpx_client): + c = PyRITApiClient(base_url="https://copyrit.example.com", auth_mode="auto", interactive=False) + fake_async_client_cls = MagicMock(return_value=mock_httpx_client) + mock_httpx_client.get.return_value = _make_response( + json_data={ + "enabled": True, + "tenantId": "tenant-id", + "clientId": "client-id", + "scopes": ["https://graph.microsoft.com/User.Read"], + } + ) + provider = MagicMock() + provider.get_token_async = AsyncMock(return_value="access-token") + provider.close_async = AsyncMock() + + with ( + patch("httpx.AsyncClient", fake_async_client_cls), + patch( + "pyrit.cli._auth.create_token_provider_async", + new_callable=AsyncMock, + return_value=provider, + ) as create_provider, + ): + async with c: + request_hook = fake_async_client_cls.call_args.kwargs["event_hooks"]["request"][0] + request = httpx.Request("GET", "https://copyrit.example.com/api/targets") + await request_hook(request) + assert request.headers["Authorization"] == "Bearer access-token" + + mock_httpx_client.get.assert_awaited_once_with("/api/auth/config") + create_provider.assert_awaited_once() + provider.close_async.assert_awaited_once() + + +async def test_context_manager_leaves_public_requests_unauthenticated(mock_httpx_client): + c = PyRITApiClient(base_url="https://copyrit.example.com", auth_mode="auto") + fake_async_client_cls = MagicMock(return_value=mock_httpx_client) + mock_httpx_client.get.return_value = _make_response( + json_data={ + "enabled": False, + "tenantId": "", + "clientId": "", + "scopes": [], + } + ) + + with patch("httpx.AsyncClient", fake_async_client_cls): + async with c: + request_hook = fake_async_client_cls.call_args.kwargs["event_hooks"]["request"][0] + request = httpx.Request("GET", "https://copyrit.example.com/api/health") + await request_hook(request) + assert "Authorization" not in request.headers + + +async def test_context_manager_accepts_legacy_backend_without_auth_endpoint(mock_httpx_client): + c = PyRITApiClient(base_url="http://legacy.example.com", auth_mode="auto") + fake_async_client_cls = MagicMock(return_value=mock_httpx_client) + mock_httpx_client.get.return_value = _make_response(status_code=404) + + with patch("httpx.AsyncClient", fake_async_client_cls): + async with c: + assert c._token_provider is None + + +async def test_context_manager_rejects_authentication_over_remote_http(mock_httpx_client): + c = PyRITApiClient(base_url="http://copyrit.example.com", auth_mode="auto") + fake_async_client_cls = MagicMock(return_value=mock_httpx_client) + mock_httpx_client.get.return_value = _make_response( + json_data={ + "enabled": True, + "tenantId": "tenant-id", + "clientId": "client-id", + "scopes": ["https://graph.microsoft.com/User.Read"], + } + ) + + with patch("httpx.AsyncClient", fake_async_client_cls): + with pytest.raises(CliAuthenticationError, match="non-HTTPS"): + await c.__aenter__() + + mock_httpx_client.aclose.assert_awaited_once() + + def test_get_client_raises_when_not_opened(): c = PyRITApiClient(base_url="http://localhost:8000") with pytest.raises(ServerNotAvailableError, match="not connected"): diff --git a/tests/unit/cli/test_auth.py b/tests/unit/cli/test_auth.py new file mode 100644 index 0000000000..cfb0b54a40 --- /dev/null +++ b/tests/unit/cli/test_auth.py @@ -0,0 +1,167 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for thin-client authentication helpers.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyrit.cli import _auth +from pyrit.cli._auth import BackendAuthConfig, CliAuthenticationError, create_token_provider_async + +_AUTH_CONFIG = BackendAuthConfig( + enabled=True, + tenant_id="tenant-id", + client_id="client-id", + scopes=("https://graph.microsoft.com/User.Read",), +) + + +def test_backend_auth_config_parses_enabled_contract() -> None: + result = BackendAuthConfig.from_payload( + { + "enabled": True, + "tenantId": " tenant-id ", + "clientId": " client-id ", + "scopes": [" https://graph.microsoft.com/User.Read "], + } + ) + + assert result == _AUTH_CONFIG + + +def test_backend_auth_config_accepts_legacy_disabled_contract() -> None: + result = BackendAuthConfig.from_payload( + { + "tenantId": "", + "clientId": "", + "allowedGroupIds": "", + } + ) + + assert result == BackendAuthConfig(enabled=False, tenant_id="", client_id="", scopes=()) + + +def test_backend_auth_config_rejects_non_graph_scope() -> None: + with pytest.raises(CliAuthenticationError, match="unsupported authentication scope"): + BackendAuthConfig.from_payload( + { + "enabled": True, + "tenantId": "tenant-id", + "clientId": "client-id", + "scopes": ["https://management.azure.com/.default"], + } + ) + + +@pytest.mark.parametrize( + "payload", + [ + {"enabled": "true"}, + {"enabled": True, "tenantId": "", "clientId": "client-id", "scopes": ["scope"]}, + {"enabled": True, "tenantId": "tenant-id", "clientId": "client-id", "scopes": []}, + {"enabled": True, "tenantId": "tenant-id", "clientId": "client-id", "scopes": "scope"}, + ], +) +def test_backend_auth_config_rejects_invalid_contract(payload: object) -> None: + with pytest.raises(CliAuthenticationError): + BackendAuthConfig.from_payload(payload) + + +async def test_create_token_provider_returns_none_when_server_auth_is_disabled() -> None: + auth_config = BackendAuthConfig(enabled=False, tenant_id="", client_id="", scopes=()) + + result = await create_token_provider_async(auth_config=auth_config, auth_mode="auto") + + assert result is None + + +async def test_create_token_provider_auto_uses_device_code() -> None: + provider = MagicMock() + provider.get_token_async = AsyncMock(return_value="token") + provider.close_async = AsyncMock() + + with patch.object(_auth, "_create_device_code_provider", return_value=provider) as create_device_code: + result = await create_token_provider_async( + auth_config=_AUTH_CONFIG, + auth_mode="auto", + interactive=True, + ) + + assert result is provider + create_device_code.assert_called_once_with(auth_config=_AUTH_CONFIG) + provider.get_token_async.assert_awaited_once() + provider.close_async.assert_not_awaited() + + +async def test_create_token_provider_auto_fails_fast_when_non_interactive() -> None: + with pytest.raises(CliAuthenticationError, match="interactive terminal"): + await create_token_provider_async( + auth_config=_AUTH_CONFIG, + auth_mode="auto", + interactive=False, + ) + + +async def test_create_token_provider_device_code_requires_interactive_terminal() -> None: + with pytest.raises(CliAuthenticationError, match="interactive terminal"): + await create_token_provider_async( + auth_config=_AUTH_CONFIG, + auth_mode="device_code", + interactive=False, + ) + + +async def test_create_token_provider_azure_cli_warns(capsys) -> None: + provider = MagicMock() + provider.get_token_async = AsyncMock(return_value="token") + provider.close_async = AsyncMock() + + with patch.object(_auth, "_create_azure_cli_provider", return_value=provider): + result = await create_token_provider_async( + auth_config=_AUTH_CONFIG, + auth_mode="azure_cli", + interactive=False, + ) + + assert result is provider + assert "permissions beyond User.Read" in capsys.readouterr().err + + +def test_create_device_code_provider_uses_persistent_cache() -> None: + cache_options = MagicMock() + credential = MagicMock() + + with ( + patch("azure.identity.TokenCachePersistenceOptions", return_value=cache_options) as cache_type, + patch("azure.identity.DeviceCodeCredential", return_value=credential) as credential_type, + ): + provider = _auth._create_device_code_provider(auth_config=_AUTH_CONFIG) + + cache_type.assert_called_once_with(name="pyrit-copyrit-client-id") + credential_type.assert_called_once_with( + tenant_id="tenant-id", + client_id="client-id", + cache_persistence_options=cache_options, + ) + assert isinstance(provider, _auth._DeviceCodeTokenProvider) + + +async def test_azure_cli_provider_caches_token_until_refresh_window() -> None: + credential = MagicMock() + credential.get_token = AsyncMock() + credential.close = AsyncMock() + credential.get_token.return_value = MagicMock(token="access-token", expires_on=2_000_000_000) + provider = _auth._AzureIdentityTokenProvider( + credential=credential, + auth_config=_AUTH_CONFIG, + mode="azure_cli", + ) + + with patch("pyrit.cli._auth.time.time", return_value=1_000_000_000): + first = await provider.get_token_async() + second = await provider.get_token_async() + + assert first == second == "access-token" + credential.get_token.assert_awaited_once_with("https://graph.microsoft.com/.default") diff --git a/tests/unit/cli/test_config_reader.py b/tests/unit/cli/test_config_reader.py index 8bdeebbd8c..816ad6783c 100644 --- a/tests/unit/cli/test_config_reader.py +++ b/tests/unit/cli/test_config_reader.py @@ -11,6 +11,7 @@ from pyrit.cli import _config_reader from pyrit.cli._config_reader import ( + DEFAULT_AUTH_MODE, DEFAULT_SERVER_STARTUP_TIMEOUT, DEFAULT_SERVER_URL, ConfigError, @@ -24,6 +25,7 @@ def test_default_server_url_constant(): assert DEFAULT_SERVER_URL == "http://localhost:8000" assert DEFAULT_SERVER_STARTUP_TIMEOUT == 120.0 + assert DEFAULT_AUTH_MODE == "auto" def test_read_server_url_returns_none_when_no_files(tmp_path): @@ -133,6 +135,19 @@ def test_read_server_settings_overlay_overrides_startup_timeout(tmp_path): ) +def test_read_server_settings_overlay_overrides_auth_mode(tmp_path): + default = tmp_path / "default.yaml" + default.write_text("server:\n url: http://default:8000\n auth_mode: auto\n", encoding="utf-8") + overlay = tmp_path / "overlay.yaml" + overlay.write_text("server:\n auth_mode: azure_cli\n", encoding="utf-8") + + with patch.object(_config_reader, "_DEFAULT_CONFIG_FILE", default): + assert read_server_settings(config_file=overlay) == ServerSettings( + url="http://default:8000", + auth_mode="azure_cli", + ) + + def test_read_server_settings_overlay_timeout_preserves_default_url(tmp_path): default = tmp_path / "default.yaml" default.write_text("server:\n url: http://default:8000\n startup_timeout: 180\n", encoding="utf-8") @@ -192,6 +207,16 @@ def test_read_server_settings_rejects_invalid_startup_timeout(tmp_path, startup_ read_server_settings(config_file=bad) +@pytest.mark.parametrize("auth_mode", ["interactive", "", 123, True]) +def test_read_server_settings_rejects_invalid_auth_mode(tmp_path, auth_mode): + bad = tmp_path / "bad.yaml" + bad.write_text(f"server:\n auth_mode: {auth_mode}\n", encoding="utf-8") + + with patch.object(_config_reader, "_DEFAULT_CONFIG_FILE", tmp_path / "missing.yaml"): + with pytest.raises(ConfigError, match="server.auth_mode"): + read_server_settings(config_file=bad) + + def test_validate_client_config_rejects_removed_scenario_block(tmp_path): cfg = tmp_path / "conf.yaml" cfg.write_text("scenario:\n name: test\n", encoding="utf-8") diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index 911d629cd4..81ee21f62f 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -187,6 +187,14 @@ def test_list_with_server_url(self): args = pyrit_scan.parse_args(["list-scenarios", "--server-url", "http://remote:9000"]) assert args.server_url == "http://remote:9000" + def test_list_with_auth_mode(self): + args = pyrit_scan.parse_args(["list-scenarios", "--auth-mode", "device_code"]) + assert args.auth_mode == "device_code" + + def test_list_rejects_invalid_auth_mode(self): + with pytest.raises(SystemExit): + pyrit_scan.parse_args(["list-scenarios", "--auth-mode", "interactive"]) + def test_global_flag_before_verb(self): args = pyrit_scan.parse_args(["--server-url", "http://remote:9000", "list-scenarios"]) assert args.command == "list-scenarios" @@ -434,6 +442,21 @@ def test_main_list_scenarios(self, mock_client_class, mock_probe): assert result == 0 mock_client.list_scenarios_async.assert_awaited_once() + @patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new_callable=AsyncMock, + return_value=True, + ) + @patch("pyrit.cli.api_client.PyRITApiClient") + def test_main_passes_auth_mode_to_client(self, mock_client_class, mock_probe): + mock_client = _mock_api_client() + mock_client_class.return_value = mock_client + + result = pyrit_scan.main(["list-scenarios", "--auth-mode", "azure_cli"]) + + assert result == 0 + assert mock_client_class.call_args.kwargs["auth_mode"] == "azure_cli" + @patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new_callable=AsyncMock, diff --git a/tests/unit/cli/test_pyrit_shell.py b/tests/unit/cli/test_pyrit_shell.py index 4168df61d6..beaf32d6d3 100644 --- a/tests/unit/cli/test_pyrit_shell.py +++ b/tests/unit/cli/test_pyrit_shell.py @@ -311,6 +311,21 @@ def test_main_parses_server_url(self): mock_shell_class.assert_called_once() assert mock_shell_class.call_args.kwargs["server_url"] == "http://remote:9000" + def test_main_parses_auth_mode(self): + with ( + patch("pyrit.cli._banner.play_animation", return_value=""), + patch("pyrit.cli.pyrit_shell.PyRITShell") as mock_shell_class, + patch( + "sys.argv", + ["pyrit_shell", "--auth-mode", "device_code", "--no-animation"], + ), + ): + mock_shell_class.return_value = MagicMock() + + pyrit_shell.main() + + assert mock_shell_class.call_args.kwargs["auth_mode"] == "device_code" + def test_main_keyboard_interrupt(self, capsys): with ( patch("pyrit.cli._banner.play_animation", return_value=""), @@ -408,6 +423,48 @@ def test_start_server_launches_when_not_running(self): assert s._api_client is mock_client assert s._start_server is False # only auto-start once + def test_authentication_failure_returns_false(self, capsys): + from pyrit.cli._auth import CliAuthenticationError + + s = pyrit_shell.PyRITShell(no_animation=True, server_url="https://copyrit.example.com") + with ( + patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new_callable=AsyncMock, + return_value=True, + ), + patch("pyrit.cli.api_client.PyRITApiClient") as mock_client_class, + ): + mock_client = MagicMock() + mock_client.__aenter__ = AsyncMock(side_effect=CliAuthenticationError("login failed")) + mock_client_class.return_value = mock_client + + assert s._ensure_client() is False + + assert s._api_client is None + assert "login failed" in capsys.readouterr().out + + def test_auth_discovery_http_failure_returns_false(self, capsys): + import httpx + + s = pyrit_shell.PyRITShell(no_animation=True, server_url="https://copyrit.example.com") + with ( + patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new_callable=AsyncMock, + return_value=True, + ), + patch("pyrit.cli.api_client.PyRITApiClient") as mock_client_class, + ): + mock_client = MagicMock() + mock_client.__aenter__ = AsyncMock(side_effect=httpx.ConnectError("discovery failed")) + mock_client_class.return_value = mock_client + + assert s._ensure_client() is False + + assert s._api_client is None + assert "discovery failed" in capsys.readouterr().out + def test_start_server_failure_returns_false(self, capsys): s = pyrit_shell.PyRITShell(no_animation=True, start_server=True) with ( @@ -767,6 +824,27 @@ def test_start_server_launch_success(self): s.do_start_server("") assert s._base_url == "http://localhost:8000" + def test_start_server_authentication_failure_stays_in_repl(self, capsys): + from pyrit.cli._auth import CliAuthenticationError + + s = pyrit_shell.PyRITShell(no_animation=True) + with ( + patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new_callable=AsyncMock, + return_value=True, + ), + patch("pyrit.cli.api_client.PyRITApiClient") as mock_client_class, + ): + mock_client = MagicMock() + mock_client.__aenter__ = AsyncMock(side_effect=CliAuthenticationError("login failed")) + mock_client_class.return_value = mock_client + + s.do_start_server("") + + assert s._api_client is None + assert "login failed" in capsys.readouterr().out + def test_start_server_launch_replaces_existing_client(self): s = pyrit_shell.PyRITShell(no_animation=True) existing = AsyncMock() diff --git a/tests/unit/infra/test_deploy_instance.py b/tests/unit/infra/test_deploy_instance.py new file mode 100644 index 0000000000..48d89b011b --- /dev/null +++ b/tests/unit/infra/test_deploy_instance.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for CoPyRIT deployment authentication configuration.""" + +import json +from unittest.mock import patch + +from infra import deploy_instance + + +def test_post_deploy_enables_spa_and_device_code_authentication() -> None: + with patch.object(deploy_instance, "run_az") as run_az: + deploy_instance.post_deploy( + app_object_id="app-object-id", + fqdn="copyrit.example.com", + ) + + args = run_az.call_args.kwargs["args"] + body = json.loads(args[args.index("--body") + 1]) + assert body == { + "spa": {"redirectUris": ["https://copyrit.example.com"]}, + "isFallbackPublicClient": True, + } From 48720fd5bdca5279304729600bf64c3da18ccff0 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 26 Aug 2026 16:44:52 -0700 Subject: [PATCH 2/3] STYLE: Apply Ruff formatting Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6f936dc9-8e9b-4295-9c3b-aefc3d7554fc --- pyrit/cli/_auth.py | 3 +-- pyrit/cli/api_client.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/pyrit/cli/_auth.py b/pyrit/cli/_auth.py index 0cdf67e281..151a58909e 100644 --- a/pyrit/cli/_auth.py +++ b/pyrit/cli/_auth.py @@ -77,8 +77,7 @@ def from_payload(cls, payload: Any) -> BackendAuthConfig: scopes = tuple(scope.strip() for scope in raw_scopes) if enabled and scopes != (_GRAPH_USER_READ_SCOPE,): raise CliAuthenticationError( - "The server requested an unsupported authentication scope. " - f"Only {_GRAPH_USER_READ_SCOPE} is allowed." + f"The server requested an unsupported authentication scope. Only {_GRAPH_USER_READ_SCOPE} is allowed." ) tenant_id = tenant_id.strip() client_id = client_id.strip() diff --git a/pyrit/cli/api_client.py b/pyrit/cli/api_client.py index 642403749f..574d10619b 100644 --- a/pyrit/cli/api_client.py +++ b/pyrit/cli/api_client.py @@ -430,8 +430,7 @@ async def _configure_authentication_async(self) -> None: auth_config = BackendAuthConfig.from_payload(payload) if auth_config.enabled and not self._uses_secure_auth_transport(): raise CliAuthenticationError( - "Refusing to send an Entra access token over a non-HTTPS connection. " - "Use HTTPS for remote backends." + "Refusing to send an Entra access token over a non-HTTPS connection. Use HTTPS for remote backends." ) self._token_provider = await create_token_provider_async( auth_config=auth_config, From a27be414b68f2f7b87d709beca050e7c3ef72f67 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Thu, 27 Aug 2026 10:05:46 -0700 Subject: [PATCH 3/3] FIX: Harden CLI device code authentication Persist account metadata for cross-process token reuse, keep prompts on stderr, and surface encrypted cache failures. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6f936dc9-8e9b-4295-9c3b-aefc3d7554fc --- pyrit/cli/_auth.py | 230 ++++++++++++++----- tests/unit/cli/test_api_client.py | 14 ++ tests/unit/cli/test_auth.py | 339 ++++++++++++++++++++++++++++- tests/unit/cli/test_pyrit_scan.py | 10 + tests/unit/cli/test_pyrit_shell.py | 10 + 5 files changed, 547 insertions(+), 56 deletions(-) diff --git a/pyrit/cli/_auth.py b/pyrit/cli/_auth.py index 151a58909e..e8e6fbd0e0 100644 --- a/pyrit/cli/_auth.py +++ b/pyrit/cli/_auth.py @@ -6,16 +6,22 @@ from __future__ import annotations import asyncio +import hashlib import sys import time from dataclasses import dataclass -from typing import Any, Literal, Protocol +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, Protocol + +if TYPE_CHECKING: + from datetime import datetime AuthMode = Literal["auto", "azure_cli", "device_code", "none"] AUTH_MODES: tuple[AuthMode, ...] = ("auto", "azure_cli", "device_code", "none") _GRAPH_USER_READ_SCOPE = "https://graph.microsoft.com/User.Read" _GRAPH_DEFAULT_SCOPE = "https://graph.microsoft.com/.default" _TOKEN_REFRESH_BUFFER_SECONDS = 300 +_CACHE_ERROR_MARKERS = ("cache encryption", "persistent cache", "libsecret", "keychain") _AZURE_CLI_WARNING = ( "Warning: azure_cli mode sends the Azure CLI application's Microsoft Graph token to the backend. " "That token can contain permissions beyond User.Read. Prefer auto or device_code." @@ -104,13 +110,12 @@ async def close_async(self) -> None: """Release credential resources.""" -class _AzureIdentityTokenProvider: - """Adapt an asynchronous Azure Identity credential to the CLI token protocol.""" +class _AzureCliTokenProvider: + """Adapt Azure CLI credentials to the CLI token protocol.""" - def __init__(self, *, credential: Any, auth_config: BackendAuthConfig, mode: AuthMode) -> None: + def __init__(self, *, credential: Any, auth_config: BackendAuthConfig) -> None: self._credential = credential self._auth_config = auth_config - self._mode = mode self._access_token: Any = None async def get_token_async(self) -> str: @@ -133,32 +138,17 @@ async def get_token_async(self) -> str: try: access_token = await self._credential.get_token(_GRAPH_DEFAULT_SCOPE) except (ClientAuthenticationError, CredentialUnavailableError) as exc: - if self._mode == "azure_cli": - hint = f"Run 'az login --tenant {self._auth_config.tenant_id}' and try again." - else: - hint = "Confirm that device-code authentication is enabled for the CoPyRIT Entra application." - raise CliAuthenticationError(f"Entra authentication failed. {hint}") from exc - - token = getattr(access_token, "token", "") - if not isinstance(token, str) or not token: - raise CliAuthenticationError("Entra authentication returned an empty access token.") + raise CliAuthenticationError( + f"Entra authentication failed. Run 'az login --tenant {self._auth_config.tenant_id}' and try again." + ) from exc + + token = _get_access_token_text(access_token) self._access_token = access_token return token def _get_current_token(self) -> str | None: """Return a cached token that is valid beyond the refresh buffer.""" - if self._access_token is None: - return None - expires_on = getattr(self._access_token, "expires_on", 0) - token = getattr(self._access_token, "token", "") - if ( - isinstance(expires_on, int | float) - and expires_on > time.time() + _TOKEN_REFRESH_BUFFER_SECONDS - and isinstance(token, str) - and token - ): - return token - return None + return _get_current_token_text(self._access_token) async def close_async(self) -> None: """Close the underlying Azure Identity credential.""" @@ -168,9 +158,18 @@ async def close_async(self) -> None: class _DeviceCodeTokenProvider: """Adapt the synchronous device-code credential without blocking the event loop.""" - def __init__(self, *, credential: Any, auth_config: BackendAuthConfig) -> None: + def __init__( + self, + *, + credential: Any, + auth_config: BackendAuthConfig, + authentication_record_path: Path, + has_authentication_record: bool, + ) -> None: self._credential = credential self._auth_config = auth_config + self._authentication_record_path = authentication_record_path + self._has_authentication_record = has_authentication_record self._access_token: Any = None async def get_token_async(self) -> str: @@ -184,40 +183,51 @@ async def get_token_async(self) -> str: CliAuthenticationError: If Entra rejects device-code authentication. """ from azure.core.exceptions import ClientAuthenticationError - from azure.identity import CredentialUnavailableError + from azure.identity import AuthenticationRequiredError, CredentialUnavailableError cached_token = self._get_current_token() if cached_token is not None: return cached_token try: - access_token = await asyncio.to_thread(self._credential.get_token, *self._auth_config.scopes) + if not self._has_authentication_record: + await self._authenticate_async() + try: + access_token = await asyncio.to_thread(self._credential.get_token, *self._auth_config.scopes) + except AuthenticationRequiredError: + await self._authenticate_async() + access_token = await asyncio.to_thread(self._credential.get_token, *self._auth_config.scopes) except (ClientAuthenticationError, CredentialUnavailableError) as exc: + if _is_persistent_cache_error(exc): + raise CliAuthenticationError( + "Encrypted token caching is unavailable. Configure the platform credential store " + "(for example, libsecret on Linux) and try again." + ) from exc raise CliAuthenticationError( "Entra authentication failed. Confirm that device-code authentication " "is enabled for the CoPyRIT Entra application." ) from exc - token = getattr(access_token, "token", "") - if not isinstance(token, str) or not token: - raise CliAuthenticationError("Entra authentication returned an empty access token.") + token = _get_access_token_text(access_token) self._access_token = access_token return token + async def _authenticate_async(self) -> None: + """Authenticate interactively and persist the resulting account metadata.""" + authentication_record = await asyncio.to_thread( + self._credential.authenticate, + scopes=self._auth_config.scopes, + ) + await asyncio.to_thread( + _save_authentication_record, + authentication_record=authentication_record, + path=self._authentication_record_path, + ) + self._has_authentication_record = True + def _get_current_token(self) -> str | None: """Return a cached token that is valid beyond the refresh buffer.""" - if self._access_token is None: - return None - expires_on = getattr(self._access_token, "expires_on", 0) - token = getattr(self._access_token, "token", "") - if ( - isinstance(expires_on, int | float) - and expires_on > time.time() + _TOKEN_REFRESH_BUFFER_SECONDS - and isinstance(token, str) - and token - ): - return token - return None + return _get_current_token_text(self._access_token) async def close_async(self) -> None: """Close the underlying synchronous Azure Identity credential.""" @@ -229,23 +239,144 @@ def _is_interactive() -> bool: return bool(sys.stdin.isatty() and sys.stderr.isatty()) +def _get_access_token_text(access_token: Any) -> str: + """ + Return validated token text from an Azure Identity access token. + + Returns: + The non-empty access token text. + + Raises: + CliAuthenticationError: If Azure Identity returns an empty token. + """ + token = getattr(access_token, "token", "") + if not isinstance(token, str) or not token: + raise CliAuthenticationError("Entra authentication returned an empty access token.") + return token + + +def _get_current_token_text(access_token: Any) -> str | None: + """Return token text when an access token remains valid beyond the refresh buffer.""" + if access_token is None: + return None + expires_on = getattr(access_token, "expires_on", 0) + token = getattr(access_token, "token", "") + if ( + isinstance(expires_on, int | float) + and expires_on > time.time() + _TOKEN_REFRESH_BUFFER_SECONDS + and isinstance(token, str) + and token + ): + return token + return None + + +def _authentication_cache_key(*, auth_config: BackendAuthConfig) -> str: + """Return a path-safe identifier for one tenant and client pair.""" + cache_identity = f"{auth_config.tenant_id}:{auth_config.client_id}" + return hashlib.sha256(cache_identity.encode()).hexdigest()[:24] + + +def _authentication_record_path(*, auth_config: BackendAuthConfig) -> Path: + """Return the local path for non-secret Azure Identity account metadata.""" + cache_key = _authentication_cache_key(auth_config=auth_config) + return Path.home() / ".pyrit" / ".pyrit_cache" / f"copyrit-auth-{cache_key}.json" + + +def _load_authentication_record(*, auth_config: BackendAuthConfig, path: Path) -> Any | None: + """ + Load account metadata required to reuse the encrypted token cache. + + Returns: + The stored Azure Identity authentication record, or ``None`` when absent. + + Raises: + CliAuthenticationError: If the record is invalid or belongs to another app. + """ + from azure.identity import AuthenticationRecord + + if not path.exists(): + return None + try: + authentication_record = AuthenticationRecord.deserialize(path.read_text(encoding="utf-8")) + except (KeyError, OSError, ValueError) as exc: + raise CliAuthenticationError( + f"Could not read the CoPyRIT authentication record at {path}: {exc}. Remove the file and try again." + ) from exc + if authentication_record.client_id != auth_config.client_id: + raise CliAuthenticationError(f"The CoPyRIT authentication record at {path} does not match the server.") + return authentication_record + + +def _save_authentication_record(*, authentication_record: Any, path: Path) -> None: + """ + Atomically store non-secret account metadata for later cache access. + + Raises: + CliAuthenticationError: If the record cannot be serialized or saved. + """ + temporary_path = path.with_suffix(".tmp") + try: + path.parent.mkdir(parents=True, exist_ok=True) + temporary_path.write_text(authentication_record.serialize(), encoding="utf-8") + temporary_path.chmod(0o600) + temporary_path.replace(path) + except (OSError, ValueError) as exc: + raise CliAuthenticationError(f"Could not save the CoPyRIT authentication record at {path}: {exc}") from exc + + +def _print_device_code_prompt(verification_uri: str, user_code: str, expires_on: datetime) -> None: + """Write device-code instructions to stderr so stdout remains redirectable.""" + print( + f"To sign in, open {verification_uri} and enter code {user_code}. " + f"The code expires at {expires_on.isoformat()}.", + file=sys.stderr, + ) + + +def _is_persistent_cache_error(exc: BaseException) -> bool: + """Return whether an exception chain identifies a platform token-cache failure.""" + current: BaseException | None = exc + while current is not None: + message = str(current).lower() + if any(marker in message for marker in _CACHE_ERROR_MARKERS): + return True + current = current.__cause__ or current.__context__ + return False + + def _create_azure_cli_provider(*, auth_config: BackendAuthConfig) -> TokenProvider: from azure.identity.aio import AzureCliCredential credential = AzureCliCredential(tenant_id=auth_config.tenant_id) - return _AzureIdentityTokenProvider(credential=credential, auth_config=auth_config, mode="azure_cli") + return _AzureCliTokenProvider(credential=credential, auth_config=auth_config) -def _create_device_code_provider(*, auth_config: BackendAuthConfig) -> TokenProvider: +async def _create_device_code_provider_async(*, auth_config: BackendAuthConfig) -> TokenProvider: from azure.identity import DeviceCodeCredential, TokenCachePersistenceOptions - cache_options = TokenCachePersistenceOptions(name=f"pyrit-copyrit-{auth_config.client_id}") + cache_key = _authentication_cache_key(auth_config=auth_config) + authentication_record_path = _authentication_record_path(auth_config=auth_config) + authentication_record = await asyncio.to_thread( + _load_authentication_record, + auth_config=auth_config, + path=authentication_record_path, + ) + cache_options = TokenCachePersistenceOptions(name=f"pyrit-copyrit-{cache_key}") credential = DeviceCodeCredential( tenant_id=auth_config.tenant_id, client_id=auth_config.client_id, + authentication_record=authentication_record, cache_persistence_options=cache_options, + disable_automatic_authentication=True, + prompt_callback=_print_device_code_prompt, + ) + return _DeviceCodeTokenProvider( + credential=credential, + auth_config=auth_config, + authentication_record_path=authentication_record_path, + has_authentication_record=authentication_record is not None, ) - return _DeviceCodeTokenProvider(credential=credential, auth_config=auth_config) async def _verify_provider_async(*, provider: TokenProvider) -> TokenProvider: @@ -299,4 +430,5 @@ async def create_token_provider_async( "Device-code authentication requires an interactive terminal. " "Use azure_cli only when you accept sending the Azure CLI Graph token to the backend." ) - return await _verify_provider_async(provider=_create_device_code_provider(auth_config=auth_config)) + provider = await _create_device_code_provider_async(auth_config=auth_config) + return await _verify_provider_async(provider=provider) diff --git a/tests/unit/cli/test_api_client.py b/tests/unit/cli/test_api_client.py index 28988ba6d9..96cc41769e 100644 --- a/tests/unit/cli/test_api_client.py +++ b/tests/unit/cli/test_api_client.py @@ -240,6 +240,20 @@ async def test_context_manager_rejects_authentication_over_remote_http(mock_http mock_httpx_client.aclose.assert_awaited_once() +async def test_context_manager_rejects_invalid_auth_config_json(mock_httpx_client): + c = PyRITApiClient(base_url="https://copyrit.example.com", auth_mode="auto") + fake_async_client_cls = MagicMock(return_value=mock_httpx_client) + response = _make_response() + response.json.side_effect = ValueError("invalid JSON") + mock_httpx_client.get.return_value = response + + with patch("httpx.AsyncClient", fake_async_client_cls): + with pytest.raises(CliAuthenticationError, match="invalid JSON"): + await c.__aenter__() + + mock_httpx_client.aclose.assert_awaited_once() + + def test_get_client_raises_when_not_opened(): c = PyRITApiClient(base_url="http://localhost:8000") with pytest.raises(ServerNotAvailableError, match="not connected"): diff --git a/tests/unit/cli/test_auth.py b/tests/unit/cli/test_auth.py index cfb0b54a40..6affe2777d 100644 --- a/tests/unit/cli/test_auth.py +++ b/tests/unit/cli/test_auth.py @@ -3,9 +3,13 @@ """Tests for thin-client authentication helpers.""" +from datetime import datetime, timezone +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import pytest +from azure.core.exceptions import ClientAuthenticationError +from azure.identity import AuthenticationRecord, AuthenticationRequiredError, CredentialUnavailableError from pyrit.cli import _auth from pyrit.cli._auth import BackendAuthConfig, CliAuthenticationError, create_token_provider_async @@ -43,6 +47,16 @@ def test_backend_auth_config_accepts_legacy_disabled_contract() -> None: assert result == BackendAuthConfig(enabled=False, tenant_id="", client_id="", scopes=()) +def test_backend_auth_config_rejects_non_mapping_payload() -> None: + with pytest.raises(CliAuthenticationError, match="unsupported authentication contract"): + BackendAuthConfig.from_payload([]) + + +def test_backend_auth_config_rejects_non_string_identifiers() -> None: + with pytest.raises(CliAuthenticationError, match="invalid Entra tenant or client"): + BackendAuthConfig.from_payload({"enabled": False, "tenantId": 1, "clientId": "client-id"}) + + def test_backend_auth_config_rejects_non_graph_scope() -> None: with pytest.raises(CliAuthenticationError, match="unsupported authentication scope"): BackendAuthConfig.from_payload( @@ -62,6 +76,12 @@ def test_backend_auth_config_rejects_non_graph_scope() -> None: {"enabled": True, "tenantId": "", "clientId": "client-id", "scopes": ["scope"]}, {"enabled": True, "tenantId": "tenant-id", "clientId": "client-id", "scopes": []}, {"enabled": True, "tenantId": "tenant-id", "clientId": "client-id", "scopes": "scope"}, + { + "enabled": True, + "tenantId": "", + "clientId": "client-id", + "scopes": ["https://graph.microsoft.com/User.Read"], + }, ], ) def test_backend_auth_config_rejects_invalid_contract(payload: object) -> None: @@ -82,7 +102,12 @@ async def test_create_token_provider_auto_uses_device_code() -> None: provider.get_token_async = AsyncMock(return_value="token") provider.close_async = AsyncMock() - with patch.object(_auth, "_create_device_code_provider", return_value=provider) as create_device_code: + with patch.object( + _auth, + "_create_device_code_provider_async", + new_callable=AsyncMock, + return_value=provider, + ) as create_device_code: result = await create_token_provider_async( auth_config=_AUTH_CONFIG, auth_mode="auto", @@ -90,7 +115,7 @@ async def test_create_token_provider_auto_uses_device_code() -> None: ) assert result is provider - create_device_code.assert_called_once_with(auth_config=_AUTH_CONFIG) + create_device_code.assert_awaited_once_with(auth_config=_AUTH_CONFIG) provider.get_token_async.assert_awaited_once() provider.close_async.assert_not_awaited() @@ -129,34 +154,57 @@ async def test_create_token_provider_azure_cli_warns(capsys) -> None: assert "permissions beyond User.Read" in capsys.readouterr().err -def test_create_device_code_provider_uses_persistent_cache() -> None: +async def test_create_device_code_provider_uses_persistent_cache(tmp_path: Path) -> None: cache_options = MagicMock() credential = MagicMock() + record_path = tmp_path / "record.json" with ( + patch.object(_auth, "_authentication_record_path", return_value=record_path), patch("azure.identity.TokenCachePersistenceOptions", return_value=cache_options) as cache_type, patch("azure.identity.DeviceCodeCredential", return_value=credential) as credential_type, ): - provider = _auth._create_device_code_provider(auth_config=_AUTH_CONFIG) + provider = await _auth._create_device_code_provider_async(auth_config=_AUTH_CONFIG) - cache_type.assert_called_once_with(name="pyrit-copyrit-client-id") + cache_type.assert_called_once_with( + name=f"pyrit-copyrit-{_auth._authentication_cache_key(auth_config=_AUTH_CONFIG)}" + ) credential_type.assert_called_once_with( tenant_id="tenant-id", client_id="client-id", + authentication_record=None, cache_persistence_options=cache_options, + disable_automatic_authentication=True, + prompt_callback=_auth._print_device_code_prompt, ) assert isinstance(provider, _auth._DeviceCodeTokenProvider) +async def test_create_device_code_provider_reuses_authentication_record(tmp_path: Path) -> None: + record_path = tmp_path / "record.json" + authentication_record = MagicMock() + + with ( + patch.object(_auth, "_authentication_record_path", return_value=record_path), + patch.object(_auth, "_load_authentication_record", return_value=authentication_record), + patch("azure.identity.TokenCachePersistenceOptions"), + patch("azure.identity.DeviceCodeCredential") as credential_type, + ): + provider = await _auth._create_device_code_provider_async(auth_config=_AUTH_CONFIG) + + assert credential_type.call_args.kwargs["authentication_record"] is authentication_record + assert isinstance(provider, _auth._DeviceCodeTokenProvider) + assert provider._has_authentication_record is True + + async def test_azure_cli_provider_caches_token_until_refresh_window() -> None: credential = MagicMock() credential.get_token = AsyncMock() credential.close = AsyncMock() credential.get_token.return_value = MagicMock(token="access-token", expires_on=2_000_000_000) - provider = _auth._AzureIdentityTokenProvider( + provider = _auth._AzureCliTokenProvider( credential=credential, auth_config=_AUTH_CONFIG, - mode="azure_cli", ) with patch("pyrit.cli._auth.time.time", return_value=1_000_000_000): @@ -165,3 +213,280 @@ async def test_azure_cli_provider_caches_token_until_refresh_window() -> None: assert first == second == "access-token" credential.get_token.assert_awaited_once_with("https://graph.microsoft.com/.default") + + +async def test_azure_cli_provider_refreshes_expired_token() -> None: + credential = MagicMock() + credential.get_token = AsyncMock( + side_effect=[ + MagicMock(token="first", expires_on=1_000_000_100), + MagicMock(token="second", expires_on=2_000_000_000), + ] + ) + provider = _auth._AzureCliTokenProvider(credential=credential, auth_config=_AUTH_CONFIG) + + with patch("pyrit.cli._auth.time.time", return_value=1_000_000_000): + assert await provider.get_token_async() == "first" + assert await provider.get_token_async() == "second" + + assert credential.get_token.await_count == 2 + + +@pytest.mark.parametrize("exception_type", [ClientAuthenticationError, CredentialUnavailableError]) +async def test_azure_cli_provider_wraps_authentication_failure(exception_type: type[Exception]) -> None: + credential = MagicMock() + credential.get_token = AsyncMock(side_effect=exception_type("failed")) + provider = _auth._AzureCliTokenProvider(credential=credential, auth_config=_AUTH_CONFIG) + + with pytest.raises(CliAuthenticationError, match="az login --tenant tenant-id"): + await provider.get_token_async() + + +async def test_azure_cli_provider_rejects_empty_token_and_closes() -> None: + credential = MagicMock() + credential.get_token = AsyncMock(return_value=MagicMock(token="", expires_on=2_000_000_000)) + credential.close = AsyncMock() + provider = _auth._AzureCliTokenProvider(credential=credential, auth_config=_AUTH_CONFIG) + + with pytest.raises(CliAuthenticationError, match="empty access token"): + await provider.get_token_async() + await provider.close_async() + + credential.close.assert_awaited_once() + + +def test_create_azure_cli_provider_uses_discovered_tenant() -> None: + credential = MagicMock() + + with patch("azure.identity.aio.AzureCliCredential", return_value=credential) as credential_type: + provider = _auth._create_azure_cli_provider(auth_config=_AUTH_CONFIG) + + credential_type.assert_called_once_with(tenant_id="tenant-id") + assert isinstance(provider, _auth._AzureCliTokenProvider) + + +async def test_device_code_provider_authenticates_once_and_persists_record(tmp_path: Path) -> None: + credential = MagicMock() + authentication_record = MagicMock() + credential.authenticate.return_value = authentication_record + credential.get_token.return_value = MagicMock(token="access-token", expires_on=2_000_000_000) + record_path = tmp_path / "record.json" + provider = _auth._DeviceCodeTokenProvider( + credential=credential, + auth_config=_AUTH_CONFIG, + authentication_record_path=record_path, + has_authentication_record=False, + ) + + with ( + patch("pyrit.cli._auth.time.time", return_value=1_000_000_000), + patch.object(_auth, "_save_authentication_record") as save_record, + ): + assert await provider.get_token_async() == "access-token" + assert await provider.get_token_async() == "access-token" + + credential.authenticate.assert_called_once_with(scopes=_AUTH_CONFIG.scopes) + save_record.assert_called_once_with(authentication_record=authentication_record, path=record_path) + credential.get_token.assert_called_once_with(*_AUTH_CONFIG.scopes) + + +async def test_device_code_provider_uses_loaded_record_without_authenticating(tmp_path: Path) -> None: + credential = MagicMock() + credential.get_token.return_value = MagicMock(token="access-token", expires_on=2_000_000_000) + provider = _auth._DeviceCodeTokenProvider( + credential=credential, + auth_config=_AUTH_CONFIG, + authentication_record_path=tmp_path / "record.json", + has_authentication_record=True, + ) + + assert await provider.get_token_async() == "access-token" + + credential.authenticate.assert_not_called() + + +async def test_device_code_provider_reauthenticates_and_updates_stale_record(tmp_path: Path) -> None: + credential = MagicMock() + authentication_record = MagicMock() + credential.authenticate.return_value = authentication_record + credential.get_token.side_effect = [ + AuthenticationRequiredError(_AUTH_CONFIG.scopes), + MagicMock(token="access-token", expires_on=2_000_000_000), + ] + record_path = tmp_path / "record.json" + provider = _auth._DeviceCodeTokenProvider( + credential=credential, + auth_config=_AUTH_CONFIG, + authentication_record_path=record_path, + has_authentication_record=True, + ) + + with patch.object(_auth, "_save_authentication_record") as save_record: + assert await provider.get_token_async() == "access-token" + + credential.authenticate.assert_called_once_with(scopes=_AUTH_CONFIG.scopes) + save_record.assert_called_once_with(authentication_record=authentication_record, path=record_path) + assert credential.get_token.call_count == 2 + + +async def test_device_code_provider_reports_cache_failure(tmp_path: Path) -> None: + cache_error = ValueError("Cache encryption is impossible because libsecret is unavailable") + auth_error = ClientAuthenticationError("Authentication failed") + auth_error.__cause__ = cache_error + credential = MagicMock() + credential.authenticate.side_effect = auth_error + provider = _auth._DeviceCodeTokenProvider( + credential=credential, + auth_config=_AUTH_CONFIG, + authentication_record_path=tmp_path / "record.json", + has_authentication_record=False, + ) + + with pytest.raises(CliAuthenticationError, match="Encrypted token caching is unavailable"): + await provider.get_token_async() + + +async def test_device_code_provider_reports_entra_failure(tmp_path: Path) -> None: + credential = MagicMock() + credential.get_token.side_effect = CredentialUnavailableError("failed") + provider = _auth._DeviceCodeTokenProvider( + credential=credential, + auth_config=_AUTH_CONFIG, + authentication_record_path=tmp_path / "record.json", + has_authentication_record=True, + ) + + with pytest.raises(CliAuthenticationError, match="device-code authentication is enabled"): + await provider.get_token_async() + + +async def test_device_code_provider_rejects_empty_token_and_closes(tmp_path: Path) -> None: + credential = MagicMock() + credential.get_token.return_value = MagicMock(token="", expires_on=2_000_000_000) + provider = _auth._DeviceCodeTokenProvider( + credential=credential, + auth_config=_AUTH_CONFIG, + authentication_record_path=tmp_path / "record.json", + has_authentication_record=True, + ) + + with pytest.raises(CliAuthenticationError, match="empty access token"): + await provider.get_token_async() + await provider.close_async() + + credential.close.assert_called_once() + + +def test_authentication_record_round_trip(tmp_path: Path) -> None: + path = tmp_path / "nested" / "record.json" + authentication_record = AuthenticationRecord( + tenant_id="tenant-id", + client_id="client-id", + authority="https://login.microsoftonline.com", + home_account_id="home-account-id", + username="user@example.com", + ) + + _auth._save_authentication_record(authentication_record=authentication_record, path=path) + result = _auth._load_authentication_record(auth_config=_AUTH_CONFIG, path=path) + + assert result.tenant_id == "tenant-id" + assert result.client_id == "client-id" + + +def test_save_authentication_record_wraps_serialization_failure(tmp_path: Path) -> None: + authentication_record = MagicMock() + authentication_record.serialize.side_effect = ValueError("invalid record") + + with pytest.raises(CliAuthenticationError, match="Could not save"): + _auth._save_authentication_record( + authentication_record=authentication_record, + path=tmp_path / "record.json", + ) + + +def test_load_authentication_record_returns_none_when_absent(tmp_path: Path) -> None: + assert _auth._load_authentication_record(auth_config=_AUTH_CONFIG, path=tmp_path / "missing.json") is None + + +def test_load_authentication_record_rejects_invalid_file(tmp_path: Path) -> None: + path = tmp_path / "record.json" + path.write_text("not json", encoding="utf-8") + + with pytest.raises(CliAuthenticationError, match="Could not read"): + _auth._load_authentication_record(auth_config=_AUTH_CONFIG, path=path) + + +def test_load_authentication_record_rejects_wrong_client(tmp_path: Path) -> None: + path = tmp_path / "record.json" + authentication_record = AuthenticationRecord( + tenant_id="tenant-id", + client_id="other-client", + authority="https://login.microsoftonline.com", + home_account_id="home-account-id", + username="user@example.com", + ) + path.write_text(authentication_record.serialize(), encoding="utf-8") + + with pytest.raises(CliAuthenticationError, match="does not match"): + _auth._load_authentication_record(auth_config=_AUTH_CONFIG, path=path) + + +def test_load_authentication_record_accepts_tenant_alias(tmp_path: Path) -> None: + path = tmp_path / "record.json" + authentication_record = AuthenticationRecord( + tenant_id="tenant-guid", + client_id="client-id", + authority="https://login.microsoftonline.com", + home_account_id="home-account-id", + username="user@example.com", + ) + path.write_text(authentication_record.serialize(), encoding="utf-8") + auth_config = BackendAuthConfig( + enabled=True, + tenant_id="contoso.onmicrosoft.com", + client_id="client-id", + scopes=_AUTH_CONFIG.scopes, + ) + + result = _auth._load_authentication_record(auth_config=auth_config, path=path) + + assert result.tenant_id == "tenant-guid" + + +def test_authentication_record_path_uses_hashed_identity() -> None: + path = _auth._authentication_record_path(auth_config=_AUTH_CONFIG) + + assert path.parent.name == ".pyrit_cache" + assert "tenant-id" not in path.name + assert "client-id" not in path.name + + +def test_device_code_prompt_uses_stderr(capsys) -> None: + expires_on = datetime(2026, 1, 1, tzinfo=timezone.utc) + + _auth._print_device_code_prompt("https://microsoft.com/devicelogin", "ABCD-EFGH", expires_on) + + captured = capsys.readouterr() + assert captured.out == "" + assert "ABCD-EFGH" in captured.err + assert expires_on.isoformat() in captured.err + + +def test_is_interactive_requires_stdin_and_stderr_ttys() -> None: + with ( + patch("sys.stdin.isatty", return_value=True), + patch("sys.stderr.isatty", return_value=False), + ): + assert _auth._is_interactive() is False + + +async def test_verify_provider_closes_after_authentication_failure() -> None: + provider = MagicMock() + provider.get_token_async = AsyncMock(side_effect=CliAuthenticationError("failed")) + provider.close_async = AsyncMock() + + with pytest.raises(CliAuthenticationError, match="failed"): + await _auth._verify_provider_async(provider=provider) + + provider.close_async.assert_awaited_once() diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index 81ee21f62f..9fd5aaaa8b 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -652,6 +652,16 @@ def test_main_failed_scenario(self, mock_client_class, mock_probe): # --------------------------------------------------------------------------- +def test_resolve_auth_mode_rejects_unsupported_programmatic_value() -> None: + parsed_args = Namespace(config_file=None, auth_mode="invalid") + + with ( + patch("pyrit.cli._config_reader.read_server_settings", return_value=MagicMock(auth_mode="auto")), + pytest.raises(ValueError, match="Unsupported authentication mode"), + ): + pyrit_scan._resolve_auth_mode(parsed_args=parsed_args) + + class TestStopServerOnPort: """Tests for stop_server_on_port helper (now lives in _server_launcher).""" diff --git a/tests/unit/cli/test_pyrit_shell.py b/tests/unit/cli/test_pyrit_shell.py index beaf32d6d3..25be9e4d4e 100644 --- a/tests/unit/cli/test_pyrit_shell.py +++ b/tests/unit/cli/test_pyrit_shell.py @@ -465,6 +465,16 @@ def test_auth_discovery_http_failure_returns_false(self, capsys): assert s._api_client is None assert "discovery failed" in capsys.readouterr().out + def test_config_failure_returns_false(self, capsys): + from pyrit.cli._config_reader import ConfigError + + s = pyrit_shell.PyRITShell(no_animation=True, server_url="https://copyrit.example.com") + with patch.object(s, "_resolve_auth_mode", side_effect=ConfigError("invalid config")): + assert s._open_client(base_url="https://copyrit.example.com") is False + + assert s._api_client is None + assert "invalid config" in capsys.readouterr().out + def test_start_server_failure_returns_false(self, capsys): s = pyrit_shell.PyRITShell(no_animation=True, start_server=True) with (