From d90f44fe8e9c0f17b04243e587f96e38e82f2eae Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 26 Mar 2024 14:01:02 -0700
Subject: [PATCH] fix(proxy_server.py): check if team scope in jwt

---
 litellm/proxy/auth/handle_jwt.py | 5 +++++
 litellm/proxy/proxy_server.py    | 6 ++++++
 2 files changed, 11 insertions(+)

diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py
index ec7f75562f..6cb67b1717 100644
--- a/litellm/proxy/auth/handle_jwt.py
+++ b/litellm/proxy/auth/handle_jwt.py
@@ -85,6 +85,11 @@ class JWTHandler:
             return True
         return False
 
+    def is_team(self, scopes: list) -> bool:
+        if self.litellm_proxy_roles.team_jwt_scope in scopes:
+            return True
+        return False
+
     def get_end_user_id(self, token: dict, default_value: Optional[str]) -> str:
         try:
             if self.litellm_proxy_roles.team_id_jwt_field is not None:
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index f405630f88..d16b297fe9 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -390,6 +390,12 @@ async def user_api_key_auth(
                         raise Exception(
                             f"Admin not allowed to access this route. Route={route}, Allowed Routes={actual_routes}"
                         )
+                # check if team in scopes
+                is_team = jwt_handler.is_team(scopes=scopes)
+                if is_team == False:
+                    raise Exception(
+                        f"Missing both Admin and Team scopes from token. Either is required. Admin Scope={jwt_handler.litellm_proxy_roles.admin_jwt_scope}, Team Scope={jwt_handler.litellm_proxy_roles.team_jwt_scope}"
+                    )
                 # get team id
                 team_id = jwt_handler.get_team_id(token=valid_token, default_value=None)