[Core] Add CAE flag to auth policies (#31012)

pvaneck · vivekmore-msft · commit 589de01e13d6 · 2023-08-09T17:26:13.000+05:30
This first adds a keyword argument to the TokenCredential protocol method `get_token`. This enables users and client SDKs to pass in a flag to denote that `get_token` requests should be requesting CAE tokens. If the underlying credential's `get_token` implementation supports this flag, then a CAE token will be requested. Otherwise, a non-CAE token will be requested. Signed-off-by: Paul Van Eck <paulvaneck@microsoft.com> (cherry picked from commit a7519f9)
diff --git a/sdk/azureadb2c/azure-mgmt-azureadb2c/azure/mgmt/azureadb2c/models.py b/sdk/azureadb2c/azure-mgmt-azureadb2c/azure/mgmt/azureadb2c/models.py
@@ -1,8 +1,14 @@
-# coding=utf-8
-# --------------------------------------------------------------------------
+# -------------------------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from .v2019_01_01_preview.models import *
-from .v2020_05_01_preview.models import *
+from azure.core.credentials import AccessToken
+
+
+class FakeTokenCredential(object):
+    def __init__(self):
+        self.token = AccessToken("Fake Token", 0)
+
+    def get_token(self, *args, **kwargs):
+        return self.token
diff --git a/sdk/communication/azure-communication-callautomation/tests/_shared/async_fake_token_credential.py b/sdk/communication/azure-communication-callautomation/tests/_shared/async_fake_token_credential.py
@@ -10,5 +10,5 @@ class AsyncFakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    async def get_token(self, *args):
+    async def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-callautomation/tests/_shared/fake_token_credential.py b/sdk/communication/azure-communication-callautomation/tests/_shared/fake_token_credential.py
@@ -10,5 +10,5 @@ class FakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    def get_token(self, *args):
+    def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-chat/tests/_shared/async_fake_token_credential.py b/sdk/communication/azure-communication-chat/tests/_shared/async_fake_token_credential.py
@@ -10,5 +10,5 @@ class AsyncFakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    async def get_token(self, *args):
+    async def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-chat/tests/_shared/fake_token_credential.py b/sdk/communication/azure-communication-chat/tests/_shared/fake_token_credential.py
@@ -10,5 +10,5 @@ class FakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    def get_token(self, *args):
+    def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-email/tests/_shared/async_fake_token_credential.py b/sdk/communication/azure-communication-email/tests/_shared/async_fake_token_credential.py
@@ -10,5 +10,5 @@ class AsyncFakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    async def get_token(self, *args):
+    async def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-email/tests/_shared/fake_token_credential.py b/sdk/communication/azure-communication-email/tests/_shared/fake_token_credential.py
@@ -10,5 +10,5 @@ class FakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    def get_token(self, *args):
+    def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-identity/tests/_shared/async_fake_token_credential.py b/sdk/communication/azure-communication-identity/tests/_shared/async_fake_token_credential.py
@@ -10,5 +10,5 @@ class AsyncFakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    async def get_token(self, *args):
+    async def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-identity/tests/_shared/fake_token_credential.py b/sdk/communication/azure-communication-identity/tests/_shared/fake_token_credential.py
@@ -10,5 +10,5 @@ class FakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    def get_token(self, *args):
+    def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-jobrouter/tests/_shared/async_fake_token_credential.py b/sdk/communication/azure-communication-jobrouter/tests/_shared/async_fake_token_credential.py
@@ -0,0 +1,14 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+from azure.core.credentials import AccessToken
+
+
+class AsyncFakeTokenCredential(object):
+    def __init__(self):
+        self.token = AccessToken("Fake Token", 0)
+
+    async def get_token(self, *args, **kwargs):
+        return self.token
diff --git a/sdk/communication/azure-communication-networktraversal/tests/_shared/async_fake_token_credential.py b/sdk/communication/azure-communication-networktraversal/tests/_shared/async_fake_token_credential.py
@@ -10,5 +10,5 @@ class AsyncFakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    async def get_token(self, *args):
+    async def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-networktraversal/tests/_shared/fake_token_credential.py b/sdk/communication/azure-communication-networktraversal/tests/_shared/fake_token_credential.py
@@ -10,5 +10,5 @@ class FakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    def get_token(self, *args):
+    def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-phonenumbers/test/_shared/async_fake_token_credential.py b/sdk/communication/azure-communication-phonenumbers/test/_shared/async_fake_token_credential.py
@@ -10,5 +10,5 @@ class AsyncFakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    async def get_token(self, *args):
+    async def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-phonenumbers/test/_shared/fake_token_credential.py b/sdk/communication/azure-communication-phonenumbers/test/_shared/fake_token_credential.py
@@ -10,5 +10,5 @@ class FakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    def get_token(self, *args):
+    def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-rooms/tests/_shared/async_fake_token_credential.py b/sdk/communication/azure-communication-rooms/tests/_shared/async_fake_token_credential.py
@@ -10,5 +10,5 @@ class AsyncFakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    async def get_token(self, *args):
+    async def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-rooms/tests/_shared/fake_token_credential.py b/sdk/communication/azure-communication-rooms/tests/_shared/fake_token_credential.py
@@ -10,5 +10,5 @@ class FakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    def get_token(self, *args):
+    def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-sms/tests/_shared/async_fake_token_credential.py b/sdk/communication/azure-communication-sms/tests/_shared/async_fake_token_credential.py
@@ -10,5 +10,5 @@ class AsyncFakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    async def get_token(self, *args):
+    async def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-sms/tests/_shared/fake_token_credential.py b/sdk/communication/azure-communication-sms/tests/_shared/fake_token_credential.py
@@ -10,5 +10,5 @@ class FakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    def get_token(self, *args):
+    def get_token(self, *args, **kwargs):
         return self.token
diff --git a/sdk/communication/azure-communication-sms/tests/test_sms_client.py b/sdk/communication/azure-communication-sms/tests/test_sms_client.py
@@ -15,7 +15,7 @@ class FakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    def get_token(self, *args):
+    def get_token(self, *args, **kwargs):
         return self.token
 
 class TestSMSClient(unittest.TestCase):
@@ -36,7 +36,7 @@ def mock_send(*_, **__):
                     }
                 ]
             })
-            
+
         sms_client = SmsClient("https://endpoint", FakeTokenCredential(), transport=Mock(send=mock_send))
 
         sms_response = None
@@ -74,7 +74,7 @@ def test_send_message_parameters(self, mock_send):
             message=msg,
             enable_delivery_report=True,
             tag=tag)
-        
+
         send_message_request = mock_send.call_args[0][0]
         self.assertEqual(phone_number, send_message_request.from_property)
         self.assertEqual(phone_number, send_message_request.sms_recipients[0].to)
diff --git a/sdk/communication/azure-communication-sms/tests/test_sms_client_async.py b/sdk/communication/azure-communication-sms/tests/test_sms_client_async.py
@@ -15,7 +15,7 @@ class FakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("Fake Token", 0)
 
-    async def get_token(self, *args):
+    async def get_token(self, *args, **kwargs):
         return self.token
 
 class TestSMSClientAsync(aiounittest.AsyncTestCase):
@@ -58,7 +58,7 @@ async def mock_send(*_, **__):
         self.assertEqual(202, sms_response.http_status_code)
         self.assertIsNotNone(sms_response.error_message)
         self.assertTrue(sms_response.successful)
-    
+
     @patch(
         "azure.communication.sms._generated.aio.operations._sms_operations.SmsOperations.send"
     )
diff --git a/sdk/containerregistry/azure-containerregistry/tests/asynctestcase.py b/sdk/containerregistry/azure-containerregistry/tests/asynctestcase.py
@@ -25,7 +25,7 @@ class AsyncFakeTokenCredential(object):
     def __init__(self):
         self.token = AccessToken("YOU SHALL NOT PASS", 0)
 
-    async def get_token(self, *args):
+    async def get_token(self, *args, **kwargs):
         return self.token
 
 
@@ -54,7 +54,7 @@ def create_anon_client(self, endpoint, **kwargs):
         authority = get_authority(endpoint)
         audience = get_audience(authority)
         return ContainerRegistryClient(endpoint=endpoint, credential=None, audience=audience, **kwargs)
-    
+
     async def upload_oci_manifest_prerequisites(self, repo, client):
         layer = "654b93f61054e4ce90ed203bb8d556a6200d5f906cf3eca0620738d6dc18cbed"
         config = "config.json"
diff --git a/sdk/core/azure-core/CHANGELOG.md b/sdk/core/azure-core/CHANGELOG.md
@@ -1,17 +1,46 @@
 # Release History
 
-## 1.27.0 (Unreleased)
+## 1.28.1 (Unreleased)
 
 ### Features Added
 
-- Added support to use sync credentials in `AsyncBearerTokenCredentialPolicy`. #30381
+- A keyword argument `enable_cae` was added to the `get_token` method of the `TokenCredential` protocol.  #31012
+- `BearerTokenCredentialPolicy` and `AsyncBearerTokenCredentialPolicy` now accept `enable_cae` keyword arguments in their constructors. This is used in determining if [Continuous Access Evaluation (CAE)](https://learn.microsoft.com/azure/active-directory/conditional-access/concept-continuous-access-evaluation) should be enabled for each `get_token` request.  #31012
 
 ### Breaking Changes
 
 ### Bugs Fixed
 
 ### Other Changes
 
+## 1.28.0 (2023-07-06)
+
+### Features Added
+
+- Added header name parameter to `RequestIdPolicy`. #30772
+- Added `SensitiveHeaderCleanupPolicy` that cleans up sensitive headers if a redirect happens and the new destination is in another domain. #28349
+
+### Other Changes
+
+- Catch aiohttp errors and translate them into azure-core errors.
+
+## 1.27.1 (2023-06-13)
+
+### Bugs Fixed
+
+- Fix url building for some complex query parameters scenarios  #30707
+
+## 1.27.0 (2023-06-01)
+
+### Features Added
+
+- Added support to use sync credentials in `AsyncBearerTokenCredentialPolicy`. #30381
+- Added "prefix" parameter to AzureKeyCredentialPolicy #29901
+
+### Bugs Fixed
+
+- Improve error message when providing the wrong credential type for AzureKeyCredential  #30380
+
 ## 1.26.4 (2023-04-06)
 
 ### Features Added
diff --git a/sdk/core/azure-core/CLIENT_LIBRARY_DEVELOPER.md b/sdk/core/azure-core/CLIENT_LIBRARY_DEVELOPER.md
@@ -122,10 +122,10 @@ For example if you would like to alter connection pool you can initialise `Reque
  adapter = requests.adapters.HTTPAdapter(pool_connections=42, pool_maxsize=42)
  session.mount('https://', adapter)
  client = FooServiceClient(endpoint, creds, transport=RequestsTransport(session=session, session_owner=False))
- 
+
  # Here we want to manage the session by ourselves. When we are done with the session, we need to close the session.
  session.close()
- 
+
  # Note: `session_owner` gives the information of ownership of the requests sessions to the transport instance, to authorize it to close on customer's behalf. If you're ok that the client closes your session on your behalf as necessary, you don't need to pass a value.
  ```
 
@@ -565,6 +565,8 @@ class TokenCredential(Protocol):
         :keyword str claims: Additional claims required in the token, such as those returned in a resource
             provider's claims challenge following an authorization failure.
         :keyword str tenant_id: Optional tenant to include in the token request.
+        :keyword bool enable_cae: Indicates whether to enable Continuous Access Evaluation (CAE) for the requested
+            token. Defaults to False.
 
         :rtype: AccessToken
         :return: An AccessToken instance containing the token string and its expiration time in Unix time.
@@ -592,14 +594,32 @@ manager, with `__aenter__`, `__aexit__`, and `close` methods.
 
 | Service/Feature | Reason |
 | --- | --- |
-| [Continuous Access Evaluation][cae_doc] | Respond to claim challenges when unexpired tokens have access revoked
+| [Continuous Access Evaluation (CAE)][cae_doc] | Respond to claim challenges when unexpired tokens have access revoked
 
 **`tenant_id`**
 
 | Service/Feature | Reason |
 | --- | --- |
 | Key Vault ([example][kv_tenant_id]) | Request access in a tenant that was discovered as part of an authentication challenge
 
+**`enable_cae`**
+
+| Service/Feature | Reason |
+| --- | --- |
+| [Continuous Access Evaluation (CAE)][cae_doc] | Indicates that handling CAE claim challenges is supported and that a CAE-enabled token should be requested
+
+### BearerTokenCredentialPolicy and AsyncBearerTokenCredentialPolicy
+
+`BearerTokenCredentialPolicy` and `AsyncBearerTokenCredentialPolicy` are HTTP policies that are used to authenticate requests to services that accept bearer tokens in their authorization headers. These credential policies take a `TokenCredential` instance and scopes as parameters in their constructors. The `TokenCredential` instance is used to get an access token for the scopes, and the policy adds the token to the request's authorization header.
+
+Both of these policies also accept an `enable_cae` keyword argument that is passed to the `TokenCredential` instance's `get_token` method if set to `True`. This argument is used to indicate that the requested token should be [CAE-enabled][cae_doc]. If an SDK's service supports CAE, it should set this value to `True` when creating the policy.
+
+```python
+from azure.core.pipeline.policies import BearerTokenCredentialPolicy
+
+authentication_policy = BearerTokenCredentialPolicy(credential, scopes, enable_cae=True)
+```
+
 ## Long-running operation (LRO) customization
 
 See [doc/dev/customize_long_running_operation.md](https://github.com/Azure/azure-sdk-for-python/blob/main/doc/dev/customize_long_running_operation.md) for more information.
diff --git a/sdk/core/azure-core/azure/core/credentials.py b/sdk/core/azure-core/azure/core/credentials.py
@@ -33,6 +33,8 @@ def get_token(
         :keyword str claims: Additional claims required in the token, such as those returned in a resource
             provider's claims challenge following an authorization failure.
         :keyword str tenant_id: Optional tenant to include in the token request.
+        :keyword bool enable_cae: Indicates whether to enable Continuous Access Evaluation (CAE) for the requested
+            token. Defaults to False.
 
         :rtype: AccessToken
         :return: An AccessToken instance containing the token string and its expiration time in Unix time.
diff --git a/sdk/core/azure-core/azure/core/credentials_async.py b/sdk/core/azure-core/azure/core/credentials_async.py
@@ -21,6 +21,8 @@ async def get_token(
         :keyword str claims: Additional claims required in the token, such as those returned in a resource
             provider's claims challenge following an authorization failure.
         :keyword str tenant_id: Optional tenant to include in the token request.
+        :keyword bool enable_cae: Indicates whether to enable Continuous Access Evaluation (CAE) for the requested
+            token. Defaults to False.
 
         :rtype: AccessToken
         :return: An AccessToken instance containing the token string and its expiration time in Unix time.
diff --git a/sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py b/sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py
diff --git a/sdk/core/azure-core/azure/core/pipeline/policies/_authentication_async.py b/sdk/core/azure-core/azure/core/pipeline/policies/_authentication_async.py
diff --git a/sdk/core/azure-core/tests/async_tests/test_authentication_async.py b/sdk/core/azure-core/tests/async_tests/test_authentication_async.py
diff --git a/sdk/core/azure-core/tests/test_authentication.py b/sdk/core/azure-core/tests/test_authentication.py
diff --git a/sdk/core/azure-mgmt-core/tests/asynctests/test_authentication_async.py b/sdk/core/azure-mgmt-core/tests/asynctests/test_authentication_async.py
diff --git a/sdk/core/azure-mgmt-core/tests/test_authentication.py b/sdk/core/azure-mgmt-core/tests/test_authentication.py
diff --git a/sdk/formrecognizer/azure-ai-formrecognizer/tests/asynctestcase.py b/sdk/formrecognizer/azure-ai-formrecognizer/tests/asynctestcase.py
diff --git a/sdk/formrecognizer/azure-ai-formrecognizer/tests/testcase.py b/sdk/formrecognizer/azure-ai-formrecognizer/tests/testcase.py
diff --git a/tools/azure-sdk-tools/devtools_testutils/fake_credentials.py b/tools/azure-sdk-tools/devtools_testutils/fake_credentials.py
