Refactor for cleanliness

Author: alukach

runtime/eoapi/stac/eoapi/stac/app.py

@@ -1,11 +1,13 @@
 """FastAPI application using PGStac."""
 
 from contextlib import asynccontextmanager
+from typing import Annotated, Any, Dict
 
+import jwt
 from eoapi.stac.config import ApiSettings, TilesApiSettings
 from eoapi.stac.extension import TiTilerExtension
 from eoapi.stac.extension import extensions_map as PgStacExtensions
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, Security, security, status
 from fastapi.responses import ORJSONResponse
 from stac_fastapi.api.app import StacApi
 from stac_fastapi.api.models import create_get_request_model, create_post_request_model
@@ -19,6 +21,8 @@
 from starlette.templating import Jinja2Templates
 from starlette_cramjam.middleware import CompressionMiddleware
 
+from .auth import KeycloakAuth
+
 try:
     from importlib.resources import files as resources_files  # type: ignore
 except ImportError:
@@ -32,6 +36,13 @@
 tiles_settings = TilesApiSettings()
 settings = Settings()
 
+keycloak = KeycloakAuth(
+    realm="eoapi",
+    client_id="stac-api",
+    host="http://localhost:8080",
+    internal_host="http://keycloak:8080",
+)
+
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -54,7 +65,15 @@ async def lifespan(app: FastAPI):
 GETModel = create_get_request_model(extensions)
 
 api = StacApi(
-    app=FastAPI(title=api_settings.name, lifespan=lifespan),
+    app=FastAPI(
+        title=api_settings.name,
+        lifespan=lifespan,
+        swagger_ui_init_oauth={
+            "appName": "eoAPI",
+            "clientId": keycloak.client_id,
+            "usePkceWithAuthorizationCodeGrant": True,
+        },
+    ),
     title=api_settings.name,
     description=api_settings.name,
     settings=settings,
@@ -84,10 +103,18 @@ async def lifespan(app: FastAPI):
 
 
 @app.get("/index.html", response_class=HTMLResponse)
-async def viewer_page(request: Request):
+async def viewer_page(
+    request: Request, token: Annotated[str, Security(keycloak.scheme)]
+):
     """Search viewer."""
     return templates.TemplateResponse(
         "stac-viewer.html",
         {"request": request, "endpoint": str(request.url).replace("/index.html", "")},
         media_type="text/html",
     )
+
+
+@app.get("/user", tags=["auth"])
+def get_user(user_token: Annotated[Dict[Any, Any], Security(keycloak.user_validator)]):
+    """View auth token."""
+    return user_token

runtime/eoapi/stac/eoapi/stac/auth.py

@@ -0,0 +1,77 @@
+from typing import Annotated, Iterable, Optional
+from dataclasses import dataclass
+from functools import cached_property
+
+import jwt
+from fastapi import FastAPI, HTTPException, Security, security, status
+
+
+@dataclass
+class KeycloakAuth:
+    realm: str
+    host: str
+    client_id: str
+    internal_host: Optional[str] = None
+
+    required_audience: Optional[str | Iterable[str]] = None
+
+    def _build_url(self, host: str):
+        return f"{host}/realms/{self.realm}/protocol/openid-connect"
+
+    @property
+    def user_validator(
+        self,
+    ):
+        def valid_user_token(
+            token_str: Annotated[str, Security(self.scheme)],
+            required_scopes: security.SecurityScopes,
+        ):
+            # Parse & validate token
+            try:
+                token = jwt.decode(
+                    token_str,
+                    self.jwks_client.get_signing_key_from_jwt(token_str).key,
+                    algorithms=["RS256"],
+                    audience=self.required_audience,
+                )
+            except jwt.exceptions.DecodeError as e:
+                raise HTTPException(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    detail="Could not validate credentials",
+                    headers={"WWW-Authenticate": "Bearer"},
+                ) from e
+
+            # Validate scopes (if required)
+            for scope in required_scopes.scopes:
+                if scope not in token["scope"]:
+                    raise HTTPException(
+                        status_code=status.HTTP_401_UNAUTHORIZED,
+                        detail="Not enough permissions",
+                        headers={
+                            "WWW-Authenticate": f'Bearer scope="{required_scopes.scope_str}"'
+                        },
+                    )
+
+            return token
+
+        return valid_user_token
+
+    @property
+    def internal_keycloak_api(self):
+        return self._build_url(self.internal_host or self.host)
+
+    @property
+    def keycloak_api(self):
+        return self._build_url(self.host)
+
+    @property
+    def scheme(self):
+        return security.OAuth2AuthorizationCodeBearer(
+            authorizationUrl=f"{self.keycloak_api}/auth",
+            tokenUrl=f"{self.keycloak_api}/token",
+            scopes={},
+        )
+
+    @cached_property
+    def jwks_client(self):
+        return jwt.PyJWKClient(f"{self.internal_keycloak_api}/certs")