Start abstracting out sticky cookie state.

3 files changed, 38 insertions, 3 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
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index ffc9b264..100a2e34 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -85,8 +85,8 @@ def parse_url(url):
         else:
             port = 80
     path = urlparse.urlunparse(('', '', path, params, query, fragment))
-    if not path:
-        path = "/"
+    if not path.startswith("/"):
+        path = "/" + path
     return scheme, host, port, path
 
 
diff --git a/test/test_flow.py b/test/test_flow.py
index e574f545..f58b1dd0 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -4,6 +4,17 @@ import utils
 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
+
+
 class uServerPlaybackState(libpry.AutoTree):
     def test_hash(self):
         s = flow.ServerPlaybackState(None)
@@ -345,6 +356,7 @@ class uFlowMaster(libpry.AutoTree):
 
 
 tests = [
+    uStickyCookieState(),
     uServerPlaybackState(),
     uFlow(),
     uState(),