{Identity} Support Cloud Shell login

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

@@ -53,6 +53,16 @@ def _delete_file(file_path):
 class Identity:
     """Class to interact with Azure Identity.
     """
+    MANAGED_IDENTITY_TENANT_ID = "tenant_id"
+    MANAGED_IDENTITY_CLIENT_ID = "client_id"
+    MANAGED_IDENTITY_OBJECT_ID = "object_id"
+    MANAGED_IDENTITY_RESOURCE_ID = "resource_id"
+    MANAGED_IDENTITY_SYSTEM_ASSIGNED = 'systemAssignedIdentity'
+    MANAGED_IDENTITY_USER_ASSIGNED = 'userAssignedIdentity'
+    MANAGED_IDENTITY_TYPE = 'type'
+    MANAGED_IDENTITY_ID_TYPE = "id_type"
+
+    CLOUD_SHELL_IDENTITY_UNIQUE_NAME = "unique_name"
 
     def __init__(self, authority=None, tenant_id=None, client_id=_CLIENT_ID, **kwargs):
         self.authority = authority
@@ -62,16 +72,6 @@ def __init__(self, authority=None, tenant_id=None, client_id=_CLIENT_ID, **kwarg
         # todo: MSAL support force encryption
         self.allow_unencrypted = True
 
-        # Initialize _msal_app for logout, since Azure Identity doesn't provide the functionality for logout
-        from msal import PublicClientApplication
-        # sdk/identity/azure-identity/azure/identity/_internal/msal_credentials.py:95
-        from azure.identity._internal.persistent_cache import load_persistent_cache
-
-        # Store for user token persistence
-        cache = load_persistent_cache(self.allow_unencrypted)
-        authority = "https://{}/organizations".format(self.authority)
-        self._msal_app = PublicClientApplication(authority=authority, client_id=_CLIENT_ID, token_cache=cache)
-
         # Store for Service principal credential persistence
         self._msal_store = MSALSecretStore(fallback_to_plaintext=self.allow_unencrypted)
 
@@ -83,6 +83,17 @@ def __init__(self, authority=None, tenant_id=None, client_id=_CLIENT_ID, **kwarg
         from azure.cli.core._debug import change_ssl_cert_verification_track2
         self.ssl_kwargs = change_ssl_cert_verification_track2()
 
+    @property
+    def _msal_app(self):
+        # Initialize _msal_app for logout, since Azure Identity doesn't provide the functionality for logout
+        from msal import PublicClientApplication
+        # sdk/identity/azure-identity/azure/identity/_internal/msal_credentials.py:95
+        from azure.identity._internal.persistent_cache import load_persistent_cache
+
+        # Store for user token persistence
+        cache = load_persistent_cache(self.allow_unencrypted)
+        return PublicClientApplication(authority=self.authority, client_id=self.client_id, token_cache=cache)
+
     def login_with_interactive_browser(self):
         # Use InteractiveBrowserCredential
         credential = InteractiveBrowserCredential(authority=self.authority,
@@ -151,16 +162,7 @@ def login_with_service_principal_certificate(self, client_id, certificate_path):
         credential = CertificateCredential(self.tenant_id, client_id, certificate_path, authority=self.authority)
         return credential
 
-    MANAGED_IDENTITY_TENANT_ID = "tenant_id"
-    MANAGED_IDENTITY_CLIENT_ID = "client_id"
-    MANAGED_IDENTITY_OBJECT_ID = "object_id"
-    MANAGED_IDENTITY_RESOURCE_ID = "resource_id"
-    MANAGED_IDENTITY_SYSTEM_ASSIGNED = 'systemAssignedIdentity'
-    MANAGED_IDENTITY_USER_ASSIGNED = 'userAssignedIdentity'
-    MANAGED_IDENTITY_TYPE = 'type'
-    MANAGED_IDENTITY_ID_TYPE = "id_type"
-
-    def login_with_managed_identity(self, identity_id, resource):
+    def login_with_managed_identity(self, resource, identity_id=None):
         from msrestazure.tools import is_valid_resource_id
         from requests import HTTPError
 
@@ -169,7 +171,8 @@ def login_with_managed_identity(self, identity_id, resource):
         if identity_id:
             # Try resource ID
             if is_valid_resource_id(identity_id):
-                credential = ManagedIdentityCredential(resource=resource, msi_res_id=identity_id)
+                # TODO: Support resource ID in Azure Identity
+                credential = ManagedIdentityCredential(resource_id=identity_id)
                 id_type = self.MANAGED_IDENTITY_RESOURCE_ID
             else:
                 authenticated = False
@@ -187,7 +190,8 @@ def login_with_managed_identity(self, identity_id, resource):
                 if not authenticated:
                     try:
                         # Try object ID
-                        credential = ManagedIdentityCredential(resource=resource, object_id=identity_id)
+                        # TODO: Support resource ID in Azure Identity
+                        credential = ManagedIdentityCredential(object_id=identity_id)
                         id_type = self.MANAGED_IDENTITY_OBJECT_ID
                         authenticated = True
                     except HTTPError as ex:
@@ -202,33 +206,54 @@ def login_with_managed_identity(self, identity_id, resource):
         else:
             credential = ManagedIdentityCredential()
 
-        # As Managed Identity doesn't have ID token, we need to get an initial access token and extract info from it
-        # The resource is only used for acquiring the initial access token
-        scope = resource.rstrip('/') + '/.default'
-        token = credential.get_token(scope)
-        from msal.oauth2cli.oidc import decode_part
-        access_token = token.token
-
-        # Access token consists of headers.claims.signature. Decode the claim part
-        decoded_str = decode_part(access_token.split('.')[1])
-        logger.debug('MSI token retrieved: %s', decoded_str)
-        decoded = json.loads(decoded_str)
+        decoded = self._decode_managed_identity_token(credential, resource)
+        resource_id = decoded.get('xms_mirid')
+        # User-assigned identity has resourceID as
+        # /subscriptions/xxx/resourcegroups/xxx/providers/Microsoft.ManagedIdentity/userAssignedIdentities/xxx
+        if resource_id and 'Microsoft.ManagedIdentity' in resource_id:
+            mi_type = self.MANAGED_IDENTITY_USER_ASSIGNED
+        else:
+            mi_type = self.MANAGED_IDENTITY_SYSTEM_ASSIGNED
 
-        resource_id = decoded['xms_mirid']
         managed_identity_info = {
-            self.MANAGED_IDENTITY_TYPE: self.MANAGED_IDENTITY_USER_ASSIGNED
-            if 'Microsoft.ManagedIdentity' in resource_id else self.MANAGED_IDENTITY_SYSTEM_ASSIGNED,
+            self.MANAGED_IDENTITY_TYPE: mi_type,
             # The type of the ID provided with --username, only valid for a user-assigned managed identity
             self.MANAGED_IDENTITY_ID_TYPE: id_type,
             self.MANAGED_IDENTITY_TENANT_ID: decoded['tid'],
             self.MANAGED_IDENTITY_CLIENT_ID: decoded['appid'],
             self.MANAGED_IDENTITY_OBJECT_ID: decoded['oid'],
-            self.MANAGED_IDENTITY_RESOURCE_ID: resource_id
+            self.MANAGED_IDENTITY_RESOURCE_ID: resource_id,
         }
         logger.warning('Using Managed Identity: %s', json.dumps(managed_identity_info))
 
         return credential, managed_identity_info
 
+    def login_in_cloud_shell(self, resource):
+        credential = ManagedIdentityCredential()
+        decoded = self._decode_managed_identity_token(credential, resource)
+
+        cloud_shell_identity_info = {
+            self.MANAGED_IDENTITY_TENANT_ID: decoded['tid'],
+            # For getting the user email in Cloud Shell, maybe 'email' can also be used
+            self.CLOUD_SHELL_IDENTITY_UNIQUE_NAME: decoded.get('unique_name', 'N/A')
+        }
+        logger.warning('Using Cloud Shell Managed Identity: %s', json.dumps(cloud_shell_identity_info))
+        return credential, cloud_shell_identity_info
+
+    def _decode_managed_identity_token(self, credential, resource):
+        # As Managed Identity doesn't have ID token, we need to get an initial access token and extract info from it
+        # The resource is only used for acquiring the initial access token
+        scope = resource.rstrip('/') + '/.default'
+        token = credential.get_token(scope)
+        from msal.oauth2cli.oidc import decode_part
+        access_token = token.token
+
+        # Access token consists of headers.claims.signature. Decode the claim part
+        decoded_str = decode_part(access_token.split('.')[1])
+        logger.debug('MSI token retrieved: %s', decoded_str)
+        decoded = json.loads(decoded_str)
+        return decoded
+
     def get_user(self, user_or_sp=None):
         return self._msal_app.get_accounts(user_or_sp)
 

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

@@ -208,7 +208,7 @@ def login_with_managed_identity(self, identity_id=None, allow_no_subscriptions=N
 
         resource = self.cli_ctx.cloud.endpoints.active_directory_resource_id
         identity = Identity()
-        credential, mi_info = identity.login_with_managed_identity(identity_id, resource)
+        credential, mi_info = identity.login_with_managed_identity(resource, identity_id)
 
         tenant = mi_info[Identity.MANAGED_IDENTITY_TENANT_ID]
         if find_subscriptions:
@@ -244,12 +244,39 @@ def login_with_managed_identity(self, identity_id=None, allow_no_subscriptions=N
 
         consolidated = self._normalize_properties(user_name, subscriptions, is_service_principal=True,
                                                   user_assigned_identity_id=legacy_base_name,
-                                                  managed_identity_client_id=client_id)
+                                                  managed_identity_info=mi_info)
+        self._set_subscriptions(consolidated)
+        return deepcopy(consolidated)
+
+    def login_in_cloud_shell(self, allow_no_subscriptions=None, find_subscriptions=True):
+        # TODO: deprecate allow_no_subscriptions
+        resource = self.cli_ctx.cloud.endpoints.active_directory_resource_id
+        identity = Identity()
+        credential, identity_info = identity.login_in_cloud_shell(resource)
+
+        tenant = identity_info[Identity.MANAGED_IDENTITY_TENANT_ID]
+        if find_subscriptions:
+            logger.info('Finding subscriptions...')
+            subscription_finder = SubscriptionFinder(self.cli_ctx)
+            subscriptions = subscription_finder.find_using_specific_tenant(tenant, credential)
+            if not subscriptions:
+                if allow_no_subscriptions:
+                    subscriptions = self._build_tenant_level_accounts([tenant])
+                else:
+                    raise CLIError('No access was configured for the VM, hence no subscriptions were found. '
+                                   "If this is expected, use '--allow-no-subscriptions' to have tenant level access.")
+        else:
+            subscriptions = self._build_tenant_level_accounts([tenant])
+
+        consolidated = self._normalize_properties(identity_info[Identity.CLOUD_SHELL_IDENTITY_UNIQUE_NAME],
+                                                  subscriptions, is_service_principal=False)
+        for s in consolidated:
+            s[_USER_ENTITY][_CLOUD_SHELL_ID] = True
         self._set_subscriptions(consolidated)
         return deepcopy(consolidated)
 
     def _normalize_properties(self, user, subscriptions, is_service_principal, cert_sn_issuer_auth=None,
-                              user_assigned_identity_id=None, home_account_id=None, managed_identity_client_id=None):
+                              user_assigned_identity_id=None, home_account_id=None, managed_identity_info=None):
         import sys
         consolidated = []
         for s in subscriptions:
@@ -285,8 +312,10 @@ def _normalize_properties(self, user, subscriptions, is_service_principal, cert_
 
             if cert_sn_issuer_auth:
                 subscription_dict[_USER_ENTITY][_SERVICE_PRINCIPAL_CERT_SN_ISSUER_AUTH] = True
-            if managed_identity_client_id:
-                subscription_dict[_USER_ENTITY]['clientId'] = managed_identity_client_id
+            if managed_identity_info:
+                subscription_dict[_USER_ENTITY]['clientId'] = managed_identity_info[Identity.MANAGED_IDENTITY_CLIENT_ID]
+                subscription_dict[_USER_ENTITY]['objectId'] = managed_identity_info[Identity.MANAGED_IDENTITY_OBJECT_ID]
+                subscription_dict[_USER_ENTITY]['resourceId'] = managed_identity_info[Identity.MANAGED_IDENTITY_RESOURCE_ID]
 
             # This will be deprecated and client_id will be the only persisted ID
             if user_assigned_identity_id:
@@ -315,33 +344,6 @@ def _new_account(self):
         s.state = StateType.enabled
         return s
 
-    def find_subscriptions_in_cloud_console(self):
-        import jwt
-
-        _, token, _ = self._get_token_from_cloud_shell(self.cli_ctx.cloud.endpoints.active_directory_resource_id)
-        logger.info('MSI: token was retrieved. Now trying to initialize local accounts...')
-        decode = jwt.decode(token, verify=False, algorithms=['RS256'])
-        tenant = decode['tid']
-
-        subscription_finder = SubscriptionFinder(self.cli_ctx, self.auth_ctx_factory, None)
-        subscriptions = subscription_finder.find_from_raw_token(tenant, token)
-        if not subscriptions:
-            raise CLIError('No subscriptions were found in the cloud shell')
-        user = decode.get('unique_name', 'N/A')
-
-        consolidated = self._normalize_properties(user, subscriptions, is_service_principal=False)
-        for s in consolidated:
-            s[_USER_ENTITY][_CLOUD_SHELL_ID] = True
-        self._set_subscriptions(consolidated)
-        return deepcopy(consolidated)
-
-    def _get_token_from_cloud_shell(self, resource):  # pylint: disable=no-self-use
-        from msrestazure.azure_active_directory import MSIAuthentication
-        auth = MSIAuthentication(resource=resource)
-        auth.set_token()
-        token_entry = auth.token
-        return (token_entry['token_type'], token_entry['access_token'], token_entry)
-
     def _set_subscriptions(self, new_subscriptions, merge=True, secondary_key_name=None):
 
         def _get_key_name(account, secondary_key_name):

src/azure-cli/azure/cli/command_modules/profile/custom.py

@@ -124,7 +124,7 @@ def login(cmd, username=None, password=None, service_principal=None, tenant=None
 
     if identity:
         if in_cloud_console():
-            return profile.find_subscriptions_in_cloud_console()
+            return profile.login_in_cloud_shell()
         return profile.login_with_managed_identity(username, allow_no_subscriptions)
     if in_cloud_console():  # tell users they might not need login
         logger.warning(_CLOUD_CONSOLE_LOGIN_WARNING)