Abandon get_accounts(username=msal.CURRENT_USER) by acquire_token_interactive(..., prompt="none")

Author: rayluo

msal/__init__.py

@@ -30,7 +30,6 @@
     ClientApplication,
     ConfidentialClientApplication,
     PublicClientApplication,
-    CURRENT_USER,
     )
 from .oauth2cli.oidc import Prompt
 from .token_cache import TokenCache, SerializableTokenCache

msal/application.py

@@ -22,16 +22,13 @@
 from .region import _detect_region
 from .throttled_http_client import ThrottledHttpClient
 from .cloudshell import _is_running_in_cloud_shell
-from .cloudshell import _acquire_token as _acquire_token_by_cloud_shell
 
 
 # The __init__.py will import this. Not the other way around.
 __version__ = "1.17.0"  # When releasing, also check and bump our dependencies's versions if needed
 
 logger = logging.getLogger(__name__)
-CURRENT_USER = "Current User"  # The value is subject to change
-_CLOUD_SHELL_USER = "current_cloud_shell_user"  # The value is subject to change
-
+_AUTHORITY_TYPE_CLOUDSHELL = "CLOUDSHELL"
 
 def extract_certs(public_cert_content):
     # Parses raw public certificate file contents and returns a list of strings
@@ -118,8 +115,6 @@ def _preferred_browser():
     return None
 
 
-
-
 class _ClientWithCcsRoutingInfo(Client):
 
     def initiate_auth_code_flow(self, **kwargs):
@@ -950,16 +945,6 @@ def get_accounts(self, username=None):
             Your app can choose to display those information to end user,
             and allow user to choose one of his/her accounts to proceed.
         """
-        cloud_shell_pseudo_account = {
-            "home_account_id": _CLOUD_SHELL_USER,
-            "environment": "",
-            "realm": "",
-            "local_account_id": _CLOUD_SHELL_USER,
-            "username": CURRENT_USER,
-            "authority_type": TokenCache.AuthorityType.MSSTS,
-            }
-        if _is_running_in_cloud_shell() and username == CURRENT_USER:
-            return [cloud_shell_pseudo_account]
         accounts = self._find_msal_accounts(environment=self.authority.instance)
         if not accounts:  # Now try other aliases of this authority instance
             for alias in self._get_authority_aliases(self.authority.instance):
@@ -978,13 +963,6 @@ def get_accounts(self, username=None):
                     "they would contain no username for filtering. "
                     "Consider calling get_accounts(username=None) instead."
                     ).format(username))
-        if _is_running_in_cloud_shell() and not username:
-            # In Cloud Shell, user already signed in w/ an account [email protected]
-            # We pretend we have that account, for acquire_token_silent() to work.
-            # Note: If user calls acquire_token_by_xyz() with same account later,
-            # the get_accounts(username=None) would return multiple accounts,
-            # with different usernames: [email protected] and CURRENT_USER.
-            accounts.insert(0, cloud_shell_pseudo_account)
         # Does not further filter by existing RTs here. It probably won't matter.
         # Because in most cases Accounts and RTs co-exist.
         # Even in the rare case when an RT is revoked and then removed,
@@ -993,6 +971,10 @@ def get_accounts(self, username=None):
         return accounts
 
     def _find_msal_accounts(self, environment):
+        interested_authority_types = [
+            TokenCache.AuthorityType.ADFS, TokenCache.AuthorityType.MSSTS]
+        if _is_running_in_cloud_shell():
+            interested_authority_types.append(_AUTHORITY_TYPE_CLOUDSHELL)
         grouped_accounts = {
             a.get("home_account_id"):  # Grouped by home tenant's id
                 {  # These are minimal amount of non-tenant-specific account info
@@ -1008,8 +990,7 @@ def _find_msal_accounts(self, environment):
             for a in self.token_cache.find(
                 TokenCache.CredentialType.ACCOUNT,
                 query={"environment": environment})
-            if a["authority_type"] in (
-                TokenCache.AuthorityType.ADFS, TokenCache.AuthorityType.MSSTS)
+            if a["authority_type"] in interested_authority_types
             }
         return list(grouped_accounts.values())
 
@@ -1069,6 +1050,21 @@ def _forget_me(self, home_account):
                 TokenCache.CredentialType.ACCOUNT, query=owned_by_home_account):
             self.token_cache.remove_account(a)
 
+    def _acquire_token_by_cloud_shell(self, scopes, data=None):
+        from .cloudshell import _acquire_token
+        response = _acquire_token(
+            self.http_client, scopes, client_id=self.client_id, data=data)
+        if "error" not in response:
+            self.token_cache.add(dict(
+                client_id=self.client_id,
+                scope=response["scope"].split() if "scope" in response else scopes,
+                token_endpoint=self.authority.token_endpoint,
+                response=response.copy(),
+                data=data or {},
+                authority_type=_AUTHORITY_TYPE_CLOUDSHELL,
+                ))
+        return response
+
     def acquire_token_silent(
             self,
             scopes,  # type: List[str]
@@ -1148,13 +1144,6 @@ def acquire_token_silent_with_error(
         """
         assert isinstance(scopes, list), "Invalid parameter type"
         self._validate_ssh_cert_input_data(kwargs.get("data", {}))
-
-        # The special code path only for _CLOUD_SHELL_USER
-        if account and account.get("home_account_id") == _CLOUD_SHELL_USER:
-            # Since we don't currently store cloud shell tokens in MSAL's cache,
-            # we can have a shortcut here, and semantically bypass all those
-            # _acquire_token_silent_from_cache_and_possibly_refresh_it()
-            return _acquire_token_by_cloud_shell(self.http_client, scopes, **kwargs)
         correlation_id = msal.telemetry._get_new_correlation_id()
         if authority:
             warnings.warn("We haven't decided how/if this method will accept authority parameter")
@@ -1209,6 +1198,7 @@ def _acquire_token_silent_from_cache_and_possibly_refresh_it(
             authority,  # This can be different than self.authority
             force_refresh=False,  # type: Optional[boolean]
             claims_challenge=None,
+            correlation_id=None,
             **kwargs):
         access_token_from_cache = None
         if not (force_refresh or claims_challenge):  # Bypass AT when desired or using claims
@@ -1247,14 +1237,13 @@ def _acquire_token_silent_from_cache_and_possibly_refresh_it(
             refresh_reason = msal.telemetry.FORCE_REFRESH  # TODO: It could also mean claims_challenge
         assert refresh_reason, "It should have been established at this point"
         try:
-            ## When/if we will store Cloud Shell tokens into MSAL's token cache,
-            # then we will add the following code snippet here.
-            #if account and account.get("home_account_id") == _CLOUD_SHELL_USER:
-            #    result = _acquire_token_by_cloud_shell(..., scopes, **kwargs)
-            #else:
+            if account and account.get("authority_type") == _AUTHORITY_TYPE_CLOUDSHELL:
+                return self._acquire_token_by_cloud_shell(
+                    scopes, data=kwargs.get("data"))
             result = _clean_up(self._acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
                 authority, self._decorate_scope(scopes), account,
                 refresh_reason=refresh_reason, claims_challenge=claims_challenge,
+                correlation_id=correlation_id,
                 **kwargs))
             if (result and "error" not in result) or (not access_token_from_cache):
                 return result
@@ -1593,6 +1582,9 @@ def acquire_token_interactive(
             - A dict containing an "error" key, when token refresh failed.
         """
         self._validate_ssh_cert_input_data(kwargs.get("data", {}))
+        if _is_running_in_cloud_shell() and prompt == "none":
+            return self._acquire_token_by_cloud_shell(
+                scopes, data=kwargs.pop("data", {}))
         claims = _merge_claims_challenge_and_capabilities(
             self._client_capabilities, claims_challenge)
         telemetry_context = self._build_telemetry_context(

msal/cloudshell.py

@@ -4,13 +4,16 @@
 # This code is licensed under the MIT License.
 
 """This module wraps Cloud Shell's IMDS-like interface inside an OAuth2-like helper"""
+import base64
 import json
 import logging
 import os
+import time
 try:  # Python 2
     from urlparse import urlparse
 except:  # Python 3
     from urllib.parse import urlparse
+from .oauth2cli.oidc import decode_part
 
 
 logger = logging.getLogger(__name__)
@@ -34,33 +37,78 @@ def _scope_to_resource(scope):  # This is an experimental reasonable-effort appr
     return scope  # There is no much else we can do here
 
 
-def _acquire_token(http_client, scopes, **kwargs):
+def _acquire_token(http_client, scopes, client_id=None, data=None):
     resp = http_client.post(
         "http://localhost:50342/oauth2/token",
         data=dict(
-            kwargs.pop("data", {}),
+            data or {},
             resource=" ".join(map(_scope_to_resource, scopes))),
-        headers=dict(kwargs.pop("headers", {}), Metadata="true"),
-        **kwargs)
+        headers={"Metadata": "true"},
+        )
     if resp.status_code >= 300:
         logger.debug("Cloud Shell IMDS error: %s", resp.text)
         cs_error = json.loads(resp.text).get("error", {})
         return {k: v for k, v in {
             "error": cs_error.get("code"),
             "error_description": cs_error.get("message"),
             }.items() if v}
-    payload = json.loads(resp.text)
+    imds_payload = json.loads(resp.text)
     oauth2_response = {
-        "access_token": payload["access_token"],
-        "expires_in": int(payload["expires_in"]),
-        "token_type": payload.get("token_type", "Bearer"),
+        "access_token": imds_payload["access_token"],
+        "expires_in": int(imds_payload["expires_in"]),
+        "token_type": imds_payload.get("token_type", "Bearer"),
         }
+    parts = imds_payload["access_token"].split(".")
+
+    # The following default values are useful in SSH Cert scenario
+    client_info = {  # Default value, in case the real value will be unavailable
+        "uid": "user",
+        "utid": "cloudshell",
+        }
+    now = time.time()
+    preferred_username = "currentuser@cloudshell"
+    oauth2_response["id_token_claims"] = {  # First 5 claims are required per OIDC
+        "iss": "cloudshell",
+        "sub": "user",
+        "aud": client_id,
+        "exp": now + 3600,
+        "iat": now,
+        "preferred_username": preferred_username,  # Useful as MSAL account's username
+        }
+
+    if len(parts) == 3:  # Probably a JWT. Use it to derive client_info and id token.
+        try:
+            # Data defined in https://docs.microsoft.com/en-us/azure/active-directory/develop/access-tokens#payload-claims
+            jwt_payload = json.loads(decode_part(parts[1]))
+            client_info = {
+                # Mimic a real home_account_id,
+                # so that this pseudo account and a real account would interop.
+                "uid": jwt_payload.get("oid", "user"),
+                "utid": jwt_payload.get("tid", "cloudshell"),
+                }
+            oauth2_response["id_token_claims"] = {
+                "iss": jwt_payload["iss"],
+                "sub": jwt_payload["sub"],  # Could use oid instead
+                "aud": client_id,
+                "exp": jwt_payload["exp"],
+                "iat": jwt_payload["iat"],
+                "preferred_username": jwt_payload.get("preferred_username")  # V2
+                    or jwt_payload.get("unique_name")  # V1
+                    or preferred_username,
+                }
+        except ValueError:
+            logger.debug("Unable to decode jwt payload: %s", parts[1])
+    oauth2_response["client_info"] = base64.b64encode(
+        # Mimic a client_info, so that MSAL would create an account
+        json.dumps(client_info).encode("utf-8")).decode("utf-8")
+    oauth2_response["id_token_claims"]["tid"] = client_info["utid"]  # TBD
+
     ## Note: Decided to not surface resource back as scope,
     ##       because they would cause the downstream OAuth2 code path to
     ##       cache the token with a different scope and won't hit them later.
-    #if payload.get("resource"):
-    #    oauth2_response["scope"] = payload["resource"]
-    if payload.get("refresh_token"):
-        oauth2_response["refresh_token"] = payload["refresh_token"]
+    #if imds_payload.get("resource"):
+    #    oauth2_response["scope"] = imds_payload["resource"]
+    if imds_payload.get("refresh_token"):
+        oauth2_response["refresh_token"] = imds_payload["refresh_token"]
     return oauth2_response
 

msal/token_cache.py

@@ -113,6 +113,7 @@ def wipe(dictionary, sensitive_fields):  # Masks sensitive info
             return self.__add(event, now=now)
         finally:
             wipe(event.get("response", {}), (  # These claims were useful during __add()
+                "id_token_claims",  # Provided by broker
                 "access_token", "refresh_token", "id_token", "username"))
             wipe(event, ["username"])  # Needed for federated ROPC
             logger.debug("event=%s", json.dumps(
@@ -150,7 +151,8 @@ def __add(self, event, now=None):
         id_token = response.get("id_token")
         id_token_claims = (
             decode_id_token(id_token, client_id=event["client_id"])
-            if id_token else {})
+            if id_token
+            else response.get("id_token_claims", {}))  # Broker would provide id_token_claims
         client_info, home_account_id = self.__parse_account(response, id_token_claims)
 
         target = ' '.join(event.get("scope") or [])  # Per schema, we don't sort it
@@ -195,9 +197,10 @@ def __add(self, event, now=None):
                         or data.get("username")  # Falls back to ROPC username
                         or event.get("username")  # Falls back to Federated ROPC username
                         or "",  # The schema does not like null
-                    "authority_type":
+                    "authority_type": event.get(
+                        "authority_type",  # Honor caller's choice of authority_type
                         self.AuthorityType.ADFS if realm == "adfs"
-                        else self.AuthorityType.MSSTS,
+                            else self.AuthorityType.MSSTS),
                     # "client_info": response.get("client_info"),  # Optional
                     }
                 self.modify(self.CredentialType.ACCOUNT, account, account)

tests/test_e2e.py

@@ -185,12 +185,14 @@ def _test_acquire_token_interactive(
             self, client_id=None, authority=None, scope=None, port=None,
             username_uri="",  # But you would want to provide one
             data=None,  # Needed by ssh-cert feature
+            prompt=None,
             **ignored):
         assert client_id and authority and scope
         self.app = msal.PublicClientApplication(
             client_id, authority=authority, http_client=MinimalHttpClient())
         result = self.app.acquire_token_interactive(
             scope,
+            prompt=prompt,
             timeout=120,
             port=port,
             welcome_template=  # This is an undocumented feature for testing
@@ -229,9 +231,6 @@ def test_ssh_cert_for_service_principal(self):
         self.assertEqual("ssh-cert", result["token_type"])
 
     @unittest.skipIf(os.getenv("TRAVIS"), "Browser automation is not yet implemented")
-    @unittest.skipIf(
-        msal.application._is_running_in_cloud_shell(),
-        "The test app does not opt in to Cloud Shell")
     def test_ssh_cert_for_user(self):
         result = self._test_acquire_token_interactive(
             client_id="04b07795-8ddb-461a-bbee-02f9e1bf7b46",  # Azure CLI is one
@@ -240,6 +239,7 @@ def test_ssh_cert_for_user(self):
             scope=self.SCOPE,
             data=self.DATA1,
             username_uri="https://msidlab.com/api/user?usertype=cloud",
+            prompt="none" if msal.application._is_running_in_cloud_shell() else None,
             )   # It already tests reading AT from cache, and using RT to refresh
                 # acquire_token_silent() would work because we pass in the same key
         self.assertIsNotNone(result.get("access_token"), "Encountered {}: {}".format(
@@ -256,31 +256,16 @@ def test_ssh_cert_for_user(self):
         self.assertEqual(refreshed_ssh_cert["token_type"], "ssh-cert")
         self.assertNotEqual(result["access_token"], refreshed_ssh_cert['access_token'])
 
-    @unittest.skipUnless(
-        msal.application._is_running_in_cloud_shell(),
-        "Manually run this test case from inside Cloud Shell")
-    def test_ssh_cert_for_user_silent_inside_cloud_shell(self):
-        app = msal.PublicClientApplication("client_id_wont_matter")
-        accounts = app.get_accounts()
-        self.assertNotEqual([], accounts)
-        result = app.acquire_token_silent_with_error(
-            self.SCOPE, account=accounts[0], data=self.DATA1)
-        self.assertEqual(
-            "ssh-cert", result.get("token_type"), "Unexpected result: %s" % result)
-        self.assertIsNotNone(result.get("access_token"))
-
 
 @unittest.skipUnless(
     msal.application._is_running_in_cloud_shell(),
     "Manually run this test case from inside Cloud Shell")
 class CloudShellTestCase(E2eTestCase):
-    app = msal.PublicClientApplication("client_id_wont_matter")
+    app = msal.PublicClientApplication("client_id")
     scope_that_requires_no_managed_device = "https://management.core.windows.net/"  # Scopes came from https://msazure.visualstudio.com/One/_git/compute-CloudShell?path=/src/images/agent/env/envconfig.PROD.json&version=GBmaster&_a=contents
     def test_access_token_should_be_obtained_for_a_supported_scope(self):
-        accounts = self.app.get_accounts(username=msal.CURRENT_USER)
-        self.assertNotEqual([], accounts)
-        result = self.app.acquire_token_silent_with_error(
-            [self.scope_that_requires_no_managed_device], account=accounts[0])
+        result = self.app.acquire_token_interactive(
+            [self.scope_that_requires_no_managed_device], prompt="none")
         self.assertEqual(
             "Bearer", result.get("token_type"), "Unexpected result: %s" % result)
         self.assertIsNotNone(result.get("access_token"))