Added cognito authentication - Fix 5

app/cabletv/utils/auth.py

@@ -2,7 +2,7 @@ import os
 import boto3
 import requests
 import jwt
-from fastapi import Depends, HTTPException, status
+from fastapi import Depends, HTTPException, status, Request
 from fastapi.security import OAuth2AuthorizationCodeBearer
 from fastapi.responses import RedirectResponse
 
@@ -10,42 +10,47 @@ 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 = f"http://localhost:8000/auth/callback"
 
-oauth2_scheme = OAuth2AuthorizationCodeBearer(
+# Remove the hardcoded REDIRECT_URI, we'll make it dynamic based on the request
+class DynamicOAuth2(OAuth2AuthorizationCodeBearer):
+    async def __call__(self, request: Request):
+        self.redirect_uri = str(request.base_url) + "auth/callback"
+        return await super().__call__(request)
+
+oauth2_scheme = DynamicOAuth2(
     authorizationUrl=f"{DOMAIN}/oauth2/authorize",
     tokenUrl=f"{DOMAIN}/oauth2/token"
 )
 
-def exchange_code_for_token(code: str):
+def exchange_code_for_token(code: str, redirect_uri: str):
     token_url = f"{DOMAIN}/oauth2/token"
     data = {
         'grant_type': 'authorization_code',
         'client_id': CLIENT_ID,
         'code': code,
-        'redirect_uri': REDIRECT_URI
+        'redirect_uri': redirect_uri
     }
     
     response = requests.post(token_url, data=data)
     if response.status_code == 200:
         return response.json()
-    print(f"Token exchange failed: {response.text}")  # Add logging
+    print(f"Token exchange failed: {response.text}")
     raise HTTPException(status_code=400, detail="Failed to exchange code for token")
 
-async def get_current_user(token: str = Depends(oauth2_scheme)):
+async def get_current_user(request: Request, token: str = Depends(oauth2_scheme)):
     if not token:
+        redirect_uri = str(request.base_url) + "auth/callback"
         return RedirectResponse(
             f"{DOMAIN}/login?client_id={CLIENT_ID}"
             f"&response_type=code"
-            f"&scope=openid+email+profile"  # Added more scopes
-            f"&redirect_uri={REDIRECT_URI}"
+            f"&scope=openid+email+profile"
+            f"&redirect_uri={redirect_uri}"
         )
     
     try:
-        # Decode JWT token instead of using get_user
         decoded = jwt.decode(
             token,
-            options={"verify_signature": False}  # We trust tokens from Cognito
+            options={"verify_signature": False}
         )
         return {
             "Username": decoded.get("email") or decoded.get("sub"),
@@ -55,7 +60,7 @@ async def get_current_user(token: str = Depends(oauth2_scheme)):
             ]
         }
     except Exception as e:
-        print(f"Token verification failed: {str(e)}")  # Add logging
+        print(f"Token verification failed: {str(e)}")
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Invalid authentication credentials",

app/main.py

@@ -1,6 +1,6 @@
-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
+from fastapi import FastAPI, Depends, HTTPException, Request
+from fastapi.responses import RedirectResponse, JSONResponse
+from app.cabletv.utils.auth import get_current_user, exchange_code_for_token
 
 app = FastAPI()
 
@@ -9,26 +9,25 @@ async def root():
     return {"message": "IPTV Updater API"}
 
 @app.get("/protected")
-async def protected_route(user = Depends(get_current_user)):
+async def protected_route(request: Request, 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(code: str):
+async def auth_callback(request: Request, code: str):
     try:
-        tokens = exchange_code_for_token(code)
+        redirect_uri = str(request.base_url)
+        tokens = exchange_code_for_token(code, redirect_uri)
         
-        # Use id_token instead of access_token
         response = JSONResponse(content={
             "message": "Authentication successful",
-            "id_token": tokens["id_token"]  # Changed from access_token
+            "id_token": tokens["id_token"]
         })
         
-        # Store id_token in cookie
         response.set_cookie(
             key="token",
-            value=tokens["id_token"],  # Changed from access_token
+            value=tokens["id_token"],
             httponly=True,
             secure=True,
             samesite="lax"

infrastructure/stack.py

@@ -125,7 +125,8 @@ class IptvUpdaterStack(Stack):
                 scopes=[cognito.OAuthScope.OPENID],
                 callback_urls=[
                    "http://localhost:8000/auth/callback",  # For local testing
-                    "https://*.amazonaws.com/auth/callback"  # Will match EC2 public DNS
+                    "http://*.amazonaws.com/auth/callback",  # EC2 public DNS
+                    "http://*.compute.amazonaws.com/auth/callback"  # EC2 full domain
                  ]
             )
         )