import os import json import io import hashlib import threading import time import logging import requests from requests.auth import HTTPBasicAuth from huggingface_hub import HfApi, hf_hub_download from utils import add_script_run_ctx TARGET_SPACE = os.getenv("TARGET_SPACE") HF_TOKEN = os.getenv("HF_TOKEN") DATASET_ID = os.getenv("DATASET_ID") class TokenManager: def __init__(self): self.api = HfApi(token=HF_TOKEN) self.config_file = "pollinations_tokens.json" self.tokens = self.load_tokens() self._migrate_data_structure() self.last_save_time = 0 self.lock = threading.Lock() self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True) add_script_run_ctx(self.monitor_thread) self.monitor_thread.start() def _migrate_data_structure(self): dirty = False cleaned_tokens = [] for t in self.tokens: if "key" in t and "token" not in t: t["token"] = t.pop("key") dirty = True if "alias" not in t: t["alias"] = "Unknown" if "status" not in t: t["status"] = "active" if "balance" not in t: t["balance"] = 0.0 if t.get("token"): cleaned_tokens.append(t) else: dirty = True self.tokens = cleaned_tokens if dirty: self.save_tokens() def load_tokens(self): try: return json.load(open(self.config_file, "r")) except: try: path = hf_hub_download(repo_id=DATASET_ID, filename=self.config_file, repo_type="dataset", token=HF_TOKEN) return json.load(open(path, "r")) except: return [] def save_tokens(self): # 限制保存频率 if time.time() - self.last_save_time < 5: return False try: with open(self.config_file, "w") as f: json.dump(self.tokens, f, indent=2) self.api.upload_file( path_or_fileobj=io.BytesIO(json.dumps(self.tokens, indent=2).encode("utf-8")), path_in_repo=self.config_file, repo_id=DATASET_ID, repo_type="dataset" ) self.last_save_time = time.time() return True except: return False def add_token(self, alias, token): with self.lock: existing = next((t for t in self.tokens if t["token"] == token), None) if existing: existing["alias"] = alias existing["status"] = "active" logging.info(f"🔄 Token updated: {alias}") else: self.tokens.append({"alias": alias, "token": token, "status": "active", "balance": 0.0}) logging.info(f"➕ New Token added: {alias}") self.save_tokens() def delete_token(self, idx): with self.lock: if 0 <= idx < len(self.tokens): del self.tokens[idx] self.save_tokens() def _check_balance_http(self, token): """ 查询余额。 返回: float: 余额数值 0.0: 明确的余额不足或Token无效(403) None: 网络错误，保留原状态 """ try: url = f"https://gen.pollinations.ai/account/balance?key={token}" headers = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" } resp = requests.get(url, headers=headers, timeout=10) if resp.status_code == 200: try: return float(resp.text.strip()) except: return float(resp.json().get('balance', 0)) # 如果是 403 Forbidden 或 401 Unauthorized，说明 Token 废了，视为余额 0 处理 if resp.status_code in [401, 403]: return 0.0 return None # 其他错误 (500等) 返回 None，不更新状态 except: return None def _monitor_loop(self): logging.info("🚀 TokenMonitor thread started.") while True: changes_detected = False # 浅拷贝列表以避免迭代时修改问题 with self.lock: tokens_to_check = list(enumerate(self.tokens)) for idx, t in tokens_to_check: token_str = t.get("token") if not token_str: continue bal = self._check_balance_http(token_str) with self.lock: # 再次确认索引有效性 if idx < len(self.tokens) and self.tokens[idx].get("token") == token_str: if bal is not None: # 更新余额显示 if abs(self.tokens[idx]["balance"] - bal) > 0.001: self.tokens[idx]["balance"] = bal changes_detected = True # [V30.38] 状态流转核心逻辑 # 规则：> 0.1 Active, <= 0.1 Exhausted current_status = self.tokens[idx]["status"] if bal > 0.1: if current_status != "active": logging.info(f"🔋 Token {self.tokens[idx]['alias']} 余额充足 ({bal})，状态激活 (Active)") self.tokens[idx]["status"] = "active" changes_detected = True else: if current_status != "exhausted": logging.warning(f"🪫 Token {self.tokens[idx]['alias']} 余额不足/无效 ({bal})，标记耗尽 (Exhausted)") self.tokens[idx]["status"] = "exhausted" changes_detected = True if changes_detected: self.save_tokens() time.sleep(60) # 60秒轮询一次 def get_active_tokens(self): """ [V30.38] 获取当前所有可用 Token 的列表 (快照) 执行层将拿着这个列表去轮询，而不依赖 Manager 每次分配 """ with self.lock: # 筛选出所有状态为 active 的 token # 按余额降序排列，尽可能先用钱多的 candidates = [t for t in self.tokens if t.get("status") == "active" and t.get("token")] candidates.sort(key=lambda x: x.get("balance", 0.0), reverse=True) return [t["token"] for t in candidates] class AuthManager: def __init__(self): self.api = HfApi(token=HF_TOKEN) self.user_file = "users_config.json" def load_users(self): try: if os.path.exists(self.user_file): mtime = os.path.getmtime(self.user_file) if time.time() - mtime < 600: return json.load(open(self.user_file, "r")) path = hf_hub_download(repo_id=DATASET_ID, filename=self.user_file, repo_type="dataset", token=HF_TOKEN, local_dir=".", force_download=True) return json.load(open(path, "r")) except: return {"admin": {"password": hashlib.sha256("admin".encode()).hexdigest(), "must_change_password": True}} def save_users(self, u): try: self.api.upload_file(path_or_fileobj=io.BytesIO(json.dumps(u, indent=2).encode("utf-8")), path_in_repo=self.user_file, repo_id=DATASET_ID, repo_type="dataset") with open(self.user_file, "w") as f: json.dump(u, f, indent=2) except: pass def login(self, u, p): d = self.load_users() if u in d and hashlib.sha256(p.encode()).hexdigest() == d[u]["password"]: return (True, d[u].get("must_change_password", False)) else: return (False, False) def update_password(self, u, np): d = self.load_users() d[u]["password"] = hashlib.sha256(np.encode()).hexdigest() d[u]["must_change_password"] = False self.save_users(d) return True class WebDAVManager: def __init__(self): self.api = HfApi(token=HF_TOKEN) self.config_file = "webdav_config.json" self.configs = self.load_configs() def load_configs(self): try: if os.path.exists(self.config_file): return json.load(open(self.config_file, "r")) path = hf_hub_download(repo_id=DATASET_ID, filename=self.config_file, repo_type="dataset", token=HF_TOKEN) return json.load(open(path, "r")) except: return [] def save_configs(self): try: self.api.upload_file(path_or_fileobj=io.BytesIO(json.dumps(self.configs, indent=2).encode("utf-8")), path_in_repo=self.config_file, repo_id=DATASET_ID, repo_type="dataset") with open(self.config_file, "w") as f: json.dump(self.configs, f, indent=2) except: pass def test_connection(self, h, u, p, r): h = h.rstrip("/") h = "https://" + h if not h.startswith("http") else h try: res = requests.request("PROPFIND", h + r, auth=HTTPBasicAuth(u, p), timeout=15) return (True, "✅") if res.status_code in [200, 207] else (False, f"HTTP {res.status_code}") except Exception as e: return False, str(e) def check_exists(self, url, auth): try: return requests.head(url, auth=auth, timeout=10).status_code == 200 except: return False def upload_file(self, conn_id, file_bytes, filename): target_conf = next((c for c in self.configs if c["id"] == conn_id), None) if not target_conf: return False, "Config Not Found" opts = target_conf["options"] base_url = ("https://" if not opts["webdav_hostname"].startswith("http") else "") + opts["webdav_hostname"].rstrip("/") try: from utils import ensure_png_bytes url = f"{base_url}{opts.get('webdav_root', '/').rstrip('/')}/{filename.lstrip('/')}" auth = HTTPBasicAuth(opts["webdav_login"], opts["webdav_password"]) res = requests.put(url, data=ensure_png_bytes(file_bytes), headers={'Content-Type': 'image/png'}, auth=auth, timeout=180) return (True, "Success") if res.status_code in [200, 201, 204] else (False, f"HTTP {res.status_code}") except Exception as e: return False, str(e) def upload_to_all_nodes(self, file_bytes, filename): results = [] for c in self.configs: if not c.get("enabled", True): continue ok, msg = self.upload_file(c["id"], file_bytes, filename) results.append(f"{'✅' if ok else '❌'} {c['name']}") return " | ".join(results) def add_connection(self, n, h, u, p, r="/", enabled=True): self.configs.append({"id": int(time.time()), "name": n, "enabled": enabled, "options": {"webdav_hostname": h, "webdav_login": u, "webdav_password": p, "webdav_root": r}}) self.save_configs() def delete_connection(self, conn_id): self.configs = [c for c in self.configs if c["id"] != conn_id] self.save_configs() def update_connection(self, conn_id, n, h, u, p, r, enabled): for c in self.configs: if c["id"] == conn_id: c["name"] = n c["enabled"] = enabled c["options"] = {"webdav_hostname": h, "webdav_login": u, "webdav_password": p, "webdav_root": r} break self.save_configs()