stac-auth-proxy 0.6.1__py3-none-any.whl

stac_auth_proxy/__init__.py

@@ -0,0 +1,12 @@
+"""
+STAC Auth Proxy package.
+
+This package contains the components for the STAC authentication and proxying system.
+It includes FastAPI routes for handling authentication, authorization, and interaction
+with some internal STAC API.
+"""
+
+from .app import create_app
+from .config import Settings
+
+__all__ = ["create_app", "Settings"]

stac_auth_proxy/__main__.py

@@ -0,0 +1,18 @@
+"""Entry point for running the module without customized code."""
+
+import uvicorn
+from uvicorn.config import LOGGING_CONFIG
+
+LOGGING_CONFIG["loggers"][__package__] = {
+    "level": "DEBUG",
+    "handlers": ["default"],
+}
+
+uvicorn.run(
+    f"{__package__}.app:create_app",
+    host="0.0.0.0",
+    port=8000,
+    log_config=LOGGING_CONFIG,
+    reload=True,
+    factory=True,
+)

stac_auth_proxy/app.py

@@ -0,0 +1,178 @@
+"""
+STAC Auth Proxy API.
+
+This module defines the FastAPI application for the STAC Auth Proxy, which handles
+authentication, authorization, and proxying of requests to some internal STAC API.
+"""
+
+import logging
+from contextlib import asynccontextmanager
+from typing import Optional
+
+from fastapi import FastAPI
+from starlette_cramjam.middleware import CompressionMiddleware
+
+from .config import Settings
+from .handlers import HealthzHandler, ReverseProxyHandler, SwaggerUI
+from .middleware import (
+    AddProcessTimeHeaderMiddleware,
+    ApplyCql2FilterMiddleware,
+    AuthenticationExtensionMiddleware,
+    BuildCql2FilterMiddleware,
+    EnforceAuthMiddleware,
+    OpenApiMiddleware,
+    ProcessLinksMiddleware,
+    RemoveRootPathMiddleware,
+)
+from .utils.lifespan import check_conformance, check_server_health
+
+logger = logging.getLogger(__name__)
+
+
+def create_app(settings: Optional[Settings] = None) -> FastAPI:
+    """FastAPI Application Factory."""
+    settings = settings or Settings()
+
+    #
+    # Application
+    #
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        assert settings
+
+        # Wait for upstream servers to become available
+        if settings.wait_for_upstream:
+            logger.info("Running upstream server health checks...")
+            urls = [settings.upstream_url, settings.oidc_discovery_internal_url]
+            for url in urls:
+                await check_server_health(url=url)
+            logger.info(
+                "Upstream servers are healthy:\n%s",
+                "\n".join([f" - {url}" for url in urls]),
+            )
+
+        # Log all middleware connected to the app
+        logger.info(
+            "Connected middleware:\n%s",
+            "\n".join([f" - {m.cls.__name__}" for m in app.user_middleware]),
+        )
+
+        if settings.check_conformance:
+            await check_conformance(
+                app.user_middleware,
+                str(settings.upstream_url),
+            )
+
+        yield
+
+    app = FastAPI(
+        openapi_url=None,  # Disable OpenAPI schema endpoint, we want to serve upstream's schema
+        lifespan=lifespan,
+        root_path=settings.root_path,
+    )
+    if app.root_path:
+        logger.debug("Mounted app at %s", app.root_path)
+
+    #
+    # Handlers (place catch-all proxy handler last)
+    #
+
+    if settings.swagger_ui_endpoint:
+        assert (
+            settings.openapi_spec_endpoint
+        ), "openapi_spec_endpoint must be set when using swagger_ui_endpoint"
+        app.add_route(
+            settings.swagger_ui_endpoint,
+            SwaggerUI(
+                openapi_url=settings.openapi_spec_endpoint,
+                init_oauth=settings.swagger_ui_init_oauth,
+            ).route,
+            include_in_schema=False,
+        )
+    if settings.healthz_prefix:
+        app.include_router(
+            HealthzHandler(upstream_url=str(settings.upstream_url)).router,
+            prefix=settings.healthz_prefix,
+        )
+
+    app.add_api_route(
+        "/{path:path}",
+        ReverseProxyHandler(
+            upstream=str(settings.upstream_url),
+            override_host=settings.override_host,
+        ).proxy_request,
+        methods=["GET", "POST", "PUT", "PATCH", "DELETE"],
+    )
+
+    #
+    # Middleware (order is important, last added = first to run)
+    #
+
+    if settings.enable_authentication_extension:
+        app.add_middleware(
+            AuthenticationExtensionMiddleware,
+            default_public=settings.default_public,
+            public_endpoints=settings.public_endpoints,
+            private_endpoints=settings.private_endpoints,
+            oidc_discovery_url=str(settings.oidc_discovery_url),
+        )
+
+    if settings.openapi_spec_endpoint:
+        app.add_middleware(
+            OpenApiMiddleware,
+            openapi_spec_path=settings.openapi_spec_endpoint,
+            oidc_discovery_url=str(settings.oidc_discovery_url),
+            public_endpoints=settings.public_endpoints,
+            private_endpoints=settings.private_endpoints,
+            default_public=settings.default_public,
+            root_path=settings.root_path,
+            auth_scheme_name=settings.openapi_auth_scheme_name,
+            auth_scheme_override=settings.openapi_auth_scheme_override,
+        )
+
+    if settings.items_filter or settings.collections_filter:
+        app.add_middleware(
+            ApplyCql2FilterMiddleware,
+        )
+        app.add_middleware(
+            BuildCql2FilterMiddleware,
+            items_filter=settings.items_filter() if settings.items_filter else None,
+            collections_filter=(
+                settings.collections_filter() if settings.collections_filter else None
+            ),
+            collections_filter_path=settings.collections_filter_path,
+            items_filter_path=settings.items_filter_path,
+        )
+
+    app.add_middleware(
+        AddProcessTimeHeaderMiddleware,
+    )
+
+    app.add_middleware(
+        EnforceAuthMiddleware,
+        public_endpoints=settings.public_endpoints,
+        private_endpoints=settings.private_endpoints,
+        default_public=settings.default_public,
+        oidc_discovery_url=settings.oidc_discovery_internal_url,
+    )
+
+    if settings.root_path or settings.upstream_url.path != "/":
+        app.add_middleware(
+            ProcessLinksMiddleware,
+            upstream_url=str(settings.upstream_url),
+            root_path=settings.root_path,
+        )
+
+    if settings.root_path:
+        app.add_middleware(
+            RemoveRootPathMiddleware,
+            root_path=settings.root_path,
+        )
+
+    if settings.enable_compression:
+        app.add_middleware(
+            CompressionMiddleware,
+        )
+
+    return app

stac_auth_proxy/config.py

@@ -0,0 +1,91 @@
+"""Configuration for the STAC Auth Proxy."""
+
+import importlib
+from typing import Any, Literal, Optional, Sequence, TypeAlias, Union
+
+from pydantic import BaseModel, Field, model_validator
+from pydantic.networks import HttpUrl
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+METHODS = Literal["GET", "POST", "PUT", "DELETE", "PATCH"]
+EndpointMethodsNoScope: TypeAlias = dict[str, Sequence[METHODS]]
+EndpointMethods: TypeAlias = dict[str, Sequence[Union[METHODS, tuple[METHODS, str]]]]
+
+_PREFIX_PATTERN = r"^/.*$"
+
+
+class ClassInput(BaseModel):
+    """Input model for dynamically loading a class or function."""
+
+    cls: str
+    args: Sequence[str] = Field(default_factory=list)
+    kwargs: dict[str, str] = Field(default_factory=dict)
+
+    def __call__(self):
+        """Dynamically load a class and instantiate it with args & kwargs."""
+        assert self.cls.count(":")
+        module_path, class_name = self.cls.rsplit(":", 1)
+        module = importlib.import_module(module_path)
+        cls = getattr(module, class_name)
+        return cls(*self.args, **self.kwargs)
+
+
+class Settings(BaseSettings):
+    """Configuration settings for the STAC Auth Proxy."""
+
+    # External URLs
+    upstream_url: HttpUrl
+    oidc_discovery_url: HttpUrl
+    oidc_discovery_internal_url: HttpUrl
+
+    root_path: str = ""
+    override_host: bool = True
+    healthz_prefix: str = Field(pattern=_PREFIX_PATTERN, default="/healthz")
+    wait_for_upstream: bool = True
+    check_conformance: bool = True
+    enable_compression: bool = True
+
+    # OpenAPI / Swagger UI
+    openapi_spec_endpoint: Optional[str] = Field(pattern=_PREFIX_PATTERN, default=None)
+    openapi_auth_scheme_name: str = "oidcAuth"
+    openapi_auth_scheme_override: Optional[dict] = None
+    swagger_ui_endpoint: Optional[str] = None
+    swagger_ui_init_oauth: dict = Field(default_factory=dict)
+
+    # Auth
+    enable_authentication_extension: bool = True
+    default_public: bool = False
+    public_endpoints: EndpointMethodsNoScope = {
+        r"^/api.html$": ["GET"],
+        r"^/api$": ["GET"],
+        r"^/docs/oauth2-redirect": ["GET"],
+        r"^/healthz": ["GET"],
+    }
+    private_endpoints: EndpointMethods = {
+        # https://github.com/stac-api-extensions/collection-transaction/blob/v1.0.0-beta.1/README.md#methods
+        r"^/collections$": ["POST"],
+        r"^/collections/([^/]+)$": ["PUT", "PATCH", "DELETE"],
+        # https://github.com/stac-api-extensions/transaction/blob/v1.0.0-rc.3/README.md#methods
+        r"^/collections/([^/]+)/items$": ["POST"],
+        r"^/collections/([^/]+)/items/([^/]+)$": ["PUT", "PATCH", "DELETE"],
+        # https://stac-utils.github.io/stac-fastapi/api/stac_fastapi/extensions/third_party/bulk_transactions/#bulktransactionextension
+        r"^/collections/([^/]+)/bulk_items$": ["POST"],
+    }
+
+    # Filters
+    items_filter: Optional[ClassInput] = None
+    items_filter_path: str = r"^(/collections/([^/]+)/items(/[^/]+)?$|/search$)"
+    collections_filter: Optional[ClassInput] = None
+    collections_filter_path: str = r"^/collections(/[^/]+)?$"
+
+    model_config = SettingsConfigDict(
+        env_nested_delimiter="_",
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def default_oidc_discovery_internal_url(cls, data: Any) -> Any:
+        """Set the internal OIDC discovery URL to the public URL if not set."""
+        if not data.get("oidc_discovery_internal_url"):
+            data["oidc_discovery_internal_url"] = data.get("oidc_discovery_url")
+        return data

stac_auth_proxy/filters/__init__.py

@@ -0,0 +1,9 @@
+"""CQL2 filter generators."""
+
+from .opa import Opa
+from .template import Template
+
+__all__ = [
+    "Opa",
+    "Template",
+]

stac_auth_proxy/filters/opa.py

@@ -0,0 +1,44 @@
+"""Integration with Open Policy Agent (OPA) to generate CQL2 filters for requests to a STAC API."""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import httpx
+
+from ..utils.cache import MemoryCache, get_value_by_path
+
+
+@dataclass
+class Opa:
+    """Call Open Policy Agent (OPA) to generate CQL2 filters from request context."""
+
+    host: str
+    decision: str
+
+    client: httpx.AsyncClient = field(init=False)
+    cache: MemoryCache = field(init=False)
+    cache_key: str = "req.headers.authorization"
+    cache_ttl: float = 5.0
+
+    def __post_init__(self):
+        """Initialize the client."""
+        self.client = httpx.AsyncClient(base_url=self.host)
+        self.cache = MemoryCache(ttl=self.cache_ttl)
+
+    async def __call__(self, context: dict[str, Any]) -> str:
+        """Generate a CQL2 filter for the request."""
+        token = get_value_by_path(context, self.cache_key)
+        try:
+            expr_str = self.cache[token]
+        except KeyError:
+            expr_str = await self._fetch(context)
+            self.cache[token] = expr_str
+        return expr_str
+
+    async def _fetch(self, context: dict[str, Any]) -> str:
+        """Fetch the CQL2 filter from OPA."""
+        response = await self.client.post(
+            f"/v1/data/{self.decision}",
+            json={"input": context},
+        )
+        return response.raise_for_status().json()["result"]

stac_auth_proxy/filters/template.py

@@ -0,0 +1,22 @@
+"""Generate CQL2 filter expressions via Jinja2 templating."""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from jinja2 import BaseLoader, Environment
+
+
+@dataclass
+class Template:
+    """Generate CQL2 filter expressions via Jinja2 templating."""
+
+    template_str: str
+    env: Environment = field(init=False)
+
+    def __post_init__(self):
+        """Initialize the Jinja2 environment."""
+        self.env = Environment(loader=BaseLoader).from_string(self.template_str)
+
+    async def __call__(self, context: dict[str, Any]) -> str:
+        """Render a CQL2 filter expression with the request and auth token."""
+        return self.env.render(**context).strip()

stac_auth_proxy/handlers/__init__.py

@@ -0,0 +1,7 @@
+"""Handlers to process requests."""
+
+from .healthz import HealthzHandler
+from .reverse_proxy import ReverseProxyHandler
+from .swagger_ui import SwaggerUI
+
+__all__ = ["ReverseProxyHandler", "HealthzHandler", "SwaggerUI"]

stac_auth_proxy/handlers/healthz.py

@@ -0,0 +1,31 @@
+"""Health check endpoints."""
+
+from dataclasses import dataclass, field
+
+from fastapi import APIRouter
+from httpx import AsyncClient
+
+
+@dataclass
+class HealthzHandler:
+    """Handler for health check endpoints."""
+
+    upstream_url: str
+    router: APIRouter = field(init=False)
+
+    def __post_init__(self):
+        """Initialize the router."""
+        self.router = APIRouter()
+        self.router.add_api_route("", self.healthz, methods=["GET"])
+        self.router.add_api_route("/upstream", self.healthz_upstream, methods=["GET"])
+
+    async def healthz(self):
+        """Return health of this API."""
+        return {"status": "ok"}
+
+    async def healthz_upstream(self):
+        """Return health of upstream STAC API."""
+        async with AsyncClient() as client:
+            response = await client.get(self.upstream_url)
+            response.raise_for_status()
+            return {"status": "ok", "code": response.status_code}

stac_auth_proxy/handlers/reverse_proxy.py

@@ -0,0 +1,101 @@
+"""Tooling to manage the reverse proxying of requests to an upstream STAC API."""
+
+import logging
+import time
+from dataclasses import dataclass, field
+
+import httpx
+from fastapi import Request
+from starlette.datastructures import MutableHeaders
+from starlette.responses import Response
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ReverseProxyHandler:
+    """Reverse proxy functionality."""
+
+    upstream: str
+    client: httpx.AsyncClient = None
+    timeout: httpx.Timeout = field(default_factory=lambda: httpx.Timeout(timeout=15.0))
+
+    proxy_name: str = "stac-auth-proxy"
+    override_host: bool = True
+    legacy_forwarded_headers: bool = False
+
+    def __post_init__(self):
+        """Initialize the HTTP client."""
+        self.client = self.client or httpx.AsyncClient(
+            base_url=self.upstream,
+            timeout=self.timeout,
+            http2=True,
+        )
+
+    def _prepare_headers(self, request: Request) -> MutableHeaders:
+        """Prepare headers for the proxied request."""
+        headers = MutableHeaders(request.headers)
+        headers.setdefault("Via", f"1.1 {self.proxy_name}")
+
+        proxy_client = request.client.host if request.client else "unknown"
+        proxy_proto = request.url.scheme
+        proxy_host = request.url.netloc
+        proxy_path = request.base_url.path
+        headers.setdefault(
+            "Forwarded",
+            f"for={proxy_client};host={proxy_host};proto={proxy_proto};path={proxy_path}",
+        )
+        if self.legacy_forwarded_headers:
+            headers.setdefault("X-Forwarded-For", proxy_client)
+            headers.setdefault("X-Forwarded-Host", proxy_host)
+            headers.setdefault("X-Forwarded-Path", proxy_path)
+            headers.setdefault("X-Forwarded-Proto", proxy_proto)
+
+        # Set host to the upstream host
+        if self.override_host:
+            headers["Host"] = self.client.base_url.netloc.decode("utf-8")
+
+        return headers
+
+    async def proxy_request(self, request: Request) -> Response:
+        """Proxy a request to the upstream STAC API."""
+        headers = self._prepare_headers(request)
+
+        # https://github.com/fastapi/fastapi/discussions/7382#discussioncomment-5136466
+        rp_req = self.client.build_request(
+            request.method,
+            url=httpx.URL(
+                path=request.url.path,
+                query=request.url.query.encode("utf-8"),
+            ),
+            headers=headers,
+            content=request.stream(),
+        )
+
+        # NOTE: HTTPX adds headers, so we need to trim them before sending request
+        for h in rp_req.headers:
+            if h not in headers:
+                del rp_req.headers[h]
+
+        logger.debug(f"Proxying request to {rp_req.url}")
+
+        start_time = time.perf_counter()
+        rp_resp = await self.client.send(rp_req, stream=True)
+        proxy_time = time.perf_counter() - start_time
+
+        logger.debug(
+            f"Received response status {rp_resp.status_code!r} from {rp_req.url} in {proxy_time:.3f}s"
+        )
+        rp_resp.headers["X-Upstream-Time"] = f"{proxy_time:.3f}"
+
+        # We read the content here to make use of HTTPX's decompression, ensuring we have
+        # non-compressed content for the middleware to work with.
+        content = await rp_resp.aread()
+        if rp_resp.headers.get("Content-Encoding"):
+            del rp_resp.headers["Content-Encoding"]
+
+        return Response(
+            content=content,
+            status_code=rp_resp.status_code,
+            headers=dict(rp_resp.headers),
+        )

stac_auth_proxy/handlers/swagger_ui.py

@@ -0,0 +1,41 @@
+"""
+In order to allow customization fo the Swagger UI's OAuth2 configuration, we support
+overriding the default handler. This is useful for adding custom parameters such as
+`usePkceWithAuthorizationCodeGrant` or `clientId`.
+
+See:
+- https://swagger.io/docs/open-source-tools/swagger-ui/usage/oauth2/
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+from fastapi.openapi.docs import get_swagger_ui_html
+from starlette.requests import Request
+from starlette.responses import HTMLResponse
+
+
+@dataclass
+class SwaggerUI:
+    """Swagger UI handler."""
+
+    openapi_url: str
+    title: Optional[str] = "STAC API"
+    init_oauth: dict = field(default_factory=dict)
+    parameters: dict = field(default_factory=dict)
+    oauth2_redirect_url: str = "/docs/oauth2-redirect"
+
+    async def route(self, req: Request) -> HTMLResponse:
+        """Route handler."""
+        root_path = req.scope.get("root_path", "").rstrip("/")
+        openapi_url = root_path + self.openapi_url
+        oauth2_redirect_url = self.oauth2_redirect_url
+        if oauth2_redirect_url:
+            oauth2_redirect_url = root_path + oauth2_redirect_url
+        return get_swagger_ui_html(
+            openapi_url=openapi_url,
+            title=f"{self.title} - Swagger UI",
+            oauth2_redirect_url=oauth2_redirect_url,
+            init_oauth=self.init_oauth,
+            swagger_ui_parameters=self.parameters,
+        )

stac_auth_proxy/middleware/AddProcessTimeHeaderMiddleware.py

@@ -0,0 +1,18 @@
+"""Middleware to add a header with the process time to the response."""
+
+import time
+
+from fastapi import Request, Response
+from starlette.middleware.base import BaseHTTPMiddleware
+
+
+class AddProcessTimeHeaderMiddleware(BaseHTTPMiddleware):
+    """Middleware to add a header with the process time to the response."""
+
+    async def dispatch(self, request: Request, call_next) -> Response:
+        """Add a header with the process time to the response."""
+        start_time = time.perf_counter()
+        response = await call_next(request)
+        process_time = time.perf_counter() - start_time
+        response.headers["X-Process-Time"] = f"{process_time:.3f}"
+        return response