Add set_cookies method to HTTPResponse

2 files changed, 30 insertions, 0 deletions

diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index da8eaa01..eb7749ea 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 0276cab7..28c63430 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):