From 658f7998ef5034569933c09db6d0c601ec7eecf4 Mon Sep 17 00:00:00 2001
From: Stefano <stefano@fiorinis.com>
Date: Fri, 16 May 2025 13:00:09 -0500
Subject: [PATCH] Switch to cognito user/password authentication. Major code
 refactor - Fix 2

---
 app/auth/dependencies.py |  6 +++--
 app/main.py              | 53 ++++++++++++++++++++++++++++++++++++++--
 app/utils/constants.py   |  2 +-
 infrastructure/stack.py  |  3 ++-
 4 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/app/auth/dependencies.py b/app/auth/dependencies.py
index 579e28e..6afcd8f 100644
--- a/app/auth/dependencies.py
+++ b/app/auth/dependencies.py
@@ -7,8 +7,10 @@ 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")
-
+oauth2_scheme = OAuth2PasswordBearer(
+    tokenUrl="signin",
+    scheme_name="Bearer"
+)
 
 def get_current_user(token: str = Depends(oauth2_scheme)) -> CognitoUser:
     """
diff --git a/app/main.py b/app/main.py
index 6229bab..c3d1bd9 100644
--- a/app/main.py
+++ b/app/main.py
@@ -1,3 +1,4 @@
+from fastapi.security import OAuth2PasswordBearer
 import uvicorn
 from fastapi import FastAPI, Depends
 from fastapi.responses import RedirectResponse
@@ -5,7 +6,54 @@ 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()
+from fastapi import FastAPI, Depends, Security
+from fastapi.security import OAuth2PasswordBearer
+from fastapi.openapi.utils import get_openapi
+
+app = FastAPI(
+    title="IPTV Updater API",
+    description="API for IPTV Updater service",
+    version="1.0.0",
+)
+
+def custom_openapi():
+    if app.openapi_schema:
+        return app.openapi_schema
+
+    openapi_schema = get_openapi(
+        title=app.title,
+        version=app.version,
+        description=app.description,
+        routes=app.routes,
+    )
+
+    # Ensure components object exists
+    if "components" not in openapi_schema:
+        openapi_schema["components"] = {}
+
+    # Add schemas if they don't exist
+    if "schemas" not in openapi_schema["components"]:
+        openapi_schema["components"]["schemas"] = {}
+
+    # Add security scheme component
+    openapi_schema["components"]["securitySchemes"] = {
+        "Bearer": {
+            "type": "http",
+            "scheme": "bearer",
+            "bearerFormat": "JWT"
+        }
+    }
+
+    # Add global security requirement
+    openapi_schema["security"] = [{"Bearer": []}]
+
+    # Set OpenAPI version explicitly
+    openapi_schema["openapi"] = "3.1.0"
+
+    app.openapi_schema = openapi_schema
+    return app.openapi_schema
+
+app.openapi = custom_openapi
 
 @app.get("/")
 async def root():
@@ -25,7 +73,8 @@ def signin(credentials: SigninRequest):
         token_type="Bearer",
     )
 
-@app.get("/protected")
+@app.get("/protected", 
+         summary="Protected endpoint for authenticated users")
 async def protected_route(user: CognitoUser = Depends(get_current_user)):
     """
     Protected endpoint that requires for all authenticated users.
diff --git a/app/utils/constants.py b/app/utils/constants.py
index a3c4183..8bc9aab 100644
--- a/app/utils/constants.py
+++ b/app/utils/constants.py
@@ -11,7 +11,7 @@ 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"
+USER_ROLE_ATTRIBUTE = "zoneinfo"
 
 IPTV_SERVER_URL = os.getenv("IPTV_SERVER_URL", "https://iptv.fiorinis.com")
 
diff --git a/infrastructure/stack.py b/infrastructure/stack.py
index 53d9cb1..dcc4139 100644
--- a/infrastructure/stack.py
+++ b/infrastructure/stack.py
@@ -152,7 +152,8 @@ class IptvUpdaterStack(Stack):
         # Adds one or more commands to the userdata object.
         userdata.add_commands(
             f'echo "COGNITO_USER_POOL_ID={user_pool.user_pool_id}" >> /etc/environment',
-            f'echo "COGNITO_CLIENT_ID={client.user_pool_client_id}" >> /etc/environment'
+            f'echo "COGNITO_CLIENT_ID={client.user_pool_client_id}" >> /etc/environment',
+            f'echo "COGNITO_CLIENT_SECRET={client.user_pool_client_secret}" >> /etc/environment'
         )
         userdata.add_commands(str(userdata_file, 'utf-8'))