Identity mfa support (#42568)

* Initial plan

* Implement claims challenge error for AzureCliCredential get_token and get_token_info methods

Co-authored-by: xiangyan99 <[email protected]>

* Refine claims challenge handling to ignore whitespace-only claims and add comprehensive tests

Co-authored-by: xiangyan99 <[email protected]>

* Fix MyPy errors in AzureCliCredential claims challenge handling

Co-authored-by: xiangyan99 <[email protected]>

* Fix pylint line-too-long and trailing-whitespace errors in AzureCliCredential

Co-authored-by: xiangyan99 <[email protected]>

* Apply black formatting to AzureCliCredential files

Co-authored-by: xiangyan99 <[email protected]>

* black

* Update error message format and docstrings for claims challenge handling

- Change error message from "Fail to get token, please run" to "Failed to get token. Run"
- Update docstrings to clarify error conditions as requested in code review
- Update test expectations to match new error message format

Co-authored-by: scottaddie <[email protected]>

* black

* Include scopes in az login command for claims challenge error messages

Co-authored-by: weikanglim <[email protected]>

* Refactor claims challenge logic to _get_token_base and add tenant support

- Move claims challenge logic from get_token and get_token_info methods to _get_token_base method to eliminate code duplication
- Pass claims through TokenRequestOptions instead of checking directly in individual methods
- Add tenant_id support in error messages when provided in options
- Update error message format to include tenant: "az login --claims-challenge {claims} [--tenant {tenant}] [--scope {scope}]"
- Add comprehensive tests for tenant functionality in both sync and async versions
- Maintains backward compatibility while providing more complete actionable commands

Co-authored-by: pvaneck <[email protected]>

* Added MFA support for developer credentials.

* Minor refactor and test updates

Signed-off-by: Paul Van Eck <[email protected]>

* Update error messages

Signed-off-by: Paul Van Eck <[email protected]>

---------

Signed-off-by: Paul Van Eck <[email protected]>
Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: xiangyan99 <[email protected]>
Co-authored-by: scottaddie <[email protected]>
Co-authored-by: weikanglim <[email protected]>
Co-authored-by: pvaneck <[email protected]>
Co-authored-by: Paul Van Eck <[email protected]>

Author: 5 people

sdk/identity/azure-identity/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- `AzureDeveloperCliCredential` now supports `claims` in `get_token` and `get_token_info`.  ([#42568](https://github.com/Azure/azure-sdk-for-python/pull/42568))
+
 ### Breaking Changes
 
 ### Bugs Fixed
@@ -23,6 +25,7 @@
 
 - `ManagedIdentityCredential` now retries IMDS 410 status responses for at least 70 seconds total duration as required by [Azure IMDS documentation](https://learn.microsoft.com/azure/virtual-machines/instance-metadata-service?tabs=windows#errors-and-debugging).  ([#42330](https://github.com/Azure/azure-sdk-for-python/pull/42330))
 - Improved `DefaultAzureCredential` diagnostics when `WorkloadIdentityCredential` initialization fails. If DAC fails to find a successful credential in the chain, the reason `WorkloadIdentityCredential` failed will be included in the error message. ([#42346](https://github.com/Azure/azure-sdk-for-python/pull/42346))
+- `AzureCliCredential` and `AzurePowerShellCredential` now raise `ClientAuthenticationError` when `claims` are provided to `get_token` or `get_token_info`, as these credentials do not support claims challenges. The error message includes instructions for handling claims authentication scenarios. ([#42568](https://github.com/Azure/azure-sdk-for-python/pull/42568))
 
 ## 1.24.0b1 (2025-07-17)
 

sdk/identity/azure-identity/azure/identity/_credentials/azd_cli.py

@@ -28,6 +28,10 @@
     "Please visit https://aka.ms/azure-dev for installation instructions and then,"
     "once installed, authenticate to your Azure account using 'azd auth login'."
 )
+UNKNOWN_CLAIMS_FLAG = (
+    "Claims challenges are not supported by the Azure Developer CLI version you are using. "
+    "Please update to version 1.18.1 or later."
+)
 COMMAND_LINE = ["auth", "token", "--output", "json", "--no-prompt"]
 EXECUTABLE_NAME = "azd"
 NOT_LOGGED_IN = "Please run 'azd auth login' from a command prompt to authenticate before using this credential."
@@ -99,7 +103,7 @@ def close(self) -> None:
     def get_token(
         self,
         *scopes: str,
-        claims: Optional[str] = None,  # pylint:disable=unused-argument
+        claims: Optional[str] = None,
         tenant_id: Optional[str] = None,
         **kwargs: Any,
     ) -> AccessToken:
@@ -111,7 +115,8 @@ def get_token(
         :param str scopes: desired scope for the access token. This credential allows only one scope per request.
             For more information about scopes, see
             https://learn.microsoft.com/entra/identity-platform/scopes-oidc.
-        :keyword str claims: not used by this credential; any value provided will be ignored.
+        :keyword str claims: additional claims required in the token, such as those returned in a resource provider's
+            claims challenge following an authorization failure.
         :keyword str tenant_id: optional tenant to include in the token request.
 
         :return: An access token with the desired scopes.
@@ -125,6 +130,8 @@ def get_token(
         options: TokenRequestOptions = {}
         if tenant_id:
             options["tenant_id"] = tenant_id
+        if claims:
+            options["claims"] = claims
 
         token_info = self._get_token_base(*scopes, options=options, **kwargs)
         return AccessToken(token_info.token, token_info.expires_on)
@@ -159,6 +166,7 @@ def _get_token_base(
             raise ValueError("Missing scope in request. \n")
 
         tenant_id = options.get("tenant_id") if options else None
+        claims = options.get("claims") if options else None
         if tenant_id:
             validate_tenant_id(tenant_id)
         for scope in scopes:
@@ -175,16 +183,23 @@ def _get_token_base(
         )
         if tenant:
             command_args += ["--tenant-id", tenant]
+        if claims:
+            command_args += ["--claims", claims]
         output = _run_command(command_args, self._process_timeout)
 
         token = parse_token(output)
         if not token:
-            sanitized_output = sanitize_output(output)
-            message = (
-                f"Unexpected output from Azure Developer CLI: '{sanitized_output}'. \n"
-                f"To mitigate this issue, please refer to the troubleshooting guidelines here at "
-                f"https://aka.ms/azsdk/python/identity/azdevclicredential/troubleshoot."
-            )
+            # Try to extract a meaningful error from azd consoleMessage JSON lines
+            extracted = extract_cli_error_message(output)
+            if extracted:
+                message = extracted
+            else:
+                sanitized_output = sanitize_output(output)
+                message = (
+                    f"Unexpected output from Azure Developer CLI: '{sanitized_output}'. \n"
+                    f"To mitigate this issue, please refer to the troubleshooting guidelines here at "
+                    f"https://aka.ms/azsdk/python/identity/azdevclicredential/troubleshoot."
+                )
             if within_dac.get():
                 raise CredentialUnavailableError(message=message)
             raise ClientAuthenticationError(message=message)
@@ -241,6 +256,54 @@ def sanitize_output(output: str) -> str:
     return re.sub(r"\"token\": \"(.*?)(\"|$)", "****", output)
 
 
+def extract_cli_error_message(output: str) -> Optional[str]:
+    """
+    Extract a single, user-friendly message from azd consoleMessage JSON output.
+
+    :param str output: The output from the Azure Developer CLI command.
+    :return: A user-friendly error message if found, otherwise None.
+    :rtype: Optional[str]
+
+    Preference order:
+    1) A message containing "Suggestion" (case-insensitive)
+    2) The second message if multiple are present
+    3) The first message if only one exists
+    Returns None if no messages can be parsed.
+    """
+    messages: List[str] = []
+    for line in output.splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            obj = json.loads(line)
+        except json.JSONDecodeError:  # not JSON -> ignore
+            continue
+        if isinstance(obj, dict):
+            data = obj.get("data")
+            if isinstance(data, dict):
+                msg = data.get("message")
+                if isinstance(msg, str) and msg.strip():
+                    messages.append(msg.strip())
+                    continue
+            msg = obj.get("message")
+            if isinstance(msg, str) and msg.strip():
+                messages.append(msg.strip())
+
+    if not messages:
+        return None
+
+    # Prefer the suggestion line if present
+    for msg in messages:
+        if "suggestion" in msg.lower():
+            return sanitize_output(msg)
+
+    # If more than one message exists, return the last one
+    if len(messages) > 1:
+        return sanitize_output(messages[-1])
+    return sanitize_output(messages[0])
+
+
 def _run_command(command_args: List[str], timeout: int) -> str:
     # Ensure executable exists in PATH first. This avoids a subprocess call that would fail anyway.
     azd_path = shutil.which(EXECUTABLE_NAME)
@@ -267,16 +330,18 @@ def _run_command(command_args: List[str], timeout: int) -> str:
         # Fallback check in case the executable is not found while executing subprocess.
         if ex.returncode == 127 or (ex.stderr is not None and ex.stderr.startswith("'azd' is not recognized")):
             raise CredentialUnavailableError(message=CLI_NOT_FOUND) from ex
-        if ex.stderr is not None and (
-            "not logged in, run `azd auth login` to login" in ex.stderr and "AADSTS" not in ex.stderr
-        ):
+        combined_text = "{}\n{}".format(ex.output or "", ex.stderr or "")
+        if "not logged in, run `azd auth login` to login" in combined_text and "AADSTS" not in combined_text:
             raise CredentialUnavailableError(message=NOT_LOGGED_IN) from ex
+        if "unknown flag: --claims" in combined_text:
+            raise CredentialUnavailableError(message=UNKNOWN_CLAIMS_FLAG) from ex
 
         # return code is from the CLI -> propagate its output
-        if ex.stderr:
-            message = sanitize_output(ex.stderr)
-        else:
-            message = "Failed to invoke Azure Developer CLI"
+        message = (
+            extract_cli_error_message(ex.output or "")
+            or extract_cli_error_message(ex.stderr or "")
+            or (sanitize_output(ex.stderr) if ex.stderr else "Failed to invoke Azure Developer CLI")
+        )
         if within_dac.get():
             raise CredentialUnavailableError(message=message) from ex
         raise ClientAuthenticationError(message=message) from ex

sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py

@@ -34,6 +34,11 @@
 COMMAND_LINE = ["account", "get-access-token", "--output", "json"]
 EXECUTABLE_NAME = "az"
 NOT_LOGGED_IN = "Please run 'az login' to set up an account"
+CLAIMS_UNSUPPORTED_ERROR = (
+    "This credential doesn't support claims challenges. To authenticate with the required "
+    "claims, please run the following command (requires Azure CLI version 2.76.0 or later): "
+    "az login --claims-challenge {claims_value}"
+)
 
 
 class AzureCliCredential:
@@ -90,7 +95,7 @@ def close(self) -> None:
     def get_token(
         self,
         *scopes: str,
-        claims: Optional[str] = None,  # pylint:disable=unused-argument
+        claims: Optional[str] = None,
         tenant_id: Optional[str] = None,
         **kwargs: Any,
     ) -> AccessToken:
@@ -102,20 +107,23 @@ def get_token(
         :param str scopes: desired scope for the access token. This credential allows only one scope per request.
             For more information about scopes, see
             https://learn.microsoft.com/entra/identity-platform/scopes-oidc.
-        :keyword str claims: not used by this credential; any value provided will be ignored.
+        :keyword str claims: additional claims required in the token. This credential does not support claims
+            challenges.
         :keyword str tenant_id: optional tenant to include in the token request.
 
         :return: An access token with the desired scopes.
         :rtype: ~azure.core.credentials.AccessToken
 
-        :raises ~azure.identity.CredentialUnavailableError: the credential was unable to invoke the Azure CLI.
+        :raises ~azure.identity.CredentialUnavailableError: the credential was either unable to invoke the Azure CLI
+          or a claims challenge was provided.
         :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked the Azure CLI but didn't
           receive an access token.
         """
-
         options: TokenRequestOptions = {}
         if tenant_id:
             options["tenant_id"] = tenant_id
+        if claims:
+            options["claims"] = claims
 
         token_info = self._get_token_base(*scopes, options=options, **kwargs)
         return AccessToken(token_info.token, token_info.expires_on)
@@ -136,7 +144,8 @@ def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] =
         :rtype: ~azure.core.credentials.AccessTokenInfo
         :return: An AccessTokenInfo instance containing information about the token.
 
-        :raises ~azure.identity.CredentialUnavailableError: the credential was unable to invoke the Azure CLI.
+        :raises ~azure.identity.CredentialUnavailableError: the credential was either unable to invoke the Azure CLI
+          or a claims challenge was provided.
         :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked the Azure CLI but didn't
           receive an access token.
         """
@@ -145,6 +154,19 @@ def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] =
     def _get_token_base(
         self, *scopes: str, options: Optional[TokenRequestOptions] = None, **kwargs: Any
     ) -> AccessTokenInfo:
+        # Check for claims challenge first
+        if options and options.get("claims"):
+            error_message = CLAIMS_UNSUPPORTED_ERROR.format(claims_value=options.get("claims"))
+
+            # Add tenant if provided in options
+            if options.get("tenant_id"):
+                error_message += f" --tenant {options.get('tenant_id')}"
+
+            # Add scope if provided
+            if scopes:
+                error_message += f" --scope {scopes[0]}"
+
+            raise CredentialUnavailableError(message=error_message)
 
         tenant_id = options.get("tenant_id") if options else None
         if tenant_id:

sdk/identity/azure-identity/azure/identity/_credentials/azure_powershell.py

@@ -13,7 +13,13 @@
 
 from .azure_cli import get_safe_working_dir
 from .. import CredentialUnavailableError
-from .._internal import _scopes_to_resource, resolve_tenant, within_dac, validate_tenant_id, validate_scope
+from .._internal import (
+    _scopes_to_resource,
+    resolve_tenant,
+    within_dac,
+    validate_tenant_id,
+    validate_scope,
+)
 from .._internal.decorators import log_get_token
 
 
@@ -24,6 +30,11 @@
 NO_AZ_ACCOUNT_MODULE = "NO_AZ_ACCOUNT_MODULE"
 POWERSHELL_NOT_INSTALLED = "PowerShell is not installed"
 RUN_CONNECT_AZ_ACCOUNT = 'Please run "Connect-AzAccount" to set up account'
+CLAIMS_UNSUPPORTED_ERROR = (
+    "This credential doesn't support claims challenges. To authenticate with the required "
+    "claims, please run the following command (requires Az.Accounts module version 5.2.0 or later): "
+    "Connect-AzAccount -ClaimsChallenge {claims_value}"
+)
 SCRIPT = """$ErrorActionPreference = 'Stop'
 [version]$minimumVersion = '2.2.0'
 
@@ -112,7 +123,7 @@ def close(self) -> None:
     def get_token(
         self,
         *scopes: str,
-        claims: Optional[str] = None,  # pylint:disable=unused-argument
+        claims: Optional[str] = None,
         tenant_id: Optional[str] = None,
         **kwargs: Any,
     ) -> AccessToken:
@@ -135,10 +146,11 @@ def get_token(
         :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked Azure PowerShell but didn't
           receive an access token
         """
-
         options: TokenRequestOptions = {}
         if tenant_id:
             options["tenant_id"] = tenant_id
+        if claims:
+            options["claims"] = claims
 
         token_info = self._get_token_base(*scopes, options=options, **kwargs)
         return AccessToken(token_info.token, token_info.expires_on)
@@ -170,6 +182,13 @@ def _get_token_base(
         self, *scopes: str, options: Optional[TokenRequestOptions] = None, **kwargs: Any
     ) -> AccessTokenInfo:
 
+        # Check if claims challenge is provided
+        if options and options.get("claims"):
+            error_message = CLAIMS_UNSUPPORTED_ERROR.format(claims_value=options.get("claims"))
+            if options.get("tenant_id"):
+                error_message += f" -Tenant {options.get('tenant_id')}"
+            raise CredentialUnavailableError(message=error_message)
+
         tenant_id = options.get("tenant_id") if options else None
         if tenant_id:
             validate_tenant_id(tenant_id)
@@ -269,7 +288,11 @@ def raise_for_error(return_code: int, stdout: str, stderr: str) -> None:
 
     if stderr:
         # stderr is too noisy to include with an exception but may be useful for debugging
-        _LOGGER.debug('%s received an error from Azure PowerShell: "%s"', AzurePowerShellCredential.__name__, stderr)
+        _LOGGER.debug(
+            '%s received an error from Azure PowerShell: "%s"',
+            AzurePowerShellCredential.__name__,
+            stderr,
+        )
     raise CredentialUnavailableError(
         message="Failed to invoke PowerShell. Enable debug logging for additional information."
     )