Added cognito authentication - Fix 2

app/main.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, Depends, HTTPException
-from fastapi.responses import RedirectResponse
-from app.cabletv.utils.auth import get_current_user, DOMAIN, CLIENT_ID
+from fastapi.responses import JSONResponse, RedirectResponse
+from app.cabletv.utils.auth import exchange_code_for_token, get_current_user, DOMAIN, CLIENT_ID
 
 app = FastAPI()
 
@@ -16,6 +16,28 @@ async def protected_route(user = Depends(get_current_user)):
 
 @app.get("/auth/callback")
 async def auth_callback(code: str):
-    # Here you would exchange the code for tokens
-    # For now, just redirect to protected route
-    return {"auth_code": code}
+    try:
+        # Exchange the authorization code for tokens
+        tokens = exchange_code_for_token(code)
+        
+        # Create a response with the access token
+        response = JSONResponse(content={
+            "message": "Authentication successful",
+            "access_token": tokens["access_token"]
+        })
+        
+        # Set the access token as a cookie
+        response.set_cookie(
+            key="access_token",
+            value=tokens["access_token"],
+            httponly=True,
+            secure=True,
+            samesite="lax"
+        )
+        
+        return response
+    except Exception as e:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Authentication failed: {str(e)}"
+        )