import json import os import re import shutil import asyncio from huggingface_hub import HfApi, hf_hub_download import config import tempfile DATA_FILE_NAME = "database.json" LOCAL_PATH = os.path.join(tempfile.gettempdir(), DATA_FILE_NAME) # In-memory database state db_state = { "admins": [], # List of admin user IDs (ints) "users": [], # List of authorized user IDs (ints) "urls": [], # List of dicts: {"url": str, "added_by": int} "user_metadata": {} # Dict: "user_id_str" -> {"username": str, "name": str} } # Helper to extract repository ID from a HF URL def extract_repo_id(url: str) -> str: if not url: return None url = url.strip() if "/" in url and not url.startswith("http"): return url # Match standard dataset url: huggingface.co/datasets/username/dataset-name match = re.search(r"huggingface\.co/datasets/([^/]+)/([^/]+)", url, re.IGNORECASE) if match: return f"{match.group(1)}/{match.group(2)}" return None repo_id = extract_repo_id(config.HF_DATASET_URL) def get_api(): if not config.HF_TOKEN: return None return HfApi(token=config.HF_TOKEN) def init_db(): global db_state if not repo_id or not config.HF_TOKEN: print("HF_DATASET_URL or HF_TOKEN not configured. Working locally.") if os.path.exists(LOCAL_PATH): try: with open(LOCAL_PATH, "r", encoding="utf-8") as f: loaded = json.load(f) # Migrate or update local state for key in db_state: if key in loaded: db_state[key] = loaded[key] print("Loaded database from local file.") except Exception as e: print(f"Error loading local database: {e}") return # Sync from Hugging Face Dataset try: api = get_api() # Ensure repository exists api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True) # Download file print(f"Downloading database from HF Dataset: {repo_id}") downloaded_path = hf_hub_download( repo_id=repo_id, filename=DATA_FILE_NAME, repo_type="dataset", token=config.HF_TOKEN ) # Copy to local path shutil.copy(downloaded_path, LOCAL_PATH) # Load JSON with open(LOCAL_PATH, "r", encoding="utf-8") as f: loaded = json.load(f) for key in db_state: if key in loaded: db_state[key] = loaded[key] print("Successfully synced database from Hugging Face.") except Exception as e: err_msg = str(e) if "404" in err_msg or "Entry Not Found" in err_msg: print("Database file not found in HF dataset. Initializing a new one.") else: print(f"Error loading database from HF: {err_msg}. Initializing a new one.") if os.path.exists(LOCAL_PATH): try: with open(LOCAL_PATH, "r", encoding="utf-8") as f: loaded = json.load(f) for key in db_state: if key in loaded: db_state[key] = loaded[key] except Exception: pass save_db() def save_db(): global db_state # Ensure types and clean duplicates db_state["admins"] = list(sorted(list(set(int(uid) for uid in db_state.get("admins", []))))) db_state["users"] = list(sorted(list(set(int(uid) for uid in db_state.get("users", []))))) unique_urls = [] seen = set() for item in db_state.get("urls", []): url = item.get("url") if url not in seen: seen.add(url) unique_urls.append({ "url": url, "added_by": int(item.get("added_by")) }) db_state["urls"] = unique_urls # Save locally try: with open(LOCAL_PATH, "w", encoding="utf-8") as f: json.dump(db_state, f, indent=2, ensure_ascii=False) print("Database saved locally.") except Exception as e: print(f"Error saving database locally: {e}") return # Upload to HF if not repo_id or not config.HF_TOKEN: return try: api = get_api() api.upload_file( path_or_fileobj=LOCAL_PATH, path_in_repo=DATA_FILE_NAME, repo_id=repo_id, repo_type="dataset", token=config.HF_TOKEN ) print("Database uploaded and synced to Hugging Face dataset.") except Exception as e: print(f"Error uploading database to Hugging Face: {e}") async def save_db_async(): await asyncio.to_thread(save_db) # Normalization of Hugging Face Space URLs def normalize_hf_url(url: str) -> str: url = url.strip() # Remove any query parameters or trailing slashes url = url.split("?")[0].rstrip("/") # Case 1: Match standard Hugging Face space URLs # e.g., https://huggingface.co/spaces/username/spacename # or huggingface.co/spaces/username/spacename/tree/main # We want to extract username and spacename. hf_co_pattern = r"(?:https?://)?(?:www\.)?huggingface\.co/spaces/([^/]+)/([^/]+)" match = re.match(hf_co_pattern, url, re.IGNORECASE) if match: username = match.group(1).lower().replace("_", "-").replace(".", "-") spacename = match.group(2).lower().replace("_", "-").replace(".", "-") return f"https://{username}-{spacename}.hf.space" # Case 2: Match direct hf.space URLs # e.g., https://username-spacename.hf.space # or username-spacename.hf.space/subpath hf_space_pattern = r"(?:https?://)?([^/.]+)\.hf\.space" match = re.match(hf_space_pattern, url, re.IGNORECASE) if match: subdomain = match.group(1).lower().replace("_", "-").replace(".", "-") return f"https://{subdomain}.hf.space" # Case 3: Match simple username/spacename form # e.g., username/spacename simple_pattern = r"^([^/]+)/([^/]+)$" match = re.match(simple_pattern, url) if match: username = match.group(1).lower().replace("_", "-").replace(".", "-") spacename = match.group(2).lower().replace("_", "-").replace(".", "-") return f"https://{username}-{spacename}.hf.space" raise ValueError("Invalid Hugging Face Space URL. Make sure it contains 'huggingface.co/spaces/...' or '.hf.space'") # Authorization helper functions def is_super_admin(user_id: int) -> bool: return config.SUPER_ADMIN_ID is not None and user_id == config.SUPER_ADMIN_ID def is_admin(user_id: int) -> bool: if is_super_admin(user_id): return True return user_id in db_state.get("admins", []) def is_authorized(user_id: int) -> bool: if is_admin(user_id): return True return user_id in db_state.get("users", []) # User metadata updates def update_user_metadata(user_id: int, username: str = None, first_name: str = None, last_name: str = None): if "user_metadata" not in db_state: db_state["user_metadata"] = {} name = first_name or "" if last_name: name += f" {last_name}" name = name.strip() user_id_str = str(user_id) # Check if there is any change to avoid redundant saves existing = db_state["user_metadata"].get(user_id_str) if not existing or existing.get("username") != username or existing.get("name") != name: db_state["user_metadata"][user_id_str] = { "username": username, "name": name } save_db() def get_user_display(user_id: int) -> str: user_id_str = str(user_id) metadata = db_state.get("user_metadata", {}).get(user_id_str) if metadata: username = metadata.get("username") name = metadata.get("name") if username: return f"{name} (@{username}) [`{user_id}`]" elif name: return f"{name} [`{user_id}`]" return f"`{user_id}`" # User administration functions def add_admin(user_id: int) -> bool: if user_id in db_state["admins"]: return False db_state["admins"].append(user_id) save_db() return True def remove_admin(user_id: int) -> bool: if user_id not in db_state["admins"]: return False db_state["admins"].remove(user_id) save_db() return True def add_user(user_id: int) -> bool: if user_id in db_state["users"]: return False db_state["users"].append(user_id) save_db() return True def remove_user(user_id: int) -> bool: if user_id not in db_state["users"]: return False db_state["users"].remove(user_id) save_db() return True def list_admins() -> list: return db_state["admins"] def list_users() -> list: return db_state["users"] # URL management functions def add_url(url: str, added_by: int) -> tuple[bool, str]: normalized = normalize_hf_url(url) # Check if already exists for item in db_state["urls"]: if item["url"] == normalized: return False, normalized db_state["urls"].append({ "url": normalized, "added_by": added_by }) save_db() return True, normalized def remove_url(url: str, user_id: int) -> tuple[bool, str]: try: normalized = normalize_hf_url(url) except ValueError: normalized = url.strip() # Find the URL target_idx = -1 for idx, item in enumerate(db_state["urls"]): if item["url"] == normalized or item["url"].rstrip("/") == normalized.rstrip("/"): target_idx = idx break if target_idx == -1: return False, normalized item = db_state["urls"][target_idx] # Check authorization # Super admins and admins can remove any URL. # Regular users can only remove their own. if not is_admin(user_id) and item["added_by"] != user_id: raise PermissionError("You can only remove URLs that you added.") db_state["urls"].pop(target_idx) save_db() return True, normalized def get_urls(user_id: int) -> list: # If admin or super admin, return all if is_admin(user_id): return db_state["urls"] # Else, return only owned by this user return [item for item in db_state["urls"] if item["added_by"] == user_id] def get_all_urls() -> list: return [item["url"] for item in db_state["urls"]]