重构机器人
This commit is contained in:
dela
290
bot/handlers/go.py
Normal file
290
bot/handlers/go.py
Normal file
@@ -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}")
|
||||
Reference in New Issue
Block a user