Update main.py

main.py

@@ -1,11 +1,11 @@
 #main.py
 
-from fastapi import FastAPI, Form, Depends, HTTPException, status
+from fastapi import FastAPI, Form, Depends, HTTPException, status, Query
 from fastapi.requests import Request
 from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse
 from fastapi.templating import Jinja2Templates
 from sqlalchemy.orm import Session
-from auth import verify_token, oauth2_scheme, auth_views, register, UserCreate, authenticate_user, get_user_by_verification_token
+from auth import verify_token, oauth2_scheme, auth_views, register, UserCreate, authenticate_user, get_user_by_verification_token, protected_route
 from database import get_db, get_user_by_email
 from datetime import timedelta
 from typing import Optional
@@ -132,23 +132,13 @@ async def verify_email(verification_token: str, db: Session = Depends(get_db)):
     # Redirect to the protected route with the token as a query parameter (or as required by your front-end/client)
     return RedirectResponse(url=f"/protected?token={access_token}")
 
-@app.get("/protected", response_class=HTMLResponse)
-async def protected_route(request: Request, token: Optional[str] = None, db: Session = Depends(get_db)):
-    # Try to get token from query parameter first
-    token = token or request.cookies.get("access_token")
-    if not token:
-        raise HTTPException(status_code=401, detail="Not authenticated")
 
-    try:
-        payload = jwt.decode(token, auth_views.SECRET_KEY, algorithms=[auth_views.ALGORITHM])
-        user_email = payload.get("sub")
-        if user_email is None:
-            raise HTTPException(status_code=401, detail="Not authenticated")
-    except jwt.PyJWTError:
-        raise HTTPException(status_code=401, detail="Not authenticated")
-
-    db_user = get_user_by_email(db, user_email)
-    if db_user is None or not db_user.is_verified:
-        raise HTTPException(status_code=401, detail="User not found or not verified in the database")
-
-    return templates.TemplateResponse("protected.html", {"request": request, "user": db_user})
+
+@app.get("/protected", response_class=HTMLResponse)
+async def get_protected(
+    request: Request,
+    token: str = Query(None),  # Accept token from query parameters
+    db: Session = Depends(get_db)
+):
+    # Now pass both the request and token to the protected_route function
+    return await protected_route(request, token, db)