cloud shell login: bring back the raw token support as the shell is not ready to switch (#5772)

Author: yugangw-msft

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

@@ -189,6 +189,52 @@ def find_subscriptions_on_login(self,
         # use deepcopy as we don't want to persist these changes to file.
         return deepcopy(consolidated)
 
+    def find_subscriptions_in_cloud_console_thru_raw_token(self, tokens):
+        from datetime import datetime, timedelta
+        import jwt
+        arm_token = tokens[0]  # cloud shell gurantees that the 1st is for ARM
+        arm_token_decoded = jwt.decode(arm_token, verify=False, algorithms=['RS256'])
+        tenant = arm_token_decoded['tid']
+        user_id = arm_token_decoded['unique_name'].split('#')[-1]
+        subscription_finder = SubscriptionFinder(self.cli_ctx, self.auth_ctx_factory, None)
+        subscriptions = subscription_finder.find_from_raw_token(tenant, arm_token)
+        consolidated = self._normalize_properties(user_id, subscriptions, is_service_principal=False)
+        self._set_subscriptions(consolidated)
+
+        # construct token entries to cache
+        decoded_tokens = [arm_token_decoded]
+        for t in tokens[1:]:
+            decoded_tokens.append(jwt.decode(t, verify=False, algorithms=['RS256']))
+        final_tokens = []
+        # Note, setting expiration time at 2700 seconds is bit arbitrary, but should not matter
+        # as shell should update us with new ones every 10~15 minutes
+        for t in decoded_tokens:
+            final_tokens.append({
+                '_clientId': _CLIENT_ID,
+                'expiresIn': '2700',
+                'expiresOn': str(datetime.now() + timedelta(seconds=2700)),
+                'userId': t['unique_name'].split('#')[-1],
+                '_authority': self.cli_ctx.cloud.endpoints.active_directory.rstrip('/') + '/' + t['tid'],
+                'resource': t['aud'],
+                'isMRRT': True,
+                'accessToken': tokens[decoded_tokens.index(t)],
+                'tokenType': 'Bearer',
+            })
+
+        # merging with existing cached ones
+        for t in final_tokens:
+            cached_tokens = [entry for _, entry in self._creds_cache.adal_token_cache.read_items()]
+            to_delete = [c for c in cached_tokens if (c['_clientId'].lower() == t['_clientId'].lower() and
+                                                      c['resource'].lower() == t['resource'].lower() and
+                                                      c['_authority'].lower() == t['_authority'].lower() and
+                                                      c['userId'].lower() == t['userId'].lower())]
+            if to_delete:
+                self._creds_cache.adal_token_cache.remove(to_delete)
+        self._creds_cache.adal_token_cache.add(final_tokens)
+        self._creds_cache.persist_cached_creds()
+
+        return deepcopy(consolidated)
+
     def _normalize_properties(self, user, subscriptions, is_service_principal):
         consolidated = []
         for s in subscriptions:
@@ -410,8 +456,8 @@ def _retrieve_token():
             identity_id, msi_port = Profile._try_parse_for_msi_port(account[_SUBSCRIPTION_NAME])
             if msi_port is not None:
                 return Profile.get_msi_token(resource, msi_port, identity_id)
-            elif in_cloud_console() and account[_USER_ENTITY].get(_CLOUD_SHELL_ID):
-                return Profile.get_msi_token(resource, _get_cloud_console_token_endpoint())
+            # elif in_cloud_console() and account[_USER_ENTITY].get(_CLOUD_SHELL_ID):
+            #     return Profile.get_msi_token(resource, _get_cloud_console_token_endpoint())
             elif user_type == _USER:
                 return self._creds_cache.retrieve_token_for_user(username_or_sp_id,
                                                                  account[_TENANT_ID], resource)
@@ -448,8 +494,8 @@ def get_raw_token(self, resource=None, subscription=None):
         identity_id, msi_port = Profile._try_parse_for_msi_port(account[_SUBSCRIPTION_NAME])
         if msi_port is not None:
             creds = Profile.get_msi_token(resource, msi_port, identity_id)
-        elif in_cloud_console() and account[_USER_ENTITY].get(_CLOUD_SHELL_ID):
-            creds = Profile.get_msi_token(resource, _get_cloud_console_token_endpoint())
+        # elif in_cloud_console() and account[_USER_ENTITY].get(_CLOUD_SHELL_ID):
+        #     creds = Profile.get_msi_token(resource, _get_cloud_console_token_endpoint())
         elif user_type == _USER:
             creds = self._creds_cache.retrieve_token_for_user(username_or_sp_id,
                                                               account[_TENANT_ID], resource)

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

@@ -14,10 +14,13 @@
 
 logger = get_logger(__name__)
 
-_CLOUD_CONSOLE_LOGOUT_WARNING = ("Logout successful. Re-login to your initial Cloud Shell identity with"
-                                 " 'az login --identity'. Login with a new identity with 'az login'.")
-_CLOUD_CONSOLE_LOGIN_WARNING = ("Cloud Shell is automatically authenticated under the initial account signed-in with."
-                                " Run 'az login' only if you need to use a different account")
+
+_CLOUD_CONSOLE_WARNING_TEMPLATE = ("Azure Cloud Shell automatically authenticates the user account it was initially"
+                                   " launched under, as a result 'az %s' is disabled.")
+# _CLOUD_CONSOLE_LOGOUT_WARNING = ("Logout successful. Re-login to your initial Cloud Shell identity with"
+#                                  " 'az login --identity'. Login with a new identity with 'az login'.")
+# _CLOUD_CONSOLE_LOGIN_WARNING = ("Cloud Shell is automatically authenticated under the initial account signed-in with."
+#                                 " Run 'az login' only if you need to use a different account")
 
 
 def _load_subscriptions(cli_ctx, all_clouds=False, refresh=False):
@@ -83,8 +86,8 @@ def set_active_subscription(cmd, subscription):
 
 def account_clear(cmd):
     """Clear all stored subscriptions. To clear individual, use 'logout'"""
-    if in_cloud_console():
-        logger.warning(_CLOUD_CONSOLE_LOGOUT_WARNING)
+    # if in_cloud_console():
+    #     logger.warning(_CLOUD_CONSOLE_LOGOUT_WARNING)
     profile = Profile(cli_ctx=cmd.cli_ctx)
     profile.logout_all()
 
@@ -94,6 +97,8 @@ def login(cmd, username=None, password=None, service_principal=None, tenant=None
           identity=False, identity_port=None,
           msi=False, msi_port=None):  # will remove msi_xxx in a future release
     """Log in to access Azure subscriptions"""
+    import os
+    import re
     from adal.adal_error import AdalError
     import requests
 
@@ -106,12 +111,22 @@ def login(cmd, username=None, password=None, service_principal=None, tenant=None
 
     profile = Profile(cli_ctx=cmd.cli_ctx, async_persist=False)
 
+    # if identity or msi:
+    #     if in_cloud_console():
+    #         return profile.find_subscriptions_in_cloud_console()
+    #     return profile.find_subscriptions_in_vm_with_msi(identity_port or msi_port or 50342, username)
+    # elif in_cloud_console():  # tell users they might not need login
+    #     logger.warning(_CLOUD_CONSOLE_LOGIN_WARNING)
+
+    if in_cloud_console():
+        console_tokens = os.environ.get('AZURE_CONSOLE_TOKENS', None)
+        if console_tokens:
+            return profile.find_subscriptions_in_cloud_console_thru_raw_token(re.split(';|,', console_tokens))
+        logger.warning(_CLOUD_CONSOLE_WARNING_TEMPLATE, 'login')
+        return
+
     if identity or msi:
-        if in_cloud_console():
-            return profile.find_subscriptions_in_cloud_console()
         return profile.find_subscriptions_in_vm_with_msi(identity_port or msi_port or 50342, username)
-    elif in_cloud_console():  # tell users they might not need login
-        logger.warning(_CLOUD_CONSOLE_LOGIN_WARNING)
 
     if username:
         if not password:
@@ -150,8 +165,11 @@ def login(cmd, username=None, password=None, service_principal=None, tenant=None
 
 def logout(cmd, username=None):
     """Log out to remove access to Azure subscriptions"""
+    # if in_cloud_console():
+    #     logger.warning(_CLOUD_CONSOLE_LOGOUT_WARNING)
     if in_cloud_console():
-        logger.warning(_CLOUD_CONSOLE_LOGOUT_WARNING)
+        logger.warning(_CLOUD_CONSOLE_WARNING_TEMPLATE, 'logout')
+        return
 
     profile = Profile(cli_ctx=cmd.cli_ctx)
     if not username: