import os import json import logging import time import threading from concurrent.futures import ThreadPoolExecutor, Future from typing import Dict, List, Optional, Any, Tuple from datetime import datetime from playwright.sync_api import sync_playwright, Browser, Page logger = logging.getLogger(__name__) class BrowserManager: """管理浏览器实例和配置文件.""" def __init__(self, settings=None): """ 初始化浏览器管理器. Args: settings: 应用程序设置 """ self.settings = settings or {} self.profiles_dir = self.settings.get( "profiles_dir", os.path.join(os.path.expanduser("~"), "douyin_imbot", "profiles"), ) self.browsers: Dict[str, Browser] = {} self.pages: Dict[str, Page] = {} self.profile_details: Dict[str, Dict] = {} self.playwright = None self.executor = settings.get("browser_executor") if settings else None self.running_tasks: Dict[str, Future] = {} self.browser_locks: Dict[str, threading.Lock] = {} self.load_profiles() # 如果没有提供线程池，创建一个 if not self.executor: max_workers = self.settings.get("browser_count", 3) self.executor = ThreadPoolExecutor(max_workers=max_workers) logger.info(f"已创建浏览器线程池，最大线程数: {max_workers}") def load_profiles(self): """加载配置文件.""" try: if not os.path.exists(self.profiles_dir): os.makedirs(self.profiles_dir, exist_ok=True) logger.info(f"创建配置文件目录: {self.profiles_dir}") profiles_index_path = os.path.join(self.profiles_dir, "profiles.json") if os.path.exists(profiles_index_path): with open(profiles_index_path, "r", encoding="utf-8") as f: self.profile_details = json.load(f) logger.info(f"已加载 {len(self.profile_details)} 个配置文件") else: self.profile_details = {} with open(profiles_index_path, "w", encoding="utf-8") as f: json.dump({}, f) logger.info("创建了空的配置文件索引") except Exception as e: logger.error(f"加载配置文件失败: {str(e)}") self.profile_details = {} def get_profiles(self): """获取所有配置文件.""" return list(self.profile_details.keys()) def get_profile_details(self, profile_name): """获取配置文件详情.""" return self.profile_details.get(profile_name, {}) def create_profile(self, profile_name, description=""): """创建新的配置文件.""" try: if profile_name in self.profile_details: logger.warning(f"配置文件 '{profile_name}' 已存在") return False # 创建配置目录 profile_dir = os.path.join(self.profiles_dir, profile_name) os.makedirs(profile_dir, exist_ok=True) # 记录配置信息 self.profile_details[profile_name] = { "user_data_dir": profile_dir, "created": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "last_used": "", "description": description, } # 保存配置索引 self._save_profiles_index() logger.info(f"创建了配置文件 '{profile_name}'") return True except Exception as e: logger.error(f"创建配置文件失败: {str(e)}") return False def delete_profile(self, profile_name): """删除配置文件.""" try: if profile_name not in self.profile_details: logger.warning(f"配置文件 '{profile_name}' 不存在") return False # 如果浏览器正在运行，先关闭 self._close_browser(profile_name) # 从列表中移除 del self.profile_details[profile_name] # 保存配置索引 self._save_profiles_index() logger.info(f"删除了配置文件 '{profile_name}'") return True except Exception as e: logger.error(f"删除配置文件失败: {str(e)}") return False def _save_profiles_index(self): """保存配置文件索引.""" try: profiles_index_path = os.path.join(self.profiles_dir, "profiles.json") with open(profiles_index_path, "w", encoding="utf-8") as f: json.dump(self.profile_details, f, ensure_ascii=False, indent=4) logger.debug("配置文件索引已更新") return True except Exception as e: logger.error(f"保存配置文件索引失败: {str(e)}") return False def _close_browser(self, profile_name): """关闭特定配置的浏览器.""" try: # 获取此配置的锁 if profile_name in self.browser_locks: with self.browser_locks[profile_name]: if profile_name in self.browsers: if profile_name in self.pages: # 先关闭页面 try: self.pages[profile_name].close() except Exception as e: logger.error(f"关闭页面失败: {str(e)}") del self.pages[profile_name] # 关闭浏览器 try: self.browsers[profile_name].close() except Exception as e: logger.error(f"关闭浏览器失败: {str(e)}") del self.browsers[profile_name] logger.info(f"已关闭配置 '{profile_name}' 的浏览器") # 删除锁 del self.browser_locks[profile_name] elif profile_name in self.browsers: # 没有锁但有浏览器实例，直接关闭 if profile_name in self.pages: try: self.pages[profile_name].close() except Exception: pass del self.pages[profile_name] try: self.browsers[profile_name].close() except Exception: pass del self.browsers[profile_name] logger.info(f"已关闭配置 '{profile_name}' 的浏览器（无锁）") # 取消正在运行的任务 if profile_name in self.running_tasks: self.running_tasks[profile_name].cancel() del self.running_tasks[profile_name] logger.info(f"已取消配置 '{profile_name}' 的运行任务") return True except Exception as e: logger.error(f"关闭浏览器 '{profile_name}' 时出错: {str(e)}") return False def launch_browser(self, profile_name, headless=False, on_complete=None): """ 启动特定配置的浏览器. Args: profile_name: 配置文件名称 headless: 是否以无头模式启动 on_complete: 浏览器启动完成后的回调函数 Returns: Future对象表示异步任务 """ if profile_name not in self.profile_details: logger.error(f"配置文件 '{profile_name}' 不存在") return None # 先检查该配置文件是否已有浏览器实例在运行 if profile_name in self.browsers: logger.warning(f"配置 '{profile_name}' 已有浏览器实例在运行") if on_complete: on_complete( self.browsers[profile_name], self.pages.get(profile_name), False ) return None # 为这个配置创建一个锁 if profile_name not in self.browser_locks: self.browser_locks[profile_name] = threading.Lock() # 提交任务到线程池 task = self.executor.submit( self._launch_browser_task, profile_name, headless, on_complete ) self.running_tasks[profile_name] = task logger.info(f"已提交浏览器启动任务: {profile_name}") return task def _launch_browser_task(self, profile_name, headless, on_complete): """在线程中执行的浏览器启动任务.""" logger.info(f"开始在线程中启动浏览器: {profile_name}") success = False browser = None page = None try: # 获取锁，确保同一配置文件只启动一个浏览器实例 with self.browser_locks[profile_name]: # 再次检查浏览器是否已启动（双重检查锁定模式） if profile_name in self.browsers: logger.warning( f"配置 '{profile_name}' 在获取锁后已有浏览器实例运行" ) browser = self.browsers[profile_name] page = self.pages.get(profile_name) success = True else: # 获取配置信息 profile_info = self.profile_details[profile_name] user_data_dir = profile_info.get("user_data_dir") if not user_data_dir or not os.path.exists( os.path.dirname(user_data_dir) ): logger.error( f"配置 '{profile_name}' 的用户数据目录无效: {user_data_dir}" ) return False # 启动浏览器 try: self.playwright = sync_playwright().start() logger.info(f"正在启动配置 '{profile_name}' 的浏览器...") browser = self.playwright.chromium.launch_persistent_context( user_data_dir=user_data_dir, headless=headless, args=[ "--disable-blink-features=AutomationControlled", "--no-sandbox", "--disable-setuid-sandbox", ], ) # 等待浏览器完全启动 time.sleep(1) # 创建或获取页面 if len(browser.pages) > 0: page = browser.pages[0] else: page = browser.new_page() # 存储浏览器和页面实例 self.browsers[profile_name] = browser self.pages[profile_name] = page # 更新最后使用时间 self.profile_details[profile_name][ "last_used" ] = datetime.now().strftime("%Y-%m-%d %H:%M:%S") self._save_profiles_index() logger.info(f"配置 '{profile_name}' 的浏览器启动成功") success = True except Exception as e: logger.error(f"启动浏览器失败: {str(e)}") if self.playwright: try: self.playwright.stop() except Exception: pass raise e except Exception as e: logger.error(f"在线程中启动浏览器时出错: {str(e)}") finally: # 如果设置了回调，通知浏览器启动结果 if on_complete: on_complete(browser, page, success) # 删除运行任务记录 if profile_name in self.running_tasks: del self.running_tasks[profile_name] return success, browser, page def navigate_to_url(self, profile_name, url): """ 导航到指定URL. Args: profile_name: 配置文件名称 url: 要导航到的URL Returns: 成功返回True，失败返回False """ try: if profile_name not in self.browsers or profile_name not in self.pages: logger.error(f"配置 '{profile_name}' 的浏览器未启动") return False page = self.pages[profile_name] logger.info(f"浏览器 '{profile_name}' 导航到: {url}") page.goto(url, wait_until="networkidle", timeout=60000) return True except Exception as e: logger.error(f"导航到URL时出错: {str(e)}") return False def launch_multiple_browsers( self, profile_names, headless=False, on_each_complete=None, on_all_complete=None ): """ 并行启动多个浏览器. Args: profile_names: 配置文件名称列表 headless: 是否以无头模式启动 on_each_complete: 每个浏览器启动完成后的回调函数 on_all_complete: 所有浏览器启动完成后的回调函数 Returns: 包含所有Future对象的字典 """ tasks = {} results = {} # 定义单个浏览器完成时的回调 def on_browser_complete(browser, page, success, profile): results[profile] = (browser, page, success) if on_each_complete: on_each_complete(browser, page, success, profile) # 检查是否所有浏览器都已启动 if len(results) == len(profile_names) and on_all_complete: on_all_complete(results) # 提交所有浏览器启动任务 for profile in profile_names: # 为每个配置文件创建特定的回调 callback = lambda b, p, s, name=profile: on_browser_complete(b, p, s, name) task = self.launch_browser(profile, headless, callback) if task: tasks[profile] = task logger.info(f"已提交 {len(tasks)} 个浏览器启动任务") return tasks def close_all_browsers(self): """关闭所有浏览器实例.""" for profile_name in list(self.browsers.keys()): self._close_browser(profile_name) # 确保playwright实例被关闭 if self.playwright: try: self.playwright.stop() self.playwright = None except Exception as e: logger.error(f"关闭Playwright实例失败: {str(e)}") logger.info("已关闭所有浏览器实例") return True def __del__(self): """析构函数，确保资源被释放.""" self.close_all_browsers()