Spaces:
Running
Running
| import json | |
| import time | |
| from huggingface_hub import HfApi, hf_hub_download | |
| from werkzeug.security import generate_password_hash, check_password_hash | |
| import gradio as gr | |
| from retry import retry # Requires `pip install retry` | |
| # Handle HfHubHTTPError import gracefully (it's in utils, and may vary by version) | |
| try: | |
| from huggingface_hub.utils import HfHubHTTPError | |
| except ImportError: | |
| from requests.exceptions import HTTPError as HfHubHTTPError # Fallback to base HTTPError for compatibility | |
| # ================== ALLOWED USERS + AUDIT LOG ================== | |
| def is_allowed_username(username, auth_globals): | |
| username = username.strip().lower() | |
| try: | |
| path = hf_hub_download(repo_id=auth_globals["user_repo"], filename="allowed_users.json", repo_type="dataset", token=auth_globals["hf_token"]) | |
| with open(path, 'r') as f: | |
| allowed = json.load(f) | |
| return username in [u.strip().lower() for u in allowed] | |
| except HfHubHTTPError as e: | |
| gr.Warning(f"Failed to check allowed users: {str(e)}") | |
| return False | |
| except Exception as e: | |
| gr.Warning(f"Unexpected error in is_allowed_username: {str(e)}") | |
| return False | |
| # Retry on transient errors | |
| def log_activity(action, username, details="", auth_globals=None): | |
| if not auth_globals or not username: | |
| return | |
| try: | |
| hf_api = HfApi(token=auth_globals["hf_token"]) | |
| user_repo = auth_globals["user_repo"] | |
| try: | |
| path = hf_hub_download(repo_id=user_repo, filename="activity_log.json", repo_type="dataset", token=auth_globals["hf_token"]) | |
| with open(path, 'r') as f: | |
| log = json.load(f) | |
| except HfHubHTTPError: | |
| log = [] # Initialize if file doesn't exist | |
| entry = { | |
| "timestamp": time.time(), | |
| "datetime": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), | |
| "action": action, | |
| "username": username.strip().lower(), | |
| "details": details | |
| } | |
| log.append(entry) | |
| hf_api.upload_file( | |
| path_or_fileobj=json.dumps(log, indent=2).encode(), | |
| path_in_repo="activity_log.json", | |
| repo_id=user_repo, | |
| repo_type="dataset", | |
| token=auth_globals["hf_token"] | |
| ) | |
| except Exception as e: | |
| gr.Warning(f"Failed to log activity: {str(e)}") # Better than silent failure | |
| def get_total_accounts(auth_globals): | |
| try: | |
| hf_api = HfApi(token=auth_globals["hf_token"]) | |
| files = hf_api.list_repo_files(repo_id=auth_globals["user_repo"], repo_type="dataset") | |
| return sum(1 for f in files if f.endswith('.json') and f not in ["allowed_users.json", "activity_log.json"]) | |
| except HfHubHTTPError as e: | |
| gr.Warning(f"Failed to get total accounts: {str(e)}") | |
| return 0 | |
| except Exception as e: | |
| gr.Warning(f"Unexpected error in get_total_accounts: {str(e)}") | |
| return 0 | |
| # ================== HELPER FOR FAILURE RETURNS ================== | |
| def failure_return(logged_in_username, message, user_chats=None, current_title=None, history=None, dropdown=None, context_tb=None, save_btn=None): | |
| # Defaults are set to None and computed here to avoid NameError on evaluation | |
| if user_chats is None: | |
| user_chats = [] | |
| if current_title is None: | |
| current_title = None | |
| if history is None: | |
| history = [] | |
| if dropdown is None: | |
| dropdown = gr.update(choices=[], value=None) | |
| if context_tb is None: | |
| context_tb = gr.update(visible=bool(logged_in_username), value="") | |
| if save_btn is None: | |
| save_btn = gr.update(visible=bool(logged_in_username)) | |
| return ( | |
| user_chats, current_title, history, dropdown, logged_in_username, f"❌ {message}", | |
| gr.update(visible=bool(logged_in_username)), gr.update(visible=False), gr.update(visible=False), context_tb, save_btn | |
| ) | |
| # ================== REGISTER ================== | |
| def register_user(username, password, logged_in_username, auth_globals): | |
| hf_token = auth_globals["hf_token"] | |
| user_repo = auth_globals["user_repo"] | |
| hf_api = HfApi(token=hf_token) | |
| username = username.strip().lower() | |
| if not username.isalnum(): # Added: Sanitize username (alphanumeric only) | |
| return failure_return(logged_in_username, "Username must be alphanumeric.") | |
| if get_total_accounts(auth_globals) >= 20: | |
| return failure_return(logged_in_username, "Maximum 20 accounts reached.") | |
| if not username or not password: | |
| return failure_return(logged_in_username, "Please fill username and password.") | |
| if not is_allowed_username(username, auth_globals): | |
| return failure_return(logged_in_username, "This username is not pre-approved by admin.") | |
| try: | |
| # Check if already exists | |
| hf_hub_download(repo_id=user_repo, filename=f"{username}.json", repo_type="dataset", token=hf_token) | |
| return failure_return(logged_in_username, "User already exists!") | |
| except HfHubHTTPError: # Expected if file doesn't exist | |
| pass | |
| except Exception as e: | |
| return failure_return(logged_in_username, f"Registration check failed: {str(e)}") | |
| try: | |
| user_data = { | |
| "password": generate_password_hash(password), | |
| "chats": [] | |
| } | |
| hf_api.upload_file( | |
| path_or_fileobj=json.dumps(user_data).encode('utf-8'), | |
| path_in_repo=f"{username}.json", | |
| repo_id=user_repo, | |
| repo_type="dataset", | |
| token=hf_token | |
| ) | |
| log_activity("REGISTER", username, "New account created", auth_globals) | |
| return [], None, [], gr.update(choices=[], value=None), username, "✅ Registration successful! Logged in.", gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True, value=""), gr.update(visible=True) | |
| except Exception as e: | |
| return failure_return(logged_in_username, f"Registration failed: {str(e)}") | |
| # ================== LOGIN ================== | |
| def login_user(username, password, logged_in_username, auth_globals): | |
| hf_token = auth_globals["hf_token"] | |
| user_repo = auth_globals["user_repo"] | |
| hf_api = HfApi(token=hf_token) | |
| username = username.strip().lower() | |
| if not username or not password: | |
| return failure_return(logged_in_username, "Please enter username and password.") | |
| if not is_allowed_username(username, auth_globals): | |
| return failure_return(logged_in_username, "This username is not pre-approved by admin.") | |
| try: | |
| user_file = hf_hub_download(repo_id=user_repo, filename=f"{username}.json", repo_type="dataset", token=hf_token) | |
| with open(user_file, 'r') as f: | |
| data = json.load(f) | |
| if check_password_hash(data["password"], password): | |
| chats = data.get("chats", []) | |
| if chats: | |
| current_title = chats[-1]["title"] | |
| history = chats[-1]["history"] | |
| context = chats[-1].get("context", "") # Added: Load context | |
| choices = [chat["title"] for chat in chats] | |
| value = current_title | |
| else: | |
| current_title = None | |
| history = [] | |
| context = "" | |
| choices = [] | |
| value = None | |
| dropdown = gr.update(choices=choices, value=value) | |
| log_activity("LOGIN", username, "Successful login", auth_globals) | |
| return chats, current_title, history, dropdown, username, "✅ Login successful!", gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True, value=context), gr.update(visible=True) | |
| else: | |
| return failure_return(logged_in_username, "Incorrect password.") | |
| except HfHubHTTPError as e: | |
| return failure_return(logged_in_username, f"Login failed: User not found ({str(e)}).") | |
| except Exception as e: | |
| return failure_return(logged_in_username, f"Login failed: {str(e)}") | |
| # ================== LOGOUT ================== | |
| def logout_user(): | |
| return [], None, [], gr.update(choices=[], value=None), None, "Logged out successfully!", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False, value=""), gr.update(visible=False) | |
| # ================== SAVE CHATS ================== | |
| # Added retry for reliability | |
| def save_chats(username, chats, auth_globals): | |
| if not username: | |
| return "No user logged in!" | |
| try: | |
| hf_token = auth_globals["hf_token"] | |
| user_repo = auth_globals["user_repo"] | |
| hf_api = HfApi(token=hf_token) | |
| user_file = hf_hub_download(repo_id=user_repo, filename=f"{username}.json", repo_type="dataset", token=hf_token) | |
| with open(user_file, 'r') as f: | |
| data = json.load(f) | |
| data["chats"] = chats # Chats now include "context" | |
| hf_api.upload_file(path_or_fileobj=json.dumps(data).encode('utf-8'), path_in_repo=f"{username}.json", repo_id=user_repo, repo_type="dataset", token=hf_token) | |
| log_activity("SAVE_CHATS", username, "Chats updated", auth_globals) # Added logging for saves | |
| return "Chats saved automatically!" | |
| except HfHubHTTPError as e: | |
| return f"Failed to save chats: {str(e)} (HF API error)" | |
| except Exception as e: | |
| return f"Failed to save chats: {str(e)}" | |
| # ================== DELETE ACCOUNT ================== | |
| def delete_account(confirm_username, logged_in_username, auth_globals): | |
| confirm_username = confirm_username.strip().lower() | |
| if not logged_in_username or confirm_username != logged_in_username: | |
| return "❌ Username does not match.", logged_in_username, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None, None, None, None, gr.update(visible=True), gr.update(visible=True) | |
| if not is_allowed_username(logged_in_username, auth_globals): | |
| return "❌ Not allowed.", logged_in_username, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None, None, None, None, gr.update(visible=True), gr.update(visible=True) | |
| hf_token = auth_globals["hf_token"] | |
| user_repo = auth_globals["user_repo"] | |
| hf_api = HfApi(token=hf_token) | |
| try: | |
| hf_api.delete_file(path_in_repo=f"{logged_in_username}.json", repo_id=user_repo, repo_type="dataset", token=hf_token) | |
| log_activity("DELETE_ACCOUNT", logged_in_username, "Account permanently deleted", auth_globals) | |
| return "✅ Account permanently deleted.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), [], None, [], gr.update(choices=[], value=None), gr.update(visible=False, value=""), gr.update(visible=False) | |
| except HfHubHTTPError as e: | |
| return f"❌ Deletion failed: {str(e)}", logged_in_username, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None, None, None, None, gr.update(visible=True), gr.update(visible=True) | |
| except Exception as e: | |
| return f"❌ Deletion failed: {str(e)}", logged_in_username, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), None, None, None, None, gr.update(visible=True), gr.update(visible=True) |