Switch to cognito user/password authentication. Major code refactor.

.vscode/settings.json

@@ -1,5 +1,6 @@
 {
     "cSpell.words": [
+        "adminpassword",
         "altinstall",
         "awscliv",
         "boto",

app.py

@@ -1,7 +1,11 @@
 #!/usr/bin/env python3
 import aws_cdk as cdk
+import uvicorn
 from infrastructure.stack import IptvUpdaterStack
 
 app = cdk.App()
 IptvUpdaterStack(app, "IptvUpdater")
 app.synth()
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)

app/auth/cognito.py

@@ -0,0 +1,81 @@
+import boto3
+from fastapi import HTTPException, status
+from app.models.auth import CognitoUser
+from app.utils.auth import calculate_secret_hash
+from app.utils.constants import (AWS_REGION, COGNITO_CLIENT_ID,
+                            COGNITO_CLIENT_SECRET, USER_ROLE_ATTRIBUTE)
+
+cognito_client = boto3.client("cognito-idp", region_name=AWS_REGION)
+
+
+def initiate_auth(username: str, password: str) -> dict:
+    """
+    Initiate AUTH flow with Cognito using USER_PASSWORD_AUTH.
+    """
+    auth_params = {
+        "USERNAME": username,
+        "PASSWORD": password
+    }
+
+    # If a client secret is required, add SECRET_HASH
+    if COGNITO_CLIENT_SECRET:
+        auth_params["SECRET_HASH"] = calculate_secret_hash(
+            username, COGNITO_CLIENT_ID, COGNITO_CLIENT_SECRET)
+
+    try:
+        response = cognito_client.initiate_auth(
+            AuthFlow="USER_PASSWORD_AUTH",
+            AuthParameters=auth_params,
+            ClientId=COGNITO_CLIENT_ID
+        )
+        return response["AuthenticationResult"]
+    except cognito_client.exceptions.NotAuthorizedException:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid username or password"
+        )
+    except cognito_client.exceptions.UserNotFoundException:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="User not found"
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"An error occurred during authentication: {str(e)}"
+        )
+
+
+def get_user_from_token(access_token: str) -> CognitoUser:
+    """
+    Verify the token by calling GetUser in Cognito and retrieve user attributes including roles.
+    """
+    try:
+        user_response = cognito_client.get_user(AccessToken=access_token)
+        username = user_response.get("Username", "")
+        attributes = user_response.get("UserAttributes", [])
+        user_roles = []
+
+        for attr in attributes:
+            if attr["Name"] == USER_ROLE_ATTRIBUTE:
+                # Assume roles are stored as a comma-separated string
+                user_roles = [r.strip()
+                              for r in attr["Value"].split(",") if r.strip()]
+                break
+
+        return CognitoUser(username=username, roles=user_roles)
+    except cognito_client.exceptions.NotAuthorizedException:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or expired token."
+        )
+    except cognito_client.exceptions.UserNotFoundException:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="User not found or invalid token."
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Token verification failed: {str(e)}"
+        )

app/auth/dependencies.py

@@ -0,0 +1,41 @@
+from functools import wraps
+from typing import Callable
+
+from fastapi import Depends, HTTPException, status
+from fastapi.security import OAuth2PasswordBearer
+
+from app.auth.cognito import get_user_from_token
+from app.models.auth import CognitoUser
+
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="signin")
+
+
+def get_current_user(token: str = Depends(oauth2_scheme)) -> CognitoUser:
+    """
+    Dependency to get the current user from the given token.
+    This will verify the token with Cognito and return the user's information.
+    """
+    return get_user_from_token(token)
+
+
+def require_roles(*required_roles: str) -> Callable:
+    """
+    Decorator for role-based access control.
+    Use on endpoints to enforce that the user possesses all required roles.
+    """
+
+    def decorator(endpoint: Callable) -> Callable:
+        @wraps(endpoint)
+        def wrapper(*args, user: CognitoUser = Depends(get_current_user), **kwargs):
+            user_roles = set(user.roles or [])
+            needed_roles = set(required_roles)
+            if not needed_roles.issubset(user_roles):
+                raise HTTPException(
+                    status_code=status.HTTP_403_FORBIDDEN,
+                    detail="You do not have the required roles to access this endpoint.",
+                )
+            return endpoint(*args, user=user, **kwargs)
+
+        return wrapper
+
+    return decorator

app/cabletv/utils/auth.py

@@ -1,62 +0,0 @@
-import os
-import requests
-import jwt
-from fastapi import Depends, HTTPException, status
-from fastapi.security import OAuth2AuthorizationCodeBearer
-from fastapi.responses import RedirectResponse
-
-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(
-    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()
-    print(f"Token exchange failed: {response.text}")  # Add logging
-    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+email+profile"  # Added more scopes
-            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
-        )
-        return {
-            "Username": decoded.get("email") or decoded.get("sub"),
-            "UserAttributes": [
-                {"Name": k, "Value": v}
-                for k, v in decoded.items()
-            ]
-        }
-    except Exception as e:
-        print(f"Token verification failed: {str(e)}")  # Add logging
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Invalid authentication credentials",
-            headers={"WWW-Authenticate": "Bearer"},
-        )

app/iptv/createEpg.py

@@ -5,7 +5,7 @@ import json
 import xml.etree.ElementTree as ET
 import requests
 import argparse
-from utils.config import IPTV_SERVER_ADMIN_PASSWORD, IPTV_SERVER_ADMIN_USER, IPTV_SERVER_URL
+from utils.constants import IPTV_SERVER_ADMIN_PASSWORD, IPTV_SERVER_ADMIN_USER, IPTV_SERVER_URL
 
 def parse_arguments():
     parser = argparse.ArgumentParser(description='EPG Grabber')

app/iptv/createPlaylist.py

@@ -6,7 +6,7 @@ import requests
 from pathlib import Path
 from datetime import datetime
 from utils.check_streams import StreamValidator
-from utils.config import EPG_URL, IPTV_SERVER_ADMIN_PASSWORD, IPTV_SERVER_ADMIN_USER, IPTV_SERVER_URL
+from utils.constants import EPG_URL, IPTV_SERVER_ADMIN_PASSWORD, IPTV_SERVER_ADMIN_USER, IPTV_SERVER_URL
 
 def parse_arguments():
     parser = argparse.ArgumentParser(description='IPTV playlist generator')

app/main.py

@@ -1,6 +1,9 @@
-from fastapi import FastAPI, Depends, HTTPException
-from fastapi.responses import RedirectResponse, JSONResponse
-from app.cabletv.utils.auth import get_current_user, exchange_code_for_token
+import uvicorn
+from fastapi import FastAPI, Depends
+from fastapi.responses import RedirectResponse
+from app.auth.cognito import initiate_auth
+from app.auth.dependencies import get_current_user, require_roles
+from app.models.auth import CognitoUser, SigninRequest, TokenResponse
 
 app = FastAPI()
 
@@ -8,35 +11,33 @@ app = FastAPI()
 async def root():
     return {"message": "IPTV Updater API"}
 
+@app.post("/signin", response_model=TokenResponse, summary="Signin Endpoint")
+def signin(credentials: SigninRequest):
+    """
+    Sign-in endpoint to authenticate the user with AWS Cognito using username and password.
+    On success, returns JWT tokens (access_token, id_token, refresh_token).
+    """
+    auth_result = initiate_auth(credentials.username, credentials.password)
+    return TokenResponse(
+        access_token=auth_result["AccessToken"],
+        id_token=auth_result["IdToken"],
+        refresh_token=auth_result.get("RefreshToken"),
+        token_type="Bearer",
+    )
+
 @app.get("/protected")
-async def protected_route(user = Depends(get_current_user)):
-    if isinstance(user, RedirectResponse):
-        return user
-    return {"message": "Protected content", "user": user['Username']}
+async def protected_route(user: CognitoUser = Depends(get_current_user)):
+    """
+    Protected endpoint that requires for all authenticated users.
+    If the user is authenticates, returns success message.
+    """
+    return {"message": f"Hello {user.username}, you have access to support resources!"}
 
-@app.get("/auth/callback")
-async def auth_callback(code: str):
-    try:
-        tokens = exchange_code_for_token(code)
-        
-        # Use id_token instead of access_token
-        response = JSONResponse(content={
-            "message": "Authentication successful",
-            "id_token": tokens["id_token"]  # Changed from access_token
-        })
-        
-        # Store id_token in cookie
-        response.set_cookie(
-            key="token",
-            value=tokens["id_token"],  # Changed from 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)}"
-        )
+@app.get("/protected_admin", summary="Protected endpoint for Admin role")
+@require_roles("admin")
+def protected_admin_endpoint(user: CognitoUser = Depends(get_current_user)):
+    """
+    Protected endpoint that requires the 'admin' role.
+    If the user has 'admin' role, returns success message.
+    """
+    return {"message": f"Hello {user.username}, you have admin privileges!"}

app/models/auth.py

@@ -0,0 +1,20 @@
+from typing import List, Optional
+from pydantic import BaseModel, Field
+
+class SigninRequest(BaseModel):
+    """Request model for the signin endpoint."""
+    username: str = Field(..., description="The user's username")
+    password: str = Field(..., description="The user's password")
+
+class TokenResponse(BaseModel):
+    """Response model for successful authentication."""
+    access_token: str = Field(..., description="Access JWT token from Cognito")
+    id_token: str = Field(..., description="ID JWT token from Cognito")
+    refresh_token: Optional[str] = Field(
+        None, description="Refresh token from Cognito")
+    token_type: str = Field(..., description="Type of the token returned")
+
+class CognitoUser(BaseModel):
+    """Model representing the user returned from token verification."""
+    username: str
+    roles: List[str]

app/utils/auth.py

@@ -0,0 +1,12 @@
+import base64
+import hashlib
+import hmac
+
+def calculate_secret_hash(username: str, client_id: str, client_secret: str) -> str:
+    """
+    Calculate the Cognito SECRET_HASH using HMAC SHA256 for secret-enabled clients.
+    """
+    msg = username + client_id
+    dig = hmac.new(client_secret.encode('utf-8'),
+                   msg.encode('utf-8'), hashlib.sha256).digest()
+    return base64.b64encode(dig).decode()

app/utils/constants.py

@@ -1,12 +1,21 @@
+# Utility functions and constants
 import os
+
 from dotenv import load_dotenv
 
 # Load environment variables from a .env file if it exists
 load_dotenv()
 
+# AWS related constants
+AWS_REGION = os.environ.get("AWS_REGION", "us-east-2")
+COGNITO_USER_POOL_ID = os.getenv("COGNITO_USER_POOL_ID")
+COGNITO_CLIENT_ID = os.getenv("COGNITO_CLIENT_ID")
+COGNITO_CLIENT_SECRET = os.environ.get("COGNITO_CLIENT_SECRET", None)
+USER_ROLE_ATTRIBUTE = "custom:role"
+
 IPTV_SERVER_URL = os.getenv("IPTV_SERVER_URL", "https://iptv.fiorinis.com")
 
-# Super iptv-server admin credentials for basic auth
+# iptv-server super admin credentials for basic auth
 # Reads from environment variables IPTV_SERVER_ADMIN_USER and IPTV_SERVER_ADMIN_PASSWORD
 IPTV_SERVER_ADMIN_USER = os.getenv("IPTV_SERVER_ADMIN_USER", "admin")
 IPTV_SERVER_ADMIN_PASSWORD = os.getenv("IPTV_SERVER_ADMIN_PASSWORD", "adminpassword")

infrastructure/stack.py

@@ -1,5 +1,6 @@
 import os
 from aws_cdk import (
+    Duration,
     Stack,
     aws_ec2 as ec2,
     aws_iam as iam,
@@ -118,17 +119,20 @@ class IptvUpdaterStack(Stack):
 
         # Add App Client with the correct callback URL
         client = user_pool.add_client("IptvUpdaterClient",
+            access_token_validity=Duration.minutes(60),
+            id_token_validity=Duration.minutes(60),
+            refresh_token_validity=Duration.days(1),
+            auth_flows=cognito.AuthFlow(
+                user_password=True
+            ),
             o_auth=cognito.OAuthSettings(
                 flows=cognito.OAuthFlows(
-                    authorization_code_grant=True
-                ),
-                scopes=[cognito.OAuthScope.OPENID],
-                callback_urls=[
-                    "http://localhost:8000/auth/callback",  # For local testing
-                    "https://*.amazonaws.com/auth/callback",  # EC2 public DNS
-                    "https://*.compute.amazonaws.com/auth/callback"  # EC2 full domain
-                 ]
-            )
+                    implicit_code_grant=True
+                )
+            ),
+            prevent_user_existence_errors=True,
+            generate_secret=True,
+            enable_token_revocation=True
         )
 
         # Add domain for hosted UI