fix(proxy_server.py): prevent user from deleting non-user owned keys when they use ui

This commit is contained in:
Krrish Dholakia 2024-03-11 12:13:30 -07:00
parent 40c9682de7
commit 4eb244c3ca
2 changed files with 44 additions and 12 deletions

View file

@ -1356,9 +1356,12 @@ class PrismaClient:
tokens: Optional[List] = None,
team_id_list: Optional[List] = None,
table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None,
user_id: Optional[str] = None,
):
"""
Allow user to delete a key(s)
Ensure user owns that key, unless admin.
"""
try:
if tokens is not None and isinstance(tokens, List):
@ -1369,15 +1372,23 @@ class PrismaClient:
else:
hashed_token = token
hashed_tokens.append(hashed_token)
await self.db.litellm_verificationtoken.delete_many(
where={"token": {"in": hashed_tokens}}
filter_query: dict = {}
if user_id is not None:
filter_query = {
"AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}]
}
else:
filter_query = {"token": {"in": hashed_tokens}}
deleted_tokens = await self.db.litellm_verificationtoken.delete_many(
where=filter_query # type: ignore
)
return {"deleted_keys": tokens}
return {"deleted_keys": deleted_tokens}
elif (
table_name == "team"
and team_id_list is not None
and isinstance(team_id_list, List)
):
# admin only endpoint -> `/team/delete`
await self.db.litellm_teamtable.delete_many(
where={"team_id": {"in": team_id_list}}
)
@ -1387,6 +1398,7 @@ class PrismaClient:
and team_id_list is not None
and isinstance(team_id_list, List)
):
# admin only endpoint -> `/team/delete`
await self.db.litellm_verificationtoken.delete_many(
where={"team_id": {"in": team_id_list}}
)