More mature sticky cookie primitive. Use it in console.py.

4 files changed, 103 insertions, 42 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 187502a6..f63115aa 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -1154,15 +1154,6 @@ class ConsoleMaster(flow.FlowMaster):
         else:
             self.state.beep = None
 
-    def set_stickycookie(self, txt):
-        if txt:
-            self.stickycookie = filt.parse(txt)
-            if not self.stickycookie:
-                return "Invalid filter expression."
-        else:
-            self.stickyhosts = {}
-            self.stickycookie = None
-
     def drawscreen(self):
         size = self.ui.get_cols_rows()
         canvas = self.view.render(size, focus=1)
@@ -1311,20 +1302,10 @@ class ConsoleMaster(flow.FlowMaster):
     def handle_request(self, r):
         f = flow.FlowMaster.handle_request(self, r)
         if f:
-            if f.match(self.stickycookie):
-                hid = (f.request.host, f.request.port)
-                if f.request.headers.has_key("cookie"):
-                    self.stickyhosts[hid] = f.request.headers["cookie"]
-                elif hid in self.stickyhosts:
-                    f.request.headers["cookie"] = self.stickyhosts[hid]
             self.process_flow(f, r)
 
     def handle_response(self, r):
         f = flow.FlowMaster.handle_response(self, r)
         if f:
-            if f.match(self.stickycookie):
-                hid = (f.request.host, f.request.port)
-                if f.response.headers.has_key("set-cookie"):
-                    self.stickyhosts[hid] = f.response.headers["set-cookie"]
             self.process_flow(f, r)
 
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 9636c3bd..9b083036 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -3,7 +3,7 @@
     with their responses, and provide filtering and interception facilities.
 """
 import subprocess, base64, sys, json, hashlib, Cookie, cookielib, copy
-import proxy, threading, netstring
+import proxy, threading, netstring, filt
 import controller
 
 class RunException(Exception):
@@ -90,26 +90,44 @@ class ServerPlaybackState:
 
 
 class StickyCookieState:
-    def __init__(self):
+    def __init__(self, flt):
+        """
+            flt: A compiled filter.
+        """
         self.jar = {}
+        self.flt = flt
 
-    def ckey(self, c):
-        c = copy.copy(c)
-        del c["expires"]
-        return str(c)
+    def ckey(self, m, f):
+        """
+            Returns a (domain, port, path) tuple. 
+        """
+        return (
+            m["domain"] or f.request.host,
+            f.request.port,
+            m["path"] or "/"
+        )
 
-    def add_cookies(self, headers):
-        for i in headers:
+    def handle_response(self, f):
+        for i in f.response.headers.get("set-cookie", []):
             c = Cookie.SimpleCookie(i)
             m = c.values()[0]
-            self.jar[self.ckey(m)] = m
+            k = self.ckey(m, f)
+            if cookielib.domain_match(f.request.host, k[0]):
+                self.jar[self.ckey(m, f)] = m
+
+    def handle_request(self, f):
+        if f.match(self.flt):
+            cs = []
+            for i in self.jar.keys():
+                match = [
+                    cookielib.domain_match(i[0], f.request.host),
+                    f.request.port == i[1],
+                    f.request.path.startswith(i[2])
+                ]
+                if all(match):
+                    l = f.request.headers.setdefault("cookie", [])
+                    l.append(self.jar[i].output(header="").strip())
 
-    def get_cookies(self, domain, path):
-        cs = []
-        for i in self.jar.values():
-            if cookielib.domain_match(domain, i["domain"]) and path.startswith(i.get("path", "/")):
-                cs.append(i)
-        return cs
 
 
 class Flow:
@@ -369,6 +387,7 @@ class FlowMaster(controller.Master):
         self.playback = None
         self.scripts = {}
         self.kill_nonreplay = False
+        self.stickycookie_state = False
 
     def _runscript(self, f, script):
         return f.run_script(script)
@@ -379,6 +398,15 @@ class FlowMaster(controller.Master):
     def set_request_script(self, s):
         self.scripts["request"] = s
 
+    def set_stickycookie(self, txt):
+        if txt:
+            flt = filt.parse(txt)
+            if not flt:
+                return "Invalid filter expression."
+            self.stickycookie_state = StickyCookieState(flt)
+        else:
+            self.stickycookie_state = None
+
     def start_playback(self, flows, kill, headers):
         """
             flows: A list of flows.
@@ -419,6 +447,8 @@ class FlowMaster(controller.Master):
 
     def handle_request(self, r):
         f = self.state.add_request(r)
+        if self.stickycookie_state:
+            self.stickycookie_state.handle_request(f)
         if "request" in self.scripts:
             self._runscript(f, self.scripts["request"])
         if self.playback:
@@ -434,6 +464,8 @@ class FlowMaster(controller.Master):
         f = self.state.add_response(r)
         if not f:
             r.ack()
+        if self.stickycookie_state:
+            self.stickycookie_state.handle_response(f)
         if "response" in self.scripts:
             self._runscript(f, self.scripts["response"])
         return f
diff --git a/libmproxy/utils.py b/libmproxy/utils.py
index afef8e63..c67b9397 100644
--- a/libmproxy/utils.py
+++ b/libmproxy/utils.py
@@ -172,6 +172,10 @@ class MultiDict:
         key = self._helper[0](key)
         return self._d.get(key, d)
 
+    def __contains__(self, key):
+        key = self._helper[0](key)
+        return self._d.__contains__(key)
+
     def __eq__(self, other):
         return dict(self) == dict(other)
 
@@ -192,6 +196,10 @@ class MultiDict:
         key = self._helper[0](key)
         return self._d.has_key(key)
 
+    def setdefault(self, key, default=None):
+        key = self._helper[0](key)
+        return self._d.setdefault(key, default)
+
     def keys(self):
         return self._d.keys()
 
diff --git a/test/test_flow.py b/test/test_flow.py
index f58b1dd0..10b445db 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -5,14 +5,33 @@ import libpry
 
 
 class uStickyCookieState(libpry.AutoTree):
-    def test_simple(self):
-        s = flow.StickyCookieState()
-        s.add_cookies(
-            ["SSID=mooo, FOO=bar; Domain=.google.com; Path=/; Expires=Wed, 13-Jan-2021 22:23:01 GMT; Secure; "]
-        )
-        assert len(s.jar) == 1
-        assert len(s.get_cookies("www.google.com", "/foo")) == 1
-        assert len(s.get_cookies("www.foo.com", "/foo")) == 0
+    def _response(self, cookie, host):
+        s = flow.StickyCookieState(filt.parse(".*"))
+        f = utils.tflow_full()
+        f.request.host = host 
+        f.response.headers["Set-Cookie"] = [cookie]
+        s.handle_response(f)
+        return s, f
+
+    def test_handle_response(self):
+        c = "SSID=mooo, 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 s.jar.keys()
+
+        s, f = self._response("SSID=mooo", "www.google.com")
+        assert s.jar.keys()[0] == ('www.google.com', 80, '/')
+
+    def test_handle_request(self):
+        s, f = self._response("SSID=mooo", "www.google.com")
+        assert "cookie" not in f.request.headers
+        s.handle_request(f)
+        assert "cookie" in f.request.headers
+
 
 
 class uServerPlaybackState(libpry.AutoTree):
@@ -353,6 +372,27 @@ class uFlowMaster(libpry.AutoTree):
         r.request.content = "gibble"
         assert not fm.do_playback(r)
 
+    def test_stickycookie(self):
+        s = flow.State()
+        fm = flow.FlowMaster(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(".*")
+        tf = utils.tflow_full()
+        tf.response.headers["set-cookie"] = ["foo=bar"]
+        fm.handle_request(tf.request)
+        f = fm.handle_response(tf.response)
+        assert fm.stickycookie_state.jar
+        assert not "cookie" in tf.request.headers
+        fm.handle_request(tf.request)
+        assert tf.request.headers["cookie"] == ["foo=bar"]
+
+
+
 
 
 tests = [