fix(proxy_server.py): fix key create logic + add unit tests

litellm/proxy/proxy_server.py

@@ -1010,8 +1010,10 @@ async def user_api_key_auth(
                         db=custom_db_client,
                     )
                 )
-            if route in LiteLLMRoutes.info_routes.value and (
+
-                not _is_user_proxy_admin(user_id_information)
+            if not _is_user_proxy_admin(user_id_information):  # if non-admin
+                if (
+                    route in LiteLLMRoutes.info_routes.value
                 ):  # check if user allowed to call an info route
                     if route == "/key/info":
                         # check if user can access this route
@@ -1051,7 +1053,7 @@ async def user_api_key_auth(
                             )
                 else:
                     raise Exception(
-                        f"Only master key can be used to generate, delete, update info for new keys/users."
+                        f"Only master key can be used to generate, delete, update info for new keys/users/teams."
                     )
 
         # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions

tests/test_keys.py

@@ -44,9 +44,13 @@ async def generate_key(
     models=["azure-models", "gpt-4", "dall-e-3"],
     max_parallel_requests: Optional[int] = None,
     user_id: Optional[str] = None,
+    calling_key="sk-1234",
 ):
     url = "http://0.0.0.0:4000/key/generate"
-    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    headers = {
+        "Authorization": f"Bearer {calling_key}",
+        "Content-Type": "application/json",
+    }
     data = {
         "models": models,
         "aliases": {"mistral-7b": "gpt-3.5-turbo"},
@@ -80,6 +84,35 @@ async def test_key_gen():
         await asyncio.gather(*tasks)
 
 
+@pytest.mark.asyncio
+async def test_key_gen_bad_key():
+    """
+    Test if you can create a key with a non-admin key, even with UI setup
+    """
+    async with aiohttp.ClientSession() as session:
+        ## LOGIN TO UI
+        form_data = {"username": "admin", "password": "sk-1234"}
+        async with session.post(
+            "http://0.0.0.0:4000/login", data=form_data
+        ) as response:
+            assert (
+                response.status == 200
+            )  # Assuming the endpoint returns a 500 status code for error handling
+            text = await response.text()
+            print(text)
+        ## create user key with admin key -> expect to work
+        key_data = await generate_key(session=session, i=0, user_id="user-1234")
+        key = key_data["key"]
+        ## create new key with user key -> expect to fail
+        try:
+            await generate_key(
+                session=session, i=0, user_id="user-1234", calling_key=key
+            )
+            pytest.fail("Expected to fail")
+        except Exception as e:
+            pass
+
+
 async def update_key(session, get_key):
     """
     Make sure only models user has access to are returned