Merge pull request #218 from openzim/healthcheck-optional-auth

perform authentication healthcheck with oauth token

Author: elfkuzco

healthcheck/src/healthcheck/context.py

@@ -1,3 +1,4 @@
+import datetime
 import os
 from typing import Any
 
@@ -29,8 +30,18 @@ class Context:
 
     cms_api_url = get_mandatory_env("CMS_API_URL")
     cms_frontend_url = get_mandatory_env("CMS_FRONTEND_URL")
-    cms_username = get_mandatory_env("CMS_USERNAME")
-    cms_password = get_mandatory_env("CMS_PASSWORD")
+    auth_mode = os.getenv("AUTH_MODE", default="local")
+    cms_username = os.getenv("CMS_USERNAME", default="")
+    cms_password = os.getenv("CMS_PASSWORD")
+    cms_oauth_issuer = os.getenv(
+        "CMS_OAUTH_ISSUER", default="https://ory.login.kiwix.org"
+    )
+    cms_oauth_client_id = os.getenv("CMS_OAUTH_CLIENT_ID", default="")
+    cms_oauth_client_secret = os.getenv("CMS_OAUTH_CLIENT_SECRET", default="")
+    cms_oauth_audience_id = os.getenv("CMS_OAUTH_AUDIENCE_ID", default="")
+    cms_token_renewal_window = datetime.timedelta(
+        seconds=parse_timespan(os.getenv("CMS_TOKEN_RENEWAL_WINDOW", default="5m"))
+    )
     cms_database_url = get_mandatory_env("CMS_DATABASE_URL")
     catalog_generation_timeout = parse_timespan(
         os.getenv("CATALOG_GENERATION_TIMEOUT", default="10s")

healthcheck/src/healthcheck/status/auth.py

@@ -1,9 +1,13 @@
 import datetime
+from http import HTTPStatus
+from typing import cast
 
+from aiohttp.helpers import BasicAuth
 from pydantic import BaseModel
 
 from healthcheck.context import Context
 from healthcheck.status import Result
+from healthcheck.status import status_logger as logger
 from healthcheck.status.requests import query_api
 
 
@@ -16,16 +20,144 @@ class Token(BaseModel):
     token_type: str = "Bearer"
 
 
+def getnow():
+    """naive UTC now"""
+    return datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+
+
+class ClientTokenProvider:
+    """Client to generate access tokens to authenticate with Zimfarm"""
+
+    def __init__(self):
+        self._access_token: str | None = None
+        self._refresh_token: str | None = None
+        self._expires_at: datetime.datetime = datetime.datetime.fromtimestamp(
+            0
+        ).replace(tzinfo=None)
+
+    async def _generate_oauth_access_token(self) -> None:
+        """Generate oauth access token and update expires_at."""
+
+        response = await query_api(
+            f"{Context.cms_oauth_issuer}/oauth2/token",
+            method="POST",
+            data={
+                "grant_type": "client_credentials",
+                "audience": Context.cms_oauth_audience_id,
+            },
+            auth=BasicAuth(
+                Context.cms_oauth_client_id, Context.cms_oauth_client_secret
+            ),
+            timeout=Context.requests_timeout,
+            check_name="zimfarm-api-authentication",
+        )
+        if response.json:
+            self._access_token = cast(str, response.json["access_token"])
+            self._expires_at = getnow() + datetime.timedelta(
+                seconds=response.json["expires_in"]
+            )
+
+    async def _generate_local_access_token(self) -> None:
+        check_name = "zimfarm-api-authentication"
+        if self._refresh_token:
+            response = await query_api(
+                f"{Context.cms_api_url}/auth/refresh",
+                method="POST",
+                payload={
+                    "refresh_token": self._refresh_token,
+                },
+                timeout=Context.requests_timeout,
+                check_name=check_name,
+            )
+        else:
+            response = await query_api(
+                f"{Context.cms_api_url}/auth/authorize",
+                method="POST",
+                payload={
+                    "username": Context.cms_username,
+                    "password": Context.cms_password,
+                },
+                timeout=Context.requests_timeout,
+                check_name=check_name,
+            )
+
+        if response.json:
+            self._access_token = cast(str, response.json["access_token"])
+            self._refresh_token = cast(str, response.json["refresh_token"])
+            self._expires_at = datetime.datetime.fromisoformat(
+                response.json["expires_time"]
+            ).replace(tzinfo=None)
+
+    async def get_access_token(self, *, force_refresh: bool = False) -> str:
+        """Retrieve or generate access token depending on if token has expired."""
+        now = getnow()
+        if (
+            force_refresh
+            or self._access_token is None
+            or now >= (self._expires_at - Context.cms_token_renewal_window)
+        ):
+            if Context.auth_mode == "oauth":
+                await self._generate_oauth_access_token()
+            elif Context.auth_mode == "local":
+                await self._generate_local_access_token()
+            else:
+                raise ValueError(
+                    f"Unknown authentication mode: {Context.auth_mode}. "
+                    "Allowed values are: 'local', 'oauth'"
+                )
+        if self._access_token is None:
+            raise ValueError("Failed to generate access token.")
+        return self._access_token
+
+    @property
+    def expires_at(self) -> datetime.datetime:
+        return self._expires_at
+
+    @property
+    def refresh_token(self) -> str | None:
+        return self._refresh_token
+
+
+_token_provider = ClientTokenProvider()
+
+
 async def authenticate() -> Result[Token]:
-    """Check if authentication is sucessful with CMS"""
-    response = await query_api(
-        f"{Context.cms_api_url}/auth/authorize",
-        method="POST",
-        payload={"username": Context.cms_username, "password": Context.cms_password},
-        check_name="cms-api-authentication",
-    )
-    return Result(
-        success=response.success,
-        status_code=response.status_code,
-        data=Token.model_validate(response.json) if response.success else None,
-    )
+    """Check if authentication is successful with CMS API"""
+    try:
+        access_token = await _token_provider.get_access_token()
+        token = Token(
+            access_token=access_token,
+            expires_time=_token_provider.expires_at,
+            refresh_token=_token_provider.refresh_token or "",
+        )
+
+        response = await query_api(
+            f"{Context.cms_api_url}/auth/me",
+            method="GET",
+            headers={"Authorization": f"Bearer {token.access_token}"},
+            check_name="zimfarm-api-authentication",
+        )
+
+        if response.success:
+            logger.debug(
+                f"Authentication successful using {Context.auth_mode} mode",
+                extra={"checkname": "zimfarm-api-authentication"},
+            )
+
+            return Result(
+                success=True,
+                status_code=HTTPStatus.OK,
+                data=token,
+            )
+        else:
+            return Result(
+                success=False,
+                status_code=HTTPStatus.UNAUTHORIZED,
+                data=None,
+            )
+    except Exception:
+        return Result(
+            success=False,
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+            data=None,
+        )

healthcheck/src/healthcheck/status/collections.py

@@ -78,7 +78,8 @@ async def check_catalog_generation() -> Result[CatalogStatus]:
         if failures:
             logger.error(
                 f"Failed to generate catalogs for the following collections: "
-                f"{','.join([failure.name for failure in failures])}"
+                f"{','.join([failure.name for failure in failures])}",
+                extra={"checkname": "cms-check-catalog-generation"},
             )
 
         return Result(

healthcheck/src/healthcheck/status/requests.py

@@ -4,6 +4,7 @@
 from typing import Any
 
 import aiohttp
+from aiohttp.helpers import BasicAuth
 
 from healthcheck.context import Context
 from healthcheck.status import status_logger as logger
@@ -28,14 +29,17 @@ async def query_api(
     payload: dict[str, Any] | None = None,
     params: dict[str, Any] | None = None,
     timeout: float = Context.requests_timeout,
+    auth: BasicAuth | None = None,
+    data: dict[str, Any] | None = None,
 ) -> Response:
     req_headers: dict[str, Any] = {}
     req_headers.update(headers if headers else {})
 
     # Log request details
     logger.debug(
         f"Sending request: method={method.upper()}, url={url}, "
-        f"headers={req_headers}, params={params}, body={payload}",
+        f"headers={req_headers}, params={params}, body={payload}, "
+        f"data={data}",
         extra={"checkname": check_name},
     )
 
@@ -48,6 +52,8 @@ async def query_api(
             json=payload,
             params=params,
             timeout=aiohttp.ClientTimeout(total=timeout),
+            auth=auth,
+            data=data,
         ) as resp:
             try:
                 text = await resp.text()