"""浏览器登录模块 - 使用 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 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: # 访问网站 self.tab.get(url) # 添加 cookies for cookie in cookies: self.tab.set.cookies(cookie) # 刷新页面 self.tab.refresh() # 等待页面加载 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}")