Update nb2.py

nb2.py

@@ -6,6 +6,10 @@ from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.common.exceptions import TimeoutException
 from selenium.webdriver.common.keys import Keys
+from webdriver_manager.chrome import ChromeDriverManager
+from fastapi import FastAPI, BackgroundTasks
+import uvicorn
+import asyncio
 import time
 import random
 import string
@@ -14,8 +18,6 @@ import requests
 import json
 import re
 import os
-from fastapi import FastAPI
-import uvicorn
 
 # 配置日志
 logging.basicConfig(
@@ -24,9 +26,6 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# FastAPI应用
-app = FastAPI()
-
 class MailGw:
     def __init__(self):
         self.base_url = "https://api.mail.gw"
@@ -132,20 +131,7 @@ class MailGw:
 
         logger.error("未能获取到验证码")
         return None
-
-def setup_driver():
-    """设置Chrome驱动"""
-    chrome_options = Options()
-    chrome_options.add_argument('--headless')
-    chrome_options.add_argument('--no-sandbox')
-    chrome_options.add_argument('--disable-dev-shm-usage')
-    chrome_options.binary_location = os.getenv('CHROME_BIN', '/usr/bin/chromium')
-    
-    service = Service(os.getenv('CHROMEDRIVER_PATH', '/usr/bin/chromedriver'))
-    driver = webdriver.Chrome(service=service, options=chrome_options)
-    return driver
-
-def get_cookie(driver):
+        def get_cookie(driver):
     try:
         cookie_script = """
         var cookieString = document.cookie;
@@ -296,36 +282,93 @@ def registration_process(driver, wait):
             cookie = get_cookie(driver)
             if cookie:
                 logger.info("成功获取cookie")
-                return cookie
+                return True
             else:
                 logger.error("获取cookie失败")
-                return None
+                return False
         else:
             logger.error("URL未变化为预期值")
-            return None
+            return False
 
     except Exception as e:
         logger.error(f"注册过程错误: {str(e)}")
-        return None
+        return False
 
+app = FastAPI()
+cookie_generator = None
+
+class CookieGenerator:
+    def __init__(self):
+        self.driver = None
+        self.running = False
+        self.cookies = []
+
+    async def start_generation(self):
+        if self.running:
+            return
+        
+        self.running = True
+        try:
+            chrome_options = Options()
+            chrome_options.add_argument('--headless')
+            chrome_options.add_argument('--no-sandbox')
+            chrome_options.add_argument('--disable-dev-shm-usage')
+            chrome_options.binary_location = os.getenv('CHROME_BIN', '/usr/bin/chromium')
+            
+            service = Service(os.getenv('CHROMEDRIVER_PATH', '/usr/bin/chromedriver'))
+            self.driver = webdriver.Chrome(service=service, options=chrome_options)
+            wait = WebDriverWait(self.driver, 20)
+
+            while self.running:
+                success = registration_process(self.driver, wait)
+                if success:
+                    logger.info("本次循环成功")
+                else:
+                    logger.error("本次循环失败")
+                await asyncio.sleep(5)
+
+        except Exception as e:
+            logger.error(f"Cookie生成错误: {str(e)}")
+        finally:
+            self.running = False
+            if self.driver:
+                self.driver.quit()
+
+@app.on_event("startup")
+async def startup_event():
+    global cookie_generator
+    cookie_generator = CookieGenerator()
+    
 @app.get("/")
-def read_root():
+async def read_root():
     return {"status": "running"}
 
-@app.get("/generate_cookie")
-async def generate_cookie():
+@app.get("/start")
+async def start_generation(background_tasks: BackgroundTasks):
+    global cookie_generator
+    if not cookie_generator.running:
+        background_tasks.add_task(cookie_generator.start_generation)
+        return {"status": "started"}
+    return {"status": "already running"}
+
+@app.get("/stop")
+async def stop_generation():
+    global cookie_generator
+    if cookie_generator.running:
+        cookie_generator.running = False
+        return {"status": "stopping"}
+    return {"status": "not running"}
+
+@app.get("/cookies")
+async def get_cookies():
     try:
-        driver = setup_driver()
-        wait = WebDriverWait(driver, 20)
-        cookie = registration_process(driver, wait)
-        driver.quit()
-        
-        if cookie:
-            return {"status": "success", "cookie": cookie}
-        else:
-            return {"status": "error", "message": "Failed to generate cookie"}
+        if os.path.exists('cookies.txt'):
+            with open('cookies.txt', 'r') as file:
+                cookies = file.readlines()
+            return {"cookies": cookies}
+        return {"cookies": []}
     except Exception as e:
-        return {"status": "error", "message": str(e)}
+        return {"error": str(e)}
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)