diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f91b768e6..5dbdd3a71 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -79,6 +79,11 @@ class LiteLLMRoutes(enum.Enum): "/v1/models", ] + # NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend + master_key_only_routes: List = [ + "/global/spend/reset", + ] + info_routes: List = [ "/key/info", "/team/info", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9b2e01bc8..30f0275b6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -589,6 +589,15 @@ async def user_api_key_auth( ) return _user_api_key_obj + + ## IF it's not a master key + ## Route should not be in master_key_only_routes + if route in LiteLLMRoutes.master_key_only_routes.value: + raise Exception( + f"Tried to access route={route}, which is only for MASTER KEY" + ) + + ## Check DB if isinstance( api_key, str ): # if generated token, make sure it starts with sk-. @@ -5927,9 +5936,7 @@ async def view_spend_logs( tags=["Budget & Spend Tracking"], dependencies=[Depends(user_api_key_auth)], ) -async def global_spend_reset( - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +async def global_spend_reset(): """ ADMIN ONLY / MASTER KEY Only Endpoint @@ -5940,7 +5947,7 @@ async def global_spend_reset( 3. LiteLLM_TeamTable spend will be set = 0 """ - global prisma_client, master_key + global prisma_client if prisma_client is None: raise ProxyException( message="Prisma Client is not initialized", @@ -5949,30 +5956,6 @@ async def global_spend_reset( code=status.HTTP_401_UNAUTHORIZED, ) - if master_key is None: - raise ProxyException( - message="Master key is not initialized, please set LITELLM_MASTER_KEY in .env", - type="internal_error", - param="None", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if user_api_key_dict.api_key is None: - raise ProxyException( - message="no api_key passed", - type="auth_error", - param="master_key", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if not secrets.compare_digest(master_key, user_api_key_dict.api_key): - raise ProxyException( - message="/global/spend/reset Route only allowed for master key", - type="auth_error", - param="master_key", - code=status.HTTP_401_UNAUTHORIZED, - ) - await prisma_client.db.litellm_verificationtoken.update_many( data={"spend": 0.0}, where={} )