File size: 13,938 Bytes
4a823b5 |
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 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 |
"""
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')}"
|