diff options
| -rw-r--r-- | mitmproxy/addons/stickycookie.py | 39 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_dumper.py | 1 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_stickycookie.py | 4 |
3 files changed, 20 insertions, 24 deletions
diff --git a/mitmproxy/addons/stickycookie.py b/mitmproxy/addons/stickycookie.py index 04d99975..e58e0a58 100644 --- a/mitmproxy/addons/stickycookie.py +++ b/mitmproxy/addons/stickycookie.py @@ -1,14 +1,14 @@ import collections from http import cookiejar +from typing import List, Tuple, Dict, Optional # noqa +from mitmproxy import http, flowfilter, ctx, exceptions from mitmproxy.net.http import cookies -from mitmproxy import exceptions -from mitmproxy import flowfilter -from mitmproxy import ctx +TOrigin = Tuple[str, int, str] -def ckey(attrs, f): +def ckey(attrs: Dict[str, str], f: http.HTTPFlow) -> TOrigin: """ Returns a (domain, port, path) tuple. """ @@ -21,18 +21,18 @@ def ckey(attrs, f): return (domain, f.request.port, path) -def domain_match(a, b): - if cookiejar.domain_match(a, b): +def domain_match(a: str, b: str) -> bool: + if cookiejar.domain_match(a, b): # type: ignore return True - elif cookiejar.domain_match(a, b.strip(".")): + elif cookiejar.domain_match(a, b.strip(".")): # type: ignore return True return False class StickyCookie: def __init__(self): - self.jar = collections.defaultdict(dict) - self.flt = None + self.jar = collections.defaultdict(dict) # type: Dict[TOrigin, Dict[str, str]] + self.flt = None # type: Optional[flowfilter.TFilter] def configure(self, updated): if "stickycookie" in updated: @@ -46,7 +46,7 @@ class StickyCookie: else: self.flt = None - def response(self, flow): + def response(self, flow: http.HTTPFlow): if self.flt: for name, (value, attrs) in flow.response.cookies.items(multi=True): # FIXME: We now know that Cookie.py screws up some cookies with @@ -63,24 +63,21 @@ class StickyCookie: if not self.jar[dom_port_path]: self.jar.pop(dom_port_path, None) else: - b = attrs.copy() - b.insert(0, name, value) - self.jar[dom_port_path][name] = b + self.jar[dom_port_path][name] = value - def request(self, flow): + def request(self, flow: http.HTTPFlow): if self.flt: - l = [] + cookie_list = [] # type: List[Tuple[str,str]] if flowfilter.match(self.flt, flow): - for domain, port, path in self.jar.keys(): + for (domain, port, path), c in self.jar.items(): match = [ domain_match(flow.request.host, domain), flow.request.port == port, flow.request.path.startswith(path) ] if all(match): - c = self.jar[(domain, port, path)] - l.extend([cookies.format_cookie_header(c[name].items(multi=True)) for name in c.keys()]) - if l: + cookie_list.extend(c.items()) + if cookie_list: # FIXME: we need to formalise this... - flow.request.stickycookie = True - flow.request.headers["cookie"] = "; ".join(l) + flow.metadata["stickycookie"] = True + flow.request.headers["cookie"] = cookies.format_cookie_header(cookie_list) diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py index d2cefe79..d8aa593b 100644 --- a/test/mitmproxy/addons/test_dumper.py +++ b/test/mitmproxy/addons/test_dumper.py @@ -68,7 +68,6 @@ def test_simple(): ctx.configure(d, flow_detail=4) flow = tflow.tflow() flow.request = tutils.treq() - flow.request.stickycookie = True flow.client_conn = mock.MagicMock() flow.client_conn.address[0] = "foo" flow.response = tutils.tresp(content=None) diff --git a/test/mitmproxy/addons/test_stickycookie.py b/test/mitmproxy/addons/test_stickycookie.py index 9092e09b..f77d019d 100644 --- a/test/mitmproxy/addons/test_stickycookie.py +++ b/test/mitmproxy/addons/test_stickycookie.py @@ -110,8 +110,8 @@ class TestStickyCookie: f.response.headers["Set-Cookie"] = c2 sc.response(f) googlekey = list(sc.jar.keys())[0] - assert len(sc.jar[googlekey].keys()) == 1 - assert list(sc.jar[googlekey]["somecookie"].items())[0][1] == "newvalue" + assert len(sc.jar[googlekey]) == 1 + assert sc.jar[googlekey]["somecookie"] == "newvalue" def test_response_delete(self): sc = stickycookie.StickyCookie() |