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

3 files changed, 55 insertions, 34 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()