Merge remote-tracking branch 'alts/encoding'

3 files changed, 48 insertions, 7 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 438e3305..67e64d37 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -543,8 +543,7 @@ class ConnectionView(WWrap):
                 conn = self.flow.response
                 e = conn.headers["content-encoding"] or ["identity"]
                 if e[0] != "identity":
-                    conn.content = encoding.decode(e[0], conn.content)
-                    conn.headers["content-encoding"] = ["identity"]
+                    conn.decode()
                 else:
                     self.master.prompt_onekey(
                         "Select encoding: ",
@@ -563,11 +562,7 @@ class ConnectionView(WWrap):
             "z": "gzip",
             "d": "deflate",
         }
-        conn.content = encoding.encode(
-            encoding_map[key],
-            conn.content
-        )
-        conn.headers["content-encoding"] = [encoding_map[key]]
+        conn.encode(encoding_map[key])
         self.master.refresh_connection(self.flow)
 
 
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index e765d8b6..e19577ef 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -428,6 +428,27 @@ class Response(controller.Msg):
         c += self.headers.replace(pattern, repl, *args, **kwargs)
         return c
 
+    def decode(self):
+        """
+            Alters Response object, decoding its content based on the current
+            Content-Encoding header and changing Content-Encoding header to
+            'identity'.
+        """
+        self.content = encoding.decode(
+            (self.headers["content-encoding"] or ["identity"])[0],
+            self.content
+        )
+        self.headers["content-encoding"] = ["identity"]
+
+    def encode(self, e):
+        """
+            Alters Response object, encoding its content with the specified
+            coding. This method should only be called on Responses with
+            Content-Encoding headers of 'identity'.
+        """
+        self.content = encoding.encode(e, self.content)
+        self.headers["content-encoding"] = [e]
+
 
 class ClientDisconnect(controller.Msg):
     def __init__(self, client_conn):
diff --git a/test/test_proxy.py b/test/test_proxy.py
index ffd29a73..1e68fd3d 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -203,6 +203,31 @@ class uResponse(libpry.AutoTree):
         assert not "foo" in r.content
         assert r.headers["boo"] == ["boo"]
 
+    def test_decodeencode(self):
+        r = tutils.tresp()
+        r.headers["content-encoding"] = ["identity"]
+        r.content = "falafel"
+        r.decode()
+        assert r.headers["content-encoding"] == ["identity"]
+        assert r.content == "falafel"
+
+        r = tutils.tresp()
+        r.headers["content-encoding"] = ["identity"]
+        r.content = "falafel"
+        r.encode("identity")
+        assert r.headers["content-encoding"] == ["identity"]
+        assert r.content == "falafel"
+
+        r = tutils.tresp()
+        r.headers["content-encoding"] = ["identity"]
+        r.content = "falafel"
+        r.encode("gzip")
+        assert r.headers["content-encoding"] == ["gzip"]
+        assert r.content != "falafel"
+        r.decode()
+        assert r.headers["content-encoding"] == ["identity"]
+        assert r.content == "falafel"
+
 
 class uError(libpry.AutoTree):
     def test_getset_state(self):