487 lines
17 KiB
Python
487 lines
17 KiB
Python
"""
|
|
1. Default permissions for members in a team - allowed to call /key/info and /key/health
|
|
- Create a team, create a member in a team (role = "user")
|
|
|
|
|
|
Invalid Permissions:
|
|
- User tries creating a key with team_id = team_id -> expect to fail. Invalid Permissions
|
|
- User tries editing a key with team_id = team_id -> expect to fail. Invalid Permissions
|
|
- User tries deleting a key with team_id = team_id -> expect to fail. Invalid Permissions
|
|
- User tries regenerating a key with team_id = team_id -> expect to fail. Invalid Permissions
|
|
|
|
Valid Permissions:
|
|
- User tries calling /key/info with team_id, expect to get valid response
|
|
|
|
|
|
|
|
2. Permissions - members allowd to edit, delete keys but not allowed to create keys
|
|
- Create a team with member_permissions = ["/key/update", "/key/delete", "/key/info"]
|
|
- Create a member in the team with role = "user"
|
|
|
|
Valid Permissions:
|
|
- User tries editing a key with team_id = team_id -> expect to pass. Valid Permissions
|
|
- User tries deleting a key with team_id = team_id -> expect to pass. Valid Permissions
|
|
|
|
|
|
- User tries creating a key with team_id = team_id -> expect to fail. Invalid Permissions
|
|
- User tries regenerating a key with team_id = team_id -> expect to fail. Invalid Permissions
|
|
- User tries calling /key/info with team_id, expect to get valid response
|
|
|
|
|
|
|
|
3. Permissions - members allowed to create keys but not allowed to edit, delete keys
|
|
- Create a team with member_permissions = ["/key/generate"]
|
|
- Create a member in the team with role = "user"
|
|
|
|
Valid Permissions:
|
|
- User tries creating a key with team_id = team_id -> expect to pass. Valid Permissions
|
|
|
|
Invalid Permissions:
|
|
- User tries editing a key with team_id = team_id -> expect to fail. Invalid Permissions
|
|
- User tries deleting a key with team_id = team_id -> expect to fail. Invalid Permissions
|
|
- User tries regenerating a key with team_id = team_id -> expect to fail. Invalid Permissions
|
|
"""
|
|
|
|
import pytest
|
|
import asyncio
|
|
import aiohttp, openai
|
|
import uuid
|
|
import json
|
|
from litellm.proxy._types import ProxyErrorTypes
|
|
from typing import Optional
|
|
LITELLM_MASTER_KEY = "sk-1234"
|
|
|
|
async def create_team(session, key, member_permissions=None):
|
|
url = "http://0.0.0.0:4000/team/new"
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data = {
|
|
"team_member_permissions": member_permissions
|
|
}
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
if status != 200:
|
|
raise Exception(response_text)
|
|
|
|
return await response.json()
|
|
|
|
async def create_user(session, key, user_id, team_id=None):
|
|
url = "http://0.0.0.0:4000/user/new"
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data = {
|
|
"user_id": user_id
|
|
}
|
|
if team_id:
|
|
data["team_id"] = team_id
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
if status != 200:
|
|
raise Exception(response_text)
|
|
|
|
return await response.json()
|
|
|
|
async def add_team_member(session, key, team_id, user_id, role="user"):
|
|
url = "http://0.0.0.0:4000/team/member_add"
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data = {
|
|
"team_id": team_id,
|
|
"member": {
|
|
"role": role,
|
|
"user_id": user_id
|
|
}
|
|
}
|
|
print("Adding team member with data: ", data)
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
if status != 200:
|
|
raise Exception(response_text)
|
|
|
|
return await response.json()
|
|
|
|
async def generate_key(session, key, team_id=None, user_id=None):
|
|
url = "http://0.0.0.0:4000/key/generate"
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data = {}
|
|
if team_id:
|
|
data["team_id"] = team_id
|
|
if user_id:
|
|
data["user_id"] = user_id
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
if status != 200:
|
|
return {"status": status, "error": response_text}
|
|
|
|
return await response.json()
|
|
|
|
async def key_info(session, key, key_id):
|
|
url = f"http://0.0.0.0:4000/key/info?key={key_id}"
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
async with session.get(url, headers=headers) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
if status != 200:
|
|
return {"status": status, "error": response_text}
|
|
|
|
return await response.json()
|
|
|
|
async def update_key(
|
|
session: aiohttp.ClientSession,
|
|
key: str,
|
|
key_id: str,
|
|
team_id: Optional[str] = None,
|
|
):
|
|
"""
|
|
Update a key
|
|
|
|
Args:
|
|
key: key to use for authentication
|
|
key_id: key to update
|
|
"""
|
|
url = "http://0.0.0.0:4000/key/update"
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data = {
|
|
"key": key_id,
|
|
"metadata": {"updated": True}
|
|
}
|
|
if team_id:
|
|
data["team_id"] = team_id
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
if status != 200:
|
|
return {"status": status, "error": response_text}
|
|
|
|
return await response.json()
|
|
|
|
async def delete_key(session, key, key_id):
|
|
url = "http://0.0.0.0:4000/key/delete"
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data = {
|
|
"keys": [key_id]
|
|
}
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
if status != 200:
|
|
return {"status": status, "error": response_text}
|
|
|
|
return await response.json()
|
|
|
|
async def regenerate_key(session, key, key_id, team_id=None):
|
|
url = "http://0.0.0.0:4000/key/regenerate"
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data = {
|
|
"key": key_id
|
|
}
|
|
if team_id:
|
|
data["team_id"] = team_id
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
if status != 200:
|
|
return {"status": status, "error": response_text}
|
|
|
|
return await response.json()
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_default_member_permissions():
|
|
"""
|
|
Test default permissions for members in a team - allowed to call /key/info and /key/health
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
master_key = LITELLM_MASTER_KEY
|
|
|
|
# Create a team
|
|
team_data = await create_team(
|
|
session=session,
|
|
key=master_key
|
|
)
|
|
team_id = team_data["team_id"]
|
|
|
|
# create a team key
|
|
team_key_data = await generate_key(
|
|
session=session,
|
|
key=master_key,
|
|
team_id=team_id
|
|
)
|
|
team_key = team_key_data["key"]
|
|
|
|
# create a user
|
|
user_data = await create_user(
|
|
session=session,
|
|
key=master_key,
|
|
user_id=f"user_{uuid.uuid4().hex[:8]}",
|
|
team_id=team_id
|
|
)
|
|
user_id = user_data["user_id"]
|
|
|
|
# Create a user key
|
|
print("New user data: ", user_data)
|
|
|
|
# Create a user key
|
|
user_key_data = await generate_key(
|
|
session=session,
|
|
key=master_key,
|
|
user_id=user_id
|
|
)
|
|
print("new user key: ", user_key_data)
|
|
user_key = user_key_data["key"]
|
|
|
|
# Test invalid permissions
|
|
# User tries creating a key with team_id
|
|
print("Regular team member trying to create a key with team_id. Expecting error.")
|
|
create_result = await generate_key(
|
|
session=session,
|
|
key=user_key,
|
|
team_id=team_id
|
|
)
|
|
print("result: ", create_result)
|
|
assert "status" in create_result and create_result["status"] == 401, "User should not be able to create keys for team"
|
|
error_data = json.loads(create_result["error"])
|
|
print("error response =", json.dumps(error_data, indent=4))
|
|
assert error_data["error"]["type"] == ProxyErrorTypes.team_member_permission_error.value, "Error should be a team member permission error"
|
|
|
|
# User tries editing a key with team_id
|
|
print("Regular team member trying to edit a key with team_id. Expecting error.")
|
|
update_result = await update_key(
|
|
session=session,
|
|
key=user_key,
|
|
key_id=team_key,
|
|
team_id="ATTACKER_TEAM_ID"
|
|
)
|
|
assert "status" in update_result and update_result["status"] == 401, "User should not be able to update keys for team"
|
|
error_data = json.loads(update_result["error"])
|
|
print("error response =", json.dumps(error_data, indent=4))
|
|
assert error_data["error"]["type"] == ProxyErrorTypes.team_member_permission_error.value, "Error should be a team member permission error"
|
|
|
|
# User tries deleting a key with team_id
|
|
print("Regular team member trying to delete a key with team_id. Expecting error.")
|
|
delete_result = await delete_key(
|
|
session=session,
|
|
key=user_key,
|
|
key_id=team_key,
|
|
)
|
|
assert "status" in delete_result and delete_result["status"] == 401, "User should not be able to delete keys for team"
|
|
error_data = json.loads(delete_result["error"])
|
|
print("error response =", json.dumps(error_data, indent=4))
|
|
assert error_data["error"]["type"] == ProxyErrorTypes.team_member_permission_error.value, "Error should be a team member permission error"
|
|
|
|
# User tries regenerating a key with team_id
|
|
print("Regular team member trying to regenerate a key with team_id. Expecting error.")
|
|
regenerate_result = await regenerate_key(
|
|
session=session,
|
|
key=user_key,
|
|
key_id=team_key,
|
|
)
|
|
assert "status" in regenerate_result and regenerate_result["status"] == 401, "User should not be able to regenerate keys for team"
|
|
error_data = json.loads(regenerate_result["error"])
|
|
print("error response =", json.dumps(error_data, indent=4))
|
|
assert error_data["error"]["type"] == ProxyErrorTypes.team_member_permission_error.value, "Error should be a team member permission error"
|
|
|
|
# Test valid permissions
|
|
# User tries calling /key/info with team_id
|
|
print("Regular team member trying to get key info with team_id. Expecting success.")
|
|
info_result = await key_info(
|
|
session=session,
|
|
key=user_key,
|
|
key_id=team_key,
|
|
)
|
|
print("info result =", info_result)
|
|
assert "status" not in info_result, "Admin should be able to get key info"
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_edit_delete_permissions():
|
|
"""
|
|
Test permissions - members allowed to edit, delete keys but not allowed to create keys
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
master_key = LITELLM_MASTER_KEY
|
|
|
|
# Create a team with specific member permissions
|
|
team_data = await create_team(
|
|
session=session,
|
|
key=master_key,
|
|
member_permissions=["/key/update", "/key/delete", "/key/info"]
|
|
)
|
|
team_id = team_data["team_id"]
|
|
|
|
# create a user in team=team_id
|
|
user_data = await create_user(
|
|
session=session,
|
|
key=master_key,
|
|
user_id=f"user_{uuid.uuid4().hex[:8]}",
|
|
team_id=team_id
|
|
)
|
|
user_id = user_data["user_id"]
|
|
|
|
# Generate an admin key for the team
|
|
admin_key_data = await generate_key(session, master_key, team_id)
|
|
key_id = admin_key_data["key"]
|
|
|
|
# Create a user key
|
|
user_key_data = await generate_key(
|
|
session=session,
|
|
key=master_key,
|
|
user_id=user_id
|
|
)
|
|
user_key = user_key_data["key"]
|
|
|
|
# Test valid permissions
|
|
# User tries editing a key with team_id
|
|
update_result = await update_key(
|
|
session=session,
|
|
key=user_key,
|
|
key_id=key_id,
|
|
team_id=team_id
|
|
)
|
|
assert "status" not in update_result, "User should be able to update keys for team"
|
|
|
|
# User tries deleting a key with team_id - test this last
|
|
delete_result = await delete_key(
|
|
session=session,
|
|
key=user_key,
|
|
key_id=key_id
|
|
)
|
|
assert "status" not in delete_result, "User should be able to delete keys for team"
|
|
|
|
# Test invalid permissions
|
|
# User tries creating a key with team_id
|
|
create_result = await generate_key(
|
|
session=session,
|
|
key=user_key,
|
|
team_id=team_id
|
|
)
|
|
assert "status" in create_result and create_result["status"] != 200, "User should not be able to create keys for team"
|
|
|
|
# User tries regenerating a key with team_id
|
|
regenerate_result = await regenerate_key(
|
|
session=session,
|
|
key=user_key,
|
|
key_id=key_id,
|
|
team_id=team_id
|
|
)
|
|
assert "status" in regenerate_result and regenerate_result["status"] != 200, "User should not be able to regenerate keys for team"
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_create_permissions():
|
|
"""
|
|
Test permissions - members allowed to create keys but not allowed to edit, delete keys
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
master_key = LITELLM_MASTER_KEY
|
|
|
|
# Create a team with specific member permissions
|
|
team_data = await create_team(
|
|
session=session,
|
|
key=master_key,
|
|
member_permissions=["/key/generate"]
|
|
)
|
|
team_id = team_data["team_id"]
|
|
|
|
# Create a user in the team
|
|
user_id = f"user_{uuid.uuid4().hex[:8]}"
|
|
await add_team_member(
|
|
session=session,
|
|
key=master_key,
|
|
team_id=team_id,
|
|
user_id=user_id,
|
|
role="user"
|
|
)
|
|
|
|
# Generate an admin key for the team
|
|
admin_key_data = await generate_key(
|
|
session=session,
|
|
key=master_key,
|
|
team_id=team_id
|
|
)
|
|
admin_key = admin_key_data["key"]
|
|
key_id = admin_key_data["key"]
|
|
|
|
# Create a user key
|
|
user_key_data = await generate_key(
|
|
session=session,
|
|
key=master_key,
|
|
user_id=user_id
|
|
)
|
|
user_key = user_key_data["key"]
|
|
|
|
# Test valid permissions
|
|
# User tries creating a key with team_id
|
|
create_result = await generate_key(
|
|
session=session,
|
|
key=user_key,
|
|
team_id=team_id
|
|
)
|
|
print("success, user created key for team=", create_result)
|
|
assert "key" in create_result, "User should be able to create keys for team"
|
|
assert create_result["team_id"] == team_id, "User should be able to create keys for team"
|
|
assert "status" not in create_result, "User should be able to create keys for team"
|
|
|
|
# Test invalid permissions
|
|
# User tries editing a key with team_id
|
|
update_result = await update_key(
|
|
session=session,
|
|
key=user_key,
|
|
key_id=key_id,
|
|
team_id=team_id
|
|
)
|
|
assert "status" in update_result and update_result["status"] != 200, "User should not be able to update keys for team"
|
|
|
|
# User tries deleting a key with team_id
|
|
delete_result = await delete_key(
|
|
session=session,
|
|
key=user_key,
|
|
key_id=key_id
|
|
)
|
|
assert "status" in delete_result and delete_result["status"] != 200, "User should not be able to delete keys for team"
|
|
|
|
# User tries regenerating a key with team_id
|
|
regenerate_result = await regenerate_key(
|
|
session=session,
|
|
key=user_key,
|
|
key_id=key_id,
|
|
team_id=team_id
|
|
)
|
|
assert "status" in regenerate_result and regenerate_result["status"] != 200, "User should not be able to regenerate keys for team" |