Merge branch 'dev' into dharshanb/brokerSupportLinux

Author: DharshanBJ

msal/managed_identity.py

@@ -2,6 +2,7 @@
 # All rights reserved.
 #
 # This code is licensed under the MIT License.
+import hashlib
 import json
 import logging
 import os
@@ -10,7 +11,7 @@
 import time
 from urllib.parse import urlparse  # Python 3+
 from collections import UserDict  # Python 3+
-from typing import Optional, Union  # Needed in Python 3.7 & 3.8
+from typing import List, Optional, Union  # Needed in Python 3.7 & 3.8
 from .token_cache import TokenCache
 from .individual_cache import _IndividualCache as IndividualCache
 from .throttled_http_client import ThrottledHttpClientBase, RetryAfterParser
@@ -162,6 +163,7 @@ def __init__(
         http_client,
         token_cache=None,
         http_cache=None,
+        client_capabilities: Optional[List[str]] = None,
     ):
         """Create a managed identity client.
 
@@ -192,6 +194,17 @@ def __init__(
             Optional. It has the same characteristics as the
             :paramref:`msal.ClientApplication.http_cache`.
 
+        :param list[str] client_capabilities: (optional)
+            Allows configuration of one or more client capabilities, e.g. ["CP1"].
+
+            Client capability is meant to inform the Microsoft identity platform
+            (STS) what this client is capable for,
+            so STS can decide to turn on certain features.
+
+            Implementation details:
+            Client capability in Managed Identity is relayed as-is
+            via ``xms_cc`` parameter on the wire.
+
         Recipe 1: Hard code a managed identity for your app::
 
             import msal, requests
@@ -238,6 +251,7 @@ def __init__(
             http_cache=http_cache,
         )
         self._token_cache = token_cache or TokenCache()
+        self._client_capabilities = client_capabilities
 
     def _get_instance(self):
         if self.__instance is None:
@@ -266,8 +280,7 @@ def acquire_token_for_client(
             and then a *claims challenge* will be returned by the target resource,
             as a `claims_challenge` directive in the `www-authenticate` header,
             even if the app developer did not opt in for the "CP1" client capability.
-            Upon receiving a `claims_challenge`, MSAL will skip a token cache read,
-            and will attempt to acquire a new token.
+            Upon receiving a `claims_challenge`, MSAL will attempt to acquire a new token.
 
         .. note::
 
@@ -278,12 +291,14 @@ def acquire_token_for_client(
             This is a service-side behavior that cannot be changed by this library.
             `Azure VM docs <https://learn.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token#get-a-token-using-http>`_
         """
+        access_token_to_refresh = None  # This could become a public parameter in the future
         access_token_from_cache = None
         client_id_in_cache = self._managed_identity.get(
             ManagedIdentity.ID, "SYSTEM_ASSIGNED_MANAGED_IDENTITY")
         now = time.time()
-        if not claims_challenge:  # Then attempt token cache search
-            matches = self._token_cache.find(
+        if True:  # Attempt cache search even if receiving claims_challenge,
+                  # because we want to locate the existing token (if any) and refresh it
+            matches = self._token_cache.search(
                 self._token_cache.CredentialType.ACCESS_TOKEN,
                 target=[resource],
                 query=dict(
@@ -297,6 +312,11 @@ def acquire_token_for_client(
                 expires_in = int(entry["expires_on"]) - now
                 if expires_in < 5*60:  # Then consider it expired
                     continue  # Removal is not necessary, it will be overwritten
+                if claims_challenge and not access_token_to_refresh:
+                    # Since caller did not pinpoint the token causing claims challenge,
+                    # we have to assume it is the first token we found in cache.
+                    access_token_to_refresh = entry["secret"]
+                    break
                 logger.debug("Cache hit an AT")
                 access_token_from_cache = {  # Mimic a real response
                     "access_token": entry["secret"],
@@ -310,7 +330,13 @@ def acquire_token_for_client(
                         break  # With a fallback in hand, we break here to go refresh
                 return access_token_from_cache  # It is still good as new
         try:
-            result = _obtain_token(self._http_client, self._managed_identity, resource)
+            result = _obtain_token(
+                self._http_client, self._managed_identity, resource,
+                access_token_sha256_to_refresh=hashlib.sha256(
+                    access_token_to_refresh.encode("utf-8")).hexdigest()
+                    if access_token_to_refresh else None,
+                client_capabilities=self._client_capabilities,
+            )
             if "access_token" in result:
                 expires_in = result.get("expires_in", 3600)
                 if "refresh_in" not in result and expires_in >= 7200:
@@ -385,8 +411,12 @@ def get_managed_identity_source():
     return DEFAULT_TO_VM
 
 
-def _obtain_token(http_client, managed_identity, resource):
-    # A unified low-level API that talks to different Managed Identity
+def _obtain_token(
+    http_client, managed_identity, resource,
+    *,
+    access_token_sha256_to_refresh: Optional[str] = None,
+    client_capabilities: Optional[List[str]] = None,
+):
     if ("IDENTITY_ENDPOINT" in os.environ and "IDENTITY_HEADER" in os.environ
             and "IDENTITY_SERVER_THUMBPRINT" in os.environ
     ):
@@ -402,6 +432,8 @@ def _obtain_token(http_client, managed_identity, resource):
             os.environ["IDENTITY_HEADER"],
             os.environ["IDENTITY_SERVER_THUMBPRINT"],
             resource,
+            access_token_sha256_to_refresh=access_token_sha256_to_refresh,
+            client_capabilities=client_capabilities,
         )
     if "IDENTITY_ENDPOINT" in os.environ and "IDENTITY_HEADER" in os.environ:
         return _obtain_token_on_app_service(
@@ -553,6 +585,9 @@ def _obtain_token_on_machine_learning(
 
 def _obtain_token_on_service_fabric(
     http_client, endpoint, identity_header, server_thumbprint, resource,
+    *,
+    access_token_sha256_to_refresh: str = None,
+    client_capabilities: Optional[List[str]] = None,
 ):
     """Obtains token for
     `Service Fabric <https://learn.microsoft.com/en-us/azure/service-fabric/>`_
@@ -563,7 +598,12 @@ def _obtain_token_on_service_fabric(
     logger.debug("Obtaining token via managed identity on Azure Service Fabric")
     resp = http_client.get(
         endpoint,
-        params={"api-version": "2019-07-01-preview", "resource": resource},
+        params={k: v for k, v in {
+            "api-version": "2019-07-01-preview",
+            "resource": resource,
+            "token_sha256_to_refresh": access_token_sha256_to_refresh,
+            "xms_cc": ",".join(client_capabilities) if client_capabilities else None,
+            }.items() if v is not None},
         headers={"Secret": identity_header},
         )
     try:
@@ -584,7 +624,7 @@ def _obtain_token_on_service_fabric(
             "ArgumentNullOrEmpty": "invalid_scope",
             }
         return {
-            "error": error_mapping.get(payload["error"]["code"], "invalid_request"),
+            "error": error_mapping.get(error.get("code"), "invalid_request"),
             "error_description": resp.text,
             }
     except json.decoder.JSONDecodeError:

tests/test_application.py

@@ -153,8 +153,8 @@ def tester(url, data=None, **kwargs):
             return MinimalResponse(status_code=400, text=error_response)
         app._acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
             app.authority, self.scopes, self.account, post=tester)
-        self.assertNotEqual([], app.token_cache.find(
-            msal.TokenCache.CredentialType.REFRESH_TOKEN, query={"secret": self.frt}),
+        self.assertIsNotNone(next(app.token_cache.search(
+            msal.TokenCache.CredentialType.REFRESH_TOKEN, query={"secret": self.frt}), None),
             "The FRT should not be removed from the cache")
 
     def test_known_orphan_app_will_skip_frt_and_only_use_its_own_rt(self):
@@ -187,11 +187,11 @@ def tester(url, data=None, **kwargs):
             app.authority, self.scopes, self.account, post=tester)
         logger.debug("%s.cache = %s", self.id(), self.cache.serialize())
         self.assertEqual("at", at.get("access_token"), "New app should get a new AT")
-        app_metadata = app.token_cache.find(
+        app_metadata = next(app.token_cache.search(
             msal.TokenCache.CredentialType.APP_METADATA,
-            query={"client_id": app.client_id})
-        self.assertNotEqual([], app_metadata, "Should record new app's metadata")
-        self.assertEqual("1", app_metadata[0].get("family_id"),
+            query={"client_id": app.client_id}), None)
+        self.assertIsNotNone(app_metadata, "Should record new app's metadata")
+        self.assertEqual("1", app_metadata.get("family_id"),
             "The new family app should be recorded as in the same family")
     # Known family app will simply use FRT, which is largely the same as this one
 
@@ -218,25 +218,25 @@ def test_family_app_remove_account(self):
         account = app.get_accounts()[0]
         mine = {"home_account_id": account["home_account_id"]}
 
-        self.assertNotEqual([], self.cache.find(
-            self.cache.CredentialType.ACCESS_TOKEN, query=mine))
-        self.assertNotEqual([], self.cache.find(
-            self.cache.CredentialType.REFRESH_TOKEN, query=mine))
-        self.assertNotEqual([], self.cache.find(
-            self.cache.CredentialType.ID_TOKEN, query=mine))
-        self.assertNotEqual([], self.cache.find(
-            self.cache.CredentialType.ACCOUNT, query=mine))
+        self.assertIsNotNone(next(self.cache.search(
+            self.cache.CredentialType.ACCESS_TOKEN, query=mine), None))
+        self.assertIsNotNone(next(self.cache.search(
+            self.cache.CredentialType.REFRESH_TOKEN, query=mine), None))
+        self.assertIsNotNone(next(self.cache.search(
+            self.cache.CredentialType.ID_TOKEN, query=mine), None))
+        self.assertIsNotNone(next(self.cache.search(
+            self.cache.CredentialType.ACCOUNT, query=mine), None))
 
         app.remove_account(account)
 
-        self.assertEqual([], self.cache.find(
-            self.cache.CredentialType.ACCESS_TOKEN, query=mine))
-        self.assertEqual([], self.cache.find(
-            self.cache.CredentialType.REFRESH_TOKEN, query=mine))
-        self.assertEqual([], self.cache.find(
-            self.cache.CredentialType.ID_TOKEN, query=mine))
-        self.assertEqual([], self.cache.find(
-            self.cache.CredentialType.ACCOUNT, query=mine))
+        self.assertIsNone(next(self.cache.search(
+            self.cache.CredentialType.ACCESS_TOKEN, query=mine), None))
+        self.assertIsNone(next(self.cache.search(
+            self.cache.CredentialType.REFRESH_TOKEN, query=mine), None))
+        self.assertIsNone(next(self.cache.search(
+            self.cache.CredentialType.ID_TOKEN, query=mine), None))
+        self.assertIsNone(next(self.cache.search(
+            self.cache.CredentialType.ACCOUNT, query=mine), None))
 
 
 class TestClientApplicationForAuthorityMigration(unittest.TestCase):
@@ -711,25 +711,26 @@ def test_remove_tokens_for_client_should_remove_client_tokens_only(self):
         cca = msal.ConfidentialClientApplication(
             "client_id", client_credential="secret",
             authority="https://login.microsoftonline.com/microsoft.onmicrosoft.com")
-        self.assertEqual(
-            0, len(cca.token_cache.find(msal.TokenCache.CredentialType.ACCESS_TOKEN)))
+        self.assertIsNone(next(cca.token_cache.search(
+            msal.TokenCache.CredentialType.ACCESS_TOKEN), None))
         cca.acquire_token_for_client(
             ["scope"],
             post=lambda url, **kwargs: MinimalResponse(
                 status_code=200, text=json.dumps({"access_token": "AT for client"})))
-        self.assertEqual(
-            1, len(cca.token_cache.find(msal.TokenCache.CredentialType.ACCESS_TOKEN)))
+        self.assertEqual(1, len(list(cca.token_cache.search(
+            msal.TokenCache.CredentialType.ACCESS_TOKEN))))
         cca.acquire_token_by_username_password(
             "johndoe", "password", ["scope"],
             post=lambda url, **kwargs: MinimalResponse(
                 status_code=200, text=json.dumps(build_response(
                     access_token=at_for_user, expires_in=3600,
                     uid="uid", utid="utid",  # This populates home_account_id
                     ))))
-        self.assertEqual(
-            2, len(cca.token_cache.find(msal.TokenCache.CredentialType.ACCESS_TOKEN)))
+        self.assertEqual(2, len(list(cca.token_cache.search(
+            msal.TokenCache.CredentialType.ACCESS_TOKEN))))
         cca.remove_tokens_for_client()
-        remaining_tokens = cca.token_cache.find(msal.TokenCache.CredentialType.ACCESS_TOKEN)
+        remaining_tokens = list(cca.token_cache.search(
+            msal.TokenCache.CredentialType.ACCESS_TOKEN))
         self.assertEqual(1, len(remaining_tokens))
         self.assertEqual(at_for_user, remaining_tokens[0].get("secret"))
 

tests/test_mi.py

@@ -1,7 +1,9 @@
+import hashlib
 import json
 import os
 import sys
 import time
+from typing import List, Optional
 import unittest
 try:
     from unittest.mock import patch, ANY, mock_open, Mock
@@ -52,15 +54,23 @@ def test_helper_class_should_be_interchangable_with_dict_which_could_be_loaded_f
 class ClientTestCase(unittest.TestCase):
     maxDiff = None
 
-    def setUp(self):
-        self.app = ManagedIdentityClient(
+    def _build_app(
+        self,
+        *,
+        client_capabilities: Optional[List[str]] = None,
+    ):
+        return ManagedIdentityClient(
             {   # Here we test it with the raw dict form, to test that
                 # the client has no hard dependency on ManagedIdentity object
                 "ManagedIdentityIdType": "SystemAssigned", "Id": None,
             },
             http_client=requests.Session(),
+            client_capabilities=client_capabilities,
             )
 
+    def setUp(self):
+        self.app = self._build_app()
+
     def test_error_out_on_invalid_input(self):
         with self.assertRaises(ManagedIdentityError):
             ManagedIdentityClient({"foo": "bar"}, http_client=requests.Session())
@@ -79,7 +89,13 @@ def assertCacheStatus(self, app):
             "Should have expected client_id")
         self.assertEqual("managed_identity", at["realm"], "Should have expected realm")
 
-    def _test_happy_path(self, app, mocked_http, expires_in, resource="R"):
+    def _test_happy_path(
+        self, app, mocked_http, expires_in, *, resource="R", claims_challenge=None,
+    ):
+        """It tests a normal token request that is expected to hit IdP,
+        a subsequent same token request that is expected to hit cache,
+        and then a request with claims_challenge that shall hit IdP again.
+        """
         result = app.acquire_token_for_client(resource=resource)
         mocked_http.assert_called()
         call_count = mocked_http.call_count
@@ -115,7 +131,8 @@ def _test_happy_path(self, app, mocked_http, expires_in, resource="R"):
                 expected_refresh_on - 5 < result["refresh_on"] <= expected_refresh_on,
                 "Should have a refresh_on time around the middle of the token's life")
 
-        result = app.acquire_token_for_client(resource=resource, claims_challenge="foo")
+        result = app.acquire_token_for_client(
+            resource=resource, claims_challenge=claims_challenge or "placeholder")
         self.assertEqual("identity_provider", result["token_source"], "Should miss cache")
 
 
@@ -132,6 +149,9 @@ def _test_happy_path(self) -> callable:
 
     def test_happy_path_of_vm(self):
         self._test_happy_path().assert_called_with(
+            # The last call contained claims_challenge
+            # but since IMDS doesn't support token_sha256_to_refresh,
+            # the request shall remain the same as before
             'http://169.254.169.254/metadata/identity/oauth2/token',
             params={'api-version': '2018-02-01', 'resource': 'R'},
             headers={'Metadata': 'true'},
@@ -244,19 +264,46 @@ def test_machine_learning_error_should_be_normalized(self):
     "IDENTITY_SERVER_THUMBPRINT": "bar",
 })
 class ServiceFabricTestCase(ClientTestCase):
+    access_token = "AT"
+    access_token_sha256 = hashlib.sha256(access_token.encode()).hexdigest()
 
-    def _test_happy_path(self, app):
+    def _test_happy_path(self, app, *, claims_challenge=None) -> callable:
         expires_in = 1234
         with patch.object(app._http_client, "get", return_value=MinimalResponse(
             status_code=200,
-            text='{"access_token": "AT", "expires_on": %s, "resource": "R", "token_type": "Bearer"}' % (
-                int(time.time()) + expires_in),
+            text='{"access_token": "%s", "expires_on": %s, "resource": "R", "token_type": "Bearer"}' % (
+                self.access_token, int(time.time()) + expires_in),
         )) as mocked_method:
             super(ServiceFabricTestCase, self)._test_happy_path(
-                app, mocked_method, expires_in)
+                app, mocked_method, expires_in, claims_challenge=claims_challenge)
+            return mocked_method
 
-    def test_happy_path(self):
-        self._test_happy_path(self.app)
+    def test_happy_path_with_client_capabilities_should_relay_capabilities(self):
+        self._test_happy_path(self._build_app(client_capabilities=["foo", "bar"])).assert_called_with(
+            'http://localhost',
+            params={
+                'api-version': '2019-07-01-preview',
+                'resource': 'R',
+                'token_sha256_to_refresh': self.access_token_sha256,
+                "xms_cc": "foo,bar",
+            },
+            headers={'Secret': 'foo'},
+        )
+
+    def test_happy_path_with_claim_challenge_should_send_sha256_to_provider(self):
+        self._test_happy_path(
+            self._build_app(client_capabilities=[]),  # Test empty client_capabilities
+            claims_challenge='{"access_token": {"nbf": {"essential": true, "value": "1563308371"}}}',
+        ).assert_called_with(
+            'http://localhost',
+            params={
+                'api-version': '2019-07-01-preview',
+                'resource': 'R',
+                'token_sha256_to_refresh': self.access_token_sha256,
+                # There is no xms_cc in this case
+            },
+            headers={'Secret': 'foo'},
+        )
 
     def test_unified_api_service_should_ignore_unnecessary_client_id(self):
         self._test_happy_path(ManagedIdentityClient(