PoC: Silent flow utilizes Cloud Shell IMDS

Author: rayluo

msal/application.py

@@ -29,6 +29,7 @@
 __version__ = "1.16.0"
 
 logger = logging.getLogger(__name__)
+_CLOUD_SHELL_USER = "current_cloud_shell_user"
 
 
 def extract_certs(public_cert_content):
@@ -967,6 +968,19 @@ def get_accounts(self, username=None):
         # 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):
@@ -1125,6 +1139,12 @@ def acquire_token_silent_with_error(
         """
         assert isinstance(scopes, list), "Invalid parameter type"
         self._validate_ssh_cert_input_data(kwargs.get("data", {}))
+        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 self._acquire_token_by_cloud_shell(
+                scopes, data=kwargs.get("data", {}))
         correlation_id = msal.telemetry._get_new_correlation_id()
         if authority:
             warnings.warn("We haven't decided how/if this method will accept authority parameter")
@@ -1217,6 +1237,11 @@ 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 = self._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,
                 refresh_reason=refresh_reason, claims_challenge=claims_challenge,
@@ -1229,6 +1254,24 @@ 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 = {

tests/test_e2e.py

@@ -229,6 +229,9 @@ 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
@@ -253,6 +256,19 @@ 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 by python -m unittest tests.test_e2e.SshCertTestCase")
+    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"))
+
 
 THIS_FOLDER = os.path.dirname(__file__)
 CONFIG = os.path.join(THIS_FOLDER, "config.json")