first commit

douyin_imbot/browser/browser_manager.py

@@ -0,0 +1,388 @@
+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()