[Identity] Add subscription support to Azure CLI cred (Azure#37994)

This allows users to specify a subscription when authenticating with the Azure CLI.

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

Author: pvaneck

sdk/identity/azure-identity/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Added `subscription` parameter to `AzureCliCredential` to specify the subscription to use when authenticating with the Azure CLI. ([#37994](https://github.com/Azure/azure-sdk-for-python/pull/37994))
+
 ### Breaking Changes
 
 ### Bugs Fixed

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

@@ -15,7 +15,14 @@
 from azure.core.exceptions import ClientAuthenticationError
 
 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,
+    validate_subscription,
+)
 from .._internal.decorators import log_get_token
 
 
@@ -31,6 +38,8 @@ class AzureCliCredential:
     This requires previously logging in to Azure via "az login", and will use the CLI's currently logged in identity.
 
     :keyword str tenant_id: Optional tenant to include in the token request.
+    :keyword str subscription: The name or ID of a subscription. Set this to acquire tokens for an account other
+        than the Azure CLI's current account.
     :keyword List[str] additionally_allowed_tenants: Specifies tenants in addition to the specified "tenant_id"
         for which the credential may acquire tokens. Add the wildcard value "*" to allow the credential to
         acquire tokens for any tenant the application can access.
@@ -50,12 +59,17 @@ def __init__(
         self,
         *,
         tenant_id: str = "",
+        subscription: Optional[str] = None,
         additionally_allowed_tenants: Optional[List[str]] = None,
         process_timeout: int = 10,
     ) -> None:
         if tenant_id:
             validate_tenant_id(tenant_id)
+        if subscription:
+            validate_subscription(subscription)
+
         self.tenant_id = tenant_id
+        self.subscription = subscription
         self._additionally_allowed_tenants = additionally_allowed_tenants or []
         self._process_timeout = process_timeout
 
@@ -144,6 +158,9 @@ def _get_token_base(
         )
         if tenant:
             command += " --tenant " + tenant
+
+        if self.subscription:
+            command += f' --subscription "{self.subscription}"'
         output = _run_command(command, self._process_timeout)
 
         token = parse_token(output)

sdk/identity/azure-identity/azure/identity/_internal/__init__.py

@@ -13,6 +13,7 @@
     normalize_authority,
     resolve_tenant,
     validate_scope,
+    validate_subscription,
     validate_tenant_id,
     within_credential_chain,
     within_dac,
@@ -49,6 +50,7 @@ def _scopes_to_resource(*scopes) -> str:
     "normalize_authority",
     "resolve_tenant",
     "validate_scope",
+    "validate_subscription",
     "within_credential_chain",
     "within_dac",
     "wrap_exceptions",

sdk/identity/azure-identity/azure/identity/_internal/utils.py

@@ -20,6 +20,7 @@
 
 VALID_TENANT_ID_CHARACTERS = frozenset(ascii_letters + digits + "-.")
 VALID_SCOPE_CHARACTERS = frozenset(ascii_letters + digits + "_-.:/")
+VALID_SUBSCRIPTION_CHARACTERS = frozenset(ascii_letters + digits + "_-. ")
 
 
 def normalize_authority(authority: str) -> str:
@@ -68,7 +69,20 @@ def validate_tenant_id(tenant_id: str) -> None:
     if not tenant_id or any(c not in VALID_TENANT_ID_CHARACTERS for c in tenant_id):
         raise ValueError(
             "Invalid tenant ID provided. You can locate your tenant ID by following the instructions here: "
-            + "https://learn.microsoft.com/partner-center/find-ids-and-domain-names"
+            "https://learn.microsoft.com/partner-center/find-ids-and-domain-names"
+        )
+
+
+def validate_subscription(subscription: str) -> None:
+    """Raise ValueError if subscription is empty or contains a character invalid for a subscription name/ID.
+
+    :param str subscription: subscription ID to validate
+    :raises: ValueError if subscription is empty or contains a character invalid for a subscription ID.
+    """
+    if not subscription or any(c not in VALID_SUBSCRIPTION_CHARACTERS for c in subscription):
+        raise ValueError(
+            "Invalid subscription provided. You can locate your subscription by following the "
+            "instructions listed here: https://learn.microsoft.com/azure/azure-portal/get-subscription-tenant-id"
         )
 
 

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

@@ -23,7 +23,14 @@
     parse_token,
     sanitize_output,
 )
-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,
+    validate_subscription,
+)
 
 
 class AzureCliCredential(AsyncContextManager):
@@ -32,6 +39,8 @@ class AzureCliCredential(AsyncContextManager):
     This requires previously logging in to Azure via "az login", and will use the CLI's currently logged in identity.
 
     :keyword str tenant_id: Optional tenant to include in the token request.
+    :keyword str subscription: The name or ID of a subscription. Set this to acquire tokens for an account other
+        than the Azure CLI's current account.
     :keyword List[str] additionally_allowed_tenants: Specifies tenants in addition to the specified "tenant_id"
         for which the credential may acquire tokens. Add the wildcard value "*" to allow the credential to
         acquire tokens for any tenant the application can access.
@@ -51,12 +60,17 @@ def __init__(
         self,
         *,
         tenant_id: str = "",
+        subscription: Optional[str] = None,
         additionally_allowed_tenants: Optional[List[str]] = None,
         process_timeout: int = 10,
     ) -> None:
         if tenant_id:
             validate_tenant_id(tenant_id)
+        if subscription:
+            validate_subscription(subscription)
+
         self.tenant_id = tenant_id
+        self.subscription = subscription
         self._additionally_allowed_tenants = additionally_allowed_tenants or []
         self._process_timeout = process_timeout
 
@@ -141,6 +155,8 @@ async def _get_token_base(
 
         if tenant:
             command += " --tenant " + tenant
+        if self.subscription:
+            command += f' --subscription "{self.subscription}"'
         output = await _run_command(command, self._process_timeout)
 
         token = parse_token(output)

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

@@ -13,6 +13,7 @@
 
 FAKE_CLIENT_ID = "fake-client-id"
 INVALID_CHARACTERS = "|\\`;{&' "
+INVALID_SUBSCRIPTION_CHARACTERS = "|\\`;{&'"
 ACCESS_TOKEN_CLASSES = (AccessToken, AccessTokenInfo)
 GET_TOKEN_METHODS = ("get_token", "get_token_info")
 

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

@@ -15,7 +15,7 @@
 import subprocess
 import pytest
 
-from helpers import mock, INVALID_CHARACTERS, GET_TOKEN_METHODS
+from helpers import mock, INVALID_CHARACTERS, INVALID_SUBSCRIPTION_CHARACTERS, GET_TOKEN_METHODS
 
 CHECK_OUTPUT = AzureCliCredential.__module__ + ".subprocess.check_output"
 
@@ -76,6 +76,40 @@ def test_invalid_scopes(get_token_method):
             getattr(AzureCliCredential(), get_token_method)("scope" + c)
 
 
+@pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
+def test_subscription(get_token_method):
+    """The credential should accept a subscription ID"""
+
+    subscription = "foo subscription"
+    credential = AzureCliCredential(subscription=subscription)
+    assert credential.subscription == subscription
+
+    def fake_check_output(command_line, **_):
+        assert f'--subscription "{subscription}"' in command_line[-1]
+        return json.dumps(
+            {
+                "expiresOn": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
+                "accessToken": "access-token",
+                "subscription": subscription,
+                "tenant": "tenant",
+                "tokenType": "Bearer",
+            }
+        )
+
+    with mock.patch("shutil.which", return_value="az"):
+        with mock.patch(CHECK_OUTPUT, fake_check_output):
+            token = getattr(credential, get_token_method)("scope")
+            assert token.token == "access-token"
+
+
+def test_invalid_subscriptons():
+    """Subscriptions with invalid characters should raise ValueErrors."""
+
+    for c in INVALID_SUBSCRIPTION_CHARACTERS:
+        with pytest.raises(ValueError):
+            AzureCliCredential(subscription="subscription" + c)
+
+
 @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
 def test_get_token(get_token_method):
     """The credential should parse the CLI's output to an AccessToken"""

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

@@ -17,7 +17,7 @@
 from azure.core.exceptions import ClientAuthenticationError
 import pytest
 
-from helpers import INVALID_CHARACTERS, GET_TOKEN_METHODS
+from helpers import INVALID_CHARACTERS, INVALID_SUBSCRIPTION_CHARACTERS, GET_TOKEN_METHODS
 from helpers_async import get_completed_future
 from test_cli_credential import TEST_ERROR_OUTPUTS
 
@@ -74,6 +74,41 @@ async def test_invalid_scopes(get_token_method):
             await getattr(AzureCliCredential(), get_token_method)("https://scope" + c)
 
 
+@pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
+async def test_subscription(get_token_method):
+    """The credential should accept a subscription ID"""
+
+    subscription = "foo subscription"
+    credential = AzureCliCredential(subscription=subscription)
+    assert credential.subscription == subscription
+
+    async def fake_exec(*args, **_):
+        assert f'--subscription "{subscription}"' in args[-1]
+        output = json.dumps(
+            {
+                "expiresOn": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
+                "accessToken": "access-token",
+                "subscription": subscription,
+                "tenant": "tenant",
+                "tokenType": "Bearer",
+            }
+        ).encode()
+        return mock.Mock(communicate=mock.Mock(return_value=get_completed_future((output, b""))), returncode=0)
+
+    with mock.patch("shutil.which", return_value="az"):
+        with mock.patch(SUBPROCESS_EXEC, fake_exec):
+            token = await getattr(credential, get_token_method)("scope")
+            assert token.token == "access-token"
+
+
+def test_invalid_subscriptons():
+    """Subscriptions with invalid characters should raise ValueErrors."""
+
+    for c in INVALID_SUBSCRIPTION_CHARACTERS:
+        with pytest.raises(ValueError):
+            AzureCliCredential(subscription="subscription" + c)
+
+
 async def test_close():
     """The credential must define close, although it's a no-op because the credential has no transport"""