PappyProxy/pappyproxy/session.py

from .http import ResponseCookie

class Session(object):

    def __init__(self, cookie_names=None, header_names=None,
                 cookie_vals=None, header_vals=None):
        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 apply_req(self, req):
        for k, v in self.cookie_vals.iteritems():
            if isinstance(v, ResponseCookie):
                req.cookies[v.key] = v.val
            else:
                req.cookies[k] = v
        for k, v in self.header_vals.iteritems():
            req.headers[k] = v

    def apply_rsp(self, rsp):
        for k, v in self.cookie_vals.iteritems():
            if isinstance(v, ResponseCookie):
                rsp.set_cookie(v)
            else:
                cookie_str = '%s=%s' % (k, v)
                rsp.set_cookie(ResponseCookie(cookie_str))
        # Don't apply headers to responses

    def get_req(self, req, cookies=None, headers=None):
        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 get_rsp(self, rsp, cookies=None):
        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 v.key in self.cookies:
                    self.cookie_vals[v.key] = v