"""OAuth / whitelist helpers used by the Gradio submission flow.""" from __future__ import annotations import logging import gradio as gr from huggingface_hub import whoami from huggingface_hub.errors import HfHubHTTPError from src.envs import SIZE_WHITELIST, WHITELIST logger = logging.getLogger(__name__) def check_intel_org(oauth_token: gr.OAuthToken | None) -> tuple[bool, bool]: """Return ``(is_authenticated, is_intel)``. Raises :class:`HfHubHTTPError` if the token is expired/invalid. """ if oauth_token is None: return False, False info = whoami(token=oauth_token.token) username = info.get("name", "unknown") orgs = [org["name"] for org in info.get("orgs", [])] logger.info("[Auth] user=%s orgs=%s", username, orgs) return True, "intel" in [o.lower() for o in orgs] def is_whitelisted(oauth_token: gr.OAuthToken | None) -> bool: """Whether the authenticated user is in the submission whitelist.""" if oauth_token is None or not oauth_token.token: return False if WHITELIST is None: # No whitelist configured → everyone with a valid token qualifies return True try: info = whoami(token=oauth_token.token) return info.get("name", "unknown") in WHITELIST except Exception: return False def is_allowed_to_resubmit_failed(username: str) -> bool: """Whether *username* may re-submit a model that previously Failed. Mirrors the :data:`SIZE_WHITELIST` semantics: * ``SIZE_WHITELIST is None`` → everyone is allowed. * Otherwise only users in the set are allowed. """ if SIZE_WHITELIST is None: return True return bool(username) and username in SIZE_WHITELIST def process_user_on_load(oauth_token: gr.OAuthToken | None): """Page-load callback: log user info and show the hardware selector for ``SIZE_WHITELIST`` users. """ show_hw = False username = "" if oauth_token is not None and oauth_token.token: try: info = whoami(token=oauth_token.token) username = info.get("name", "") orgs = [org["name"] for org in info.get("orgs", [])] logger.info("[Auth][PageLoad] user=%s orgs=%s", username, orgs) if SIZE_WHITELIST is not None: show_hw = username in SIZE_WHITELIST except Exception as e: logger.warning("[Auth][PageLoad] whoami failed: %s", e) if username: user_html = ( f'' f'👤 {username}' ) user_html_quant = user_html user_html_eval = user_html else: user_html_quant = ( '' 'Unauthorized — ' 'Click to authorize' ) user_html_eval = ( '' 'Unauthorized — ' 'Click to authorize' ) return ( gr.update(visible=show_hw), # hw_col_quant gr.update(visible=show_hw), # hw_col_eval username, # username_state gr.update(value=user_html_quant), # current_user_quant gr.update(value=user_html_eval), # current_user_eval ) def check_auth_for_submit(oauth_token: gr.OAuthToken | None) -> str: """Pre-check auth before submission. Returns ``"TRIGGER"`` or ``"OK"``.""" if oauth_token is None or not oauth_token.token: return "TRIGGER" try: whoami(token=oauth_token.token) return "OK" except (HfHubHTTPError, Exception): logger.warning("[Auth] token check failed (likely expired)", exc_info=True) return "TRIGGER" def sync_username(oauth_token: gr.OAuthToken | None = None): """Refresh ``username_state`` after a successful auth. Returns :func:`gr.update` (no change) if the token is unavailable or the call fails. """ if oauth_token is None or not oauth_token.token: return gr.update() try: info = whoami(token=oauth_token.token) return info.get("name", "") except Exception: return gr.update()