chaos-mcp: let customers choose user-principal vs managed-identity auth at runtime

copilot-cli-plugin/mcp/README.md

@@ -3,10 +3,49 @@
 MCP server exposing Azure Chaos Studio v2 operations as agent-callable tools.
 
 The server relies on the user's local `az` CLI session for authentication
-rather than managing tokens itself, keeping the server stateless. See the
+rather than managing tokens itself, keeping the server stateless. It can also
+be pointed at an Azure Managed Identity for unattended runs — see
+[Authentication](#authentication). See the
 [copilot-cli-plugin README](../README.md) for the full tool table and how
 this package fits into the Chaos Studio Copilot CLI plugin.
 
+## Authentication
+
+By default every tool borrows the operator's local `az login` session — the
+**user principal** — via `az account get-access-token`.
+
+### Choose the mode during a session (no config change)
+
+The customer can switch between the user principal and a **Managed Identity**
+mid-conversation by asking the agent to call the `chaos_set_auth_mode` tool —
+no config edit or server restart needed. The choice applies to every subsequent
+tool call for the life of the session.
+
+| Tool | Purpose |
+|---|---|
+| `chaos_set_auth_mode(mode, msi_client_id?)` | `mode` = `cli` (user principal) or `managed-identity` (aliases `msi`/`mi`); `msi_client_id` optionally pins a user-assigned identity. |
+| `chaos_get_auth_mode()` | Report the effective `{mode, msiClientId, source}`. |
+
+> **Attribution & approval.** In managed-identity mode every ARM action — including the Reader role assignments made during workspace creation — is recorded in the Azure activity log as the **identity**, not the human operator. Review your audit/attribution requirements before switching. Because `chaos_set_auth_mode` is agent-callable, your MCP client's tool-approval prompt is the control that stops a session from being silently flipped onto a privileged host identity — do **not** blanket-auto-approve this tool.
+
+### Startup default (optional env vars)
+
+For unattended hosts (CI, containers, AKS, VMs) you can also set the initial
+mode via env vars so no interactive `az login` is required. A runtime
+`chaos_set_auth_mode` call always takes precedence over these.
+
+| Env var | Values | Effect |
+|---|---|---|
+| `CHAOS_MCP_AUTH_MODE` | `cli` (default), `managed-identity` (aliases: `msi`, `mi`) | Selects the token source. |
+| `CHAOS_MCP_MSI_CLIENT_ID` | client id (optional) | Pins a **user-assigned** identity; omit to use the **system-assigned** identity. |
+
+In `managed-identity` mode the server acquires tokens from the Azure identity
+endpoint — the App Service / Container Apps / Functions endpoint
+(`IDENTITY_ENDPOINT` + `IDENTITY_HEADER`) when present, otherwise IMDS
+(`169.254.169.254`) for VMs / VMSS / AKS. Grant the identity the same ARM roles
+you would grant the user (managing `Microsoft.Chaos` workspaces and querying
+Azure Monitor).
+
 ## Install
 
 ```bash
@@ -27,7 +66,8 @@ Register the server in your MCP client config (see
 - **Claude Code**: `claude mcp add chaos-studio -- chaos-mcp`
 - **Cursor**: add the block above to `.cursor/mcp.json`.
 
-The server requires an active `az login` session; tools return a structured
+The server requires an active `az login` session (or a managed identity — see
+[Authentication](#authentication)); tools return a structured
 `{"ok": false, "errorType": ...}` envelope (rather than raising) when auth or
 permissions are missing, so agents can remediate and retry.
 

copilot-cli-plugin/mcp/chaos_mcp/azure.py

@@ -1,12 +1,19 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
 """Thin wrappers around `az` CLI + ARM REST calls.
 
-The MCP server intentionally relies on the user's local `az` session for auth
+By default the MCP server relies on the operator's local `az` session for auth
 rather than managing tokens itself. This matches the skill's auth model and
-keeps the server stateless.
+keeps the server stateless. Setting ``CHAOS_MCP_AUTH_MODE=managed-identity``
+flips a lever so the server instead acquires tokens from an Azure Managed
+Identity — letting the skills run unattended (CI, containers, AKS, VMs) with
+no interactive `az login`.
 """
 from __future__ import annotations
 
 import json
+import os
 import shutil
 import subprocess
 import time
@@ -25,6 +32,127 @@
 # Production callers leave it None; pytest sets it to an httpx.MockTransport.
 _TEST_TRANSPORT: Any = None
 
+# -----------------------------------------------------------------------------
+# Authentication mode (the "user principal vs managed identity" lever)
+# -----------------------------------------------------------------------------
+# The server can authenticate either as the operator's local `az login` session
+# (the user principal, default) or as an Azure Managed Identity — so the tools
+# can run unattended (CI, containers, AKS, VMs) with no `az` session.
+#
+# The mode is chosen at three levels, highest precedence first:
+#   1. A runtime override set *during the session* via `set_auth_mode(...)`
+#      (surfaced to agents as the `chaos_set_auth_mode` MCP tool). This lets the
+#      customer flip between user principal and MI mid-conversation, no config
+#      edit or server restart required.
+#   2. The CHAOS_MCP_AUTH_MODE / CHAOS_MCP_MSI_CLIENT_ID env vars (a startup
+#      default, e.g. for headless deployments).
+#   3. The built-in default: `cli` (user principal).
+AUTH_MODE_ENV = "CHAOS_MCP_AUTH_MODE"
+MSI_CLIENT_ID_ENV = "CHAOS_MCP_MSI_CLIENT_ID"
+
+AUTH_MODE_CLI = "cli"
+AUTH_MODE_MANAGED_IDENTITY = "managed-identity"
+
+_MANAGED_IDENTITY_ALIASES = frozenset(
+    {"managed-identity", "managed_identity", "managedidentity", "msi", "mi", "identity"}
+)
+_CLI_ALIASES = frozenset({"cli", "az", "user", "user-principal", "userprincipal", "default"})
+
+# IMDS token endpoint (VMs, VMSS, AKS).
+IMDS_TOKEN_ENDPOINT = "http://169.254.169.254/metadata/identity/oauth2/token"
+IMDS_API_VERSION = "2018-02-01"
+# App Service / Container Apps / Functions inject IDENTITY_ENDPOINT +
+# IDENTITY_HEADER instead of exposing IMDS.
+APP_SERVICE_IDENTITY_API_VERSION = "2019-08-01"
+
+# Session-scoped runtime override (set via the chaos_set_auth_mode tool). None
+# means "not overridden — fall back to env / default". This is process-global
+# in-memory state: it lives for the life of the MCP server process (i.e. the
+# customer's session) and is never persisted to disk.
+_auth_mode_override: str | None = None
+_msi_client_id_override: str | None = None
+
+
+def _normalize_mode(value: str) -> str:
+    """Map a user-supplied mode string to a canonical mode, or raise.
+
+    Accepts the managed-identity aliases (managed-identity/msi/mi/...) and the
+    cli aliases (cli/az/user-principal/...).
+    """
+    raw = (value or "").strip().lower()
+    if raw in _MANAGED_IDENTITY_ALIASES:
+        return AUTH_MODE_MANAGED_IDENTITY
+    if raw in _CLI_ALIASES:
+        return AUTH_MODE_CLI
+    raise AzureError(
+        f"Unknown auth mode '{value}'. Use 'cli' (user principal) or "
+        "'managed-identity'."
+    )
+
+
+def set_auth_mode(mode: str, msi_client_id: str | None = None) -> dict:
+    """Set the session-scoped auth mode override and return the effective config.
+
+    `mode` is 'cli' (user principal) or 'managed-identity' (aliases accepted).
+    `msi_client_id` optionally pins a user-assigned identity when mode is
+    managed-identity; pass None/empty to use the system-assigned identity.
+    """
+    global _auth_mode_override, _msi_client_id_override
+    normalized = _normalize_mode(mode)
+    _auth_mode_override = normalized
+    if normalized == AUTH_MODE_MANAGED_IDENTITY:
+        _msi_client_id_override = (msi_client_id or "").strip() or None
+    else:
+        _msi_client_id_override = None
+    return get_auth_config()
+
+
+def reset_auth_mode() -> dict:
+    """Clear the runtime override so mode falls back to env vars / default."""
+    global _auth_mode_override, _msi_client_id_override
+    _auth_mode_override = None
+    _msi_client_id_override = None
+    return get_auth_config()
+
+
+def get_auth_config() -> dict:
+    """Return the effective auth configuration and where it came from."""
+    source = "override" if _auth_mode_override is not None else (
+        "env" if os.environ.get(AUTH_MODE_ENV) else "default"
+    )
+    return {
+        "mode": _auth_mode(),
+        "msiClientId": _msi_client_id(),
+        "source": source,
+    }
+
+
+def _auth_mode() -> str:
+    """Return the effective auth mode: 'managed-identity' or 'cli' (default).
+
+    Precedence: runtime override > CHAOS_MCP_AUTH_MODE env var > 'cli'.
+    """
+    if _auth_mode_override is not None:
+        return _auth_mode_override
+    raw = (os.environ.get(AUTH_MODE_ENV) or "cli").strip().lower()
+    return AUTH_MODE_MANAGED_IDENTITY if raw in _MANAGED_IDENTITY_ALIASES else AUTH_MODE_CLI
+
+
+def _msi_client_id() -> str | None:
+    """Return the effective user-assigned MI client id, or None.
+
+    Only meaningful in managed-identity mode (returns None otherwise). Within MI
+    mode the precedence is: explicit runtime override > CHAOS_MCP_MSI_CLIENT_ID
+    env var > None (the system-assigned identity). A mode switch that doesn't
+    name a client id therefore keeps any env-pinned identity rather than
+    silently dropping it.
+    """
+    if _auth_mode() != AUTH_MODE_MANAGED_IDENTITY:
+        return None
+    if _msi_client_id_override is not None:
+        return _msi_client_id_override
+    return (os.environ.get(MSI_CLIENT_ID_ENV) or "").strip() or None
+
 
 class AzureError(RuntimeError):
     """Raised when an ARM call or `az` invocation fails."""
@@ -65,12 +193,23 @@ def az_show_account() -> AzContext:
 
 
 def _get_token(resource: str = ARM_ENDPOINT) -> str:
-    """Acquire an access token for the given audience via the local `az` session.
+    """Acquire an access token for the given audience.
 
     `resource` is the token audience URL (e.g., `https://management.azure.com`
     for ARM or `https://api.loganalytics.io` for Log Analytics queries).
-    Tokens are not cached — each call invokes `az account get-access-token`.
+
+    Uses the local `az` session (user principal) by default; when
+    ``CHAOS_MCP_AUTH_MODE`` selects a managed identity, tokens come from the
+    Azure identity endpoint instead. Tokens are not cached — each call acquires
+    a fresh token.
     """
+    if _auth_mode() == "managed-identity":
+        return _get_token_via_managed_identity(resource)
+    return _get_token_via_cli(resource)
+
+
+def _get_token_via_cli(resource: str) -> str:
+    """Acquire a token from the local `az` session (the user principal)."""
     proc = subprocess.run(
         [_az_path(), "account", "get-access-token", "--resource", resource, "-o", "json"],
         capture_output=True,
@@ -83,6 +222,64 @@ def _get_token(resource: str = ARM_ENDPOINT) -> str:
     return json.loads(proc.stdout)["accessToken"]
 
 
+def _get_token_via_managed_identity(resource: str) -> str:
+    """Acquire a token from an Azure Managed Identity (no `az` session needed).
+
+    Honors the App Service / Container Apps / Functions identity endpoint
+    (``IDENTITY_ENDPOINT`` + ``IDENTITY_HEADER``) when present, otherwise falls
+    back to the IMDS endpoint used by VMs, VMSS and AKS. The user-assigned
+    identity (if any) comes from the runtime override or
+    ``CHAOS_MCP_MSI_CLIENT_ID``.
+    """
+    client_id = _msi_client_id()
+    identity_endpoint = os.environ.get("IDENTITY_ENDPOINT")
+    identity_header = os.environ.get("IDENTITY_HEADER")
+
+    if identity_endpoint and identity_header:
+        url = identity_endpoint
+        params: dict[str, str] = {
+            "resource": resource,
+            "api-version": APP_SERVICE_IDENTITY_API_VERSION,
+        }
+        headers = {"X-IDENTITY-HEADER": identity_header}
+    else:
+        url = IMDS_TOKEN_ENDPOINT
+        params = {"resource": resource, "api-version": IMDS_API_VERSION}
+        headers = {"Metadata": "true"}
+    if client_id:
+        params["client_id"] = client_id
+
+    try:
+        # trust_env=False: never route the IMDS / identity-endpoint request
+        # through HTTP(S)_PROXY. 169.254.169.254 must be reached directly, and
+        # proxying would also leak the App Service X-IDENTITY-HEADER secret. The
+        # Azure SDKs bypass proxies for IMDS for the same reasons.
+        resp = httpx.get(url, params=params, headers=headers, timeout=10.0, trust_env=False)
+    except httpx.HTTPError as e:
+        raise AzureError(
+            f"Failed to reach the managed-identity token endpoint for {resource}: {e}. "
+            "Is a managed identity available on this host?"
+        ) from e
+
+    if resp.status_code != 200:
+        raise AzureError(
+            f"Managed-identity token request for {resource} failed with HTTP "
+            f"{resp.status_code}: {resp.text.strip()}"
+        )
+
+    try:
+        token = resp.json().get("access_token")
+    except Exception as e:  # noqa: BLE001
+        raise AzureError(
+            f"Malformed managed-identity token response for {resource}: {resp.text.strip()}"
+        ) from e
+    if not token:
+        raise AzureError(
+            f"Managed-identity token response for {resource} contained no access_token."
+        )
+    return token
+
+
 def az_get_arm_token() -> str:
     """Back-compat shim — acquire a token scoped to ARM."""
     return _get_token(ARM_ENDPOINT)

copilot-cli-plugin/mcp/chaos_mcp/server.py

@@ -1,3 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
 """MCP server exposing Azure Chaos Studio v2 operations as agent-callable tools.
 
 Mirrors the Copilot CLI plugin's PowerShell skills (create-workspace,
@@ -81,6 +84,47 @@ def _grant_reader(scope: str, principal_id: str) -> dict[str, Any]:
 # ---------------------------------------------------------------------------
 
 
+@mcp.tool()
+def chaos_set_auth_mode(
+    mode: str,
+    msi_client_id: str | None = None,
+) -> dict[str, Any]:
+    """Choose how the Chaos Studio tools authenticate to Azure, for the rest of
+    this session.
+
+    Call this when the customer wants the tools to act as an Azure **Managed
+    Identity** instead of their signed-in **user principal** (the default), or
+    to switch back. The choice is applied immediately to every subsequent tool
+    call and persists for the life of this MCP session — no config edit or
+    server restart needed.
+
+    Args:
+        mode: 'cli' to use the local `az login` session (the user principal), or
+            'managed-identity' to use an Azure Managed Identity (aliases 'msi',
+            'mi' are accepted).
+        msi_client_id: Optional client id of a user-assigned managed identity to
+            use when mode is 'managed-identity'. Omit to keep any identity pinned
+            by CHAOS_MCP_MSI_CLIENT_ID, or to fall back to the system-assigned
+            identity when none is pinned. Ignored in 'cli' mode.
+
+    Returns the effective auth configuration ({mode, msiClientId, source}).
+    """
+    try:
+        return _ok(az.set_auth_mode(mode, msi_client_id))
+    except az.AzureError as e:
+        return _err(e)
+
+
+@mcp.tool()
+def chaos_get_auth_mode() -> dict[str, Any]:
+    """Report how the Chaos Studio tools are currently authenticating.
+
+    Returns {mode, msiClientId, source} where `source` is 'override' (set this
+    session via chaos_set_auth_mode), 'env' (from CHAOS_MCP_AUTH_MODE), or
+    'default' (the built-in `cli` / user-principal default)."""
+    return _ok(az.get_auth_config())
+
+
 @mcp.tool()
 def chaos_create_workspace(
     subscription_id: str,