modify auth to allow terra login (#25)

src/__init__.py

@@ -7,13 +7,18 @@
 from google.cloud import firestore
 
 from src import cli, signaling, status
-from src.utils import custom_logging
+from src.utils import constants, custom_logging
 from src.web import web, participants, study
 
 logger = custom_logging.setup_logging(__name__)
 
 
 def create_app() -> Quart:
+    if constants.TERRA:
+        logger.info("Creating app - on Terra")
+    else:
+        logger.info("Creating app - NOT on Terra")
+
     initialize_firebase_admin()
 
     app = Quart(__name__)

src/api_utils.py

@@ -5,7 +5,7 @@
 from google.cloud.firestore_v1 import FieldFilter
 from quart import current_app
 
-from src.utils import custom_logging
+from src.utils import constants, custom_logging
 
 logger = custom_logging.setup_logging(__name__)
 
@@ -66,17 +66,20 @@ async def get_display_names() -> dict:
 
 
 async def add_user_to_db(decoded_token: dict) -> None:
-    logger.info(f"Creating user {decoded_token['sub']}")
+    user_id = decoded_token['id'] if constants.TERRA else decoded_token['sub']
+    logger.info(f"Creating user {user_id}")
     db = current_app.config["DATABASE"]
     try:
-        await db.collection("users").document(decoded_token["sub"]).set({"about": "", "notifications": []})
-        display_name = decoded_token["sub"]
-        if "given_name" in decoded_token:
+        await db.collection("users").document(user_id).set({"about": "", "notifications": []})
+        display_name = user_id
+        if constants.TERRA and "email" in decoded_token:
+            display_name = decoded_token["email"]
+        elif "given_name" in decoded_token:
             display_name = decoded_token["given_name"]
             if "family_name" in decoded_token:
                 display_name += " " + decoded_token["family_name"]
         await db.collection("users").document("display_names").set(
-            {decoded_token["sub"]: display_name}, merge=True
+            {user_id: display_name}, merge=True
         )
     except Exception as e:
         raise RuntimeError({"error": "Failed to create user", "details": str(e)}) from e

src/auth.py

@@ -1,6 +1,7 @@
 from functools import wraps
 
 from google.cloud import firestore
+import httpx
 import jwt
 from jwt import algorithms
 import requests
@@ -20,9 +21,39 @@
 
 
 async def get_user_id(request) -> str:
-    return (await verify_token(request.headers.get("Authorization").split(" ")[1]))["sub"]
+    auth_header: str = request.headers.get("Authorization")
+    bearer, token = auth_header.split(" ")
+    if bearer != "Bearer":
+        raise ValueError("Invalid Authorization header")
+    return (await verify_token(token))["id"] if constants.TERRA else (await verify_token(token))["sub"]
+
 
 async def verify_token(token):
+    if constants.TERRA:
+        return await verify_token_terra(token)
+    else:
+        return await verify_token_azure(token)
+
+
+async def verify_token_terra(token):
+    async with httpx.AsyncClient() as client:
+        headers = {
+            "accept": "application/json",
+            "Authorization": f"Bearer {token}",
+        }
+        response = await client.get(f"{constants.SAM_API_URL}/api/users/v2/self", headers=headers)
+
+    if response.status_code != 200:
+        raise ValueError("Token is invalid")
+
+    db: firestore.AsyncClient = current_app.config["DATABASE"]
+    if not (await db.collection("users").document(response.json()["id"]).get()).exists:
+        await add_user_to_db(response.json())
+
+    return response.json()
+
+
+async def verify_token_azure(token):
     headers = jwt.get_unverified_header(token)
     kid = headers["kid"]
 
@@ -59,7 +90,7 @@ async def decorated_function(*args, **kwargs):
         if not auth_header or "Bearer" not in auth_header:
             return jsonify({"message": "Authentication token required"}), 401
 
-        token = auth_header.split(" ")[1]
+        bearer, token = auth_header.split(" ")
 
         try:
             await verify_token(token)

src/utils/constants.py

@@ -1,6 +1,8 @@
 from copy import deepcopy
 import os
 
+TERRA = os.getenv("TERRA", "")
+SAM_API_URL = os.getenv("SAM_API_URL", "https://sam.dsde-dev.broadinstitute.org")
 APP_VERSION = os.getenv('APP_VERSION', '')
 BUILD_VERSION = os.getenv('BUILD_VERSION', '')
 SERVER_GCP_PROJECT = "broad-cho-priv1"

src/web/web.py

@@ -8,7 +8,7 @@
 from quart import Blueprint, Response, current_app, jsonify, request, send_file
 
 from src.api_utils import get_display_names, get_studies, is_valid_uuid
-from src.auth import authenticate, verify_token
+from src.auth import authenticate, get_user_id, verify_token
 from src.utils import custom_logging
 from src.utils.generic_functions import add_notification, remove_notification
 from src.utils.google_cloud.google_cloud_secret_manager import get_firebase_api_key
@@ -29,13 +29,12 @@
 @bp.route("/createCustomToken", methods=["POST"])
 @authenticate
 async def create_custom_token() -> Response:
-    user = await verify_token(request.headers.get("Authorization").split(" ")[1])
-    microsoft_user_id = user["sub"]
+    user_id = get_user_id(request)
     try:
         # Use the thread executor to run the blocking function
         loop = asyncio.get_event_loop()
         custom_token = await loop.run_in_executor(
-            None, firebase_auth.create_custom_token, microsoft_user_id
+            None, firebase_auth.create_custom_token, user_id
         )
         firebase_api_key = await get_firebase_api_key()
         return (