From 1a57875d24d26bb42e44f29b51ec0d82d3d78d6a Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Thu, 24 Jul 2025 16:40:40 -0700 Subject: [PATCH] Proxy - specify `key_type` - allows specifying if key can call LLM API routes vs. Management routes only (#12909) * feat(key_management_endpoints.py): Support new 'key_type' field allow user to specify if key should be 'management' or 'llm api' key Security fix * test(test_route_checks.py): add unit tests * fix(create_key_button.tsx): add ui component to select key type allows specifying if key can call llm api vs. management routes * feat(create_key_button.tsx): add specifying key type to ui * fix(route_checks.py): add sensitive data masker for user id on not allowed error message prevent leaking sensitive information --- litellm/proxy/_types.py | 19 ++- litellm/proxy/auth/route_checks.py | 38 +++++- .../key_management_endpoints.py | 23 +++- .../proxy/auth/test_route_checks.py | 116 ++++++++++++++++ .../src/components/create_key_button.tsx | 127 ++++++++++++++---- 5 files changed, 289 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3fbf887475..3353081109 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -696,9 +696,24 @@ class KeyRequestBase(GenerateRequestBase): allowed_routes: Optional[list] = [] +class LiteLLMKeyType(str, enum.Enum): + """ + Enum for key types that determine what routes a key can access + """ + + LLM_API = "llm_api" # Can call LLM API routes (chat/completions, embeddings, etc.) + MANAGEMENT = "management" # Can call management routes (user/team/key management) + READ_ONLY = "read_only" # Can only call info/read routes + DEFAULT = "default" # Uses default allowed routes + + class GenerateKeyRequest(KeyRequestBase): soft_budget: Optional[float] = None send_invite_email: Optional[bool] = None + key_type: Optional[LiteLLMKeyType] = Field( + default=LiteLLMKeyType.DEFAULT, + description="Type of key that determines default allowed routes.", + ) class GenerateKeyResponse(KeyRequestBase): @@ -1766,16 +1781,16 @@ class UserAPIKeyAuth( if JWTHandler.is_jwt(token=api_key): return f"hashed-jwt-{hash_token(token=api_key)}" return api_key - @classmethod def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth": """ Returns a `UserAPIKeyAuth` object for the litellm internal health check service account. - + This is used to track number of requests/spend for health check calls. """ from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME + return cls( api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, team_id=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 76c8d2f7cd..e483ac6a66 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -54,6 +54,19 @@ class RouteChecks: if route in valid_token.allowed_routes: return True + ## check if 'allowed_route' is a field name in LiteLLMRoutes + if any( + allowed_route in LiteLLMRoutes._member_names_ + for allowed_route in valid_token.allowed_routes + ): + for allowed_route in valid_token.allowed_routes: + if allowed_route in LiteLLMRoutes._member_names_: + if RouteChecks.check_route_access( + route=route, + allowed_routes=LiteLLMRoutes._member_map_[allowed_route].value, + ): + return True + # check if wildcard pattern is allowed for allowed_route in valid_token.allowed_routes: if RouteChecks._route_matches_wildcard_pattern( @@ -65,6 +78,27 @@ class RouteChecks: f"Virtual key is not allowed to call this route. Only allowed to call routes: {valid_token.allowed_routes}. Tried to call route: {route}" ) + @staticmethod + def _mask_user_id(user_id: str) -> str: + """ + Mask user_id to prevent leaking sensitive information in error messages + + Args: + user_id (str): The user_id to mask + + Returns: + str: Masked user_id showing only first 2 and last 2 characters + """ + from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker + + if not user_id or len(user_id) <= 4: + return "***" + + # Use SensitiveDataMasker with custom configuration for user_id + masker = SensitiveDataMasker(visible_prefix=6, visible_suffix=2, mask_char="*") + + return masker._mask_value(user_id) + @staticmethod def non_proxy_admin_allowed_routes_check( user_obj: Optional[LiteLLM_UserTable], @@ -182,8 +216,10 @@ class RouteChecks: if user_obj is not None: user_role = user_obj.user_role or "unknown" user_id = user_obj.user_id or "unknown" + + masked_user_id = RouteChecks._mask_user_id(user_id) raise Exception( - f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={user_id}" + f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={masked_user_id}" ) @staticmethod diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8998d96e48..de8de75d07 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -46,8 +46,8 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) from litellm.proxy.management_helpers.object_permission_utils import ( - handle_update_object_permission_common, attach_object_permission_to_dict, + handle_update_object_permission_common, ) from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, @@ -331,6 +331,21 @@ def common_key_access_checks( router = APIRouter() +def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: + """ + Handle the key type. + """ + key_type = data.key_type + data_json.pop("key_type", None) + if key_type == LiteLLMKeyType.LLM_API: + data_json["allowed_routes"] = ["llm_api_routes"] + elif key_type == LiteLLMKeyType.MANAGEMENT: + data_json["allowed_routes"] = ["management_routes"] + elif key_type == LiteLLMKeyType.READ_ONLY: + data_json["allowed_routes"] = ["info_routes"] + return data_json + + async def _common_key_generation_helper( # noqa: PLR0915 data: GenerateKeyRequest, user_api_key_dict: UserAPIKeyAuth, @@ -455,6 +470,7 @@ async def _common_key_generation_helper( # noqa: PLR0915 data_json = data.model_dump(exclude_unset=True, exclude_none=True) # type: ignore + data_json = handle_key_type(data, data_json) # if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users if "max_budget" in data_json: data_json["key_max_budget"] = data_json.pop("max_budget", None) @@ -2600,7 +2616,7 @@ async def _list_key_helper( key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object else: _token = key_dict.get("token") - key_list.append(_token) # Return only the token + key_list.append(cast(str, _token)) # Return only the token return KeyListResponseObject( keys=key_list, @@ -3099,6 +3115,3 @@ def validate_model_max_budget(model_max_budget: Optional[Dict]) -> None: raise ValueError( f"Invalid model_max_budget: {str(e)}. Example of valid model_max_budget: https://docs.litellm.ai/docs/proxy/users" ) - - - diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 42838f02f6..b42b362b41 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -89,3 +89,119 @@ def test_proxy_admin_viewer_config_update_route_rejected(): assert exc_info.value.status_code == 403 assert "user not allowed to access this route" in str(exc_info.value.detail) assert "role= proxy_admin_viewer" in str(exc_info.value.detail) + + +def test_virtual_key_allowed_routes_with_litellm_routes_member_name_allowed(): + """Test that virtual key is allowed to call routes when allowed_routes contains LiteLLMRoutes member name""" + + # Create a UserAPIKeyAuth with allowed_routes containing a LiteLLMRoutes member name + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["openai_routes"], # This is a member name in LiteLLMRoutes enum + ) + + # Test that a route from the openai_routes group is allowed + result = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/chat/completions", # This is in LiteLLMRoutes.openai_routes.value + valid_token=valid_token, + ) + + assert result is True + + +def test_virtual_key_allowed_routes_with_litellm_routes_member_name_denied(): + """Test that virtual key is denied when route is not in the allowed LiteLLMRoutes group""" + + # Create a UserAPIKeyAuth with allowed_routes containing a LiteLLMRoutes member name + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["info_routes"], # This is a member name in LiteLLMRoutes enum + ) + + # Test that a route NOT in the info_routes group raises an exception + with pytest.raises(Exception) as exc_info: + RouteChecks.is_virtual_key_allowed_to_call_route( + route="/chat/completions", # This is NOT in LiteLLMRoutes.info_routes.value + valid_token=valid_token, + ) + + # Verify the exception message + assert "Virtual key is not allowed to call this route" in str(exc_info.value) + assert "Only allowed to call routes: ['info_routes']" in str(exc_info.value) + assert "Tried to call route: /chat/completions" in str(exc_info.value) + + +def test_virtual_key_allowed_routes_with_multiple_litellm_routes_member_names(): + """Test that virtual key works with multiple LiteLLMRoutes member names in allowed_routes""" + + # Create a UserAPIKeyAuth with multiple LiteLLMRoutes member names + valid_token = UserAPIKeyAuth( + user_id="test_user", allowed_routes=["openai_routes", "info_routes"] + ) + + # Test that routes from both groups are allowed + result1 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/chat/completions", valid_token=valid_token # This is in openai_routes + ) + + result2 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/user/info", valid_token=valid_token # This is in info_routes + ) + + assert result1 is True + assert result2 is True + + +def test_virtual_key_allowed_routes_with_mixed_member_names_and_explicit_routes(): + """Test that virtual key works with both LiteLLMRoutes member names and explicit routes""" + + # Create a UserAPIKeyAuth with both member names and explicit routes + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=[ + "info_routes", + "/custom/route", + ], # Mix of member name and explicit route + ) + + # Test that both info routes and explicit custom route are allowed + result1 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/user/info", valid_token=valid_token # This is in info_routes + ) + + result2 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/custom/route", valid_token=valid_token # This is explicitly listed + ) + + assert result1 is True + assert result2 is True + + +def test_virtual_key_allowed_routes_with_no_member_names_only_explicit(): + """Test that virtual key works when allowed_routes contains only explicit routes (no member names)""" + + # Create a UserAPIKeyAuth with only explicit routes (no LiteLLMRoutes member names) + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["/chat/completions", "/custom/route"], # Only explicit routes + ) + + # Test that explicit routes are allowed + result1 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/chat/completions", valid_token=valid_token + ) + + result2 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/custom/route", valid_token=valid_token + ) + + assert result1 is True + assert result2 is True + + # Test that non-allowed route raises exception + with pytest.raises(Exception) as exc_info: + RouteChecks.is_virtual_key_allowed_to_call_route( + route="/user/info", valid_token=valid_token # Not in allowed routes + ) + + assert "Virtual key is not allowed to call this route" in str(exc_info.value) diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx index e54b6ba1a2..d44b1b0ba5 100644 --- a/ui/litellm-dashboard/src/components/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/create_key_button.tsx @@ -1,6 +1,6 @@ "use client"; import React, { useState, useEffect, useRef, useCallback } from "react"; -import { Button, TextInput, Grid, Col } from "@tremor/react"; +import { Button, TextInput, Grid, Col, Select as TremorSelect, SelectItem } from "@tremor/react"; import { Card, Metric, @@ -182,12 +182,14 @@ const CreateKey: React.FC = ({ const [mcpAccessGroups, setMcpAccessGroups] = useState([]); const [mcpAccessGroupsLoaded, setMcpAccessGroupsLoaded] = useState(false); const [disabledCallbacks, setDisabledCallbacks] = useState([]); + const [keyType, setKeyType] = useState("default"); const handleOk = () => { setIsModalVisible(false); form.resetFields(); setLoggingSettings([]); setDisabledCallbacks([]); + setKeyType("default"); }; const handleCancel = () => { @@ -197,6 +199,7 @@ const CreateKey: React.FC = ({ form.resetFields(); setLoggingSettings([]); setDisabledCallbacks([]); + setKeyType("default"); }; useEffect(() => { @@ -591,31 +594,40 @@ const CreateKey: React.FC = ({ > - - - Models{' '} - - - - - } - name="models" - rules={[{ required: true, message: "Please select a model" }]} - help="required" - className="mt-4" - > - { + if (values.includes("all-team-models")) { + form.setFieldsValue({ models: ["all-team-models"] }); + } + }} + > @@ -626,6 +638,69 @@ const CreateKey: React.FC = ({ ))} + + + + Key Type{' '} + + + + + } + name="key_type" + initialValue="default" + className="mt-4" + > + + + )}