checkout

main.py

@@ -22,9 +22,11 @@ from datetime import datetime
 
 from core.session import OAISession
 from core.flow import RegisterFlow
+from core.checkout import CheckoutFlow
 from config import load_config
 from utils.logger import logger, setup_logger
 import random
+import json
 
 
 async def login_account(
@@ -96,6 +98,173 @@ async def login_account(
                 logger.warning(f"[Login {task_id}] Error closing session: {e}")
 
 
+async def get_checkout_session(
+    config,
+    email: str,
+    password: str,
+    plan_name: str = "chatgptplusplan",
+    task_id: int = 1
+) -> Dict[str, Any]:
+    """
+    获取账号的支付账单链接
+
+    参数:
+        config: AppConfig 配置对象
+        email: 登录邮箱
+        password: 登录密码
+        plan_name: 订阅计划名称 (chatgptplusplan/chatgptproplan)
+        task_id: 任务 ID（用于日志标识）
+
+    返回:
+        支付会话结果字典
+    """
+    proxy = config.proxy.get_next_proxy()
+    if proxy:
+        logger.info(f"[Checkout {task_id}] Using proxy: {_mask_proxy(proxy)}")
+    else:
+        logger.info(f"[Checkout {task_id}] No proxy configured, using direct connection")
+
+    session = None
+    try:
+        session = OAISession(
+            proxy=proxy,
+            impersonate=config.tls_impersonate
+        )
+
+        # 先登录
+        logger.info(f"[Checkout {task_id}] Logging in as {email}...")
+        login_result = await session.login(email, password)
+
+        if login_result.get("status") != "success":
+            logger.error(f"[Checkout {task_id}] ❌ Login failed: {login_result.get('error')}")
+            return {
+                "task_id": task_id,
+                "email": email,
+                "status": "failed",
+                "error": f"Login failed: {login_result.get('error')}"
+            }
+
+        # 获取支付会话
+        logger.info(f"[Checkout {task_id}] Creating checkout session for {plan_name}...")
+        checkout = CheckoutFlow(session, plan_name=plan_name)
+        result = await checkout.create_checkout_session()
+
+        result["task_id"] = task_id
+        result["email"] = email
+        result["proxy"] = _mask_proxy(proxy) if proxy else "none"
+
+        if result.get("status") == "success":
+            logger.success(
+                f"[Checkout {task_id}] ✅ Checkout session created for {email}"
+            )
+            # 保存结果
+            await save_checkout_result(result)
+        else:
+            logger.error(
+                f"[Checkout {task_id}] ❌ Checkout failed: {result.get('error', 'Unknown error')}"
+            )
+
+        return result
+
+    except Exception as e:
+        logger.exception(f"[Checkout {task_id}] Unexpected error")
+        return {
+            "task_id": task_id,
+            "email": email,
+            "status": "failed",
+            "error": str(e),
+            "message": f"Task exception: {type(e).__name__}"
+        }
+
+    finally:
+        if session:
+            try:
+                session.close()
+            except Exception as e:
+                logger.warning(f"[Checkout {task_id}] Error closing session: {e}")
+
+
+async def save_checkout_result(result: Dict[str, Any]):
+    """
+    保存支付会话结果到文件
+
+    参数:
+        result: 支付会话结果字典
+    """
+    output_file = Path("output/checkouts.json")
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+    # 读取现有数据
+    checkouts = []
+    if output_file.exists():
+        try:
+            with open(output_file, "r", encoding="utf-8") as f:
+                checkouts = json.load(f)
+        except (json.JSONDecodeError, Exception):
+            checkouts = []
+
+    # 添加新记录
+    checkout_data = {
+        "email": result.get("email", "unknown"),
+        "plan_name": result.get("plan_name"),
+        "checkout_session_id": result.get("checkout_session_id"),
+        "client_secret": result.get("client_secret"),
+        "publishable_key": result.get("publishable_key"),
+        "payment_status": result.get("payment_status"),
+        "url": result.get("url"),
+        "timestamp": timestamp
+    }
+    checkouts.append(checkout_data)
+
+    # 写入文件
+    async with asyncio.Lock():
+        with open(output_file, "w", encoding="utf-8") as f:
+            json.dump(checkouts, f, indent=2, ensure_ascii=False)
+
+    logger.info(f"Checkout saved to checkouts.json: {result.get('email')}")
+
+
+async def run_batch_checkout(config, accounts: List[Dict], plan_name: str = "chatgptplusplan"):
+    """
+    批量获取账号的支付链接
+
+    参数:
+        config: AppConfig 配置对象
+        accounts: 账号列表 [{"email": "...", "password": "..."}, ...]
+        plan_name: 订阅计划
+
+    返回:
+        结果列表
+    """
+    logger.info(f"Starting batch checkout: {len(accounts)} accounts")
+
+    all_results = []
+    for idx, acc in enumerate(accounts, 1):
+        email = acc.get("email", "")
+        password = acc.get("password", "")
+
+        if not email or not password:
+            logger.warning(f"[Checkout {idx}] Skipping - missing email or password")
+            continue
+
+        result = await get_checkout_session(
+            config,
+            email=email,
+            password=password,
+            plan_name=plan_name,
+            task_id=idx
+        )
+        all_results.append(result)
+
+        # 延迟避免速率限制
+        if idx < len(accounts):
+            delay = random.uniform(2, 4)
+            logger.info(f"Waiting {delay:.1f}s before next checkout...")
+            await asyncio.sleep(delay)
+
+    return all_results
+
+
 async def save_token(result: Dict[str, Any], config):
     """
     保存登录 token 到文件
@@ -104,7 +273,7 @@ async def save_token(result: Dict[str, Any], config):
         result: 登录结果字典
         config: 配置对象
     """
-    output_path = Path("tokens.txt")
+    output_path = Path("output/tokens.txt")
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
     email = result.get("email", "unknown")
@@ -119,7 +288,7 @@ async def save_token(result: Dict[str, Any], config):
             f.write(line)
 
     # 同时保存完整 token 到单独文件
-    token_file = Path(f"tokens/{email.replace('@', '_at_')}.txt")
+    token_file = Path(f"output/tokens/{email.replace('@', '_at_')}.txt")
     token_file.parent.mkdir(parents=True, exist_ok=True)
     with open(token_file, "w", encoding="utf-8") as f:
         f.write(f"Email: {email}\n")
@@ -185,7 +354,28 @@ async def register_account(
             if login_result.get("status") == "success":
                 result["access_token"] = login_result.get("access_token")
                 result["session_token"] = login_result.get("session_token")
+                # 更新 session 对象的 token（供 CheckoutFlow 使用）
+                session.access_token = login_result.get("access_token")
+                session.session_token = login_result.get("session_token")
                 logger.success(f"[Task {task_id}] ✅ Token obtained!")
+
+                # 获取支付链接
+                logger.info(f"[Task {task_id}] Getting checkout session...")
+                checkout = CheckoutFlow(session, plan_name="chatgptplusplan")
+                checkout_result = await checkout.create_checkout_session()
+
+                if checkout_result.get("status") == "success":
+                    result["checkout_session_id"] = checkout_result.get("checkout_session_id")
+                    result["client_secret"] = checkout_result.get("client_secret")
+                    result["publishable_key"] = checkout_result.get("publishable_key")
+                    result["payment_status"] = checkout_result.get("payment_status")
+                    logger.success(f"[Task {task_id}] ✅ Checkout session created!")
+
+                    # 保存 checkout 结果
+                    checkout_result["email"] = result["email"]
+                    await save_checkout_result(checkout_result)
+                else:
+                    logger.warning(f"[Task {task_id}] ⚠️ Checkout failed: {checkout_result.get('error')}")
             else:
                 logger.warning(f"[Task {task_id}] ⚠️ Login failed: {login_result.get('error')}")
 
@@ -247,7 +437,7 @@ async def save_account(result: Dict[str, Any], output_file: str):
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
     # 保存到 JSON 文件
-    json_file = Path("accounts.json")
+    json_file = Path("output/accounts.json")
 
     # 读取现有数据
     accounts = []
@@ -410,7 +600,7 @@ async def main():
             return
 
     # 确保必要的目录存在
-    Path("logs").mkdir(exist_ok=True)
+    Path("output/logs").mkdir(parents=True, exist_ok=True)
     Path(config.accounts_output_file).parent.mkdir(parents=True, exist_ok=True)
 
     # 询问要注册的账号数量
@@ -504,7 +694,7 @@ async def main():
     # 保存详细结果到 JSON（可选）
     if config.log_to_file:
         import json
-        result_file = f"logs/results_{int(time.time())}.json"
+        result_file = f"output/logs/results_{int(time.time())}.json"
         with open(result_file, "w", encoding="utf-8") as f:
             json.dump(all_results, f, indent=2, ensure_ascii=False)
         logger.info(f"Detailed results saved to: {result_file}")