From 3f8abe2754bcfb135723d46ecae05829b2df8c58 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 16 Apr 2024 11:39:52 -0700
Subject: [PATCH] fix(proxy_server.py): secure `/team/info` endpoint

make sure user requesting team info is part of team or admin
---
 litellm/proxy/proxy_server.py |  9 +++++++++
 tests/test_team.py            | 30 ++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 324cbd433..0f01f4e76 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -1040,6 +1040,15 @@ async def user_api_key_auth(
                 elif route == "/model/info":
                     # /model/info just shows models user has access to
                     pass
+                elif route == "/team/info":
+                    # check if key can access this team's info
+                    query_params = request.query_params
+                    team_id = query_params.get("team_id")
+                    if team_id != valid_token.team_id:
+                        raise HTTPException(
+                            status_code=status.HTTP_403_FORBIDDEN,
+                            detail="key not allowed to access this team's info",
+                        )
                 else:
                     raise Exception(
                         f"Only master key can be used to generate, delete, update info for new keys/users."
diff --git a/tests/test_team.py b/tests/test_team.py
index 7bff7b36b..2cc384a74 100644
--- a/tests/test_team.py
+++ b/tests/test_team.py
@@ -260,7 +260,18 @@ async def get_team_info(session, get_team, call_key):
 
 @pytest.mark.asyncio
 async def test_team_info():
+    """
+    Scenario 1:
+    - test with admin key -> expect to work
+    Scenario 2:
+    - test with team key -> expect to work
+    Scenario 3:
+    - test with non-team key -> expect to fail
+    """
     async with aiohttp.ClientSession() as session:
+        """
+        Scenario 1 - as admin
+        """
         new_team_data = await new_team(
             session,
             0,
@@ -268,6 +279,25 @@ async def test_team_info():
         team_id = new_team_data["team_id"]
         ## as admin ##
         await get_team_info(session=session, get_team=team_id, call_key="sk-1234")
+        """
+        Scenario 2 - as team key
+        """
+        key_gen = await generate_key(session=session, i=0, team_id=team_id)
+        key = key_gen["key"]
+
+        await get_team_info(session=session, get_team=team_id, call_key=key)
+
+        """
+        Scenario 3 - as non-team key
+        """
+        key_gen = await generate_key(session=session, i=0)
+        key = key_gen["key"]
+
+        try:
+            await get_team_info(session=session, get_team=team_id, call_key=key)
+            pytest.fail(f"Expected call to fail")
+        except Exception as e:
+            pass
 
 
 """