diff --git a/browser_manager.py b/browser_manager.py new file mode 100644 index 0000000..82181a2 --- /dev/null +++ b/browser_manager.py @@ -0,0 +1,226 @@ +"""浏览器管理模块""" +import time +from DrissionPage import ChromiumPage, ChromiumOptions +from logger import get_logger + +logger = get_logger(__name__) + + +class BrowserManager: + """浏览器管理器""" + + def __init__(self, browser_type: str = 'edge', headless: bool = False): + """ + 初始化浏览器管理器 + + Args: + browser_type: 浏览器类型 (edge 或 chrome) + headless: 是否无头模式 + """ + self.browser_type = browser_type + self.headless = headless + self.page = None + + def init_browser(self): + """初始化浏览器""" + co = ChromiumOptions() + if self.browser_type == 'edge': + co.set_browser_path(r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe") + co.set_argument('--no-sandbox') + co.set_argument('--disable-dev-shm-usage') + if self.headless: + co.set_argument('--headless') + self.page = ChromiumPage(addr_or_opts=co) + logger.info("浏览器初始化成功") + + def wait_page_loaded(self, timeout=30): + """等待页面加载完成""" + if self.page: + self.page.wait.doc_loaded(timeout=timeout) + self.page.ele('tag:body', timeout=timeout) + + def ensure_connection(self): + """确保浏览器连接正常""" + try: + if self.page: + self.page.get("about:blank") + return True + except: + pass + + try: + self.init_browser() + self.page.get("about:blank") + return True + except Exception as e: + logger.error(f"无法建立浏览器连接: {e}") + return False + + def import_cookies(self, cookies: list, domain: str): + """ + 导入cookies到浏览器 + + Args: + cookies: cookie列表 + domain: 域名 + + Returns: + 成功导入的cookie数量 + """ + if not self.ensure_connection(): + return 0 + + success_count = 0 + for cookie in cookies: + try: + if not cookie.get('name') or not cookie.get('value'): + continue + + cookie_dict = { + 'name': cookie.get('name'), + 'value': cookie.get('value'), + 'domain': cookie.get('domain', domain), + 'path': cookie.get('path', '/') + } + + self.page.set.cookies(cookie_dict) + success_count += 1 + except Exception as e: + logger.error(f"导入cookie失败: {cookie.get('name')} - {e}") + + logger.info(f"成功导入 {success_count} 个 cookies") + return success_count + + def import_local_storage(self, local_storage: dict): + """ + 导入localStorage到浏览器 + + Args: + local_storage: localStorage字典 + + Returns: + 成功导入的项目数量 + """ + if not self.ensure_connection(): + return 0 + + ls_count = 0 + for domain, items in local_storage.items(): + try: + for key, value in items.items(): + try: + self.page.set.local_storage(key, value) + ls_count += 1 + except Exception as e: + logger.error(f"导入localStorage失败: {key} - {e}") + except Exception as e: + logger.error(f"导入localStorage域失败: {domain} - {e}") + + logger.info(f"成功导入 {ls_count} 个 localStorage 项目") + return ls_count + + def verify_login(self, url: str, check_selector: str = "", success_text: str = "") -> bool: + """ + 验证登录状态 + + Args: + url: 网站地址 + check_selector: 登录检测选择器 + success_text: 登录成功时显示的文本 + + Returns: + 是否已登录 + """ + if not self.ensure_connection(): + return False + + try: + self.page.get(url) + self.wait_page_loaded() + + for i in range(3): + time.sleep(1) + self.page.refresh() + self.wait_page_loaded() + + if not check_selector and not success_text: + return False + + if check_selector: + try: + element = self.page.ele(check_selector, timeout=10) + if element: + element_text = element.text or "" + if not success_text: + return True + if success_text in element_text: + return True + except: + pass + + if success_text: + try: + page_text = self.page.html or "" + if success_text in page_text: + return True + except: + pass + + return False + + except Exception as e: + logger.error(f"验证登录状态失败: {e}") + return False + + def get_cookies(self) -> list: + """ + 获取当前浏览器的所有cookies + + Returns: + cookie列表 + """ + if not self.page: + return [] + + try: + return self.page.cookies() + except Exception as e: + logger.error(f"获取cookies失败: {e}") + return [] + + def get_local_storage(self) -> dict: + """ + 获取当前浏览器的localStorage + + Returns: + localStorage字典 + """ + if not self.page: + return {} + + try: + import json + js_code = """ + var result = {}; + for (var i = 0; i < localStorage.length; i++) { + var key = localStorage.key(i); + result[key] = localStorage.getItem(key); + } + JSON.stringify(result); + """ + result = self.page.run_js(js_code) + return json.loads(result) if result else {} + except Exception as e: + logger.error(f"获取localStorage失败: {e}") + return {} + + def close(self): + """关闭浏览器""" + if self.page: + try: + self.page.quit() + except Exception as e: + logger.error(f"关闭浏览器失败: {e}") + finally: + self.page = None + logger.info("浏览器已关闭")