From e2ee48a0ee8e6f6426686f8f7b06570cea20b236 Mon Sep 17 00:00:00 2001
From: Thomas Kriechbaumer <thomas@kriechbaumer.name>
Date: Mon, 9 May 2016 16:43:46 -0500
Subject: [PATCH] replace SimpleCookie with our own parser lib

---
 mitmproxy/flow.py           | 32 ++++++++++++++++++--------------
 netlib/http/cookies.py      |  2 +-
 test/mitmproxy/test_flow.py | 17 ++++++++++++++++-
 3 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 9292e76ab..5d46d79e9 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -13,9 +13,9 @@ from six.moves import http_cookies, http_cookiejar, urllib
 import os
 import re
 
-from netlib import wsgi
+from netlib import wsgi, odict
 from netlib.exceptions import HttpException
-from netlib.http import Headers, http1
+from netlib.http import Headers, http1, cookies
 from . import controller, tnetstring, filt, script, version, flow_format_compat
 from .onboarding import app
 from .proxy.config import HostMatcher
@@ -313,15 +313,17 @@ class StickyCookieState:
         self.jar = defaultdict(dict)
         self.flt = flt
 
-    def ckey(self, m, f):
+    def ckey(self, attrs, f):
         """
             Returns a (domain, port, path) tuple.
         """
-        return (
-            m["domain"] or f.request.host,
-            f.request.port,
-            m["path"] or "/"
-        )
+        domain = f.request.host
+        path = "/"
+        if attrs["domain"]:
+            domain = attrs["domain"][-1]
+        if attrs["path"]:
+            path = attrs["path"][-1]
+        return (domain, f.request.port, path)
 
     def domain_match(self, a, b):
         if http_cookiejar.domain_match(a, b):
@@ -334,11 +336,12 @@ class StickyCookieState:
         for i in f.response.headers.get_all("set-cookie"):
             # FIXME: We now know that Cookie.py screws up some cookies with
             # valid RFC 822/1123 datetime specifications for expiry. Sigh.
-            c = http_cookies.SimpleCookie(str(i))
-            for m in c.values():
-                k = self.ckey(m, f)
-                if self.domain_match(f.request.host, k[0]):
-                    self.jar[k][m.key] = m
+            name, value, attrs = cookies.parse_set_cookie_header(str(i))
+            a = self.ckey(attrs, f)
+            if self.domain_match(f.request.host, a[0]):
+                b = attrs.lst
+                b.insert(0, [name, value])
+                self.jar[a][name] = odict.ODictCaseless(b)
 
     def handle_request(self, f):
         l = []
@@ -350,7 +353,8 @@ class StickyCookieState:
                     f.request.path.startswith(i[2])
                 ]
                 if all(match):
-                    l.extend([m.output(header="").strip() for m in self.jar[i].values()])
+                    c = self.jar[i]
+                    l.extend([cookies.format_cookie_header(c[name]) for name in c.keys()])
         if l:
             f.request.stickycookie = True
             f.request.headers["cookie"] = "; ".join(l)
diff --git a/netlib/http/cookies.py b/netlib/http/cookies.py
index 2d5c18ca5..4451f1daf 100644
--- a/netlib/http/cookies.py
+++ b/netlib/http/cookies.py
@@ -209,7 +209,7 @@ def refresh_set_cookie_header(c, delta):
         raise ValueError("Invalid Cookie")
 
     if "expires" in attrs:
-        e = parsedate_tz(attrs["expires"][0])
+        e = parsedate_tz(attrs["expires"][-1])
         if e:
             f = mktime_tz(e) + delta
             attrs["expires"] = [formatdate(f)]
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 145e91cf2..4af588b28 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -76,6 +76,21 @@ class TestStickyCookieState:
         googlekey = s.jar.keys()[0]
         assert len(s.jar[googlekey].keys()) == 2
 
+        # Test setting of weird cookie keys
+        s = flow.StickyCookieState(filt.parse(".*"))
+        f = tutils.tflow(req=netlib.tutils.treq(host="www.google.com", port=80), resp=True)
+        cs = [
+            "foo/bar=hello",
+            "foo:bar=world",
+            "foo@bar=fizz",
+            "foo,bar=buzz",
+        ]
+        for c in cs:
+            f.response.headers["Set-Cookie"] = c
+            s.handle_response(f)
+        googlekey = s.jar.keys()[0]
+        assert len(s.jar[googlekey].keys()) == len(cs)
+
         # Test overwriting of a cookie value
         c1 = "somecookie=helloworld; Path=/"
         c2 = "somecookie=newvalue; Path=/"
@@ -84,7 +99,7 @@ class TestStickyCookieState:
         s.handle_response(f)
         googlekey = s.jar.keys()[0]
         assert len(s.jar[googlekey].keys()) == 1
-        assert s.jar[googlekey]["somecookie"].value == "newvalue"
+        assert s.jar[googlekey]["somecookie"].items()[0][1] == "newvalue"
 
     def test_handle_request(self):
         s, f = self._response("SSID=mooo", "www.google.com")