zsglpt-pc/utils/dependency_installer.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
依赖检测和安装模块
检测 wkhtmltoimage 和 Playwright Chromium 是否已安装
如果缺失则从直链下载安装
"""

import os
import shutil
import subprocess
import tempfile
import zipfile
from pathlib import Path
from typing import Optional, Callable, Dict, Tuple

# 直链下载地址
DOWNLOAD_URLS = {
    "wkhtmltopdf": "http://by.haory.cn/3/1223/wkhtmltox.exe",
    "chromium": "http://by.haory.cn/3/1223/chromium-win64.zip",
}

# Playwright chromium 版本目录名（与 playwright 1.57.0 对应）
CHROMIUM_VERSION_DIR = "chromium-1200"


def check_wkhtmltoimage() -> bool:
    """检测 wkhtmltoimage 是否已安装"""
    # 先检查 PATH
    if shutil.which("wkhtmltoimage"):
        return True

    # 检查默认安装路径
    default_paths = [
        r"C:\Program Files\wkhtmltopdf\bin\wkhtmltoimage.exe",
        r"C:\Program Files (x86)\wkhtmltopdf\bin\wkhtmltoimage.exe",
    ]
    for p in default_paths:
        if os.path.exists(p):
            return True

    return False


def check_playwright_chromium() -> bool:
    """检测 Playwright Chromium 浏览器是否已安装"""
    local_appdata = os.environ.get("LOCALAPPDATA", "")
    if not local_appdata:
        return False

    chromium_path = Path(local_appdata) / "ms-playwright" / CHROMIUM_VERSION_DIR

    # 检查目录是否存在且包含chrome.exe
    if chromium_path.exists():
        chrome_exe = chromium_path / "chrome-win" / "chrome.exe"
        if chrome_exe.exists():
            return True

    return False


def get_missing_dependencies() -> Dict[str, bool]:
    """获取缺失的依赖列表"""
    return {
        "wkhtmltoimage": not check_wkhtmltoimage(),
        "chromium": not check_playwright_chromium(),
    }


def download_file(url: str, dest_path: str, progress_callback: Optional[Callable] = None) -> bool:
    """
    下载文件

    Args:
        url: 下载地址
        dest_path: 保存路径
        progress_callback: 进度回调 (downloaded_bytes, total_bytes)

    Returns:
        是否成功
    """
    import urllib.request

    try:
        def report_progress(block_num, block_size, total_size):
            if progress_callback:
                downloaded = block_num * block_size
                progress_callback(downloaded, total_size)

        urllib.request.urlretrieve(url, dest_path, reporthook=report_progress)
        return True
    except Exception as e:
        print(f"下载失败: {e}")
        return False


def install_wkhtmltopdf(installer_path: str, log_callback: Optional[Callable] = None) -> bool:
    """
    安装 wkhtmltopdf

    Args:
        installer_path: 安装包路径
        log_callback: 日志回调

    Returns:
        是否成功
    """
    def log(msg):
        if log_callback:
            log_callback(msg)

    try:
        log("正在启动 wkhtmltopdf 安装程序...")
        # 启动安装程序（需要用户手动完成安装向导）
        subprocess.Popen([installer_path], shell=True)
        log("请在安装向导中完成安装（建议保持默认路径）")
        return True
    except Exception as e:
        log(f"启动安装程序失败: {e}")
        return False


def install_playwright_chromium(zip_path: str, log_callback: Optional[Callable] = None) -> bool:
    """
    安装 Playwright Chromium 浏览器

    Args:
        zip_path: zip包路径
        log_callback: 日志回调

    Returns:
        是否成功
    """
    def log(msg):
        if log_callback:
            log_callback(msg)

    try:
        local_appdata = os.environ.get("LOCALAPPDATA", "")
        if not local_appdata:
            log("无法获取 LOCALAPPDATA 路径")
            return False

        # 目标目录
        playwright_dir = Path(local_appdata) / "ms-playwright"
        chromium_dir = playwright_dir / CHROMIUM_VERSION_DIR

        # 创建目录
        playwright_dir.mkdir(parents=True, exist_ok=True)

        log("正在解压 Chromium 浏览器...")

        # 解压
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(chromium_dir)

        # 创建标记文件（Playwright需要这些文件来确认安装完成）
        (chromium_dir / "INSTALLATION_COMPLETE").touch()
        (chromium_dir / "DEPENDENCIES_VALIDATED").touch()

        log("Chromium 浏览器安装成功")
        return True
    except Exception as e:
        log(f"安装 Chromium 失败: {e}")
        return False


def download_and_install_dependency(
    dep_name: str,
    progress_callback: Optional[Callable] = None,
    log_callback: Optional[Callable] = None
) -> Tuple[bool, str]:
    """
    下载并安装指定依赖

    Args:
        dep_name: 依赖名称 ("wkhtmltoimage" 或 "chromium")
        progress_callback: 下载进度回调
        log_callback: 日志回调

    Returns:
        (是否成功, 错误信息)
    """
    def log(msg):
        if log_callback:
            log_callback(msg)

    # 获取下载地址
    if dep_name == "wkhtmltoimage":
        url = DOWNLOAD_URLS["wkhtmltopdf"]
        filename = "wkhtmltox_installer.exe"
    elif dep_name == "chromium":
        url = DOWNLOAD_URLS["chromium"]
        filename = "chromium-win64.zip"
    else:
        return False, f"未知依赖: {dep_name}"

    # 下载到临时目录
    temp_dir = tempfile.gettempdir()
    download_path = os.path.join(temp_dir, filename)

    log(f"正在下载 {dep_name}...")

    if not download_file(url, download_path, progress_callback):
        return False, "下载失败"

    log("下载完成，正在安装...")

    # 安装
    if dep_name == "wkhtmltoimage":
        success = install_wkhtmltopdf(download_path, log_callback)
        if success:
            return True, "安装程序已启动，请完成安装向导"
        else:
            return False, "启动安装程序失败"
    elif dep_name == "chromium":
        success = install_playwright_chromium(download_path, log_callback)
        # 清理下载的zip文件
        try:
            os.remove(download_path)
        except:
            pass
        if success:
            return True, ""
        else:
            return False, "安装失败"

    return False, "未知错误"