squirrels 0.5.0b3__py3-none-any.whl → 0.6.0.post0__py3-none-any.whl

squirrels/__init__.py

@@ -2,6 +2,8 @@ from ._version import __version__
 
 from .arguments import *
 
+from .auth import *
+
 from .connections import *
 
 from .parameter_options import *
@@ -15,3 +17,5 @@ from .dashboards import *
 from .types import *
 
 from ._project import SquirrelsProject
+
+__all__ = ["SquirrelsProject"]

squirrels/_api_routes/__init__.py

@@ -0,0 +1,5 @@
+"""
+API Routes Package
+
+This package contains modular route definitions for the Squirrels API server.
+""" 

squirrels/_api_routes/auth.py

@@ -0,0 +1,337 @@
+"""
+Authentication and user management routes
+"""
+from datetime import datetime, timezone
+import secrets
+from typing import Annotated, Literal
+from urllib.parse import quote
+from fastapi import FastAPI, Depends, Request, Response, Form, APIRouter
+from fastapi.responses import RedirectResponse, HTMLResponse
+from fastapi.security import HTTPBearer
+from pydantic import BaseModel, Field
+from authlib.integrations.starlette_client import OAuth
+
+from .. import _utils as u
+from .._schemas import response_models as rm
+from .._exceptions import InvalidInputError
+from .._schemas.auth_models import AbstractUser, RegisteredUser, GuestUser, UserFieldsModel, ApiKey
+from .._manifest import AuthStrategy
+from .base import RouteBase
+
+
+class AuthRoutes(RouteBase):
+    """Authentication and user management routes"""
+    
+    def __init__(self, get_bearer_token: HTTPBearer, project, no_cache: bool = False):
+        super().__init__(get_bearer_token, project, no_cache)
+
+    def setup_routes(self, app: FastAPI) -> None:
+        """Setup all authentication routes"""
+
+        auth_path = "/auth"
+        auth_router = APIRouter(prefix=auth_path)
+        user_management_router = APIRouter(prefix=auth_path + "/user-management")
+        
+        auth_strategy = self.manifest_cfg.project_variables.auth_strategy
+        is_external = (auth_strategy == AuthStrategy.EXTERNAL)
+
+        # Get expiry configuration
+        expiry_mins = self.env_vars.auth_token_expire_minutes
+        
+        # Create user models
+        class CustomFieldsModel(self.authenticator.CustomUserFields):
+            pass
+
+        class UpdateUserModel(BaseModel):
+            access_level: Literal["admin", "member"] = Field(description="The access level of the user. Admins have more permissions such as creating and updating users.")
+            custom_fields: CustomFieldsModel = Field(description="User fields that are specific to this Squirrels project")
+
+        class UserInfoModel(UpdateUserModel):
+            username: str
+
+        class AddUserModel(UserInfoModel):
+            password: str
+        
+        class UserSessionModel(BaseModel):
+            user: UserInfoModel
+            session_expiry_timestamp: float | None
+
+        # Setup OAuth2 login providers
+        oauth = OAuth()
+
+        for provider in self.authenticator.auth_providers:
+            oauth.register(
+                name=provider.name,
+                server_metadata_url=provider.provider_configs.server_metadata_url,
+                client_id=provider.provider_configs.client_id,
+                client_secret=provider.provider_configs.client_secret,
+                client_kwargs=provider.provider_configs.client_kwargs
+            )
+
+        # User info endpoint
+        user_session_path = '/user-session'
+
+        @auth_router.get(user_session_path, description="Get the authenticated user's fields", tags=["Authentication"])
+        async def get_user_session(
+            request: Request, user: RegisteredUser | GuestUser = Depends(self.get_current_user)
+        ) -> UserSessionModel:
+            if isinstance(user, GuestUser):
+                raise InvalidInputError(401, "invalid_authorization_token", "Invalid authorization token, no user info found")
+            
+            expiry = request.session.get("access_token_expiry")
+            if expiry is None:
+                expiry = getattr(request.state, "access_token_expiry", None)
+            
+            user_session = UserSessionModel(
+                user=user.model_dump(mode='json'), 
+                session_expiry_timestamp=float(expiry) if expiry is not None else None
+            )
+            return user_session
+
+        # Login endpoint
+        if not is_external:
+            login_path = '/login'
+            
+            @auth_router.post(login_path, tags=["Authentication"], description="Authenticate with username & password. Returns user information if no redirect_url is provided, otherwise redirects to the specified URL.")
+            async def login(
+                request: Request, username: Annotated[str, Form()], password: Annotated[str, Form()]
+            ) -> UserSessionModel:
+                user = self.authenticator.get_user(username, password)
+                
+                access_token, expiry = self.authenticator.create_access_token(user, expiry_minutes=expiry_mins)
+                expiry_timestamp = expiry.timestamp()
+                request.session["access_token"] = access_token
+                request.session["access_token_expiry"] = expiry_timestamp
+
+                user_session = UserSessionModel(user=user.model_dump(mode='json'), session_expiry_timestamp=expiry_timestamp)
+                return user_session
+        
+        # Provider authentication endpoints
+        providers_path = '/providers'
+        provider_login_path = '/providers/{provider_name}/login'
+        provider_callback_path = '/providers/{provider_name}/callback'
+
+        @auth_router.get(providers_path, tags=["Authentication"])
+        async def get_providers(request: Request) -> list[rm.ProviderResponse]:
+            """Get list of available authentication providers"""
+            _, root_path = self._get_base_url_for_current_app(request)
+
+            def get_icon_url(icon: str) -> str:
+                if icon.startswith("/public/"):
+                    core_url = root_path.split("/api/")[0]
+                    return core_url + icon
+                return icon
+
+            return [
+                rm.ProviderResponse(
+                    name=provider.name,
+                    label=provider.label,
+                    icon=get_icon_url(provider.icon),
+                    login_url=f"{root_path}{auth_path}/providers/{quote(provider.name)}/login",
+                )
+                for provider in self.authenticator.auth_providers
+            ]
+
+        @auth_router.get(provider_login_path, tags=["Authentication"], responses={
+            307: {"description": "Redirect to sign in with provider"},
+        })
+        async def provider_login(request: Request, provider_name: str, redirect_url: str | None = None) -> RedirectResponse:
+            """
+            Redirect to the login URL for the OAuth provider. 
+            
+            If login is successful, this endpoint redirects to the specified `redirect_url`. If no `redirect_url` is provided, it returns the user information of the Squirrels project's user.
+            """
+            client = oauth.create_client(provider_name)
+            if client is None:
+                raise InvalidInputError(status_code=404, error="provider_not_found", error_description=f"Provider {provider_name} not found or configured.")
+
+            origin, root_path = self._get_base_url_for_current_app(request)
+            callback_uri = f"{origin}{root_path}{auth_path}/providers/{quote(provider_name)}/callback"
+            request.session["redirect_url"] = redirect_url
+            
+            # OIDC best practice: include a nonce when requesting an id_token.
+            # Not all providers will use it, but major OIDC providers support it.
+            nonce = secrets.token_urlsafe(24)
+            request.session[f"oidc_nonce:{provider_name}"] = nonce
+
+            # PKCE: Some providers (e.g. Keycloak) require the authorization request to include
+            # `code_challenge_method=S256`. We also store the verifier so we can send it when
+            # exchanging the authorization code for tokens.
+            code_verifier = secrets.token_urlsafe(64)  # ~86 chars; within 43-128 PKCE range
+            request.session[f"pkce_verifier:{provider_name}"] = code_verifier
+            code_challenge = u.generate_pkce_challenge(code_verifier)
+
+            return await client.authorize_redirect(
+                request,
+                callback_uri,
+                nonce=nonce,
+                code_challenge=code_challenge,
+                code_challenge_method="S256",
+            )
+
+        @auth_router.get(provider_callback_path, tags=["Authentication"], responses={
+            200: {"description": "HTML page indicating successful login"},
+            307: {"description": "Redirect to redirect_url provided from provider login"},
+        })
+        async def provider_callback(request: Request, provider_name: str):
+            """Handle OAuth callback from provider"""
+            client = oauth.create_client(provider_name)
+            if client is None:
+                raise InvalidInputError(status_code=404, error="provider_not_found", error_description=f"Provider {provider_name} not found or configured.")
+
+            try:
+                code_verifier = request.session.pop(f"pkce_verifier:{provider_name}", None)
+                if code_verifier is None:
+                    token_details: dict = await client.authorize_access_token(request)
+                else:
+                    token_details = await client.authorize_access_token(request, code_verifier=code_verifier)
+            except Exception as e:
+                raise InvalidInputError(status_code=400, error="provider_authorization_failed", error_description=f"Could not authorize with provider for access token: {str(e)}")
+            
+            if is_external:
+                # Prefer id_token (JWT) for session auth if available. Many providers (e.g. Google)
+                # issue opaque access tokens that do not contain an issuer, which breaks provider
+                # auto-detection for session-based auth.
+                access_token = token_details.get("access_token")
+                id_token = token_details.get("id_token")
+                if isinstance(id_token, str) and id_token and id_token.count(".") == 2:
+                    access_token = id_token
+
+                if not isinstance(access_token, str) or not access_token:
+                    raise InvalidInputError(400, "provider_missing_access_token", f"Provider token not found for {provider_name}")
+
+                expires_in = token_details.get("expires_in")
+                if expires_in is None:
+                    # Fallback for providers that only return absolute expiry
+                    expiry_timestamp = token_details.get("expires_at")
+                    if expiry_timestamp is None:
+                        raise InvalidInputError(400, "provider_missing_expiry", f"Provider expiry timestamp not found for {provider_name}")
+                else:
+                    expiry_timestamp = datetime.now(timezone.utc).timestamp() + float(expires_in)
+            else:
+                expected_nonce = request.session.pop(f"oidc_nonce:{provider_name}", None)
+                user_info = self.authenticator.get_user_info_from_token_details(
+                    provider_name, token_details, expected_nonce=expected_nonce
+                )
+                user = self.authenticator.create_or_get_user_from_provider(provider_name, user_info)
+                access_token, expiry = self.authenticator.create_access_token(user, expiry_minutes=expiry_mins)
+                expiry_timestamp = expiry.timestamp()
+            
+            request.session["access_token"] = access_token
+            request.session["access_token_expiry"] = expiry_timestamp
+
+            redirect_url = request.session.pop("redirect_url", None)
+            if redirect_url:
+                return RedirectResponse(url=redirect_url)
+
+            template = self.templates.get_template("login_successful.html")
+            return HTMLResponse(content=template.render({"request": request}), status_code=200)
+
+        # Logout endpoint
+        logout_path = '/logout'
+        
+        @auth_router.post(logout_path, tags=["Authentication"])
+        async def logout(request: Request):
+            """Logout the current user by clearing the access token and expiry from the session"""
+            request.session.pop("access_token", None)
+            request.session.pop("access_token_expiry", None)
+            return Response(status_code=200)
+        
+        if not is_external:
+            # Change password endpoint
+            change_password_path = '/password'
+
+            class ChangePasswordRequest(BaseModel):
+                old_password: str
+                new_password: str
+
+            @auth_router.put(change_password_path, description="Change the password for the current user", tags=["Authentication"])
+            async def change_password(request: ChangePasswordRequest, user: RegisteredUser | GuestUser = Depends(self.get_current_user)) -> None:
+                if isinstance(user, GuestUser):
+                    raise InvalidInputError(401, "invalid_authorization_token", "Invalid authorization token")
+                self.authenticator.change_password(user.username, request.old_password, request.new_password)
+
+            # API Key endpoints
+            api_key_path = '/api-keys'
+
+            class ApiKeyRequestBody(BaseModel):
+                title: str = Field(description="The title of the API key")
+                expiry_minutes: int | None = Field(
+                    default=None,
+                    description="The number of minutes the API key is valid for (or valid indefinitely if not provided)."
+                )
+
+            @auth_router.post(api_key_path, description="Create a new API key for the user", tags=["Authentication"])
+            async def create_api_key(body: ApiKeyRequestBody, user: RegisteredUser | GuestUser = Depends(self.get_current_user)) -> rm.ApiKeyResponse:
+                if isinstance(user, GuestUser):
+                    raise InvalidInputError(401, "invalid_authorization_token", "Invalid authorization token, cannot create API key")
+                
+                api_key, _ = self.authenticator.create_access_token(user, expiry_minutes=body.expiry_minutes, title=body.title)
+                return rm.ApiKeyResponse(api_key=api_key)
+            
+            @auth_router.get(api_key_path, description="Get all API keys with title for the current user", tags=["Authentication"])
+            async def get_all_api_keys(user: RegisteredUser | GuestUser = Depends(self.get_current_user)) -> list[ApiKey]:
+                if isinstance(user, GuestUser):
+                    raise InvalidInputError(401, "invalid_authorization_token", "Invalid authorization token, cannot get API keys")
+                return self.authenticator.get_all_api_keys(user.username)
+            
+            revoke_api_key_path = '/api-keys/{key_id}'
+
+            @auth_router.delete(revoke_api_key_path, description="Revoke an API key", tags=["Authentication"], responses={
+                204: { "description": "API key revoked successfully" }
+            })
+            async def revoke_api_key(key_id: str, user: RegisteredUser | GuestUser = Depends(self.get_current_user)) -> Response:
+                if isinstance(user, GuestUser):
+                    raise InvalidInputError(401, "invalid_authorization_token", "Invalid authorization token, cannot revoke API key")
+                self.authenticator.revoke_api_key(user.username, key_id)
+                return Response(status_code=204)
+
+        app.include_router(auth_router)
+
+        # User management endpoints (disabled if external auth only)
+        if is_external: 
+            return
+
+        user_fields_path = '/user-fields'
+
+        @user_management_router.get(user_fields_path, description="Get details of the user fields", tags=["User Management"])
+        async def get_user_fields() -> UserFieldsModel:
+            return self.authenticator.user_fields
+        
+        list_or_add_users_path = '/users'
+        update_or_delete_user_path = '/users/{username}'
+
+        @user_management_router.get(list_or_add_users_path, tags=["User Management"])
+        async def list_all_users(user: AbstractUser = Depends(self.get_current_user)) -> list[UserInfoModel]:
+            if user.access_level != "admin":
+                raise InvalidInputError(403, "unauthorized_to_list_users", "Current user does not have permission to list users")
+            return self.authenticator.get_all_users()
+        
+        @user_management_router.post(list_or_add_users_path, description="Add a new user by providing details for username, password, and user fields", tags=["User Management"])
+        async def add_user(
+            new_user: AddUserModel, user: AbstractUser = Depends(self.get_current_user)
+        ) -> UserInfoModel:
+            if user.access_level != "admin":
+                raise InvalidInputError(403, "unauthorized_to_add_user", "Current user does not have permission to add new users")
+            return self.authenticator.add_user(new_user.username, new_user.model_dump(mode='json', exclude={"username"}))
+
+        @user_management_router.put(update_or_delete_user_path, description="Update the user of the given username given the new user details", tags=["User Management"])
+        async def update_user(
+            username: str, updated_user: UpdateUserModel, user: AbstractUser = Depends(self.get_current_user)
+        ) -> UserInfoModel:
+            if user.access_level != "admin":
+                raise InvalidInputError(403, "unauthorized_to_update_user", "Current user does not have permission to update users")
+            return self.authenticator.add_user(username, updated_user.model_dump(mode='json'), update_user=True)
+
+        @user_management_router.delete(update_or_delete_user_path, tags=["User Management"], responses={
+            204: { "description": "User deleted successfully" }
+        })
+        async def delete_user(username: str, user: AbstractUser = Depends(self.get_current_user)) -> Response:
+            if user.access_level != "admin":
+                raise InvalidInputError(403, "unauthorized_to_delete_user", "Current user cannot delete users")
+            if username == user.username:
+                raise InvalidInputError(403, "cannot_delete_own_user", "Cannot delete your own user")
+            self.authenticator.delete_user(username)
+            return Response(status_code=204)
+
+        app.include_router(user_management_router)

squirrels/_api_routes/base.py

@@ -0,0 +1,196 @@
+"""
+Base utilities and dependencies for API routes
+"""
+from typing import Any, Mapping, TypeVar, Callable, Coroutine, Literal
+from textwrap import dedent
+from fastapi import Request, Response, Depends, Header
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from fastapi.templating import Jinja2Templates
+from cachetools import TTLCache
+from pathlib import Path
+from datetime import datetime, timezone
+
+from .. import _utils as u
+from .._exceptions import InvalidInputError
+from .._project import SquirrelsProject
+from .._schemas.auth_models import AbstractUser
+from .._dataset_types import DatasetResultFormat
+from .._manifest import AuthType, AuthStrategy
+
+T = TypeVar('T')
+
+
+class RouteBase:
+    """Base class for route modules providing common functionality"""
+    
+    def __init__(self, get_bearer_token: HTTPBearer, project: SquirrelsProject, no_cache: bool = False):
+        self.project = project
+        self.no_cache = no_cache
+        self.logger = project._logger
+        self.env_vars = project._env_vars
+        self.manifest_cfg = project._manifest_cfg
+        self.authenticator = project._auth
+        self.param_cfg_set = project._param_cfg_set
+        
+        # Setup templates
+        template_dir = Path(__file__).parent.parent / "_package_data" / "templates"
+        self.templates = Jinja2Templates(directory=str(template_dir))
+    
+        # Authorization dependency for current user
+        def get_token_from_session(request: Request) -> str | None:
+            access_token = request.session.get("access_token")
+            if access_token is None:
+                return None
+            
+            expiry = request.session.get("access_token_expiry")
+            datetime_now = datetime.now(timezone.utc).timestamp()
+            if expiry and expiry > datetime_now:
+                return access_token
+            
+            if self.manifest_cfg.project_variables.auth_type == AuthType.REQUIRED:
+                raise InvalidInputError(401, "session_expired", "Login session expired. Please login again.")
+            else:
+                return None
+        
+        def get_user_from_headers(api_key: str | None, bearer_token: str | None) -> tuple[AbstractUser, float | None]:
+            auth_strategy = self.manifest_cfg.project_variables.auth_strategy
+            if (auth_strategy == AuthStrategy.EXTERNAL):
+                if not bearer_token:
+                    return self.project._guest_user, None
+                
+                user, expiry = self.authenticator.get_user_from_external_token(bearer_token)
+            else:
+                final_token = api_key if api_key else bearer_token
+                user, expiry = self.authenticator.get_user_from_token(final_token)
+            
+            if user is None:
+                user = self.project._guest_user
+            
+            return user, expiry
+        
+        async def get_current_user(
+            request: Request, response: Response, 
+            x_api_key: str | None = Header(None, description="API key for authentication (alternative to Authorization header)"), 
+            auth: HTTPAuthorizationCredentials = Depends(get_bearer_token)
+        ) -> AbstractUser:
+            token = auth.credentials if auth and auth.scheme == "Bearer" else None
+            access_token = token if token else get_token_from_session(request)
+            user, expiry = get_user_from_headers(x_api_key, access_token)
+            response.headers["Applied-Username"] = user.username
+            request.state.access_token_expiry = expiry
+            return user
+
+        self.get_user_from_headers = get_user_from_headers
+        self.get_current_user = get_current_user
+        
+    @staticmethod
+    def _get_base_url_for_current_app(request: Request) -> tuple[str, str]:
+        """
+        Build the absolute base URL for the *current* mounted app, including `root_path`.
+
+        We avoid `request.url_for(...)` because route names can collide when multiple Squirrels
+        FastAPI apps are mounted into the same root app.
+        """
+        origin = f"{request.url.scheme}://{request.url.netloc}"
+        root_path = str(request.scope.get("root_path") or "").rstrip("/")
+        return origin, root_path
+    
+    @property
+    def _parameters_description(self) -> str:
+        """Get the standard parameters description"""
+        return dedent("""
+            Selections of one parameter may cascade the available options in another parameter. 
+            
+            For example, if the dataset has parameters for 'country' and 'city', available options for 'city' would depend on the selected option 'country'. 
+            
+            If a parameter has `"trigger_refresh": true` and its selection changes, provide the parameter selection to this endpoint to refresh the parameter options of children parameters.
+        """).strip()
+    
+    def get_selections_as_immutable(self, params: Mapping, uncached_keys: set[str]) -> tuple[tuple[str, Any], ...]:
+        """Convert selections into a cachable tuple of pairs"""
+        selections = list()
+        for key, val in params.items():
+            if key in uncached_keys or val is None:
+                continue
+            if isinstance(val, (list, tuple)):
+                val = tuple(val)
+            selections.append((u.normalize_name(key), val))
+        return tuple(selections)
+
+    async def do_cachable_action(self, cache: TTLCache, action: Callable[..., Coroutine[Any, Any, T]], *args) -> T:
+        """Execute a cachable action"""
+        cache_key = tuple(args)
+        result = cache.get(cache_key)
+        if result is None:
+            result = await action(*args)
+            cache[cache_key] = result
+        return result
+    
+    def get_name_from_path_section(self, request: Request, section: int) -> str:
+        """Extract name from request path section"""
+        url_path: str = request.url.path
+        name_raw = url_path.split('/')[section]
+        return u.normalize_name(name_raw)
+
+    def get_configurables_from_headers(self, headers: Mapping[str, str]) -> tuple[tuple[str, str], ...]:
+        """Extract configurables from request headers with prefix 'x-config-'."""
+        prefix = "x-config-"
+        cfg_pairs: list[tuple[str, str]] = []
+        seen_configurables: dict[str, str] = {}  # normalized_name -> header_name
+
+        for key, value in headers.items():
+            key_lower = str(key).lower()
+            if key_lower.startswith(prefix):
+                cfg_name_raw = key_lower[len(prefix):]
+                cfg_name_normalized = u.normalize_name(cfg_name_raw)  # Convert to underscore convention
+
+                # Check if we've already seen this configurable (with different header format)
+                if cfg_name_normalized in seen_configurables:
+                    existing_header = seen_configurables[cfg_name_normalized]
+                    raise InvalidInputError(
+                        400, "duplicate_configurable_header",
+                        f"Only one header format is allowed for configurable '{cfg_name_normalized}'. "
+                        f"Both '{existing_header}' and '{key_lower}' were provided."
+                    )
+
+                seen_configurables[cfg_name_normalized] = key_lower
+                cfg_pairs.append((cfg_name_normalized, str(value)))
+
+        configurables = [k for k, _ in cfg_pairs]
+        self.logger.info(f"Configurables specified: {configurables}", data={"configurables_specified": configurables})
+        return tuple(cfg_pairs)
+
+    @staticmethod
+    def extract_orientation_offset_and_limit(
+        params: Mapping[str, Any], *,
+        key_prefix: str = "x_",
+        default_orientation: Literal["records", "rows", "columns"] = "records",
+        default_offset: int = 0, default_limit: int = 1000
+    ) -> DatasetResultFormat:
+        """
+        Extract orientation, offset, and limit from query parameters.
+
+        Args:
+            params: Query parameters
+
+        Returns:
+            Tuple of (orientation, offset, limit)
+        """
+        # Handle orientation
+        orientation = str(params.get(f"{key_prefix}orientation", default_orientation)).lower()
+
+        if orientation not in ["records", "rows", "columns"]:
+            raise InvalidInputError(400, "invalid_orientation", f"Orientation must be 'records', 'rows', or 'columns'. Invalid orientation provided: {orientation}")
+
+        # Handle limit and offset
+        offset = int(params.get(f"{key_prefix}offset", default_offset))
+        limit = int(params.get(f"{key_prefix}limit", default_limit))
+
+        if offset < 0:
+            raise InvalidInputError(400, "invalid_offset", "Offset must be non-negative")
+
+        if limit < 0:
+            raise InvalidInputError(400, "invalid_limit", "Limit must be non-negative")
+        
+        return DatasetResultFormat(orientation, offset, limit)
+