diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py index da8eaa010..eb7749ea6 100644 --- a/libmproxy/protocol/http.py +++ b/libmproxy/protocol/http.py @@ -918,6 +918,25 @@ class HTTPResponse(HTTPMessage): ret.append([name, [value, attrs]]) return odict.ODict(ret) + def set_cookies(self, odict): + """ + Set the Set-Cookie headers on this response, over-writing existing + headers. + + Accepts an ODict of the same format as that returned by get_cookies. + """ + values = [] + for i in odict.lst: + values.append( + http_cookies.format_set_cookie_header( + i[0], + i[1][0], + i[1][1] + ) + ) + self.headers["Set-Cookie"] = values + + class HTTPFlow(Flow): """ diff --git a/test/test_protocol_http.py b/test/test_protocol_http.py index 0276cab77..28c634308 100644 --- a/test/test_protocol_http.py +++ b/test/test_protocol_http.py @@ -285,6 +285,17 @@ class TestHTTPResponse: assert "othercookie" in result assert result["othercookie"][0] == ["othervalue", odict.ODict()] + def test_set_cookies(self): + h = odict.ODictCaseless() + resp = tutils.tresp() + v = resp.get_cookies() + v.add("foo", ["bar", odict.ODictCaseless()]) + resp.set_cookies(v) + + v = resp.get_cookies() + assert len(v) == 1 + assert v["foo"] == [["bar", odict.ODictCaseless()]] + class TestHTTPFlow(object):