Use keycloak to handle auth

bats_ai/core/views/nabat/nabat_recording.py

@@ -68,6 +68,14 @@ def admin_auth(request):
 """
 
 
+def get_auth_header(request: HttpRequest):
+    auth_header = request.headers.get("Authorization")
+    if not auth_header:
+        return None
+    auth_header_parts = auth_header.split(" ")
+    return auth_header_parts[1] if len(auth_header_parts) > 1 else None
+
+
 def decode_jwt(token):
     # Split the token into parts
     parts = token.split(".")
@@ -89,7 +97,6 @@ def decode_jwt(token):
 
 def get_email_if_authorized(  # noqa: PLR0911
     request: HttpRequest,
-    api_token: str,
     recording_id: int | None = None,
     recording_pk: int | None = None,
 ) -> str | JsonResponse:
@@ -106,6 +113,7 @@ def get_email_if_authorized(  # noqa: PLR0911
     if request.user and request.user.is_authenticated and request.user.is_superuser:
         return request.user.email or "superuser@nabat.org"
     # Decode JWT token
+    api_token = get_auth_header(request)
     try:
         payload = decode_jwt(api_token)
         email = payload.get("email")
@@ -165,11 +173,17 @@ class NABatRecordingSchema(Schema):
 
 
 class NABatRecordingGenerateSchema(Schema):
-    apiToken: str
     recordingId: int
     surveyEventId: int
 
 
+class NABatAuthorizationSchema(Schema):
+    recordingId: int
+    surveyEventId: int
+    iss: str
+    code: str
+
+
 def update_nabat_species(species_id: int, api_token: str, recording_id: int, survey_event_id: int):
     """
     Update the species for a NABat recording using the NABat API.
@@ -200,6 +214,43 @@ def update_nabat_species(species_id: int, api_token: str, recording_id: int, sur
     return "NABat species updated successfully."
 
 
+def _reconstruct_redirect_url(survey_event_id, recording_id):
+    url_root = settings.BATAI_WEB_URL.rstrip("/")
+    return f"{url_root}/nabat/auth/?recordingId={recording_id}&surveyEventId={survey_event_id}"
+
+
+@router.post("/authorize", auth=None)
+def authorize_nabat_requests(request: HttpRequest, payload: Form[NABatAuthorizationSchema]):
+    if payload.iss != settings.BATAI_NABAT_OIDC_ISSUER:
+        return JsonResponse({"error": "Unexpected issuer"}, status=400)
+
+    try:
+        redirect_url = _reconstruct_redirect_url(payload.surveyEventId, payload.recordingId)
+        response = requests.post(
+            f"{settings.BATAI_NABAT_OIDC_BASE_URL}/protocol/openid-connect/token",
+            data={
+                "grant_type": "authorization_code",
+                "client_id": settings.BATAI_NABAT_OIDC_CLIENT_ID,
+                "client_secret": settings.BATAI_NABAT_OIDC_CLIENT_SECRET,
+                "redirect_uri": redirect_url,
+                "code": payload.code,
+            },
+            timeout=30,
+        )
+        response_json = response.json()
+        if response.status_code != 200:
+            logger.error(
+                "Keycloak token exchange rejected: %s - %s", response.status_code, response.text
+            )
+            return JsonResponse(
+                {"error": "Keycloak token exchange failed."}, status=response.status_code
+            )
+        return JsonResponse(response_json, status=200)
+    except Exception:
+        logger.exception("Keycloak token exchange failed")
+        return JsonResponse({"error": "Keycloak token exchange failed"}, status=500)
+
+
 @router.post("/", auth=None)
 def generate_nabat_recording(  # noqa: PLR0911
     request: HttpRequest,
@@ -221,10 +272,13 @@ def generate_nabat_recording(  # noqa: PLR0911
         )
 
     nabat_recording = NABatRecording.objects.filter(recording_id=payload.recordingId)
+    api_token = get_auth_header(request)
+    if not api_token:
+        return JsonResponse({"error": "Missing or invalid Authorization header"}, status=401)
     if not nabat_recording.exists():
         # use a task to start downloading the file using the API key and generate the spectrograms
         task = nabat_recording_initialize.delay(
-            payload.recordingId, payload.surveyEventId, payload.apiToken
+            payload.recordingId, payload.surveyEventId, api_token
         )
         with transaction.atomic():
             ProcessingTask.objects.create(
@@ -237,9 +291,8 @@ def generate_nabat_recording(  # noqa: PLR0911
                 celery_id=task.id,
             )
         return {"taskId": task.id}
-    # we want to check the apiToken and make sure the user has access to the file
+    # we want to check the api token and make sure the user has access to the file
     # before returning it
-    api_token = payload.apiToken
     recording_id = payload.recordingId
     headers = {"Authorization": f"Bearer {api_token}", "Content-Type": "application/json"}
     batch_query = QUERY % {
@@ -308,11 +361,10 @@ def get_spectrogram(request: HttpRequest, pk: int):
 def get_spectrogram_compressed(
     request: HttpRequest,
     pk: int,
-    apiToken: str,  # noqa: N803
 ):
     nabat_recording = get_object_or_404(NABatRecording, pk=pk)
 
-    email_or_response = get_email_if_authorized(request, apiToken, nabat_recording.recording_id)
+    email_or_response = get_email_if_authorized(request, nabat_recording.recording_id)
     if isinstance(email_or_response, JsonResponse):
         return email_or_response
 
@@ -395,16 +447,14 @@ class NABatCreateRecordingAnnotationSchema(Schema):
     comments: str = None
     model: str = None
     confidence: float
-    apiToken: str
 
 
 @router.get("/{nabat_recording_id}/recording-annotations", auth=admin_auth)
 def get_nabat_recording_annotation(
     request: HttpRequest,
     nabat_recording_id: int,
-    apiToken: str | None = None,  # noqa: N803
 ):
-    email_or_response = get_email_if_authorized(request, apiToken, recording_pk=nabat_recording_id)
+    email_or_response = get_email_if_authorized(request, recording_pk=nabat_recording_id)
     if isinstance(email_or_response, JsonResponse):
         return email_or_response
     user_email = email_or_response  # safe to use
@@ -431,9 +481,8 @@ def get_nabat_recording_annotation(
 def get_recording_annotation(
     request: HttpRequest,
     pk: int,
-    apiToken: str,  # noqa: N803
 ):
-    email_or_response = get_email_if_authorized(request, apiToken, recording_pk=pk)
+    email_or_response = get_email_if_authorized(request, recording_pk=pk)
     if isinstance(email_or_response, JsonResponse):
         return email_or_response
     user_email = email_or_response  # safe to use
@@ -454,9 +503,8 @@ def get_recording_annotation(
 def get_recording_annotation_details(
     request: HttpRequest,
     pk: int,
-    apiToken: str,  # noqa: N803
 ):
-    email_or_response = get_email_if_authorized(request, apiToken, recording_pk=pk)
+    email_or_response = get_email_if_authorized(request, recording_pk=pk)
     if isinstance(email_or_response, JsonResponse):
         return email_or_response
     user_email = email_or_response  # safe to use
@@ -470,14 +518,13 @@ def get_recording_annotation_details(
 
 @router.put("recording-annotation", auth=None, response={200: str})
 def create_recording_annotation(request: HttpRequest, data: NABatCreateRecordingAnnotationSchema):
-    email_or_response = get_email_if_authorized(
-        request, data.apiToken, recording_pk=data.recordingId
-    )
+    email_or_response = get_email_if_authorized(request, recording_pk=data.recordingId)
     if isinstance(email_or_response, JsonResponse):
         return email_or_response
     user_email = email_or_response  # safe to use
 
-    token_data = decode_jwt(data.apiToken)
+    api_token = get_auth_header(request)
+    token_data = decode_jwt(api_token)
     user_id = token_data["sub"]
 
     recording = get_object_or_404(NABatRecording, pk=data.recordingId)
@@ -512,9 +559,7 @@ def update_recording_annotation(
     it clears the existing species and adds the new ones. The endpoint returns a
     success message or an error message if the recording or species are not found.
     """
-    email_or_response = get_email_if_authorized(
-        request, data.apiToken, recording_pk=data.recordingId
-    )
+    email_or_response = get_email_if_authorized(request, recording_pk=data.recordingId)
     if isinstance(email_or_response, JsonResponse):
         return email_or_response
     user_email = email_or_response  # safe to use
@@ -544,9 +589,7 @@ def update_nabat_recording_annotation(
     request: HttpRequest, pk: int, data: NABatCreateRecordingAnnotationSchema
 ):
     """Update an existing recording annotation in NABat."""
-    email_or_response = get_email_if_authorized(
-        request, data.apiToken, recording_pk=data.recordingId
-    )
+    email_or_response = get_email_if_authorized(request, recording_pk=data.recordingId)
     if isinstance(email_or_response, JsonResponse):
         return email_or_response
     user_email = email_or_response  # safe to use
@@ -569,9 +612,12 @@ def update_nabat_recording_annotation(
             status=400,
         )
     species_id = data.species[0]
+    # We can pull the Authorization token out of the headers because
+    # if it didn't exist we would have returned already
+    api_token = get_auth_header(request)
     return update_nabat_species(
         species_id,
-        data.apiToken,
+        api_token,
         annotation.nabat_recording.recording_id,
         annotation.nabat_recording.survey_event_id,
     )
@@ -582,10 +628,9 @@ def update_nabat_recording_annotation(
 def delete_recording_annotation(
     request: HttpRequest,
     pk: int,
-    apiToken: str,  # noqa: N803
     recordingId: str,  # noqa: N803
 ):
-    email_or_response = get_email_if_authorized(request, apiToken, recording_pk=recordingId)
+    email_or_response = get_email_if_authorized(request, recording_pk=recordingId)
     if isinstance(email_or_response, JsonResponse):
         return email_or_response
     user_email = email_or_response  # safe to use
@@ -646,10 +691,10 @@ def linestring_to_list(ls):
 
 
 @router.get("/{pk}/pulse_contours", auth=None)
-def get_pulse_contours(request: HttpRequest, pk: int, api_token: str):
+def get_pulse_contours(request: HttpRequest, pk: int):
     recording = get_object_or_404(NABatRecording, pk=pk)
 
-    email_or_response = get_email_if_authorized(request, api_token, recording.recording_id)
+    email_or_response = get_email_if_authorized(request, recording.recording_id)
     if isinstance(email_or_response, JsonResponse):
         return email_or_response
 
@@ -660,10 +705,10 @@ def get_pulse_contours(request: HttpRequest, pk: int, api_token: str):
 
 
 @router.get("/{pk}/pulse_metadata", auth=None)
-def get_pulse_data(request: HttpRequest, pk: int, api_token: str):
+def get_pulse_data(request: HttpRequest, pk: int):
     recording = get_object_or_404(NABatRecording, pk=pk)
 
-    email_or_response = get_email_if_authorized(request, api_token, recording.recording_id)
+    email_or_response = get_email_if_authorized(request, recording.recording_id)
     if isinstance(email_or_response, JsonResponse):
         return email_or_response
 

bats_ai/settings/development.py

@@ -57,3 +57,14 @@
 SHELL_PLUS_IMPORTS = [
     "from bats_ai.core import tasks",
 ]
+
+BATAI_NABAT_OIDC_CLIENT_ID: str = env.str("DJANGO_BATAI_NABAT_OIDC_CLIENT_ID", default="batai")
+BATAI_NABAT_OIDC_CLIENT_SECRET: str = env.str(
+    "DJANGO_BATAI_NABAT_OIDC_CLIENT_SECRET", default="batai-local-dev-secret"
+)
+BATAI_NABAT_OIDC_ISSUER: str = env.str(
+    "DJANGO_BATAI_NABAT_OIDC_ISSUER", default="http://localhost:8081/auth/realms/NABAT"
+)
+BATAI_NABAT_OIDC_BASE_URL: str = env.str(
+    "DJANGO_BATAI_NABAT_OIDC_BASE_URL", default="http://localhost:8081/auth/realms/NABAT"
+)

bats_ai/settings/nabat_production.py

@@ -27,3 +27,9 @@
     FORCE_SCRIPT_NAME = _proxy_subpath
     # Work around https://code.djangoproject.com/ticket/36653
     STORAGES["staticfiles"].setdefault("OPTIONS", {})["base_url"] = f"{_proxy_subpath}/{STATIC_URL}"
+
+
+BATAI_NABAT_OIDC_CLIENT_ID: str = env.str("DJANGO_BATAI_NABAT_OIDC_CLIENT_ID")
+BATAI_NABAT_OIDC_CLIENT_SECRET: str = env.str("DJANGO_BATAI_NABAT_OIDC_CLIENT_SECRET")
+BATAI_NABAT_OIDC_ISSUER: str = env.str("DJANGO_BATAI_NABAT_OIDC_ISSUER")
+BATAI_NABAT_OIDC_BASE_URL: str = env.str("DJANGO_BATAI_NABAT_OIDC_BASE_URL")