Client capabilities (#240)

Author: abhidnya13

msal/application.py

@@ -79,6 +79,17 @@ def extract_certs(public_cert_content):
     return [public_cert_content.strip()]
 
 
+def _merge_claims_challenge_and_capabilities(capabilities, claims_challenge):
+    # Represent capabilities as {"access_token": {"xms_cc": {"values": capabilities}}}
+    # and then merge/add it into incoming claims
+    if not capabilities:
+        return claims_challenge
+    claims_dict = json.loads(claims_challenge) if claims_challenge else {}
+    for key in ["access_token"]:  # We could add "id_token" if we'd decide to
+        claims_dict.setdefault(key, {}).update(xms_cc={"values": capabilities})
+    return json.dumps(claims_dict)
+
+
 class ClientApplication(object):
 
     ACQUIRE_TOKEN_SILENT_ID = "84"
@@ -97,7 +108,8 @@ def __init__(
             token_cache=None,
             http_client=None,
             verify=True, proxies=None, timeout=None,
-            client_claims=None, app_name=None, app_version=None):
+            client_claims=None, app_name=None, app_version=None,
+            client_capabilities=None):
         """Create an instance of application.
 
         :param str client_id: Your app has a client_id after you register it on AAD.
@@ -179,10 +191,16 @@ def __init__(
         :param app_version: (optional)
             You can provide your application version for Microsoft telemetry purposes.
             Default value is None, means it will not be passed to Microsoft.
+        :param list[str] client_capabilities: (optional)
+            Allows configuration of one or more client capabilities, e.g. ["CP1"].
+            MSAL will combine them into
+            `claims parameter <https://openid.net/specs/openid-connect-core-1_0-final.html#ClaimsParameter`_
+            which you will later provide via one of the acquire-token request.
         """
         self.client_id = client_id
         self.client_credential = client_credential
         self.client_claims = client_claims
+        self._client_capabilities = client_capabilities
         if http_client:
             self.http_client = http_client
         else:
@@ -261,6 +279,7 @@ def get_authorization_request_url(
             prompt=None,
             nonce=None,
             domain_hint=None,  # type: Optional[str]
+            claims_challenge=None,
             **kwargs):
         """Constructs a URL for you to start a Authorization Code Grant.
 
@@ -289,6 +308,12 @@ def get_authorization_request_url(
             More information on possible values
             `here <https://docs.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-auth-code-flow#request-an-authorization-code>`_ and
             `here <https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-oapx/86fb452d-e34a-494e-ac61-e526e263b6d8>`_.
+        :param claims_challenge:
+             The claims_challenge parameter requests specific claims requested by the resource provider
+             in the form of a claims_challenge directive in the www-authenticate header to be
+             returned from the UserInfo Endpoint and/or in the ID Token and/or Access Token.
+             It is a string of a JSON object which contains lists of claims being requested from these locations.
+
         :return: The authorization url as a string.
         """
         """ # TBD: this would only be meaningful in a new acquire_token_interactive()
@@ -321,6 +346,8 @@ def get_authorization_request_url(
             scope=decorate_scope(scopes, self.client_id),
             nonce=nonce,
             domain_hint=domain_hint,
+            claims=_merge_claims_challenge_and_capabilities(
+                self._client_capabilities, claims_challenge),
             )
 
     def acquire_token_by_authorization_code(
@@ -332,6 +359,7 @@ def acquire_token_by_authorization_code(
                 # authorization request as described in Section 4.1.1, and their
                 # values MUST be identical.
             nonce=None,
+            claims_challenge=None,
             **kwargs):
         """The second half of the Authorization Code Grant.
 
@@ -357,6 +385,12 @@ def acquire_token_by_authorization_code(
             same nonce should also be provided here, so that we'll validate it.
             An exception will be raised if the nonce in id token mismatches.
 
+        :param claims_challenge:
+            The claims_challenge parameter requests specific claims requested by the resource provider
+            in the form of a claims_challenge directive in the www-authenticate header to be
+            returned from the UserInfo Endpoint and/or in the ID Token and/or Access Token.
+            It is a string of a JSON object which contains lists of claims being requested from these locations.
+
         :return: A dict representing the json response from AAD:
 
             - A successful response would contain "access_token" key,
@@ -377,6 +411,10 @@ def acquire_token_by_authorization_code(
                 CLIENT_CURRENT_TELEMETRY: _build_current_telemetry_request_header(
                     self.ACQUIRE_TOKEN_BY_AUTHORIZATION_CODE_ID),
                 },
+            data=dict(
+                kwargs.pop("data", {}),
+                claims=_merge_claims_challenge_and_capabilities(
+                    self._client_capabilities, claims_challenge)),
             nonce=nonce,
             **kwargs)
 
@@ -479,6 +517,7 @@ def acquire_token_silent(
             account,  # type: Optional[Account]
             authority=None,  # See get_authorization_request_url()
             force_refresh=False,  # type: Optional[boolean]
+            claims_challenge=None,
             **kwargs):
         """Acquire an access token for given account, without user interaction.
 
@@ -493,14 +532,21 @@ def acquire_token_silent(
 
         Internally, this method calls :func:`~acquire_token_silent_with_error`.
 
+        :param claims_challenge:
+            The claims_challenge parameter requests specific claims requested by the resource provider
+            in the form of a claims_challenge directive in the www-authenticate header to be
+            returned from the UserInfo Endpoint and/or in the ID Token and/or Access Token.
+            It is a string of a JSON object which contains lists of claims being requested from these locations.
+
         :return:
             - A dict containing no "error" key,
               and typically contains an "access_token" key,
               if cache lookup succeeded.
             - None when cache lookup does not yield a token.
         """
         result = self.acquire_token_silent_with_error(
-            scopes, account, authority, force_refresh, **kwargs)
+            scopes, account, authority, force_refresh,
+            claims_challenge=claims_challenge, **kwargs)
         return result if result and "error" not in result else None
 
     def acquire_token_silent_with_error(
@@ -509,6 +555,7 @@ def acquire_token_silent_with_error(
             account,  # type: Optional[Account]
             authority=None,  # See get_authorization_request_url()
             force_refresh=False,  # type: Optional[boolean]
+            claims_challenge=None,
             **kwargs):
         """Acquire an access token for given account, without user interaction.
 
@@ -529,6 +576,11 @@ def acquire_token_silent_with_error(
         :param force_refresh:
             If True, it will skip Access Token look-up,
             and try to find a Refresh Token to obtain a new Access Token.
+        :param claims_challenge:
+            The claims_challenge parameter requests specific claims requested by the resource provider
+            in the form of a claims_challenge directive in the www-authenticate header to be
+            returned from the UserInfo Endpoint and/or in the ID Token and/or Access Token.
+            It is a string of a JSON object which contains lists of claims being requested from these locations.
         :return:
             - A dict containing no "error" key,
               and typically contains an "access_token" key,
@@ -547,6 +599,7 @@ def acquire_token_silent_with_error(
         #     ) if authority else self.authority
         result = self._acquire_token_silent_from_cache_and_possibly_refresh_it(
             scopes, account, self.authority, force_refresh=force_refresh,
+            claims_challenge=claims_challenge,
             correlation_id=correlation_id,
             **kwargs)
         if result and "error" not in result:
@@ -567,6 +620,7 @@ def acquire_token_silent_with_error(
                 validate_authority=False)
             result = self._acquire_token_silent_from_cache_and_possibly_refresh_it(
                 scopes, account, the_authority, force_refresh=force_refresh,
+                claims_challenge=claims_challenge,
                 correlation_id=correlation_id,
                 **kwargs)
             if result:
@@ -589,8 +643,9 @@ def _acquire_token_silent_from_cache_and_possibly_refresh_it(
             account,  # type: Optional[Account]
             authority,  # This can be different than self.authority
             force_refresh=False,  # type: Optional[boolean]
+            claims_challenge=None,
             **kwargs):
-        if not force_refresh:
+        if not (force_refresh or claims_challenge):  # Bypass AT when desired or using claims
             query={
                     "client_id": self.client_id,
                     "environment": authority.instance,
@@ -617,7 +672,7 @@ def _acquire_token_silent_from_cache_and_possibly_refresh_it(
                     }
         return self._acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
                 authority, decorate_scope(scopes, self.client_id), account,
-                force_refresh=force_refresh, **kwargs)
+                force_refresh=force_refresh, claims_challenge=claims_challenge, **kwargs)
 
     def _acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
             self, authority, scopes, account, **kwargs):
@@ -666,7 +721,7 @@ def _get_app_metadata(self, environment):
     def _acquire_token_silent_by_finding_specific_refresh_token(
             self, authority, scopes, query,
             rt_remover=None, break_condition=lambda response: False,
-            force_refresh=False, correlation_id=None, **kwargs):
+            force_refresh=False, correlation_id=None, claims_challenge=None, **kwargs):
         matches = self.token_cache.find(
             self.token_cache.CredentialType.REFRESH_TOKEN,
             # target=scopes,  # AAD RTs are scope-independent
@@ -686,6 +741,10 @@ def _acquire_token_silent_by_finding_specific_refresh_token(
                     CLIENT_CURRENT_TELEMETRY: _build_current_telemetry_request_header(
                         self.ACQUIRE_TOKEN_SILENT_ID, force_refresh=force_refresh),
                     },
+                data=dict(
+                    kwargs.pop("data", {}),
+                    claims=_merge_claims_challenge_and_capabilities(
+                        self._client_capabilities, claims_challenge)),
                 **kwargs)
             if "error" not in response:
                 return response
@@ -780,14 +839,19 @@ def initiate_device_flow(self, scopes=None, **kwargs):
         flow[self.DEVICE_FLOW_CORRELATION_ID] = correlation_id
         return flow
 
-    def acquire_token_by_device_flow(self, flow, **kwargs):
+    def acquire_token_by_device_flow(self, flow, claims_challenge=None, **kwargs):
         """Obtain token by a device flow object, with customizable polling effect.
 
         :param dict flow:
             A dict previously generated by :func:`~initiate_device_flow`.
             By default, this method's polling effect  will block current thread.
             You can abort the polling loop at any time,
             by changing the value of the flow's "expires_at" key to 0.
+        :param claims_challenge:
+            The claims_challenge parameter requests specific claims requested by the resource provider
+            in the form of a claims_challenge directive in the www-authenticate header to be
+            returned from the UserInfo Endpoint and/or in the ID Token and/or Access Token.
+            It is a string of a JSON object which contains lists of claims being requested from these locations.
 
         :return: A dict representing the json response from AAD:
 
@@ -796,10 +860,14 @@ def acquire_token_by_device_flow(self, flow, **kwargs):
         """
         return self.client.obtain_token_by_device_flow(
             flow,
-            data=dict(kwargs.pop("data", {}), code=flow["device_code"]),
-                # 2018-10-4 Hack:
-                # during transition period,
-                # service seemingly need both device_code and code parameter.
+            data=dict(
+                kwargs.pop("data", {}),
+                code=flow["device_code"],  # 2018-10-4 Hack:
+                    # during transition period,
+                    # service seemingly need both device_code and code parameter.
+                claims=_merge_claims_challenge_and_capabilities(
+                    self._client_capabilities, claims_challenge),
+                ),
             headers={
                 CLIENT_REQUEST_ID:
                     flow.get(self.DEVICE_FLOW_CORRELATION_ID) or _get_new_correlation_id(),
@@ -809,7 +877,7 @@ def acquire_token_by_device_flow(self, flow, **kwargs):
             **kwargs)
 
     def acquire_token_by_username_password(
-            self, username, password, scopes, **kwargs):
+            self, username, password, scopes, claims_challenge=None, **kwargs):
         """Gets a token for a given resource via user credentials.
 
         See this page for constraints of Username Password Flow.
@@ -819,6 +887,11 @@ def acquire_token_by_username_password(
         :param str password: The password.
         :param list[str] scopes:
             Scopes requested to access a protected API (a resource).
+        :param claims_challenge:
+            The claims_challenge parameter requests specific claims requested by the resource provider
+            in the form of a claims_challenge directive in the www-authenticate header to be
+            returned from the UserInfo Endpoint and/or in the ID Token and/or Access Token.
+            It is a string of a JSON object which contains lists of claims being requested from these locations.
 
         :return: A dict representing the json response from AAD:
 
@@ -831,16 +904,22 @@ def acquire_token_by_username_password(
             CLIENT_CURRENT_TELEMETRY: _build_current_telemetry_request_header(
                 self.ACQUIRE_TOKEN_BY_USERNAME_PASSWORD_ID),
             }
+        data = dict(
+            kwargs.pop("data", {}),
+            claims=_merge_claims_challenge_and_capabilities(
+                self._client_capabilities, claims_challenge))
         if not self.authority.is_adfs:
             user_realm_result = self.authority.user_realm_discovery(
                 username, correlation_id=headers[CLIENT_REQUEST_ID])
             if user_realm_result.get("account_type") == "Federated":
                 return self._acquire_token_by_username_password_federated(
                     user_realm_result, username, password, scopes=scopes,
+                    data=data,
                     headers=headers, **kwargs)
         return self.client.obtain_token_by_username_password(
                 username, password, scope=scopes,
                 headers=headers,
+                data=data,
                 **kwargs)
 
     def _acquire_token_by_username_password_federated(
@@ -882,11 +961,16 @@ def _acquire_token_by_username_password_federated(
 
 class ConfidentialClientApplication(ClientApplication):  # server-side web app
 
-    def acquire_token_for_client(self, scopes, **kwargs):
+    def acquire_token_for_client(self, scopes, claims_challenge=None, **kwargs):
         """Acquires token for the current confidential client, not for an end user.
 
         :param list[str] scopes: (Required)
             Scopes requested to access a protected API (a resource).
+        :param claims_challenge:
+            The claims_challenge parameter requests specific claims requested by the resource provider
+            in the form of a claims_challenge directive in the www-authenticate header to be
+            returned from the UserInfo Endpoint and/or in the ID Token and/or Access Token.
+            It is a string of a JSON object which contains lists of claims being requested from these locations.
 
         :return: A dict representing the json response from AAD:
 
@@ -901,9 +985,13 @@ def acquire_token_for_client(self, scopes, **kwargs):
                 CLIENT_CURRENT_TELEMETRY: _build_current_telemetry_request_header(
                     self.ACQUIRE_TOKEN_FOR_CLIENT_ID),
                 },
+            data=dict(
+                kwargs.pop("data", {}),
+                claims=_merge_claims_challenge_and_capabilities(
+                    self._client_capabilities, claims_challenge)),
             **kwargs)
 
-    def acquire_token_on_behalf_of(self, user_assertion, scopes, **kwargs):
+    def acquire_token_on_behalf_of(self, user_assertion, scopes, claims_challenge=None, **kwargs):
         """Acquires token using on-behalf-of (OBO) flow.
 
         The current app is a middle-tier service which was called with a token
@@ -918,6 +1006,11 @@ def acquire_token_on_behalf_of(self, user_assertion, scopes, **kwargs):
 
         :param str user_assertion: The incoming token already received by this app
         :param list[str] scopes: Scopes required by downstream API (a resource).
+        :param claims_challenge:
+            The claims_challenge parameter requests specific claims requested by the resource provider
+            in the form of a claims_challenge directive in the www-authenticate header to be
+            returned from the UserInfo Endpoint and/or in the ID Token and/or Access Token.
+            It is a string of a JSON object which contains lists of claims being requested from these locations..
 
         :return: A dict representing the json response from AAD:
 
@@ -935,7 +1028,11 @@ def acquire_token_on_behalf_of(self, user_assertion, scopes, **kwargs):
                 # 2. Requesting an IDT (which would otherwise be unavailable)
                 #    so that the calling app could use id_token_claims to implement
                 #    their own cache mapping, which is likely needed in web apps.
-            data=dict(kwargs.pop("data", {}), requested_token_use="on_behalf_of"),
+            data=dict(
+                kwargs.pop("data", {}),
+                requested_token_use="on_behalf_of",
+                claims=_merge_claims_challenge_and_capabilities(
+                    self._client_capabilities, claims_challenge)),
             headers={
                 CLIENT_REQUEST_ID: _get_new_correlation_id(),
                 CLIENT_CURRENT_TELEMETRY: _build_current_telemetry_request_header(