Merge pull request #1352 from cortesi/stickycookie

StickyCookies to addon

11 files changed, 204 insertions, 208 deletions

diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py
index b5419378..75326712 100644
--- a/mitmproxy/builtins/__init__.py
+++ b/mitmproxy/builtins/__init__.py
@@ -3,6 +3,7 @@ from __future__ import absolute_import, print_function, division
 from mitmproxy.builtins import anticache
 from mitmproxy.builtins import anticomp
 from mitmproxy.builtins import stickyauth
+from mitmproxy.builtins import stickycookie
 
 
 def default_addons():
@@ -10,4 +11,5 @@ def default_addons():
         anticache.AntiCache(),
         anticomp.AntiComp(),
         stickyauth.StickyAuth(),
+        stickycookie.StickyCookie(),
     ]
diff --git a/mitmproxy/builtins/stickycookie.py b/mitmproxy/builtins/stickycookie.py
new file mode 100644
index 00000000..5913976c
--- /dev/null
+++ b/mitmproxy/builtins/stickycookie.py
@@ -0,0 +1,70 @@
+import collections
+from six.moves import http_cookiejar
+from netlib.http import cookies
+
+from mitmproxy import exceptions
+from mitmproxy import filt
+
+
+def ckey(attrs, f):
+    """
+        Returns a (domain, port, path) tuple.
+    """
+    domain = f.request.host
+    path = "/"
+    if "domain" in attrs:
+        domain = attrs["domain"]
+    if "path" in attrs:
+        path = attrs["path"]
+    return (domain, f.request.port, path)
+
+
+def domain_match(a, b):
+    if http_cookiejar.domain_match(a, b):
+        return True
+    elif http_cookiejar.domain_match(a, b.strip(".")):
+        return True
+    return False
+
+
+class StickyCookie:
+    def __init__(self):
+        self.jar = collections.defaultdict(dict)
+        self.flt = None
+
+    def configure(self, options):
+        if options.stickycookie:
+            flt = filt.parse(options.stickycookie)
+            if not flt:
+                raise exceptions.OptionsError(
+                    "stickycookie: invalid filter expression: %s" % options.stickycookie
+                )
+            self.flt = flt
+
+    def response(self, flow):
+        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
+                # valid RFC 822/1123 datetime specifications for expiry. Sigh.
+                a = ckey(attrs, flow)
+                if domain_match(flow.request.host, a[0]):
+                    b = attrs.with_insert(0, name, value)
+                    self.jar[a][name] = b
+
+    def request(self, flow):
+        if self.flt:
+            l = []
+            if flow.match(self.flt):
+                for domain, port, path in self.jar.keys():
+                    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:
+                # FIXME: we need to formalise this...
+                flow.request.stickycookie = True
+                flow.request.headers["cookie"] = "; ".join(l)
diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py
index 605b0e23..5c015033 100644
--- a/mitmproxy/console/master.py
+++ b/mitmproxy/console/master.py
@@ -239,11 +239,6 @@ class ConsoleMaster(flow.FlowMaster):
         if options.limit:
             self.set_limit(options.limit)
 
-        r = self.set_stickycookie(options.stickycookie)
-        if r:
-            print("Sticky cookies error: {}".format(r), file=sys.stderr)
-            sys.exit(1)
-
         self.set_stream_large_bodies(options.stream_large_bodies)
 
         self.refresh_server_playback = options.refresh_server_playback
diff --git a/mitmproxy/console/options.py b/mitmproxy/console/options.py
index c76a058f..d363ba74 100644
--- a/mitmproxy/console/options.py
+++ b/mitmproxy/console/options.py
@@ -126,7 +126,7 @@ class Options(urwid.WidgetWrap):
                 select.Option(
                     "Sticky Cookies",
                     "t",
-                    lambda: master.stickycookie_txt,
+                    lambda: master.options.stickycookie,
                     self.sticky_cookie
                 ),
             ]
@@ -161,11 +161,11 @@ class Options(urwid.WidgetWrap):
         self.master.set_ignore_filter([])
         self.master.set_tcp_filter([])
         self.master.scripts = []
-        self.master.set_stickycookie(None)
 
         self.master.options.anticache = False
         self.master.options.anticomp = False
         self.master.options.stickyauth = None
+        self.master.options.stickycookie = None
 
         self.master.state.default_body_view = contentviews.get("Auto")
 
@@ -271,8 +271,8 @@ class Options(urwid.WidgetWrap):
     def sticky_cookie(self):
         signals.status_prompt.send(
             prompt = "Sticky cookie filter",
-            text = self.master.stickycookie_txt,
-            callback = self.master.set_stickycookie
+            text = self.master.options.stickycookie,
+            callback = self.master.options.setter("stickycookie")
         )
 
     def palette(self):
diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py
index 1357d7ca..fc41869c 100644
--- a/mitmproxy/console/statusbar.py
+++ b/mitmproxy/console/statusbar.py
@@ -173,10 +173,10 @@ class StatusBar(urwid.WidgetWrap):
             r.append("[")
             r.append(("heading_key", "Marked Flows"))
             r.append("]")
-        if self.master.stickycookie_txt:
+        if self.master.options.stickycookie:
             r.append("[")
             r.append(("heading_key", "t"))
-            r.append(":%s]" % self.master.stickycookie_txt)
+            r.append(":%s]" % self.master.options.stickycookie)
         if self.master.options.stickyauth:
             r.append("[")
             r.append(("heading_key", "u"))
diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py
index b953d131..d9142a4d 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -82,9 +82,6 @@ class DumpMaster(flow.FlowMaster):
         else:
             self.filt = None
 
-        if options.stickycookie:
-            self.set_stickycookie(options.stickycookie)
-
         if options.outfile:
             err = self.start_stream_to_path(
                 options.outfile[0],
@@ -230,7 +227,9 @@ class DumpMaster(flow.FlowMaster):
 
     def _echo_request_line(self, flow):
         if flow.request.stickycookie:
-            stickycookie = click.style("[stickycookie] ", fg="yellow", bold=True)
+            stickycookie = click.style(
+                "[stickycookie] ", fg="yellow", bold=True
+            )
         else:
             stickycookie = ""
 
diff --git a/mitmproxy/flow/__init__.py b/mitmproxy/flow/__init__.py
index 2a5d9b38..4c3bb828 100644
--- a/mitmproxy/flow/__init__.py
+++ b/mitmproxy/flow/__init__.py
@@ -5,7 +5,7 @@ from mitmproxy.flow.io import FlowWriter, FilteredFlowWriter, FlowReader, read_f
 from mitmproxy.flow.master import FlowMaster
 from mitmproxy.flow.modules import (
     AppRegistry, ReplaceHooks, SetHeaders, StreamLargeBodies, ClientPlaybackState,
-    ServerPlaybackState, StickyCookieState
+    ServerPlaybackState
 )
 from mitmproxy.flow.state import State, FlowView
 
@@ -16,5 +16,5 @@ __all__ = [
     "FlowWriter", "FilteredFlowWriter", "FlowReader", "read_flows_from_paths",
     "FlowMaster",
     "AppRegistry", "ReplaceHooks", "SetHeaders", "StreamLargeBodies", "ClientPlaybackState",
-    "ServerPlaybackState", "StickyCookieState", "State", "FlowView",
+    "ServerPlaybackState", "State", "FlowView",
 ]
diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py
index 06e1b460..d67ee7cc 100644
--- a/mitmproxy/flow/master.py
+++ b/mitmproxy/flow/master.py
@@ -8,7 +8,6 @@ from typing import List, Optional, Set  # noqa
 import netlib.exceptions
 from mitmproxy import controller
 from mitmproxy import exceptions
-from mitmproxy import filt
 from mitmproxy import models
 from mitmproxy import script
 from mitmproxy.flow import io
@@ -39,9 +38,6 @@ class FlowMaster(controller.Master):
         self.scripts = []  # type: List[script.Script]
         self.pause_scripts = False
 
-        self.stickycookie_state = None  # type: Optional[modules.StickyCookieState]
-        self.stickycookie_txt = None
-
         self.stream_large_bodies = None  # type: Optional[modules.StreamLargeBodies]
         self.refresh_server_playback = False
         self.replacehooks = modules.ReplaceHooks()
@@ -115,17 +111,6 @@ class FlowMaster(controller.Master):
     def set_tcp_filter(self, host_patterns):
         self.server.config.check_tcp = HostMatcher(host_patterns)
 
-    def set_stickycookie(self, txt):
-        if txt:
-            flt = filt.parse(txt)
-            if not flt:
-                return "Invalid filter expression."
-            self.stickycookie_state = modules.StickyCookieState(flt)
-            self.stickycookie_txt = txt
-        else:
-            self.stickycookie_state = None
-            self.stickycookie_txt = None
-
     def set_stream_large_bodies(self, max_size):
         if max_size is not None:
             self.stream_large_bodies = modules.StreamLargeBodies(max_size)
@@ -309,18 +294,11 @@ class FlowMaster(controller.Master):
             raise exceptions.FlowReadException(v.strerror)
 
     def process_new_request(self, f):
-        if self.stickycookie_state:
-            self.stickycookie_state.handle_request(f)
-
         if self.server_playback:
             pb = self.do_server_playback(f)
             if not pb and self.kill_nonreplay:
                 f.kill(self)
 
-    def process_new_response(self, f):
-        if self.stickycookie_state:
-            self.stickycookie_state.handle_response(f)
-
     def replay_request(self, f, block=False, run_scripthooks=True):
         """
             Returns None if successful, or error message if not.
@@ -431,7 +409,6 @@ class FlowMaster(controller.Master):
         if not f.reply.acked:
             if self.client_playback:
                 self.client_playback.clear(f)
-        self.process_new_response(f)
         if self.stream:
             self.stream.add(f)
         return f
diff --git a/mitmproxy/flow/modules.py b/mitmproxy/flow/modules.py
index 83228a04..2ad514f0 100644
--- a/mitmproxy/flow/modules.py
+++ b/mitmproxy/flow/modules.py
@@ -1,10 +1,8 @@
 from __future__ import absolute_import, print_function, division
 
-import collections
 import hashlib
 import re
 
-from six.moves import http_cookiejar
 from six.moves import urllib
 
 from mitmproxy import controller
@@ -12,7 +10,6 @@ from mitmproxy import filt
 from netlib import wsgi
 from netlib import version
 from netlib import strutils
-from netlib.http import cookies
 from netlib.http import http1
 
 
@@ -287,66 +284,3 @@ class ServerPlaybackState:
             return l[0]
         else:
             return l.pop(0)
-
-
-class StickyCookieState:
-    def __init__(self, flt):
-        """
-            flt: Compiled filter.
-        """
-        self.jar = collections.defaultdict(dict)
-        self.flt = flt
-
-    def ckey(self, attrs, f):
-        """
-            Returns a (domain, port, path) tuple.
-        """
-        domain = f.request.host
-        path = "/"
-        if "domain" in attrs:
-            domain = attrs["domain"]
-        if "path" in attrs:
-            path = attrs["path"]
-        return (domain, f.request.port, path)
-
-    def domain_match(self, a, b):
-        if http_cookiejar.domain_match(a, b):
-            return True
-        elif http_cookiejar.domain_match(a, b.strip(".")):
-            return True
-        return False
-
-    def handle_response(self, f):
-        for name, (value, attrs) in f.response.cookies.items(multi=True):
-            # FIXME: We now know that Cookie.py screws up some cookies with
-            # valid RFC 822/1123 datetime specifications for expiry. Sigh.
-            dom_port_path = self.ckey(attrs, f)
-
-            if self.domain_match(f.request.host, dom_port_path[0]):
-                if cookies.is_expired(attrs):
-                    # Remove the cookie from jar
-                    self.jar[dom_port_path].pop(name, None)
-
-                    # If all cookies of a dom_port_path have been removed
-                    # then remove it from the jar itself
-                    if not self.jar[dom_port_path]:
-                        self.jar.pop(dom_port_path, None)
-                else:
-                    b = attrs.with_insert(0, name, value)
-                    self.jar[dom_port_path][name] = b
-
-    def handle_request(self, f):
-        l = []
-        if f.match(self.flt):
-            for domain, port, path in self.jar.keys():
-                match = [
-                    self.domain_match(f.request.host, domain),
-                    f.request.port == port,
-                    f.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:
-            f.request.stickycookie = True
-            f.request.headers["cookie"] = "; ".join(l)
diff --git a/test/mitmproxy/builtins/test_stickycookie.py b/test/mitmproxy/builtins/test_stickycookie.py
new file mode 100644
index 00000000..e64ecb5b
--- /dev/null
+++ b/test/mitmproxy/builtins/test_stickycookie.py
@@ -0,0 +1,121 @@
+from .. import tutils, mastertest
+from mitmproxy.builtins import stickycookie
+from mitmproxy.flow import master
+from mitmproxy.flow import state
+from mitmproxy import options
+from netlib import tutils as ntutils
+
+
+def test_domain_match():
+    assert stickycookie.domain_match("www.google.com", ".google.com")
+    assert stickycookie.domain_match("google.com", ".google.com")
+
+
+class TestStickyCookie(mastertest.MasterTest):
+    def mk(self):
+        s = state.State()
+        m = master.FlowMaster(options.Options(stickycookie = ".*"), None, s)
+        sc = stickycookie.StickyCookie()
+        m.addons.add(sc)
+        return s, m, sc
+
+    def test_config(self):
+        sc = stickycookie.StickyCookie()
+        tutils.raises(
+            "invalid filter",
+            sc.configure,
+            options.Options(stickycookie = "~b")
+        )
+
+    def test_simple(self):
+        s, m, sc = self.mk()
+        m.addons.add(sc)
+
+        f = tutils.tflow(resp=True)
+        f.response.headers["set-cookie"] = "foo=bar"
+        self.invoke(m, "request", f)
+
+        f.reply.acked = False
+        self.invoke(m, "response", f)
+
+        assert sc.jar
+        assert "cookie" not in f.request.headers
+
+        f = f.copy()
+        f.reply.acked = False
+        self.invoke(m, "request", f)
+        assert f.request.headers["cookie"] == "foo=bar"
+
+    def _response(self, s, m, sc, cookie, host):
+        f = tutils.tflow(req=ntutils.treq(host=host, port=80), resp=True)
+        f.response.headers["Set-Cookie"] = cookie
+        self.invoke(m, "response", f)
+        return f
+
+    def test_response(self):
+        s, m, sc = self.mk()
+
+        c = "SSID=mooo; domain=.google.com, FOO=bar; Domain=.google.com; Path=/; " \
+            "Expires=Wed, 13-Jan-2021 22:23:01 GMT; Secure; "
+
+        self._response(s, m, sc, c, "host")
+        assert not sc.jar.keys()
+
+        self._response(s, m, sc, c, "www.google.com")
+        assert sc.jar.keys()
+
+        sc.jar.clear()
+        self._response(
+            s, m, sc, "SSID=mooo", "www.google.com"
+        )
+        assert list(sc.jar.keys())[0] == ('www.google.com', 80, '/')
+
+    def test_response_multiple(self):
+        s, m, sc = self.mk()
+
+        # Test setting of multiple cookies
+        c1 = "somecookie=test; Path=/"
+        c2 = "othercookie=helloworld; Path=/"
+        f = self._response(s, m, sc, c1, "www.google.com")
+        f.response.headers["Set-Cookie"] = c2
+        self.invoke(m, "response", f)
+        googlekey = list(sc.jar.keys())[0]
+        assert len(sc.jar[googlekey].keys()) == 2
+
+    def test_response_weird(self):
+        s, m, sc = self.mk()
+
+        # Test setting of weird cookie keys
+        f = tutils.tflow(req=ntutils.treq(host="www.google.com", port=80), resp=True)
+        cs = [
+            "foo/bar=hello",
+            "foo:bar=world",
+            "foo@bar=fizz",
+            "foo,bar=buzz",
+        ]
+        for c in cs:
+            f.response.headers["Set-Cookie"] = c
+            self.invoke(m, "response", f)
+        googlekey = list(sc.jar.keys())[0]
+        assert len(sc.jar[googlekey].keys()) == len(cs)
+
+    def test_response_overwrite(self):
+        s, m, sc = self.mk()
+
+        # Test overwriting of a cookie value
+        c1 = "somecookie=helloworld; Path=/"
+        c2 = "somecookie=newvalue; Path=/"
+        f = self._response(s, m, sc, c1, "www.google.com")
+        f.response.headers["Set-Cookie"] = c2
+        self.invoke(m, "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"
+
+    def test_request(self):
+        s, m, sc = self.mk()
+
+        f = self._response(s, m, sc, "SSID=mooo", "www.google.com")
+        assert "cookie" not in f.request.headers
+        self.invoke(m, "request", f)
+        assert "cookie" in f.request.headers
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 585dbf93..91342e58 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -40,86 +40,6 @@ def test_app_registry():
     assert ar.get(r)
 
 
-class TestStickyCookieState:
-
-    def _response(self, cookie, host):
-        s = flow.StickyCookieState(filt.parse(".*"))
-        f = tutils.tflow(req=netlib.tutils.treq(host=host, port=80), resp=True)
-        f.response.headers["Set-Cookie"] = cookie
-        s.handle_response(f)
-        return s, f
-
-    def test_domain_match(self):
-        s = flow.StickyCookieState(filt.parse(".*"))
-        assert s.domain_match("www.google.com", ".google.com")
-        assert s.domain_match("google.com", ".google.com")
-
-    def test_response(self):
-        c = (
-            "SSID=mooo; domain=.google.com, FOO=bar; Domain=.google.com; Path=/; "
-            "Expires=Wed, 13-Jan-2021 22:23:01 GMT; Secure; "
-        )
-
-        s, f = self._response(c, "host")
-        assert not s.jar.keys()
-
-        s, f = self._response(c, "www.google.com")
-        assert list(s.jar.keys())[0] == ('.google.com', 80, '/')
-
-        s, f = self._response("SSID=mooo", "www.google.com")
-        assert list(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 = list(s.jar.keys())[0]
-        assert len(s.jar[googlekey].keys()) == 2
-
-        # Test setting of weird cookie keys
-        s = flow.StickyCookieState(filt.parse(".*"))
-        f = tutils.tflow(req=netlib.tutils.treq(host="www.google.com", port=80), resp=True)
-        cs = [
-            "foo/bar=hello",
-            "foo:bar=world",
-            "foo@bar=fizz",
-            "foo,bar=buzz",
-        ]
-        for c in cs:
-            f.response.headers["Set-Cookie"] = c
-            s.handle_response(f)
-        googlekey = list(s.jar.keys())[0]
-        assert len(s.jar[googlekey]) == len(cs)
-
-        # 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 = list(s.jar.keys())[0]
-        assert len(s.jar[googlekey]) == 1
-        assert list(s.jar[googlekey]["somecookie"].values())[0] == "newvalue"
-
-    def test_response_delete(self):
-        c = "duffer=zafar; Path=/", "www.google.com"
-
-        # Test that a cookie is be deleted
-        # by setting the expire time in the past
-        s, f = self._response(*c)
-        f.response.headers["Set-Cookie"] = "duffer=; Expires=Thu, 01-Jan-1970 00:00:00 GMT"
-        s.handle_response(f)
-        assert not s.jar.keys()
-
-    def test_request(self):
-        s, f = self._response("SSID=mooo", b"www.google.com")
-        assert "cookie" not in f.request.headers
-        s.handle_request(f)
-        assert "cookie" in f.request.headers
-
-
 class TestClientPlaybackState:
 
     def test_tick(self):
@@ -967,28 +887,6 @@ class TestFlowMaster:
         fm.process_new_request(f)
         assert "killed" in f.error.msg
 
-    def test_stickycookie(self):
-        s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
-        assert "Invalid" in fm.set_stickycookie("~h")
-        fm.set_stickycookie(".*")
-        assert fm.stickycookie_state
-        fm.set_stickycookie(None)
-        assert not fm.stickycookie_state
-
-        fm.set_stickycookie(".*")
-        f = tutils.tflow(resp=True)
-        f.response.headers["set-cookie"] = "foo=bar"
-        fm.request(f)
-        f.reply.acked = False
-        fm.response(f)
-        assert fm.stickycookie_state.jar
-        assert "cookie" not in f.request.headers
-        f = f.copy()
-        f.reply.acked = False
-        fm.request(f)
-        assert f.request.headers["cookie"] == "foo=bar"
-
     def test_stream(self):
         with tutils.tmpdir() as tdir:
             p = os.path.join(tdir, "foo")