mi-arc

Author: jiasli

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

@@ -216,48 +216,56 @@ def login(self,
         return deepcopy(consolidated)
 
     def login_with_managed_identity(self, identity_id=None, allow_no_subscriptions=None):
-        import jwt
-        from azure.mgmt.core.tools import is_valid_resource_id
-        from azure.cli.core.auth.adal_authentication import MSIAuthenticationWrapper
-        resource = self.cli_ctx.cloud.endpoints.active_directory_resource_id
-
-        if identity_id:
-            if is_valid_resource_id(identity_id):
-                msi_creds = MSIAuthenticationWrapper(resource=resource, msi_res_id=identity_id)
-                identity_type = MsiAccountTypes.user_assigned_resource_id
-            else:
-                authenticated = False
-                from azure.cli.core.azclierror import AzureResponseError
-                try:
-                    msi_creds = MSIAuthenticationWrapper(resource=resource, client_id=identity_id)
-                    identity_type = MsiAccountTypes.user_assigned_client_id
-                    authenticated = True
-                except AzureResponseError as ex:
-                    if 'http error: 400, reason: Bad Request' in ex.error_msg:
-                        logger.info('Sniff: not an MSI client id')
-                    else:
-                        raise
-
-                if not authenticated:
+        identity_type = None
+        if _on_azure_arc():
+            identity_type = MsiAccountTypes.system_assigned
+            from .auth.msal_authentication import ManagedIdentityCredential
+            cred = ManagedIdentityCredential()
+            token = cred.get_token(self._arm_scope).token
+        else:
+            import jwt
+            from azure.mgmt.core.tools import is_valid_resource_id
+            from azure.cli.core.auth.adal_authentication import MSIAuthenticationWrapper
+            resource = self.cli_ctx.cloud.endpoints.active_directory_resource_id
+
+            if identity_id:
+                if is_valid_resource_id(identity_id):
+                    msi_creds = MSIAuthenticationWrapper(resource=resource, msi_res_id=identity_id)
+                    identity_type = MsiAccountTypes.user_assigned_resource_id
+                else:
+                    authenticated = False
+                    from azure.cli.core.azclierror import AzureResponseError
                     try:
-                        identity_type = MsiAccountTypes.user_assigned_object_id
-                        msi_creds = MSIAuthenticationWrapper(resource=resource, object_id=identity_id)
+                        msi_creds = MSIAuthenticationWrapper(resource=resource, client_id=identity_id)
+                        identity_type = MsiAccountTypes.user_assigned_client_id
                         authenticated = True
                     except AzureResponseError as ex:
                         if 'http error: 400, reason: Bad Request' in ex.error_msg:
-                            logger.info('Sniff: not an MSI object id')
+                            logger.info('Sniff: not an MSI client id')
                         else:
                             raise
 
-                if not authenticated:
-                    raise CLIError('Failed to connect to MSI, check your managed service identity id.')
+                    if not authenticated:
+                        try:
+                            identity_type = MsiAccountTypes.user_assigned_object_id
+                            msi_creds = MSIAuthenticationWrapper(resource=resource, object_id=identity_id)
+                            authenticated = True
+                        except AzureResponseError as ex:
+                            if 'http error: 400, reason: Bad Request' in ex.error_msg:
+                                logger.info('Sniff: not an MSI object id')
+                            else:
+                                raise
 
-        else:
-            identity_type = MsiAccountTypes.system_assigned
-            msi_creds = MSIAuthenticationWrapper(resource=resource)
+                    if not authenticated:
+                        raise CLIError('Failed to connect to MSI, check your managed service identity id.')
+
+            else:
+                identity_type = MsiAccountTypes.system_assigned
+                msi_creds = MSIAuthenticationWrapper(resource=resource)
+
+            token_entry = msi_creds.token
+            token = token_entry['access_token']
 
-        token_entry = msi_creds.token
-        token = token_entry['access_token']
         logger.info('MSI: token was retrieved. Now trying to initialize local accounts...')
         decode = jwt.decode(token, algorithms=['RS256'], options={"verify_signature": False})
         tenant = decode['tid']
@@ -364,7 +372,11 @@ def get_login_credentials(self, resource=None, client_id=None, subscription_id=N
                                      resource=resource)
         else:
             # managed identity
-            cred = MsiAccountTypes.msi_auth_factory(managed_identity_type, managed_identity_id, resource)
+            if _on_azure_arc():
+                from .auth.msal_authentication import ManagedIdentityCredential
+                cred = ManagedIdentityCredential()
+            else:
+                cred = MsiAccountTypes.msi_auth_factory(managed_identity_type, managed_identity_id, resource)
         return (cred,
                 str(account[_SUBSCRIPTION_ID]),
                 str(account[_TENANT_ID]))
@@ -390,9 +402,12 @@ def get_raw_token(self, resource=None, scopes=None, subscription=None, tenant=No
             if tenant:
                 raise CLIError("Tenant shouldn't be specified for managed identity account")
             from .auth.util import scopes_to_resource
-            msi_creds = MsiAccountTypes.msi_auth_factory(identity_type, identity_id,
-                                                         scopes_to_resource(scopes))
-            sdk_token = msi_creds.get_token(*scopes)
+            if _on_azure_arc():
+                from .auth.msal_authentication import ManagedIdentityCredential
+                cred = ManagedIdentityCredential()
+            else:
+                cred = MsiAccountTypes.msi_auth_factory(identity_type, identity_id, scopes_to_resource(scopes))
+            sdk_token = cred.get_token(*scopes)
         elif in_cloud_console() and account[_USER_ENTITY].get(_CLOUD_SHELL_ID):
             # Cloud Shell, which is just a system-assigned managed identity.
             if tenant:
@@ -900,3 +915,7 @@ def _create_identity_instance(cli_ctx, *args, **kwargs):
     return Identity(*args, encrypt=encrypt, use_msal_http_cache=use_msal_http_cache,
                     enable_broker_on_windows=enable_broker_on_windows,
                     instance_discovery=instance_discovery, **kwargs)
+
+
+def _on_azure_arc():
+    return "IDENTITY_ENDPOINT" in os.environ and "IMDS_ENDPOINT" in os.environ

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

@@ -18,7 +18,8 @@
 
 from knack.log import get_logger
 from knack.util import CLIError
-from msal import PublicClientApplication, ConfidentialClientApplication
+from msal import (PublicClientApplication, ConfidentialClientApplication,
+                  ManagedIdentityClient, SystemAssignedManagedIdentity)
 
 from .util import check_result, build_sdk_access_token
 
@@ -138,3 +139,20 @@ def get_token(self, *scopes, **kwargs):
         result = self.acquire_token_for_client(scopes, **kwargs)
         check_result(result)
         return build_sdk_access_token(result)
+
+
+class ManagedIdentityCredential:
+    """Currently, only Azure Arc's system-assigned managed identity is supported.
+    """
+
+    def __init__(self):
+        import requests
+        self._msal_client = ManagedIdentityClient(SystemAssignedManagedIdentity(), http_client=requests.Session())
+
+    def get_token(self, *scopes, **kwargs):
+        logger.debug("ManagedIdentityCredential.get_token: scopes=%r, kwargs=%r", scopes, kwargs)
+
+        from .util import scopes_to_resource
+        result = self._msal_client.acquire_token_for_client(resource=scopes_to_resource(scopes))
+        check_result(result)
+        return build_sdk_access_token(result)