Start abstracting out sticky cookie state.

1 files changed, 24 insertions, 1 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 42870f17..9636c3bd 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -2,7 +2,7 @@
     This module provides more sophisticated flow tracking. These match requests
     with their responses, and provide filtering and interception facilities.
 """
-import subprocess, base64, sys, json, hashlib
+import subprocess, base64, sys, json, hashlib, Cookie, cookielib, copy
 import proxy, threading, netstring
 import controller
 
@@ -89,6 +89,29 @@ class ServerPlaybackState:
         return l.pop(0)
 
 
+class StickyCookieState:
+    def __init__(self):
+        self.jar = {}
+
+    def ckey(self, c):
+        c = copy.copy(c)
+        del c["expires"]
+        return str(c)
+
+    def add_cookies(self, headers):
+        for i in headers:
+            c = Cookie.SimpleCookie(i)
+            m = c.values()[0]
+            self.jar[self.ckey(m)] = m
+
+    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:
     def __init__(self, request):
         self.request = request