Added cognito authentication - Fix 9

app/cabletv/utils/auth.py

@@ -4,7 +4,6 @@ import requests
 import jwt
 from fastapi import Depends, HTTPException, status, Request
 from fastapi.security import OAuth2AuthorizationCodeBearer
-from fastapi.security.utils import get_authorization_scheme_param
 from fastapi.responses import RedirectResponse
 
 REGION = "us-east-2"
@@ -12,40 +11,13 @@ 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"
 
-class BrowserAwareOAuth2(OAuth2AuthorizationCodeBearer):
-    async def __call__(self, request: Request) -> str:
-        # Check if this is a browser request
-        is_browser = "text/html" in request.headers.get("accept", "")
-        
-        # Try to get token from cookie first, then header
-        authorization = request.cookies.get("token")
-        if not authorization:
-            authorization = request.headers.get("Authorization")
-        
-        scheme, param = get_authorization_scheme_param(authorization)
-        
-        if not authorization or scheme.lower() != "bearer":
-            if is_browser:
-                redirect_uri = str(request.base_url) + "auth/callback"
-                # Return redirect for browser requests
-                return RedirectResponse(
-                    f"{DOMAIN}/login?client_id={CLIENT_ID}"
-                    f"&response_type=code"
-                    f"&scope=openid+email+profile"
-                    f"&redirect_uri={redirect_uri}",
-                    status_code=302
-                )
-            # Return 401 for API requests
-            raise HTTPException(
-                status_code=401,
-                detail="Not authenticated",
-                headers={"WWW-Authenticate": "Bearer"},
-            )
-                
-        return param
+# 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)
 
-# Update the oauth2_scheme to use our custom class
-oauth2_scheme = BrowserAwareOAuth2(
+oauth2_scheme = DynamicOAuth2(
     authorizationUrl=f"{DOMAIN}/oauth2/authorize",
     tokenUrl=f"{DOMAIN}/oauth2/token"
 )

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"