diff --git a/src/sentry/identity/__init__.py b/src/sentry/identity/__init__.py index bbbf8ebf183d92..d428083794c119 100644 --- a/src/sentry/identity/__init__.py +++ b/src/sentry/identity/__init__.py @@ -11,6 +11,7 @@ def _register_providers() -> None: from .bitbucket.provider import BitbucketIdentityProvider + from .datadog.provider import DatadogIdentityProvider from .discord.provider import DiscordIdentityProvider from .gcp.provider import GCPIdentityProvider from .github.provider import GitHubIdentityProvider @@ -34,6 +35,7 @@ def _register_providers() -> None: register(GitlabIdentityProvider) register(GoogleIdentityProvider) register(DiscordIdentityProvider) + register(DatadogIdentityProvider) register(GCPIdentityProvider) diff --git a/src/sentry/identity/datadog/provider.py b/src/sentry/identity/datadog/provider.py index 150c23c20a9fdb..6f9caef6a7798c 100644 --- a/src/sentry/identity/datadog/provider.py +++ b/src/sentry/identity/datadog/provider.py @@ -3,6 +3,7 @@ import base64 import hashlib import secrets +from typing import Any import orjson from django.http.request import HttpRequest @@ -10,23 +11,67 @@ from requests import ConnectionError, HTTPError, Response from requests.exceptions import SSLError +from sentry.auth.exceptions import IdentityNotValid from sentry.http import safe_urlopen, safe_urlread from sentry.identity.oauth2 import ( OAuth2CallbackView, OAuth2LoginView, + OAuth2Provider, _redirect_url, record_event, ) from sentry.identity.pipeline import IdentityPipeline +from sentry.identity.services.identity.model import RpcIdentity +from sentry.integrations.types import IntegrationProviderSlug from sentry.integrations.utils.metrics import IntegrationPipelineViewType +from sentry.pipeline.views.base import PipelineView from sentry.users.models.identity import Identity from sentry.utils.http import absolute_uri +MCP_REGISTER_PATH = "/api/unstable/mcp-server/register" +MCP_AUTHORIZE_PATH = "/api/unstable/mcp-server/authorize" +MCP_TOKEN_PATH = "/api/unstable/mcp-server/token" +MCP_ENDPOINT_PATH = "/api/unstable/mcp-server/mcp" + def _basic_auth_header(client_id: str, client_secret: str) -> str: return "Basic " + base64.b64encode(f"{client_id}:{client_secret}".encode()).decode("ascii") +def get_user_info(access_token: str, site: str) -> dict[str, Any]: + """Fetch the current Datadog user via the MCP ``datadog://mcp/whoami`` resource.""" + url = f"https://mcp.{site}{MCP_ENDPOINT_PATH}" + headers = {"Authorization": f"Bearer {access_token}", "Content-Type": "application/json"} + + init_resp = safe_urlopen( + url, + method="POST", + headers=headers, + json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}, + ) + init_resp.raise_for_status() + headers["Mcp-Session-Id"] = init_resp.headers["mcp-session-id"] + + resp = safe_urlopen( + url, + method="POST", + headers=headers, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "resources/read", + "params": {"uri": "datadog://mcp/whoami"}, + }, + ) + resp.raise_for_status() + + try: + body = orjson.loads(safe_urlread(resp)) + return orjson.loads(body["result"]["contents"][0]["text"]) + except (KeyError, IndexError, orjson.JSONDecodeError) as e: + raise IdentityNotValid("MCP whoami returned an unexpected response") from e + + def generate_pkce_code_verifier() -> str: return secrets.token_urlsafe(96) @@ -196,3 +241,113 @@ def get_access_token(self, pipeline: IdentityPipeline, code: str) -> Response: return safe_urlopen( self.access_token_url, data=data, headers=headers, verify_ssl=verify_ssl ) + + +class DatadogIdentityProvider(OAuth2Provider): + key = IntegrationProviderSlug.DATADOG + name = "Datadog" + + oauth_scopes: tuple[str, ...] = ( + "mcp_read", + "apm_read", + "error_tracking_read", + "events_read", + "hosts_read", + "incident_read", + "logs_read_data", + "metrics_read", + "monitors_read", + ) + + def _get_mcp_base_url(self) -> str: + return f"https://mcp.{self._get_oauth_parameter('site')}" + + def get_oauth_authorize_url(self) -> str: + return self._get_mcp_base_url() + MCP_AUTHORIZE_PATH + + def get_oauth_access_token_url(self) -> str: + return self._get_mcp_base_url() + MCP_TOKEN_PATH + + def get_pipeline_views(self) -> list[PipelineView[IdentityPipeline]]: + return [ + DatadogOAuth2DCRView( + register_url=self._get_mcp_base_url() + MCP_REGISTER_PATH, + ), + DatadogOAuth2LoginView( + authorize_url=self.get_oauth_authorize_url(), + scope=" ".join(self.get_oauth_scopes()), + resource=self._get_mcp_base_url(), + ), + DatadogOAuth2CallbackView( + access_token_url=self.get_oauth_access_token_url(), + ), + ] + + def get_oauth_data(self, payload: dict[str, Any]) -> dict[str, Any]: + data = super().get_oauth_data(payload) + if "scope" in payload: + data["scope"] = payload["scope"] + return data + + def build_identity(self, data: dict[str, Any]) -> dict[str, Any]: + token_data = data["data"] + access_token = token_data.get("access_token") + if not access_token: + raise ValueError("Datadog token exchange did not return an access_token") + + user = get_user_info(access_token, self._get_oauth_parameter("site")) + + oauth_data = self.get_oauth_data(token_data) + + # Persist DCR credentials and site so refresh_identity can access them outside a pipeline context. + client_id = data.get("dcr_client_id") + client_secret = data.get("dcr_client_secret") + if not client_id or not client_secret: + raise IdentityNotValid("Missing DCR credentials") + oauth_data["client_id"] = client_id + oauth_data["client_secret"] = client_secret + oauth_data["site"] = self._get_oauth_parameter("site") + + return { + "type": IntegrationProviderSlug.DATADOG, + "id": user["user_uuid"], + "email": user.get("user_email"), + "name": user.get("user_name"), + "scopes": [], + "data": oauth_data, + } + + def get_refresh_token_url(self) -> str: + return self.get_oauth_access_token_url() + + def get_refresh_token_headers(self) -> dict[str, str]: + client_id = self.config.get("client_id") + client_secret = self.config.get("client_secret") + + if not client_id or not client_secret: + raise IdentityNotValid("Missing DCR credentials") + + return {"Authorization": _basic_auth_header(client_id, client_secret)} + + def get_refresh_token_params( + self, refresh_token: str, identity: Identity | RpcIdentity, **kwargs: Any + ) -> dict[str, str | None]: + return {"grant_type": "refresh_token", "refresh_token": refresh_token} + + def refresh_identity(self, identity: Identity | RpcIdentity, **kwargs: Any) -> None: + # Add site and client credentials to config so get_refresh_token_headers + # and get_refresh_token_url can access them. + + site = identity.data.get("site") + if not site: + raise IdentityNotValid("Missing Datadog site") + self.config["site"] = site + + client_id = identity.data.get("client_id") + client_secret = identity.data.get("client_secret") + if not client_id or not client_secret: + raise IdentityNotValid("Missing DCR credentials") + self.config["client_id"] = client_id + self.config["client_secret"] = client_secret + + super().refresh_identity(identity, **kwargs) diff --git a/src/sentry/integrations/types.py b/src/sentry/integrations/types.py index 63d6562b05a323..e13fa84fdbdef1 100644 --- a/src/sentry/integrations/types.py +++ b/src/sentry/integrations/types.py @@ -46,6 +46,7 @@ class IntegrationProviderSlug(StrEnum): PAGERDUTY = "pagerduty" OPSGENIE = "opsgenie" PERFORCE = "perforce" + DATADOG = "datadog" GCP = "gcp" diff --git a/tests/sentry/identity/datadog/test_provider.py b/tests/sentry/identity/datadog/test_provider.py index 102f8558e2373e..e84c96863c76bb 100644 --- a/tests/sentry/identity/datadog/test_provider.py +++ b/tests/sentry/identity/datadog/test_provider.py @@ -3,10 +3,12 @@ import base64 import hashlib from functools import cached_property +from typing import Any from unittest.mock import MagicMock, patch from urllib.parse import parse_qs, parse_qsl, urlparse import orjson +import pytest import responses from django.contrib.messages.storage.fallback import FallbackStorage from django.contrib.sessions.backends.base import SessionBase @@ -15,7 +17,9 @@ from requests.exceptions import SSLError import sentry.identity +from sentry.auth.exceptions import IdentityNotValid from sentry.identity.datadog.provider import ( + DatadogIdentityProvider, DatadogOAuth2CallbackView, DatadogOAuth2DCRView, DatadogOAuth2LoginView, @@ -26,11 +30,12 @@ from sentry.testutils.asserts import assert_failure_metric from sentry.testutils.cases import TestCase from sentry.testutils.silo import control_silo_test -from sentry.users.models.identity import IdentityProvider +from sentry.users.models.identity import Identity REGISTER_URL = "https://mcp.datadoghq.com/api/unstable/mcp-server/register" AUTHORIZE_URL = "https://mcp.datadoghq.com/api/unstable/mcp-server/authorize" TOKEN_URL = "https://mcp.datadoghq.com/api/unstable/mcp-server/token" +MCP_URL = "https://mcp.datadoghq.com/api/unstable/mcp-server/mcp" RESOURCE = "https://mcp.datadoghq.com" @@ -369,7 +374,7 @@ def setUp(self) -> None: sentry.identity.register(DatadogTestProvider) super().setUp() self.request = self._make_request() - self.identity_provider = IdentityProvider.objects.create(type="datadog-test") + self.identity_provider = self.create_identity_provider(type="datadog-test") def tearDown(self) -> None: super().tearDown() @@ -441,3 +446,167 @@ def test_pipeline_views(self, mock_record: MagicMock) -> None: f"Basic {base64.b64encode(b'dcr-client-id:dcr-client-secret').decode()}" ) assert exchange_token_request.headers["Authorization"] == expected_auth_header + + +@control_silo_test +class DatadogIdentityProviderTest(TestCase): + def setUp(self) -> None: + super().setUp() + self.provider = DatadogIdentityProvider() + self.provider.config = {"site": "datadoghq.com"} + self.identity_provider = self.create_identity_provider(type="datadog") + + @patch("sentry.identity.datadog.provider.get_user_info") + def test_build_identity(self, mock_get_user_info: MagicMock) -> None: + mock_get_user_info.return_value = { + "user_uuid": "dd-user-123", + "user_email": "user@example.com", + "user_name": "Test User", + } + + result = self.provider.build_identity( + { + "data": { + "access_token": "token-abc", + "refresh_token": "refresh-xyz", + "expires_in": 3600, + "token_type": "Bearer", + "scope": "mcp_read", + }, + "dcr_client_id": "dcr-client-id", + "dcr_client_secret": "dcr-client-secret", + } + ) + + assert result["id"] == "dd-user-123" + assert result["email"] == "user@example.com" + assert result["name"] == "Test User" + assert result["type"] == "datadog" + assert result["data"]["access_token"] == "token-abc" + assert result["data"]["refresh_token"] == "refresh-xyz" + assert "expires" in result["data"] + assert result["data"]["token_type"] == "Bearer" + assert result["data"]["scope"] == "mcp_read" + assert result["data"]["client_id"] == "dcr-client-id" + assert result["data"]["client_secret"] == "dcr-client-secret" + assert result["data"]["site"] == "datadoghq.com" + mock_get_user_info.assert_called_once_with("token-abc", "datadoghq.com") + + @patch("sentry.identity.datadog.provider.get_user_info") + def test_build_identity_missing_access_token(self, mock_get_user_info: MagicMock) -> None: + with pytest.raises(ValueError, match="did not return an access_token"): + self.provider.build_identity( + { + "data": {}, + "dcr_client_id": "dcr-id", + "dcr_client_secret": "dcr-secret", + } + ) + mock_get_user_info.assert_not_called() + + @patch("sentry.identity.datadog.provider.get_user_info") + def test_build_identity_missing_dcr_credentials(self, mock_get_user_info: MagicMock) -> None: + mock_get_user_info.return_value = {"user_uuid": "dd-user-456"} + + with pytest.raises(IdentityNotValid, match="Missing DCR credentials"): + self.provider.build_identity({"data": {"access_token": "token"}}) + + @patch("sentry.identity.datadog.provider.get_user_info") + def test_build_identity_missing_user_attributes(self, mock_get_user_info: MagicMock) -> None: + mock_get_user_info.return_value = {"user_uuid": "dd-user-456"} + + result = self.provider.build_identity( + { + "data": {"access_token": "token"}, + "dcr_client_id": "dcr-id", + "dcr_client_secret": "dcr-secret", + } + ) + + assert result["id"] == "dd-user-456" + assert result["email"] is None + assert result["name"] is None + + @responses.activate + def test_build_identity_malformed_user_info(self) -> None: + responses.add( + responses.POST, + MCP_URL, + json={"jsonrpc": "2.0", "id": 1, "result": {}}, + headers={"mcp-session-id": "session-id"}, + ) + responses.add(responses.POST, MCP_URL, body="not json", status=200) + + with pytest.raises(IdentityNotValid, match="unexpected response"): + self.provider.build_identity( + { + "data": {"access_token": "token"}, + "dcr_client_id": "dcr-id", + "dcr_client_secret": "dcr-secret", + } + ) + + def _make_identity(self, **data_overrides: Any) -> Identity: + data = { + "access_token": "old-token", + "refresh_token": "old-refresh", + "client_id": "dcr-client-id", + "client_secret": "dcr-client-secret", + "site": "datadoghq.com", + } + data.update(data_overrides) + return self.create_identity( + user=self.user, + identity_provider=self.identity_provider, + external_id="dd-user-123", + data=data, + ) + + @responses.activate + def test_refresh_identity_success(self) -> None: + responses.add( + responses.POST, + TOKEN_URL, + json={"access_token": "new-token", "refresh_token": "new-refresh", "expires_in": 3600}, + ) + + identity = self._make_identity() + self.provider.refresh_identity(identity) + + auth_header = responses.calls[0].request.headers["Authorization"] + assert ( + auth_header == f"Basic {base64.b64encode(b'dcr-client-id:dcr-client-secret').decode()}" + ) + + data = dict(parse_qsl(responses.calls[0].request.body)) + assert data["grant_type"] == "refresh_token" + assert data["refresh_token"] == "old-refresh" + assert "client_id" not in data + assert "client_secret" not in data + + assert responses.calls[0].request.url == TOKEN_URL + + identity.refresh_from_db() + assert identity.data["access_token"] == "new-token" + assert identity.data["refresh_token"] == "new-refresh" + assert identity.data["client_id"] == "dcr-client-id" + assert identity.data["client_secret"] == "dcr-client-secret" + assert identity.data["site"] == "datadoghq.com" + + def test_refresh_identity_missing_site(self) -> None: + identity = self._make_identity(site=None) + + with pytest.raises(IdentityNotValid, match="Missing Datadog site"): + self.provider.refresh_identity(identity) + + def test_refresh_identity_missing_dcr_credentials(self) -> None: + identity = self._make_identity(client_id=None, client_secret=None) + + with pytest.raises(IdentityNotValid, match="Missing DCR credentials"): + self.provider.refresh_identity(identity) + + def test_refresh_identity_missing_refresh_token(self) -> None: + identity = self._make_identity(refresh_token=None) + + with pytest.raises(IdentityNotValid, match="Missing refresh token"): + self.provider.refresh_identity(identity)