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")
    
    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
        
        # Get token from cookies
        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
        """
        @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)

# Authentication UI components
def create_login_ui():
    """Create the login UI components.
    
    Returns:
        tuple: (login_button, logout_button, token_input, user_info)
    """
    with gr.Row():
        with gr.Column(scale=3):
            token_input = gr.Textbox(
                placeholder="Enter your HuggingFace token",
                label="HuggingFace Token",
                type="password",
                visible=True,
                info="Your token is only stored temporarily in browser session cookies and is never saved permanently"
            )
            login_button = gr.Button("Login")
            logout_button = gr.Button("Logout", visible=False)
        
        with gr.Column(scale=2):
            user_info = gr.Markdown("Not logged in")
    
    return login_button, logout_button, token_input, user_info

def login_handler(token, auth_manager):
    """Handle login button click.
    
    Args:
        token: HuggingFace token
        auth_manager: Authentication manager instance
        
    Returns:
        tuple: Updated UI components visibility and user info
    """
    if not token:
        return gr.update(visible=True), gr.update(visible=False), "Please enter your HuggingFace token"
    
    user = auth_manager.login_user(token)
    
    if user:
        # Set cookie in JavaScript with session-only flag (no persistent storage)
        # Cookie will expire when browser is closed
        js = f"""
        document.cookie = "hf_token={token}; path=/; SameSite=Strict";
        """
        
        # Return updated UI components
        return (
            gr.update(visible=False),  # Hide token input
            gr.update(visible=True),   # Show logout button
            f"Logged in as {user['username']}"  # Update user info
        )
    else:
        return (
            gr.update(visible=True),   # Keep token input visible
            gr.update(visible=False),  # Hide logout button
            "Login failed. Please check your token and try again."  # Update user info
        )

def logout_handler():
    """Handle logout button click.
    
    Returns:
        tuple: Updated UI components visibility and user info
    """
    # Clear cookie in JavaScript
    js = """
    document.cookie = "hf_token=; path=/; max-age=0; SameSite=Strict";
    """
    
    # Return updated UI components
    return (
        gr.update(visible=True),   # Show token input
        gr.update(visible=False),  # Hide logout button
        "Logged out"               # Update user info
    )

def setup_auth_handlers(login_button, logout_button, token_input, user_info, auth_manager):
    """Set up event handlers for authentication UI components.
    
    Args:
        login_button: Login button component
        logout_button: Logout button component
        token_input: Token input component
        user_info: User info component
        auth_manager: Authentication manager instance
    """
    login_button.click(
        fn=lambda token: login_handler(token, auth_manager),
        inputs=[token_input],
        outputs=[token_input, logout_button, user_info]
    )
    
    logout_button.click(
        fn=logout_handler,
        inputs=[],
        outputs=[token_input, logout_button, user_info]
    )