First pass at setting up an RBAC system (still need to actually validate permissions on org and role endpoints)

Author: chriscarrollsmith

README.md

@@ -18,12 +18,25 @@ Set your desired database name, username, and password in the .env file.
 
 `docker compose up -d`
 
+## Create database tables and default permissions/roles
+
+`poetry run python migrations/set_up_db.py --drop`
+
 ## Run the development server
 
+Make sure the development database is running and tables and default permissions/roles are created first.
+
 `uvicorn main:app --host 0.0.0.0 --port 8000 --reload`
 
 Navigate to http://localhost:8000/
 
+## To do
+
+- Implement password recovery
+- Implement role/org system
+- Implement user profile page
+- Add payments/billing system?
+
 ## License
 
 This project is licensed under the GPLv3 License. See the LICENSE file for more details.

main.py

@@ -5,24 +5,20 @@
 from fastapi.responses import RedirectResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
-from sqlmodel import SQLModel, create_engine
 from fastapi.exceptions import RequestValidationError, StarletteHTTPException
 from routers import auth, organization, role, user
 from utils.auth import get_authenticated_user, get_optional_user, NeedsNewTokens
-from utils.db import User, get_connection_url
+from utils.db import User
 
 
 logger = logging.getLogger("uvicorn.error")
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Startup logic
-    engine = create_engine(get_connection_url())
-    SQLModel.metadata.create_all(engine)
-    engine.dispose()
+    # Optional startup logic
     yield
-    # Shutdown logic
+    # Optional shutdown logic
 
 
 app = FastAPI(lifespan=lifespan)

migrations/set_up_db.py

@@ -0,0 +1,30 @@
+import sys
+import os
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+if __name__ == "__main__":
+    import argparse
+    import os
+    from dotenv import load_dotenv
+    from utils.db import set_up_db
+
+    load_dotenv()
+
+    parser = argparse.ArgumentParser(
+        description="Optionally drop and recreate all tables in the database and set up default roles and permissions"
+    )
+    parser.add_argument(
+        "--drop",
+        action="store_true",
+        help="Drop all tables first"
+    )
+
+    args = parser.parse_args()
+
+    set_up_db(args.drop)
+
+    print(
+        f"Set up database {os.getenv('DB_NAME')}"
+    )

routers/__init__.py

@@ -1,3 +1,108 @@
-from fastapi import APIRouter
+from fastapi import APIRouter, Depends, HTTPException
+from pydantic import BaseModel, ConfigDict
+from sqlmodel import Session, select
+from utils.db import Role, RolePermissionLink, ValidPermissions, get_session
+from typing import List
+from utils.db import utc_time
+from datetime import datetime
 
-router = APIRouter(prefix="/role", tags=["role"])
+router = APIRouter(prefix="/roles", tags=["roles"])
+
+
+class RoleCreate(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    name: str
+    permission_ids: List[int]
+
+
+class RoleRead(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    id: int
+    name: str
+    created_at: datetime
+    updated_at: datetime
+    deleted: bool
+    permissions: List[str]
+
+
+class RoleUpdate(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    name: str
+    permission_ids: List[int]
+
+
+@router.post("/", response_model=RoleRead)
+def create_role(role: RoleCreate, session: Session = Depends(get_session)):
+    db_role = session.exec(select(Role).where(Role.name == role.name)).first()
+    if db_role:
+        raise HTTPException(status_code=400, detail="Role already exists")
+    db_role = Role(name=role.name)
+    session.add(db_role)
+    session.commit()
+    session.refresh(db_role)
+
+    for permission_id in role.permission_ids:
+        db_role_permission_link = RolePermissionLink(
+            role_id=db_role.id, permission_id=permission_id)
+        session.add(db_role_permission_link)
+
+    session.commit()
+    session.refresh(db_role)
+    return db_role
+
+
+@router.get("/{role_id}", response_model=RoleRead)
+def read_role(role_id: int, session: Session = Depends(get_session)):
+    db_role = session.get(Role, role_id)
+    if not db_role:
+        raise HTTPException(status_code=404, detail="Role not found")
+    permissions = [
+        link.permission.name for link in db_role.role_permission_links]
+    return RoleRead(
+        id=db_role.id,
+        name=db_role.name,
+        created_at=db_role.created_at,
+        updated_at=db_role.updated_at,
+        deleted=db_role.deleted,
+        permissions=permissions
+    )
+
+
+@router.put("/{role_id}", response_model=RoleRead)
+def update_role(role_id: int, role: RoleUpdate, session: Session = Depends(get_session)):
+    db_role = session.get(Role, role_id)
+    if not db_role:
+        raise HTTPException(status_code=404, detail="Role not found")
+    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()
+
+    # Update RolePermissionLinks
+    session.exec(select(RolePermissionLink).where(
+        RolePermissionLink.role_id == role_id)).delete()
+    for permission_id in role.permission_ids:
+        db_role_permission_link = RolePermissionLink(
+            role_id=db_role.id, permission_id=permission_id)
+        session.add(db_role_permission_link)
+
+    session.commit()
+    session.refresh(db_role)
+    return db_role
+
+
+@router.delete("/{role_id}")
+def delete_role(role_id: int, session: Session = Depends(get_session)):
+    db_role = session.get(Role, role_id)
+    if not db_role:
+        raise HTTPException(status_code=404, detail="Role not found")
+    db_role.deleted = True
+    db_role.updated_at = utc_time()
+    session.add(db_role)
+    session.commit()
+    return {"message": "Role deleted successfully"}

routers/role.py

@@ -1,11 +1,12 @@
 import os
 import logging
+from enum import Enum
 from typing import Optional, List
 from datetime import datetime, UTC
 from dotenv import load_dotenv
 from sqlalchemy.engine import URL
-from sqlmodel import create_engine, Session, SQLModel, Field, Relationship
-
+from sqlmodel import create_engine, Session, SQLModel, Field, Relationship, select
+from sqlalchemy import Column, Enum as SQLAlchemyEnum
 
 load_dotenv()
 
@@ -41,13 +42,71 @@ def get_session():
         yield session
 
 
+def set_up_db(drop: bool = False):
+    engine = create_engine(get_connection_url())
+    if drop:
+        SQLModel.metadata.drop_all(engine)
+    SQLModel.metadata.create_all(engine)
+    with Session(engine) as session:
+        roles_in_db = []
+        # Create default roles
+        for role_name in default_roles:
+            db_role = session.exec(select(Role).where(
+                Role.name == role_name)).first()
+            if not db_role:
+                db_role = Role(name=role_name)
+                session.add(db_role)
+                roles_in_db.append(db_role)
+            else:
+                roles_in_db.append(db_role)
+
+        session.commit()  # Commit after adding roles
+
+        # Create default permissions
+        for permission in [permission for permission in ValidPermissions]:
+            db_permission = session.exec(select(Permission).where(
+                Permission.name == permission)).first()
+            if not db_permission:
+                db_permission = Permission(name=permission)
+                session.add(db_permission)
+
+            # Create RolePermissionLink for Owner and Administrator
+            # Assuming first two roles are Owner and Administrator
+            for role in roles_in_db[:2]:
+                db_role_permission_link = session.exec(select(RolePermissionLink).where(
+                    RolePermissionLink.role_id == role.id,
+                    RolePermissionLink.permission_id == db_permission.id)).first()
+                if not db_role_permission_link:
+                    if not (permission == ValidPermissions.DELETE_ORGANIZATION and role.name == "Administrator"):
+                        role_permission_link = RolePermissionLink(
+                            role_id=role.id, permission_id=db_permission.id)
+                        session.add(role_permission_link)
+
+        session.commit()  # Commit after adding permissions and links
+    engine.dispose()
+
+
 # --- Models ---
 
 
 def utc_time():
     return datetime.now(UTC)
 
 
+default_roles = ["Owner", "Administrator", "Member"]
+
+
+class ValidPermissions(Enum):
+    DELETE_ORGANIZATION = "Delete Organization"
+    EDIT_ORGANIZATION = "Edit Organization"
+    INVITE_USER = "Invite User"
+    REMOVE_USER = "Remove User"
+    EDIT_USER_ROLE = "Edit User Role"
+    CREATE_ROLE = "Create Role"
+    DELETE_ROLE = "Delete Role"
+    EDIT_ROLE = "Edit Role"
+
+
 class Organization(SQLModel, table=True):
     id: Optional[int] = Field(default=None, primary_key=True)
     name: str
@@ -58,6 +117,45 @@ class Organization(SQLModel, table=True):
     users: List["User"] = Relationship(back_populates="organization")
 
 
+class Role(SQLModel, table=True):
+    id: Optional[int] = Field(default=None, primary_key=True)
+    name: str
+    organization_id: Optional[int] = Field(
+        default=None, foreign_key="organization.id")
+    created_at: datetime = Field(default_factory=utc_time)
+    updated_at: datetime = Field(default_factory=utc_time)
+    deleted: bool = Field(default=False)
+
+    users: List["User"] = Relationship(back_populates="role")
+    role_permission_links: List["RolePermissionLink"] = Relationship(
+        back_populates="role")
+
+
+class Permission(SQLModel, table=True):
+    id: Optional[int] = Field(default=None, primary_key=True)
+    name: ValidPermissions = Field(
+        sa_column=Column(SQLAlchemyEnum(ValidPermissions)))
+    created_at: datetime = Field(default_factory=utc_time)
+    updated_at: datetime = Field(default_factory=utc_time)
+    deleted: bool = Field(default=False)
+
+    role_permission_links: List["RolePermissionLink"] = Relationship(
+        back_populates="permission")
+
+
+class RolePermissionLink(SQLModel, table=True):
+    id: Optional[int] = Field(default=None, primary_key=True)
+    role_id: Optional[int] = Field(
+        default=None, foreign_key="role.id")
+    permission_id: Optional[int] = Field(
+        default=None, foreign_key="permission.id")
+
+    role: Optional["Role"] = Relationship(
+        back_populates="role_permission_links")
+    permission: Optional["Permission"] = Relationship(
+        back_populates="role_permission_links")
+
+
 class User(SQLModel, table=True):
     id: Optional[int] = Field(default=None, primary_key=True)
     name: str
@@ -66,9 +164,19 @@ class User(SQLModel, table=True):
     avatar_url: Optional[str] = None
     organization_id: Optional[int] = Field(
         default=None, foreign_key="organization.id")
+    role_id: Optional[int] = Field(default=None, foreign_key="role.id")
     created_at: datetime = Field(default_factory=utc_time)
     updated_at: datetime = Field(default_factory=utc_time)
     deleted: bool = Field(default=False)
 
     organization: Optional["Organization"] = Relationship(
         back_populates="users")
+    role: Optional["Role"] = Relationship(back_populates="users")
+
+
+class UserOrganizationLink(SQLModel, table=True):
+    id: Optional[int] = Field(default=None, primary_key=True)
+    user_id: Optional[int] = Field(
+        default=None, foreign_key="user.id")
+    organization_id: Optional[int] = Field(
+        default=None, foreign_key="organization.id")