diff --git a/src/py/mat3ra/notebooks_utils/auth.py b/src/py/mat3ra/notebooks_utils/auth.py index b8c3cb634..433f36490 100644 --- a/src/py/mat3ra/notebooks_utils/auth.py +++ b/src/py/mat3ra/notebooks_utils/auth.py @@ -1,28 +1,65 @@ +"""Notebook authentication. + +Two paths depending on environment: + +1. JupyterLite embedded in WebApp — the host app injects a token via + postMessage and authenticate_jupyterlite() stores it in the environment. + +2. Standalone / IDE / local dev — OIDC device-code flow. Tokens are cached + (file on disk or IndexedDB in Pyodide) so the user only logs in once + per expiry window. +""" + import inspect import os from mat3ra.api_client import ACCESS_TOKEN_ENV_VAR -from .core.api.auth import authenticate_oidc +from .core.api.auth import authenticate_oidc, get_oidc_base_url, store_token_data_in_environment from .io import get_data from .ipython.ui import show_device_flow_popup -from .primitive.environment import ENVIRONMENT, EnvironmentsEnum +from .primitive.environment import is_pyodide_environment from .pyodide.api.auth import authenticate_jupyterlite +from .token_store import load_token, save_token REFRESH_TOKEN_ENV_VAR = "OIDC_REFRESH_TOKEN" +async def _authenticate_oidc_with_cache(force=False): + oidc_url = get_oidc_base_url() + cached = None if force else await load_token(oidc_url) + + if cached: + store_token_data_in_environment(cached) + return + + token_data = await authenticate_oidc(show_popup=show_device_flow_popup) + await save_token(oidc_url, token_data) + + async def authenticate(force=False, globals_dict=None): + """ + Authenticate the current notebook session with the Mat3ra platform. + Follow the link to a new browser window to complete the authentication process. + The token is cached in the environment for future use. + + Args: + force: Re-authenticate even if a valid cached token exists. + globals_dict: Caller's global namespace. Auto-detected if not provided. + """ if globals_dict is None: frame = inspect.currentframe() try: globals_dict = frame.f_back.f_globals # type: ignore finally: del frame - if ENVIRONMENT == EnvironmentsEnum.PYODIDE: + + if is_pyodide_environment(): get_data("data_from_host", globals_dict) + data_from_host = globals_dict.get("data_from_host") + if data_from_host: await authenticate_jupyterlite(data_from_host) elif ACCESS_TOKEN_ENV_VAR not in os.environ or force: - await authenticate_oidc(show_popup=show_device_flow_popup) + await _authenticate_oidc_with_cache(force) diff --git a/src/py/mat3ra/notebooks_utils/core/api/token_store.py b/src/py/mat3ra/notebooks_utils/core/api/token_store.py new file mode 100644 index 000000000..17b9f8725 --- /dev/null +++ b/src/py/mat3ra/notebooks_utils/core/api/token_store.py @@ -0,0 +1,20 @@ +"""File-based OIDC token cache storage.""" + +import json +import os + +_FILE_PATH = os.path.join(os.path.expanduser("~"), ".mat3ra", "oidc_token_cache.json") + + +async def read() -> dict: + try: + with open(_FILE_PATH) as f: + return json.load(f) + except (FileNotFoundError, json.JSONDecodeError): + return {} + + +async def write(cache: dict) -> None: + os.makedirs(os.path.dirname(_FILE_PATH), exist_ok=True) + with open(_FILE_PATH, "w") as f: + json.dump(cache, f) diff --git a/src/py/mat3ra/notebooks_utils/pyodide/api/token_store.py b/src/py/mat3ra/notebooks_utils/pyodide/api/token_store.py new file mode 100644 index 000000000..7e5414a5b --- /dev/null +++ b/src/py/mat3ra/notebooks_utils/pyodide/api/token_store.py @@ -0,0 +1,45 @@ +"""IndexedDB-based OIDC token cache storage for Pyodide Web Workers.""" + +import json + +from pyodide.code import run_js # type: ignore + +_idb = run_js( + """ +(function() { + const open = () => new Promise(resolve => { + const req = indexedDB.open("mat3ra", 1); + req.onupgradeneeded = () => req.result.createObjectStore("tokens"); + req.onsuccess = () => resolve(req.result); + req.onerror = () => resolve(null); + }); + return { + get: async () => { + const db = await open(); + if (!db) return null; + return new Promise(resolve => { + const g = db.transaction("tokens").objectStore("tokens").get("oidc_cache"); + g.onsuccess = () => { db.close(); resolve(g.result || null); }; + g.onerror = () => { db.close(); resolve(null); }; + }); + }, + set: async (data) => { + const db = await open(); + if (!db) return; + const tx = db.transaction("tokens", "readwrite"); + tx.objectStore("tokens").put(data, "oidc_cache"); + return new Promise(resolve => { tx.oncomplete = () => { db.close(); resolve(); }; }); + } + }; +})() +""" +) + + +async def read() -> dict: + result = await _idb.get() + return json.loads(str(result)) if result else {} + + +async def write(cache: dict) -> None: + await _idb.set(json.dumps(cache)) diff --git a/src/py/mat3ra/notebooks_utils/token_store.py b/src/py/mat3ra/notebooks_utils/token_store.py new file mode 100644 index 000000000..73eebd6a8 --- /dev/null +++ b/src/py/mat3ra/notebooks_utils/token_store.py @@ -0,0 +1,33 @@ +"""Persist OIDC tokens across notebook kernels. + +Switches storage between file (Python) and IndexedDB (Pyodide). +""" + +import time +from typing import Optional + +from .primitive.environment import is_pyodide_environment + +if is_pyodide_environment(): + from .pyodide.api.token_store import read as _read + from .pyodide.api.token_store import write as _write +else: + from .core.api.token_store import read as _read + from .core.api.token_store import write as _write + +_EXPIRY_BUFFER = 60 * 5 # to avoid using soon-to-expire tokens + + +async def save_token(oidc_url: str, token_data: dict) -> None: + entry = dict(token_data) + entry["expires_at"] = time.time() + entry.get("expires_in", 3600) + cache = await _read() + cache[oidc_url] = entry + await _write(cache) + + +async def load_token(oidc_url: str) -> Optional[dict]: + entry = (await _read()).get(oidc_url) + if not entry or entry.get("expires_at", 0) <= time.time() + _EXPIRY_BUFFER: + return None + return entry