Introduce get_accounts(username=msal.CURRENT_USER)

Author: rayluo

msal/__init__.py

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

msal/application.py

@@ -23,13 +23,16 @@
 import msal.telemetry
 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.16.0"
 
 logger = logging.getLogger(__name__)
-_CLOUD_SHELL_USER = "current_cloud_shell_user"
+CURRENT_USER = "Current User"  # The value is subject to change
+_CLOUD_SHELL_USER = "current_cloud_shell_user"  # The value is subject to change
 
 
 def extract_certs(public_cert_content):
@@ -113,8 +116,6 @@ def _preferred_browser():
     return None
 
 
-def _is_running_in_cloud_shell():
-    return os.environ.get("AZUREPS_HOST_ENVIRONMENT", "").startswith("cloud-shell")
 
 
 class _ClientWithCcsRoutingInfo(Client):
@@ -945,6 +946,16 @@ 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):
@@ -963,24 +974,18 @@ 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() would return multiple accounts to calling app,
+            # 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,
         # acquire_token_silent() would then yield no result,
         # apps would fall back to other acquire methods. This is the standard pattern.
-        if _is_running_in_cloud_shell():
-            # In Cloud Shell, user already signed in with an account.
-            # We pretend we have that account, for acquire_token_silent() to work.
-            # Note: If user acquire_token_by_xyz() using that account in MSAL later,
-            # the get_accounts() would return multiple accounts to calling app.
-            accounts.insert(0, {
-                "home_account_id": _CLOUD_SHELL_USER,
-                "environment": "",
-                "realm": "",
-                "local_account_id": _CLOUD_SHELL_USER,
-                "username": "Current Cloud Shell User",
-                "authority_type": TokenCache.AuthorityType.MSSTS,
-                })
         return accounts
 
     def _find_msal_accounts(self, environment):
@@ -1158,8 +1163,7 @@ def acquire_token_silent_with_error(
             # 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 self._acquire_token_by_cloud_shell(
-                scopes, data=kwargs.get("data", {}))
+            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")
@@ -1255,7 +1259,7 @@ def _acquire_token_silent_from_cache_and_possibly_refresh_it(
             ## 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 = self._acquire_token_by_cloud_shell(scopes, **kwargs)
+            #    result = _acquire_token_by_cloud_shell(..., scopes, **kwargs)
             #else:
             result = _clean_up(self._acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
                 authority, self._decorate_scope(scopes), account,
@@ -1269,24 +1273,6 @@ def _acquire_token_silent_from_cache_and_possibly_refresh_it(
                 raise  # We choose to bubble up the exception
         return access_token_from_cache
 
-    def _acquire_token_by_cloud_shell(self, scopes, **kwargs):
-        kwargs.pop("correlation_id", None)  # IMDS does not use correlation_id
-        resp = self.http_client.post(
-            "http://localhost:50342/oauth2/token",
-            data=dict(kwargs.pop("data", {}), resource=" ".join(scopes)),
-            headers=dict(kwargs.pop("headers", {}), Metadata="true"),
-            **kwargs)
-        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}
-        else:
-            # Skip token cache, for now. Cloud Shell IMDS has its own cache anyway.
-            return json.loads(resp.text)
-
     def _acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
             self, authority, scopes, account, **kwargs):
         query = {

msal/cloudshell.py

@@ -0,0 +1,78 @@
+# Copyright (c) Microsoft Corporation.
+# All rights reserved.
+#
+# This code is licensed under the MIT License.
+
+"""This module wraps Cloud Shell's IMDS-like interface inside an OAuth2-like helper"""
+import json
+import logging
+import os
+
+
+logger = logging.getLogger(__name__)
+
+
+def _is_running_in_cloud_shell():
+    return os.environ.get("AZUREPS_HOST_ENVIRONMENT", "").startswith("cloud-shell")
+
+
+def _scope_to_resource(scope):
+    cloud_shell_supported_audiences = [  # Came from https://msazure.visualstudio.com/One/_git/compute-CloudShell?path=/src/images/agent/env/envconfig.PROD.json
+        "https://management.core.windows.net/",
+        "https://management.azure.com/",
+        "https://graph.windows.net/",
+        "https://vault.azure.net",
+        "https://datalake.azure.net/",
+        "https://outlook.office365.com/",
+        "https://graph.microsoft.com/",
+        "https://batch.core.windows.net/",
+        "https://analysis.windows.net/powerbi/api",
+        "https://storage.azure.com/",
+        "https://rest.media.azure.net",
+        "https://api.loganalytics.io",
+        "https://ossrdbms-aad.database.windows.net",
+        "https://www.yammer.com",
+        "https://digitaltwins.azure.net",
+        "0b07f429-9f4b-4714-9392-cc5e8e80c8b0",
+        "822c8694-ad95-4735-9c55-256f7db2f9b4",
+        "https://dev.azuresynapse.net",
+        "https://database.windows.net",
+        "https://quantum.microsoft.com",
+        "https://iothubs.azure.net",
+        "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d",
+        "https://azuredatabricks.net/",
+        "ce34e7e5-485f-4d76-964f-b3d2b16d1e4f",
+        "https://azure-devices-provisioning.net"
+        ]
+    for a in cloud_shell_supported_audiences:
+        if scope.startswith(a):  # This is an experimental approach
+            return a
+    return scope  # Some scope would work as-is, such as the SSH Cert scope
+
+
+def _acquire_token(http_client, scopes, **kwargs):
+    kwargs.pop("correlation_id", None)  # IMDS does not use correlation_id
+    resp = http_client.post(
+        "http://localhost:50342/oauth2/token",
+        data=dict(
+            kwargs.pop("data", {}),
+            resource=" ".join(map(_scope_to_resource, scopes))),
+        headers=dict(kwargs.pop("headers", {}), Metadata="true"),
+        **kwargs)
+    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)
+    oauth2_response = {
+        "access_token": payload["access_token"],
+        "expires_in": int(payload["expires_in"]),
+        "token_type": payload.get("token_type", "Bearer"),
+        }
+    if payload.get("refresh_token"):
+        oauth2_response["refresh_token"] = payload["refresh_token"]
+    return oauth2_response
+

tests/test_e2e.py

@@ -275,10 +275,9 @@ def test_ssh_cert_for_user_silent_inside_cloud_shell(self):
     "Manually run this test case from inside Cloud Shell")
 class CloudShellTestCase(E2eTestCase):
     app = msal.PublicClientApplication("client_id_wont_matter")
-    # Scopes came from https://msazure.visualstudio.com/One/_git/compute-CloudShell?path=/src/images/agent/env/envconfig.PROD.json&version=GBmaster&_a=contents
-    scope_that_requires_no_managed_device = "https://management.core.windows.net/"
+    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()
+        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])