549 lines
18 KiB
Python
549 lines
18 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_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
|
|
|
|
|