# -*- coding: utf-8 -*- """ Created on 2022/9/7 4:11 PM --------- @summary: --------- @author: Boris @email: boris_liu@foxmail.com """ import json import os import re from collections import defaultdict from typing import Union, List try: from typing import Literal # python >= 3.8 except ImportError: # python <3.8 from typing_extensions import Literal from playwright.sync_api import Page, BrowserContext, ViewportSize, ProxySettings from playwright.sync_api import Playwright, Browser from playwright.sync_api import Response from playwright.sync_api import sync_playwright from feapder.utils import tools from feapder.utils.log import log from feapder.utils.webdriver.webdirver import * class PlaywrightDriver(WebDriver): def __init__( self, *, page_on_event_callback: dict = None, storage_state_path: str = None, driver_type: Literal["chromium", "firefox", "webkit"] = "chromium", url_regexes: list = None, save_all: bool = False, **kwargs ): """ Args: page_on_event_callback: page.on() 事件的回调 如 page_on_event_callback={"dialog": lambda dialog: dialog.accept()} storage_state_path: 保存浏览器状态的路径 driver_type: 浏览器类型 chromium, firefox, webkit url_regexes: 拦截接口，支持正则，数组类型 save_all: 是否保存所有拦截的接口, 默认只保存最后一个 **kwargs: """ super(PlaywrightDriver, self).__init__(**kwargs) self.driver: Playwright = None self.browser: Browser = None self.context: BrowserContext = None self.page: Page = None self.url = None self.storage_state_path = storage_state_path self._driver_type = driver_type or "chromium" self._page_on_event_callback = page_on_event_callback self._url_regexes = url_regexes self._save_all = save_all if self._save_all and self._url_regexes: log.warning( "获取完拦截的数据后, 请主动调用PlaywrightDriver的clear_cache()方法清空拦截的数据，否则数据会一直累加，导致内存溢出" ) self._cache_data = defaultdict(list) else: self._cache_data = {} self._setup() def _setup(self): # 处理参数 if self._proxy: proxy = self._proxy() if callable(self._proxy) else self._proxy proxy = self.format_context_proxy(proxy) else: proxy = None user_agent = ( self._user_agent() if callable(self._user_agent) else self._user_agent ) view_size = ViewportSize( width=self._window_size[0], height=self._window_size[1] ) # 初始化浏览器对象 self.driver = sync_playwright().start() self.browser = getattr(self.driver, self._driver_type).launch( headless=self._headless, args=["--no-sandbox"], proxy=proxy, executable_path=self._executable_path, downloads_path=self._download_path, ) if self.storage_state_path and os.path.exists(self.storage_state_path): self.context = self.browser.new_context( user_agent=user_agent, screen=view_size, viewport=view_size, proxy=proxy, storage_state=self.storage_state_path, ) else: self.context = self.browser.new_context( user_agent=user_agent, screen=view_size, viewport=view_size, proxy=proxy, ) if self._use_stealth_js: path = os.path.join(os.path.dirname(__file__), "../js/stealth.min.js") self.context.add_init_script(path=path) self.page = self.context.new_page() self.page.set_default_timeout(self._timeout * 1000) if self._page_on_event_callback: for event, callback in self._page_on_event_callback.items(): self.page.on(event, callback) if self._url_regexes: self.page.on("response", self.on_response) def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): if exc_val: log.error(exc_val) self.quit() return True def format_context_proxy(self, proxy) -> ProxySettings: """ Args: proxy: username:password@ip:port / ip:port Returns: { "server": "ip:port" "username": username, "password": password, } server: http://ip:port or socks5://ip:port. Short form ip:port is considered an HTTP proxy. """ if "@" in proxy: certification, _proxy = proxy.split("@") username, password = certification.split(":") context_proxy = ProxySettings( server=_proxy, username=username, password=password, ) else: context_proxy = ProxySettings(server=proxy) return context_proxy def save_storage_stage(self): if self.storage_state_path: os.makedirs(os.path.dirname(self.storage_state_path), exist_ok=True) self.context.storage_state(path=self.storage_state_path) def quit(self): self.page.close() self.context.close() self.browser.close() self.driver.stop() @property def domain(self): return tools.get_domain(self.url or self.page.url) @property def cookies(self): cookies_json = {} for cookie in self.page.context.cookies(): cookies_json[cookie["name"]] = cookie["value"] return cookies_json @cookies.setter def cookies(self, val: Union[dict, List[dict]]): """ 设置cookie Args: val: List[{name: str, value: str, url: Union[str, NoneType], domain: Union[str, NoneType], path: Union[str, NoneType], expires: Union[float, NoneType], httpOnly: Union[bool, NoneType], secure: Union[bool, NoneType], sameSite: Union["Lax", "None", "Strict", NoneType]}] Returns: """ if isinstance(val, list): self.page.context.add_cookies(val) else: cookies = [] for key, value in val.items(): cookies.append( {"name": key, "value": value, "url": self.url or self.page.url} ) self.page.context.add_cookies(cookies) @property def user_agent(self): return self.page.evaluate("() => navigator.userAgent") def on_response(self, response: Response): for regex in self._url_regexes: if re.search(regex, response.request.url): intercept_request = InterceptRequest( url=response.request.url, headers=response.request.headers, data=response.request.post_data, ) intercept_response = InterceptResponse( request=intercept_request, url=response.url, headers=response.headers, content=response.body(), status_code=response.status, ) if self._save_all: self._cache_data[regex].append(intercept_response) else: self._cache_data[regex] = intercept_response def get_response(self, url_regex) -> InterceptResponse: if self._save_all: response_list = self._cache_data.get(url_regex) if response_list: return response_list[-1] return self._cache_data.get(url_regex) def get_all_response(self, url_regex) -> List[InterceptResponse]: """ 获取所有匹配的响应, 仅在save_all=True时有效 Args: url_regex: Returns: """ response_list = self._cache_data.get(url_regex, []) if not isinstance(response_list, list): return [response_list] return response_list def get_text(self, url_regex): return ( self.get_response(url_regex).content.decode() if self.get_response(url_regex) else None ) def get_all_text(self, url_regex): """ 获取所有匹配的响应文本, 仅在save_all=True时有效 Args: url_regex: Returns: """ return [ response.content.decode() for response in self.get_all_response(url_regex) ] def get_json(self, url_regex): return ( json.loads(self.get_text(url_regex)) if self.get_response(url_regex) else None ) def get_all_json(self, url_regex): """ 获取所有匹配的响应json, 仅在save_all=True时有效 Args: url_regex: Returns: """ return [json.loads(text) for text in self.get_all_text(url_regex)] def clear_cache(self): self._cache_data = defaultdict(list)