Added cognito authentication - Fix 1

app/cabletv/utils/auth.py

@@ -0,0 +1,36 @@
+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"
+
+oauth2_scheme = OAuth2AuthorizationCodeBearer(
+    authorizationUrl=f"{DOMAIN}/oauth2/authorize",
+    tokenUrl=f"{DOMAIN}/oauth2/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"
+        )
+    
+    try:
+        cognito = boto3.client('cognito-idp', region_name=REGION)
+        response = cognito.get_user(AccessToken=token)
+        return response
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid authentication credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )

app/main.py

@@ -1,46 +1,21 @@
-import os
-import json
-import boto3
-from jose import jwt
 from fastapi import FastAPI, Depends, HTTPException
-from fastapi.security import OAuth2AuthorizationCodeBearer
+from fastapi.responses import RedirectResponse
+from app.cabletv.utils.auth import get_current_user, DOMAIN, CLIENT_ID
 
 app = FastAPI()
 
-# Get Cognito info from environment (set by userdata.sh)
-REGION = "us-east-2"
-USER_POOL_ID = os.getenv("COGNITO_USER_POOL_ID")
-CLIENT_ID = os.getenv("COGNITO_CLIENT_ID")
-
-# OAuth2 scheme for authorization code flow
-oauth2_scheme = OAuth2AuthorizationCodeBearer(
-    authorizationUrl=f"https://cognito-idp.{REGION}.amazonaws.com/{USER_POOL_ID}/oauth2/authorize",
-    tokenUrl=f"https://cognito-idp.{REGION}.amazonaws.com/{USER_POOL_ID}/oauth2/token"
-)
-
-async def get_current_user(token: str = Depends(oauth2_scheme)):
-    try:
-        # Verify the JWT token with Cognito
-        cognito_idp = boto3.client('cognito-idp', region_name=REGION)
-        response = cognito_idp.get_user(
-            AccessToken=token
-        )
-        return response
-    except Exception as e:
-        raise HTTPException(
-            status_code=401,
-            detail="Invalid authentication credentials",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-
 @app.get("/")
 async def root():
     return {"message": "IPTV Updater API"}
 
-@app.get("/health")
-async def health():
-    return {"status": "healthy"}
-
 @app.get("/protected")
 async def protected_route(user = Depends(get_current_user)):
-    return {"message": "This is a protected route", "user": user['Username']}
+    if isinstance(user, RedirectResponse):
+        return user
+    return {"message": "Protected content", "user": user['Username']}
+
+@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}