"""浏览器登录模块 - 使用 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_cookies_dict(self) -> Dict[str, List[Dict]]: """ 获取当前浏览器的所有 cookie，按域名分组 Returns: 按域名分组的 cookie 字典 """ cookies = self.tab.cookies.as_dict # 按域名分组 cookie_dict = {} for cookie in cookies: domain = cookie.get('domain', '') # 统一域名格式（去掉开头的点） if domain.startswith('.'): domain = domain[1:] if domain not in cookie_dict: cookie_dict[domain] = [] cookie_dict[domain].append(cookie) return cookie_dict 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") for cookie in cookies: try: 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') 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 = self.tab.cookies.as_dict return cookies except Exception as e: logger.error(f"刷新页面并保存 cookie 失败: {e}") return [] def close(self): """关闭浏览器""" self._close_browser() class CookieManager: """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, 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: return json.load(f) 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_cookies(self, user_name: str, website_name: str) -> List[Dict]: """ 获取指定用户和网站的 cookie Args: user_name: 用户名 website_name: 网站名 Returns: cookie 列表 """ if user_name not in self.cookies: return [] if website_name not in self.cookies[user_name]: return [] return self.cookies[user_name][website_name] def save_cookies(self, user_name: str, website_name: str, cookies: List[Dict]): """ 保存指定用户和网站的 cookie Args: user_name: 用户名 website_name: 网站名 cookies: cookie 列表 """ if user_name not in self.cookies: self.cookies[user_name] = {} self.cookies[user_name][website_name] = cookies self._save_cookies() def get_cookies_for_domain(self, domain: str) -> List[Dict]: """ 获取指定域名的 cookie（从所有用户的 cookie 中查找） Args: domain: 域名 Returns: cookie 列表 """ all_cookies = [] for user_name, websites in self.cookies.items(): for website_name, cookies in websites.items(): for cookie in 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: all_cookies.append(cookie) return all_cookies if __name__ == "__main__": # 测试代码 manager = CookieManager("cookies.json") # 保存测试 cookie test_cookies = [ {"name": "test", "value": "123", "domain": "example.com", "path": "/"} ] manager.save_cookies("测试用户", "测试网站", test_cookies) # 获取 cookie cookies = manager.get_cookies("测试用户", "测试网站") logger.info(f"获取到的 cookie: {cookies}")