class SessionStore(requests.Session): """A ``requests.Session`` subclass implementing a file-based session store. """ def __init__(self, cookie_file=None): super(SessionStore, self).__init__() if cookie_file is None: cookie_file = conf.DEFAULT_COOKIE_FILE cookie_dir = os.path.dirname(cookie_file) self.cookies = MozillaCookieJar(cookie_file) # Create the $HOME/.slipstream dir if it doesn't exist if not os.path.isdir(cookie_dir): os.mkdir(cookie_dir, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) # Load existing cookies if the cookies.txt exists if os.path.isfile(cookie_file): self.cookies.load(ignore_discard=True) self.cookies.clear_expired_cookies() def request(self, *args, **kwargs): response = super(SessionStore, self).request(*args, **kwargs) self.cookies.save(ignore_discard=True) return response def clear(self, domain): """Clear cookies for the specified domain.""" try: self.cookies.clear(domain) self.cookies.save() except KeyError: pass
class SessionStore(requests.Session): """A ``requests.Session`` subclass implementing a file-based session store.""" def __init__(self, endpoint, reauthenticate, cookie_file=None): super(SessionStore, self).__init__() self.session_base_url = '{}/api/session'.format(endpoint) self.reauthenticate = reauthenticate self.login_params = None if cookie_file is None: cookie_file = DEFAULT_COOKIE_FILE cookie_dir = os.path.dirname(cookie_file) self.cookies = MozillaCookieJar(cookie_file) # Create the $HOME/.slipstream dir if it doesn't exist if not os.path.isdir(cookie_dir): os.mkdir(cookie_dir, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) # Load existing cookies if the cookies.txt exists if os.path.isfile(cookie_file): self.cookies.load(ignore_discard=True) self.cookies.clear_expired_cookies() def need_to_login(self, accessed_url, status_code): return self.reauthenticate and status_code in [401, 403] and accessed_url != self.session_base_url def request(self, *args, **kwargs): super_request = super(SessionStore, self).request response = super_request(*args, **kwargs) if not self.verify and response.cookies: self._unsecure_cookie(args[1], response) self.cookies.save(ignore_discard=True) url = args[1] if self.need_to_login(url, response.status_code): login_response = self.cimi_login(self.login_params) if login_response is not None and login_response.status_code == 201: # retry the call after reauthentication response = super_request(*args, **kwargs) return response def cimi_login(self, login_params): self.login_params = login_params if self.login_params: return self.request('POST', self.session_base_url, headers={'Content-Type': 'application/json', 'Accept': 'application/json'}, json={'sessionTemplate': login_params}) else: return None def _unsecure_cookie(self, url_str, response): url = urlparse(url_str) if url.scheme == 'http': for cookie in response.cookies: cookie.secure = False self.cookies.set_cookie_if_ok(cookie, MockRequest(response.request)) def clear(self, domain): """Clear cookies for the specified domain.""" try: self.cookies.clear(domain) self.cookies.save() except KeyError: pass