Merge pull request #9 from alts/toggle_encodings

Toggle encodings

4 files changed, 68 insertions, 31 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 3eef4611..f74fffde 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -40,7 +40,7 @@ def highlight_key(s, k):
 
 def format_keyvals(lst, key="key", val="text", space=5, indent=0):
     """
-        Format a list of (key, value) tuples. 
+        Format a list of (key, value) tuples.
 
         If key is None, it's treated specially:
             - We assume a sub-value, and add an extra indent.
@@ -311,8 +311,9 @@ class ConnectionView(WWrap):
     def _conn_text(self, conn, viewmode):
         if conn:
             e = conn.headers["content-encoding"]
-            if e:
+            if e and conn.should_autodecode:
                 e = e[0]
+                conn.should_autodecode = False
             else:
                 e = "identity"
             return self.master._cached_conn_text(
@@ -530,6 +531,20 @@ class ConnectionView(WWrap):
             self.master.view_next_flow(self.flow)
         elif key == "|":
             self.master.path_prompt("Script: ", self.state.last_script, self.run_script)
+        elif key == "g":
+            if self.state.view_flow_mode == VIEW_FLOW_RESPONSE:
+                conn = self.flow.response
+                e = conn.headers["content-encoding"]
+                if e:
+                    if conn.last_encoding:
+                        conn.content = encoding.encode(
+                            conn.last_encoding,
+                            encoding.decode(e[0], conn.content)
+                        )
+                        conn.last_encoding, conn.headers["content-encoding"] = e[0], [conn.last_encoding]
+                    else:
+                        conn.last_encoding = "identity"
+                    self.master.refresh_connection(self.flow)
         return key
 
     def run_script(self, path):
@@ -1245,37 +1260,37 @@ class ConsoleMaster(flow.FlowMaster):
             ("L", "load saved flows"),
 
             ("m", "change body display mode"),
-            (None, 
+            (None,
                 highlight_key("raw", "r") +
                 [("text", ": raw data")]
             ),
-            (None, 
+            (None,
                 highlight_key("pretty", "p") +
                 [("text", ": pretty-print XML, HTML and JSON")]
             ),
-            (None, 
+            (None,
                 highlight_key("hex", "h") +
                 [("text", ": hex dump")]
             ),
 
             ("o", "toggle options:"),
-            (None, 
+            (None,
                 highlight_key("anticache", "a") +
                 [
-                    
+
                     ("text", ": modify requests to prevent cached responses")
-                    
+
                 ]
             ),
-            (None, 
+            (None,
                 highlight_key("anticomp", "c") +
                 [("text", ": modify requests to try to prevent compressed responses")]
             ),
-            (None, 
+            (None,
                 highlight_key("killextra", "k") +
                 [("text", ": kill requests not part of server replay")]
             ),
-            (None, 
+            (None,
                 highlight_key("norefresh", "n") +
                 [("text", ": disable server replay response refresh")]
             ),
@@ -1307,6 +1322,7 @@ class ConsoleMaster(flow.FlowMaster):
         keys = [
             ("b", "save request/response body"),
             ("e", "edit request/response"),
+            ("g", "switch response encoding"),
             ("p", "previous flow"),
             ("v", "view body in external viewer"),
             ("|", "run script"),
diff --git a/libmproxy/encoding.py b/libmproxy/encoding.py
index b56c03a6..aab4141b 100644
--- a/libmproxy/encoding.py
+++ b/libmproxy/encoding.py
@@ -10,13 +10,21 @@ ENCODINGS = set(["identity", "gzip", "deflate"])
 
 def decode(encoding, content):
     encoding_map = {
-        "identity": decode_identity,
+        "identity": identity,
         "gzip": decode_gzip,
         "deflate": decode_deflate,
     }
-    return encoding_map.get(encoding, decode_identity)(content)
+    return encoding_map.get(encoding, identity)(content)
 
-def decode_identity(content):
+def encode(encoding, content):
+    encoding_map = {
+        "identity": identity,
+        "gzip": encode_gzip,
+        "deflate": encode_deflate,
+    }
+    return encoding_map.get(encoding, identity)(content)
+
+def identity(content):
     """
         Returns content unchanged. Identity is the default value of
         Accept-Encoding headers.
@@ -30,9 +38,16 @@ def decode_gzip(content):
     except IOError:
         return None
 
+def encode_gzip(content):
+    s = cStringIO.StringIO()
+    gf = gzip.GzipFile(fileobj=s, mode='wb')
+    gf.write(content)
+    gf.close()
+    return s.getvalue()
+
 def decode_deflate(content):
     """
-        Returns decompress data for DEFLATE. Some servers may respond with
+        Returns decompressed data for DEFLATE. Some servers may respond with
         compressed data without a zlib header or checksum. An undocumented
         feature of zlib permits the lenient decompression of data missing both
         values.
@@ -46,3 +61,9 @@ def decode_deflate(content):
             return zlib.decompress(content, -15)
     except zlib.error:
         return None
+
+def encode_deflate(content):
+    """
+        Returns compressed content, always including zlib header and checksum.
+    """
+    return zlib.compress(content)
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index a7cc31e8..4ff96c4f 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -290,6 +290,8 @@ class Response(controller.Msg):
         self.timestamp = timestamp or utils.timestamp()
         controller.Msg.__init__(self)
         self.replay = False
+        self.last_encoding = None
+        self.should_autodecode = True
 
     def _refresh_cookie(self, c, delta):
         """
diff --git a/test/test_encoding.py b/test/test_encoding.py
index 00632c0e..dff146ee 100644
--- a/test/test_encoding.py
+++ b/test/test_encoding.py
@@ -4,30 +4,28 @@ import libpry
 import cStringIO
 import gzip, zlib
 
-class udecode_identity(libpry.AutoTree):
-    def test_decode(self):
-        assert 'string' == encoding.decode('identity', 'string')
+class uidentity(libpry.AutoTree):
+    def test_simple(self):
+        assert "string" == encoding.decode("identity", "string")
+        assert "string" == encoding.encode("identity", "string")
 
     def test_fallthrough(self):
-        assert 'string' == encoding.decode('nonexistent encoding', 'string')
+        assert "string" == encoding.decode("nonexistent encoding", "string")
+        assert "string" == encoding.encode("nonexistent encoding", "string")
 
-class udecode_gzip(libpry.AutoTree):
+class ugzip(libpry.AutoTree):
     def test_simple(self):
-        s = cStringIO.StringIO()
-        gf = gzip.GzipFile(fileobj=s, mode='wb')
-        gf.write('string')
-        gf.close()
-        assert 'string' == encoding.decode('gzip', s.getvalue())
+        assert "string" == encoding.decode("gzip", encoding.encode("gzip", "string"))
         assert None == encoding.decode("gzip", "bogus")
 
-class udecode_deflate(libpry.AutoTree):
+class udeflate(libpry.AutoTree):
     def test_simple(self):
-        assert 'string' == encoding.decode('deflate', zlib.compress('string'))
-        assert 'string' == encoding.decode('deflate', zlib.compress('string')[2:-4])
+        assert "string" == encoding.decode("deflate", encoding.encode("deflate", "string"))
+        assert "string" == encoding.decode("deflate", encoding.encode("deflate", "string")[2:-4])
         assert None == encoding.decode("deflate", "bogus")
 
 tests = [
-    udecode_identity(),
-    udecode_gzip(),
-    udecode_deflate()
+    uidentity(),
+    ugzip(),
+    udeflate()
 ]