forked from phoenix/litellm-mirror
fix security for global_spend_reset
This commit is contained in:
parent
3686789c36
commit
7e56e27226
2 changed files with 16 additions and 28 deletions
|
@ -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",
|
||||
|
|
|
@ -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={}
|
||||
)
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue