added method to get claims

Author: NiveditJain

api-server/app/auth/controllers/create_token.py

@@ -5,6 +5,7 @@
 
 from ..models.token_request import TokenRequest
 from ..models.token_response import TokenResponse
+from ..models.token_claims import TokenClaims
 
 from app.singletons.logs_manager import LogsManager
 
@@ -38,17 +39,17 @@ async def create_token(request: TokenRequest, x_exosphere_request_id: str) -> To
 
         logger.info("User is a super admin", x_exosphere_request_id=x_exosphere_request_id)
 
-        token_claims = {
-            "user_id": str(user.id),
-            "user_name": user.name,
-            "user_type": user.type,
-            "verification_status": user.verification_status,
-            "status": user.status,
-            "exp": datetime.now() + timedelta(seconds=JWT_EXPIRES_IN)
-        }
+        token_claims = TokenClaims(
+            user_id=str(user.id),
+            user_name=user.name,
+            user_type=user.type,
+            verification_status=user.verification_status,
+            status=user.status,
+            exp=int((datetime.now() + timedelta(seconds=JWT_EXPIRES_IN)).timestamp())
+        )
 
         return TokenResponse(
-            access_token=jwt.encode(token_claims, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)
+            access_token=jwt.encode(token_claims.model_dump(), JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)
         )
 
     except Exception as e:

api-server/app/auth/middlewares/get_token_claims.py

@@ -0,0 +1,43 @@
+import jwt
+import os
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from app.singletons.logs_manager import LogsManager
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from ..models.token_claims import TokenClaims
+
+logger = LogsManager().get_logger()
+
+JWT_SECRET_KEY = os.getenv("JWT_SECRET_KEY")
+if not JWT_SECRET_KEY:
+    raise ValueError("JWT_SECRET_KEY environment variable is not set or is empty.")
+JWT_ALGORITHM = "HS256"
+
+
+class GetTokenClaimsMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+
+        token = request.headers.get("Authorization")
+
+        if token:
+            if not token.startswith("Bearer "):
+                logger.error("Invalid token format", x_exosphere_request_id=getattr(request.state, 'x_exosphere_request_id', None))
+                return JSONResponse(status_code=401, content={"message": "Invalid token format", "success": False})
+            
+            try:
+                token_claims = jwt.decode(token.split(" ")[1], JWT_SECRET_KEY, algorithms=[JWT_ALGORITHM])
+                request.state.token_claims = TokenClaims(**token_claims)
+                logger.info("Token claims decoded", x_exosphere_request_id=getattr(request.state, 'x_exosphere_request_id', None), user_id=request.state.token_claims.user_id)
+
+            except Exception as e:
+                logger.error("Error decoding token", error=e, x_exosphere_request_id=getattr(request.state, 'x_exosphere_request_id', None))
+
+                return JSONResponse(status_code=401, content={"message": "Invalid token", "success": False})
+                
+        else:
+            logger.error("No token provided", x_exosphere_request_id=getattr(request.state, 'x_exosphere_request_id', None))
+            request.state.token_claims = None
+
+        return await call_next(request)

api-server/app/auth/models/token_claims.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+class TokenClaims(BaseModel):
+    user_id: str
+    user_name: str
+    user_type: str
+    verification_status: str
+    status: str
+    exp: int

api-server/app/main.py

@@ -16,6 +16,7 @@
     UnhandledExceptionsMiddleware,
 )
 from .middlewares.request_id_middleware import RequestIdMiddleware
+from .auth.middlewares.get_token_claims import GetTokenClaimsMiddleware
 
 # injecting databases
 from .user.models.user_database_model import User
@@ -49,13 +50,16 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 
+
 # this middleware should be the first one
+app.add_middleware(GetTokenClaimsMiddleware)
+
 app.add_middleware(RequestIdMiddleware)
 
-# this middleware should be the last one
 app.add_middleware(UnhandledExceptionsMiddleware)
 
 
+
 @app.get("/health-check")
 def health() -> dict:
     return {"message": "OK"}