Spaces:
Sleeping
Sleeping
| """ | |
| Authentication module for Dynamic Highscores system. | |
| This module handles user authentication with HuggingFace, | |
| user session management, and access control. | |
| """ | |
| import os | |
| import json | |
| import time | |
| import requests | |
| import gradio as gr | |
| from huggingface_hub import HfApi, login | |
| from functools import wraps | |
| class HuggingFaceAuth: | |
| """Authentication manager for HuggingFace integration.""" | |
| def __init__(self, db_manager): | |
| """Initialize the authentication manager. | |
| Args: | |
| db_manager: Database manager instance for user storage | |
| """ | |
| self.db_manager = db_manager | |
| self.hf_api = HfApi() | |
| self.admin_username = os.environ.get("ADMIN_USERNAME", "Quazim0t0") | |
| self.running_in_space = 'SPACE_ID' in os.environ | |
| def login_user(self, token): | |
| """Log in a user with their HuggingFace token. | |
| Args: | |
| token: HuggingFace API token | |
| Returns: | |
| dict: User information if login successful, None otherwise | |
| """ | |
| try: | |
| # Validate token with HuggingFace | |
| login(token=token, add_to_git_credential=False) | |
| # Get user info from HuggingFace | |
| user_info = self.hf_api.whoami(token=token) | |
| if not user_info: | |
| return None | |
| # Check if user exists in our database, create if not | |
| username = user_info.get("name", user_info.get("fullname", "")) | |
| hf_user_id = user_info.get("id", "") | |
| if not hf_user_id: | |
| return None | |
| # Check if this is the admin account | |
| is_admin = (username == self.admin_username) | |
| # Add or get user from database | |
| user_id = self.db_manager.add_user(username, hf_user_id, is_admin) | |
| # Get complete user info from database | |
| user = self.db_manager.get_user(hf_user_id) | |
| if user: | |
| # Add token to user info for session only (not stored in database) | |
| user['token'] = token | |
| return user | |
| return None | |
| except Exception as e: | |
| print(f"Login error: {e}") | |
| return None | |
| def check_login(self, request: gr.Request): | |
| """Check if a user is logged in from a Gradio request. | |
| Args: | |
| request: Gradio request object | |
| Returns: | |
| dict: User information if logged in, None otherwise | |
| """ | |
| if not request: | |
| return None | |
| # First, check if we're in a HuggingFace Space with OAuth | |
| if self.running_in_space: | |
| # Check for HF-User header from Space OAuth | |
| username = request.headers.get("HF-User") | |
| if username: | |
| # Check if user exists in our database, create if not | |
| user = self.db_manager.get_user_by_username(username) | |
| if not user: | |
| # Create a new user | |
| is_admin = (username == self.admin_username) | |
| user_id = self.db_manager.add_user(username, username, is_admin) | |
| user = self.db_manager.get_user_by_username(username) | |
| return user | |
| # Fallback to token-based auth for local development | |
| token = request.cookies.get("hf_token") | |
| if not token: | |
| return None | |
| try: | |
| # Validate token with HuggingFace | |
| user_info = self.hf_api.whoami(token=token) | |
| if not user_info: | |
| return None | |
| # Get user from database | |
| hf_user_id = user_info.get("id", "") | |
| user = self.db_manager.get_user(hf_user_id) | |
| if user: | |
| # Add token to user info for session only (not stored in database) | |
| user['token'] = token | |
| return user | |
| return None | |
| except Exception as e: | |
| print(f"Check login error: {e}") | |
| return None | |
| def require_login(self, func): | |
| """Decorator to require login for a function. | |
| Args: | |
| func: Function to decorate | |
| Returns: | |
| Function: Decorated function that requires login | |
| """ | |
| def wrapper(*args, **kwargs): | |
| # Find the request argument | |
| request = None | |
| for arg in args: | |
| if isinstance(arg, gr.Request): | |
| request = arg | |
| break | |
| if not request and 'request' in kwargs: | |
| request = kwargs['request'] | |
| if not request: | |
| return "Please log in to access this feature." | |
| # Check if user is logged in | |
| user = self.check_login(request) | |
| if not user: | |
| return "Please log in to access this feature." | |
| # Add user to kwargs | |
| kwargs['user'] = user | |
| # Call the original function | |
| return func(*args, **kwargs) | |
| return wrapper | |
| def require_admin(self, func): | |
| """Decorator to require admin privileges for a function. | |
| Args: | |
| func: Function to decorate | |
| Returns: | |
| Function: Decorated function that requires admin privileges | |
| """ | |
| def wrapper(*args, **kwargs): | |
| # Find the request argument | |
| request = None | |
| for arg in args: | |
| if isinstance(arg, gr.Request): | |
| request = arg | |
| break | |
| if not request and 'request' in kwargs: | |
| request = kwargs['request'] | |
| if not request: | |
| return "Admin access required." | |
| # Check if user is logged in | |
| user = self.check_login(request) | |
| if not user: | |
| return "Admin access required." | |
| # Check if user is admin | |
| if not user.get('is_admin', False): | |
| return "Admin access required." | |
| # Add user to kwargs | |
| kwargs['user'] = user | |
| # Call the original function | |
| return func(*args, **kwargs) | |
| return wrapper | |
| def can_submit_benchmark(self, user_id): | |
| """Check if a user can submit a benchmark today. | |
| Args: | |
| user_id: User ID to check | |
| Returns: | |
| bool: True if user can submit, False otherwise | |
| """ | |
| return self.db_manager.can_submit_today(user_id) | |
| def update_submission_date(self, user_id): | |
| """Update the last submission date for a user. | |
| Args: | |
| user_id: User ID to update | |
| """ | |
| self.db_manager.update_submission_date(user_id) | |
| # Authentication UI components | |
| def create_login_ui(): | |
| """Create the login UI components. | |
| Returns: | |
| tuple: (login_button, logout_button, user_info) | |
| """ | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| # If running in a HuggingFace Space, use their OAuth | |
| if 'SPACE_ID' in os.environ: | |
| login_button = gr.Button("Login with HuggingFace", visible=False) | |
| logout_button = gr.Button("Logout", visible=False) | |
| else: | |
| # For local development, use token-based login | |
| login_button = gr.Button("Login with HuggingFace Token") | |
| logout_button = gr.Button("Logout", visible=False) | |
| with gr.Column(scale=2): | |
| user_info = gr.Markdown("Checking login status...") | |
| return login_button, logout_button, user_info | |
| def login_handler(auth_manager): | |
| """Handle login button click. | |
| Args: | |
| auth_manager: Authentication manager instance | |
| Returns: | |
| tuple: JS to redirect to login and updated UI visibility | |
| """ | |
| # This is only used for local development | |
| # For HuggingFace Spaces, the built-in OAuth is used | |
| return ( | |
| gr.update(visible=False), # Hide login button | |
| gr.update(visible=True), # Show logout button | |
| "Redirecting to login...", | |
| """ | |
| <script> | |
| // Open a popup window for token entry | |
| function promptForToken() { | |
| const token = prompt("Enter your HuggingFace token:"); | |
| if (token) { | |
| // Set the token as a cookie | |
| document.cookie = "hf_token=" + token + "; path=/; SameSite=Strict"; | |
| // Reload the page to apply the token | |
| window.location.reload(); | |
| } | |
| } | |
| // Call the function | |
| promptForToken(); | |
| </script> | |
| """ | |
| ) | |
| def logout_handler(): | |
| """Handle logout button click. | |
| Returns: | |
| tuple: Updated UI components visibility and user info | |
| """ | |
| # Clear token cookie in JavaScript | |
| return ( | |
| gr.update(visible=True), # Show login button | |
| gr.update(visible=False), # Hide logout button | |
| "Logged out", | |
| """ | |
| <script> | |
| // Clear the token cookie | |
| document.cookie = "hf_token=; path=/; max-age=0; SameSite=Strict"; | |
| // Reload the page | |
| window.location.reload(); | |
| </script> | |
| """ | |
| ) | |
| def setup_auth_handlers(login_button, logout_button, user_info, auth_manager): | |
| """Set up event handlers for authentication UI components. | |
| Args: | |
| login_button: Login button component | |
| logout_button: Logout button component | |
| user_info: User info component | |
| auth_manager: Authentication manager instance | |
| """ | |
| # Only add event handlers if not running in a HuggingFace Space | |
| if 'SPACE_ID' not in os.environ: | |
| login_button.click( | |
| fn=lambda: login_handler(auth_manager), | |
| inputs=[], | |
| outputs=[login_button, logout_button, user_info, gr.HTML()] | |
| ) | |
| logout_button.click( | |
| fn=logout_handler, | |
| inputs=[], | |
| outputs=[login_button, logout_button, user_info, gr.HTML()] | |
| ) |