"""浏览器登录模块 - 使用 DrissionPage 控制浏览器进行登录检测和 cookie 管理""" from typing import Dict, List, Optional import json import time from DrissionPage import Chromium, ChromiumOptions from applogger import get_logger logger = get_logger(__name__) class BrowserLogin: """浏览器登录客户端""" def __init__(self, browser_type: str = "edge", headless: bool = False): """ 初始化浏览器登录客户端 Args: browser_type: 浏览器类型 (chrome 或 edge) headless: 是否无头模式 """ self.browser_type = browser_type.lower() self.headless = headless self.browser = None self.tab = None def _create_browser(self): """创建浏览器实例""" co = ChromiumOptions() # 设置浏览器路径 if self.browser_type == "edge": co.set_browser_path("C:\\Program Files (x86)\\Microsoft\\Edge\\Application\\msedge.exe") elif self.browser_type == "chrome": co.set_browser_path("chrome") # 设置无头模式 if self.headless: co.headless = True # 设置其他选项 co.set_argument('--no-sandbox') co.set_argument('--disable-dev-shm-usage') # 创建浏览器 self.browser = Chromium(co) self.tab = self.browser.new_tab() def _close_browser(self): """关闭浏览器""" if self.browser: try: self.browser.quit() except Exception as e: logger.error(f"关闭浏览器失败: {e}") finally: self.browser = None self.tab = None def _get_domain_from_url(self, url: str) -> str: """ 从 URL 中提取域名 Args: url: 网站地址 Returns: 域名 """ from urllib.parse import urlparse parsed = urlparse(url) return parsed.netloc def get_all_cookies(self) -> List[Dict]: """ 获取当前浏览器实例的所有 cookie Returns: 所有 cookie 的列表 """ # 先获取 cookies 对象，然后调用 as_dict 方法 cookies_obj = self.tab.cookies cookies = cookies_obj() if callable(cookies_obj) else cookies_obj cookies = cookies.as_dict() if hasattr(cookies, 'as_dict') else cookies # 确保返回完整的 cookie 列表，包含所有属性 formatted_cookies = [] for cookie in cookies: formatted_cookie = { 'name': cookie.get('name', ''), 'value': cookie.get('value', ''), 'domain': cookie.get('domain', ''), 'path': cookie.get('path', '/'), 'expires': cookie.get('expirationDate', None), 'secure': cookie.get('secure', False), 'httpOnly': cookie.get('httpOnly', False), 'sameSite': cookie.get('sameSite', 'Lax') } formatted_cookies.append(formatted_cookie) return formatted_cookies def check_login(self, url: str, check_selector: str, success_text: str) -> bool: """ 检查是否已登录 Args: url: 网站地址 check_selector: 登录检测选择器 success_text: 登录成功时显示的文本 Returns: 是否已登录 """ if not self.browser or not self.tab: self._create_browser() try: # 访问网站 self.tab.get(url) # 等待页面加载 time.sleep(2) # 检查登录状态 element = self.tab.ele(check_selector, timeout=10) if element and element.text and success_text in element.text: return True return False except Exception as e: logger.error(f"检查登录状态失败: {e}") return False def import_all_cookies(self, cookies: List[Dict]) -> bool: """ 一次性导入所有 cookies 到浏览器 Args: cookies: cookie 列表 Returns: 是否成功 """ if not self.browser or not self.tab: self._create_browser() success_count = 0 fail_count = 0 logger.info(f"开始导入 {len(cookies)} 个 cookie") try: # 准备所有 cookie 数据 formatted_cookies = [] for cookie in cookies: cookie_dict = { 'name': cookie.get('name', ''), 'value': cookie.get('value', ''), 'domain': cookie.get('domain', ''), 'path': cookie.get('path', '/'), } if cookie.get('secure'): cookie_dict['secure'] = True if cookie.get('httpOnly'): cookie_dict['httpOnly'] = True if cookie.get('sameSite'): cookie_dict['sameSite'] = cookie.get('sameSite') formatted_cookies.append(cookie_dict) # 根据 DrissionPage 文档，set.cookies() 方法可以直接接收列表格式的多个 cookie logger.info("使用一次性注入方式导入 cookie") self.tab.set.cookies(formatted_cookies) success_count = len(formatted_cookies) except Exception as e: logger.error(f"导入 cookie 失败: {e}") # 回退到逐个注入 logger.info("回退到逐个注入方式导入 cookie") for cookie_dict in formatted_cookies: try: self.tab.set.cookies(cookie_dict) success_count += 1 except Exception as e: fail_count += 1 logger.info(f"Cookie 导入完成: 成功 {success_count}, 失败 {fail_count}") return fail_count == 0 def verify_login(self, url: str, check_selector: str = "", success_text: str = "") -> bool: """ 验证登录状态（cookies已预先导入） Args: url: 网站地址 check_selector: 登录检测选择器（可选） success_text: 登录成功时显示的文本（可选） Returns: 是否已登录 """ if not self.browser or not self.tab: self._create_browser() try: self.tab.get(url) time.sleep(5) # 如果没有指定选择器，只检查页面是否成功加载 if not check_selector: logger.info(" 未指定登录检测选择器，跳过登录验证") return True # 尝试查找元素 try: element = self.tab.ele(check_selector, timeout=10) if element: element_text = element.text or "" logger.info(f" 找到元素，文本内容: {element_text[:50]}...") if not success_text: logger.info(" 未指定成功文本，找到元素即认为登录成功") return True if success_text in element_text: logger.info(f" 检测到成功文本: {success_text}") return True else: logger.warning(f" 未检测到成功文本 '{success_text}'，元素文本: {element_text[:100]}") return False else: logger.warning(f" 未找到元素: {check_selector}") return False except Exception as ele_error: logger.warning(f" 查找元素失败: {ele_error}") return False except Exception as e: logger.error(f"验证登录状态失败: {e}") return False def login_with_cookies( self, url: str, cookies: List[Dict], check_selector: str, success_text: str ) -> bool: """ 使用 cookie 登录 Args: url: 网站地址 cookies: cookie 列表 check_selector: 登录检测选择器 success_text: 登录成功时显示的文本 Returns: 是否登录成功 """ if not self.browser or not self.tab: self._create_browser() try: from urllib.parse import urlparse parsed = urlparse(url) base_url = f"{parsed.scheme}://{parsed.netloc}" self.tab.get(base_url) time.sleep(1) logger.info(f"开始设置 {len(cookies)} 个 cookie") for cookie in cookies: try: cookie_dict = { 'name': cookie.get('name', ''), 'value': cookie.get('value', ''), 'domain': cookie.get('domain', parsed.netloc), 'path': cookie.get('path', '/'), } if cookie.get('secure'): cookie_dict['secure'] = True if cookie.get('httpOnly'): cookie_dict['httpOnly'] = True if cookie.get('sameSite'): cookie_dict['sameSite'] = cookie.get('sameSite') self.tab.set.cookies(cookie_dict) logger.info(f"设置 cookie: {cookie_dict['name']}") except Exception as e: logger.error(f"设置 cookie 失败 {cookie.get('name')}: {e}") self.tab.get(url) time.sleep(3) element = self.tab.ele(check_selector, timeout=10) if element and element.text and success_text in element.text: return True return False except Exception as e: logger.error(f"使用 cookie 登录失败: {e}") return False def refresh_and_save_cookies(self, url: str) -> List[Dict]: """ 刷新页面并保存新的 cookie Args: url: 网站地址 Returns: cookie 列表 """ if not self.browser or not self.tab: self._create_browser() try: # 刷新页面 self.tab.refresh() # 等待页面加载 time.sleep(3) # 获取 cookies cookies_obj = self.browser.cookies cookies = cookies_obj() if callable(cookies_obj) else cookies_obj cookies = cookies.as_dict() if hasattr(cookies, 'as_dict') else cookies # 确保返回的是列表 if isinstance(cookies, list): return cookies elif isinstance(cookies, dict): # 如果是字典，转换为列表 return [cookies] else: # 如果是其他类型，返回空列表 return [] except Exception as e: logger.error(f"刷新页面并保存 cookie 失败: {e}") return [] def close(self): """关闭浏览器""" self._close_browser() class CookieManager: """Cookie 管理器 - 实现浏览器实例级别的 Cookie 共享""" def __init__(self, cookie_file: str): """ 初始化 Cookie 管理器 Args: cookie_file: cookie 文件路径 """ self.cookie_file = cookie_file self.cookies = self._load_cookies() def _load_cookies(self) -> Dict[str, List[Dict]]: """加载 cookie 文件""" import os if os.path.exists(self.cookie_file): try: with open(self.cookie_file, 'r', encoding='utf-8') as f: data = json.load(f) # 兼容旧格式 if isinstance(data, dict) and any(isinstance(v, dict) for v in data.values()): # 转换旧格式为新格式 new_data = {} for user, websites in data.items(): all_cookies = [] for site_cookies in websites.values(): all_cookies.extend(site_cookies) new_data[user] = all_cookies return new_data return data except Exception as e: logger.error(f"加载 cookie 文件失败: {e}") return {} def _save_cookies(self): """保存 cookie 文件""" try: with open(self.cookie_file, 'w', encoding='utf-8') as f: json.dump(self.cookies, f, indent=2, ensure_ascii=False) except Exception as e: logger.error(f"保存 cookie 文件失败: {e}") def get_all_cookies(self, user_name: str) -> List[Dict]: """ 获取指定用户的所有 cookie Args: user_name: 用户名 Returns: cookie 列表 """ return self.cookies.get(user_name, []) def save_all_cookies(self, user_name: str, cookies: List[Dict]): """ 保存指定用户的所有 cookie Args: user_name: 用户名 cookies: cookie 列表 """ self.cookies[user_name] = cookies self._save_cookies() def save_cookies(self, user_name: str, website_name: str, cookies: List[Dict]): """ 保存指定用户和网站的 cookie（兼容旧接口） Args: user_name: 用户名 website_name: 网站名（未使用，保持兼容） cookies: cookie 列表 """ # 保持向后兼容，将网站的cookie添加到用户的所有cookie中 if user_name not in self.cookies: self.cookies[user_name] = [] # 合并新cookie到现有cookie中 existing_cookies = self.cookies[user_name] # 创建现有cookie的索引，用于快速查找 cookie_index = {} for i, cookie in enumerate(existing_cookies): if isinstance(cookie, dict): key = (cookie.get('name'), cookie.get('domain'), cookie.get('path')) cookie_index[key] = i # 添加或更新cookie if isinstance(cookies, list): valid_cookies = [] for cookie in cookies: if isinstance(cookie, dict): valid_cookies.append(cookie) for cookie in valid_cookies: key = (cookie.get('name'), cookie.get('domain'), cookie.get('path')) if key in cookie_index: # 更新现有cookie existing_cookies[cookie_index[key]] = cookie else: # 添加新cookie existing_cookies.append(cookie) if valid_cookies: logger.info(f" 已保存 {len(valid_cookies)} 个新 cookie") self._save_cookies() def add_cookie(self, user_name: str, cookie: Dict): """ 为指定用户添加单个 cookie Args: user_name: 用户名 cookie: cookie 字典 """ if user_name not in self.cookies: self.cookies[user_name] = [] # 检查是否已存在相同的 cookie existing_index = -1 for i, existing_cookie in enumerate(self.cookies[user_name]): if (existing_cookie.get('name') == cookie.get('name') and existing_cookie.get('domain') == cookie.get('domain') and existing_cookie.get('path') == cookie.get('path')): existing_index = i break if existing_index >= 0: # 更新现有 cookie self.cookies[user_name][existing_index] = cookie else: # 添加新 cookie self.cookies[user_name].append(cookie) self._save_cookies() def remove_cookie(self, user_name: str, cookie_name: str, domain: str): """ 从指定用户中移除指定的 cookie Args: user_name: 用户名 cookie_name: cookie 名称 domain: cookie 域名 """ if user_name in self.cookies: self.cookies[user_name] = [ cookie for cookie in self.cookies[user_name] if not (cookie.get('name') == cookie_name and cookie.get('domain') == domain) ] self._save_cookies() def clear_cookies(self, user_name: str): """ 清空指定用户的所有 cookie Args: user_name: 用户名 """ if user_name in self.cookies: self.cookies[user_name] = [] self._save_cookies() def get_cookies_for_domain(self, user_name: str, domain: str) -> List[Dict]: """ 获取指定用户的指定域名的 cookie Args: user_name: 用户名 domain: 域名 Returns: cookie 列表 """ user_cookies = self.get_all_cookies(user_name) filtered_cookies = [] for cookie in user_cookies: cookie_domain = cookie.get('domain', '') # 统一域名格式 if cookie_domain.startswith('.'): cookie_domain = cookie_domain[1:] if domain in cookie_domain or cookie_domain in domain: filtered_cookies.append(cookie) return filtered_cookies