4

api_register.py

@@ -20,6 +20,20 @@ except ImportError:
 import requests
 
 
+def _is_shutdown_requested():
+    """检查是否收到停止请求"""
+    try:
+        import run
+        return run._shutdown_requested
+    except Exception:
+        return False
+
+
+class ShutdownRequested(Exception):
+    """用户请求停止异常"""
+    pass
+
+
 def log_status(step, message):
     """日志输出"""
     timestamp = time.strftime("%H:%M:%S")
@@ -330,12 +344,20 @@ def get_verification_code_api(target_email: str, mail_api_base: str, mail_api_to
         
     Returns:
         str: 验证码，失败返回空字符串
+        
+    Raises:
+        ShutdownRequested: 用户请求停止时抛出
     """
     log_status("API监听", "正在监听邮件...")
     headers = {"Authorization": mail_api_token, "Content-Type": "application/json"}
     start_time = time.time()
     
     for i in range(max_retries):
+        # 检查停止请求
+        if _is_shutdown_requested():
+            log_status("停止", "[!] 检测到停止请求，中断邮件监听")
+            raise ShutdownRequested("用户请求停止")
+        
         elapsed = int(time.time() - start_time)
         try:
             url = f"{mail_api_base}/api/public/emailList"
@@ -389,6 +411,9 @@ def api_register_flow(
         
     Returns:
         ChatGPTAPIRegister: 成功返回 reg 对象，失败返回 None
+        
+    Raises:
+        ShutdownRequested: 用户请求停止时抛出
     """
     def log_cb(msg):
         if progress_callback:
@@ -396,57 +421,72 @@ def api_register_flow(
         else:
             log_progress(msg)
     
+    def check_shutdown():
+        """检查停止请求"""
+        if _is_shutdown_requested():
+            log_cb("[!] 检测到停止请求")
+            raise ShutdownRequested("用户请求停止")
+    
     reg = ChatGPTAPIRegister(proxy=proxy)
     
     try:
+        check_shutdown()
         log_status("API注册", "初始化会话...")
         if not reg.init_session():
             log_cb("[X] 初始化失败")
             return None
         log_cb("[OK] 会话初始化成功")
         
+        check_shutdown()
         log_status("API注册", "获取授权 URL...")
         if not reg.get_authorize_url(email):
             log_cb("[X] 获取授权 URL 失败")
             return None
         log_cb("[OK] 授权 URL 获取成功")
         
+        check_shutdown()
         log_status("API注册", "开始授权流程...")
         if not reg.start_authorize():
             log_cb("[X] 授权流程启动失败")
             return None
         log_cb("[OK] 授权流程已启动")
         
+        check_shutdown()
         log_status("API注册", "注册账户...")
         if not reg.register(email, password):
             log_cb("[X] 注册失败")
             return None
         log_cb("[OK] 账户注册成功")
         
+        check_shutdown()
         log_status("API注册", "发送验证邮件...")
         if not reg.send_verification_email():
             log_cb("[X] 发送验证邮件失败")
             return None
         log_cb("[OK] 验证邮件已发送")
         
+        check_shutdown()
         # 获取验证码
         otp_code = get_verification_code_api(email, mail_api_base, mail_api_token)
         if not otp_code:
             log_cb("[X] 未能获取验证码")
             return None
         
+        check_shutdown()
         log_status("API注册", f"验证 OTP: {otp_code}")
         if not reg.validate_otp(otp_code):
             log_cb("[X] OTP 验证失败")
             return None
         log_cb("[OK] OTP 验证成功")
         
+        check_shutdown()
         # 创建账户（带重试）
         log_status("API注册", "创建账户...")
         create_success = reg.create_account(real_name, birthdate)
         
         # 如果创建失败，重新获取验证码再试一次
         if not create_success:
+            check_shutdown()
             log_cb("[!] 创建账户失败，尝试重新验证...")
             
             # 重新发送验证邮件
@@ -456,6 +496,7 @@ def api_register_flow(
                 return None
             log_cb("[OK] 验证邮件已重新发送")
             
+            check_shutdown()
             # 重新获取验证码
             time.sleep(2)  # 等待新邮件
             otp_code = get_verification_code_api(email, mail_api_base, mail_api_token)
@@ -463,12 +504,14 @@ def api_register_flow(
                 log_cb("[X] 未能获取新验证码")
                 return None
             
+            check_shutdown()
             log_status("API注册", f"重新验证 OTP: {otp_code}")
             if not reg.validate_otp(otp_code):
                 log_cb("[X] OTP 重新验证失败")
                 return None
             log_cb("[OK] OTP 重新验证成功")
             
+            check_shutdown()
             # 再次尝试创建账户
             log_status("API注册", "重新创建账户...")
             if not reg.create_account(real_name, birthdate):
@@ -477,6 +520,7 @@ def api_register_flow(
         
         log_cb("[OK] 账户创建成功")
         
+        check_shutdown()
         # 验证 session 是否有效
         token = reg.get_session_token()
         if token:
@@ -486,6 +530,8 @@ def api_register_flow(
         
         return reg
         
+    except ShutdownRequested:
+        raise  # 重新抛出停止请求异常
     except Exception as e:
         log_status("错误", f"注册异常: {e}")
         return None