Massively imporved/fixed role.py

Author: chriscarrollsmith

routers/organization.py

@@ -4,7 +4,7 @@
 from pydantic import BaseModel, ConfigDict, field_validator
 from sqlmodel import Session, select
 from utils.db import get_session
-from utils.auth import get_authenticated_user, get_user_with_relations
+from utils.auth import get_authenticated_user, get_user_with_relations, InsufficientPermissionsError
 from utils.models import Organization, User, Role, utc_time, default_roles, ValidPermissions
 from datetime import datetime
 
@@ -37,14 +37,6 @@ def __init__(self):
         )
 
 
-class InsufficientPermissionsError(HTTPException):
-    def __init__(self):
-        super().__init__(
-            status_code=403,
-            detail="You don't have permission to perform this action"
-        )
-
-
 router = APIRouter(prefix="/organizations", tags=["organizations"])
 
 
@@ -173,7 +165,11 @@ def delete_organization(
     # Check if user has permission to delete organization
     organization: Organization | None = next(
         (org for org in user.organizations if org.id == org_id), None)
-    if not organization or not any(role.permissions.DELETE_ORGANIZATION for role in organization.roles):
+    if not organization or not any(
+        p.name == ValidPermissions.DELETE_ORGANIZATION
+        for role in organization.roles
+        for p in role.permissions
+    ):
         raise InsufficientPermissionsError()
 
     # Delete organization

routers/role.py

@@ -1,13 +1,15 @@
-from typing import List
-from datetime import datetime
+# TODO: User with permission to create/edit roles can only assign permissions
+# they themselves have.
+from typing import List, Sequence, Optional
 from logging import getLogger
 from fastapi import APIRouter, Depends, Form, HTTPException
 from fastapi.responses import RedirectResponse
 from pydantic import BaseModel, ConfigDict, field_validator
-from sqlmodel import Session, select
+from sqlmodel import Session, select, col
+from sqlalchemy.orm import selectinload
 from utils.db import get_session
-from utils.auth import get_authenticated_user
-from utils.models import Role, RolePermissionLink, ValidPermissions, utc_time, User
+from utils.auth import get_authenticated_user, InsufficientPermissionsError
+from utils.models import Role, Permission, ValidPermissions, utc_time, User, DataIntegrityError
 
 logger = getLogger("uvicorn.error")
 
@@ -17,6 +19,16 @@
 # -- Custom Exceptions --
 
 
+class InvalidPermissionError(HTTPException):
+    """Raised when a user attempts to assign an invalid permission to a role"""
+
+    def __init__(self, permission: ValidPermissions):
+        super().__init__(
+            status_code=400,
+            detail=f"Invalid permission: {permission}"
+        )
+
+
 class RoleAlreadyExistsError(HTTPException):
     """Raised when attempting to create a role with a name that already exists"""
 
@@ -31,53 +43,46 @@ def __init__(self):
         super().__init__(status_code=404, detail="Role not found")
 
 
+class RoleHasUsersError(HTTPException):
+    """Raised when a requested role to be deleted has users"""
+
+    def __init__(self):
+        super().__init__(
+            status_code=400,
+            detail="Role cannot be deleted until users with that role are reassigned"
+        )
+
+
 # -- Server Request Models --
 
 class RoleCreate(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
     name: str
+    organization_id: int
     permissions: List[ValidPermissions]
 
-    @field_validator("name")
-    @classmethod
-    def validate_unique_name(cls, name: str, info):
-        # Note: This requires passing session as a dependency to as_form
-        session = info.context.get("session")
-        if session and session.exec(select(Role).where(Role.name == name)).first():
-            raise RoleAlreadyExistsError()
-        return name
-
     @classmethod
     async def as_form(
         cls,
         name: str = Form(...),
-        permissions: List[ValidPermissions] = Form(...),
-        session: Session = Depends(get_session)
+        organization_id: int = Form(...),
+        permissions: List[ValidPermissions] = Form(...)
     ):
         # Pass session to validator context
         return cls(
             name=name,
-            permissions=permissions,
-            context={"session": session}
+            organization_id=organization_id,
+            permissions=permissions
         )
 
 
-class RoleRead(BaseModel):
-    model_config = ConfigDict(from_attributes=True)
-
-    id: int
-    name: str
-    created_at: datetime
-    updated_at: datetime
-    permissions: List[ValidPermissions]
-
-
 class RoleUpdate(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
     id: int
     name: str
+    organization_id: int
     permissions: List[ValidPermissions]
 
     @field_validator("id")
@@ -95,17 +100,32 @@ async def as_form(
         cls,
         id: int = Form(...),
         name: str = Form(...),
-        permissions: List[ValidPermissions] = Form(...),
-        session: Session = Depends(get_session)
+        organization_id: int = Form(...),
+        permissions: List[ValidPermissions] = Form(...)
     ):
         return cls(
             id=id,
             name=name,
-            permissions=permissions,
-            context={"session": session}
+            organization_id=organization_id,
+            permissions=permissions
         )
 
 
+class RoleDelete(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    id: int
+    organization_id: int
+
+    @classmethod
+    async def as_form(
+        cls,
+        id: int = Form(...),
+        organization_id: int = Form(...)
+    ):
+        return cls(id=id, organization_id=organization_id)
+
+
 # -- Routes --
 
 
@@ -115,19 +135,34 @@ def create_role(
     user: User = Depends(get_authenticated_user),
     session: Session = Depends(get_session)
 ) -> RedirectResponse:
-    # Create role and permissions in a single transaction
+    # Check that the user-selected role name is unique for the organization
+    if session.exec(
+        select(Role).where(
+            Role.name == role.name,
+            Role.organization_id == role.organization_id
+        )
+    ).first():
+        raise RoleAlreadyExistsError()
+
+    # Check that the user is authorized to create roles in the organization
+    if not user.has_permission(ValidPermissions.CREATE_ROLE, role.organization_id):
+        raise InsufficientPermissionsError()
+
+    # Create role
     db_role = Role(
         name=role.name,
-        organization_id=user.organization_id  # Add organization ID to role
+        organization_id=role.organization_id
     )
+    session.add(db_role)
 
-    # Create RolePermissionLink objects and associate them with the role
-    db_role.permissions = [
-        RolePermissionLink(permission_id=permission.name)
-        for permission in role.permissions
-    ]
+    # Select Permission records corresponding to the user-selected permissions
+    # and associate them with the newly created role
+    permissions: Sequence[Permission] = session.exec(
+        select(Permission).where(col(Permission.name).in_(role.permissions))
+    ).all()
+    db_role.permissions.extend(permissions)
 
-    session.add(db_role)
+    # Commit transaction
     session.commit()
 
     return RedirectResponse(url="/profile", status_code=303)
@@ -139,39 +174,85 @@ def update_role(
     user: User = Depends(get_authenticated_user),
     session: Session = Depends(get_session)
 ) -> RedirectResponse:
-    db_role: Role | None = session.get(Role, role.id)
-    role_data = role.model_dump(exclude_unset=True)
-    for key, value in role_data.items():
-        setattr(db_role, key, value)
-    db_role.updated_at = utc_time()
-    session.add(db_role)
-    session.commit()
+    # Check that the user is authorized to update the role
+    if not user.has_permission(ValidPermissions.EDIT_ROLE, role.organization_id):
+        raise InsufficientPermissionsError()
+
+    # Select db_role to update, along with its permissions, by ID
+    db_role: Optional[Role] = session.exec(
+        select(Role).where(Role.id == role.id).options(
+            selectinload(Role.permissions))
+    ).first()
+
+    if not db_role:
+        raise RoleNotFoundError()
 
-    # Correctly delete RolePermissionLinks for the role
-    session.delete(RolePermissionLink.role_id == role.id)
+    # If any user-selected permissions are not valid, raise an error
+    for permission in role.permissions:
+        if permission not in ValidPermissions:
+            raise InvalidPermissionError(permission)
 
+    # Add any user-selected permissions that are not already associated with the role
     for permission in role.permissions:
-        db_role_permission_link = RolePermissionLink(
-            role_id=db_role.id,
-            permission_id=permission.name
+        if permission not in [p.name for p in db_role.permissions]:
+            db_permission: Optional[Permission] = session.exec(
+                select(Permission).where(Permission.name == permission)
+            ).first()
+            if db_permission:
+                db_role.permissions.append(db_permission)
+            else:
+                raise DataIntegrityError(resource=f"Permission: {permission}")
+
+    # Remove any permissions that are not user-selected
+    for db_permission in db_role.permissions:
+        if db_permission.name not in role.permissions:
+            db_role.permissions.remove(db_permission)
+
+    # Check that no existing organization role has the same name but a different ID
+    if session.exec(
+        select(Role).where(
+            Role.name == role.name,
+            Role.organization_id == role.organization_id,
+            Role.id != role.id
         )
-        session.add(db_role_permission_link)
+    ).first():
+        raise RoleAlreadyExistsError()
+
+    # Update role name and updated_at timestamp
+    db_role.name = role.name
+    db_role.updated_at = utc_time()
 
     session.commit()
     session.refresh(db_role)
     return RedirectResponse(url="/profile", status_code=303)
 
 
-# TODO: Reject role deletion if anyone in the organization has that role
 @router.delete("/{role_id}", response_class=RedirectResponse)
 def delete_role(
-    role_id: int,
+    role: RoleDelete = Depends(RoleDelete.as_form),
     user: User = Depends(get_authenticated_user),
     session: Session = Depends(get_session)
 ) -> RedirectResponse:
-    db_role = session.get(Role, role_id)
+    # Check that the user is authorized to delete the role
+    if not user.has_permission(ValidPermissions.DELETE_ROLE, role.organization_id):
+        raise InsufficientPermissionsError()
+
+    # Select the role to delete by ID, along with its users
+    db_role: Role | None = session.exec(
+        select(Role).where(Role.id == role.id).options(
+            selectinload(Role.users)
+        )
+    ).first()
+
     if not db_role:
-        raise HTTPException(status_code=404, detail="Role not found")
+        raise RoleNotFoundError()
+
+    # Check that no users have the role
+    if db_role.users:
+        raise RoleHasUsersError()
+
+    # Delete the role
     session.delete(db_role)
     session.commit()
+
     return RedirectResponse(url="/profile", status_code=303)