File size: 12,328 Bytes
9a46619 acc68d6 9a46619 acc68d6 9a46619 acc68d6 9a46619 acc68d6 9a46619 acc68d6 9a46619 acc68d6 |
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 |
"""
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()]
)
|