Update nb2.py

nb2.py

@@ -85,7 +85,6 @@ class MailGw:
         except Exception as e:
             logger.error(f"创建邮箱时发生错误: {str(e)}")
             return None
-
     def get_verification_code(self, max_retries=10, delay=2):
         if not self.token:
             logger.error("Token不存在")
@@ -131,8 +130,9 @@ class MailGw:
 
         logger.error("未能获取到验证码")
         return None
+
 def get_cookie(driver):
-    try:    
+    try:
         cookie_script = """
         var cookieString = document.cookie;
         var essentialCookies = cookieString.split('; ').filter(function(cookie) {
@@ -190,7 +190,6 @@ def wait_for_email_input(driver, timeout=10):
         time.sleep(0.5)
     logger.error("邮箱输入框未出现")
     return False
-
 def input_email(driver, email):
     """使用JavaScript输入邮箱"""
     try:
@@ -292,4 +291,101 @@ def registration_process(driver, wait):
 
     except Exception as e:
         logger.error(f"注册过程错误: {str(e)}")
-        return False
+        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=new')
+            chrome_options.add_argument('--no-sandbox')
+            chrome_options.add_argument('--disable-dev-shm-usage')
+            chrome_options.add_argument('--disable-gpu')
+            chrome_options.add_argument('--disable-software-rasterizer')
+            chrome_options.add_argument('--disable-extensions')
+            chrome_options.add_argument('--disable-setuid-sandbox')
+            chrome_options.add_argument('--disable-web-security')
+            chrome_options.add_argument('--no-first-run')
+            chrome_options.add_argument('--no-default-browser-check')
+            chrome_options.add_argument('--allow-running-insecure-content')
+            chrome_options.add_argument('--disable-blink-features=AutomationControlled')
+            chrome_options.add_argument('--window-size=1920,1080')
+            chrome_options.add_argument('--user-data-dir=/tmp/chrome-data')
+            chrome_options.add_argument('--disable-features=VizDisplayCompositor')
+            chrome_options.add_argument('--disable-dev-tools')
+            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)}")
+            if hasattr(e, 'msg'):
+                logger.error(f"错误消息: {e.msg}")
+        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("/")
+async def read_root():
+    return {"status": "running"}
+
+@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:
+        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 {"error": str(e)}
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)