Initial commit

browser_login.py

@@ -0,0 +1,320 @@
+"""浏览器登录模块 - 使用 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}")