diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e935a5d9093..24dd26df65d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -6081,8 +6081,6 @@ async def list_keys( # an exact user_id (e.g. an integration scoping to one user with an admin # key) receive other users' keys (user_id="alice" -> "alice2"). Exact by # default restores the prior behavior; the dashboard opts in explicitly. - # user_id substring stays admin-only: non-admins are scoped to their own - # user_id below. key_alias is a global AND filter, so it only narrows. use_substring_matching: Final = substring_matching and is_proxy_admin use_key_alias_substring_matching: Final = substring_matching diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 71a3b169798..bb1f5108d70 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -18,6 +18,7 @@ from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LiteLLM_VerificationToken, @@ -9385,34 +9386,6 @@ async def test_build_key_filter_admin_substring_matching(): assert {"key_alias": {"contains": key_alias, "mode": "insensitive"}} in where["AND"] -def test_build_key_filter_key_alias_substring_keeps_user_id_exact(): - """A non-admin searching a team's keys by partial alias gets a substring - key_alias filter while their own-user scoping stays exact, so alias search - can never widen visibility to another user (user_id="alice" -> "alice2").""" - from litellm.proxy.management_endpoints.key_management_endpoints import ( - _build_key_filter_conditions, - ) - - where = _build_key_filter_conditions( - user_id="alice", - team_id="team-a", - organization_id=None, - key_alias="first", - key_hash=None, - exclude_team_id=None, - admin_team_ids=None, - member_team_ids=["team-a"], - include_created_by_keys=False, - use_substring_matching=False, - use_key_alias_substring_matching=True, - ) - - assert {"key_alias": {"contains": "first", "mode": "insensitive"}} in where["AND"] - assert {"key_alias": "first"} not in where["AND"] - assert json.dumps({"user_id": "alice"}) in json.dumps(where) - assert '"contains": "alice"' not in json.dumps(where) - - @pytest.mark.asyncio async def test_build_key_filter_non_admin_exact_matching(): """ @@ -15188,33 +15161,104 @@ async def test_list_keys_non_admin_cannot_opt_into_substring(): assert kwargs["user_id"] == "alice" -@pytest.mark.asyncio +def _prisma_where_matches(row, where): + for field, expected in where.items(): + if field == "AND": + if not all(_prisma_where_matches(row, child) for child in expected): + return False + elif field == "OR": + if not any(_prisma_where_matches(row, child) for child in expected): + return False + elif isinstance(expected, dict): + value = getattr(row, field) + if "in" in expected and value not in expected["in"]: + return False + if "not" in expected and value == expected["not"]: + return False + if "contains" in expected: + haystack, needle = value or "", expected["contains"] + if expected.get("mode") == "insensitive": + haystack, needle = haystack.lower(), needle.lower() + if needle not in haystack: + return False + elif getattr(row, field) != expected: + return False + return True + + +class _InMemoryVerificationTokenTable: + def __init__(self, rows): + self.rows = rows + + async def find_many(self, where, **kwargs): + return [row for row in self.rows if _prisma_where_matches(row, where)] + + async def count(self, where): + return len(await self.find_many(where)) + + +def _team_key(token, key_alias, user_id): + return LiteLLM_VerificationToken(token=token, key_alias=key_alias, user_id=user_id, team_id="team-a") + + +_TEAM_A_KEYS = ( + _team_key("tok-alice-first", "app_llmhub_first.last", "alice"), + _team_key("tok-alice-other", "alice_other_key", "alice"), + _team_key("tok-bob-first", "bob_First_key", "bob"), + _team_key("tok-svc-first", "service_first_key", None), +) + + +def _list_team_a_keys_as(user_role, members_with_roles, query): + from fastapi import FastAPI + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.management_endpoints.key_management_endpoints import router + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken = _InMemoryVerificationTokenTable(_TEAM_A_KEYS) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="alice", teams=["team-a"], organization_memberships=[]) + ) + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[LiteLLM_TeamTable(team_id="team-a", members_with_roles=members_with_roles)] + ) + test_app = FastAPI() + test_app.include_router(router) + test_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=user_role, user_id="alice") + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + response = TestClient(test_app).get( + f"/key/list?team_id=team-a&include_team_keys=true&include_created_by_keys=true&{query}" + ) + assert response.status_code == 200, response.text + return sorted(response.json()["keys"]) + + +_ALICE_TEAM_ADMIN = [Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")] +_ALICE_TEAM_MEMBER = [Member(user_id="alice", role="user"), Member(user_id="bob", role="user")] + + @pytest.mark.parametrize( "user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, LitellmUserRoles.TEAM], ) -async def test_list_keys_non_admin_key_alias_substring_is_honored(user_role): - """Team admins and internal users searching a team's keys by a partial alias - (key_alias=first for app_llmhub_first.last) must get substring matching on - key_alias, while user_id substring matching stays admin-only.""" - user = UserAPIKeyAuth(user_role=user_role, user_id="alice") - kwargs = await _list_keys_capture_helper_kwargs( - user, user_id=None, key_alias="first", team_id="team-a", substring_matching=True - ) - assert kwargs["use_key_alias_substring_matching"] is True - assert kwargs["use_substring_matching"] is False - assert kwargs["key_alias"] == "first" - assert kwargs["user_id"] == "alice" +def test_list_keys_team_admin_key_alias_substring_returns_every_matching_team_key(user_role): + keys = _list_team_a_keys_as(user_role, _ALICE_TEAM_ADMIN, "key_alias=first&substring_matching=true") + assert keys == ["tok-alice-first", "tok-bob-first", "tok-svc-first"] -@pytest.mark.asyncio -async def test_list_keys_key_alias_substring_defaults_off(): - """Without substring_matching, key_alias stays an exact filter for every role.""" - user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice") - kwargs = await _list_keys_capture_helper_kwargs( - user, user_id=None, key_alias="first", substring_matching=False +def test_list_keys_team_member_key_alias_substring_stays_within_own_visibility(): + keys = _list_team_a_keys_as( + LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_MEMBER, "key_alias=first&substring_matching=true" ) - assert kwargs["use_key_alias_substring_matching"] is False + assert keys == ["tok-alice-first", "tok-svc-first"] + + +def test_list_keys_key_alias_stays_exact_without_substring_matching(): + assert _list_team_a_keys_as(LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_ADMIN, "key_alias=first") == [] + assert _list_team_a_keys_as( + LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_ADMIN, "key_alias=app_llmhub_first.last" + ) == ["tok-alice-first"] @pytest.mark.asyncio