diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 20cd592ba..bb8bf2eab 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -22,6 +22,7 @@ from .proxy.config import HostMatcher from .protocol.http_replay import RequestReplayThread from .exceptions import Kill from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest +from collections import defaultdict class AppRegistry: @@ -309,7 +310,7 @@ class StickyCookieState: """ flt: Compiled filter. """ - self.jar = {} + self.jar = defaultdict(dict) self.flt = flt def ckey(self, m, f): @@ -337,7 +338,7 @@ class StickyCookieState: for m in c.values(): k = self.ckey(m, f) if self.domain_match(f.request.host, k[0]): - self.jar[k] = m + self.jar[k][m.key] = m def handle_request(self, f): l = [] @@ -349,10 +350,10 @@ class StickyCookieState: f.request.path.startswith(i[2]) ] if all(match): - l.append(self.jar[i].output(header="").strip()) + l.extend([m.output(header="").strip() for m in self.jar[i].values()]) if l: f.request.stickycookie = True - f.request.headers.set_all("cookie", l) + f.request.headers["cookie"] = "; ".join(l) class StickyAuthState: diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 60f6b1a92..8729cc779 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -69,6 +69,25 @@ class TestStickyCookieState: s, f = self._response("SSID=mooo", "www.google.com") assert s.jar.keys()[0] == ('www.google.com', 80, '/') + # Test setting of multiple cookies + c1 = "somecookie=test; Path=/" + c2 = "othercookie=helloworld; Path=/" + s, f = self._response(c1, "www.google.com") + f.response.headers["Set-Cookie"] = c2 + s.handle_response(f) + googlekey = s.jar.keys()[0] + assert len(s.jar[googlekey].keys()) == 2 + + # Test overwriting of a cookie value + c1 = "somecookie=helloworld; Path=/" + c2 = "somecookie=newvalue; Path=/" + s, f = self._response(c1, "www.google.com") + f.response.headers["Set-Cookie"] = c2 + s.handle_response(f) + googlekey = s.jar.keys()[0] + assert len(s.jar[googlekey].keys()) == 1 + assert s.jar[googlekey]["somecookie"].value == "newvalue" + def test_handle_request(self): s, f = self._response("SSID=mooo", "www.google.com") assert "cookie" not in f.request.headers