diff options
| -rw-r--r-- | libmproxy/protocol/http.py | 25 | ||||
| -rw-r--r-- | 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 94077e42..cd9458f2 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 6b949ce3..11ab503b 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: |