Refactor Security and User Role Management

- Organized imports in private_ai_keys.py and users.py for better readability.
- Updated role validation logic to utilize UserRole.get_all_roles() for consistency across user role checks.
- Removed unused role definitions and hierarchy from security.py, streamlining the codebase.
- Enhanced role management by centralizing role-related functionality in the roles module.

Author: PhiRho

app/api/private_ai_keys.py

@@ -13,9 +13,22 @@
 from app.db.postgres import PostgresManager
 from app.db.models import DBPrivateAIKey, DBRegion, DBUser, DBTeam
 from app.services.litellm import LiteLLMService
-from app.core.security import get_current_user_from_auth, get_role_min_key_creator, get_role_min_team_admin, get_private_ai_access, UserRole, check_system_admin
+from app.core.security import (
+    get_current_user_from_auth,
+    get_role_min_team_admin,
+    get_private_ai_access,
+    check_system_admin
+)
+from app.core.roles import UserRole
 from app.core.config import settings
-from app.core.resource_limits import check_key_limits, check_vector_db_limits, get_token_restrictions, DEFAULT_KEY_DURATION, DEFAULT_MAX_SPEND, DEFAULT_RPM_PER_KEY
+from app.core.resource_limits import (
+    check_key_limits,
+    check_vector_db_limits,
+    get_token_restrictions,
+    DEFAULT_KEY_DURATION,
+    DEFAULT_MAX_SPEND,
+    DEFAULT_RPM_PER_KEY
+)
 
 router = APIRouter(
     tags=["private-ai-keys"]

app/api/users.py

@@ -1,12 +1,13 @@
 from fastapi import APIRouter, Depends, HTTPException, status
 from sqlalchemy.orm import Session
-from typing import List, get_args
+from typing import List
 from app.core.config import settings
 from app.core.resource_limits import check_team_user_limit
 from app.db.database import get_db
 from app.schemas.models import User, UserUpdate, UserCreate, TeamOperation, UserRoleUpdate
 from app.db.models import DBUser, DBTeam
-from app.core.security import get_password_hash, check_system_admin, get_current_user_from_auth, UserRole, get_role_min_team_admin
+from app.core.security import get_password_hash, check_system_admin, get_current_user_from_auth, get_role_min_team_admin
+from app.core.roles import UserRole
 from datetime import datetime, UTC
 
 router = APIRouter(
@@ -77,10 +78,10 @@ async def create_user(
         check_team_user_limit(db, user.team_id)
 
     # Validate role if provided
-    if user.role and user.role not in get_args(UserRole):
+    if user.role and user.role not in UserRole.get_all_roles():
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
-            detail=f"Invalid role. Must be one of: {', '.join(get_args(UserRole))}"
+            detail=f"Invalid role. Must be one of: {', '.join(UserRole.get_all_roles())}"
         )
 
     # Default to the lowest permissions for a user in a team
@@ -256,10 +257,10 @@ async def update_user_role(
     Update a user's role. Accessible by admin users or team admins for their team members.
     """
     # Validate role
-    if role_update.role not in get_args(UserRole):
+    if role_update.role not in UserRole.get_all_roles():
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
-            detail=f"Invalid role. Must be one of: {', '.join(get_args(UserRole))}"
+            detail=f"Invalid role. Must be one of: {', '.join(UserRole.get_all_roles())}"
         )
 
     # Get the user to update

app/core/security.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timedelta, UTC
-from typing import Optional, Literal, Dict
+from typing import Optional
 from jose import JWTError, jwt
 from passlib.context import CryptContext
 from fastapi import Depends, HTTPException, status, Cookie, Header, Request
@@ -9,14 +9,11 @@
 from app.db.database import get_db
 from sqlalchemy.orm import Session
 from app.db.models import DBUser, DBAPIToken
-from app.core.roles import UserRole as RoleClass
 from app.core.rbac import (
     require_system_admin,
     require_team_admin,
     require_key_creator_or_higher,
-    require_read_only_or_higher,
     require_sales_or_higher,
-    require_any_role
 )
 
 logger = logging.getLogger(__name__)
@@ -27,19 +24,6 @@
 # Custom bearer scheme
 bearer_scheme = HTTPBearer(auto_error=False)
 
-# Define valid user roles as a Literal type - MUST match existing values exactly
-UserRole = Literal["admin", "key_creator", "read_only", "user", "system_admin", "sales"]
-
-# Define a hierarchy for roles - updated to include new roles
-user_role_hierarchy: Dict[UserRole, int] = {
-    "system_admin": 0,
-    "admin": 1,
-    "user": 2,
-    "key_creator": 3,
-    "read_only": 4,
-    "sales": 5,
-}
-
 def verify_password(plain_password: str, hashed_password: str) -> bool:
     """Verify a password against its hash."""
     return pwd_context.verify(plain_password, hashed_password)
@@ -152,16 +136,6 @@ async def check_system_admin(current_user: DBUser = Depends(get_current_user_fro
     dependency = require_system_admin()
     return dependency.check_access(current_user)
 
-def get_user_role(minimum_role: UserRole, current_user: DBUser):
-    if current_user.is_admin:
-        return "system_admin"
-    elif user_role_hierarchy[current_user.role] > user_role_hierarchy[minimum_role]:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail="Not authorized to perform this action"
-        )
-    return current_user.role
-
 async def get_role_min_team_admin(current_user: DBUser = Depends(get_current_user_from_auth)):
     """Require team admin role or higher."""
     dependency = require_team_admin()