Added cognito authentication - Fix 2

app/cabletv/utils/auth.py

@@ -1,27 +1,43 @@
+import os
+import boto3
+import requests
 from fastapi import Depends, HTTPException, status
 from fastapi.security import OAuth2AuthorizationCodeBearer
 from fastapi.responses import RedirectResponse
 from typing import Optional
-import os
-import boto3
 
 REGION = "us-east-2"
 USER_POOL_ID = os.getenv("COGNITO_USER_POOL_ID")
 CLIENT_ID = os.getenv("COGNITO_CLIENT_ID")
 DOMAIN = f"https://iptv-updater.auth.{REGION}.amazoncognito.com"
+REDIRECT_URI = "http://localhost:8000/auth/callback"
 
 oauth2_scheme = OAuth2AuthorizationCodeBearer(
     authorizationUrl=f"{DOMAIN}/oauth2/authorize",
     tokenUrl=f"{DOMAIN}/oauth2/token"
 )
 
+def exchange_code_for_token(code: str):
+    token_url = f"{DOMAIN}/oauth2/token"
+    data = {
+        'grant_type': 'authorization_code',
+        'client_id': CLIENT_ID,
+        'code': code,
+        'redirect_uri': REDIRECT_URI
+    }
+    
+    response = requests.post(token_url, data=data)
+    if response.status_code == 200:
+        return response.json()
+    raise HTTPException(status_code=400, detail="Failed to exchange code for token")
+
 async def get_current_user(token: str = Depends(oauth2_scheme)):
     if not token:
         return RedirectResponse(
             f"{DOMAIN}/login?client_id={CLIENT_ID}"
             f"&response_type=code"
             f"&scope=openid"
-            f"&redirect_uri=http://localhost:8000/auth/callback"
+            f"&redirect_uri={REDIRECT_URI}"
         )
     
     try:

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)}"
+        )