320 lines
8.8 KiB
Python
320 lines
8.8 KiB
Python
"""浏览器登录模块 - 使用 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}") |