[Authorization] Check user role before operations (QuivrHQ#588)

* feat(security): add RBAC on /explore/*

* feat(security): add RBAC on /brains/*

backend/models/brains.py

@@ -62,6 +62,19 @@ def get_user_brains(self, user_id):
         )
         return [item["brains"] for item in response.data]
 
+    def get_brain_for_user(self, user_id):
+        response = (
+            self.commons["supabase"]
+            .from_("brains_users")
+            .select("id:brain_id, rights, brains (id: brain_id, name)")
+            .filter("user_id", "eq", user_id)
+            .filter("brain_id", "eq", self.id)
+            .execute()
+        )
+        if len(response.data) == 0:
+            return None
+        return response.data[0]
+
     def get_brain_details(self):
         response = (
             self.commons["supabase"]

backend/routes/authorizations/brain_authorization.py

@@ -0,0 +1,60 @@
+from functools import wraps
+from typing import Optional
+from uuid import UUID
+
+from fastapi import HTTPException, status
+from models.brains import Brain
+from models.users import User
+
+
+def has_brain_authorization(required_role: str = "Owner"):
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(current_user: User, *args, **kwargs):
+            brain_id: Optional[UUID] = kwargs.get("brain_id")
+            user_id = current_user.id
+
+            if brain_id is None:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail="Missing brain ID",
+                )
+
+            validate_brain_authorization(
+                brain_id, user_id=user_id, required_role=required_role
+            )
+
+            return await func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def validate_brain_authorization(
+    brain_id: UUID,
+    user_id: UUID,
+    required_role: Optional[str] = "Owner",
+):
+    if required_role is None:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Missing required role",
+        )
+
+    brain = Brain(id=brain_id)
+    user_brain = brain.get_brain_for_user(user_id)
+
+    if user_brain is None:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="You don't have permission for this brain",
+        )
+
+    # TODO: Update this logic when we have more roles
+    # Eg: Owner > Admin > User ... this should be taken into account
+    if user_brain.get("rights") != required_role:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="You don't have the required role for this brain",
+        )

backend/routes/brain_routes.py

@@ -9,6 +9,8 @@
 from models.users import User
 from pydantic import BaseModel
 
+from routes.authorizations.brain_authorization import has_brain_authorization
+
 logger = get_logger(__name__)
 
 brain_router = APIRouter()
@@ -73,9 +75,16 @@ async def get_default_brain_endpoint(current_user: User = Depends(get_current_us
 
 # get one brain
 @brain_router.get(
-    "/brains/{brain_id}/", dependencies=[Depends(AuthBearer())], tags=["Brain"]
+    "/brains/{brain_id}/",
+    dependencies=[
+        Depends(AuthBearer()),
+        Depends(has_brain_authorization),
+    ],
+    tags=["Brain"],
 )
-async def get_brain_endpoint(brain_id: UUID):
+async def get_brain_endpoint(
+    brain_id: UUID,
+):
     """
     Retrieve details of a specific brain by brain ID.
 
@@ -99,7 +108,12 @@ async def get_brain_endpoint(brain_id: UUID):
 
 # delete one brain
 @brain_router.delete(
-    "/brains/{brain_id}/", dependencies=[Depends(AuthBearer())], tags=["Brain"]
+    "/brains/{brain_id}/",
+    dependencies=[
+        Depends(AuthBearer()),
+        Depends(has_brain_authorization),
+    ],
+    tags=["Brain"],
 )
 async def delete_brain_endpoint(
     brain_id: UUID,
@@ -167,11 +181,19 @@ async def create_brain_endpoint(
 
 # update existing brain
 @brain_router.put(
-    "/brains/{brain_id}/", dependencies=[Depends(AuthBearer())], tags=["Brain"]
+    "/brains/{brain_id}/",
+    dependencies=[
+        Depends(
+            AuthBearer(),
+        ),
+        Depends(has_brain_authorization),
+    ],
+    tags=["Brain"],
 )
 async def update_brain_endpoint(
     brain_id: UUID,
     input_brain: Brain,
+    current_user: User = Depends(get_current_user),
 ):
     """
     Update an existing brain with new brain parameters/files.

backend/routes/explore_routes.py

@@ -6,13 +6,17 @@
 from models.settings import common_dependencies
 from models.users import User
 
+from routes.authorizations.brain_authorization import (
+    has_brain_authorization,
+    validate_brain_authorization,
+)
+
 explore_router = APIRouter()
 
 
 @explore_router.get("/explore/", dependencies=[Depends(AuthBearer())], tags=["Explore"])
 async def explore_endpoint(
     brain_id: UUID = Query(..., description="The ID of the brain"),
-    current_user: User = Depends(get_current_user),
 ):
     """
     Retrieve and explore unique user data vectors.
@@ -25,7 +29,12 @@ async def explore_endpoint(
 
 
 @explore_router.delete(
-    "/explore/{file_name}/", dependencies=[Depends(AuthBearer())], tags=["Explore"]
+    "/explore/{file_name}/",
+    dependencies=[
+        Depends(AuthBearer()),
+        Depends(has_brain_authorization),
+    ],
+    tags=["Explore"],
 )
 async def delete_endpoint(
     file_name: str,
@@ -66,4 +75,16 @@ async def download_endpoint(
         .execute()
     )
     documents = response.data
+
+    if len(documents) == 0:
+        return {"documents": []}
+
+    related_brain_id = UUID(documents[0]["brain_id"])
+    if related_brain_id is None:
+        raise Exception(
+            f"File {file_name} has no brain_id associated with it. Please contact support."
+        )
+
+    validate_brain_authorization(brain_id=related_brain_id, user_id=current_user.id)
+
     return {"documents": documents}