zsglpt/temp_fix_screenshot.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
临时修复截图问题的脚本
提供三个选项：安装wkhtmltoimage、修改为Playwright、或临时禁用截图
"""

import os
import sys
import subprocess


def check_wkhtmltoimage():
    """检查wkhtmltoimage是否已安装"""
    try:
        result = subprocess.run(["wkhtmltoimage", "--version"], capture_output=True, text=True, timeout=5)
        return result.returncode == 0
    except:
        return False


def check_playwright():
    """检查Playwright是否已安装"""
    try:
        from playwright.sync_api import sync_playwright

        return True
    except ImportError:
        return False


def option1_install_wkhtmltoimage():
    """选项1: 指导安装wkhtmltoimage"""
    print("\n" + "=" * 60)
    print("选项 1: 安装 wkhtmltoimage (推荐)")
    print("=" * 60)

    if check_wkhtmltoimage():
        print("✓ wkhtmltoimage 已经安装")
        return True

    print("wkhtmltoimage 未安装，需要手动安装")
    print("\n安装步骤:")
    print("1. 访问: https://wkhtmltopdf.org/downloads.html")
    print("2. 下载Windows版本 (.msi)")
    print("3. 运行安装程序")
    print("4. 将安装路径添加到系统PATH")
    print("5. 重启命令行验证: wkhtmltoimage --version")

    return False


def option2_modify_to_playwright():
    """选项2: 修改为使用Playwright"""
    print("\n" + "=" * 60)
    print("选项 2: 修改为使用 Playwright")
    print("=" * 60)

    if not check_playwright():
        print("❌ Playwright 未安装")
        return False

    print("✓ Playwright 已安装")
    print("正在修改截图实现为Playwright...")

    # 备份原文件
    original_file = "services/screenshots.py"
    backup_file = "services/screenshots.py.wkhtmltoimage.backup"

    try:
        # 读取原文件
        with open(original_file, "r", encoding="utf-8") as f:
            content = f.read()

        # 创建备份
        with open(backup_file, "w", encoding="utf-8") as f:
            f.write(content)

        print(f"✓ 已备份原文件为: {backup_file}")

        # 修改实现（简化版本）
        playwright_content = '''#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
截图服务 - Playwright版本
临时替换wkhtmltoimage实现
"""

import os
from playwright.sync_api import sync_playwright

def take_screenshot_playwright(url, output_path, width=1920, height=1080):
    """使用Playwright截图"""
    try:
        with sync_playwright() as p:
            browser = p.chromium.launch(headless=True)
            page = browser.new_page()
            page.set_viewport_size({"width": width, "height": height})
            page.goto(url, timeout=30000)
            page.wait_for_timeout(3000)  # 等待页面加载
            page.screenshot(path=output_path, full_page=True)
            browser.close()
            return True
    except Exception as e:
        print(f"截图失败: {e}")
        return False

def take_screenshot_for_account(account, target_url, browse_type, user_id, account_id):
    """为账号截图"""
    screenshot_filename = f"account_{account_id}_{browse_type}.png"
    screenshot_path = os.path.join("screenshots", screenshot_filename)
    
    os.makedirs("screenshots", exist_ok=True)
    
    success = take_screenshot_playwright(target_url, screenshot_path)
    
    if success:
        return {"success": True, "screenshot_path": screenshot_path}
    else:
        return {"success": False, "error": "截图失败"}
'''

        # 写入新实现
        with open(original_file, "w", encoding="utf-8") as f:
            f.write(playwright_content)

        print("✓ 已修改为Playwright实现")
        print("✓ 重启应用后生效")
        return True

    except Exception as e:
        print(f"❌ 修改失败: {e}")
        return False


def option3_disable_screenshot():
    """选项3: 临时禁用截图"""
    print("\n" + "=" * 60)
    print("选项 3: 临时禁用截图功能")
    print("=" * 60)

    # 设置环境变量禁用截图
    os.environ["ENABLE_SCREENSHOT"] = "0"
    print("✓ 已设置环境变量: ENABLE_SCREENSHOT=0")
    print("✓ 重启应用后截图功能将被跳过")

    # 检查tasks.py中是否有截图调用
    try:
        with open("services/tasks.py", "r", encoding="utf-8") as f:
            content = f.read()

        if "take_screenshot_for_account" in content:
            print("⚠️  发现tasks.py中有截图调用，建议注释掉:")
            print("   查找: take_screenshot_for_account")
            print("   临时注释: # take_screenshot_for_account(...)")

    except Exception as e:
        print(f"检查tasks.py失败: {e}")

    return True


def main():
    print("🔧 截图问题修复工具")
    print("=" * 60)

    # 检查当前状态
    print("📊 当前状态:")
    print(f"  wkhtmltoimage: {'✓ 已安装' if check_wkhtmltoimage() else '❌ 未安装'}")
    print(f"  Playwright: {'✓ 已安装' if check_playwright() else '❌ 未安装'}")

    while True:
        print("\n请选择修复方案:")
        print("1. 安装 wkhtmltoimage (推荐)")
        print("2. 修改为使用 Playwright")
        print("3. 临时禁用截图功能")
        print("4. 查看状态")
        print("5. 退出")

        choice = input("\n请输入选项 (1-5): ").strip()

        if choice == "1":
            if option1_install_wkhtmltoimage():
                print("\n🎉 wkhtmltoimage安装完成！重启应用即可。")
        elif choice == "2":
            option2_modify_to_playwright()
        elif choice == "3":
            option3_disable_screenshot()
        elif choice == "4":
            print("\n📊 当前状态:")
            print(f"  wkhtmltoimage: {'✓ 已安装' if check_wkhtmltoimage() else '❌ 未安装'}")
            print(f"  Playwright: {'✓ 已安装' if check_playwright() else '❌ 未安装'}")
        elif choice == "5":
            print("👋 再见！")
            break
        else:
            print("❌ 无效选项，请重新输入")


if __name__ == "__main__":
    main()