[azure-mgmt-core] update ARMChallengeAuthenticationPolicy (Azure#38703)

* update cae policy

* update test

* update changelog

* update version

* format

* update test

* fix pylint

Author: msyyc

sdk/core/azure-mgmt-core/CHANGELOG.md

@@ -1,5 +1,11 @@
 # Release History
 
+## 1.5.1 (2023-XX-XX)
+
+### Other Changes
+
+- `ARMChallengeAuthenticationPolicy` adopt `on_challenge` in `BearerTokenCredentialPolicy` of `azure-core` to support complete CAE challenges.
+
 ## 1.5.0 (2024-10-31)
 
 ### Features Added

sdk/core/azure-mgmt-core/azure/mgmt/core/_version.py

@@ -9,4 +9,4 @@
 # regenerated.
 # --------------------------------------------------------------------------
 
-VERSION = "1.5.0"
+VERSION = "1.5.1"

sdk/core/azure-mgmt-core/azure/mgmt/core/policies/_authentication.py

@@ -23,14 +23,13 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
-import base64
 import time
 from typing import Optional, Union, MutableMapping, List, Any, Sequence, TypeVar, Generic
 
 from azure.core.credentials import AccessToken, TokenCredential
 from azure.core.credentials_async import AsyncTokenCredential
 from azure.core.pipeline.policies import BearerTokenCredentialPolicy, SansIOHTTPPolicy
-from azure.core.pipeline import PipelineRequest, PipelineResponse
+from azure.core.pipeline import PipelineRequest
 from azure.core.exceptions import ServiceRequestError
 from azure.core.pipeline.transport import (
     HttpRequest as LegacyHttpRequest,
@@ -49,33 +48,8 @@ class ARMChallengeAuthenticationPolicy(BearerTokenCredentialPolicy):
 
     This policy internally handles Continuous Access Evaluation (CAE) challenges. When it can't complete a challenge,
     it will return the 401 (unauthorized) response from ARM.
-
-    :param ~azure.core.credentials.TokenCredential credential: credential for authorizing requests
-    :param str scopes: required authentication scopes
     """
 
-    def on_challenge(
-        self,
-        request: PipelineRequest[HTTPRequestType],
-        response: PipelineResponse[HTTPRequestType, HTTPResponseType],
-    ) -> bool:
-        """Authorize request according to an ARM authentication challenge
-
-        :param ~azure.core.pipeline.PipelineRequest request: the request which elicited an authentication challenge
-        :param ~azure.core.pipeline.PipelineResponse response: ARM's response
-        :returns: a bool indicating whether the policy should send the request
-        :rtype: bool
-        """
-
-        challenge = response.http_response.headers.get("WWW-Authenticate")
-        if challenge:
-            claims = _parse_claims_challenge(challenge)
-            if claims:
-                self.authorize_request(request, *self._scopes, claims=claims)
-                return True
-
-        return False
-
 
 # pylint:disable=too-few-public-methods
 class _AuxiliaryAuthenticationPolicyBase(Generic[TokenCredentialType]):
@@ -150,33 +124,3 @@ def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
             self._aux_tokens = self._get_auxiliary_tokens(*self._scopes)
 
         self._update_headers(request.http_request.headers)
-
-
-def _parse_claims_challenge(challenge: str) -> Optional[str]:
-    """Parse the "claims" parameter from an authentication challenge
-
-    Example challenge with claims:
-        Bearer authorization_uri="https://login.windows-ppe.net/", error="invalid_token",
-        error_description="User session has been revoked",
-        claims="eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwgInZhbHVlIjoiMTYwMzc0MjgwMCJ9fX0="
-
-    :param str challenge: The authentication challenge
-    :return: the challenge's "claims" parameter or None, if it doesn't contain that parameter
-    """
-    encoded_claims = None
-    for parameter in challenge.split(","):
-        if "claims=" in parameter:
-            if encoded_claims:
-                # multiple claims challenges, e.g. for cross-tenant auth, would require special handling
-                return None
-            encoded_claims = parameter[parameter.index("=") + 1 :].strip(" \"'")
-
-    if not encoded_claims:
-        return None
-
-    padding_needed = -len(encoded_claims) % 4
-    try:
-        decoded_claims = base64.urlsafe_b64decode(encoded_claims + "=" * padding_needed).decode()
-        return decoded_claims
-    except Exception:  # pylint:disable=broad-except
-        return None

sdk/core/azure-mgmt-core/azure/mgmt/core/policies/_authentication_async.py

@@ -23,7 +23,7 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
-from typing import cast, Awaitable, Optional, List, Union, Any
+from typing import Awaitable, Optional, List, Union, Any
 import inspect
 
 from azure.core.pipeline.policies import (
@@ -40,7 +40,7 @@
 from azure.core.credentials_async import AsyncTokenCredential
 
 
-from ._authentication import _parse_claims_challenge, _AuxiliaryAuthenticationPolicyBase
+from ._authentication import _AuxiliaryAuthenticationPolicyBase
 
 
 HTTPRequestType = Union[LegacyHttpRequest, HttpRequest]
@@ -66,33 +66,8 @@ class AsyncARMChallengeAuthenticationPolicy(AsyncBearerTokenCredentialPolicy):
 
     This policy internally handles Continuous Access Evaluation (CAE) challenges. When it can't complete a challenge,
     it will return the 401 (unauthorized) response from ARM.
-
-    :param ~azure.core.credentials.TokenCredential credential: credential for authorizing requests
-    :param str scopes: required authentication scopes
     """
 
-    # pylint:disable=unused-argument
-    async def on_challenge(
-        self,
-        request: PipelineRequest[HTTPRequestType],
-        response: PipelineResponse[HTTPRequestType, AsyncHTTPResponseType],
-    ) -> bool:
-        """Authorize request according to an ARM authentication challenge
-
-        :param ~azure.core.pipeline.PipelineRequest request: the request which elicited an authentication challenge
-        :param ~azure.core.pipeline.PipelineResponse response: the resource provider's response
-        :returns: a bool indicating whether the policy should send the request
-        :rtype: bool
-        """
-        # Casting, as the code seems to be certain that on_challenge this header will be present
-        challenge: str = cast(str, response.http_response.headers.get("WWW-Authenticate"))
-        claims = _parse_claims_challenge(challenge)
-        if claims:
-            await self.authorize_request(request, *self._scopes, claims=claims)
-            return True
-
-        return False
-
 
 class AsyncAuxiliaryAuthenticationPolicy(
     _AuxiliaryAuthenticationPolicyBase[AsyncTokenCredential],

sdk/core/azure-mgmt-core/setup.py

@@ -69,7 +69,7 @@
         "pytyped": ["py.typed"],
     },
     install_requires=[
-        "azure-core>=1.31.0",
+        "azure-core>=1.32.0",
     ],
     python_requires=">=3.8",
 )

sdk/core/azure-mgmt-core/tests/asynctests/test_authentication_async.py

@@ -50,7 +50,7 @@ async def test_claims_challenge():
     expected_claims = '{"access_token": {"essential": "true"}'
     expected_scope = "scope"
 
-    challenge = 'Bearer authorization_uri="https://localhost", error=".", error_description=".", claims="{}"'.format(
+    challenge = 'Bearer authorization_uri="https://localhost", error="insufficient_claims", error_description=".", claims="{}"'.format(
         base64.b64encode(expected_claims.encode()).decode()
     )
     responses = (
@@ -97,34 +97,68 @@ async def get_token(*scopes, **kwargs):
 
 
 async def test_multiple_claims_challenges():
-    """ARMChallengeAuthenticationPolicy should not attempt to handle a response having multiple claims challenges"""
+    """ARMChallengeAuthenticationPolicy handle a response having multiple claims challenges"""
+    first_token = AccessToken("first", int(time.time()) + 3600)
+    second_token = AccessToken("second", int(time.time()) + 3600)
+    tokens = (t for t in (first_token, second_token))
+
+    expected_claims = '{"access_token": {"essential": "true"}'
+    expected_scope = "scope"
+
+    claims = base64.b64encode(expected_claims.encode()).decode()
 
     expected_header = ",".join(
         (
-            'Bearer realm="", authorization_uri="https://login.microsoftonline.com/common/oauth2/authorize", client_id="00000003-0000-0000-c000-000000000000", error="insufficient_claims", claims="eyJhY2Nlc3NfdG9rZW4iOiB7ImZvbyI6ICJiYXIifX0="',
-            'Bearer authorization_uri="https://login.windows-ppe.net/", error="invalid_token", error_description="User session has been revoked", claims="eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwgInZhbHVlIjoiMTYwMzc0MjgwMCJ9fX0="',
+            'Bearer realm="", authorization_uri="https://localhost", client_id="00", error="insufficient_claims", claims="{}"'.format(
+                claims
+            ),
+            'Bearer authorization_uri="https://login.windows-ppe.net/", error="invalid_token", error_description="User session has been revoked", claims="{}"'.format(
+                claims
+            ),
+        )
+    )
+
+    responses = (
+        r
+        for r in (
+            Mock(status_code=401, headers={"WWW-Authenticate": expected_header}),
+            Mock(status_code=200),
         )
     )
 
     async def send(request):
-        return Mock(status_code=401, headers={"WWW-Authenticate": expected_header})
+        res = next(responses)
+        if res.status_code == 401:
+            expected_token = first_token.token
+        else:
+            expected_token = second_token.token
+        assert request.headers["Authorization"] == "Bearer " + expected_token
 
-    async def get_token(*_, **__):
-        return AccessToken("***", 42)
+        return res
+
+    async def get_token(*scopes, **kwargs):
+        assert scopes == (expected_scope,)
+        return next(tokens)
 
-    transport = Mock(send=Mock(wraps=send))
     credential = Mock(spec_set=["get_token"], get_token=Mock(wraps=get_token))
-    policies = [AsyncARMChallengeAuthenticationPolicy(credential, "scope")]
+    transport = Mock(send=Mock(wraps=send))
+    policies = [AsyncARMChallengeAuthenticationPolicy(credential, expected_scope)]
     pipeline = AsyncPipeline(transport=transport, policies=policies)
 
     response = await pipeline.run(HttpRequest("GET", "https://localhost"))
 
-    assert transport.send.call_count == 1
-    assert credential.get_token.call_count == 1
+    assert response.http_response.status_code == 200
+    assert transport.send.call_count == 2
+    assert credential.get_token.call_count == 2
+
+    args, kwargs = credential.get_token.call_args
+    assert expected_scope in args
+    assert kwargs["claims"] == expected_claims
 
-    # the policy should have returned the error response because it was unable to handle the challenge
-    assert response.http_response.status_code == 401
-    assert response.http_response.headers["WWW-Authenticate"] == expected_header
+    with pytest.raises(StopIteration):
+        next(tokens)
+    with pytest.raises(StopIteration):
+        next(responses)
 
 
 async def test_auxiliary_authentication_policy():

sdk/core/azure-mgmt-core/tests/test_authentication.py

@@ -29,7 +29,6 @@
 from azure.core.credentials import AccessToken
 from azure.core.pipeline import Pipeline
 from azure.mgmt.core.policies._authentication import (
-    _parse_claims_challenge,
     ARMChallengeAuthenticationPolicy,
     AuxiliaryAuthenticationPolicy,
 )
@@ -48,36 +47,6 @@
 CLAIM_IP = base64.b64encode(ip_claim).decode()[:-2]  # Trim off padding = characters
 
 
-@pytest.mark.parametrize(
-    "challenge,expected_claims",
-    (
-        # CAE - insufficient claims
-        (
-            f'Bearer realm="", authorization_uri="https://login.microsoftonline.com/common/oauth2/authorize", client_id="00000003-0000-0000-c000-000000000000", error="insufficient_claims", claims="{CLAIM_TOKEN}"',
-            '{"access_token": {"foo": "bar"}}',
-        ),
-        # CAE - sessions revoked
-        (
-            f'Bearer authorization_uri="https://login.windows-ppe.net/", error="invalid_token", error_description="User session has been revoked", claims={CLAIM_NBF}',
-            '{"access_token":{"nbf":{"essential":true, "value":"1603742800"}}}',
-        ),
-        # CAE - IP policy
-        (
-            f'Bearer authorization_uri="https://login.windows.net/", error="invalid_token", error_description="Tenant IP Policy validate failed.", claims={CLAIM_IP}',
-            '{"access_token":{"nbf":{"essential":true,"value":"1610563006"},"xms_rp_ipaddr":{"value":"1.2.3.4"}}}',
-        ),
-        # ARM
-        (
-            'Bearer authorization_uri="https://login.windows.net/", error="invalid_token", error_description="The authentication failed because of missing \'Authorization\' header."',
-            None,
-        ),
-    ),
-)
-def test_challenge_parsing(challenge, expected_claims):
-    claims = _parse_claims_challenge(challenge)
-    assert claims == expected_claims
-
-
 def test_auxiliary_authentication_policy():
     """The auxiliary authentication policy should add a header containing a token from its credential"""
     first_token = AccessToken("first", int(time.time()) + 3600)
@@ -119,7 +88,7 @@ def test_claims_challenge():
     expected_claims = '{"access_token": {"essential": "true"}'
     expected_scope = "scope"
 
-    challenge = 'Bearer authorization_uri="https://localhost", error=".", error_description=".", claims="{}"'.format(
+    challenge = 'Bearer authorization_uri="https://localhost", error="insufficient_claims", error_description=".", claims="{}"'.format(
         base64.b64encode(expected_claims.encode()).decode()
     )
     responses = (
@@ -166,28 +135,66 @@ def get_token(*scopes, **kwargs):
 
 
 def test_multiple_claims_challenges():
-    """ARMChallengeAuthenticationPolicy should not attempt to handle a response having multiple claims challenges"""
+    """ARMChallengeAuthenticationPolicy handle a response having multiple claims challenges"""
+
+    first_token = AccessToken("first", int(time.time()) + 3600)
+    second_token = AccessToken("second", int(time.time()) + 3600)
+    tokens = (t for t in (first_token, second_token))
+
+    expected_claims = '{"access_token": {"essential": "true"}'
+    expected_scope = "scope"
+
+    claims = base64.b64encode(expected_claims.encode()).decode()
 
     expected_header = ",".join(
         (
-            'Bearer realm="", authorization_uri="https://login.microsoftonline.com/common/oauth2/authorize", client_id="00000003-0000-0000-c000-000000000000", error="insufficient_claims", claims="eyJhY2Nlc3NfdG9rZW4iOiB7ImZvbyI6ICJiYXIifX0="',
-            'Bearer authorization_uri="https://login.windows-ppe.net/", error="invalid_token", error_description="User session has been revoked", claims="eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwgInZhbHVlIjoiMTYwMzc0MjgwMCJ9fX0="',
+            'Bearer realm="", authorization_uri="https://localhost", client_id="00", error="insufficient_claims", claims="{}"'.format(
+                claims
+            ),
+            'Bearer authorization_uri="https://login.windows-ppe.net/", error="invalid_token", error_description="User session has been revoked", claims="{}"'.format(
+                claims
+            ),
+        )
+    )
+
+    responses = (
+        r
+        for r in (
+            Mock(status_code=401, headers={"WWW-Authenticate": expected_header}),
+            Mock(status_code=200),
         )
     )
 
     def send(request):
-        return Mock(status_code=401, headers={"WWW-Authenticate": expected_header})
+        res = next(responses)
+        if res.status_code == 401:
+            expected_token = first_token.token
+        else:
+            expected_token = second_token.token
+        assert request.headers["Authorization"] == "Bearer " + expected_token
+
+        return res
 
+    def get_token(*scopes, **kwargs):
+        assert scopes == (expected_scope,)
+        return next(tokens)
+
+    credential = Mock(spec_set=["get_token"], get_token=Mock(wraps=get_token))
     transport = Mock(send=Mock(wraps=send))
-    credential = FakeTokenCredential()
-    policies = [ARMChallengeAuthenticationPolicy(credential, "scope")]
+    policies = [ARMChallengeAuthenticationPolicy(credential, expected_scope)]
     pipeline = Pipeline(transport=transport, policies=policies)
 
     response = pipeline.run(HttpRequest("GET", "https://localhost"))
 
-    assert transport.send.call_count == 1
-    assert credential.get_token_count == 1
+    assert response.http_response.status_code == 200
+    assert transport.send.call_count == 2
+    assert credential.get_token.call_count == 2
 
-    # the policy should have returned the error response because it was unable to handle the challenge
-    assert response.http_response.status_code == 401
-    assert response.http_response.headers["WWW-Authenticate"] == expected_header
+    args, kwargs = credential.get_token.call_args
+    assert expected_scope in args
+    assert kwargs["claims"] == expected_claims
+
+    with pytest.raises(StopIteration):
+        next(tokens)
+    with pytest.raises(StopIteration):
+        next(responses)