Added ref login

backend/app/core/auth.py

@@ -6,7 +6,7 @@ from fastapi.security import OAuth2PasswordBearer
 import jwt
 from jwt.exceptions import PyJWTError
 
-from .config import SECRET_KEY, ALGORITHM, ADMIN_USER
+from .config import SECRET_KEY, ALGORITHM
 
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="auth/token")
 oauth2_scheme_optional = OAuth2PasswordBearer(tokenUrl="auth/token", auto_error=False)
@@ -20,43 +20,47 @@ def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
         expire = datetime.now(timezone.utc) + timedelta(minutes=15)
 
     to_encode.update({"exp": expire})
-
     encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
     return encoded_jwt
 
 
-async def get_current_user(token: str = Depends(oauth2_scheme)):
+async def get_authenticated_user(token: str = Depends(oauth2_scheme)):
+    """Allows both Admins and Refs"""
     credentials_exception = HTTPException(
         status_code=status.HTTP_401_UNAUTHORIZED,
         detail="Could not validate credentials",
         headers={"WWW-Authenticate": "Bearer"},
     )
-
     try:
         payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
-        username = payload.get("sub")
-
-        if username is None or username != ADMIN_USER:
+        role = payload.get("role")
+        if role not in ["admin", "ref"]:
             raise credentials_exception
-
+        return payload
     except PyJWTError:
         raise credentials_exception
 
-    return username
+
+async def get_admin_user(token: str = Depends(oauth2_scheme)):
+    """Strictly allows ONLY Admins"""
+    user = await get_authenticated_user(token)
+    if user.get("role") != "admin":
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN, detail="Admin privileges required"
+        )
+    return user
 
 
 async def get_optional_user(
     token: Optional[str] = Depends(oauth2_scheme_optional),
-) -> Optional[str]:
+) -> Optional[dict]:
+    """Returns user payload if valid token exists, else None"""
     if not token:
         return None
-
     try:
         payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
-        username = payload.get("sub")
-        if username == ADMIN_USER:
-            return username
+        if payload.get("role") in ["admin", "ref"]:
+            return payload
     except Exception:
         pass
-
     return None