fix(seer-infra-telemetry): Add Datadog site allowlist to identity provider

src/sentry/api/endpoints/organization_monitoring_providers.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import logging
+
+from django.http import HttpResponseRedirect
+from rest_framework.request import Request
+from rest_framework.response import Response
+
+from sentry import features
+from sentry.api.api_owners import ApiOwner
+from sentry.api.api_publish_status import ApiPublishStatus
+from sentry.api.base import control_silo_endpoint
+from sentry.api.bases.organization import (
+    ControlSiloOrganizationEndpoint,
+    OrganizationPermission,
+)
+from sentry.identity.datadog.provider import DATADOG_VALID_SITES
+from sentry.identity.pipeline import IdentityPipeline
+from sentry.organizations.services.organization.model import RpcOrganization
+from sentry.users.models.identity import Identity, IdentityProvider
+
+logger = logging.getLogger(__name__)
+
+MONITORING_PROVIDERS: dict[str, dict[str, str]] = {
+    "datadog": {"name": "Datadog"},
+    "gcp": {"name": "Google Cloud Platform"},
+}
+
+MONITORING_PROVIDER_FEATURE = "organizations:seer-infra-telemetry"
+
+
+class MonitoringProviderPermission(OrganizationPermission):
+    scope_map = {
+        "GET": ["org:read", "org:write", "org:admin"],
+        "POST": ["org:write", "org:admin"],
+        "DELETE": ["org:write", "org:admin"],
+    }
+
+
+@control_silo_endpoint
+class OrganizationMonitoringProviderIndexEndpoint(ControlSiloOrganizationEndpoint):
+    owner = ApiOwner.CODING_WORKFLOWS
+    publish_status = {
+        "GET": ApiPublishStatus.PRIVATE,
+    }
+    permission_classes = (MonitoringProviderPermission,)
+
+    def get(self, request: Request, organization: RpcOrganization, **kwargs: object) -> Response:
+        if not features.has(MONITORING_PROVIDER_FEATURE, organization, actor=request.user):
+            return Response(status=404)
+
+        user_id = request.user.id
+        if user_id is None:
+            return Response(status=401)
+
+        connected_identities = {
+            identity.idp.type: identity
+            for identity in Identity.objects.filter(
+                idp__type__in=MONITORING_PROVIDERS.keys(),
+                user_id=user_id,
+            ).select_related("idp")
+        }
+
+        providers = []
+        for key, meta in MONITORING_PROVIDERS.items():
+            identity = connected_identities.get(key)
+            providers.append(
+                {
+                    "provider": key,
+                    "name": meta["name"],
+                    "connected": identity is not None,
+                }
+            )
+
+        return Response({"providers": providers})
+
+
+@control_silo_endpoint
+class OrganizationMonitoringProviderDetailsEndpoint(ControlSiloOrganizationEndpoint):
+    owner = ApiOwner.CODING_WORKFLOWS
+    publish_status = {
+        "POST": ApiPublishStatus.PRIVATE,
+        "DELETE": ApiPublishStatus.PRIVATE,
+    }
+    permission_classes = (MonitoringProviderPermission,)
+
+    def post(
+        self, request: Request, organization: RpcOrganization, provider_key: str, **kwargs: object
+    ) -> Response:
+        if not features.has(MONITORING_PROVIDER_FEATURE, organization, actor=request.user):
+            return Response(status=404)
+
+        if request.user.id is None:
+            return Response(status=401)
+
+        if provider_key not in MONITORING_PROVIDERS:
+            return Response({"detail": "Unknown monitoring provider."}, status=400)
+
+        config: dict[str, str] = {}
+        if provider_key == "datadog":
+            site = request.data.get("site")
+            if not site:
+                return Response(
+                    {"detail": "Datadog requires a 'site' parameter (e.g. 'datadoghq.com')."},
+                    status=400,
+                )
+            elif site not in DATADOG_VALID_SITES:
+                return Response({"detail": f"Invalid Datadog site: {site}"}, status=400)
+            config["site"] = site
+
+        # Datadog: the IdentityProvider is auto-created during the pipeline
+        idp: IdentityProvider | None = None
+        if provider_key != "datadog":
+            idp, _ = IdentityProvider.objects.get_or_create(type=provider_key, external_id="")
+
+        pipeline = IdentityPipeline(
+            request=request._request,
+            provider_key=provider_key,
+            organization=organization,
+            provider_model=idp,
+            config=config,
+        )
+        pipeline.initialize()
+
+        response = pipeline.current_step()
+
+        if isinstance(response, HttpResponseRedirect):
+            return Response({"redirectUrl": response.url})
+
+        logger.error(
+            "monitoring_provider.connect.unexpected_response",
+            extra={"provider": provider_key, "response_type": type(response).__name__},
+        )
+        return Response({"detail": "Failed to start OAuth flow."}, status=500)
+
+    def delete(
+        self, request: Request, organization: RpcOrganization, provider_key: str, **kwargs: object
+    ) -> Response:
+        if not features.has(MONITORING_PROVIDER_FEATURE, organization, actor=request.user):
+            return Response(status=404)
+
+        user_id = request.user.id
+        if user_id is None:
+            return Response(status=401)
+
+        if provider_key not in MONITORING_PROVIDERS:
+            return Response({"detail": "Unknown monitoring provider."}, status=400)
+
+        identities = list(
+            Identity.objects.filter(
+                idp__type=provider_key,
+                user_id=user_id,
+            )
+        )
+
+        if not identities:
+            return Response({"detail": "Not connected to this provider."}, status=404)
+
+        for identity in identities:
+            identity.delete()
+
+        return Response(status=204)

src/sentry/identity/datadog/provider.py

@@ -28,6 +28,19 @@
 from sentry.users.models.identity import Identity
 from sentry.utils.http import absolute_uri
 
+DATADOG_VALID_SITES = frozenset(
+    {
+        "datadoghq.com",
+        "us3.datadoghq.com",
+        "us5.datadoghq.com",
+        "datadoghq.eu",
+        "ddog-gov.com",
+        "us2.ddog-gov.com",
+        "ap1.datadoghq.com",
+        "ap2.datadoghq.com",
+    }
+)
+
 MCP_REGISTER_PATH = "/api/unstable/mcp-server/register"
 MCP_AUTHORIZE_PATH = "/api/unstable/mcp-server/authorize"
 MCP_TOKEN_PATH = "/api/unstable/mcp-server/token"
@@ -38,9 +51,9 @@ 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]:
+def get_user_info(access_token: str, mcp_base_url: str) -> dict[str, Any]:
     """Fetch the current Datadog user via the MCP ``datadog://mcp/whoami`` resource."""
-    url = f"https://mcp.{site}{MCP_ENDPOINT_PATH}"
+    url = f"{mcp_base_url}{MCP_ENDPOINT_PATH}"
     headers = {"Authorization": f"Bearer {access_token}", "Content-Type": "application/json"}
 
     init_resp = safe_urlopen(
@@ -267,7 +280,10 @@ def get_pipeline_config(self, data: dict[str, Any]) -> dict[str, str]:
         return {"site": site}
 
     def _get_mcp_base_url(self) -> str:
-        return f"https://mcp.{self._get_oauth_parameter('site')}"
+        site = self._get_oauth_parameter("site")
+        if site not in DATADOG_VALID_SITES:
+            raise ValueError(f"Invalid Datadog site: {site}")
+        return f"https://mcp.{site}"
 
     def get_oauth_authorize_url(self) -> str:
         return self._get_mcp_base_url() + MCP_AUTHORIZE_PATH
@@ -302,14 +318,13 @@ def build_identity(self, data: dict[str, Any]) -> dict[str, Any]:
         if not access_token:
             raise ValueError("Datadog token exchange did not return an access_token")
 
-        site = self._get_oauth_parameter("site")
-        user = get_user_info(access_token, site)
-
+        user = get_user_info(access_token, self._get_mcp_base_url())
         if "user_uuid" not in user or "org_uuid" not in user:
             raise IdentityNotValid(
                 "User info response missing required fields (user_uuid, org_uuid)"
             )
 
+        site = 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.
@@ -356,6 +371,8 @@ def refresh_identity(self, identity: Identity | RpcIdentity, **kwargs: Any) -> N
         site = identity.data.get("site")
         if not site:
             raise IdentityNotValid("Missing Datadog site")
+        elif site not in DATADOG_VALID_SITES:
+            raise IdentityNotValid(f"Invalid Datadog site: {site}")
         self.config["site"] = site
 
         client_id = identity.data.get("client_id")