feat: handle popup dialog and save screenshot on successful login

browser/navigation.py

@@ -7,23 +7,45 @@ from utils.common import ensure_dir
 class KeepAliveError(Exception):
     pass
 
-def handle_untrusted_dialog(page: Page, logger=None):
+def handle_popup_dialog(page: Page, logger=None):
     """
-    检查并处理 "Last modified by..." 的弹窗。
-    如果弹窗出现，则点击 "OK" 按钮。
+    检查并处理弹窗。
+    交替点击 Got it 和 Continue to the app 按钮直到没有弹窗。
     """
-    ok_button_locator = page.get_by_role("button", name="OK")
-
+    logger.info("开始处理弹窗...")
+    
+    # 定义需要查找的按钮列表
+    button_names = ["Got it", "Continue to the app"]
+    max_iterations = 10  # 最多尝试10轮，防止死循环
+    total_clicks = 0
+    
     try:
-        if ok_button_locator.is_visible(timeout=10000): # 等待最多10秒
-            logger.info(f"检测到弹窗，正在点击 'OK' 按钮...")
+        for iteration in range(max_iterations):
+            clicked_in_round = False
+            
+            # 等待页面稳定
+            time.sleep(1)
+            
+            # 每轮交替尝试点击所有按钮
+            for button_name in button_names:
+                try:
+                    button_locator = page.locator(f'button:visible:has-text("{button_name}")')
+                    if button_locator.count() > 0 and button_locator.first.is_visible():
+                        # logger.info(f"检测到弹窗: 点击 '{button_name}'")
+                        button_locator.first.click(force=True, timeout=2000)
+                        total_clicks += 1
+                        clicked_in_round = True
+                        time.sleep(1)
+                except:
+                    pass
             
-            ok_button_locator.click(force=True)
-            logger.info(f"'OK' 按钮已点击")
-            expect(ok_button_locator).to_be_hidden(timeout=1000)
-            logger.info(f"弹窗已确认关闭")
+            if not clicked_in_round:
+                break
+        
+        if total_clicks > 0:
+            logger.info(f"弹窗处理完成, 共点击 {total_clicks} 次")
         else:
-            logger.info(f"在10秒内未检测到弹窗，继续执行...")
+            logger.info("未检测到弹窗")
     except Exception as e:
         logger.info(f"检查弹窗时发生意外：{e}，将继续执行...")
 
@@ -34,8 +56,20 @@ def handle_successful_navigation(page: Page, logger, cookie_file_config, shutdow
     logger.info("已成功到达目标页面")
     page.click('body') # 给予页面焦点
 
-    # 检查并处理 "Last modified by..." 的弹窗
-    handle_untrusted_dialog(page, logger=logger)
+    # 检查并处理弹窗
+    handle_popup_dialog(page, logger=logger)
+
+    # 保存登录成功截图
+    try:
+        from datetime import datetime
+        screenshot_dir = logs_dir()
+        ensure_dir(screenshot_dir)
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        screenshot_path = os.path.join(screenshot_dir, f"SUCCESS_{cookie_file_config}_{timestamp}.png")
+        page.screenshot(path=screenshot_path)
+        logger.info(f"已保存登录成功截图: {screenshot_path}")
+    except Exception as e:
+        logger.warning(f"保存截图失败: {e}")
 
     if cookie_validator:
         logger.info("Cookie验证器已创建，将定期验证Cookie有效性")

main.py

@@ -5,6 +5,21 @@ import signal
 import sys
 import time
 
+# 加载 .env 文件（仅在非 Docker 环境且文件存在时）
+def load_env_file():
+    """加载 .env 文件，不影响已存在的环境变量"""
+    if os.environ.get("DOCKER_ENV") or os.path.exists("/.dockerenv"):
+        return  # Docker 环境，跳过加载
+    try:
+        from dotenv import load_dotenv
+        env_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env")
+        if os.path.exists(env_path):
+            load_dotenv(env_path, override=False)  # override=False 不覆盖已有环境变量
+    except ImportError:
+        pass  # python-dotenv 未安装，跳过
+
+load_env_file()
+
 from browser.instance import run_browser_instance
 from utils.logger import setup_logging
 from utils.paths import cookies_dir, logs_dir

requirements.txt

@@ -30,4 +30,5 @@ typing_extensions==4.14.0
 ua-parser==1.0.1
 ua-parser-builtins==0.18.0.post1
 urllib3==2.4.0
-yarl==1.20.1
+yarl==1.20.1
+python-dotenv>=1.0.0