StickyCookies to addon

7 files changed, 81 insertions, 38 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/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