File size: 9,677 Bytes
6b3e7b5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 |
"""
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]
)
|