auth.py

"""
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
        
        # OAuth configuration
        self.client_id = os.environ.get("OAUTH_CLIENT_ID", "")
        self.client_secret = os.environ.get("OAUTH_CLIENT_SECRET", "")
        self.oauth_scopes = os.environ.get("OAUTH_SCOPES", "openid profile")
        self.openid_provider_url = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co")
        
        print(f"Running in Space: {self.running_in_space}")
        if self.running_in_space:
            print(f"OAuth Client ID: {self.client_id}")
            print(f"OAuth Scopes: {self.oauth_scopes}")
            print(f"OpenID Provider URL: {self.openid_provider_url}")
    
    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:
                print(f"Found HF-User header: {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:
            # Try to get token from localStorage via headers
            token = request.headers.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
        """
        @wraps(func)
        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
        """
        @wraps(func)
        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)
    
    def get_oauth_login_url(self, redirect_uri=None):
        """Get the OAuth login URL for HuggingFace.
        
        Args:
            redirect_uri: Redirect URI after login (optional)
            
        Returns:
            str: OAuth login URL
        """
        if not self.client_id:
            return None
        
        if not redirect_uri:
            space_host = os.environ.get("SPACE_HOST", "localhost:7860")
            redirect_uri = f"https://{space_host}"
        
        # Create a random state for security
        state = os.urandom(16).hex()
        
        # Build the OAuth URL
        scopes = self.oauth_scopes.replace(" ", "%20")
        auth_url = f"{self.openid_provider_url}/oauth/authorize?client_id={self.client_id}&redirect_uri={redirect_uri}&scope={scopes}&state={state}&response_type=code"
        
        return auth_url

# 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";
        // Clear localStorage
        localStorage.removeItem("hf_user");
        localStorage.removeItem("hf_token");
        // 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()]
        )