fix security for global_spend_reset

This commit is contained in:
Ishaan Jaff 2024-05-14 11:04:50 -07:00
parent 3686789c36
commit 7e56e27226
2 changed files with 16 additions and 28 deletions

View file

@ -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",

View file

@ -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={}
)