diff --git a/pyrit/backend/middleware/auth.py b/pyrit/backend/middleware/auth.py index 81594fc66e..c2c24acd94 100644 --- a/pyrit/backend/middleware/auth.py +++ b/pyrit/backend/middleware/auth.py @@ -19,6 +19,7 @@ import os from dataclasses import dataclass from typing import Any, ClassVar +from urllib.parse import urlparse import httpx import jwt @@ -51,6 +52,13 @@ class EntraAuthMiddleware(BaseHTTPMiddleware): "/api/media", } + # Hosts the user's Bearer token may be forwarded to + _GRAPH_HOSTS: ClassVar[set[str]] = {"graph.microsoft.com"} + + # Trusted URL for resolving group membership; built from a constant (not the + # token-supplied endpoint) so token data cannot control the request destination. + _GRAPH_MEMBER_OBJECTS_URL: ClassVar[str] = "https://graph.microsoft.com/v1.0/me/getMemberObjects" + def __init__(self, app: ASGIApp) -> None: """Initialize the middleware with Entra ID configuration from environment variables.""" super().__init__(app) @@ -167,34 +175,34 @@ async def _resolve_excess_groups_async(self, claims: dict[str, Any], token: str) """ Resolve group membership via Microsoft Graph when user is in >200 groups. - When a user is in >200 groups, Entra ID replaces the `groups` claim with - `_claim_sources` containing a Graph API endpoint. This method calls the - Microsoft Graph `getMemberObjects` endpoint to retrieve transitive group - memberships, using the user's access token. + When a user is in >200 groups, Entra ID replaces the `groups` claim with a + `_claim_names` / `_claim_sources` overage pointer. This method confirms the + overage is present, then calls the Microsoft Graph `getMemberObjects` endpoint + to retrieve transitive group memberships, using the user's access token. + + The Graph URL is built from a trusted constant (``_GRAPH_MEMBER_OBJECTS_URL``) + rather than the token-supplied endpoint, so token data cannot control the + request host, path, or port. Pagination links from the Graph response are + additionally checked against the Graph allowlist (see ``_is_trusted_graph_url``). Args: - claims: The decoded JWT claims containing _claim_sources. + claims: The decoded JWT claims containing the group overage pointer. token: The raw Bearer token to forward to Graph API. Returns: - List of group IDs the user belongs to, or empty list on failure. + List of group IDs the user belongs to, or empty list on failure or + when no group overage is present. """ try: - claim_sources = claims.get("_claim_sources", {}) - src = claim_sources.get("src1", {}) - endpoint = src.get("endpoint", "") - - if not endpoint: - logger.debug("No group resolution endpoint found in _claim_sources") + # Entra signals a groups overage via `_claim_names` -> source key -> `_claim_sources`. + # We only confirm the overage is present; the endpoint it supplies is intentionally + # ignored in favor of a trusted constant so token data cannot control the request. + claim_source_key = claims.get("_claim_names", {}).get("groups") + if not claim_source_key or claim_source_key not in claims.get("_claim_sources", {}): + logger.debug("No group overage claim source found") return [] - # The _claim_sources endpoint may be a legacy graph.windows.net URL. - # Rewrite to Microsoft Graph (graph.microsoft.com) which is the - # supported API. The legacy Azure AD Graph was retired in 2023. - if "graph.windows.net" in endpoint: - # Legacy format: https://graph.windows.net/{tenant}/users/{oid}/getMemberObjects - # Graph format: https://graph.microsoft.com/v1.0/me/getMemberObjects - endpoint = "https://graph.microsoft.com/v1.0/me/getMemberObjects" + endpoint = self._GRAPH_MEMBER_OBJECTS_URL all_group_ids: list[str] = [] async with httpx.AsyncClient() as client: @@ -222,6 +230,9 @@ async def _resolve_excess_groups_async(self, claims: dict[str, Any], token: str) # Handle pagination — Graph may return @odata.nextLink for large results next_link = data.get("@odata.nextLink") while next_link: + if not self._is_trusted_graph_url(next_link): + logger.warning("Refusing to follow untrusted pagination link: %s", next_link) + break response = await client.get( next_link, headers={"Authorization": f"Bearer {token}"}, @@ -241,6 +252,16 @@ async def _resolve_excess_groups_async(self, claims: dict[str, Any], token: str) logger.warning("Failed to resolve group memberships: %s", e) return [] + def _is_trusted_graph_url(self, url: str) -> bool: + """ + Whether *url* is an HTTPS Microsoft Graph URL safe to forward the token to. + + Returns: + True if *url* uses HTTPS and its host is in the Graph allowlist, False otherwise. + """ + parsed = urlparse(url) + return parsed.scheme == "https" and parsed.hostname in self._GRAPH_HOSTS + def _validate_token(self, token: str) -> tuple[AuthenticatedUser | None, dict[str, Any]]: """ Validate a JWT against Entra ID JWKS. diff --git a/tests/unit/backend/test_auth_middleware.py b/tests/unit/backend/test_auth_middleware.py index 6a49d7660e..fdadc199e5 100644 --- a/tests/unit/backend/test_auth_middleware.py +++ b/tests/unit/backend/test_auth_middleware.py @@ -5,11 +5,18 @@ Tests for the Entra ID auth middleware. """ -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest from pyrit.backend.middleware.auth import EntraAuthMiddleware +def _make_middleware() -> EntraAuthMiddleware: + with patch.dict("os.environ", {"ENTRA_TENANT_ID": "", "ENTRA_CLIENT_ID": ""}, clear=False): + return EntraAuthMiddleware(MagicMock()) + + def test_validate_token_returns_none_when_jwks_client_is_none(): """Test that _validate_token returns (None, {}) when _jwks_client is None.""" mock_app = MagicMock() @@ -27,3 +34,84 @@ def test_validate_token_returns_none_when_jwks_client_is_none(): assert user is None assert claims == {} + + +@pytest.mark.parametrize( + "url, expected", + [ + ("https://graph.microsoft.com/v1.0/me/getMemberObjects", True), + ("https://graph.microsoft.com/v1.0/me/memberOf?$skiptoken=abc", True), + ("http://graph.microsoft.com/v1.0/me/getMemberObjects", False), # not https + ("https://evil.com/v1.0/me/getMemberObjects", False), # wrong host + ("https://graph.microsoft.com.evil.com/x", False), # suffix spoof + ("https://graph.microsoft.com@evil.com/x", False), # userinfo spoof + ("", False), + ], +) +def test_is_trusted_graph_url(url, expected): + """Only HTTPS Microsoft Graph hosts are trusted to receive the forwarded token.""" + assert _make_middleware()._is_trusted_graph_url(url) is expected + + +async def test_resolve_excess_groups_no_overage_returns_empty(): + """Claims without a groups overage pointer return [] and make no HTTP request.""" + middleware = _make_middleware() + claims = {"_claim_sources": {"src1": {"endpoint": "https://graph.microsoft.com/x"}}} # no _claim_names.groups + + with patch("pyrit.backend.middleware.auth.httpx.AsyncClient") as mock_client: + result = await middleware._resolve_excess_groups_async(claims, "the-token") + + assert result == [] + mock_client.assert_not_called() + + +async def test_resolve_excess_groups_ignores_token_endpoint(): + """The Graph request uses the trusted constant URL, never the token-supplied endpoint.""" + middleware = _make_middleware() + claims = { + "_claim_names": {"groups": "src1"}, + "_claim_sources": {"src1": {"endpoint": "https://evil.com/steal"}}, + } + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"value": ["group-1", "group-2"]} + + mock_client = AsyncMock() + mock_client.post.return_value = mock_response + mock_client_cm = MagicMock() + mock_client_cm.__aenter__ = AsyncMock(return_value=mock_client) + mock_client_cm.__aexit__ = AsyncMock(return_value=False) + + with patch("pyrit.backend.middleware.auth.httpx.AsyncClient", return_value=mock_client_cm): + result = await middleware._resolve_excess_groups_async(claims, "the-token") + + assert result == ["group-1", "group-2"] + posted_url = mock_client.post.call_args.args[0] + assert posted_url == EntraAuthMiddleware._GRAPH_MEMBER_OBJECTS_URL + assert "evil.com" not in posted_url + + +async def test_resolve_excess_groups_stops_on_untrusted_pagination_link(): + """An untrusted @odata.nextLink halts pagination without issuing the follow-up GET.""" + middleware = _make_middleware() + claims = { + "_claim_names": {"groups": "src1"}, + "_claim_sources": {"src1": {"endpoint": "https://graph.microsoft.com/x"}}, + } + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"value": ["group-1"], "@odata.nextLink": "https://evil.com/next"} + + mock_client = AsyncMock() + mock_client.post.return_value = mock_response + mock_client_cm = MagicMock() + mock_client_cm.__aenter__ = AsyncMock(return_value=mock_client) + mock_client_cm.__aexit__ = AsyncMock(return_value=False) + + with patch("pyrit.backend.middleware.auth.httpx.AsyncClient", return_value=mock_client_cm): + result = await middleware._resolve_excess_groups_async(claims, "the-token") + + assert result == ["group-1"] + mock_client.get.assert_not_called()