from .http import ResponseCookie class Session(object): """ A class used to maintain a session over multiple requests. Can remember cookies and apply a specific header to requests. It is also possible to give the session a list of cookie names and it will only save those cookies. """ def __init__(self, cookie_names=None, header_names=None, cookie_vals=None, header_vals=None): """ Session(self, cookie_names=None, header_names=None, cookie_vals=None, header_vals=None) Constructor :param cookie_names: A whitelist for cookies that should be saved from :func:`~pappyproxy.session.Session.save_req` and :func:`~pappyproxy.session.Session.save_rsp` in the session. If no values are given, all cookies will be saved. :param header_names: A whitelist for headers that should be saved from :func:`~pappyproxy.session.Session.save_req` in the session. If no values are given, no headers will be saved. :param cookie_vals: A dictionary of cookies to populate the session session with. The key should be the cookie name, and the value can be either a string or a :class:`~pappyproxy.http.ResponseCookie`. If a :class:`~pappyproxy.http.ResponseCookie` is given, its flags will be used in :func:`~pappyproxy.session.Session.apply_rsp`. :param header_vals: A dictionary of header values to populate the session with. The key should be the header name and the value should be a string which should be the header value. """ self.cookies = cookie_names or [] self.headers = header_names or [] self.cookie_vals = cookie_vals or {} self.header_vals = header_vals or {} if cookie_vals: for k, v in cookie_vals.iteritems(): if k not in self.cookies: self.cookies.append(k) if header_vals: for k, v in header_vals.iteritems(): if k not in self.headers: self.headers.append(k) def _cookie_obj(self, k, v): """ Returns the value as a cookie object regardless of if the cookie is a string or a ResponseCookie. """ if isinstance(v, ResponseCookie): return v else: cookie_str = '%s=%s' % (k, v) return ResponseCookie(cookie_str) def _cookie_val(self, v): """ Returns the value of the cookie regardless of if the value is a string or a ResponseCookie """ if isinstance(v, ResponseCookie): return v.val else: return v def apply_req(self, req): """ apply_req(request) Apply saved headers and cookies to the request """ for k, v in self.cookie_vals.iteritems(): req.cookies[k] = self._cookie_val(v) for k, v in self.header_vals.iteritems(): req.headers[k] = v def apply_rsp(self, rsp): """ apply_rsp(response) Will add a Set-Cookie header for each saved cookie. Will not apply any saved headers. If the cookie was added from a call to :func:`~pappyproxy.session.Session.save_rsp`, the Set-Cookie flags will be the same as the original response. """ for k, v in self.cookie_vals.iteritems(): val = self._cookie_obj(k, v) rsp.set_cookie(val) # Don't apply headers to responses def save_req(self, req, cookies=None, headers=None): """ save_req(req, cookies=None, headers=None) Updates the state of the session from the given request. Cookie and headers can be added to their whitelists by passing in a list for either ``cookies`` or ``headers``. """ if cookies: for c in cookies: if c not in self.cookies: self.cookies.append(c) if headers: for h in headers: if h not in self.headers: self.headers.append(h) if cookies: for cookie in cookies: if cookie in req.cookies: if cookie not in self.cookies: self.cookies.append(cookie) cookie_str = '%s=%s' % (cookie, req.cookies[cookie]) self.cookie_vals[cookie] = ResponseCookie(cookie_str) else: for k, v in req.cookies.all_pairs(): if k in self.cookies: cookie_str = '%s=%s' % (k, v) self.cookie_vals[cookie] = ResponseCookie(cookie_str) if headers: for header in headers: if header in self.headers: self.header_vals[header] = req.headers[header] def save_rsp(self, rsp, cookies=None, save_all=False): """ save_rsp(rsp, cookies=None) Update the state of the session from the response. Only cookies can be updated from a response. Additional values can be added to the whitelist by passing in a list of values for the ``cookies`` parameter. If save_all is given, all set cookies will be added to the session. """ if cookies: for c in cookies: if c not in self.cookies: self.cookies.append(c) if cookies: for cookie in cookies: if cookie in rsp.cookies: if cookie not in self.cookies: self.cookies.append(cookie) self.cookie_vals[cookie] = rsp.cookies[cookie] else: for k, v in rsp.cookies.all_pairs(): if save_all: self.cookie_vals[v.key] = v if not v.key in self.cookies: self.cookies.append(v.key) elif v.key in self.cookies: self.cookie_vals[v.key] = v def set_cookie(key, val): """ set_cookie(key, val) Set a cookie in the session. ``val`` can be either a string or a :class:`~pappyproxy.http.ResponseCookie`. If a :class:`~pappyproxy.http.ResponseCookie` is used, make sure its ``key`` value is the same as the key passed in to the function. """ self.cookie_vals[key] = val def get_cookie(key): """ get_cookie(key) Returns a string with the value of the cookie with the given string, even if the value is a :class:`~pappyproxy.http.ResponseCookie`. If you want to get a :class:`~pappyproxy.http.ResponseCookie`, use :func:`~pappyproxy.session.Session.get_rsp_cookie`. """ if not key in self.cookie_vals: raise KeyError('Cookie is not stored in session.') v = self.cookie_vals[key] return self._cookie_val(v) def get_rsp_cookie(key): """ get_rsp_cookie(key) Returns the :class:`~pappyproxy.http.ResponseCookie` associated with the key regardless of if the value is stored as a string or a :class:`~pappyproxy.http.ResponseCookie`. """ if not key in self.cookie_vals: raise KeyError('Cookie is not stored in session.') v = self.cookie_vals[key] return self._cookie_obj(key, v)