From 9b7ecb7b80e6b004a828d3e6f51a48af62219b83 Mon Sep 17 00:00:00 2001 From: dela Date: Fri, 30 Jan 2026 10:48:56 +0800 Subject: [PATCH] =?UTF-8?q?=E9=87=8D=E6=9E=84=E6=9C=BA=E5=99=A8=E4=BA=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 3 +- bot/__init__.py | 9 ++ bot/app.py | 63 ++++++++ bot/handlers/__init__.py | 7 + bot/handlers/go.py | 290 +++++++++++++++++++++++++++++++++++ bot/handlers/settings.py | 96 ++++++++++++ bot/handlers/start.py | 57 +++++++ bot/handlers/status.py | 112 ++++++++++++++ bot/middlewares/__init__.py | 7 + bot/middlewares/auth.py | 111 ++++++++++++++ bot/services/__init__.py | 7 + bot/services/file_sender.py | 101 ++++++++++++ bot/services/task_manager.py | 228 +++++++++++++++++++++++++++ bot_main.py | 104 +++++++++++++ config.py | 18 +++ pyproject.toml | 1 + uv.lock | 57 +++++++ 17 files changed, 1270 insertions(+), 1 deletion(-) create mode 100644 bot/__init__.py create mode 100644 bot/app.py create mode 100644 bot/handlers/__init__.py create mode 100644 bot/handlers/go.py create mode 100644 bot/handlers/settings.py create mode 100644 bot/handlers/start.py create mode 100644 bot/handlers/status.py create mode 100644 bot/middlewares/__init__.py create mode 100644 bot/middlewares/auth.py create mode 100644 bot/services/__init__.py create mode 100644 bot/services/file_sender.py create mode 100644 bot/services/task_manager.py create mode 100644 bot_main.py diff --git a/.gitignore b/.gitignore index 8199c15..5e97db4 100644 --- a/.gitignore +++ b/.gitignore @@ -38,4 +38,5 @@ Thumbs.db /.gemini /docs -/reference \ No newline at end of file +/reference +/outputs \ No newline at end of file diff --git a/bot/__init__.py b/bot/__init__.py new file mode 100644 index 0000000..3c70dbd --- /dev/null +++ b/bot/__init__.py @@ -0,0 +1,9 @@ +""" +Telegram Bot 模块 + +提供 Telegram Bot 接口来控制 OpenAI 账号注册系统 +""" + +from bot.app import create_application + +__all__ = ["create_application"] diff --git a/bot/app.py b/bot/app.py new file mode 100644 index 0000000..4657c65 --- /dev/null +++ b/bot/app.py @@ -0,0 +1,63 @@ +""" +Telegram Bot Application 创建模块 +""" + +from telegram.ext import Application, CommandHandler + +from config import AppConfig +from bot.handlers import start, go, status, settings + + +def create_application(config: AppConfig) -> Application: + """ + 创建并配置 Telegram Bot Application + + 参数: + config: AppConfig 配置对象 + + 返回: + 配置好的 Application 实例 + """ + if not config.telegram_bot_token: + raise ValueError("TELEGRAM_BOT_TOKEN is not configured") + + # 创建 Application + application = Application.builder().token(config.telegram_bot_token).build() + + # 存储配置到 bot_data 供 handlers 使用 + application.bot_data["config"] = config + + # 解析允许的用户列表 + allowed_users = set() + if config.telegram_allowed_users: + for uid in config.telegram_allowed_users.split(","): + uid = uid.strip() + if uid.isdigit(): + allowed_users.add(int(uid)) + + admin_users = set() + if config.telegram_admin_users: + for uid in config.telegram_admin_users.split(","): + uid = uid.strip() + if uid.isdigit(): + admin_users.add(int(uid)) + + application.bot_data["allowed_users"] = allowed_users + application.bot_data["admin_users"] = admin_users + + # 初始化用户设置 + application.bot_data["user_settings"] = {} + + # 初始化任务管理器 + from bot.services.task_manager import TaskManager + + application.bot_data["task_manager"] = TaskManager() + + # 注册命令处理器 + application.add_handler(CommandHandler("start", start.start_command)) + application.add_handler(CommandHandler("help", start.help_command)) + application.add_handler(CommandHandler("go", go.go_command)) + application.add_handler(CommandHandler("status", status.status_command)) + application.add_handler(CommandHandler("set", settings.set_command)) + + return application diff --git a/bot/handlers/__init__.py b/bot/handlers/__init__.py new file mode 100644 index 0000000..333604f --- /dev/null +++ b/bot/handlers/__init__.py @@ -0,0 +1,7 @@ +""" +Bot 命令处理器模块 +""" + +from bot.handlers import start, go, status, settings + +__all__ = ["start", "go", "status", "settings"] diff --git a/bot/handlers/go.py b/bot/handlers/go.py new file mode 100644 index 0000000..d7262a2 --- /dev/null +++ b/bot/handlers/go.py @@ -0,0 +1,290 @@ +""" +/go 命令处理器 - 生成支付链接并发送 JSON 文件 +""" + +import asyncio +from telegram import Update +from telegram.ext import ContextTypes + +from bot.middlewares.auth import require_auth +from bot.services.task_manager import TaskStatus +from bot.services.file_sender import send_results_as_json +from bot.handlers.settings import get_user_settings, DEFAULT_WORKERS + + +PLAN_MAPPING = { + "plus": "chatgptplusplan", + "pro": "chatgptproplan", + "chatgptplusplan": "chatgptplusplan", + "chatgptproplan": "chatgptproplan", +} + + +@require_auth +async def go_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """ + 处理 /go 命令 - 注册账号并生成支付链接 + + 用法: + /go - 生成 1 个 Plus 支付链接 + /go 5 - 生成 5 个 Plus 支付链接 + /go 5 plus - 生成 5 个 Plus 支付链接 + /go 3 pro - 生成 3 个 Pro 支付链接 + """ + # 解析参数 + num_accounts = 1 + plan_input = "plus" + + if context.args: + # 第一个参数：数量 + try: + num_accounts = int(context.args[0]) + if num_accounts < 1: + num_accounts = 1 + elif num_accounts > 10: + await update.message.reply_text( + "⚠️ 单次最多处理 10 个账号，已自动调整为 10" + ) + num_accounts = 10 + except ValueError: + # 可能第一个参数是 plan + if context.args[0].lower() in PLAN_MAPPING: + plan_input = context.args[0].lower() + else: + await update.message.reply_text( + "❌ 无效参数\n\n" + "用法: `/go [数量] [plan]`\n" + "示例:\n" + "• `/go` - 1个 Plus 计划\n" + "• `/go 5` - 5个 Plus 计划\n" + "• `/go 3 pro` - 3个 Pro 计划", + parse_mode="Markdown", + ) + return + + # 第二个参数：plan + if len(context.args) > 1: + plan_input = context.args[1].lower() + + # 验证 plan + plan_name = PLAN_MAPPING.get(plan_input) + if not plan_name: + await update.message.reply_text( + f"❌ 无效的订阅计划: `{plan_input}`\n\n" f"可选值: `plus`, `pro`", + parse_mode="Markdown", + ) + return + + plan_display = "Plus" if "plus" in plan_name else "Pro" + + # 获取配置和任务管理器 + config = context.bot_data["config"] + task_manager = context.bot_data["task_manager"] + user_id = update.effective_user.id + + # 获取用户设置的并发数 + user_settings = get_user_settings(context, user_id) + workers = user_settings.get("workers", DEFAULT_WORKERS) + + # 创建任务 + task_id = task_manager.create_task( + user_id=user_id, + task_type="go", + total=num_accounts, + description=f"{num_accounts} x {plan_display} 支付链接", + ) + + await update.message.reply_text( + f"🚀 **开始生成支付链接**\n\n" + f"📦 计划: {plan_display}\n" + f"📊 数量: {num_accounts}\n" + f"👷 并发: {workers}\n" + f"📋 任务 ID: `{task_id}`\n\n" + f"使用 `/status {task_id}` 查看进度", + parse_mode="Markdown", + ) + + # 在后台执行工作流 + asyncio.create_task( + _run_go_workflow( + config, task_manager, task_id, num_accounts, plan_name, workers, update, context + ) + ) + + +async def _run_go_workflow( + config, + task_manager, + task_id: str, + num_accounts: int, + plan_name: str, + workers: int, + update: Update, + context: ContextTypes.DEFAULT_TYPE, +): + """ + 后台执行工作流：并发注册 + 登录 + 获取支付链接 + """ + plan_display = "Plus" if "plus" in plan_name else "Pro" + semaphore = asyncio.Semaphore(workers) + + async def process_one(index: int): + async with semaphore: + account_num = index + 1 + try: + task_manager.update_progress( + task_id, + current_item=f"账号 {account_num}/{num_accounts}", + ) + + result = await _register_with_checkout(config, account_num, plan_name) + return result + except Exception as e: + return { + "status": "failed", + "failed_stage": "exception", + "error": str(e), + } + + # 并发执行所有任务 + tasks = [process_one(i) for i in range(num_accounts)] + results = await asyncio.gather(*tasks) + + # 更新进度 + success_count = sum(1 for r in results if r.get("status") == "success") + failed_count = len(results) - success_count + + task_manager.update_progress(task_id, completed=num_accounts) + + # 完成任务 + task_manager.complete_task( + task_id, + status=TaskStatus.COMPLETED if failed_count == 0 else TaskStatus.PARTIAL, + result={ + "success": success_count, + "failed": failed_count, + "total": num_accounts, + "plan": plan_display, + }, + ) + + # 发送 JSON 文件结果 + await send_results_as_json( + update=update, + context=context, + task_id=task_id, + plan=plan_display, + results=results, + ) + + +async def _register_with_checkout(config, task_id: int, plan_name: str): + """ + 执行完整的注册 + 登录 + checkout 流程 + + 返回包含所有信息的结果字典 + """ + from core.session import OAISession + from core.flow import RegisterFlow + from core.login_flow import LoginFlow + from core.checkout import CheckoutFlow + from utils.logger import logger + import re + + def _mask_proxy(proxy: str) -> str: + return re.sub(r"://([^:]+):([^@]+)@", r"://***:***@", proxy) + + # 选择代理 + proxy = config.proxy.get_next_proxy() + if proxy: + logger.info(f"[Go {task_id}] Using proxy: {_mask_proxy(proxy)}") + + session = None + try: + session = OAISession(proxy=proxy, impersonate=config.tls_impersonate) + + # Step 1: 注册 + logger.info(f"[Go {task_id}] Step 1: Registering...") + flow = RegisterFlow(session, config) + reg_result = await flow.run() + + if reg_result.get("status") != "success": + return { + "status": "failed", + "failed_stage": "register", + "error": reg_result.get("error", "Registration failed"), + } + + email = reg_result.get("email") + password = reg_result.get("password") + logger.info(f"[Go {task_id}] Registered: {email}") + + # Step 2: 登录获取 Token + logger.info(f"[Go {task_id}] Step 2: Logging in...") + login_flow = LoginFlow(session, email, password) + login_result = await login_flow.run() + + if login_result.get("status") != "success": + return { + "status": "partial", + "failed_stage": "login", + "email": email, + "password": password, + "error": login_result.get("error", "Login failed"), + } + + session.access_token = login_result.get("access_token") + session.session_token = login_result.get("session_token") + logger.info(f"[Go {task_id}] Logged in, token obtained") + + # Step 3: 获取支付链接 + logger.info(f"[Go {task_id}] Step 3: Creating checkout session for {plan_name}...") + checkout = CheckoutFlow(session, plan_name=plan_name) + checkout_result = await checkout.create_checkout_session() + + result = { + "status": "success", + "email": email, + "password": password, + "access_token": login_result.get("access_token"), + "plan_name": plan_name, + } + + if checkout_result.get("status") == "success": + result["checkout_session_id"] = checkout_result.get("checkout_session_id") + result["checkout_url"] = checkout_result.get("url", "") + result["client_secret"] = checkout_result.get("client_secret") + logger.info(f"[Go {task_id}] Checkout session created") + + # 保存结果 + from main import save_account, save_checkout_result + + await save_account(result, config.accounts_output_file) + checkout_result["email"] = email + await save_checkout_result(checkout_result) + else: + # Checkout 失败但账号创建成功 + result["status"] = "partial" + result["checkout_error"] = checkout_result.get("error") + logger.warning(f"[Go {task_id}] Checkout failed: {checkout_result.get('error')}") + + from main import save_account + + await save_account(result, config.accounts_output_file) + + return result + + except Exception as e: + logger.exception(f"[Go {task_id}] Unexpected error") + return { + "status": "failed", + "failed_stage": "exception", + "error": str(e), + } + + finally: + if session: + try: + session.close() + except Exception as e: + logger.warning(f"[Go {task_id}] Error closing session: {e}") diff --git a/bot/handlers/settings.py b/bot/handlers/settings.py new file mode 100644 index 0000000..a295bc0 --- /dev/null +++ b/bot/handlers/settings.py @@ -0,0 +1,96 @@ +""" +/set 命令处理器 - 用户设置 +""" + +from telegram import Update +from telegram.ext import ContextTypes + +from bot.middlewares.auth import require_auth + + +DEFAULT_WORKERS = 2 +MIN_WORKERS = 1 +MAX_WORKERS = 5 + + +def get_user_settings(context: ContextTypes.DEFAULT_TYPE, user_id: int) -> dict: + """获取用户设置""" + if "user_settings" not in context.bot_data: + context.bot_data["user_settings"] = {} + + if user_id not in context.bot_data["user_settings"]: + context.bot_data["user_settings"][user_id] = { + "workers": DEFAULT_WORKERS, + } + + return context.bot_data["user_settings"][user_id] + + +def set_user_setting( + context: ContextTypes.DEFAULT_TYPE, user_id: int, key: str, value: any +) -> None: + """设置用户配置""" + settings = get_user_settings(context, user_id) + settings[key] = value + + +@require_auth +async def set_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """ + 处理 /set 命令 - 用户设置 + + 用法: + /set - 查看当前设置 + /set workers <数量> - 设置并发数 (1-5) + """ + user_id = update.effective_user.id + settings = get_user_settings(context, user_id) + + # 无参数：显示当前设置 + if not context.args: + await update.message.reply_text( + f"⚙️ **当前设置**\n\n" + f"👷 并发数: {settings['workers']}\n\n" + f"使用 `/set workers <数量>` 修改并发数 ({MIN_WORKERS}-{MAX_WORKERS})", + parse_mode="Markdown", + ) + return + + # 解析设置项 + setting_name = context.args[0].lower() + + if setting_name == "workers": + if len(context.args) < 2: + await update.message.reply_text( + f"❌ 请指定并发数\n\n" + f"用法: `/set workers <数量>`\n" + f"范围: {MIN_WORKERS}-{MAX_WORKERS}", + parse_mode="Markdown", + ) + return + + try: + workers = int(context.args[1]) + if workers < MIN_WORKERS: + workers = MIN_WORKERS + elif workers > MAX_WORKERS: + workers = MAX_WORKERS + + set_user_setting(context, user_id, "workers", workers) + + await update.message.reply_text( + f"✅ 并发数已设置为 **{workers}**", + parse_mode="Markdown", + ) + except ValueError: + await update.message.reply_text( + f"❌ 无效的数值: `{context.args[1]}`", + parse_mode="Markdown", + ) + else: + await update.message.reply_text( + f"❌ 未知设置项: `{setting_name}`\n\n" + f"可用设置:\n" + f"• `workers` - 并发数 ({MIN_WORKERS}-{MAX_WORKERS})", + parse_mode="Markdown", + ) diff --git a/bot/handlers/start.py b/bot/handlers/start.py new file mode 100644 index 0000000..95e81c9 --- /dev/null +++ b/bot/handlers/start.py @@ -0,0 +1,57 @@ +""" +/start 和 /help 命令处理器 +""" + +from telegram import Update +from telegram.ext import ContextTypes + +from bot.middlewares.auth import require_auth + + +@require_auth +async def start_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """ + 处理 /start 命令 - 欢迎信息 + """ + user = update.effective_user + welcome_text = f""" +👋 你好, {user.first_name}! + +🔗 **支付链接生成器** + +📋 命令: +• `/go [数量] [plan]` - 生成支付链接 +• `/status [task_id]` - 查看任务状态 +• `/set workers <数量>` - 设置并发数 + +💡 使用 /help 查看详细帮助 +""" + await update.message.reply_text(welcome_text, parse_mode="Markdown") + + +@require_auth +async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """ + 处理 /help 命令 - 详细帮助信息 + """ + help_text = """ +📖 **命令说明** + +**生成支付链接** +`/go` - 生成 1 个 Plus 支付链接 +`/go 5` - 生成 5 个 Plus 支付链接 +`/go 3 pro` - 生成 3 个 Pro 支付链接 + +可选 plan: `plus` (默认) 或 `pro` + +**查看任务状态** +`/status` - 查看所有任务 +`/status ` - 查看指定任务 + +**设置** +`/set` - 查看当前设置 +`/set workers 3` - 设置并发数 (1-5) + +📎 结果将以 JSON 文件形式发送 +""" + await update.message.reply_text(help_text, parse_mode="Markdown") diff --git a/bot/handlers/status.py b/bot/handlers/status.py new file mode 100644 index 0000000..2720cac --- /dev/null +++ b/bot/handlers/status.py @@ -0,0 +1,112 @@ +""" +/status 命令处理器 - 查看任务状态 +""" + +from telegram import Update +from telegram.ext import ContextTypes + +from bot.middlewares.auth import require_auth +from bot.services.task_manager import TaskStatus + + +STATUS_EMOJI = { + TaskStatus.PENDING: "⏳", + TaskStatus.RUNNING: "🔄", + TaskStatus.COMPLETED: "✅", + TaskStatus.FAILED: "❌", + TaskStatus.PARTIAL: "⚠️", +} + + +@require_auth +async def status_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """ + 处理 /status 命令 - 查看任务状态 + + 用法: + /status - 查看所有进行中的任务 + /status - 查看指定任务详情 + """ + task_manager = context.bot_data["task_manager"] + user_id = update.effective_user.id + + # 检查是否指定了任务 ID + if context.args: + task_id = context.args[0] + task = task_manager.get_task(task_id) + + if not task: + await update.message.reply_text(f"❌ 未找到任务: `{task_id}`", parse_mode="Markdown") + return + + # 检查是否是自己的任务（或管理员） + admin_users = context.bot_data.get("admin_users", set()) + if task.user_id != user_id and user_id not in admin_users: + await update.message.reply_text("❌ 你没有权限查看此任务") + return + + # 显示任务详情 + emoji = STATUS_EMOJI.get(task.status, "❓") + progress_pct = (task.completed / task.total * 100) if task.total > 0 else 0 + + text = ( + f"{emoji} **任务详情**\n\n" + f"🆔 ID: `{task.task_id}`\n" + f"📋 类型: {task.task_type}\n" + f"📝 描述: {task.description}\n" + f"📊 状态: {task.status.value}\n" + f"📈 进度: {task.completed}/{task.total} ({progress_pct:.0f}%)\n" + ) + + if task.current_item: + text += f"🔄 当前: {task.current_item}\n" + + if task.result: + text += f"\n📦 结果:\n```\n{task.result}\n```" + + await update.message.reply_text(text, parse_mode="Markdown") + return + + # 获取用户的所有任务 + admin_users = context.bot_data.get("admin_users", set()) + is_admin = user_id in admin_users + + if is_admin: + # 管理员可以看所有任务 + tasks = task_manager.get_all_tasks() + else: + # 普通用户只能看自己的 + tasks = task_manager.get_user_tasks(user_id) + + if not tasks: + await update.message.reply_text( + "📭 没有找到任务记录\n\n" + "使用 `/register` 开始注册账号", + parse_mode="Markdown" + ) + return + + # 构建任务列表 + text = "📋 **任务列表**\n\n" + + # 先显示进行中的任务 + running_tasks = [t for t in tasks if t.status in [TaskStatus.PENDING, TaskStatus.RUNNING]] + completed_tasks = [t for t in tasks if t.status not in [TaskStatus.PENDING, TaskStatus.RUNNING]] + + if running_tasks: + text += "**进行中:**\n" + for task in running_tasks[:5]: + emoji = STATUS_EMOJI.get(task.status, "❓") + progress_pct = (task.completed / task.total * 100) if task.total > 0 else 0 + text += f"{emoji} `{task.task_id}` - {task.description} ({progress_pct:.0f}%)\n" + text += "\n" + + if completed_tasks: + text += "**已完成:**\n" + for task in completed_tasks[:5]: + emoji = STATUS_EMOJI.get(task.status, "❓") + text += f"{emoji} `{task.task_id}` - {task.description}\n" + + text += f"\n💡 使用 `/status ` 查看详情" + + await update.message.reply_text(text, parse_mode="Markdown") diff --git a/bot/middlewares/__init__.py b/bot/middlewares/__init__.py new file mode 100644 index 0000000..70bc342 --- /dev/null +++ b/bot/middlewares/__init__.py @@ -0,0 +1,7 @@ +""" +Bot 中间件模块 +""" + +from bot.middlewares.auth import require_auth, AuthMiddleware + +__all__ = ["require_auth", "AuthMiddleware"] diff --git a/bot/middlewares/auth.py b/bot/middlewares/auth.py new file mode 100644 index 0000000..b533e7c --- /dev/null +++ b/bot/middlewares/auth.py @@ -0,0 +1,111 @@ +""" +用户鉴权中间件 + +检查用户是否在白名单中 +""" + +from functools import wraps +from typing import Callable, Any + +from telegram import Update +from telegram.ext import ContextTypes + + +class AuthMiddleware: + """ + 鉴权中间件类 + """ + + def __init__(self, allowed_users: set, admin_users: set = None): + """ + 初始化鉴权中间件 + + 参数: + allowed_users: 允许使用的用户 ID 集合 + admin_users: 管理员用户 ID 集合 + """ + self.allowed_users = allowed_users or set() + self.admin_users = admin_users or set() + + def is_allowed(self, user_id: int) -> bool: + """检查用户是否允许使用""" + # 如果白名单为空，允许所有人 + if not self.allowed_users: + return True + return user_id in self.allowed_users or user_id in self.admin_users + + def is_admin(self, user_id: int) -> bool: + """检查用户是否是管理员""" + return user_id in self.admin_users + + +def require_auth(func: Callable) -> Callable: + """ + 鉴权装饰器 + + 用于命令处理函数，检查用户是否在白名单中 + + 用法: + @require_auth + async def my_command(update: Update, context: ContextTypes.DEFAULT_TYPE): + ... + """ + + @wraps(func) + async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs) -> Any: + user = update.effective_user + if not user: + return + + user_id = user.id + allowed_users = context.bot_data.get("allowed_users", set()) + admin_users = context.bot_data.get("admin_users", set()) + + # 如果白名单为空，允许所有人使用 + if not allowed_users and not admin_users: + return await func(update, context, *args, **kwargs) + + # 检查是否在白名单或管理员列表中 + if user_id in allowed_users or user_id in admin_users: + return await func(update, context, *args, **kwargs) + + # 未授权 + await update.message.reply_text( + f"⛔ 你没有权限使用此 Bot\n\n" + f"你的用户 ID: `{user_id}`\n\n" + f"请联系管理员将你的 ID 添加到白名单", + parse_mode="Markdown" + ) + return None + + return wrapper + + +def require_admin(func: Callable) -> Callable: + """ + 管理员权限装饰器 + + 用于需要管理员权限的命令 + + 用法: + @require_admin + async def admin_command(update: Update, context: ContextTypes.DEFAULT_TYPE): + ... + """ + + @wraps(func) + async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs) -> Any: + user = update.effective_user + if not user: + return + + user_id = user.id + admin_users = context.bot_data.get("admin_users", set()) + + if user_id in admin_users: + return await func(update, context, *args, **kwargs) + + await update.message.reply_text("⛔ 此命令需要管理员权限") + return None + + return wrapper diff --git a/bot/services/__init__.py b/bot/services/__init__.py new file mode 100644 index 0000000..55c0280 --- /dev/null +++ b/bot/services/__init__.py @@ -0,0 +1,7 @@ +""" +Bot 服务模块 +""" + +from bot.services.task_manager import TaskManager, Task, TaskStatus + +__all__ = ["TaskManager", "Task", "TaskStatus"] diff --git a/bot/services/file_sender.py b/bot/services/file_sender.py new file mode 100644 index 0000000..d8f1e3a --- /dev/null +++ b/bot/services/file_sender.py @@ -0,0 +1,101 @@ +""" +JSON 文件生成与发送服务 +""" + +import json +import os +import tempfile +from datetime import datetime +from typing import Any + +from telegram import Update +from telegram.ext import ContextTypes + + +async def send_results_as_json( + update: Update, + context: ContextTypes.DEFAULT_TYPE, + task_id: str, + plan: str, + results: list[dict[str, Any]], +) -> None: + """ + 将结果生成为 JSON 文件并发送给用户 + + 参数: + update: Telegram Update 对象 + context: Telegram Context 对象 + task_id: 任务 ID + plan: 订阅计划名称 + results: 结果列表 + """ + # 统计 + success_count = sum(1 for r in results if r.get("status") == "success") + failed_count = len(results) - success_count + + # 构建 JSON 数据 + now = datetime.now() + data = { + "task_id": task_id, + "plan": plan, + "created_at": now.strftime("%Y-%m-%d %H:%M:%S"), + "results": [ + { + "email": r.get("email", ""), + "password": r.get("password", ""), + "url": r.get("checkout_url", ""), + "checkout_session_id": r.get("checkout_session_id", ""), + "status": r.get("status", "failed"), + } + for r in results + ], + "summary": { + "total": len(results), + "success": success_count, + "failed": failed_count, + }, + } + + # 生成文件名 + timestamp = now.strftime("%Y%m%d_%H%M%S") + filename = f"checkout_{timestamp}.json" + + # 写入临时文件并发送 + with tempfile.NamedTemporaryFile( + mode="w", + suffix=".json", + prefix="checkout_", + delete=False, + encoding="utf-8", + ) as f: + json.dump(data, f, indent=2, ensure_ascii=False) + temp_path = f.name + + try: + # 发送消息（不含敏感信息） + message_text = ( + f"✅ **任务完成**\n\n" + f"📦 Plan: {plan}\n" + f"✅ 成功: {success_count}\n" + f"❌ 失败: {failed_count}\n" + ) + + await context.bot.send_message( + chat_id=update.effective_chat.id, + text=message_text, + parse_mode="Markdown", + ) + + # 发送 JSON 文件 + with open(temp_path, "rb") as f: + await context.bot.send_document( + chat_id=update.effective_chat.id, + document=f, + filename=filename, + caption=f"📎 任务 `{task_id}` 结果", + parse_mode="Markdown", + ) + finally: + # 清理临时文件 + if os.path.exists(temp_path): + os.remove(temp_path) diff --git a/bot/services/task_manager.py b/bot/services/task_manager.py new file mode 100644 index 0000000..55a7f4c --- /dev/null +++ b/bot/services/task_manager.py @@ -0,0 +1,228 @@ +""" +后台任务管理器 + +管理异步执行的注册、登录等任务的状态和进度 +""" + +import uuid +from datetime import datetime +from enum import Enum +from typing import Dict, List, Optional, Any +from dataclasses import dataclass, field + + +class TaskStatus(Enum): + """任务状态枚举""" + PENDING = "pending" # 等待执行 + RUNNING = "running" # 执行中 + COMPLETED = "completed" # 已完成 + FAILED = "failed" # 失败 + PARTIAL = "partial" # 部分成功 + + +@dataclass +class Task: + """任务数据类""" + task_id: str + user_id: int + task_type: str # register, login, checkout + status: TaskStatus + description: str + total: int = 0 # 总任务数 + completed: int = 0 # 已完成数 + current_item: str = "" # 当前处理项 + result: Optional[Any] = None + error: Optional[str] = None + created_at: datetime = field(default_factory=datetime.now) + updated_at: datetime = field(default_factory=datetime.now) + + +class TaskManager: + """ + 任务管理器 + + 管理后台任务的创建、状态更新和查询 + """ + + def __init__(self, max_history: int = 100): + """ + 初始化任务管理器 + + 参数: + max_history: 保留的历史任务数量 + """ + self._tasks: Dict[str, Task] = {} + self._max_history = max_history + + def create_task( + self, + user_id: int, + task_type: str, + total: int = 1, + description: str = "" + ) -> str: + """ + 创建新任务 + + 参数: + user_id: 用户 ID + task_type: 任务类型 (register, login, checkout) + total: 总任务数 + description: 任务描述 + + 返回: + 任务 ID + """ + task_id = self._generate_task_id() + + task = Task( + task_id=task_id, + user_id=user_id, + task_type=task_type, + status=TaskStatus.RUNNING, + description=description or f"{task_type} task", + total=total, + ) + + self._tasks[task_id] = task + self._cleanup_old_tasks() + + return task_id + + def update_progress( + self, + task_id: str, + completed: Optional[int] = None, + current_item: Optional[str] = None + ) -> bool: + """ + 更新任务进度 + + 参数: + task_id: 任务 ID + completed: 已完成数量 + current_item: 当前处理项描述 + + 返回: + 是否更新成功 + """ + task = self._tasks.get(task_id) + if not task: + return False + + if completed is not None: + task.completed = completed + if current_item is not None: + task.current_item = current_item + + task.updated_at = datetime.now() + return True + + def complete_task( + self, + task_id: str, + status: TaskStatus = TaskStatus.COMPLETED, + result: Any = None, + error: str = None + ) -> bool: + """ + 完成任务 + + 参数: + task_id: 任务 ID + status: 最终状态 + result: 任务结果 + error: 错误信息（如果失败） + + 返回: + 是否更新成功 + """ + task = self._tasks.get(task_id) + if not task: + return False + + task.status = status + task.result = result + task.error = error + task.completed = task.total + task.current_item = "" + task.updated_at = datetime.now() + + return True + + def get_task(self, task_id: str) -> Optional[Task]: + """ + 获取任务 + + 参数: + task_id: 任务 ID + + 返回: + Task 对象，如果不存在返回 None + """ + return self._tasks.get(task_id) + + def get_user_tasks(self, user_id: int, limit: int = 10) -> List[Task]: + """ + 获取用户的任务列表 + + 参数: + user_id: 用户 ID + limit: 返回数量限制 + + 返回: + 任务列表（按创建时间倒序） + """ + user_tasks = [ + task for task in self._tasks.values() + if task.user_id == user_id + ] + user_tasks.sort(key=lambda t: t.created_at, reverse=True) + return user_tasks[:limit] + + def get_all_tasks(self, limit: int = 20) -> List[Task]: + """ + 获取所有任务（管理员用） + + 参数: + limit: 返回数量限制 + + 返回: + 任务列表（按创建时间倒序） + """ + all_tasks = list(self._tasks.values()) + all_tasks.sort(key=lambda t: t.created_at, reverse=True) + return all_tasks[:limit] + + def get_running_tasks(self) -> List[Task]: + """ + 获取正在运行的任务 + + 返回: + 运行中的任务列表 + """ + return [ + task for task in self._tasks.values() + if task.status in [TaskStatus.PENDING, TaskStatus.RUNNING] + ] + + def _generate_task_id(self) -> str: + """生成唯一任务 ID""" + return uuid.uuid4().hex[:8] + + def _cleanup_old_tasks(self): + """清理旧任务，保持历史记录在限制内""" + if len(self._tasks) <= self._max_history: + return + + # 获取已完成的任务，按时间排序 + completed_tasks = [ + (task_id, task) for task_id, task in self._tasks.items() + if task.status not in [TaskStatus.PENDING, TaskStatus.RUNNING] + ] + completed_tasks.sort(key=lambda x: x[1].created_at) + + # 删除最旧的任务 + to_remove = len(self._tasks) - self._max_history + for task_id, _ in completed_tasks[:to_remove]: + del self._tasks[task_id] diff --git a/bot_main.py b/bot_main.py new file mode 100644 index 0000000..517f94d --- /dev/null +++ b/bot_main.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +""" +Telegram Bot 启动入口 + +使用方法: + python bot_main.py + +配置 (.env): + TELEGRAM_BOT_TOKEN=your_bot_token_here + TELEGRAM_ALLOWED_USERS=123456789,987654321 + TELEGRAM_ADMIN_USERS=123456789 +""" + +import asyncio +import sys +from pathlib import Path + +# 确保项目根目录在 Python 路径中 +sys.path.insert(0, str(Path(__file__).parent)) + +from config import load_config +from utils.logger import logger, setup_logger + + +def main(): + """主函数""" + print("=" * 60) + print(" OpenAI 账号自动注册系统 - Telegram Bot") + print("=" * 60) + print() + + # 加载配置 + config = load_config() + setup_logger(config.log_level) + + # 检查配置 + if not config.telegram_bot_enabled: + logger.error("Telegram Bot is disabled in config") + print("❌ Telegram Bot 已禁用，请设置 TELEGRAM_BOT_ENABLED=true") + sys.exit(1) + + if not config.telegram_bot_token: + logger.error("TELEGRAM_BOT_TOKEN is not configured") + print("❌ 请在 .env 文件中配置 TELEGRAM_BOT_TOKEN") + print() + print("步骤:") + print("1. 在 Telegram 中找到 @BotFather") + print("2. 发送 /newbot 创建新 Bot") + print("3. 将获得的 Token 填入 .env 文件:") + print(" TELEGRAM_BOT_TOKEN=your_token_here") + sys.exit(1) + + # 打印配置信息 + allowed_count = len([x for x in config.telegram_allowed_users.split(",") if x.strip()]) + admin_count = len([x for x in config.telegram_admin_users.split(",") if x.strip()]) + + logger.info(f"Bot Token: {config.telegram_bot_token[:10]}...") + logger.info(f"Allowed users: {allowed_count} configured") + logger.info(f"Admin users: {admin_count} configured") + + if not config.telegram_allowed_users and not config.telegram_admin_users: + logger.warning("⚠️ No user whitelist configured - Bot will be accessible to everyone!") + print() + print("⚠️ 警告: 未配置用户白名单，所有人都可以使用此 Bot") + print(" 建议在 .env 中添加:") + print(" TELEGRAM_ALLOWED_USERS=your_user_id") + print() + + # 确保输出目录存在 + Path("output/logs").mkdir(parents=True, exist_ok=True) + Path("output/tokens").mkdir(parents=True, exist_ok=True) + + # 创建并运行 Bot + try: + from bot.app import create_application + + logger.info("Creating Telegram Bot application...") + application = create_application(config) + + logger.info("Starting Bot in polling mode...") + print() + print("🤖 Bot 已启动! 在 Telegram 中发送 /start 开始使用") + print(" 按 Ctrl+C 停止") + print() + + # 运行 Bot + application.run_polling( + allowed_updates=["message"], + drop_pending_updates=True # 忽略 Bot 离线期间的消息 + ) + + except KeyboardInterrupt: + print() + logger.info("Bot stopped by user") + print("👋 Bot 已停止") + + except Exception as e: + logger.exception("Failed to start Bot") + print(f"❌ 启动失败: {e}") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/config.py b/config.py index fe4e475..cb402ba 100644 --- a/config.py +++ b/config.py @@ -287,6 +287,24 @@ class AppConfig(BaseSettings): description="HTTP 请求超时时间（秒）" ) + # ========== Telegram Bot 配置 ========== + telegram_bot_token: Optional[str] = Field( + default=None, + description="Telegram Bot Token (从 @BotFather 获取)" + ) + telegram_allowed_users: str = Field( + default="", + description="允许使用 Bot 的用户 ID（逗号分隔）" + ) + telegram_admin_users: str = Field( + default="", + description="管理员用户 ID（逗号分隔，拥有更高权限）" + ) + telegram_bot_enabled: bool = Field( + default=True, + description="是否启用 Telegram Bot" + ) + @field_validator("proxy_pool") @classmethod def validate_proxy_pool(cls, v: str) -> str: diff --git a/pyproject.toml b/pyproject.toml index 2dd4b0c..a95a366 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,4 +11,5 @@ dependencies = [ "loguru>=0.7.2", # 日志系统 "python-dotenv>=1.0.0", # 环境变量加载 "httpx>=0.25.0", # 异步 HTTP 客户端（可选，用于外部 API） + "python-telegram-bot[job-queue]>=20.7", # Telegram Bot ] diff --git a/uv.lock b/uv.lock index 3fe2f5a..22d3ecd 100644 --- a/uv.lock +++ b/uv.lock @@ -1,6 +1,10 @@ version = 1 revision = 3 requires-python = ">=3.12" +resolution-markers = [ + "python_full_version >= '3.14'", + "python_full_version < '3.14'", +] [[package]] name = "annotated-types" @@ -24,6 +28,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/38/0e/27be9fdef66e72d64c0cdc3cc2823101b80585f8119b5c112c2e8f5f7dab/anyio-4.12.1-py3-none-any.whl", hash = "sha256:d405828884fc140aa80a3c667b8beed277f1dfedec42ba031bd6ac3db606ab6c", size = 113592, upload-time = "2026-01-06T11:45:19.497Z" }, ] +[[package]] +name = "apscheduler" +version = "3.11.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "tzlocal" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/07/12/3e4389e5920b4c1763390c6d371162f3784f86f85cd6d6c1bfe68eef14e2/apscheduler-3.11.2.tar.gz", hash = "sha256:2a9966b052ec805f020c8c4c3ae6e6a06e24b1bf19f2e11d91d8cca0473eef41", size = 108683, upload-time = "2025-12-22T00:39:34.884Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/64/2e54428beba8d9992aa478bb8f6de9e4ecaa5f8f513bcfd567ed7fb0262d/apscheduler-3.11.2-py3-none-any.whl", hash = "sha256:ce005177f741409db4e4dd40a7431b76feb856b9dd69d57e0da49d6715bfd26d", size = 64439, upload-time = "2025-12-22T00:39:33.303Z" }, +] + [[package]] name = "certifi" version = "2026.1.4" @@ -133,6 +149,7 @@ dependencies = [ { name = "pydantic" }, { name = "pydantic-settings" }, { name = "python-dotenv" }, + { name = "python-telegram-bot", extra = ["job-queue"] }, ] [package.metadata] @@ -143,6 +160,7 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.5.0" }, { name = "pydantic-settings", specifier = ">=2.1.0" }, { name = "python-dotenv", specifier = ">=1.0.0" }, + { name = "python-telegram-bot", extras = ["job-queue"], specifier = ">=20.7" }, ] [[package]] @@ -322,6 +340,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/14/1b/a298b06749107c305e1fe0f814c6c74aea7b2f1e10989cb30f544a1b3253/python_dotenv-1.2.1-py3-none-any.whl", hash = "sha256:b81ee9561e9ca4004139c6cbba3a238c32b03e4894671e181b671e8cb8425d61", size = 21230, upload-time = "2025-10-26T15:12:09.109Z" }, ] +[[package]] +name = "python-telegram-bot" +version = "22.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpcore", marker = "python_full_version >= '3.14'" }, + { name = "httpx" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cd/9b/8df90c85404166a6631e857027866263adb27440d8af1dbeffbdc4f0166c/python_telegram_bot-22.6.tar.gz", hash = "sha256:50ae8cc10f8dff01445628687951020721f37956966b92a91df4c1bf2d113742", size = 1503761, upload-time = "2026-01-24T13:57:00.269Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/13/97/7298f0e1afe3a1ae52ff4c5af5087ed4de319ea73eb3b5c8c4dd4e76e708/python_telegram_bot-22.6-py3-none-any.whl", hash = "sha256:e598fe171c3dde2dfd0f001619ee9110eece66761a677b34719fb18934935ce0", size = 737267, upload-time = "2026-01-24T13:56:58.06Z" }, +] + +[package.optional-dependencies] +job-queue = [ + { name = "apscheduler" }, +] + [[package]] name = "typing-extensions" version = "4.15.0" @@ -343,6 +379,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, ] +[[package]] +name = "tzdata" +version = "2025.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5e/a7/c202b344c5ca7daf398f3b8a477eeb205cf3b6f32e7ec3a6bac0629ca975/tzdata-2025.3.tar.gz", hash = "sha256:de39c2ca5dc7b0344f2eba86f49d614019d29f060fc4ebc8a417896a620b56a7", size = 196772, upload-time = "2025-12-13T17:45:35.667Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/b0/003792df09decd6849a5e39c28b513c06e84436a54440380862b5aeff25d/tzdata-2025.3-py2.py3-none-any.whl", hash = "sha256:06a47e5700f3081aab02b2e513160914ff0694bce9947d6b76ebd6bf57cfc5d1", size = 348521, upload-time = "2025-12-13T17:45:33.889Z" }, +] + +[[package]] +name = "tzlocal" +version = "5.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "tzdata", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8b/2e/c14812d3d4d9cd1773c6be938f89e5735a1f11a9f184ac3639b93cef35d5/tzlocal-5.3.1.tar.gz", hash = "sha256:cceffc7edecefea1f595541dbd6e990cb1ea3d19bf01b2809f362a03dd7921fd", size = 30761, upload-time = "2025-03-05T21:17:41.549Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/14/e2a54fabd4f08cd7af1c07030603c3356b74da07f7cc056e600436edfa17/tzlocal-5.3.1-py3-none-any.whl", hash = "sha256:eb1a66c3ef5847adf7a834f1be0800581b683b5608e74f86ecbcef8ab91bb85d", size = 18026, upload-time = "2025-03-05T21:17:39.857Z" }, +] + [[package]] name = "win32-setctime" version = "1.2.0"