CAE support for azure-mgmt-core (#19365)

* update azure-core requirement

* package metadata

* tests

* add challenge auth policies

* better calculation of base64 padding

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

@@ -1,8 +1,22 @@
-
 # Release History
 
-## 1.2.3 (Unreleased)
+## 1.3.0b3 (2021-06-07)
+
+### Changed
+
+- Updated required `azure-core` version
+
+## 1.3.0b2 (2021-05-13)
+
+### Changed
+
+- Updated required `azure-core` version
+
+## 1.3.0b1 (2021-03-10)
+
+### Features
 
+- ARMChallengeAuthenticationPolicy supports bearer token authorization and CAE challenges
 
 ## 1.2.2 (2020-11-09)
 

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

@@ -9,4 +9,4 @@
 # regenerated.
 # --------------------------------------------------------------------------
 
-VERSION = "1.2.3"
+VERSION = "1.3.0b3"

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

@@ -25,6 +25,7 @@
 # --------------------------------------------------------------------------
 
 from azure.core.pipeline.policies import HttpLoggingPolicy
+from ._authentication import ARMChallengeAuthenticationPolicy
 from ._base import ARMAutoResourceProviderRegistrationPolicy
 
 
@@ -48,13 +49,13 @@ class ARMHttpLoggingPolicy(HttpLoggingPolicy):
     ])
 
 
-__all__ = ["ARMAutoResourceProviderRegistrationPolicy", "ARMHttpLoggingPolicy"]
+__all__ = ["ARMAutoResourceProviderRegistrationPolicy", "ARMChallengeAuthenticationPolicy", "ARMHttpLoggingPolicy"]
 
 try:
-    from ._base_async import (  # pylint: disable=unused-import
-        AsyncARMAutoResourceProviderRegistrationPolicy,
-    )
+    # pylint: disable=unused-import
+    from ._authentication_async import AsyncARMChallengeAuthenticationPolicy
+    from ._base_async import AsyncARMAutoResourceProviderRegistrationPolicy
 
-    __all__.extend(["AsyncARMAutoResourceProviderRegistrationPolicy"])
+    __all__.extend(["AsyncARMAutoResourceProviderRegistrationPolicy", "AsyncARMChallengeAuthenticationPolicy"])
 except (ImportError, SyntaxError):
     pass  # Async not supported

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

@@ -0,0 +1,92 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+import base64
+from typing import TYPE_CHECKING
+
+from azure.core.pipeline.policies import BearerTokenCredentialPolicy
+
+if TYPE_CHECKING:
+    from typing import Optional
+    from azure.core.pipeline import PipelineRequest, PipelineResponse
+
+
+class ARMChallengeAuthenticationPolicy(BearerTokenCredentialPolicy):
+    """Adds a bearer token Authorization header to requests.
+
+    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, response):  # pylint:disable=unused-argument
+        # type: (PipelineRequest, PipelineResponse) -> 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
+        """
+
+        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
+
+
+def _parse_claims_challenge(challenge):
+    # type: (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="
+
+    :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

@@ -0,0 +1,61 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+from typing import TYPE_CHECKING
+
+from azure.core.pipeline.policies import AsyncBearerTokenCredentialPolicy
+
+from ._authentication import _parse_claims_challenge
+
+if TYPE_CHECKING:
+    from azure.core.pipeline import PipelineRequest, PipelineResponse
+
+
+class AsyncARMChallengeAuthenticationPolicy(AsyncBearerTokenCredentialPolicy):
+    """Adds a bearer token Authorization header to requests.
+
+    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", response: "PipelineResponse") -> 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
+        """
+
+        challenge = 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

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

@@ -45,7 +45,7 @@
     author_email='azpysdkhelp@microsoft.com',
     url='https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/core/azure-mgmt-core',
     classifiers=[
-        "Development Status :: 5 - Production/Stable",
+        "Development Status :: 4 - Beta",
         'Programming Language :: Python',
         'Programming Language :: Python :: 2',
         'Programming Language :: Python :: 2.7',
@@ -68,7 +68,7 @@
         'pytyped': ['py.typed'],
     },
     install_requires=[
-        "azure-core<2.0.0,>=1.13.0",
+        "azure-core<2.0.0,>=1.15.0",
     ],
     extras_require={
         ":python_version<'3.0'": ['azure-mgmt-nspkg'],

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

@@ -0,0 +1,114 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+import base64
+import time
+from unittest.mock import Mock
+
+from azure.core.credentials import AccessToken
+from azure.core.pipeline import AsyncPipeline
+from azure.mgmt.core.policies import AsyncARMChallengeAuthenticationPolicy
+from azure.core.pipeline.transport import HttpRequest
+
+import pytest
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_claims_challenge():
+    """AsyncAsyncARMChallengeAuthenticationPolicy should pass claims from an authentication challenge to its credential"""
+
+    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"
+
+    challenge = 'Bearer authorization_uri="https://localhost", error=".", error_description=".", claims="{}"'.format(
+        base64.b64encode(expected_claims.encode()).decode()
+    )
+    responses = (r for r in (Mock(status_code=401, headers={"WWW-Authenticate": challenge}), Mock(status_code=200)))
+
+    async def send(request):
+        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
+
+    async def get_token(*scopes, **kwargs):
+        assert scopes == (expected_scope,)
+        return next(tokens)
+
+    credential = Mock(get_token=Mock(wraps=get_token))
+    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 response.http_response.status_code == 200
+    assert transport.send.call_count == 2
+    assert credential.get_token.call_count == 2
+    credential.get_token.assert_called_with(expected_scope, claims=expected_claims)
+    with pytest.raises(StopIteration):
+        next(tokens)
+    with pytest.raises(StopIteration):
+        next(responses)
+
+
+async def test_multiple_claims_challenges():
+    """ARMChallengeAuthenticationPolicy should not attempt to handle a response having multiple claims challenges"""
+
+    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="',
+        )
+    )
+
+    async def send(request):
+        return Mock(status_code=401, headers={"WWW-Authenticate": expected_header})
+
+    async def get_token(*_, **__):
+        return AccessToken("***", 42)
+
+    transport = Mock(send=Mock(wraps=send))
+    credential = Mock(get_token=Mock(wraps=get_token))
+    policies = [AsyncARMChallengeAuthenticationPolicy(credential, "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
+
+    # 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