Update app.py

app.py

@@ -405,28 +405,33 @@ def register_user(user_data: UserCreate, db: Session):
 
 @app.get("/protected", response_class=HTMLResponse)
 async def get_protected(
-    request: Request, 
+    request: Request,
     db: Session = Depends(get_db),
-    token: Optional[str] = None  # token is Optional because it may come from the cookie
+    authorization: Optional[str] = Header(None)  # token from Authorization header
 ):
-    print(token)
-    # Try to get the token from the query parameter first, then fall back to the cookie
-    access_token = token or request.cookies.get("access_token")
-    print(access_token)
-    if not access_token:
-        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Not authenticated xxx")
-    
-    # Here verify_token is used directly in the endpoint
-    # If the token is invalid, verify_token will raise an HTTPException and the following lines will not be executed
-    user_email = verify_token(access_token)  # Assuming that verify_token returns the user's email if the token is valid
-    
-    # Get the user from the database
-    db_user = get_user_by_email(db, user_email)
-    if db_user is None or not db_user.is_verified:
-        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="User not found or not verified in the database")
+    # Try to get the token from the Authorization header
+    if authorization:
+        scheme, _, token = authorization.partition(' ')
+        if scheme.lower() != 'bearer':
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid authentication scheme")
+    else:
+        # Fall back to the cookie
+        token = request.cookies.get("access_token")
+
+    if not token:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Not authenticated")
 
-    # Render a template response
-    return templates.TemplateResponse("protected.html", {"request": request, "user": db_user.username})
+    # Verify the token and get user
+    try:
+        user_email = verify_token(token)  # Implement your token verification logic
+        db_user = get_user_by_email(db, user_email)
+        if db_user is None or not db_user.is_verified:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="User not found or not verified in the database")
+
+        # Render a template response
+        return templates.TemplateResponse("protected.html", {"request": request, "user": db_user.username})
+    except Exception as e:  # Replace with specific exceptions as per your verification logic
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=str(e))
     
 def verify_email(verification_token: str, db: Session = Depends(get_db)):
     # Verify the email using the token