优化SessionPage和WebPage的cookies设置

This commit is contained in:
g1879 2022-11-15 01:33:44 +08:00
parent e1f3b729e6
commit dd762fccb9
2 changed files with 66 additions and 40 deletions

View File

@ -15,18 +15,60 @@ from tldextract import extract
from DownloadKit import DownloadKit
from .base import BasePage
from .config import _cookie_to_dict
from .config import _cookie_to_dict, SessionOptions, _cookies_to_tuple
from .session_element import SessionElement, make_session_ele
class SessionPage(BasePage):
"""SessionPage封装了页面操作的常用功能使用requests来获取、解析网页"""
def __init__(self, session: Session, timeout: float = 10):
def __init__(self, session_or_options: Union[Session, SessionOptions] = None,
timeout: float = 10):
"""初始化函数"""
super().__init__(timeout)
self._session = session
self._response = None
self._create_session(session_or_options)
def _create_session(self, Session_or_Options) -> None:
if Session_or_Options is None or isinstance(Session_or_Options, SessionOptions):
options = Session_or_Options or SessionOptions()
self._set_session(options.as_dict())
elif isinstance(Session_or_Options, Session):
self._session = Session_or_Options
def _set_session(self, data: dict) -> None:
"""根据传入字典对session进行设置 \n
:param data: session配置字典
:return: None
"""
if self._session is None:
self._session = Session()
if 'headers' in data:
self._session.headers = CaseInsensitiveDict(data['headers'])
if 'cookies' in data:
self.set_cookies(data['cookies'])
attrs = ['auth', 'proxies', 'hooks', 'params', 'verify',
'cert', 'stream', 'trust_env', 'max_redirects'] # , 'adapters'
for i in attrs:
if i in data:
self._session.__setattr__(i, data[i])
def set_cookies(self, cookies):
cookies = _cookies_to_tuple(cookies)
for cookie in cookies:
if cookie['value'] is None:
cookie['value'] = ''
kwargs = {x: cookie[x] for x in cookie
if x.lower() in ('version', 'port', 'domain', 'path', 'secure',
'expires', 'discard', 'comment', 'comment_url', 'rest')}
if 'expiry' in cookie:
kwargs['expires'] = cookie['expiry']
self.session.cookies.set(cookie['name'], cookie['value'], **kwargs)
def __call__(self,
loc_or_str: Union[Tuple[str, str], str, SessionElement],

View File

@ -4,7 +4,6 @@ from typing import Union, Tuple, List
from DownloadKit import DownloadKit
from pychrome import Tab
from requests import Session, Response
from requests.structures import CaseInsensitiveDict
from tldextract import extract
from .chromium_element import ChromiumElement
@ -210,7 +209,7 @@ class WebPage(SessionPage, ChromiumPage, BasePage):
return
self._mode = 's' if self._mode == 'd' else 'd'
print(self._mode)
# s模式转d模式
if self._mode == 'd':
if not self._has_driver:
@ -246,7 +245,7 @@ class WebPage(SessionPage, ChromiumPage, BasePage):
selenium_user_agent = self.run_script("navigator.userAgent;")
self.session.headers.update({"User-Agent": selenium_user_agent})
self.set_cookies(super(SessionPage, self).get_cookies(as_dict=True), set_session=True)
self.set_cookies(self._get_driver_cookies(as_dict=True), set_session=True)
def cookies_to_driver(self) -> None:
"""把session对象的cookies复制到driver对象"""
@ -270,28 +269,32 @@ class WebPage(SessionPage, ChromiumPage, BasePage):
if self._mode == 's':
return super().get_cookies(as_dict, all_domains)
elif self._mode == 'd':
return super(SessionPage, self).get_cookies(as_dict)
return self._get_driver_cookies(as_dict)
def _get_driver_cookies(self, as_dict: bool = False):
cookies = super(SessionPage, self).driver.Network.getCookies()['cookies']
if as_dict:
return {cookie['name']: cookie['value'] for cookie in cookies}
else:
return cookies
def set_cookies(self, cookies, set_session: bool = False, set_driver: bool = False):
# 添加cookie到driver
if set_driver:
super(SessionPage, self).set_cookies(cookies)
cookies = _cookies_to_tuple(cookies)
result_cookies = []
for cookie in cookies:
if not cookie.get('domain', None):
continue
c = {'value': '' if cookie['value'] is None else cookie['value'],
'name': cookie['name'],
'domain': cookie['domain']}
result_cookies.append(c)
super(SessionPage, self).driver.Network.setCookies(cookies=result_cookies)
# 添加cookie到session
if set_session:
cookies = _cookies_to_tuple(cookies)
for cookie in cookies:
if cookie['value'] is None:
cookie['value'] = ''
kwargs = {x: cookie[x] for x in cookie
if x.lower() in ('version', 'port', 'domain', 'path', 'secure',
'expires', 'discard', 'comment', 'comment_url', 'rest')}
if 'expiry' in cookie:
kwargs['expires'] = cookie['expiry']
self.session.cookies.set(cookie['name'], cookie['value'], **kwargs)
super().set_cookies(cookies)
def check_page(self, by_requests: bool = False) -> Union[bool, None]:
"""d模式时检查网页是否符合预期 \n
@ -368,25 +371,6 @@ class WebPage(SessionPage, ChromiumPage, BasePage):
elif self._mode == 'd':
return super(SessionPage, self)._ele(loc_or_ele, timeout=timeout, single=single)
def _set_session(self, data: dict) -> None:
"""根据传入字典对session进行设置 \n
:param data: session配置字典
:return: None
"""
if self._session is None:
self._session = Session()
if 'headers' in data:
self._session.headers = CaseInsensitiveDict(data['headers'])
if 'cookies' in data:
self.set_cookies(data['cookies'], set_session=True)
attrs = ['auth', 'proxies', 'hooks', 'params', 'verify',
'cert', 'stream', 'trust_env', 'max_redirects'] # , 'adapters'
for i in attrs:
if i in data:
self._session.__setattr__(i, data[i])
def _set_driver_options(self, Tab_or_Options):
"""处理driver设置"""
if Tab_or_Options is None: