obi-auth 1.2.1__py3-none-any.whl

obi_auth/__init__.py

@@ -0,0 +1,13 @@
+"""obi_auth."""
+
+from obi_auth.client import get_token, get_token_info, get_user_info
+from obi_auth.typedef import AuthMode, DeploymentEnvironment, TokenProvider
+
+__all__ = [
+    "get_token",
+    "get_token_info",
+    "get_user_info",
+    "DeploymentEnvironment",
+    "AuthMode",
+    "TokenProvider",
+]

obi_auth/cache.py

@@ -0,0 +1,129 @@
+"""Token cache module."""
+
+import time
+
+import jwt
+from cryptography.fernet import Fernet, InvalidToken
+
+from obi_auth.config import settings
+from obi_auth.storage import Storage
+from obi_auth.typedef import (
+    AuthManagerTokenInfo,
+    CachedAuthManagerTokenInfo,
+    CachedTokenInfo,
+    KeycloakTokenInfo,
+)
+from obi_auth.util import derive_fernet_key
+
+
+class BaseTokenCache:
+    """Shared Fernet helpers for token caches."""
+
+    def __init__(self):
+        """Initialize the token cache."""
+        self._cipher = Fernet(key=derive_fernet_key())
+
+    def _encrypt_access_token(self, access_token: str) -> tuple[bytes, int]:
+        creation_time, time_to_live = _get_token_times(access_token)
+        fernet_token = self._cipher.encrypt_at_time(
+            data=access_token.encode(encoding="utf-8"),
+            current_time=creation_time,
+        )
+        return fernet_token, time_to_live
+
+    def _decrypt_access_token(self, token: bytes, ttl: int) -> str:
+        return self._cipher.decrypt_at_time(
+            token=token,
+            ttl=ttl,
+            current_time=_now(),
+        ).decode()
+
+
+class TokenCache(BaseTokenCache):
+    """Cache for Keycloak access tokens."""
+
+    def get(self, storage: Storage) -> KeycloakTokenInfo | None:
+        """Get a cached Keycloak token if still valid, else None."""
+        if not (cached_token_info := storage.read()):
+            return None
+        if not isinstance(cached_token_info, CachedTokenInfo):
+            storage.clear()
+            return None
+        try:
+            access_token = self._decrypt_access_token(
+                cached_token_info.token, cached_token_info.ttl
+            )
+        except InvalidToken:
+            storage.clear()
+            return None
+        return KeycloakTokenInfo(access_token=access_token)
+
+    def set(self, token_info: KeycloakTokenInfo, storage: Storage) -> None:
+        """Store a Keycloak token in the cache."""
+        fernet_token, time_to_live = self._encrypt_access_token(token_info.access_token)
+        storage.write(CachedTokenInfo(token=fernet_token, ttl=time_to_live))
+
+
+class AuthManagerTokenCache(BaseTokenCache):
+    """Cache for auth-manager access tokens and persistent token ids."""
+
+    def get(self, storage: Storage) -> AuthManagerTokenInfo | None:
+        """Get a cached auth-manager token.
+
+        Returns a valid access token when available. If the access token has expired
+        but a persistent token id is stored, returns ``AuthManagerTokenInfo`` with an
+        empty ``access_token`` so the caller can remint.
+        """
+        if not (cached_token_info := storage.read()):
+            return None
+        if not isinstance(cached_token_info, CachedAuthManagerTokenInfo):
+            storage.clear()
+            return None
+
+        try:
+            persistent_token_id = self._cipher.decrypt(
+                cached_token_info.persistent_token_id
+            ).decode()
+        except InvalidToken:
+            storage.clear()
+            return None
+
+        try:
+            access_token = self._decrypt_access_token(
+                cached_token_info.token, cached_token_info.ttl
+            )
+        except InvalidToken:
+            return AuthManagerTokenInfo(access_token=None, persistent_token_id=persistent_token_id)
+
+        return AuthManagerTokenInfo(
+            access_token=access_token, persistent_token_id=persistent_token_id
+        )
+
+    def set(self, token_info: AuthManagerTokenInfo, storage: Storage) -> None:
+        """Store an auth-manager token and persistent id in the cache."""
+        if token_info.access_token is None:
+            msg = "Cannot cache AuthManagerTokenInfo without an access_token"
+            raise ValueError(msg)
+        fernet_token, time_to_live = self._encrypt_access_token(token_info.access_token)
+        persistent_token_id = self._cipher.encrypt(
+            token_info.persistent_token_id.encode(encoding="utf-8")
+        )
+        storage.write(
+            CachedAuthManagerTokenInfo(
+                token=fernet_token,
+                ttl=time_to_live,
+                persistent_token_id=persistent_token_id,
+            )
+        )
+
+
+def _now() -> int:
+    """Return UTC timestamp now."""
+    return int(time.time())
+
+
+def _get_token_times(token: str) -> tuple[int, int]:
+    """Get the creation time and time to live of a token."""
+    info = jwt.decode(token.encode(), options={"verify_signature": False})
+    effective_ttl = info["exp"] - info["iat"] - settings.EPSILON_TOKEN_TTL_SECONDS
+    return info["iat"], effective_ttl

obi_auth/cli.py

@@ -0,0 +1,151 @@
+"""CLI for obi-auth."""
+
+import json
+import logging
+import sys
+
+try:
+    import click
+except ImportError as exc:
+    raise SystemExit(
+        "The obi-auth CLI requires optional dependencies that are not installed.\n"
+        "Install them with:\n"
+        "  pip install 'obi-auth[cli]'"
+    ) from exc
+
+import obi_auth
+from obi_auth.typedef import AuthMode, DeploymentEnvironment, TokenProvider
+
+
+@click.group()
+@click.option(
+    "--log-level",
+    type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], case_sensitive=False),
+    default="WARNING",
+    show_default=True,
+    help="Logging level",
+)
+def main(log_level: str):
+    """CLI for obi-auth."""
+    logging.basicConfig(level=log_level)
+
+
+@main.command()
+@click.option(
+    "--environment",
+    "-e",
+    type=click.Choice(DeploymentEnvironment),
+    default=DeploymentEnvironment.staging,
+    show_default=True,
+    help="Target environment",
+)
+@click.option(
+    "--auth-mode",
+    "-m",
+    type=click.Choice(AuthMode),
+    default=AuthMode.pkce,
+    show_default=True,
+    help="Authentication method",
+)
+@click.option(
+    "--token-provider",
+    "-p",
+    type=click.Choice(TokenProvider),
+    default=TokenProvider.keycloak,
+    show_default=True,
+    help="Token issuer: keycloak or auth_manager",
+)
+@click.option(
+    "--persistent-token-id",
+    default=None,
+    help="Persistent token id (required when auth-mode is persistent_token)",
+)
+@click.option(
+    "--force-refresh",
+    help="Clear the cached token and authenticate again",
+    is_flag=True,
+    default=False,
+)
+def get_token(
+    environment: DeploymentEnvironment,
+    auth_mode: AuthMode,
+    token_provider: TokenProvider,
+    persistent_token_id: str | None,
+    force_refresh: bool,
+):
+    """Authenticate, print the token to stdout."""
+    access_token = obi_auth.get_token(
+        environment=environment,
+        auth_mode=auth_mode,
+        token_provider=token_provider,
+        persistent_token_id=persistent_token_id,
+        force_refresh=force_refresh,
+    )
+    print(access_token)
+
+
+@main.command()
+@click.argument("access_token", required=False)
+def decode_token(access_token: str | None):
+    """Decode token from argument or stdin."""
+    if not access_token:
+        access_token = sys.stdin.read().strip()
+
+    if not access_token:
+        raise click.ClickException("No access token provided via argument or stdin")
+
+    print(json.dumps(obi_auth.get_token_info(access_token), indent=2))
+
+
+@main.command()
+@click.option(
+    "--environment",
+    "-e",
+    type=click.Choice(DeploymentEnvironment),
+    default=DeploymentEnvironment.staging,
+    show_default=True,
+    help="Target environment",
+)
+@click.option(
+    "--auth-mode",
+    "-m",
+    type=click.Choice(AuthMode),
+    default=AuthMode.pkce,
+    show_default=True,
+    help="Authentication method",
+)
+@click.option(
+    "--token-provider",
+    "-p",
+    type=click.Choice(TokenProvider),
+    default=TokenProvider.keycloak,
+    show_default=True,
+    help="Token issuer: keycloak or auth_manager",
+)
+@click.option(
+    "--persistent-token-id",
+    default=None,
+    help="Persistent token id (required when auth-mode is persistent_token)",
+)
+@click.option(
+    "--force-refresh",
+    help="Clear the cached token and authenticate again",
+    is_flag=True,
+    default=False,
+)
+def get_user_info(
+    environment: DeploymentEnvironment,
+    auth_mode: AuthMode,
+    token_provider: TokenProvider,
+    persistent_token_id: str | None,
+    force_refresh: bool,
+):
+    """Show user info information."""
+    access_token = obi_auth.get_token(
+        environment=environment,
+        auth_mode=auth_mode,
+        token_provider=token_provider,
+        persistent_token_id=persistent_token_id,
+        force_refresh=force_refresh,
+    )
+    print(json.dumps(obi_auth.get_user_info(access_token, environment=environment), indent=2))

obi_auth/client.py

@@ -0,0 +1,236 @@
+"""This module provides a client for the obi_auth service."""
+
+import logging
+from collections.abc import Callable
+
+import jwt
+
+from obi_auth.cache import AuthManagerTokenCache, TokenCache
+from obi_auth.config import settings
+from obi_auth.exception import AuthFlowError, ClientError, ConfigError, LocalServerError
+from obi_auth.flows.auth_manager import (
+    auth_manager_exchange_token,
+    auth_manager_mint_access_token,
+)
+from obi_auth.flows.daf import daf_authenticate
+from obi_auth.flows.pkce import pkce_authenticate
+from obi_auth.request import user_info
+from obi_auth.server import AuthServer
+from obi_auth.storage import Storage
+from obi_auth.typedef import (
+    AuthManagerTokenInfo,
+    AuthMode,
+    DeploymentEnvironment,
+    KeycloakTokenInfo,
+    TokenProvider,
+)
+
+L = logging.getLogger(__name__)
+
+
+_TOKEN_CACHE = TokenCache()
+_AUTH_MANAGER_TOKEN_CACHE = AuthManagerTokenCache()
+
+
+def get_token(
+    *,
+    environment: DeploymentEnvironment = DeploymentEnvironment.staging,
+    auth_mode: AuthMode = AuthMode.pkce,
+    token_provider: TokenProvider = TokenProvider.keycloak,
+    force_refresh: bool = False,
+    persistent_token_id: str | None = None,
+) -> str:
+    """Get token.
+
+    Args:
+        environment: Target deployment environment.
+        auth_mode: How to authenticate with Keycloak (``pkce`` or ``daf``), or
+            ``persistent_token`` to mint from an existing auth-manager id
+            (backwards-compatible public API for callers such as obi-one).
+        token_provider: Who issues the returned access token when ``auth_mode`` is
+            ``pkce`` or ``daf``. ``keycloak`` returns the Keycloak token;
+            ``auth_manager`` exchanges it and mints an auth-manager access token.
+            Ignored when ``auth_mode`` is ``persistent_token``.
+        force_refresh: Clear the cached token and authenticate again.
+        persistent_token_id: Required when ``auth_mode`` is ``persistent_token``.
+    """
+    auth_mode = AuthMode(auth_mode)
+    token_provider = TokenProvider(token_provider)
+
+    L.debug("Using %s as the config dir", settings.config_dir)
+
+    # Public API compat: get_token(auth_mode="persistent_token", persistent_token_id=...)
+    if auth_mode == AuthMode.persistent_token:
+        if not persistent_token_id:
+            raise ClientError("persistent_token_id is required when auth_mode is persistent_token.")
+        return _get_persistent_token(
+            environment=environment,
+            persistent_token_id=persistent_token_id,
+            force_refresh=force_refresh,
+        )
+
+    storage = Storage(
+        config_dir=settings.config_dir,
+        environment=environment,
+        key=f"{auth_mode}_{token_provider}",
+    )
+
+    if token_provider == TokenProvider.auth_manager:
+        return _get_auth_manager_token(
+            storage=storage,
+            environment=environment,
+            auth_mode=auth_mode,
+            force_refresh=force_refresh,
+        )
+    return _get_keycloak_token(
+        storage=storage,
+        environment=environment,
+        auth_mode=auth_mode,
+        force_refresh=force_refresh,
+    )
+
+
+def _get_persistent_token(
+    *,
+    environment: DeploymentEnvironment,
+    persistent_token_id: str,
+    force_refresh: bool,
+) -> str:
+    """Mint from a known persistent id using the auth-manager cache/mint helpers."""
+    storage = Storage(
+        config_dir=settings.config_dir,
+        environment=environment,
+        key=persistent_token_id,
+    )
+    if force_refresh:
+        L.debug("Forcing token refresh, clearing cached token")
+        storage.clear()
+    elif token_info := _AUTH_MANAGER_TOKEN_CACHE.get(storage):
+        if token_info.access_token:
+            L.debug("Using cached token")
+            return token_info.access_token
+        if refreshed := _refresh_auth_manager_token(
+            token_info.persistent_token_id, storage=storage, environment=environment
+        ):
+            if refreshed.access_token is None:
+                raise ClientError("Authentication process failed.")
+            return refreshed.access_token
+
+    try:
+        token_info = auth_manager_mint_access_token(persistent_token_id, environment=environment)
+    except AuthFlowError as e:
+        raise ClientError("Authentication process failed.") from e
+
+    _AUTH_MANAGER_TOKEN_CACHE.set(token_info, storage)
+    if token_info.access_token is None:
+        raise ClientError("Authentication process failed.")
+    return token_info.access_token
+
+
+def _get_keycloak_token(
+    *,
+    storage: Storage,
+    environment: DeploymentEnvironment,
+    auth_mode: AuthMode,
+    force_refresh: bool,
+) -> str:
+    if force_refresh:
+        L.debug("Forcing token refresh, clearing cached token")
+        storage.clear()
+    elif token_info := _TOKEN_CACHE.get(storage):
+        L.debug("Using cached token")
+        return token_info.access_token
+
+    auth_method = _get_auth_method(auth_mode)
+    token_info = auth_method(environment=environment)
+    _TOKEN_CACHE.set(token_info, storage)
+    return token_info.access_token
+
+
+def _get_auth_manager_token(
+    *,
+    storage: Storage,
+    environment: DeploymentEnvironment,
+    auth_mode: AuthMode,
+    force_refresh: bool,
+) -> str:
+    if force_refresh:
+        L.debug("Forcing token refresh, clearing cached token")
+        storage.clear()
+    elif token_info := _AUTH_MANAGER_TOKEN_CACHE.get(storage):
+        if token_info.access_token:
+            L.debug("Using cached token")
+            return token_info.access_token
+        if refreshed := _refresh_auth_manager_token(
+            token_info.persistent_token_id, storage=storage, environment=environment
+        ):
+            if refreshed.access_token is None:
+                raise ClientError("Authentication process failed.")
+            return refreshed.access_token
+
+    auth_method = _get_auth_method(auth_mode)
+    keycloak_token = auth_method(environment=environment)
+    try:
+        token_info = auth_manager_exchange_token(keycloak_token, environment=environment)
+    except AuthFlowError as e:
+        raise ClientError("Authentication process failed.") from e
+
+    _AUTH_MANAGER_TOKEN_CACHE.set(token_info, storage)
+    if token_info.access_token is None:
+        raise ClientError("Authentication process failed.")
+    return token_info.access_token
+
+
+def _refresh_auth_manager_token(
+    persistent_token_id: str, *, storage: Storage, environment: DeploymentEnvironment
+) -> AuthManagerTokenInfo | None:
+    """Mint a new access token from a cached persistent token id."""
+    L.debug("Cached access token expired, minting a new one from persistent token id")
+    try:
+        token_info = auth_manager_mint_access_token(persistent_token_id, environment=environment)
+    except AuthFlowError:
+        L.debug("Failed to mint access token from persistent token id, clearing cache")
+        storage.clear()
+        return None
+
+    _AUTH_MANAGER_TOKEN_CACHE.set(token_info, storage)
+    return token_info
+
+
+def _get_auth_method(auth_mode: AuthMode) -> Callable[..., KeycloakTokenInfo]:
+    methods: dict[AuthMode, Callable[..., KeycloakTokenInfo]] = {
+        AuthMode.pkce: _pkce_authenticate,
+        AuthMode.daf: _daf_authenticate,
+    }
+    return methods[auth_mode]
+
+
+def _pkce_authenticate(*, environment: DeploymentEnvironment) -> KeycloakTokenInfo:
+    try:
+        with AuthServer().run() as local_server:
+            return pkce_authenticate(server=local_server, environment=environment)
+    except AuthFlowError as e:
+        raise ClientError("Authentication process failed.") from e
+    except LocalServerError as e:
+        raise ClientError("Local server failed to authenticate.") from e
+    except ConfigError as e:
+        raise ClientError("There is a mistake with configuration settings.") from e
+
+
+def _daf_authenticate(*, environment: DeploymentEnvironment) -> KeycloakTokenInfo:
+    try:
+        return daf_authenticate(environment=environment)
+    except AuthFlowError as e:
+        raise ClientError("Authentication process failed.") from e
+
+
+def get_token_info(token: str) -> dict:
+    """Decode token information."""
+    return jwt.decode(token, options={"verify_signature": False})
+
+
+def get_user_info(
+    token: str, environment: DeploymentEnvironment = DeploymentEnvironment.staging
+) -> dict:
+    """Get user info from valid token."""
+    return user_info(token, environment=environment).json()

obi_auth/config.py

@@ -0,0 +1,101 @@
+"""This module provides a config for the obi_auth service."""
+
+from pathlib import Path
+from typing import Annotated
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from obi_auth.exception import ConfigError
+from obi_auth.typedef import DeploymentEnvironment, KeycloakRealm
+from obi_auth.util import get_config_dir
+
+
+class Settings(BaseSettings):
+    """Environment settings for this library."""
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_prefix="OBI_AUTH_",
+        env_file_encoding="utf-8",
+        case_sensitive=False,
+        validate_default=False,
+        extra="ignore",
+    )
+
+    config_dir: Annotated[
+        Path,
+        Field(
+            description="Directory to store the token.",
+            default_factory=get_config_dir,
+        ),
+    ]
+
+    KEYCLOAK_ENV: DeploymentEnvironment = DeploymentEnvironment.staging
+    KEYCLOAK_REALM: KeycloakRealm = KeycloakRealm.sbo
+    KEYCLOAK_CLIENT_ID: str = "obi-entitysdk-auth"
+
+    EPSILON_TOKEN_TTL_SECONDS: int = 60
+
+    LOCAL_SERVER_TIMEOUT: int = 60
+
+    def _get_domain_url(self, override_env: DeploymentEnvironment) -> str:
+        """Return domain url based on environment."""
+        match env := override_env or self.KEYCLOAK_ENV:
+            case DeploymentEnvironment.staging:
+                return "https://staging.cell-a.openbraininstitute.org"
+            case DeploymentEnvironment.production:
+                return "https://cell-a.openbraininstitute.org"
+            case _:
+                raise ConfigError(f"Unknown deployment environment {env}")
+
+    def get_keycloak_url(self, override_env: DeploymentEnvironment | None = None):
+        """Return keycloak url."""
+        url = self._get_domain_url(override_env or self.KEYCLOAK_ENV)
+        return f"{url}/auth/realms/{self.KEYCLOAK_REALM}"
+
+    def get_keycloak_token_endpoint(self, override_env: DeploymentEnvironment | None = None) -> str:
+        """Return keycloak token endpoint."""
+        base_url = self.get_keycloak_url(override_env=override_env)
+        return f"{base_url}/protocol/openid-connect/token"
+
+    def get_keycloak_auth_endpoint(self, override_env: DeploymentEnvironment | None = None) -> str:
+        """Return keycloak auth endpoint."""
+        base_url = self.get_keycloak_url(override_env=override_env)
+        return f"{base_url}/protocol/openid-connect/auth"
+
+    def get_keycloak_device_auth_endpoint(
+        self, override_env: DeploymentEnvironment | None = None
+    ) -> str:
+        """Return keycloack device auth endpoint."""
+        base_url = self.get_keycloak_url(override_env=override_env)
+        return f"{base_url}/protocol/openid-connect/auth/device"
+
+    def get_keycloak_user_info_endpoint(
+        self, override_env: DeploymentEnvironment | None = None
+    ) -> str:
+        """Return keycloak user info endpoint."""
+        base_url = self.get_keycloak_url(override_env=override_env)
+        return f"{base_url}/protocol/openid-connect/userinfo"
+
+    def get_auth_manager_url(self, override_env: DeploymentEnvironment | None = None) -> str:
+        """Return auth manager url."""
+        url = self._get_domain_url(override_env or self.KEYCLOAK_ENV)
+        return f"{url}/api/auth-manager/v1"
+
+    def get_auth_manager_access_token_endpoint(
+        self, override_env: DeploymentEnvironment | None = None
+    ) -> str:
+        """Return auth-manager token endpoint."""
+        base_url = self.get_auth_manager_url(override_env=override_env)
+        return f"{base_url}/access-token"
+
+    def get_auth_manager_token_exchange_endpoint(
+        self, override_env: DeploymentEnvironment | None = None
+    ) -> str:
+        """Return auth-manager token exchange endpoint."""
+        base_url = self.get_auth_manager_url(override_env=override_env)
+        return f"{base_url}/token-exchange"
+
+
+settings = Settings()

obi_auth/exception.py

@@ -0,0 +1,21 @@
+"""Exceptions module."""
+
+
+class ObiAuthError(Exception):
+    """Generic obi-auth error."""
+
+
+class AuthFlowError(ObiAuthError):
+    """Authenticatin flow error."""
+
+
+class ClientError(ObiAuthError):
+    """Client related error."""
+
+
+class LocalServerError(ObiAuthError):
+    """Local server related error."""
+
+
+class ConfigError(ObiAuthError):
+    """Configuration settings error."""

obi_auth/flows/__init__.py

@@ -0,0 +1 @@
+"""Authentication flows."""