From 34215222bf0a067f81f86a4a2b8a8ca1b81dc65a Mon Sep 17 00:00:00 2001 From: kyx236 Date: Fri, 13 Feb 2026 04:06:42 +0800 Subject: [PATCH] feat: Implement account pooling for concurrent scheduling and introduce a new permissions module. --- account_store.py | 55 ++++++ bot.py | 466 ++++++++++++++++++++++++++++++++++++++------ config.py | 39 ++++ config.toml.example | 15 ++ permissions.py | 97 +++++++++ 5 files changed, 607 insertions(+), 65 deletions(-) create mode 100644 permissions.py diff --git a/account_store.py b/account_store.py index 40091a5..6d0072a 100644 --- a/account_store.py +++ b/account_store.py @@ -12,6 +12,61 @@ _ACCOUNTS_FILE = Path(__file__).parent / "accounts.txt" _STATS_FILE = Path(__file__).parent / "stats.json" _lock = threading.Lock() +# ====== 账号池（并发调度）====== +# _busy 记录当前被占用的账号行（原始字符串） +_busy: set[str] = set() + + +def acquire(n: int = 0) -> list[str]: + """从空闲账号中获取最多 n 个，标记为占用。 + + Args: + n: 需要的账号数。0 表示尽量获取所有空闲账号（至少 1 个）。 + + Returns: + 被获取的账号行列表（可能少于 n），为空表示没有空闲账号。 + """ + with _lock: + try: + with open(_ACCOUNTS_FILE, "r", encoding="utf-8") as f: + all_lines = [line.strip() for line in f if line.strip()] + except FileNotFoundError: + return [] + + free = [line for line in all_lines if line not in _busy] + if not free: + return [] + + if n <= 0: + # 获取全部空闲 + acquired = free + else: + acquired = free[:n] + + _busy.update(acquired) + return acquired + + +def release(lines: list[str]) -> None: + """释放指定账号，标记为空闲。""" + with _lock: + for line in lines: + _busy.discard(line) + + +def pool_status() -> dict: + """返回账号池状态。""" + with _lock: + try: + with open(_ACCOUNTS_FILE, "r", encoding="utf-8") as f: + all_lines = [line.strip() for line in f if line.strip()] + except FileNotFoundError: + all_lines = [] + + total = len(all_lines) + busy = sum(1 for line in all_lines if line in _busy) + return {"total": total, "busy": busy, "free": total - busy} + # ====== 账号操作 ====== diff --git a/bot.py b/bot.py index df0fdd4..a0c79b7 100644 --- a/bot.py +++ b/bot.py @@ -27,6 +27,9 @@ from telegram.ext import ( from config import ( TG_BOT_TOKEN, TG_ALLOWED_USERS, + TG_USER_PERMISSIONS, + ADMIN_USERS, + get_merged_permissions, MAIL_SYSTEMS, ) from mail_service import MailPool, extract_magic_link @@ -57,13 +60,38 @@ _stop_event = threading.Event() # ============================================================ def restricted(func): - """权限控制装饰器：仅允许 TG_ALLOWED_USERS 中的用户使用""" + """权限控制装饰器：检查用户是否有权使用当前命令。 + + 权限检查规则： + 1. 如果配置了 roles，按角色的 commands 列表检查；"*" 表示全部权限 + 2. 如果没有配置 roles，回退到 allowed_users 全量放行 + 3. 如果两者都没配，所有用户都可以使用 + """ + # 从函数名提取命令名（cmd_register -> register） + cmd_name = func.__name__.removeprefix("cmd_").removeprefix("handle_") + @wraps(func) async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs): user_id = update.effective_user.id - if TG_ALLOWED_USERS and user_id not in TG_ALLOWED_USERS: - await update.message.reply_text("⛔ 你没有权限使用此 Bot。") - return + + if TG_USER_PERMISSIONS or True: # 总是检查合并权限 + # 实时合并 config.toml + permissions.json + merged = get_merged_permissions() + if merged: + user_cmds = merged.get(user_id) + if user_cmds is None: + # 用户不在任何权限表中 + if TG_ALLOWED_USERS or TG_USER_PERMISSIONS: + # 有限制配置但用户不在其中 + await update.message.reply_text("⛔ 你没有权限使用此 Bot。") + return + # 没有任何限制配置：放行 + elif "*" not in user_cmds and cmd_name not in user_cmds: + await update.message.reply_text( + f"⛔ 你没有权限使用 /{cmd_name} 命令。" + ) + return + return await func(update, context, *args, **kwargs) return wrapper @@ -116,7 +144,6 @@ def _progress_bar(current: int, total: int, width: int = 12) -> str: # 命令处理 # ============================================================ -@restricted async def cmd_start(update: Update, context: ContextTypes.DEFAULT_TYPE): """/start — 欢迎信息""" welcome = ( @@ -138,7 +165,6 @@ async def cmd_start(update: Update, context: ContextTypes.DEFAULT_TYPE): await update.message.reply_text(welcome, parse_mode="HTML") -@restricted async def cmd_help(update: Update, context: ContextTypes.DEFAULT_TYPE): """/help — 命令列表""" text = ( @@ -159,7 +185,12 @@ async def cmd_help(update: Update, context: ContextTypes.DEFAULT_TYPE): " /proxy on|off — 开关代理\n" " /proxytest — 测试代理\n" " /proxystatus — 代理池状态\n" - " /mailstatus — 邮件系统状态\n" + " /mailstatus — 邮件系统状态\n\n" + "🔐 用户管理（仅管理员）\n" + " /adduser [cmds] — 添加用户\n" + " /removeuser — 移除用户\n" + " /setperm — 修改权限\n" + " /users — 查看用户列表\n" ) await update.message.reply_text(text, parse_mode="HTML") @@ -167,14 +198,25 @@ async def cmd_help(update: Update, context: ContextTypes.DEFAULT_TYPE): @restricted async def cmd_status(update: Update, context: ContextTypes.DEFAULT_TYPE): """/status — 查看任务状态""" + lines = [] + + # 全局任务（注册等） with _task_lock: if _task_running: - await update.message.reply_text( - f"⏳ 当前任务：{_task_name}", - parse_mode="HTML", - ) - else: - await update.message.reply_text("✅ 当前无运行中的任务。") + lines.append(f"⏳ 当前任务：{_task_name}") + + # 账号池状态 + ps = account_store.pool_status() + if ps["total"] > 0: + lines.append( + f"🗄 账号池：共 {ps['total']} | " + f"🟢 空闲 {ps['free']} | 🔴 占用 {ps['busy']}" + ) + + if not lines: + await update.message.reply_text("✅ 当前无运行中的任务。") + else: + await update.message.reply_text("\n".join(lines), parse_mode="HTML") @restricted @@ -281,6 +323,192 @@ async def cmd_proxystatus(update: Update, context: ContextTypes.DEFAULT_TYPE): await update.message.reply_text(text, parse_mode="HTML") +# ============================================================ +# 用户权限管理（仅管理员） +# ============================================================ + +import permissions + +def _admin_only(func): + """管理员专属装饰器：只有 config.toml 中 commands=["*"] 的用户可使用""" + @wraps(func) + async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs): + user_id = update.effective_user.id + if user_id not in ADMIN_USERS: + await update.message.reply_text("⛔ 此命令仅管理员可用。") + return + return await func(update, context, *args, **kwargs) + return wrapper + + +@_admin_only +async def cmd_adduser(update: Update, context: ContextTypes.DEFAULT_TYPE): + """/adduser [cmd1,cmd2,...|*] — 添加用户并设置权限""" + if not context.args or len(context.args) < 1: + await update.message.reply_text( + "❌ 用法：\n" + " /adduser 123456789 * — 全部权限\n" + " /adduser 123456789 accounts,verify,stats — 指定命令\n\n" + "💡 可用命令列表：\n" + f" {', '.join(sorted(permissions.ALL_COMMANDS))}", + parse_mode="HTML", + ) + return + + try: + target_uid = int(context.args[0]) + except ValueError: + await update.message.reply_text("❌ 用户 ID 必须是数字。") + return + + # 解析权限 + if len(context.args) >= 2: + cmd_arg = context.args[1] + if cmd_arg == "*": + cmds = ["*"] + else: + cmds = [c.strip() for c in cmd_arg.split(",") if c.strip()] + # 验证命令名 + invalid = [c for c in cmds if c not in permissions.ALL_COMMANDS] + if invalid: + await update.message.reply_text( + f"❌ 未知命令: {', '.join(invalid)}\n\n" + f"可用命令：\n{', '.join(sorted(permissions.ALL_COMMANDS))}", + parse_mode="HTML", + ) + return + else: + # 默认给基础查看权限 + cmds = ["start", "help", "accounts", "verify", "stats", "status"] + + permissions.add_user(target_uid, cmds) + + cmd_display = "全部命令" if cmds == ["*"] else ", ".join(cmds) + await update.message.reply_text( + f"✅ 已添加用户 {target_uid}\n" + f"🔑 权限: {cmd_display}", + parse_mode="HTML", + ) + + +@_admin_only +async def cmd_removeuser(update: Update, context: ContextTypes.DEFAULT_TYPE): + """/removeuser — 移除用户""" + if not context.args: + await update.message.reply_text("❌ 用法: /removeuser 123456789", parse_mode="HTML") + return + + try: + target_uid = int(context.args[0]) + except ValueError: + await update.message.reply_text("❌ 用户 ID 必须是数字。") + return + + # 不能移除 config.toml 中的管理员 + if target_uid in ADMIN_USERS: + await update.message.reply_text("❌ 不能移除 config.toml 中定义的管理员。") + return + + if permissions.remove_user(target_uid): + await update.message.reply_text(f"🗑 已移除用户 {target_uid}", parse_mode="HTML") + else: + await update.message.reply_text(f"❌ 用户 {target_uid} 不在运行时权限列表中。", parse_mode="HTML") + + +@_admin_only +async def cmd_setperm(update: Update, context: ContextTypes.DEFAULT_TYPE): + """/setperm — 修改用户权限""" + if not context.args or len(context.args) < 2: + await update.message.reply_text( + "❌ 用法：\n" + " /setperm 123456789 *\n" + " /setperm 123456789 accounts,verify,stats\n\n" + f"可用命令：\n{', '.join(sorted(permissions.ALL_COMMANDS))}", + parse_mode="HTML", + ) + return + + try: + target_uid = int(context.args[0]) + except ValueError: + await update.message.reply_text("❌ 用户 ID 必须是数字。") + return + + # 不能修改 config.toml 管理员 + if target_uid in ADMIN_USERS: + await update.message.reply_text("❌ 不能修改 config.toml 中定义的管理员权限。") + return + + cmd_arg = context.args[1] + if cmd_arg == "*": + cmds = ["*"] + else: + cmds = [c.strip() for c in cmd_arg.split(",") if c.strip()] + invalid = [c for c in cmds if c not in permissions.ALL_COMMANDS] + if invalid: + await update.message.reply_text( + f"❌ 未知命令: {', '.join(invalid)}", + parse_mode="HTML", + ) + return + + # 如果用户不存在，自动添加 + user = permissions.get_user(target_uid) + if user: + permissions.set_commands(target_uid, cmds) + else: + permissions.add_user(target_uid, cmds) + + cmd_display = "全部命令" if cmds == ["*"] else ", ".join(cmds) + await update.message.reply_text( + f"✅ 用户 {target_uid} 权限已更新\n" + f"🔑 权限: {cmd_display}", + parse_mode="HTML", + ) + + +@_admin_only +async def cmd_users(update: Update, context: ContextTypes.DEFAULT_TYPE): + """/users — 查看所有用户及权限""" + merged = get_merged_permissions() + + if not merged: + await update.message.reply_text("📭 没有任何用户。") + return + + text = "👥 用户权限列表\n\n" + + # config.toml 管理员 + text += "📌 管理员（config.toml）\n" + for uid in sorted(ADMIN_USERS): + text += f" • {uid} — ✨ 全部权限\n" + + # config.toml 非管理员角色 + static_non_admin = { + uid: cmds for uid, cmds in TG_USER_PERMISSIONS.items() + if uid not in ADMIN_USERS + } + if static_non_admin: + text += "\n📌 静态角色（config.toml）\n" + for uid, cmds in sorted(static_non_admin.items()): + cmd_str = ", ".join(sorted(cmds)) if "*" not in cmds else "全部命令" + text += f" • {uid} — {cmd_str}\n" + + # 运行时添加的用户 + runtime_users = permissions.list_users() + if runtime_users: + text += "\n🔧 运行时用户（Bot 添加）\n" + for uid, info in sorted(runtime_users.items()): + cmds = info.get("commands", []) + cmd_str = ", ".join(sorted(cmds)) if "*" not in cmds else "全部命令" + text += f" • {uid} — {cmd_str}\n" + + if len(text) > 4000: + text = text[:4000] + "\n...（已截断）" + + await update.message.reply_text(text, parse_mode="HTML") + + @restricted async def cmd_proxy(update: Update, context: ContextTypes.DEFAULT_TYPE): """/proxy on|off — 开关代理""" @@ -710,47 +938,80 @@ def _register_worker(loop: asyncio.AbstractEventLoop, status_msg, count: int): @restricted async def cmd_check(update: Update, context: ContextTypes.DEFAULT_TYPE): - """/check — CC 检查""" + """/check — CC 检查（支持多行批量）""" if not context.args: await update.message.reply_text( "❌ 用法：/check 卡号|月|年|CVC\n" - "示例：/check 4111111111111111|12|2025|123", + "支持一次粘贴多行：\n" + "/check\n" + "4111111111111111|12|25|123\n" + "5200000000000007|01|26|456", parse_mode="HTML", ) return - card_line = context.args[0] - parts = card_line.split("|") - if len(parts) != 4: - await update.message.reply_text("❌ 格式错误，需要：卡号|月|年|CVC", parse_mode="HTML") + # 解析所有行，提取有效卡片 + cards = [] + for arg in context.args: + arg = arg.strip() + if not arg: + continue + parts = arg.split("|") + if len(parts) == 4: + cards.append(arg) + + if not cards: + await update.message.reply_text( + "❌ 没有找到有效卡片。\n格式：卡号|月|年|CVC", + parse_mode="HTML", + ) return - # 读取可用账号 - acc_lines = account_store.read_lines() + # 从账号池获取空闲账号 + if len(cards) == 1: + acquired = account_store.acquire(1) + else: + acquired = account_store.acquire() # 尽量获取所有空闲 - if not acc_lines: - await update.message.reply_text("❌ 没有可用账号，请先 /register 注册一个。") + if not acquired: + ps = account_store.pool_status() + if ps["total"] == 0: + await update.message.reply_text("❌ 没有可用账号，请先 /register 注册一个。") + else: + await update.message.reply_text( + f"⏳ 所有账号繁忙（{ps['busy']}/{ps['total']}），请稍后再试。" + ) return - if not _set_task("CC 检查"): - await update.message.reply_text(f"⚠️ 已有任务在运行：{_task_name}") - return - - status_msg = await update.message.reply_text( - f"🔍 正在检查卡片：{parts[0][:4]}****{parts[0][-4:]}", - parse_mode="HTML", - ) - - loop = asyncio.get_event_loop() - threading.Thread( - target=_check_worker, - args=(loop, status_msg, card_line, acc_lines[-1]), - daemon=True, - ).start() + if len(cards) == 1: + # 单卡检查 + parts = cards[0].split("|") + status_msg = await update.message.reply_text( + f"🔍 正在检查卡片：{parts[0][:4]}****{parts[0][-4:]}", + parse_mode="HTML", + ) + loop = asyncio.get_event_loop() + threading.Thread( + target=_check_worker, + args=(loop, status_msg, cards[0], acquired[0]), + daemon=True, + ).start() + else: + # 多卡批量检查 + status_msg = await update.message.reply_text( + f"📋 解析到 {len(cards)} 张卡片，{len(acquired)} 个账号就绪，开始批量检查...", + parse_mode="HTML", + ) + loop = asyncio.get_event_loop() + threading.Thread( + target=_batch_check_worker, + args=(loop, status_msg, cards, acquired), + daemon=True, + ).start() def _check_worker(loop: asyncio.AbstractEventLoop, status_msg, card_line: str, account_line: str): - """CC 检查工作线程""" + """CC 检查工作线程（完成后自动释放账号）""" try: cc, mm, yy, cvc = card_line.split("|") acc_parts = account_line.split("|") @@ -815,7 +1076,7 @@ def _check_worker(loop: asyncio.AbstractEventLoop, status_msg, card_line: str, a loop, ) finally: - _clear_task() + account_store.release([account_line]) # ============================================================ @@ -837,15 +1098,17 @@ async def handle_document(update: Update, context: ContextTypes.DEFAULT_TYPE): await update.message.reply_text("❌ 文件太大（最大 1MB）") return - # 读取可用账号 - acc_lines = account_store.read_lines() + # 从账号池获取空闲账号 + acquired = account_store.acquire() # 获取所有空闲 - if not acc_lines: - await update.message.reply_text("❌ 没有可用账号，请先 /register 注册一个。") - return - - if not _set_task("批量 CC 检查"): - await update.message.reply_text(f"⚠️ 已有任务在运行：{_task_name}") + if not acquired: + ps = account_store.pool_status() + if ps["total"] == 0: + await update.message.reply_text("❌ 没有可用账号，请先 /register 注册一个。") + else: + await update.message.reply_text( + f"⏳ 所有账号繁忙（{ps['busy']}/{ps['total']}），请稍后再试。" + ) return # 下载文件 @@ -857,7 +1120,7 @@ async def handle_document(update: Update, context: ContextTypes.DEFAULT_TYPE): content = file_bytes.decode("utf-8", errors="ignore") except Exception as e: await _edit_or_send(status_msg, f"💥 下载文件失败：{e}") - _clear_task() + account_store.release(acquired) return # 解析卡片 @@ -871,7 +1134,7 @@ async def handle_document(update: Update, context: ContextTypes.DEFAULT_TYPE): if not cards: await _edit_or_send(status_msg, "❌ 文件中没有找到有效卡片。\n格式：卡号|月|年|CVC") - _clear_task() + account_store.release(acquired) return await _edit_or_send( @@ -883,40 +1146,54 @@ async def handle_document(update: Update, context: ContextTypes.DEFAULT_TYPE): loop = asyncio.get_event_loop() threading.Thread( target=_batch_check_worker, - args=(loop, status_msg, cards, acc_lines[-1]), + args=(loop, status_msg, cards, acquired), daemon=True, ).start() -def _batch_check_worker(loop: asyncio.AbstractEventLoop, status_msg, cards: list, account_line: str): - """批量 CC 检查工作线程""" +def _batch_check_worker(loop: asyncio.AbstractEventLoop, status_msg, cards: list, acc_lines: list): + """批量 CC 检查工作线程（round-robin 轮询账号）""" from models import ClaudeAccount from identity import random_ua results = [] total = len(cards) + num_accounts = len(acc_lines) + + # 预构建所有账号对象 + accounts = [] + for line in acc_lines: + parts = line.split("|") + email, session_key, org_uuid = parts[0], parts[1], parts[2] + acc = ClaudeAccount(email, session_key, org_uuid, random_ua()) + accounts.append(acc) + + # 为每个账号创建对应的 tokenizer 和 checker + tokenizers = [StripeTokenizer(acc.user_agent) for acc in accounts] + checkers = [GiftChecker(acc) for acc in accounts] try: - acc_parts = account_line.split("|") - email, session_key, org_uuid = acc_parts[0], acc_parts[1], acc_parts[2] - account = ClaudeAccount(email, session_key, org_uuid, random_ua()) - - tokenizer = StripeTokenizer(account.user_agent) - checker = GiftChecker(account) - for i, card_line in enumerate(cards): if _is_stopped(): break + # Round-robin 选择账号 + idx = i % num_accounts + current_acc = accounts[idx] + tokenizer = tokenizers[idx] + checker = checkers[idx] + acc_label = current_acc.email.split("@")[0] # 简短标识 + cc, mm, yy, cvc = card_line.split("|") masked = f"{cc[:4]}****{cc[-4:]}" # 更新进度 recent = "\n".join(results[-5:]) # 显示最近 5 条结果 + acc_info = f" 👤 账号 {idx + 1}/{num_accounts}" if num_accounts > 1 else "" asyncio.run_coroutine_threadsafe( _edit_or_send( status_msg, - f"🔍 [{i + 1}/{total}] 检查中：{masked}\n\n" + f"🔍 [{i + 1}/{total}] 检查中：{masked}{acc_info}\n\n" + recent, ), loop, @@ -939,7 +1216,13 @@ def _batch_check_worker(loop: asyncio.AbstractEventLoop, status_msg, cards: list "DEAD": "💀", "ERROR": "⚠️", } icon = result_icons.get(result, "❓") - results.append(f"{icon} {masked} → {result}") + + # 仅成功结果显示完整卡号 + 对应邮箱 + live_results = {"LIVE", "INSUFFICIENT_FUNDS", "CCN_LIVE"} + if result in live_results: + results.append(f"{icon} {card_line} → {result}\n 📧 {current_acc.email}") + else: + results.append(f"{icon} {masked} → {result}") # 间隔防限流 if i < total - 1: @@ -950,9 +1233,10 @@ def _batch_check_worker(loop: asyncio.AbstractEventLoop, status_msg, cards: list dead = sum(1 for r in results if "DEAD" in r or "DECLINED" in r or "Stripe" in r) other = total - live - dead + acc_note = f" | 👤 {num_accounts} 个账号轮询" if num_accounts > 1 else "" report = ( f"🏁 批量 CC 检查完成\n\n" - f"📊 共 {total} 张 | 💰 有效 {live} | 💀 无效 {dead} | ❓ 其他 {other}\n\n" + f"📊 共 {total} 张 | 💰 有效 {live} | 💀 无效 {dead} | ❓ 其他 {other}{acc_note}\n\n" ) report += "\n".join(results) @@ -972,7 +1256,7 @@ def _batch_check_worker(loop: asyncio.AbstractEventLoop, status_msg, cards: list loop, ) finally: - _clear_task() + account_store.release(acc_lines) # ============================================================ @@ -1014,7 +1298,7 @@ async def callback_export_json(update: Update, context: ContextTypes.DEFAULT_TYP # ============================================================ async def post_init(application: Application): - """Bot 启动后设置命令菜单""" + """Bot 启动后设置命令菜单 + 推送 help""" commands = [ BotCommand("start", "欢迎信息"), BotCommand("register", "注册 Claude 账号 [数量]"), @@ -1028,12 +1312,60 @@ async def post_init(application: Application): BotCommand("proxy", "代理开关 on/off"), BotCommand("proxytest", "测试代理"), BotCommand("proxystatus", "代理池状态"), + BotCommand("adduser", "添加用户"), + BotCommand("removeuser", "移除用户"), + BotCommand("setperm", "修改用户权限"), + BotCommand("users", "用户列表"), BotCommand("status", "任务状态"), BotCommand("help", "命令帮助"), ] await application.bot.set_my_commands(commands) logger.info("Bot 命令菜单已设置") + # 向所有已知用户推送启动通知 + help + merged = get_merged_permissions() + user_ids = set(merged.keys()) | set(TG_ALLOWED_USERS) + if not user_ids: + logger.info("未配置用户，跳过启动通知") + return + + startup_msg = ( + "🤖 autoClaude Bot 已启动\n\n" + "📖 命令说明\n\n" + "📝 注册与账号\n" + " /register [N] — 注册 N 个账号\n" + " /accounts — 查看已注册账号\n" + " /delete <序号|邮箱> — 删除账号\n" + " /verify — 验证 SK 有效性\n\n" + "💳 CC 检查\n" + " /check — 单张检查\n" + " 📎 发送 .txt 文件 — 批量检查\n\n" + "🛠 工具与状态\n" + " /stop — 中断当前任务\n" + " /stats — 统计面板\n" + " /status — 任务状态\n\n" + "🌐 代理与邮件\n" + " /proxy on|off — 开关代理\n" + " /proxytest — 测试代理\n" + " /proxystatus — 代理池状态\n" + " /mailstatus — 邮件系统状态\n\n" + "🔐 用户管理（仅管理员）\n" + " /adduser [cmds] — 添加用户\n" + " /removeuser — 移除用户\n" + " /setperm — 修改权限\n" + " /users — 查看用户列表\n" + ) + + for uid in user_ids: + try: + await application.bot.send_message( + chat_id=uid, text=startup_msg, parse_mode="HTML", + ) + except Exception as e: + logger.debug(f"启动通知发送失败 (uid={uid}): {e}") + + logger.info(f"启动通知已推送给 {len(user_ids)} 个用户") + def main(): """启动 Bot""" @@ -1058,6 +1390,10 @@ def main(): app.add_handler(CommandHandler("proxystatus", cmd_proxystatus)) app.add_handler(CommandHandler("proxy", cmd_proxy)) app.add_handler(CommandHandler("status", cmd_status)) + app.add_handler(CommandHandler("adduser", cmd_adduser)) + app.add_handler(CommandHandler("removeuser", cmd_removeuser)) + app.add_handler(CommandHandler("setperm", cmd_setperm)) + app.add_handler(CommandHandler("users", cmd_users)) app.add_handler(CallbackQueryHandler(callback_export_json, pattern="^export_accounts_json$")) app.add_handler(MessageHandler(filters.Document.ALL, handle_document)) diff --git a/config.py b/config.py index 4c4ea1c..b6c18a5 100644 --- a/config.py +++ b/config.py @@ -29,6 +29,45 @@ PRODUCT_ID: str = _cfg["stripe"]["product_id"] TG_BOT_TOKEN: str = _cfg["telegram"]["bot_token"] TG_ALLOWED_USERS: list[int] = _cfg["telegram"].get("allowed_users", []) +# --- 角色权限 --- +# 静态权限来自 config.toml，运行时权限来自 permissions.json +# 每次调用 get_merged_permissions() 合并两者 +_roles: list[dict] = _cfg["telegram"].get("roles", []) + +# 静态权限映射（来自 config.toml） +_STATIC_PERMISSIONS: dict[int, set[str]] = {} + +# 检查 roles 是否实际配置了用户（至少一个 role 的 users 非空） +_roles_active = _roles and any(role.get("users") for role in _roles) + +if _roles_active: + for role in _roles: + cmds = set(role.get("commands", [])) + for uid in role.get("users", []): + _STATIC_PERMISSIONS.setdefault(uid, set()).update(cmds) +else: + # roles 未配置或 users 全为空 → 回退到 allowed_users 全量放行 + for uid in TG_ALLOWED_USERS: + _STATIC_PERMISSIONS[uid] = {"*"} + +# config.toml 中拥有 "*" 权限的用户 = 超级管理员 +ADMIN_USERS: set[int] = { + uid for uid, cmds in _STATIC_PERMISSIONS.items() if "*" in cmds +} + + +def get_merged_permissions() -> dict[int, set[str]]: + """合并 config.toml 静态权限 + permissions.json 运行时权限""" + import permissions as perm_mod + merged = dict(_STATIC_PERMISSIONS) + for uid, cmds in perm_mod.get_permissions_map().items(): + merged.setdefault(uid, set()).update(cmds) + return merged + + +# 向后兼容：初始静态权限 +TG_USER_PERMISSIONS = _STATIC_PERMISSIONS + # --- 邮箱系统 --- MAIL_SYSTEMS: list[dict] = _cfg.get("mail", []) diff --git a/config.toml.example b/config.toml.example index 19166d6..9a500e0 100644 --- a/config.toml.example +++ b/config.toml.example @@ -17,6 +17,21 @@ product_id = "prod_TXU4hGh2EDxASl" bot_token = "your_bot_token_here" # @BotFather 获取 allowed_users = [] # 允许使用的用户ID列表（空=不限制） +# --- 角色权限控制 --- +# 每个 [[telegram.roles]] 定义一个角色，包含用户列表和允许的命令 +# commands = ["*"] 表示全部命令 +# 如果不配置 roles，所有 allowed_users 拥有全部权限 + +[[telegram.roles]] +name = "admin" +users = [] # 管理员用户 ID +commands = ["*"] # 全部命令 + +[[telegram.roles]] +name = "user" +users = [] # 普通用户 ID +commands = ["accounts", "verify", "stats", "status", "help", "start"] + # --- 邮箱系统（轮询使用，API 接口相同）--- # 可添加多个 [[mail]] 块 # api_token: 直接配置 API Token，无需管理员账号密码 diff --git a/permissions.py b/permissions.py new file mode 100644 index 0000000..9f56942 --- /dev/null +++ b/permissions.py @@ -0,0 +1,97 @@ +""" +运行时权限管理模块 +管理员可通过 Bot 命令动态添加/删除用户和设置权限。 +持久化存储在 permissions.json 中。 +""" + +import json +import threading +from pathlib import Path + +_PERM_FILE = Path(__file__).parent / "permissions.json" +_lock = threading.Lock() + +# 所有可用命令列表（用于验证输入） +ALL_COMMANDS = { + "start", "help", "register", "stop", "check", + "accounts", "delete", "verify", "stats", "status", + "mailstatus", "proxy", "proxytest", "proxystatus", + "document", # 文件上传 + "adduser", "removeuser", "setperm", "users", # 管理命令 +} + + +def _load() -> dict: + """加载权限数据""" + try: + with open(_PERM_FILE, "r", encoding="utf-8") as f: + return json.load(f) + except (FileNotFoundError, json.JSONDecodeError): + return {"users": {}} + + +def _save(data: dict): + """保存权限数据""" + with open(_PERM_FILE, "w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, indent=2) + + +def add_user(user_id: int, commands: list[str]) -> None: + """添加用户或更新已有用户权限""" + with _lock: + data = _load() + data["users"][str(user_id)] = { + "commands": commands, + } + _save(data) + + +def remove_user(user_id: int) -> bool: + """删除用户，返回是否成功""" + with _lock: + data = _load() + key = str(user_id) + if key in data["users"]: + del data["users"][key] + _save(data) + return True + return False + + +def set_commands(user_id: int, commands: list[str]) -> bool: + """设置用户权限，返回是否成功（用户必须存在）""" + with _lock: + data = _load() + key = str(user_id) + if key not in data["users"]: + return False + data["users"][key]["commands"] = commands + _save(data) + return True + + +def get_user(user_id: int) -> dict | None: + """获取单个用户信息""" + with _lock: + data = _load() + return data["users"].get(str(user_id)) + + +def list_users() -> dict[int, dict]: + """列出所有运行时用户""" + with _lock: + data = _load() + return {int(k): v for k, v in data["users"].items()} + + +def get_permissions_map() -> dict[int, set[str]]: + """ + 返回运行时权限映射：user_id → set[command_name] + 用于与 config.toml 的静态权限合并 + """ + with _lock: + data = _load() + result = {} + for uid_str, info in data["users"].items(): + result[int(uid_str)] = set(info.get("commands", [])) + return result