Merge pull request #221 from Point72/tkp/scope

Allow scoping authentication middlewares, allow multiple middlewares

Author: timkpaine

.vscode/settings.json

@@ -1,3 +1,7 @@
 {
-    "eslint.workingDirectories": ["./js"]
+    "eslint.workingDirectories": [
+        "./js"
+    ],
+    "python-envs.defaultEnvManager": "ms-python.python:system",
+    "python-envs.pythonProjects": []
 }

csp_gateway/server/gateway/csp/module.py

@@ -1,7 +1,6 @@
-import abc
-import typing
+from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Any, Dict, Generic, List, Optional, Set, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, Generic, List, Optional, Set, Type, Union
 
 from ccflow import BaseModel
 from pydantic import Field, TypeAdapter, model_validator
@@ -11,11 +10,11 @@
 
 from .channels import ChannelsType
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from csp_gateway.server import GatewaySettings, GatewayWebApp
 
 
-class Module(BaseModel, Generic[ChannelsType], abc.ABC):
+class Module(BaseModel, Generic[ChannelsType], ABC):
     model_config = {"arbitrary_types_allowed": True}
 
     requires: Optional[ChannelSelection] = None
@@ -28,14 +27,14 @@ class Module(BaseModel, Generic[ChannelsType], abc.ABC):
         """,
     )
 
-    @abc.abstractmethod
+    @abstractmethod
     def connect(self, Channels: ChannelsType) -> None: ...
 
     def rest(self, app: "GatewayWebApp") -> None: ...
 
     def info(self, settings: "GatewaySettings") -> Optional[str]: ...
 
-    @abc.abstractmethod
+    @abstractmethod
     def shutdown(self) -> None: ...
 
     def dynamic_keys(self) -> Optional[Dict[str, List[Any]]]: ...

csp_gateway/server/middleware/api_key.py

@@ -37,22 +37,18 @@ def info(self, settings: GatewaySettings) -> str:
             return f"\tUI: {url}?token={self.api_key}"
         return f"\tAPI: {url}/openapi.json?token={self.api_key}"
 
-    def rest(self, app: GatewayWebApp) -> None:
-        # reinitialize header
-        api_key_query = APIKeyQuery(name=self.api_key_name, auto_error=False)
-        api_key_header = APIKeyHeader(name=self.api_key_name, auto_error=False)
-        api_key_cookie = APIKeyCookie(name=self.api_key_name, auto_error=False)
-
-        # routers
-        auth_router: APIRouter = app.get_router("auth")
-
-        # now mount middleware
-        async def get_api_key(
-            api_key_query: str = Security(api_key_query),
-            api_key_header: str = Security(api_key_header),
-            api_key_cookie: str = Security(api_key_cookie),
-        ):
-            # Support both single string and list of valid API keys
+    def validate(self):
+        """Return a FastAPI dependency function for API key validation."""
+        api_key_query_security = Security(APIKeyQuery(name=self.api_key_name, auto_error=False))
+        api_key_header_security = Security(APIKeyHeader(name=self.api_key_name, auto_error=False))
+        api_key_cookie_security = Security(APIKeyCookie(name=self.api_key_name, auto_error=False))
+
+        async def validate_credentials(
+            api_key_query: str = api_key_query_security,
+            api_key_header: str = api_key_header_security,
+            api_key_cookie: str = api_key_cookie_security,
+        ) -> str:
+            """Validate API key from query, header, or cookie."""
             valid_keys = self.api_key if isinstance(self.api_key, list) else [self.api_key]
             for provided_key in (api_key_query, api_key_header, api_key_cookie):
                 if provided_key in valid_keys:
@@ -62,8 +58,15 @@ async def get_api_key(
                 detail=self.unauthorized_status_message,
             )
 
+        return validate_credentials
+
+    def rest(self, app: GatewayWebApp) -> None:
+        # routers
+        auth_router: APIRouter = app.get_router("auth")
+        check = self.get_check_dependency()
+
         @auth_router.get("/login")
-        async def route_login_and_add_cookie(api_key: str = Depends(get_api_key)):
+        async def route_login_and_add_cookie(api_key: str = Depends(check)):
             response = RedirectResponse(url="/")
             response.set_cookie(
                 self.api_key_name,
@@ -81,10 +84,10 @@ async def route_logout_and_remove_cookie():
             response.delete_cookie(self.api_key_name, domain=self.domain)
             return response
 
-        self._setup_public_routes(app, get_api_key)
+        self._setup_public_routes(app)
 
-    def _setup_public_routes(self, app: GatewayWebApp, get_api_key) -> None:
-        """Setup public routes, middleware, and exception handler. Shared by subclasses."""
+    def _setup_public_routes(self, app: GatewayWebApp) -> None:
+        """Setup public routes, middleware, and exception handler. Shared by subclasses.""" ""
         public_router: APIRouter = app.get_router("public")
 
         @public_router.get("/login", response_class=HTMLResponse, include_in_schema=False)
@@ -102,7 +105,7 @@ async def get_logout_page(request: Request = None):
             return app.templates.TemplateResponse("logout.html.j2", {"request": request})
 
         # add auth to all other routes
-        app.add_middleware(Depends(get_api_key))
+        app.add_middleware(Depends(self.get_check_dependency()))
 
         @app.app.exception_handler(403)
         async def custom_403_handler(request: Request = None, *args):

csp_gateway/server/middleware/api_key_external.py

@@ -40,24 +40,24 @@ def _invoke_external(self, api_key: str, settings: GatewaySettings, module=None)
             return None
         return self.external_validator.object(api_key, settings, module)
 
-    def _get_api_key_dependency(self, app: GatewayWebApp):
-        """Returns the get_api_key dependency for external validation."""
-        api_key_query = APIKeyQuery(name=self.api_key_name, auto_error=False)
-        api_key_header = APIKeyHeader(name=self.api_key_name, auto_error=False)
-        api_key_cookie = APIKeyCookie(name=self.api_key_name, auto_error=False)
-
-        async def get_api_key(
-            api_key_query: str = Security(api_key_query),
-            api_key_header: str = Security(api_key_header),
-            api_key_cookie: str = Security(api_key_cookie),
-        ):
+    def validate(self):
+        """Return a FastAPI dependency function for external API key validation."""
+        api_key_query_security = Security(APIKeyQuery(name=self.api_key_name, auto_error=False))
+        api_key_header_security = Security(APIKeyHeader(name=self.api_key_name, auto_error=False))
+        api_key_cookie_security = Security(APIKeyCookie(name=self.api_key_name, auto_error=False))
+
+        async def validate_credentials(
+            api_key_query: str = api_key_query_security,
+            api_key_header: str = api_key_header_security,
+            api_key_cookie: str = api_key_cookie_security,
+        ) -> str:
+            """Validate API key using external validator and return a session UUID."""
             try:
                 for provided_key in (api_key_query, api_key_header, api_key_cookie):
-                    identity = self._invoke_external(provided_key, app.settings, app)
+                    identity = self._invoke_external(provided_key, self._app_settings, self._app_module)
                     if identity and isinstance(identity, dict):
                         user_uuid = str(uuid4())
                         while user_uuid in self._identity_store:
-                            # Should never happen, but just in case of a uuid collision, generate a new one
                             user_uuid = str(uuid4())
                         self._identity_store[user_uuid] = identity
                         return user_uuid
@@ -66,20 +66,23 @@ async def get_api_key(
                     status_code=HTTP_403_FORBIDDEN,
                     detail=self.unauthorized_status_message,
                 ) from e
-            # No valid key found
             raise HTTPException(
                 status_code=HTTP_403_FORBIDDEN,
                 detail=self.unauthorized_status_message,
             )
 
-        return get_api_key
+        return validate_credentials
 
     def rest(self, app: GatewayWebApp) -> None:
+        # Store app references for use in check()
+        self._app_settings = app.settings
+        self._app_module = app
+
         auth_router: APIRouter = app.get_router("auth")
-        get_api_key = self._get_api_key_dependency(app)
+        check = self.get_check_dependency()
 
         @auth_router.get("/login")
-        async def route_login_and_add_cookie(api_key: str = Depends(get_api_key)):
+        async def route_login_and_add_cookie(api_key: str = Depends(check)):
             response = RedirectResponse(url="/")
             if api_key in self._identity_store:
                 response.set_cookie(
@@ -102,4 +105,4 @@ async def route_logout_and_remove_cookie(request: Request = None):
             return response
 
         # Call parent to set up public routes, middleware, and exception handler
-        self._setup_public_routes(app, get_api_key)
+        self._setup_public_routes(app)

csp_gateway/server/middleware/base.py

@@ -1,9 +1,61 @@
+from fnmatch import fnmatch
+from typing import Callable, List, Optional, Union
+
+from ccflow import PyObjectPath
+from fastapi import Request
+from starlette.middleware.base import RequestResponseEndpoint
+from starlette.responses import Response
+
 from csp_gateway.server import GatewayChannels, GatewayModule
 
 __all__ = ("AuthenticationMiddleware",)
 
 
 class AuthenticationMiddleware(GatewayModule):
+    scope: Optional[Union[str, List[str]]] = "*"
+    check: Optional[Union[PyObjectPath, Callable]] = None
+
+    def get_check_callable(self) -> Optional[Callable]:
+        """Return the check callable from PyObjectPath or direct callable."""
+        if self.check is None:
+            return None
+        return self.check if callable(self.check) else self.check.object
+
+    def _matches_scope(self, path: str) -> bool:
+        """Check if path matches any of the scope glob patterns."""
+        if self.scope is None:
+            return True
+        patterns = self.scope if isinstance(self.scope, list) else [self.scope]
+        return any(fnmatch(path, pattern) for pattern in patterns)
+
+    def validate(self) -> Callable:
+        """Return a FastAPI dependency function for credential validation.
+
+        Subclasses must implement this method. The returned function should:
+        - Accept credentials (extracted via Security dependencies)
+        - Return a validated identity/token on success
+        - Raise HTTPException on failure
+
+        Note: Scope checking via _matches_scope() is available but not automatically
+        applied in validate() due to WebSocket route compatibility constraints.
+        """
+        raise NotImplementedError("Subclasses must implement validate()")
+
+    def _skip_if_out_of_scope(self, request: Request) -> bool:
+        """Check if request is out of scope. Returns True if should skip auth."""
+        return not self._matches_scope(request.url.path)
+
+    def get_check_dependency(self) -> Callable:
+        """Return the validate() dependency. Scope checking is handled in validate()."""
+        return self.validate()
+
+    async def check_scope(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        """Check that request path is valid in the scope/s. Returns True if in scope."""
+        if self._matches_scope(request.url.path):
+            return await call_next(request)
+        # Path not in scope, skip authentication middleware
+        return await call_next(request)
+
     def connect(self, channels: GatewayChannels) -> None:
         # NO-OP
         ...