Spaces:
Paused
Paused
| """ | |
| OAuth 2.0 / OpenID Connect Authentication Backend | |
| Supports Google, GitHub, and generic OIDC providers for single sign-on. | |
| Uses Authlib for OAuth flow management, token exchange, and OIDC discovery. | |
| """ | |
| import logging | |
| import os | |
| from typing import Any, Dict, List, Optional | |
| from urllib.parse import urljoin | |
| from authlib.integrations.flask_client import OAuth | |
| from potato.authentication import AuthBackend | |
| logger = logging.getLogger(__name__) | |
| # Well-known OAuth provider configurations | |
| PROVIDER_CONFIGS = { | |
| "google": { | |
| "display_name": "Google", | |
| "server_metadata_url": "https://accounts.google.com/.well-known/openid-configuration", | |
| "client_kwargs": {"scope": "openid email profile"}, | |
| "icon_class": "fab fa-google", | |
| "button_class": "oauth-btn-google", | |
| }, | |
| "github": { | |
| "display_name": "GitHub", | |
| "api_base_url": "https://api.github.com/", | |
| "access_token_url": "https://github.com/login/oauth/access_token", | |
| "authorize_url": "https://github.com/login/oauth/authorize", | |
| "client_kwargs": {"scope": "user:email"}, | |
| "icon_class": "fab fa-github", | |
| "button_class": "oauth-btn-github", | |
| }, | |
| "huggingface": { | |
| "display_name": "HuggingFace", | |
| "server_metadata_url": "https://huggingface.co/.well-known/openid-configuration", | |
| "client_kwargs": {"scope": "openid profile email"}, | |
| "icon_class": "fas fa-robot", | |
| "button_class": "oauth-btn-huggingface", | |
| }, | |
| } | |
| class OAuthBackend(AuthBackend): | |
| """Authentication backend using OAuth 2.0 / OIDC providers. | |
| This backend delegates authentication to external identity providers | |
| (Google, GitHub, or any OIDC-compliant provider). After a user | |
| authenticates with the provider, they are mapped to a Potato user | |
| identity based on the configured identity field (email, username, etc.). | |
| """ | |
| def __init__(self, auth_config: Dict[str, Any]): | |
| """Initialize the OAuth backend from authentication config. | |
| Args: | |
| auth_config: The 'authentication' section of the Potato YAML config. | |
| Must contain 'providers' dict with at least one provider. | |
| """ | |
| self.users = {} # username -> user profile data | |
| self.providers_config = auth_config.get("providers", {}) | |
| self.user_identity_field = auth_config.get("user_identity_field", "email") | |
| self.auto_register = auth_config.get("auto_register", True) | |
| self.allow_local_login = auth_config.get("allow_local_login", False) | |
| self._oauth = None # Initialized later via init_oauth() | |
| self._provider_metadata = {} # Cached display info per provider | |
| if not self.providers_config: | |
| raise ValueError( | |
| "OAuth authentication requires at least one provider in " | |
| "'authentication.providers'" | |
| ) | |
| # Build provider metadata for login page rendering | |
| for name, pconfig in self.providers_config.items(): | |
| well_known = PROVIDER_CONFIGS.get(name, {}) | |
| self._provider_metadata[name] = { | |
| "name": name, | |
| "display_name": pconfig.get( | |
| "display_name", well_known.get("display_name", name.title()) | |
| ), | |
| "icon_class": well_known.get("icon_class", "fas fa-sign-in-alt"), | |
| "button_class": well_known.get("button_class", "oauth-btn-generic"), | |
| } | |
| logger.info( | |
| "OAuth backend initialized with %d provider(s)", | |
| len(self.providers_config), | |
| ) | |
| def init_oauth(self, app): | |
| """Register OAuth providers with the Flask app. | |
| Must be called after the Flask app is created but before any | |
| OAuth routes are used. | |
| Args: | |
| app: The Flask application instance. | |
| """ | |
| self._oauth = OAuth(app) | |
| for name, pconfig in self.providers_config.items(): | |
| well_known = PROVIDER_CONFIGS.get(name, {}) | |
| # Build the registration kwargs | |
| reg_kwargs = { | |
| "client_id": self._resolve_env(pconfig.get("client_id", "")), | |
| "client_secret": self._resolve_env(pconfig.get("client_secret", "")), | |
| } | |
| # For well-known providers, use built-in config | |
| if name == "google": | |
| reg_kwargs["server_metadata_url"] = well_known["server_metadata_url"] | |
| reg_kwargs["client_kwargs"] = well_known["client_kwargs"] | |
| elif name == "github": | |
| reg_kwargs["api_base_url"] = well_known["api_base_url"] | |
| reg_kwargs["access_token_url"] = well_known["access_token_url"] | |
| reg_kwargs["authorize_url"] = well_known["authorize_url"] | |
| # GitHub needs read:org scope if allowed_org is set | |
| scope = "user:email" | |
| if pconfig.get("allowed_org"): | |
| scope = "user:email read:org" | |
| reg_kwargs["client_kwargs"] = {"scope": scope} | |
| elif name == "oidc" or "discovery_url" in pconfig: | |
| # Generic OIDC provider | |
| discovery_url = pconfig.get("discovery_url", "") | |
| if not discovery_url: | |
| raise ValueError( | |
| f"OIDC provider '{name}' requires 'discovery_url'" | |
| ) | |
| reg_kwargs["server_metadata_url"] = discovery_url | |
| scopes = pconfig.get("scopes", ["openid", "email", "profile"]) | |
| reg_kwargs["client_kwargs"] = {"scope": " ".join(scopes)} | |
| self._oauth.register(name, **reg_kwargs) | |
| logger.info("Registered OAuth provider: %s", name) | |
| def get_oauth_client(self, provider_name: str): | |
| """Get the Authlib OAuth client for a provider. | |
| Args: | |
| provider_name: The provider key (e.g. 'google', 'github', 'oidc'). | |
| Returns: | |
| The Authlib OAuth client, or None if provider not found. | |
| """ | |
| if self._oauth is None: | |
| logger.error("OAuth not initialized — call init_oauth(app) first") | |
| return None | |
| client = getattr(self._oauth, provider_name, None) | |
| if client is None: | |
| logger.error("Unknown OAuth provider: %s", provider_name) | |
| return client | |
| def get_login_providers(self) -> List[Dict[str, str]]: | |
| """Return display metadata for all configured providers. | |
| Used by the login template to render SSO buttons. | |
| Returns: | |
| List of dicts with keys: name, display_name, icon_class, button_class, login_url | |
| """ | |
| providers = [] | |
| for name, meta in self._provider_metadata.items(): | |
| providers.append({ | |
| **meta, | |
| "login_url": f"/auth/login/{name}", | |
| }) | |
| return providers | |
| def extract_user_id(self, profile: Dict[str, Any], provider_name: str = None) -> str: | |
| """Extract the Potato user identity from an OAuth profile. | |
| Args: | |
| profile: The user profile dict from the OAuth provider. | |
| provider_name: The provider name (for fallback logic). | |
| Returns: | |
| The user identity string. | |
| Raises: | |
| ValueError: If no usable identity field is found. | |
| """ | |
| field = self.user_identity_field | |
| # Try the configured field first | |
| if field == "email" and profile.get("email"): | |
| return profile["email"] | |
| if field == "username": | |
| # GitHub uses 'login', others may use 'preferred_username' | |
| for key in ("login", "preferred_username", "username"): | |
| if profile.get(key): | |
| return profile[key] | |
| if field == "sub" and profile.get("sub"): | |
| return str(profile["sub"]) | |
| if field == "name" and profile.get("name"): | |
| return profile["name"] | |
| # Fallback chain: email -> login -> sub | |
| if profile.get("email"): | |
| return profile["email"] | |
| if profile.get("login"): | |
| return profile["login"] | |
| if profile.get("sub"): | |
| return str(profile["sub"]) | |
| raise ValueError( | |
| f"Cannot extract user identity from OAuth profile. " | |
| f"Configured field '{field}' not found, and no fallback available. " | |
| f"Profile keys: {list(profile.keys())}" | |
| ) | |
| def check_restrictions(self, provider_name: str, profile: Dict[str, Any]) -> tuple: | |
| """Check domain/org restrictions for a provider. | |
| Args: | |
| provider_name: The provider key. | |
| profile: The user profile dict. | |
| Returns: | |
| Tuple of (allowed: bool, reason: str). reason is empty if allowed. | |
| """ | |
| pconfig = self.providers_config.get(provider_name, {}) | |
| # Check Google domain restriction | |
| allowed_domain = pconfig.get("allowed_domain") | |
| if allowed_domain: | |
| email = profile.get("email", "") | |
| domain = email.split("@")[-1] if "@" in email else "" | |
| if domain.lower() != allowed_domain.lower(): | |
| return False, ( | |
| f"Access restricted to {allowed_domain} accounts. " | |
| f"Your email domain ({domain}) is not allowed." | |
| ) | |
| # GitHub org restriction is checked separately via API (see routes) | |
| # We store the config here for routes to use | |
| # Note: allowed_org check requires an API call with the user's token, | |
| # which happens in the route handler, not here. | |
| return True, "" | |
| # --- AuthBackend interface --- | |
| def authenticate(self, username: str, password: Optional[str]) -> bool: | |
| """Authenticate an OAuth user. | |
| For OAuth, authentication happens via the provider redirect flow, | |
| not via username/password. This method returns True if the user | |
| exists (was previously authenticated via OAuth). | |
| """ | |
| return username in self.users | |
| def add_user(self, username: str, password: Optional[str], **kwargs) -> str: | |
| """Register an OAuth-authenticated user.""" | |
| if username in self.users: | |
| # Update profile data | |
| self.users[username].update(kwargs) | |
| return "Success" | |
| self.users[username] = kwargs | |
| return "Success" | |
| def is_valid_username(self, username: str) -> bool: | |
| """Check if a username was registered via OAuth.""" | |
| return username in self.users | |
| def update_password(self, username: str, new_password: str) -> bool: | |
| """Not supported for OAuth - passwords are managed by providers.""" | |
| raise NotImplementedError("Password management is handled by OAuth providers") | |
| def get_all_users(self) -> list: | |
| """Return all registered OAuth usernames.""" | |
| return list(self.users.keys()) | |
| def get_allowed_org(self, provider_name: str) -> Optional[str]: | |
| """Get the allowed_org restriction for a provider, if any.""" | |
| return self.providers_config.get(provider_name, {}).get("allowed_org") | |
| # --- Helpers --- | |
| def _resolve_env(value: str) -> str: | |
| """Resolve ${ENV_VAR} references in a config value.""" | |
| if not isinstance(value, str): | |
| return value | |
| if value.startswith("${") and value.endswith("}"): | |
| env_name = value[2:-1] | |
| resolved = os.environ.get(env_name, "") | |
| if not resolved: | |
| logger.warning( | |
| "Environment variable referenced in OAuth config is not set" | |
| ) | |
| return resolved | |
| return value | |