Add decoding/encoding for requests.

3 files changed, 71 insertions, 42 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 67e64d37..6fd8c2ca 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -541,23 +541,25 @@ class ConnectionView(WWrap):
         elif key == "z":
             if self.state.view_flow_mode == VIEW_FLOW_RESPONSE:
                 conn = self.flow.response
-                e = conn.headers["content-encoding"] or ["identity"]
-                if e[0] != "identity":
-                    conn.decode()
-                else:
-                    self.master.prompt_onekey(
-                        "Select encoding: ",
-                        (
-                            ("gzip", "z"),
-                            ("deflate", "d"),
-                        ),
-                        self.encode_response_callback
-                    )
-                self.master.refresh_connection(self.flow)
+            else:
+                conn = self.flow.request
+            e = conn.headers["content-encoding"] or ["identity"]
+            if e[0] != "identity":
+                conn.decode()
+            else:
+                self.master.prompt_onekey(
+                    "Select encoding: ",
+                    (
+                        ("gzip", "z"),
+                        ("deflate", "d"),
+                    ),
+                    self.encode_callback,
+                    conn
+                )
+            self.master.refresh_connection(self.flow)
         return key
 
-    def encode_response_callback(self, key):
-        conn = self.flow.response
+    def encode_callback(self, key, conn):
         encoding_map = {
             "z": "gzip",
             "d": "deflate",
@@ -1347,8 +1349,8 @@ class ConsoleMaster(flow.FlowMaster):
             ("e", "edit request/response"),
             ("p", "previous flow"),
             ("v", "view body in external viewer"),
-            ("z", "switch response encoding"),
-            ("tab", "toggle response/request view"),
+            ("z", "encode/decode a request/response"),
+            ("tab", "toggle request/response view"),
             ("space", "next flow"),
         ]
         text.extend(format_keyvals(keys, key="key", val="text", indent=4))
@@ -1412,7 +1414,7 @@ class ConsoleMaster(flow.FlowMaster):
         self.view.set_focus("footer")
         self.prompting = (callback, [])
 
-    def prompt_onekey(self, prompt, keys, callback):
+    def prompt_onekey(self, prompt, keys, callback, *args):
         """
             Keys are a set of (word, key) tuples. The appropriate key in the
             word is highlighted.
@@ -1426,7 +1428,7 @@ class ConsoleMaster(flow.FlowMaster):
         prompt.extend(mkup)
         prompt.append(")? ")
         self.onekey = "".join([i[1] for i in keys])
-        self.prompt(prompt, "", callback)
+        self.prompt(prompt, "", callback, *args)
 
     def prompt_done(self):
         self.prompting = False
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index e19577ef..a2805069 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -125,9 +125,32 @@ def parse_request_line(request):
     if major != 1:
         raise ProxyError(400, "Unsupported protocol")
     return method, scheme, host, port, path, minor
+    
+
+class HTTPMsg(controller.Msg):
+    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 Request(controller.Msg):
+class Request(HTTPMsg):
     FMT = '%s %s HTTP/1.1\r\n%s\r\n%s'
     FMT_PROXY = '%s %s://%s:%s%s HTTP/1.1\r\n%s\r\n%s'
     def __init__(self, client_conn, host, port, scheme, method, path, headers, content, timestamp=None):
@@ -293,7 +316,7 @@ class Request(controller.Msg):
         return c
 
 
-class Response(controller.Msg):
+class Response(HTTPMsg):
     FMT = '%s\r\n%s\r\n%s'
     def __init__(self, request, code, msg, headers, content, timestamp=None):
         self.request = request
@@ -428,27 +451,6 @@ 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 1e68fd3d..340a68d7 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -138,6 +138,31 @@ class uRequest(libpry.AutoTree):
         assert not "foo" in r.content
         assert r.headers["boo"] == ["boo"]
 
+    def test_decodeencode(self):
+        r = tutils.treq()
+        r.headers["content-encoding"] = ["identity"]
+        r.content = "falafel"
+        r.decode()
+        assert r.headers["content-encoding"] == ["identity"]
+        assert r.content == "falafel"
+
+        r = tutils.treq()
+        r.headers["content-encoding"] = ["identity"]
+        r.content = "falafel"
+        r.encode("identity")
+        assert r.headers["content-encoding"] == ["identity"]
+        assert r.content == "falafel"
+
+        r = tutils.treq()
+        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 uResponse(libpry.AutoTree):
     def test_simple(self):