From b6fc114a14b1856c1f6b014debf9090c01d83cc3 Mon Sep 17 00:00:00 2001
From: BalconyJH <balconyjh@gmail.com>
Date: Mon, 12 Aug 2024 15:38:37 +0800
Subject: [PATCH] =?UTF-8?q?:children=5Fcrossing:=20=E6=9B=B4=E6=96=B0=20Ch?=
 =?UTF-8?q?romium=20=E5=AE=89=E8=A3=85=E9=80=BB=E8=BE=91,=20=E6=B7=BB?=
 =?UTF-8?q?=E5=8A=A0=20Playwright=20=E4=BE=9D=E8=B5=96=E6=A3=80=E6=9F=A5?=
 =?UTF-8?q?=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 zhenxun/utils/browser.py | 59 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 52 insertions(+), 7 deletions(-)

diff --git a/zhenxun/utils/browser.py b/zhenxun/utils/browser.py
index c7d89727..432c28ed 100644
--- a/zhenxun/utils/browser.py
+++ b/zhenxun/utils/browser.py
@@ -1,5 +1,10 @@
+import os
+import sys
+
 from nonebot import get_driver
+from playwright.__main__ import main
 from playwright.async_api import Browser, Playwright, async_playwright
+from zhenxun.configs.config import SYSTEM_PROXY
 
 from zhenxun.services.log import logger
 
@@ -13,6 +18,8 @@ _browser: Browser | None = None
 async def start_browser():
     global _playwright
     global _browser
+    install()
+    await check_playwright_env()
     _playwright = await async_playwright().start()
     _browser = await _playwright.chromium.launch()
 
@@ -33,13 +40,51 @@ def get_browser() -> Browser:
 
 def install():
     """自动安装、更新 Chromium"""
-    logger.info("正在检查 Chromium 更新")
-    import sys
 
-    from playwright.__main__ import main
+    def set_env_variables():
+        os.environ[
+            "PLAYWRIGHT_DOWNLOAD_HOST"] = "https://npmmirror.com/mirrors/playwright/"
+        if SYSTEM_PROXY:
+            os.environ["HTTPS_PROXY"] = SYSTEM_PROXY
 
-    sys.argv = ["", "install", "chromium"]
+    def restore_env_variables():
+        os.environ.pop("PLAYWRIGHT_DOWNLOAD_HOST", None)
+        if SYSTEM_PROXY:
+            os.environ.pop("HTTPS_PROXY", None)
+        if original_proxy is not None:
+            os.environ["HTTPS_PROXY"] = original_proxy
+
+    def try_install_chromium():
+        try:
+            sys.argv = ["", "install", "chromium"]
+            main()
+        except SystemExit as e:
+            return e.code == 0
+        return False
+
+    logger.info("检查 Chromium 更新")
+
+    original_proxy = os.environ.get("HTTPS_PROXY")
+    set_env_variables()
+
+    success = try_install_chromium()
+
+    if not success:
+        logger.info("Chromium 更新失败，尝试从原始仓库下载，速度较慢")
+        os.environ["PLAYWRIGHT_DOWNLOAD_HOST"] = ""
+        success = try_install_chromium()
+
+    restore_env_variables()
+
+    if not success:
+        raise RuntimeError("未知错误，Chromium 下载失败")
+
+
+async def check_playwright_env():
+    """检查 Playwright 依赖"""
+    logger.info("检查 Playwright 依赖")
     try:
-        main()
-    except SystemExit:
-        pass
+        async with async_playwright() as p:
+            await p.chromium.launch()
+    except Exception as e:
+        raise ImportError("加载失败，Playwright 依赖不全，") from e