token_auth.py

"""
Authentication module for Dynamic Highscores system with token-based authentication.

This module handles user authentication with HuggingFace using direct token input,
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 HuggingFaceTokenAuth:
    """Authentication manager for HuggingFace integration using tokens."""
    
    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
        
        print(f"Running in Space: {self.running_in_space}")
    
    def validate_token(self, token):
        """Validate a HuggingFace token.
        
        Args:
            token: HuggingFace API token
            
        Returns:
            dict: User information if token is valid, None otherwise
        """
        try:
            # Validate token with HuggingFace
            user_info = self.hf_api.whoami(token=token)
            return user_info
        except Exception as e:
            print(f"Token validation error: {e}")
            return None
    
    def login_user(self, read_token, write_token=None):
        """Log in a user with their HuggingFace tokens.
        
        Args:
            read_token: HuggingFace read API token
            write_token: HuggingFace write API token (optional)
            
        Returns:
            dict: User information if login successful, None otherwise
        """
        try:
            # Validate read token with HuggingFace
            user_info = self.validate_token(read_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 tokens to user info for session only (not stored in database)
                user['read_token'] = read_token
                if write_token:
                    user['write_token'] = write_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:
            # Try to get token from 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['read_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 enter your HuggingFace tokens to access this feature."
            
            # Check if user is logged in
            user = self.check_login(request)
            
            if not user:
                return "Please enter your HuggingFace tokens 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)

# Token input UI components
def create_token_input_ui():
    """Create the token input UI components.
    
    Returns:
        tuple: (read_token, write_token, save_button, clear_button, token_status, token_state, token_js)
    """
    with gr.Row():
        with gr.Column():
            gr.Markdown("### HuggingFace Token Authentication")
            gr.Markdown("""
            Enter your HuggingFace tokens to use this application. 
            You can find your tokens in your [HuggingFace settings](https://huggingface.co/settings/tokens).
            
            - **Read Token**: Required for accessing models and datasets
            - **Write Token**: Required for submitting evaluation results
            
            Your tokens are stored only in your browser's local storage and are not saved on the server.
            """)
            
            read_token = gr.Textbox(
                label="Read Token", 
                placeholder="Enter your HuggingFace read token", 
                type="password"
            )
            write_token = gr.Textbox(
                label="Write Token", 
                placeholder="Enter your HuggingFace write token", 
                type="password"
            )
            save_button = gr.Button("Save Tokens")
            clear_button = gr.Button("Clear Tokens")
            token_status = gr.Markdown("Not authenticated")
            
            # Hidden field to store the token status
            token_state = gr.State(None)
    
    # JavaScript to handle token storage
    token_js = """
    <script>
    // Function to save tokens to localStorage
    function saveTokens() {
        const readToken = document.querySelector('input[placeholder="Enter your HuggingFace read token"]').value;
        const writeToken = document.querySelector('input[placeholder="Enter your HuggingFace write token"]').value;
        
        if (readToken && writeToken) {
            localStorage.setItem("hf_read_token", readToken);
            localStorage.setItem("hf_write_token", writeToken);
            
            // Set token in cookie for server-side access
            document.cookie = "hf_token=" + readToken + "; path=/; SameSite=Strict";
            
            // Update status
            const statusElement = document.querySelector('div[data-testid="markdown"] p');
            if (statusElement) {
                statusElement.textContent = "Authenticated with tokens";
                statusElement.style.color = "green";
            }
            
            // Reload page to apply tokens
            setTimeout(() => window.location.reload(), 1000);
        } else {
            alert("Please enter both read and write tokens");
        }
    }
    
    // Function to clear tokens from localStorage
    function clearTokens() {
        localStorage.removeItem("hf_read_token");
        localStorage.removeItem("hf_write_token");
        
        // Clear token cookie
        document.cookie = "hf_token=; path=/; max-age=0; SameSite=Strict";
        
        // Update status
        const statusElement = document.querySelector('div[data-testid="markdown"] p');
        if (statusElement) {
            statusElement.textContent = "Not authenticated";
            statusElement.style.color = "red";
        }
        
        // Clear input fields
        document.querySelector('input[placeholder="Enter your HuggingFace read token"]').value = "";
        document.querySelector('input[placeholder="Enter your HuggingFace write token"]').value = "";
        
        // Reload page to apply changes
        setTimeout(() => window.location.reload(), 1000);
    }
    
    // Function to load tokens from localStorage
    function loadTokens() {
        const readToken = localStorage.getItem("hf_read_token");
        const writeToken = localStorage.getItem("hf_write_token");
        
        if (readToken && writeToken) {
            document.querySelector('input[placeholder="Enter your HuggingFace read token"]').value = readToken;
            document.querySelector('input[placeholder="Enter your HuggingFace write token"]').value = writeToken;
            
            // Update status
            const statusElement = document.querySelector('div[data-testid="markdown"] p');
            if (statusElement) {
                statusElement.textContent = "Authenticated with tokens";
                statusElement.style.color = "green";
            }
            
            // Set token in cookie for server-side access if not already set
            if (!document.cookie.includes("hf_token=")) {
                document.cookie = "hf_token=" + readToken + "; path=/; SameSite=Strict";
            }
        }
    }
    
    // Add event listeners once DOM is loaded
    document.addEventListener("DOMContentLoaded", function() {
        // Load tokens from localStorage
        loadTokens();
        
        // Add event listeners to buttons
        const saveButton = document.querySelector('button:nth-of-type(1)');
        const clearButton = document.querySelector('button:nth-of-type(2)');
        
        if (saveButton) {
            saveButton.addEventListener("click", saveTokens);
        }
        
        if (clearButton) {
            clearButton.addEventListener("click", clearTokens);
        }
    });
    </script>
    """
    
    return read_token, write_token, save_button, clear_button, token_status, token_state, token_js

def validate_tokens(auth_manager, read_token, write_token):
    """Validate HuggingFace tokens.
    
    Args:
        auth_manager: Authentication manager instance
        read_token: HuggingFace read token
        write_token: HuggingFace write token
        
    Returns:
        str: Status message
    """
    if not read_token or not write_token:
        return "Please enter both read and write tokens"
    
    # Validate read token
    read_user_info = auth_manager.validate_token(read_token)
    if not read_user_info:
        return "Invalid read token"
    
    # Validate write token
    write_user_info = auth_manager.validate_token(write_token)
    if not write_user_info:
        return "Invalid write token"
    
    # Check if tokens belong to the same user
    if read_user_info.get("id") != write_user_info.get("id"):
        return "Tokens must belong to the same user"
    
    # Login user with tokens
    user = auth_manager.login_user(read_token, write_token)
    if not user:
        return "Failed to authenticate user"
    
    return f"Authenticated as {user.get('username')}"