From e17eacd8d77c78daa88d8f89ace990463378d98d Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Tue, 14 Apr 2015 13:45:38 +1200
Subject: [PATCH] New get_cookie and set_cookie implementations for HTTPRequest

---
 libmproxy/protocol/http.py | 25 ++++++++++++++++---------
 test/test_protocol_http.py | 22 +++++++++++++++-------
 2 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 94077e427..cd9458f20 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -6,7 +6,7 @@ import time
 import copy
 from email.utils import parsedate_tz, formatdate, mktime_tz
 import threading
-from netlib import http, tcp, http_status
+from netlib import http, tcp, http_status, http_cookies
 import netlib.utils
 from netlib import odict
 from .tcp import TCPHandler
@@ -670,15 +670,22 @@ class HTTPRequest(HTTPMessage):
         self.scheme, self.host, self.port, self.path = parts
 
     def get_cookies(self):
-        cookie_headers = self.headers.get("cookie")
-        if not cookie_headers:
-            return None
+        """
 
-        cookies = []
-        for header in cookie_headers:
-            pairs = [pair.partition("=") for pair in header.split(';')]
-            cookies.extend((pair[0], (pair[2], {})) for pair in pairs)
-        return dict(cookies)
+            Returns a possibly empty netlib.odict.ODict object.
+        """
+        ret = odict.ODict()
+        for i in self.headers["cookie"]:
+            ret.extend(http_cookies.parse_cookie_header(i))
+        return ret
+
+    def set_cookies(self, odict):
+        """
+            Takes an netlib.odict.ODict object. Over-writes any existing Cookie
+            headers.
+        """
+        v = http_cookies.format_cookie_header(odict)
+        self.headers["Cookie"] = [v]
 
     def replace(self, pattern, repl, *args, **kwargs):
         """
diff --git a/test/test_protocol_http.py b/test/test_protocol_http.py
index 6b949ce3e..11ab503be 100644
--- a/test/test_protocol_http.py
+++ b/test/test_protocol_http.py
@@ -161,7 +161,7 @@ class TestHTTPRequest:
         h = odict.ODictCaseless()
         r = tutils.treq()
         r.headers = h
-        assert r.get_cookies() is None
+        assert len(r.get_cookies()) == 0
 
     def test_get_cookies_single(self):
         h = odict.ODictCaseless()
@@ -170,7 +170,7 @@ class TestHTTPRequest:
         r.headers = h
         result = r.get_cookies()
         assert len(result)==1
-        assert result['cookiename']==('cookievalue',{})
+        assert result['cookiename']==['cookievalue']
 
     def test_get_cookies_double(self):
         h = odict.ODictCaseless()
@@ -179,8 +179,8 @@ class TestHTTPRequest:
         r.headers = h
         result = r.get_cookies()
         assert len(result)==2
-        assert result['cookiename']==('cookievalue',{})
-        assert result['othercookiename']==('othercookievalue',{})
+        assert result['cookiename']==['cookievalue']
+        assert result['othercookiename']==['othercookievalue']
 
     def test_get_cookies_withequalsign(self):
         h = odict.ODictCaseless()
@@ -189,10 +189,18 @@ class TestHTTPRequest:
         r.headers = h
         result = r.get_cookies()
         assert len(result)==2
-        assert result['cookiename']==('coo=kievalue',{})
-        assert result['othercookiename']==('othercookievalue',{})
-
+        assert result['cookiename']==['coo=kievalue']
+        assert result['othercookiename']==['othercookievalue']
 
+    def test_set_cookies(self):
+        h = odict.ODictCaseless()
+        h["Cookie"] = ["cookiename=cookievalue"]
+        r = tutils.treq()
+        r.headers = h
+        result = r.get_cookies()
+        result["cookiename"] = ["foo"]
+        r.set_cookies(result)
+        assert r.get_cookies()["cookiename"] == ["foo"]
 
 
 class TestHTTPResponse: