389 lines
15 KiB
Python
389 lines
15 KiB
Python
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()
|