zsglpt/fix_screenshot.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import re

NEW_FUNC = '''def take_screenshot_for_account(user_id, account_id, browse_type="应读", source="manual", task_start_time=None, browse_result=None):
    """为账号任务完成后截图(使用独立进程，避免eventlet冲突)"""
    import subprocess
    import json
    import uuid

    if user_id not in user_accounts or account_id not in user_accounts[user_id]:
        return

    account = user_accounts[user_id][account_id]

    # 使用截图信号量控制并发
    semaphore, max_concurrent = get_screenshot_semaphore()
    screenshot_acquired = semaphore.acquire(blocking=True, timeout=300)

    if not screenshot_acquired:
        log_to_client(f"截图资源获取超时,跳过截图", user_id, account_id)
        if account_id in task_status:
            del task_status[account_id]
        return

    max_retries = 2
    screenshot_success = False

    for attempt in range(1, max_retries + 1):
        try:
            if account_id in task_status:
                task_status[account_id]["detail_status"] = f"正在截图{f' (第{attempt}次)' if attempt > 1 else ''}"

            if attempt > 1:
                log_to_client(f"🔄 第 {attempt} 次截图尝试...", user_id, account_id)

            beijing_tz = pytz.timezone('Asia/Shanghai')
            now_beijing = datetime.now(beijing_tz)
            timestamp = now_beijing.strftime('%Y%m%d_%H%M%S')

            user_info = database.get_user_by_id(user_id)
            username_prefix = user_info['username'] if user_info else f"user{user_id}"
            login_account = account.username
            actual_browse_type = account.last_browse_type or browse_type
            screenshot_filename = f"{username_prefix}_{login_account}_{actual_browse_type}_{timestamp}.jpg"
            screenshot_path = os.path.join(SCREENSHOTS_DIR, screenshot_filename)

            config = {
                'username': account.username,
                'password': account.password,
                'remember': account.remember if hasattr(account, 'remember') else '',
                'browse_type': actual_browse_type,
                'screenshot_path': screenshot_path,
                'proxy_config': account.proxy_config if hasattr(account, 'proxy_config') else None
            }

            config_file = f"/tmp/screenshot_config_{uuid.uuid4().hex}.json"
            with open(config_file, 'w', encoding='utf-8') as f:
                json.dump(config, f, ensure_ascii=False)

            log_to_client(f"启动截图进程...", user_id, account_id)

            try:
                result = subprocess.run(
                    ['python3', '/www/wwwroot/zsglpt/screenshot_worker.py', config_file],
                    capture_output=True,
                    text=True,
                    timeout=120
                )

                output = result.stdout
                if "===RESULT===" in output:
                    result_json = output.split("===RESULT===")[1].strip()
                    result_data = json.loads(result_json)

                    if result_data.get('success'):
                        log_to_client(f"✓ 截图成功: {screenshot_filename}", user_id, account_id)
                        screenshot_success = True
                        break
                    else:
                        log_to_client(f"截图失败: {result_data.get('message', '未知错误')}", user_id, account_id)
                else:
                    err_msg = result.stderr[:200] if result.stderr else '无输出'
                    log_to_client(f"截图进程异常: {err_msg}", user_id, account_id)

            except subprocess.TimeoutExpired:
                log_to_client(f"截图超时", user_id, account_id)
            except Exception as e:
                log_to_client(f"截图进程错误: {str(e)}", user_id, account_id)

            try:
                if os.path.exists(config_file):
                    os.remove(config_file)
            except:
                pass

            if attempt < max_retries:
                log_to_client(f"将重试...", user_id, account_id)
                time.sleep(2)

        except Exception as e:
            log_to_client(f"截图出错: {str(e)}", user_id, account_id)
            if attempt < max_retries:
                log_to_client(f"将重试...", user_id, account_id)
                time.sleep(2)

    if not screenshot_success:
        log_to_client(f"❌ 截图失败: 已重试{max_retries}次", user_id, account_id)

    account.status = "未开始"
    socketio.emit('account_update', account.to_dict(), room=f'user_{user_id}')

    semaphore.release()
    if account_id in task_status:
        del task_status[account_id]

    if task_start_time and browse_result:
        import time as time_module
        total_duration = int(time_module.time() - task_start_time)
        database.create_task_log(
            user_id=user_id,
            account_id=account_id,
            username=account.username if account else "",
            browse_type=actual_browse_type,
            status='success',
            total_items=browse_result.get('total_items', 0),
            total_attachments=browse_result.get('total_attachments', 0),
            error_message='' if screenshot_success else '截图失败',
            duration=total_duration,
            source=source
        )
        log_to_client(f"任务完成！总耗时 {total_duration} 秒", user_id, account_id)

'''

def main():
    with open('/www/wwwroot/zsglpt/app.py', 'r', encoding='utf-8') as f:
        content = f.read()
    
    lines = content.split('\n')
    start_idx = None
    end_idx = None
    
    for i, line in enumerate(lines):
        if line.startswith('def take_screenshot_for_account('):
            start_idx = i
        elif start_idx is not None and line.startswith('def ') and i > start_idx:
            end_idx = i
            break
    
    if start_idx is None:
        print("未找到函数")
        return
    
    if end_idx is None:
        end_idx = len(lines)
    
    new_lines = lines[:start_idx] + NEW_FUNC.split('\n') + [''] + lines[end_idx:]
    
    with open('/www/wwwroot/zsglpt/app.py', 'w', encoding='utf-8') as f:
        f.write('\n'.join(new_lines))
    
    print(f"已替换函数 (行 {start_idx+1} 到 {end_idx})")

if __name__ == '__main__':
    main()