replace screenshot pipeline and update admin

services/screenshots.py

@@ -3,15 +3,16 @@
 from __future__ import annotations
 
 import os
+import shutil
+import subprocess
 import time
 
 import database
 import email_service
+from api_browser import APIBrowser, get_cookie_jar_path, is_cookie_jar_fresh
 from app_config import get_config
 from app_logger import get_logger
 from browser_pool_worker import get_browser_worker_pool
-from playwright_automation import PlaywrightAutomation
-from services.browser_manager import get_browser_manager
 from services.client_log import log_to_client
 from services.runtime import get_socketio
 from services.state import safe_get_account, safe_remove_task_status, safe_update_task_status
@@ -24,6 +25,93 @@ config = get_config()
 SCREENSHOTS_DIR = config.SCREENSHOTS_DIR
 os.makedirs(SCREENSHOTS_DIR, exist_ok=True)
 
+_WKHTMLTOIMAGE_TIMEOUT_SECONDS = int(os.environ.get("WKHTMLTOIMAGE_TIMEOUT_SECONDS", "60"))
+_WKHTMLTOIMAGE_JS_DELAY_MS = int(os.environ.get("WKHTMLTOIMAGE_JS_DELAY_MS", "3000"))
+_WKHTMLTOIMAGE_WIDTH = int(os.environ.get("WKHTMLTOIMAGE_WIDTH", "1920"))
+_WKHTMLTOIMAGE_QUALITY = int(os.environ.get("WKHTMLTOIMAGE_QUALITY", "95"))
+_WKHTMLTOIMAGE_UA = os.environ.get(
+    "WKHTMLTOIMAGE_USER_AGENT",
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
+)
+
+
+def _resolve_wkhtmltoimage_path() -> str | None:
+    return os.environ.get("WKHTMLTOIMAGE_PATH") or shutil.which("wkhtmltoimage")
+
+
+def _ensure_login_cookies(account, proxy_config, log_callback) -> bool:
+    """确保有可用的登录 cookies（通过 API 登录刷新）"""
+    try:
+        with APIBrowser(log_callback=log_callback, proxy_config=proxy_config) as api_browser:
+            if not api_browser.login(account.username, account.password):
+                return False
+            return api_browser.save_cookies_for_screenshot(account.username)
+    except Exception:
+        return False
+
+
+def take_screenshot_wkhtmltoimage(
+    url: str,
+    output_path: str,
+    cookies_path: str | None = None,
+    proxy_server: str | None = None,
+    log_callback=None,
+) -> bool:
+    wkhtmltoimage_path = _resolve_wkhtmltoimage_path()
+    if not wkhtmltoimage_path:
+        if log_callback:
+            log_callback("wkhtmltoimage 未安装或不在 PATH 中")
+        return False
+
+    ext = os.path.splitext(output_path)[1].lower()
+    image_format = "jpg" if ext in (".jpg", ".jpeg") else "png"
+
+    cmd = [
+        wkhtmltoimage_path,
+        "--format",
+        image_format,
+        "--width",
+        str(_WKHTMLTOIMAGE_WIDTH),
+        "--disable-smart-width",
+        "--javascript-delay",
+        str(_WKHTMLTOIMAGE_JS_DELAY_MS),
+        "--load-error-handling",
+        "ignore",
+        "--enable-local-file-access",
+        "--encoding",
+        "utf-8",
+        "--user-agent",
+        _WKHTMLTOIMAGE_UA,
+    ]
+
+    if image_format in ("jpg", "jpeg"):
+        cmd.extend(["--quality", str(_WKHTMLTOIMAGE_QUALITY)])
+
+    if cookies_path:
+        cmd.extend(["--cookie-jar", cookies_path])
+
+    if proxy_server:
+        cmd.extend(["--proxy", proxy_server])
+
+    cmd.extend([url, output_path])
+
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=_WKHTMLTOIMAGE_TIMEOUT_SECONDS)
+        if result.returncode != 0:
+            if log_callback:
+                err_msg = (result.stderr or result.stdout or "").strip()
+                log_callback(f"wkhtmltoimage 截图失败: {err_msg[:200]}")
+            return False
+        return True
+    except subprocess.TimeoutExpired:
+        if log_callback:
+            log_callback("wkhtmltoimage 截图超时")
+        return False
+    except Exception as e:
+        if log_callback:
+            log_callback(f"wkhtmltoimage 截图异常: {e}")
+        return False
+
 
 def _emit(event: str, data: object, *, room: str | None = None) -> None:
     try:
@@ -42,7 +130,7 @@ def take_screenshot_for_account(
     task_start_time=None,
     browse_result=None,
 ):
-    """为账号任务完成后截图（使用工作线程池，真正的浏览器复用）"""
+    """为账号任务完成后截图（使用截图线程池并发执行）"""
     account = safe_get_account(user_id, account_id)
     if not account:
         return
@@ -63,9 +151,11 @@ def take_screenshot_for_account(
             _emit("account_update", acc.to_dict(), room=f"user_{user_id}")
 
         max_retries = 3
+        proxy_config = account.proxy_config if hasattr(account, "proxy_config") else None
+        proxy_server = proxy_config.get("server") if proxy_config else None
+        cookie_path = get_cookie_jar_path(account.username)
 
         for attempt in range(1, max_retries + 1):
-            automation = None
             try:
                 safe_update_task_status(
                     account_id,
@@ -75,100 +165,39 @@ def take_screenshot_for_account(
                 if attempt > 1:
                     log_to_client(f"🔄 第 {attempt} 次截图尝试...", user_id, account_id)
 
+                worker_id = browser_instance.get("worker_id", "?") if isinstance(browser_instance, dict) else "?"
+                use_count = browser_instance.get("use_count", 0) if isinstance(browser_instance, dict) else 0
                 log_to_client(
-                    f"使用Worker-{browser_instance['worker_id']}的浏览器（已使用{browser_instance['use_count']}次）",
+                    f"使用Worker-{worker_id}执行截图（已执行{use_count}次）",
                     user_id,
                     account_id,
                 )
 
-                proxy_config = account.proxy_config if hasattr(account, "proxy_config") else None
-                automation = PlaywrightAutomation(get_browser_manager(), account_id, proxy_config=proxy_config)
-                automation.playwright = browser_instance["playwright"]
-                automation.browser = browser_instance["browser"]
-
                 def custom_log(message: str):
                     log_to_client(message, user_id, account_id)
 
-                automation.log = custom_log
-
-                log_to_client("登录中...", user_id, account_id)
-                login_result = automation.quick_login(account.username, account.password, account.remember)
-                if not login_result["success"]:
-                    error_message = login_result.get("message", "截图登录失败")
-                    log_to_client(f"截图登录失败: {error_message}", user_id, account_id)
-                    if attempt < max_retries:
-                        log_to_client("将重试...", user_id, account_id)
-                        time.sleep(2)
-                        continue
-                    log_to_client("❌ 截图失败: 登录失败", user_id, account_id)
-                    return {"success": False, "error": "登录失败"}
+                if not is_cookie_jar_fresh(cookie_path) or attempt > 1:
+                    log_to_client("正在刷新登录态...", user_id, account_id)
+                    if not _ensure_login_cookies(account, proxy_config, custom_log):
+                        log_to_client("截图登录失败", user_id, account_id)
+                        if attempt < max_retries:
+                            log_to_client("将重试...", user_id, account_id)
+                            time.sleep(2)
+                            continue
+                        log_to_client("❌ 截图失败: 登录失败", user_id, account_id)
+                        return {"success": False, "error": "登录失败"}
 
                 log_to_client(f"导航到 '{browse_type}' 页面...", user_id, account_id)
 
-                # 截图场景：优先用 bz 参数直达页面（更稳定，避免页面按钮点击失败导致截图跑偏）
-                navigated = False
-                try:
-                    from urllib.parse import urlsplit
+                from urllib.parse import urlsplit
 
-                    parsed = urlsplit(config.ZSGL_LOGIN_URL)
-                    base = f"{parsed.scheme}://{parsed.netloc}"
-                    if "注册前" in str(browse_type):
-                        bz = 0
-                    else:
-                        bz = 2  # 应读
-                    target_url = f"{base}/admin/center.aspx?bz={bz}"
-                    # 目标：保留外层框架（左侧菜单/顶部栏），仅在 mainframe 内部导航到目标内容页
-                    iframe = None
-                    try:
-                        iframe = automation.get_iframe_safe(retry=True, max_retries=5)
-                    except Exception:
-                        iframe = None
-
-                    if iframe:
-                        iframe.goto(target_url, timeout=60000)
-                        current_url = getattr(iframe, "url", "") or ""
-                        if "center.aspx" not in current_url:
-                            raise RuntimeError(f"unexpected_iframe_url:{current_url}")
-                        try:
-                            iframe.wait_for_load_state("networkidle", timeout=10000)
-                        except Exception:
-                            pass
-                        try:
-                            iframe.wait_for_selector("table.ltable", timeout=5000)
-                        except Exception:
-                            pass
-                    else:
-                        # 兜底：若获取不到 iframe，则退回到主页面直达
-                        automation.main_page.goto(target_url, timeout=60000)
-                        current_url = getattr(automation.main_page, "url", "") or ""
-                        if "center.aspx" not in current_url:
-                            raise RuntimeError(f"unexpected_url:{current_url}")
-                        try:
-                            automation.main_page.wait_for_load_state("networkidle", timeout=10000)
-                        except Exception:
-                            pass
-                        try:
-                            automation.main_page.wait_for_selector("table.ltable", timeout=5000)
-                        except Exception:
-                            pass
-                    navigated = True
-                except Exception as nav_error:
-                    log_to_client(f"直达页面失败，将尝试按钮切换: {str(nav_error)[:120]}", user_id, account_id)
-
-                # 兼容兜底：若直达失败，则回退到原有按钮切换方式
-                if not navigated:
-                    result = automation.browse_content(
-                        navigate_only=True,
-                        browse_type=browse_type,
-                        auto_next_page=False,
-                        auto_view_attachments=False,
-                        interval=0,
-                        should_stop_callback=None,
-                    )
-                    if not result.success and result.error_message:
-                        log_to_client(f"导航警告: {result.error_message}", user_id, account_id)
-
-                time.sleep(2)
+                parsed = urlsplit(config.ZSGL_LOGIN_URL)
+                base = f"{parsed.scheme}://{parsed.netloc}"
+                if "注册前" in str(browse_type):
+                    bz = 0
+                else:
+                    bz = 2  # 应读
+                target_url = f"{base}/admin/center.aspx?bz={bz}"
 
                 timestamp = get_beijing_now().strftime("%Y%m%d_%H%M%S")
 
@@ -178,7 +207,13 @@ def take_screenshot_for_account(
                 screenshot_filename = f"{username_prefix}_{login_account}_{browse_type}_{timestamp}.jpg"
                 screenshot_path = os.path.join(SCREENSHOTS_DIR, screenshot_filename)
 
-                if automation.take_screenshot(screenshot_path):
+                if take_screenshot_wkhtmltoimage(
+                    target_url,
+                    screenshot_path,
+                    cookies_path=cookie_path if is_cookie_jar_fresh(cookie_path) else None,
+                    proxy_server=proxy_server,
+                    log_callback=custom_log,
+                ):
                     if os.path.exists(screenshot_path) and os.path.getsize(screenshot_path) > 1000:
                         log_to_client(f"✓ 截图成功: {screenshot_filename}", user_id, account_id)
                         return {"success": True, "filename": screenshot_filename}
@@ -197,15 +232,6 @@ def take_screenshot_for_account(
                 if attempt < max_retries:
                     log_to_client("将重试...", user_id, account_id)
                     time.sleep(2)
-            finally:
-                if automation:
-                    try:
-                        if automation.context:
-                            automation.context.close()
-                        automation.context = None
-                        automation.page = None
-                    except Exception as e:
-                        logger.debug(f"关闭context时出错: {e}")
 
         return {"success": False, "error": "截图失败，已重试3次"}