class SessionStore(requests.Session):
"""A ``requests.Session`` subclass implementing a file-based session store."""
def __init__(self, endpoint, persist_cookie, cookie_file, reauthenticate,
login_params, authn_header=None, debug=False):
super(SessionStore, self).__init__()
self.session_base_url = '{0}/api/session'.format(endpoint)
self.reauthenticate = reauthenticate
self.persist_cookie = persist_cookie
self.login_params = login_params
self.authn_header = authn_header
self._debug = debug
if persist_cookie:
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/.nuvla 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):
kwargs.setdefault('timeout', DEFAULT_TIMEOUT)
if self.authn_header is not None:
headers = kwargs.get('headers', {})
headers['nuvla-authn-info'] = self.authn_header
kwargs['headers'] = headers
return super(SessionStore, self).request(*args, **kwargs)
def request(self, *args, **kwargs):
response = self._request(*args, **kwargs)
if not self.verify and response.cookies:
self._unsecure_cookie(args[1], response)
if self.persist_cookie and 'Set-Cookie' in response.headers:
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 = self._request(*args, **kwargs)
return response
def cimi_login(self, login_params):
self.login_params = login_params
if self.login_params:
method = 'POST'
endpoint = self.session_base_url
headers = {'Content-Type': 'application/json',
'Accept': 'application/json'}
json = {'template': login_params}
if self._debug:
_request_debug(method, endpoint, doc=json, headers=headers)
response = self.request(method, endpoint, headers=headers, json=json)
if self._debug:
_response_debug(response)
return response
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)
if self.persist_cookie:
self.cookies.save()
except KeyError:
pass
