ArticleReplaceBatch/setup_driver.py

"""
下载并准备 chromedriver
自动根据系统架构下载对应版本
"""
import os
import sys
import requests
import zipfile
import platform
from pathlib import Path
import shutil


def get_chrome_version():
    """获取本地 Chrome 版本"""
    import subprocess
    
    try:
        result = subprocess.run(
            ['reg', 'query', 'HKEY_CURRENT_USER\\Software\\Google\\Chrome\\BLBeacon', '/v', 'version'],
            capture_output=True, text=True, creationflags=subprocess.CREATE_NO_WINDOW
        )
        output = result.stdout.strip()
        parts = output.split()
        if len(parts) >= 3:
            return parts[-1]
    except Exception:
        pass
    
    return None


def get_latest_driver_version(chrome_version):
    """根据 Chrome 版本获取对应的 chromedriver 版本"""
    main_version = chrome_version.split('.')[0] if chrome_version else ""
    
    if not main_version:
        print("无法获取主版本号，使用默认版本")
        return "145.0.7632.117"
    
    url = f"https://chromedriver.storage.googleapis.com/LATEST_RELEASE_{main_version}"
    try:
        response = requests.get(url, timeout=10)
        if response.status_code == 200:
            return response.text.strip()
    except Exception:
        pass
    
    return "145.0.7632.117"


def download_driver(driver_version, is_64bit, drivers_dir):
    """下载 chromedriver"""
    if is_64bit:
        archive_name = "chromedriver-win64"
        download_url = f"https://storage.googleapis.com/chrome-for-testing-public/{driver_version}/win64/{archive_name}.zip"
    else:
        archive_name = "chromedriver-win32"
        download_url = f"https://storage.googleapis.com/chrome-for-testing-public/{driver_version}/win32/{archive_name}.zip"
    
    print(f"下载链接: {download_url}")
    
    zip_path = drivers_dir / "chromedriver.zip"
    
    response = requests.get(download_url, timeout=60, stream=True)
    response.raise_for_status()
    
    with open(zip_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
    
    return zip_path, archive_name


def extract_driver(zip_path, archive_name, drivers_dir, is_64bit):
    """解压驱动文件"""
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(drivers_dir)
    
    extracted_dir = drivers_dir / archive_name
    if not extracted_dir.exists():
        for item in drivers_dir.iterdir():
            if item.is_dir() and item.name.startswith("chromedriver-"):
                extracted_dir = item
                break
    
    driver_file = extracted_dir / "chromedriver.exe"
    if not driver_file.exists():
        raise FileNotFoundError(f"未找到 chromedriver.exe 在 {extracted_dir}")
    
    if is_64bit:
        target_name = "chromedriver64.exe"
    else:
        target_name = "chromedriver32.exe"
    
    driver_file.rename(drivers_dir / target_name)
    
    shutil.rmtree(extracted_dir)
    zip_path.unlink()
    
    return drivers_dir / target_name


def main():
    """主函数"""
    print("=" * 60)
    print("ChromeDriver 自动下载工具")
    print("=" * 60)
    
    drivers_dir = Path("drivers")
    drivers_dir.mkdir(exist_ok=True)
    
    is_64bit = platform.machine().endswith('64')
    driver_name = "chromedriver64.exe" if is_64bit else "chromedriver32.exe"
    existing_driver = drivers_dir / driver_name
    
    if existing_driver.exists():
        print(f"检测到已有驱动: {existing_driver}")
        while True:
            choice = input("是否重新下载? (y/n): ").strip().lower()
            if choice == 'n':
                print("使用现有驱动")
                return 0
            elif choice == 'y':
                break
            else:
                print("请输入 y 或 n")
    
    print("\n正在检测 Chrome 版本...")
    chrome_version = get_chrome_version()
    
    if chrome_version:
        print(f"检测到 Chrome 版本: {chrome_version}")
    else:
        print("未检测到 Chrome 版本，使用默认版本")
        chrome_version = "123.0.0.0"
    
    driver_version = get_latest_driver_version(chrome_version)
    print(f"匹配的 chromedriver 版本: {driver_version}")
    
    print("\n正在下载驱动...")
    zip_path, archive_name = download_driver(driver_version, is_64bit, drivers_dir)
    
    print("正在解压...")
    driver_path = extract_driver(zip_path, archive_name, drivers_dir, is_64bit)
    
    print(f"\n[OK] 驱动已准备完成: {driver_path}")
    print("=" * 60)
    
    return 0


if __name__ == "__main__":
    sys.exit(main())