Identity token cred env

sdk/identity/azure-identity/CHANGELOG.md

@@ -1,9 +1,13 @@
 # Release History
 
-## 1.22.1 (Unreleased)
+## 1.23.0 (Unreleased)
 
 ### Features Added
 
+- Added `AZURE_TOKEN_CREDENTIALS` environment variable to `DefaultAzureCredential` to allow for choosing groups of credentials.
+  - `prod` for `EnvironmentCredential`, `WorkloadIdentityCredential`,  and `ManagedIdentityCredential`.
+  - `dev` for `SharedTokenCacheCredential`, `AzureCliCredential`, `AzurePowershellCredential`, and `AzureDeveloperCliCredential`.
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/identity/azure-identity/azure/identity/_constants.py

@@ -64,4 +64,5 @@ class EnvironmentVariables:
     AZURE_REGIONAL_AUTHORITY_NAME = "AZURE_REGIONAL_AUTHORITY_NAME"
 
     AZURE_FEDERATED_TOKEN_FILE = "AZURE_FEDERATED_TOKEN_FILE"
+    AZURE_TOKEN_CREDENTIALS = "AZURE_TOKEN_CREDENTIALS"
     WORKLOAD_IDENTITY_VARS = (AZURE_AUTHORITY_HOST, AZURE_TENANT_ID, AZURE_FEDERATED_TOKEN_FILE)

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

@@ -133,6 +133,7 @@ def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statement
 
         process_timeout = kwargs.pop("process_timeout", 10)
 
+        token_credentials_env = os.environ.get(EnvironmentVariables.AZURE_TOKEN_CREDENTIALS, "").strip().lower()
         exclude_workload_identity_credential = kwargs.pop("exclude_workload_identity_credential", False)
         exclude_environment_credential = kwargs.pop("exclude_environment_credential", False)
         exclude_managed_identity_credential = kwargs.pop("exclude_managed_identity_credential", False)
@@ -143,6 +144,26 @@ def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statement
         exclude_interactive_browser_credential = kwargs.pop("exclude_interactive_browser_credential", True)
         exclude_powershell_credential = kwargs.pop("exclude_powershell_credential", False)
 
+        if token_credentials_env == "dev":
+            # In dev mode, use only developer credentials
+            exclude_environment_credential = True
+            exclude_managed_identity_credential = True
+            exclude_workload_identity_credential = True
+        elif token_credentials_env == "prod":
+            # In prod mode, use only production credentials
+            exclude_shared_token_cache_credential = True
+            exclude_visual_studio_code_credential = True
+            exclude_cli_credential = True
+            exclude_developer_cli_credential = True
+            exclude_powershell_credential = True
+            exclude_interactive_browser_credential = True
+        elif token_credentials_env != "":
+            # If the environment variable is set to something other than dev or prod, raise an error
+            raise ValueError(
+                f"Invalid value for {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS}: {token_credentials_env}. "
+                "Valid values are 'dev' or 'prod'."
+            )
+
         credentials: List[SupportsTokenInfo] = []
         within_dac.set(True)
         if not exclude_environment_credential:
@@ -155,15 +176,15 @@ def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statement
                         client_id=cast(str, client_id),
                         tenant_id=workload_identity_tenant_id,
                         token_file_path=os.environ[EnvironmentVariables.AZURE_FEDERATED_TOKEN_FILE],
-                        **kwargs
+                        **kwargs,
                     )
                 )
         if not exclude_managed_identity_credential:
             credentials.append(
                 ManagedIdentityCredential(
                     client_id=managed_identity_client_id,
                     _exclude_workload_identity_credential=exclude_workload_identity_credential,
-                    **kwargs
+                    **kwargs,
                 )
             )
         if not exclude_shared_token_cache_credential and SharedTokenCacheCredential.supported():

sdk/identity/azure-identity/azure/identity/_version.py

@@ -2,4 +2,4 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-VERSION = "1.22.1"
+VERSION = "1.23.0"

sdk/identity/azure-identity/azure/identity/aio/_credentials/default.py

@@ -125,6 +125,7 @@ def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statement
 
         process_timeout = kwargs.pop("process_timeout", 10)
 
+        token_credentials_env = os.environ.get(EnvironmentVariables.AZURE_TOKEN_CREDENTIALS, "").strip().lower()
         exclude_workload_identity_credential = kwargs.pop("exclude_workload_identity_credential", False)
         exclude_visual_studio_code_credential = kwargs.pop("exclude_visual_studio_code_credential", True)
         exclude_developer_cli_credential = kwargs.pop("exclude_developer_cli_credential", False)
@@ -134,6 +135,25 @@ def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statement
         exclude_shared_token_cache_credential = kwargs.pop("exclude_shared_token_cache_credential", False)
         exclude_powershell_credential = kwargs.pop("exclude_powershell_credential", False)
 
+        if token_credentials_env == "dev":
+            # In dev mode, use only developer credentials
+            exclude_environment_credential = True
+            exclude_managed_identity_credential = True
+            exclude_workload_identity_credential = True
+        elif token_credentials_env == "prod":
+            # In prod mode, use only production credentials
+            exclude_shared_token_cache_credential = True
+            exclude_visual_studio_code_credential = True
+            exclude_cli_credential = True
+            exclude_developer_cli_credential = True
+            exclude_powershell_credential = True
+        elif token_credentials_env != "":
+            # If the environment variable is set to something other than dev or prod, raise an error
+            raise ValueError(
+                f"Invalid value for {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS}: {token_credentials_env}. "
+                "Valid values are 'dev' or 'prod'."
+            )
+
         credentials: List[AsyncSupportsTokenInfo] = []
         within_dac.set(True)
         if not exclude_environment_credential:
@@ -146,15 +166,15 @@ def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statement
                         client_id=cast(str, client_id),
                         tenant_id=workload_identity_tenant_id,
                         token_file_path=os.environ[EnvironmentVariables.AZURE_FEDERATED_TOKEN_FILE],
-                        **kwargs
+                        **kwargs,
                     )
                 )
         if not exclude_managed_identity_credential:
             credentials.append(
                 ManagedIdentityCredential(
                     client_id=managed_identity_client_id,
                     _exclude_workload_identity_credential=exclude_workload_identity_credential,
-                    **kwargs
+                    **kwargs,
                 )
             )
         if not exclude_shared_token_cache_credential and SharedTokenCacheCredential.supported():

sdk/identity/azure-identity/tests/test_token_credentials_env.py

@@ -0,0 +1,118 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import os
+from unittest.mock import patch
+
+import pytest
+
+from azure.identity import (
+    AzureCliCredential,
+    AzureDeveloperCliCredential,
+    AzurePowerShellCredential,
+    DefaultAzureCredential,
+    EnvironmentCredential,
+    ManagedIdentityCredential,
+    SharedTokenCacheCredential,
+    WorkloadIdentityCredential,
+)
+from azure.identity._constants import EnvironmentVariables
+
+
+def test_token_credentials_env_dev():
+    """With AZURE_TOKEN_CREDENTIALS=dev, DefaultAzureCredential should use only developer credentials"""
+
+    prod_credentials = {EnvironmentCredential, WorkloadIdentityCredential, ManagedIdentityCredential}
+
+    with patch.dict("os.environ", {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS: "dev"}, clear=False):
+        credential = DefaultAzureCredential()
+
+        # Get the actual credential classes in the chain
+        actual_classes = {c.__class__ for c in credential.credentials}
+
+        # All dev credentials should be present (if supported)
+        if SharedTokenCacheCredential.supported():
+            assert SharedTokenCacheCredential in actual_classes
+
+        # Other developer credentials should be present
+        assert AzureCliCredential in actual_classes
+        assert AzureDeveloperCliCredential in actual_classes
+        assert AzurePowerShellCredential in actual_classes
+
+        # Production credentials should NOT be present
+        for cred_class in prod_credentials:
+            if cred_class == WorkloadIdentityCredential:
+                # Skip this check unless env vars are set
+                if not all(os.environ.get(var) for var in EnvironmentVariables.WORKLOAD_IDENTITY_VARS):
+                    continue
+            assert cred_class not in actual_classes
+
+
+def test_token_credentials_env_prod():
+    """With AZURE_TOKEN_CREDENTIALS=prod, DefaultAzureCredential should use only production credentials"""
+
+    dev_credentials = {
+        SharedTokenCacheCredential,
+        AzureCliCredential,
+        AzureDeveloperCliCredential,
+        AzurePowerShellCredential,
+    }
+
+    with patch.dict("os.environ", {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS: "prod"}, clear=False):
+        # Print to verify the environment variable is set in the test
+        print(f"AZURE_TOKEN_CREDENTIALS={os.environ.get(EnvironmentVariables.AZURE_TOKEN_CREDENTIALS)}")
+
+        credential = DefaultAzureCredential()
+
+        # Get the actual credential classes in the chain
+        actual_classes = {c.__class__ for c in credential.credentials}
+
+        # Print which credentials are actually in the chain
+        print("Credentials in chain:")
+        for cls in actual_classes:
+            print(f" - {cls.__name__}")
+
+        # Production credentials should be present
+        assert EnvironmentCredential in actual_classes
+        assert ManagedIdentityCredential in actual_classes
+
+        # Check WorkloadIdentityCredential only if env vars are set
+        if all(os.environ.get(var) for var in EnvironmentVariables.WORKLOAD_IDENTITY_VARS):
+            assert WorkloadIdentityCredential in actual_classes
+
+        # Developer credentials should NOT be present
+        for cred_class in dev_credentials:
+            assert cred_class not in actual_classes
+
+
+def test_token_credentials_env_case_insensitive():
+    """AZURE_TOKEN_CREDENTIALS should be case insensitive"""
+
+    with patch.dict("os.environ", {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS: "DeV"}, clear=False):
+        credential = DefaultAzureCredential()
+
+        # Get the actual credential classes in the chain
+        actual_classes = {c.__class__ for c in credential.credentials}
+
+        # EnvironmentCredential (prod) should not be present
+        assert EnvironmentCredential not in actual_classes
+
+        # AzureCliCredential (dev) should be present
+        assert AzureCliCredential in actual_classes
+
+
+def test_token_credentials_env_invalid():
+    """Invalid AZURE_TOKEN_CREDENTIALS value should raise an error"""
+
+    with patch.dict("os.environ", {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS: "invalid"}, clear=False):
+        with pytest.raises(ValueError):
+            credential = DefaultAzureCredential()
+
+
+def test_token_credentials_env_with_exclude():
+    with patch.dict("os.environ", {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS: "prod"}, clear=False):
+        credential = DefaultAzureCredential(exclude_environment_credential=True)
+        actual_classes = {c.__class__ for c in credential.credentials}
+
+        assert EnvironmentCredential not in actual_classes

sdk/identity/azure-identity/tests/test_token_credentials_env_async.py

@@ -0,0 +1,118 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import os
+from unittest.mock import patch
+
+import pytest
+
+from azure.identity.aio import (
+    AzureCliCredential,
+    AzureDeveloperCliCredential,
+    AzurePowerShellCredential,
+    DefaultAzureCredential,
+    EnvironmentCredential,
+    ManagedIdentityCredential,
+    SharedTokenCacheCredential,
+    WorkloadIdentityCredential,
+)
+from azure.identity._constants import EnvironmentVariables
+
+
+def test_token_credentials_env_dev():
+    """With AZURE_TOKEN_CREDENTIALS=dev, DefaultAzureCredential should use only developer credentials"""
+
+    prod_credentials = {EnvironmentCredential, WorkloadIdentityCredential, ManagedIdentityCredential}
+
+    with patch.dict("os.environ", {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS: "dev"}, clear=False):
+        credential = DefaultAzureCredential()
+
+        # Get the actual credential classes in the chain
+        actual_classes = {c.__class__ for c in credential.credentials}
+
+        # All dev credentials should be present (if supported)
+        if SharedTokenCacheCredential.supported():
+            assert SharedTokenCacheCredential in actual_classes
+
+        # Other developer credentials should be present
+        assert AzureCliCredential in actual_classes
+        assert AzureDeveloperCliCredential in actual_classes
+        assert AzurePowerShellCredential in actual_classes
+
+        # Production credentials should NOT be present
+        for cred_class in prod_credentials:
+            if cred_class == WorkloadIdentityCredential:
+                # Skip this check unless env vars are set
+                if not all(os.environ.get(var) for var in EnvironmentVariables.WORKLOAD_IDENTITY_VARS):
+                    continue
+            assert cred_class not in actual_classes
+
+
+def test_token_credentials_env_prod():
+    """With AZURE_TOKEN_CREDENTIALS=prod, DefaultAzureCredential should use only production credentials"""
+
+    dev_credentials = {
+        SharedTokenCacheCredential,
+        AzureCliCredential,
+        AzureDeveloperCliCredential,
+        AzurePowerShellCredential,
+    }
+
+    with patch.dict("os.environ", {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS: "prod"}, clear=False):
+        # Print to verify the environment variable is set in the test
+        print(f"AZURE_TOKEN_CREDENTIALS={os.environ.get(EnvironmentVariables.AZURE_TOKEN_CREDENTIALS)}")
+
+        credential = DefaultAzureCredential()
+
+        # Get the actual credential classes in the chain
+        actual_classes = {c.__class__ for c in credential.credentials}
+
+        # Print which credentials are actually in the chain
+        print("Credentials in chain:")
+        for cls in actual_classes:
+            print(f" - {cls.__name__}")
+
+        # Production credentials should be present
+        assert EnvironmentCredential in actual_classes
+        assert ManagedIdentityCredential in actual_classes
+
+        # Check WorkloadIdentityCredential only if env vars are set
+        if all(os.environ.get(var) for var in EnvironmentVariables.WORKLOAD_IDENTITY_VARS):
+            assert WorkloadIdentityCredential in actual_classes
+
+        # Developer credentials should NOT be present
+        for cred_class in dev_credentials:
+            assert cred_class not in actual_classes
+
+
+def test_token_credentials_env_case_insensitive():
+    """AZURE_TOKEN_CREDENTIALS should be case insensitive"""
+
+    with patch.dict("os.environ", {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS: "DeV"}, clear=False):
+        credential = DefaultAzureCredential()
+
+        # Get the actual credential classes in the chain
+        actual_classes = {c.__class__ for c in credential.credentials}
+
+        # EnvironmentCredential (prod) should not be present
+        assert EnvironmentCredential not in actual_classes
+
+        # AzureCliCredential (dev) should be present
+        assert AzureCliCredential in actual_classes
+
+
+def test_token_credentials_env_invalid():
+    """Invalid AZURE_TOKEN_CREDENTIALS value should raise an error"""
+
+    with patch.dict("os.environ", {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS: "invalid"}, clear=False):
+        with pytest.raises(ValueError):
+            credential = DefaultAzureCredential()
+
+
+def test_token_credentials_env_with_exclude():
+    with patch.dict("os.environ", {EnvironmentVariables.AZURE_TOKEN_CREDENTIALS: "prod"}, clear=False):
+        credential = DefaultAzureCredential(exclude_environment_credential=True)
+        actual_classes = {c.__class__ for c in credential.credentials}
+
+        assert EnvironmentCredential not in actual_classes