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 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={REDIRECT_URI}" ) 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"}, )