From 8053b43c09ce2556cbea7d9fdde66d8df9ac48e0 Mon Sep 17 00:00:00 2001 From: jsong468 Date: Mon, 13 Jul 2026 18:17:08 -0700 Subject: [PATCH 1/3] docstring --- pyrit/backend/middleware/auth.py | 27 +++++++++++++++- tests/unit/backend/test_auth_middleware.py | 36 ++++++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/pyrit/backend/middleware/auth.py b/pyrit/backend/middleware/auth.py index 81594fc66e..a532517e91 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,9 @@ class EntraAuthMiddleware(BaseHTTPMiddleware): "/api/media", } + # Hosts the user's Bearer token may be forwarded to + _GRAPH_HOSTS: ClassVar[set[str]] = {"graph.microsoft.com"} + def __init__(self, app: ASGIApp) -> None: """Initialize the middleware with Entra ID configuration from environment variables.""" super().__init__(app) @@ -172,12 +176,16 @@ async def _resolve_excess_groups_async(self, claims: dict[str, Any], token: str) Microsoft Graph `getMemberObjects` endpoint to retrieve transitive group memberships, using the user's access token. + The token is only forwarded to hosts in the Graph allowlist (see + ``_is_trusted_graph_url``); untrusted endpoints and pagination links are refused. + Args: claims: The decoded JWT claims containing _claim_sources. 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 the resolution endpoint is not a trusted Graph host. """ try: claim_sources = claims.get("_claim_sources", {}) @@ -196,6 +204,10 @@ async def _resolve_excess_groups_async(self, claims: dict[str, Any], token: str) # Graph format: https://graph.microsoft.com/v1.0/me/getMemberObjects endpoint = "https://graph.microsoft.com/v1.0/me/getMemberObjects" + if not self._is_trusted_graph_url(endpoint): + logger.warning("Refusing to forward token to untrusted group resolution endpoint: %s", endpoint) + return [] + all_group_ids: list[str] = [] async with httpx.AsyncClient() as client: response = await client.post( @@ -222,6 +234,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 +256,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..03507ccb47 100644 --- a/tests/unit/backend/test_auth_middleware.py +++ b/tests/unit/backend/test_auth_middleware.py @@ -7,9 +7,16 @@ from unittest.mock import 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,32 @@ 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_refuses_untrusted_endpoint(): + """An untrusted _claim_sources endpoint returns [] without making any HTTP request.""" + middleware = _make_middleware() + claims = {"_claim_sources": {"src1": {"endpoint": "https://evil.com/steal"}}} + + 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() From 4a9781ba06c199685efb044cfe64f034afcc3c2a Mon Sep 17 00:00:00 2001 From: jsong468 Date: Tue, 14 Jul 2026 13:25:26 -0700 Subject: [PATCH 2/3] pr feedback change --- pyrit/backend/middleware/auth.py | 46 ++++++++++------------ tests/unit/backend/test_auth_middleware.py | 35 ++++++++++++++-- 2 files changed, 52 insertions(+), 29 deletions(-) diff --git a/pyrit/backend/middleware/auth.py b/pyrit/backend/middleware/auth.py index a532517e91..c2c24acd94 100644 --- a/pyrit/backend/middleware/auth.py +++ b/pyrit/backend/middleware/auth.py @@ -55,6 +55,10 @@ class EntraAuthMiddleware(BaseHTTPMiddleware): # 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) @@ -171,42 +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 token is only forwarded to hosts in the Graph allowlist (see - ``_is_trusted_graph_url``); untrusted endpoints and pagination links are refused. + 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 or - when the resolution endpoint is not a trusted Graph host. + 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" - - if not self._is_trusted_graph_url(endpoint): - logger.warning("Refusing to forward token to untrusted group resolution endpoint: %s", endpoint) - return [] + endpoint = self._GRAPH_MEMBER_OBJECTS_URL all_group_ids: list[str] = [] async with httpx.AsyncClient() as client: diff --git a/tests/unit/backend/test_auth_middleware.py b/tests/unit/backend/test_auth_middleware.py index 03507ccb47..db362670c7 100644 --- a/tests/unit/backend/test_auth_middleware.py +++ b/tests/unit/backend/test_auth_middleware.py @@ -5,7 +5,7 @@ Tests for the Entra ID auth middleware. """ -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -53,13 +53,40 @@ def test_is_trusted_graph_url(url, expected): assert _make_middleware()._is_trusted_graph_url(url) is expected -async def test_resolve_excess_groups_refuses_untrusted_endpoint(): - """An untrusted _claim_sources endpoint returns [] without making any HTTP request.""" +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://evil.com/steal"}}} + 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 From d65a380fd933394c2fd3ff9717b5f10bd97964e8 Mon Sep 17 00:00:00 2001 From: jsong468 Date: Tue, 14 Jul 2026 14:19:36 -0700 Subject: [PATCH 3/3] address test coverage --- tests/unit/backend/test_auth_middleware.py | 25 ++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/unit/backend/test_auth_middleware.py b/tests/unit/backend/test_auth_middleware.py index db362670c7..fdadc199e5 100644 --- a/tests/unit/backend/test_auth_middleware.py +++ b/tests/unit/backend/test_auth_middleware.py @@ -90,3 +90,28 @@ async def test_resolve_excess_groups_ignores_token_endpoint(): 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()