Update app.py

app.py

@@ -1,17 +1,15 @@
-from fastapi import FastAPI, HTTPException
+import gradio as gr
 from selenium import webdriver
-from selenium.webdriver.chrome.service import Service
 from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.common.by import By
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
+from selenium.common.exceptions import TimeoutException
 import time
 import random
 import string
-import re
-from pydantic import BaseModel
-from typing import Optional
 import logging
+import re
 
 # 配置日志
 logging.basicConfig(
@@ -20,51 +18,35 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-app = FastAPI()
-
-class EmailResponse(BaseModel):
-    email: str
-    verification_code: Optional[str] = None
-    error: Optional[str] = None
-
 def setup_driver():
+    """设置 Chrome 驱动"""
     chrome_options = Options()
     chrome_options.add_argument('--no-sandbox')
-    chrome_options.add_argument('--headless')
     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('--headless')  # 无头模式
+    chrome_options.add_argument('--disable-blink-features=AutomationControlled')
+    chrome_options.add_experimental_option('excludeSwitches', ['enable-automation'])
+    chrome_options.add_experimental_option('useAutomationExtension', False)
     
-    # 明确指定 ChromeDriver 路径
-    service = Service('/usr/local/bin/chromedriver')
-    
-    try:
-        driver = webdriver.Chrome(service=service, options=chrome_options)
-        return driver
-    except Exception as e:
-        logger.error(f"创建 Chrome driver 失败: {str(e)}")
-        raise
-
-def generate_email():
-    """生成随机邮箱地址"""
-    random_prefix = ''.join(random.choices(string.ascii_letters + string.digits, k=8))
-    return f"{random_prefix}@youxiang.dev"
+    driver = webdriver.Chrome(options=chrome_options)
+    return driver
 
 def get_verification_code(driver, email, max_retries=10, delay=2):
+    """获取验证码"""
+    email_url = f'https://linshiyou.com/#/{email}'
+    
     try:
-        email_url = f'https://linshiyou.com/#/{email}'
-        
-        # 打开新标签页并访问邮箱
-        driver.execute_script(f"window.open('{email_url}', '_blank')")
-        driver.switch_to.window(driver.window_handles[-1])
-        logger.info(f"已打开邮箱页面: {email}")
+        driver.get(email_url)
         
         for attempt in range(max_retries):
             try:
-                # 刷新页面
+                # 刷新页面获取最新邮件
                 driver.refresh()
-                time.sleep(2)
+                time.sleep(delay)
                 
-                # 获取邮件内容
+                # 尝试获取邮件内容
                 body_content = driver.find_element(By.CLASS_NAME, "body").get_attribute('outerHTML')
                 logger.info(f"获取到邮件内容: {body_content}")
                 
@@ -72,68 +54,68 @@ def get_verification_code(driver, email, max_retries=10, delay=2):
                 match = re.search(r'\b\d{6}\b', body_content)
                 if match:
                     verification_code = match.group(0)
-                    logger.info(f"成功获取验证码: {verification_code}")
-                    
-                    # 关闭邮箱标签页
-                    driver.close()
-                    driver.switch_to.window(driver.window_handles[0])
-                    
+                    logger.info(f"获取到验证码: {verification_code}")
                     return verification_code
                     
             except Exception as e:
-                logger.error(f"第 {attempt + 1} 次尝试失败: {str(e)}")
-                if attempt == max_retries - 1:
-                    raise
+                logger.error(f"第 {attempt + 1} 次获取验证码失败: {str(e)}")
                 
-            time.sleep(delay)
             logger.info(f"等待验证码，重试第 {attempt + 1} 次")
-        
-        # 如果获取失败，确保回到主标签页
-        try:
-            driver.close()
-            driver.switch_to.window(driver.window_handles[0])
-        except:
-            pass
+            time.sleep(delay)
             
-        return None
-        
     except Exception as e:
-        logger.error(f"获取验证码时发生错误: {str(e)}")
-        # 确保清理标签页
-        try:
-            driver.close()
-            driver.switch_to.window(driver.window_handles[0])
-        except:
-            pass
-        raise
-
-@app.get("/")
-def read_root():
-    return {"message": "临时邮箱服务运行中"}
+        logger.error(f"打开邮箱页面失败: {str(e)}")
+        
+    return None
 
-@app.post("/generate-email")
-def create_email():
-    email = generate_email()
-    logger.info(f"生成新邮箱: {email}")
-    return EmailResponse(email=email)
+def generate_random_email():
+    """生成随机邮箱地址"""
+    random_prefix = ''.join(random.choices(string.ascii_letters + string.digits, k=8))
+    email = f"{random_prefix}@youxiang.dev"
+    logger.info(f"生成随机邮箱: {email}")
+    return {"email": email}
 
-@app.get("/get-verification-code/{email}")
-def get_code(email: str):
-    driver = None
+def get_email_code(email: str):
+    """获取指定邮箱的验证码"""
+    logger.info(f"开始获取邮箱验证码: {email}")
     try:
-        logger.info(f"开始获取邮箱验证码: {email}")
         driver = setup_driver()
         code = get_verification_code(driver, email)
+        driver.quit()
+        
         if code:
-            return EmailResponse(email=email, verification_code=code)
-        raise HTTPException(status_code=404, detail="未找到验证码")
+            logger.info(f"成功获取验证码: {code}")
+            return {"code": code}
+        else:
+            logger.error("未能获取验证码")
+            return {"error": "未能获取验证码"}
+            
     except Exception as e:
-        logger.error(f"获取验证码失败: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"获取验证码过程出错: {str(e)}")
+        return {"error": str(e)}
     finally:
-        if driver:
+        try:
             driver.quit()
+        except:
+            pass
+
+# Gradio 接口
+demo = gr.Interface(
+    [generate_random_email, get_email_code],
+    [
+        None,  # generate_random_email 不需要输入
+        gr.Textbox(label="邮箱地址")  # get_email_code 需要邮箱输入
+    ],
+    [
+        gr.JSON(label="生成的邮箱"),  # 返回 JSON 格式的邮箱
+        gr.JSON(label="验证码结果")   # 返回 JSON 格式的验证码
+    ],
+    title="邮箱验证码服务",
+    description="提供随机邮箱生成和验证码获取服务"
+)
+
+# 启用队列处理请求
+demo.queue()
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)