Add AAD support for EG

sdk/eventgrid/azure-eventgrid/CHANGELOG.md

@@ -1,11 +1,13 @@
 # Release History
 
-## 4.3.1 (Unreleased)
+## 4.4.0 (Unreleased)
 
 - Bumped `msrest` dependency to `0.6.21` to align with mgmt package.
 
 ### Features Added
 
+- `EventGridPublisherClient` now supports Azure Active Directory (AAD) for authentication.
+
 ### Breaking Changes
 
 ### Key Bugs Fixed

sdk/eventgrid/azure-eventgrid/README.md

@@ -38,6 +38,34 @@ az eventgrid domain --create --location <location> --resource-group <resource-gr
 In order to interact with the Event Grid service, you will need to create an instance of a client.
 An **endpoint** and **credential** are necessary to instantiate the client object.
 
+#### Using Azure Active Directory (AAD)
+
+Azure Event Grid provides integration with Azure Active Directory (Azure AD) for identity-based authentication of requests. With Azure AD, you can use role-based access control (RBAC) to grant access to your Azure Event Grid resources to users, groups, or applications.
+
+To send events to a topic or domain with a `TokenCredential`, the authenticated identity should have the "EventGrid Data Sender" role assigned.
+
+With the `azure-identity` package, you can seamlessly authorize requests in both development and production environments. To learn more about Azure Active Directory, see the [`azure-identity` README](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/identity/azure-identity/README.md).
+
+For example, you can use `DefaultAzureCredential` to construct a client which will authenticate using Azure Active Directory:
+
+```Python
+from azure.identity import DefaultAzureCredential
+from azure.eventgrid import EventGridPublisherClient, EventGridEvent
+
+event = EventGridEvent(
+    data={"team": "azure-sdk"},
+    subject="Door1",
+    event_type="Azure.Sdk.Demo",
+    data_version="2.0"
+)
+
+credential = DefaultAzureCredential()
+endpoint = os.environ["EG_TOPIC_HOSTNAME"]
+client = EventGridPublisherClient(endpoint, credential)
+
+client.send(event)
+```
+
 #### Looking up the endpoint
 You can find the topic endpoint within the Event Grid Topic resource on the Azure portal. This will look like:
 `"https://<event-grid-topic-name>.<topic-location>.eventgrid.azure.net/api/events"`

sdk/eventgrid/azure-eventgrid/azure/eventgrid/_constants.py

@@ -3,6 +3,7 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+DEFAULT_EVENTGRID_SCOPE = "https://eventgrid.azure.net/.default"
 EVENTGRID_KEY_HEADER = "aeg-sas-key"
 EVENTGRID_TOKEN_HEADER = "aeg-sas-token"
 DEFAULT_API_VERSION = "2018-01-01"

sdk/eventgrid/azure-eventgrid/azure/eventgrid/_helpers.py

@@ -16,7 +16,7 @@
 
 from msrest import Serializer
 from azure.core.pipeline.transport import HttpRequest
-from azure.core.pipeline.policies import AzureKeyCredentialPolicy
+from azure.core.pipeline.policies import AzureKeyCredentialPolicy, BearerTokenCredentialPolicy
 from azure.core.credentials import AzureKeyCredential, AzureSasCredential
 from ._signature_credential_policy import EventGridSasCredentialPolicy
 from . import _constants as constants
@@ -28,7 +28,6 @@
 if TYPE_CHECKING:
     from datetime import datetime
 
-
 def generate_sas(endpoint, shared_access_key, expiration_date_utc, **kwargs):
     # type: (str, str, datetime, Any) -> str
     """Helper method to generate shared access signature given hostname, key, and expiration date.
@@ -70,9 +69,14 @@ def _generate_hmac(key, message):
     return base64.b64encode(hmac_new)
 
 
-def _get_authentication_policy(credential):
+def _get_authentication_policy(credential, bearer_token_policy=BearerTokenCredentialPolicy):
     if credential is None:
         raise ValueError("Parameter 'self._credential' must not be None.")
+    if hasattr(credential, "get_token"):
+        return bearer_token_policy(
+            credential,
+            constants.DEFAULT_EVENTGRID_SCOPE
+        )
     if isinstance(credential, AzureKeyCredential):
         return AzureKeyCredentialPolicy(
             credential=credential, name=constants.EVENTGRID_KEY_HEADER
@@ -82,7 +86,7 @@ def _get_authentication_policy(credential):
             credential=credential, name=constants.EVENTGRID_TOKEN_HEADER
         )
     raise ValueError(
-        "The provided credential should be an instance of AzureSasCredential or AzureKeyCredential"
+        "The provided credential should be an instance of a TokenCredential, AzureSasCredential or AzureKeyCredential"
     )
 
 

sdk/eventgrid/azure-eventgrid/azure/eventgrid/_publisher_client.py

@@ -40,7 +40,7 @@
 
 if TYPE_CHECKING:
     # pylint: disable=unused-import,ungrouped-imports
-    from azure.core.credentials import AzureKeyCredential, AzureSasCredential
+    from azure.core.credentials import AzureKeyCredential, AzureSasCredential, TokenCredential
 
     SendType = Union[
         CloudEvent,
@@ -60,8 +60,9 @@ class EventGridPublisherClient(object):
 
     :param str endpoint: The topic endpoint to send the events to.
     :param credential: The credential object used for authentication which
-     implements SAS key authentication or SAS token authentication.
-    :type credential: ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials.AzureSasCredential
+     implements SAS key authentication or SAS token authentication or a TokenCredential.
+    :type credential: ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials.AzureSasCredential or
+     ~azure.core.credentials.TokenCredential
     :rtype: None
 
     .. admonition:: Example:
@@ -82,15 +83,15 @@ class EventGridPublisherClient(object):
     """
 
     def __init__(self, endpoint, credential, **kwargs):
-        # type: (str, Union[AzureKeyCredential, AzureSasCredential], Any) -> None
+        # type: (str, Union[AzureKeyCredential, AzureSasCredential, TokenCredential], Any) -> None
         self._endpoint = endpoint
         self._client = EventGridPublisherClientImpl(
             policies=EventGridPublisherClient._policies(credential, **kwargs), **kwargs
         )
 
     @staticmethod
     def _policies(credential, **kwargs):
-        # type: (Union[AzureKeyCredential, AzureSasCredential], Any) -> List[Any]
+        # type: (Union[AzureKeyCredential, AzureSasCredential, TokenCredential], Any) -> List[Any]
         auth_policy = _get_authentication_policy(credential)
         sdk_moniker = "eventgrid/{}".format(VERSION)
         policies = [

sdk/eventgrid/azure-eventgrid/azure/eventgrid/_version.py

@@ -9,4 +9,4 @@
 # regenerated.
 # --------------------------------------------------------------------------
 
-VERSION = "4.3.1"
+VERSION = "4.4.0"

sdk/eventgrid/azure-eventgrid/azure/eventgrid/aio/_publisher_client_async.py

@@ -6,7 +6,7 @@
 # Changes may cause incorrect behavior and will be lost if the code is regenerated.
 # --------------------------------------------------------------------------
 
-from typing import Any, Union, List, Dict, cast
+from typing import Any, Union, List, Dict, TYPE_CHECKING, cast
 from azure.core.credentials import AzureKeyCredential, AzureSasCredential
 from azure.core.tracing.decorator_async import distributed_trace_async
 from azure.core.messaging import CloudEvent
@@ -22,20 +22,24 @@
     DistributedTracingPolicy,
     HttpLoggingPolicy,
     UserAgentPolicy,
+    AsyncBearerTokenCredentialPolicy
 )
 from .._policies import CloudEventDistributedTracingPolicy
 from .._models import EventGridEvent
 from .._helpers import (
-    _get_authentication_policy,
     _is_cloud_event,
     _is_eventgrid_event,
     _eventgrid_data_typecheck,
     _build_request,
     _cloud_event_to_generated,
+    _get_authentication_policy
 )
 from .._generated.aio import EventGridPublisherClient as EventGridPublisherClientAsync
 from .._version import VERSION
 
+if TYPE_CHECKING:
+    from azure.core.credentials_async import AsyncTokenCredential
+
 SendType = Union[
     CloudEvent, EventGridEvent, Dict, List[CloudEvent], List[EventGridEvent], List[Dict]
 ]
@@ -49,8 +53,9 @@ class EventGridPublisherClient:
 
     :param str endpoint: The topic endpoint to send the events to.
     :param credential: The credential object used for authentication which implements
-     SAS key authentication or SAS token authentication.
-    :type credential: ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials.AzureSasCredential
+     SAS key authentication or SAS token authentication or an AsyncTokenCredential.
+    :type credential: ~azure.core.credentials.AzureKeyCredential or ~azure.core.credentials.AzureSasCredential or
+     ~azure.core.credentials_async.AsyncTokenCredential
     :rtype: None
 
     .. admonition:: Example:
@@ -73,7 +78,7 @@ class EventGridPublisherClient:
     def __init__(
         self,
         endpoint: str,
-        credential: Union[AzureKeyCredential, AzureSasCredential],
+        credential: Union["AsyncTokenCredential", AzureKeyCredential, AzureSasCredential],
         **kwargs: Any
     ) -> None:
         self._client = EventGridPublisherClientAsync(
@@ -83,9 +88,9 @@ def __init__(
 
     @staticmethod
     def _policies(
-        credential: Union[AzureKeyCredential, AzureSasCredential], **kwargs: Any
+        credential: Union[AzureKeyCredential, AzureSasCredential, "AsyncTokenCredential"], **kwargs: Any
     ) -> List[Any]:
-        auth_policy = _get_authentication_policy(credential)
+        auth_policy = _get_authentication_policy(credential, AsyncBearerTokenCredentialPolicy)
         sdk_moniker = "eventgridpublisherclient/{}".format(VERSION)
         policies = [
             RequestIdPolicy(**kwargs),

sdk/eventgrid/azure-eventgrid/samples/async_samples/sample_authentication_async.py

@@ -38,3 +38,20 @@
 credential = AzureSasCredential(signature)
 client = EventGridPublisherClient(endpoint, credential)
 # [END client_auth_with_sas_cred_async]
+
+# [START client_auth_with_token_cred_async]
+from azure.identity.aio import DefaultAzureCredential
+from azure.eventgrid.aio import EventGridPublisherClient
+from azure.eventgrid import EventGridEvent
+
+event = EventGridEvent(
+    data={"team": "azure-sdk"},
+    subject="Door1",
+    event_type="Azure.Sdk.Demo",
+    data_version="2.0"
+)
+
+credential = DefaultAzureCredential()
+endpoint = os.environ["EG_TOPIC_HOSTNAME"]
+client = EventGridPublisherClient(endpoint, credential)
+# [END client_auth_with_token_cred_async]

sdk/eventgrid/azure-eventgrid/samples/sync_samples/sample_authentication.py

@@ -38,3 +38,12 @@
 credential = AzureSasCredential(signature)
 client = EventGridPublisherClient(endpoint, credential)
 # [END client_auth_with_sas_cred]
+
+# [START client_auth_with_token_cred]
+from azure.identity import DefaultAzureCredential
+from azure.eventgrid import EventGridPublisherClient, EventGridEvent
+
+credential = DefaultAzureCredential()
+endpoint = os.environ["EG_TOPIC_HOSTNAME"]
+client = EventGridPublisherClient(endpoint, credential)
+# [END client_auth_with_token_cred]

sdk/eventgrid/azure-eventgrid/tests/asynctestcase.py

@@ -0,0 +1,40 @@
+
+# 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.
+# --------------------------------------------------------------------------
+import asyncio
+import os
+import functools
+from azure_devtools.scenario_tests.utilities import trim_kwargs_from_test_function
+from azure.core.credentials import AccessToken
+from eventgrid_preparer import EventGridTest
+
+
+class AsyncFakeTokenCredential(object):
+    """Protocol for classes able to provide OAuth tokens.
+    :param str scopes: Lets you specify the type of access needed.
+    """
+    def __init__(self):
+        self.token = AccessToken("YOU SHALL NOT PASS", 0)
+
+    async def get_token(self, *args):
+        return self.token
+
+
+class AsyncEventGridTest(EventGridTest):
+
+    def generate_oauth_token(self):
+        if self.is_live:
+            from azure.identity.aio import ClientSecretCredential
+            return ClientSecretCredential(
+                os.getenv("AZURE_TENANT_ID"),
+                os.getenv("AZURE_CLIENT_ID"),
+                os.getenv("AZURE_CLIENT_SECRET"),
+            )
+        return self.generate_fake_token()
+
+    def generate_fake_token(self):
+        return AsyncFakeTokenCredential()

sdk/eventgrid/azure-eventgrid/tests/eventgrid_preparer.py

@@ -3,12 +3,14 @@
 import os
 from collections import namedtuple
 
+from azure_devtools.scenario_tests import ReplayableTest
+from azure.core.credentials import AccessToken
 from azure.mgmt.eventgrid import EventGridManagementClient
 from azure.mgmt.eventgrid.models import Topic, InputSchema, JsonInputSchemaMapping, JsonField, JsonFieldWithDefault
 from azure_devtools.scenario_tests.exceptions import AzureTestError
 
 from devtools_testutils import (
-    ResourceGroupPreparer, AzureMgmtPreparer, FakeResource
+    ResourceGroupPreparer, AzureMgmtPreparer, FakeResource, AzureMgmtTestCase
 )
 
 from devtools_testutils.resource_testcase import RESOURCE_GROUP_PARAM
@@ -25,6 +27,15 @@
 DATA_VERSION_JSON_FIELD_WITH_DEFAULT = JsonFieldWithDefault(source_field='customDataVersion', default_value='')
 CUSTOM_JSON_INPUT_SCHEMA_MAPPING = JsonInputSchemaMapping(id=ID_JSON_FIELD, topic=TOPIC_JSON_FIELD, event_time=EVENT_TIME_JSON_FIELD, event_type=EVENT_TYPE_JSON_FIELD_WITH_DEFAULT, subject=SUBJECT_JSON_FIELD_WITH_DEFAULT, data_version=DATA_VERSION_JSON_FIELD_WITH_DEFAULT)
 
+class FakeTokenCredential(object):
+    """Protocol for classes able to provide OAuth tokens.
+    :param str scopes: Lets you specify the type of access needed.
+    """
+    def __init__(self):
+        self.token = AccessToken("YOU SHALL NOT PASS", 0)
+
+    def get_token(self, *args):
+        return self.token
 
 class EventGridTopicPreparer(AzureMgmtPreparer):
     def __init__(self,
@@ -94,4 +105,28 @@ def _get_resource_group(self, **kwargs):
                        'decorator @{} in front of this event grid topic preparer.'
             raise AzureTestError(template.format(ResourceGroupPreparer.__name__))
 
+
 CachedEventGridTopicPreparer = functools.partial(EventGridTopicPreparer, use_cache=True)
+
+
+class EventGridTest(AzureMgmtTestCase):
+    FILTER_HEADERS = ReplayableTest.FILTER_HEADERS + ['aeg-sas-key', 'aeg-sas-token']
+
+    def __init__(self, method_name):
+        super(EventGridTest, self).__init__(method_name)
+
+    def get_oauth_endpoint(self):
+        return os.getenv("EG_TOPIC_HOSTNAME")
+
+    def generate_oauth_token(self):
+        if self.is_live:
+            from azure.identity import ClientSecretCredential
+            return ClientSecretCredential(
+                os.getenv("AZURE_TENANT_ID"),
+                os.getenv("AZURE_CLIENT_ID"),
+                os.getenv("AZURE_CLIENT_SECRET"),
+            )
+        return self.generate_fake_token()
+
+    def generate_fake_token(self):
+        return FakeTokenCredential()

sdk/eventgrid/azure-eventgrid/tests/test_eg_publisher_client.py

@@ -29,10 +29,11 @@
 from azure.eventgrid._helpers import _cloud_event_to_generated
 
 from eventgrid_preparer import (
-    CachedEventGridTopicPreparer
+    CachedEventGridTopicPreparer,
+    EventGridTest
 )
 
-class EventGridPublisherClientTests(AzureMgmtTestCase):
+class EventGridPublisherClientTests(EventGridTest):
     FILTER_HEADERS = ReplayableTest.FILTER_HEADERS + ['aeg-sas-key', 'aeg-sas-token']
 
     @CachedResourceGroupPreparer(name_prefix='eventgridtest')
@@ -343,5 +344,19 @@ def test_send_custom_schema_event_as_list(self, resource_group, eventgrid_topic,
 
     def test_send_throws_with_bad_credential(self):
         bad_credential = "I am a bad credential"
-        with pytest.raises(ValueError, match="The provided credential should be an instance of AzureSasCredential or AzureKeyCredential"):
+        with pytest.raises(ValueError, match="The provided credential should be an instance of a TokenCredential, AzureSasCredential or AzureKeyCredential"):
             client = EventGridPublisherClient("eventgrid_endpoint", bad_credential)
+
+    @pytest.mark.live_test_only
+    @CachedResourceGroupPreparer(name_prefix='eventgridtest')
+    @CachedEventGridTopicPreparer(name_prefix='eventgridtest')
+    def test_send_token_credential(self, resource_group, eventgrid_topic, eventgrid_topic_primary_key, eventgrid_topic_endpoint):
+        credential = self.generate_oauth_token()
+        client = EventGridPublisherClient(self.get_oauth_endpoint(), credential)
+        eg_event = EventGridEvent(
+                subject="sample", 
+                data={"sample": "eventgridevent"}, 
+                event_type="Sample.EventGrid.Event",
+                data_version="2.0"
+                )
+        client.send(eg_event)