rough role based auth decorator

Author: mmiqball

backend/app/middleware/auth.py

@@ -4,6 +4,9 @@
 import firebase_admin.auth
 from ..services.implementations.user_service import UserService
 from ..utilities.db_utils import get_db
+from ..schemas.user import UserRole
+from functools import wraps
+from typing import Set
 
 security = HTTPBearer()
 logger = logging.getLogger(__name__)
@@ -13,7 +16,7 @@ async def get_current_user(
     db = Depends(get_db)
 ):
     """
-    Validates the authorization token and returns the current user
+    Validates the authorization token and returns the current user and token
     """
     try:
         # Remove 'Bearer ' prefix
@@ -34,18 +37,46 @@ async def get_current_user(
                 status_code=401,
                 detail="User not found in database"
             )
-        logger.info("workedd boss")
-        return user
+
+        return {"user": user, "token": token}
 
     except firebase_admin.auth.InvalidIdTokenError as e:
         logger.error(f"Invalid token: {str(e)}")
         raise HTTPException(
-            status_code=401,
+            status_code=401, 
             detail=f"Invalid token: {str(e)}"
         )
     except Exception as e:
         logger.error(f"Authentication error: {str(e)}")
         raise HTTPException(
             status_code=401,
             detail=str(e)
-        )
+        )
+
+def require_roles(allowed_roles: Set[str]):
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args, current_user=Depends(get_current_user), **kwargs):
+            try:
+                # Get user role using the token from current_user
+                user_service = UserService(kwargs.get('db'))
+                user_role = user_service.get_user_role_by_auth_id(
+                    firebase_admin.auth.verify_id_token(current_user["token"])["uid"]
+                )
+                
+                # Check if user's role is allowed
+                if user_role not in allowed_roles:
+                    raise HTTPException(
+                        status_code=403,
+                        detail=f"Access denied: role '{user_role}' not authorized"
+                    )
+                
+                return await func(*args, current_user=current_user, **kwargs)
+                
+            except Exception as e:
+                raise HTTPException(
+                    status_code=403,
+                    detail=str(e)
+                )
+        return wrapper
+    return decorator

backend/app/middleware/role_auth.py

@@ -0,0 +1,29 @@
+from functools import wraps
+from typing import Set
+from fastapi import Depends, HTTPException
+from ..schemas.user import UserRole
+from ..services.implementations.auth_service import AuthService
+from ..utilities.db_utils import get_db
+from .auth import get_current_user
+
+def get_auth_service(db = Depends(get_db)):
+    return AuthService(logger=logging.getLogger(__name__), user_service=UserService(db))
+
+def require_role_authorization(roles: Set[str]):
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(
+            *args,
+            current_user_data = Depends(get_current_user),
+            auth_service: AuthService = Depends(get_auth_service),
+            **kwargs
+        ):
+            if not auth_service.is_authorized_by_role(current_user_data["token"], roles):
+                raise HTTPException(
+                    status_code=403,
+                    detail=f"User does not have required role(s): {', '.join(roles)}"
+                )
+            kwargs["current_user"] = current_user_data["user"]
+            return await func(*args, **kwargs)
+        return wrapper
+    return decorator