mi-msal

Author: jiasli

src/azure-cli-core/azure/cli/core/_profile.py

@@ -227,9 +227,10 @@ def login(self,
 
     def login_with_managed_identity(self, identity_id=None, client_id=None, object_id=None, resource_id=None,
                                     allow_no_subscriptions=None):
-        if _on_azure_arc():
-            return self.login_with_managed_identity_azure_arc(
-                identity_id=identity_id, allow_no_subscriptions=allow_no_subscriptions)
+        if _use_msal_managed_identity(self.cli_ctx):
+            return self.login_with_managed_identity_msal(
+                client_id=client_id, object_id=object_id, resource_id=resource_id,
+                allow_no_subscriptions=allow_no_subscriptions)
 
         import jwt
         from azure.mgmt.core.tools import is_valid_resource_id
@@ -310,13 +311,31 @@ def login_with_managed_identity(self, identity_id=None, client_id=None, object_i
         self._set_subscriptions(consolidated)
         return deepcopy(consolidated)
 
-    def login_with_managed_identity_azure_arc(self, identity_id=None, allow_no_subscriptions=None):
+    def login_with_managed_identity_msal(self, client_id=None, object_id=None, resource_id=None,
+                                         allow_no_subscriptions=None):
         import jwt
-        identity_type = MsiAccountTypes.system_assigned
-        from .auth.msal_credentials import ManagedIdentityCredential
         from .auth.constants import ACCESS_TOKEN
 
-        cred = ManagedIdentityCredential()
+        id_arg_count = len([arg for arg in (client_id, object_id, resource_id) if arg])
+        if id_arg_count > 1:
+            raise CLIError('Usage error: Provide only one of --client-id, --object-id, --resource-id.')
+
+        identity_type = None
+        identity_id = None
+        if id_arg_count == 0:
+            identity_type = MsiAccountTypes.system_assigned
+            identity_id = None
+        elif client_id:
+            identity_type = MsiAccountTypes.user_assigned_client_id
+            identity_id = client_id
+        elif object_id:
+            identity_type = MsiAccountTypes.user_assigned_object_id
+            identity_id = object_id
+        elif resource_id:
+            identity_type = MsiAccountTypes.user_assigned_resource_id
+            identity_id = resource_id
+
+        cred = MsiAccountTypes.msal_credential_factory(identity_type, identity_id)
         token = cred.acquire_token(self._arm_scope)[ACCESS_TOKEN]
         logger.info('Managed identity: token was retrieved. Now trying to initialize local accounts...')
         decode = jwt.decode(token, algorithms=['RS256'], options={"verify_signature": False})
@@ -405,10 +424,10 @@ def get_login_credentials(self, subscription_id=None, aux_subscriptions=None, au
 
         elif managed_identity_type:
             # managed identity
-            if _on_azure_arc():
-                from .auth.msal_credentials import ManagedIdentityCredential
+            if _use_msal_managed_identity(self.cli_ctx):
                 # The credential must be wrapped by CredentialAdaptor so that it can work with Track 1 SDKs.
-                sdk_cred = CredentialAdaptor(ManagedIdentityCredential())
+                cred = MsiAccountTypes.msal_credential_factory(managed_identity_type, managed_identity_id)
+                sdk_cred = CredentialAdaptor(cred)
             else:
                 # The resource is merely used by msrestazure to get the first access token.
                 # It is not actually used in an API invocation.
@@ -466,9 +485,9 @@ def get_raw_token(self, resource=None, scopes=None, subscription=None, tenant=No
             # managed identity
             if tenant:
                 raise CLIError("Tenant shouldn't be specified for managed identity account")
-            if _on_azure_arc():
-                from .auth.msal_credentials import ManagedIdentityCredential
-                sdk_cred = CredentialAdaptor(ManagedIdentityCredential())
+            if _use_msal_managed_identity(self.cli_ctx):
+                cred = MsiAccountTypes.msal_credential_factory(managed_identity_type, managed_identity_id)
+                sdk_cred = CredentialAdaptor(cred)
             else:
                 from .auth.util import scopes_to_resource
                 sdk_cred = MsiAccountTypes.msi_auth_factory(managed_identity_type, managed_identity_id,
@@ -816,6 +835,18 @@ def msi_auth_factory(cli_account_name, identity, resource):
             return MSIAuthenticationWrapper(resource=resource, msi_res_id=identity)
         raise ValueError("unrecognized msi account name '{}'".format(cli_account_name))
 
+    @staticmethod
+    def msal_credential_factory(id_type, id_value):
+        from azure.cli.core.auth.msal_credentials import ManagedIdentityCredential
+        if id_type == MsiAccountTypes.system_assigned:
+            return ManagedIdentityCredential()
+        if id_type == MsiAccountTypes.user_assigned_client_id:
+            return ManagedIdentityCredential(client_id=id_value)
+        if id_type == MsiAccountTypes.user_assigned_object_id:
+            return ManagedIdentityCredential(object_id=id_value)
+        if id_type == MsiAccountTypes.user_assigned_resource_id:
+            return ManagedIdentityCredential(resource_id=id_value)
+        raise ValueError("Unrecognized managed identity ID type '{}'".format(id_type))
 
 class SubscriptionFinder:
     # An ARM client. It finds subscriptions for a user or service principal. It shouldn't do any
@@ -982,7 +1013,9 @@ def _create_identity_instance(cli_ctx, authority, tenant_id=None, client_id=None
                     instance_discovery=instance_discovery)
 
 
-def _on_azure_arc():
+def _use_msal_managed_identity(cli_ctx):
     # This indicates an Azure Arc-enabled server
     from msal.managed_identity import get_managed_identity_source, AZURE_ARC
-    return get_managed_identity_source() == AZURE_ARC
+    # PREVIEW: Use core.use_msal_managed_identity=true to enable managed identity authentication with MSAL
+    use_msal_managed_identity = cli_ctx.config.getboolean('core', 'use_msal_managed_identity', fallback=False)
+    return use_msal_managed_identity or get_managed_identity_source() == AZURE_ARC

src/azure-cli-core/azure/cli/core/auth/msal_credentials.py

@@ -10,7 +10,7 @@
 from knack.log import get_logger
 from knack.util import CLIError
 from msal import (PublicClientApplication, ConfidentialClientApplication,
-                  ManagedIdentityClient, SystemAssignedManagedIdentity)
+                  ManagedIdentityClient, SystemAssignedManagedIdentity, UserAssignedManagedIdentity)
 
 from .constants import AZURE_CLI_CLIENT_ID
 from .util import check_result
@@ -131,9 +131,14 @@ class ManagedIdentityCredential:  # pylint: disable=too-few-public-methods
     Currently, only Azure Arc's system-assigned managed identity is supported.
     """
 
-    def __init__(self):
+    def __init__(self, client_id=None, resource_id=None, object_id=None):
         import requests
-        self._msal_client = ManagedIdentityClient(SystemAssignedManagedIdentity(), http_client=requests.Session())
+        if client_id or resource_id or object_id:
+            managed_identity = UserAssignedManagedIdentity(
+                client_id=client_id, resource_id=resource_id, object_id=object_id)
+        else:
+            managed_identity = SystemAssignedManagedIdentity()
+        self._msal_client = ManagedIdentityClient(managed_identity, http_client=requests.Session())
 
     def acquire_token(self, scopes, **kwargs):
         logger.debug("ManagedIdentityCredential.acquire_token: scopes=%r, kwargs=%r", scopes, kwargs)