first commit

This commit is contained in:
jianxin committed 2025-04-02 21:01:25 +08:00
commit 4890793ce5
1029 files changed
+5992

No files matched your search

+388
View File
@@ -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()
+223
View File
@@ -0,0 +1,223 @@
import os
import shutil
import logging
import uuid
from datetime import datetime
from typing import List, Dict, Optional, Any
logger = logging.getLogger(__name__)
class BrowserProfileManager:
"""浏览器用户配置文件管理器."""
def __init__(self, settings: Dict[str, Any]):
"""
初始化浏览器配置文件管理器.
Args:
settings: 应用程序设置字典
"""
self.settings = settings
self.base_cache_dir = settings.get("profiles_dir", "")
# 确保基础缓存目录存在
if self.base_cache_dir and not os.path.exists(self.base_cache_dir):
os.makedirs(self.base_cache_dir, exist_ok=True)
def get_all_profiles(self) -> List[Dict[str, Any]]:
"""
获取所有浏览器配置文件.
Returns:
List[Dict]: 配置文件列表
"""
return self.settings.get("browser_profiles", [])
def get_profile_by_id(self, profile_id: str) -> Optional[Dict[str, Any]]:
"""
通过ID获取浏览器配置文件.
Args:
profile_id: 配置文件ID
Returns:
Dict or None: 找到的配置文件或None
"""
profiles = self.get_all_profiles()
for profile in profiles:
if profile.get("profile_id") == profile_id:
return profile
return None
def create_profile(
self, profile_name: str, douyin_account: Optional[str] = None
) -> Dict[str, Any]:
"""
创建新的浏览器配置文件.
Args:
profile_name: 配置文件名称
douyin_account: 关联的抖音账号(可选)
Returns:
Dict: 新创建的配置文件
"""
if not self.base_cache_dir:
raise ValueError("未设置浏览器缓存目录")
profile_id = str(uuid.uuid4())
profile_dir = os.path.join(self.base_cache_dir, f"Profile_{profile_id}")
# 创建配置文件目录
os.makedirs(profile_dir, exist_ok=True)
new_profile = {
"profile_id": profile_id,
"profile_name": profile_name,
"profile_dir": profile_dir,
"enabled": True,
"douyin_account": douyin_account,
"auto_reply_enabled": True,
"creation_time": datetime.now().isoformat(),
"last_used_time": None,
}
# 添加到配置文件列表
profiles = self.get_all_profiles()
profiles.append(new_profile)
self.settings["browser_profiles"] = profiles
# 保存设置
from douyin_imbot.config.settings import save_settings
save_settings(self.settings)
logger.info(f"已创建新的浏览器配置文件: {profile_name}")
return new_profile
def update_profile(self, profile_id: str, **kwargs) -> Optional[Dict[str, Any]]:
"""
更新浏览器配置文件.
Args:
profile_id: 配置文件ID
**kwargs: 要更新的字段
Returns:
Dict or None: 更新后的配置文件或None
"""
profiles = self.get_all_profiles()
for i, profile in enumerate(profiles):
if profile.get("profile_id") == profile_id:
# 更新字段
for key, value in kwargs.items():
if key in profile and key != "profile_id": # 不允许更新profile_id
profile[key] = value
# 如果更新了名称,可能需要更新目录
if "profile_name" in kwargs:
logger.info(f"已更新配置文件名称: {kwargs['profile_name']}")
# 保存更新后的列表
self.settings["browser_profiles"] = profiles
from douyin_imbot.config.settings import save_settings
save_settings(self.settings)
return profile
logger.warning(f"未找到ID为{profile_id}的配置文件")
return None
def delete_profile(self, profile_id: str, delete_files: bool = False) -> bool:
"""
删除浏览器配置文件.
Args:
profile_id: 配置文件ID
delete_files: 是否同时删除配置文件目录
Returns:
bool: 删除成功返回True,否则返回False
"""
profiles = self.get_all_profiles()
for i, profile in enumerate(profiles):
if profile.get("profile_id") == profile_id:
# 如果需要,删除配置文件目录
if delete_files and "profile_dir" in profile:
profile_dir = profile["profile_dir"]
if os.path.exists(profile_dir):
try:
shutil.rmtree(profile_dir)
logger.info(f"已删除配置文件目录: {profile_dir}")
except Exception as e:
logger.error(f"删除配置文件目录时出错: {str(e)}")
# 从列表中删除
removed_profile = profiles.pop(i)
self.settings["browser_profiles"] = profiles
# 保存设置
from douyin_imbot.config.settings import save_settings
save_settings(self.settings)
logger.info(f"已删除配置文件: {removed_profile.get('profile_name')}")
return True
logger.warning(f"未找到ID为{profile_id}的配置文件")
return False
def set_base_cache_dir(self, directory: str) -> bool:
"""
设置浏览器缓存基础目录.
Args:
directory: 新的缓存目录路径
Returns:
bool: 设置成功返回True,否则返回False
"""
if not directory:
logger.error("缓存目录不能为空")
return False
# 确保目录存在
try:
os.makedirs(directory, exist_ok=True)
self.base_cache_dir = directory
self.settings["browser_cache_dir"] = directory
# 保存设置
from douyin_imbot.config.settings import save_settings
save_settings(self.settings)
logger.info(f"已设置浏览器缓存目录: {directory}")
return True
except Exception as e:
logger.error(f"设置缓存目录时出错: {str(e)}")
return False
def update_last_used_time(self, profile_id: str) -> None:
"""
更新配置文件的最后使用时间.
Args:
profile_id: 配置文件ID
"""
self.update_profile(profile_id, last_used_time=datetime.now().isoformat())
def clear_cache(self):
"""清除配置文件缓存"""
try:
# 清除所有缓存的配置文件信息
self._cached_profiles = None
# 强制下次请求重新从磁盘读取
self._last_refresh_time = 0
logger.debug("配置文件管理器缓存已清除")
return True
except Exception as e:
logger.error(f"清除配置文件管理器缓存失败: {str(e)}")
return False