Added cognito authentication - Fix 9

app/main.py

@@ -1,6 +1,6 @@
-from fastapi import FastAPI, Depends, HTTPException, Request
-from fastapi.responses import RedirectResponse
-from app.cabletv.utils.auth import get_current_user, exchange_code_for_token
+from fastapi import FastAPI, Depends, HTTPException
+from fastapi.responses import JSONResponse, RedirectResponse
+from app.cabletv.utils.auth import exchange_code_for_token, get_current_user, DOMAIN, CLIENT_ID
 
 app = FastAPI()
 
@@ -9,22 +9,26 @@ async def root():
     return {"message": "IPTV Updater API"}
 
 @app.get("/protected")
-async def protected_route(request: Request, user = Depends(get_current_user)):
+async def protected_route(user = Depends(get_current_user)):
+    if isinstance(user, RedirectResponse):
+        return user
     return {"message": "Protected content", "user": user['Username']}
 
 @app.get("/auth/callback")
-async def auth_callback(request: Request, code: str):
+async def auth_callback(code: str):
     try:
-        redirect_uri = str(request.base_url)
-        tokens = exchange_code_for_token(code, redirect_uri)
+        tokens = exchange_code_for_token(code)
         
-        # Create redirect response to protected route
-        response = RedirectResponse(url="/protected", status_code=302)
+        # Use id_token
+        response = JSONResponse(content={
+            "message": "Authentication successful",
+            "id_token": tokens["id_token"]  # Changed from access_token
+        })
         
-        # Set token cookie
+        # Store id_token in cookie
         response.set_cookie(
             key="token",
-            value=tokens["id_token"],
+            value=tokens["id_token"],  # Changed from access_token
             httponly=True,
             secure=True,
             samesite="lax"