Integrated encode/decoder for brotli

mitmproxy/console/flowview.py

@@ -713,6 +713,7 @@ class FlowView(tabs.Tabs):
                     keys = (
                         ("gzip", "z"),
                         ("deflate", "d"),
+                        ("brotli", "b"),
                     ),
                     callback = self.encode_callback,
                     args = (conn,)
@@ -726,6 +727,7 @@ class FlowView(tabs.Tabs):
         encoding_map = {
             "z": "gzip",
             "d": "deflate",
+            "b": "brotli",
         }
         conn.encode(encoding_map[key])
         signals.flow_change.send(self, flow = self.flow)

netlib/encoding.py

@@ -8,6 +8,7 @@ import collections
 from io import BytesIO
 import gzip
 import zlib
+import brotli
 
 from typing import Union  # noqa
 
@@ -45,7 +46,7 @@ def decode(encoded, encoding, errors='strict'):
             decoded = custom_decode[encoding](encoded)
         except KeyError:
             decoded = codecs.decode(encoded, encoding, errors)
-        if encoding in ("gzip", "deflate"):
+        if encoding in ("gzip", "deflate", "br"):
             _cache = CachedDecode(encoded, encoding, errors, decoded)
         return decoded
     except Exception as e:
@@ -81,7 +82,7 @@ def encode(decoded, encoding, errors='strict'):
             encoded = custom_encode[encoding](decoded)
         except KeyError:
             encoded = codecs.encode(decoded, encoding, errors)
-        if encoding in ("gzip", "deflate"):
+        if encoding in ("gzip", "deflate", "br"):
             _cache = CachedDecode(encoded, encoding, errors, decoded)
         return encoded
     except Exception as e:
@@ -113,6 +114,14 @@ def encode_gzip(content):
     return s.getvalue()
 
 
+def decode_brotli(content):
+    return brotli.decompress(content)
+
+
+def encode_brotli(content):
+    return brotli.compress(content)
+
+
 def decode_deflate(content):
     """
         Returns decompressed data for DEFLATE. Some servers may respond with
@@ -139,11 +148,13 @@ custom_decode = {
     "identity": identity,
     "gzip": decode_gzip,
     "deflate": decode_deflate,
+    "br": decode_brotli,
 }
 custom_encode = {
     "identity": identity,
     "gzip": encode_gzip,
     "deflate": encode_deflate,
+    "br": encode_brotli,
 }
 
 __all__ = ["encode", "decode"]

netlib/http/message.py

@@ -248,7 +248,7 @@ class Message(basetypes.Serializable):
 
     def encode(self, e):
         """
-        Encodes body with the encoding e, where e is "gzip", "deflate" or "identity".
+        Encodes body with the encoding e, where e is "gzip", "deflate", "identity", or "br".
         Any existing content-encodings are overwritten,
         the content is not decoded beforehand.
 

netlib/http/request.py

@@ -337,7 +337,7 @@ class Request(message.Message):
             self.headers["accept-encoding"] = (
                 ', '.join(
                     e
-                    for e in {"gzip", "identity", "deflate"}
+                    for e in {"gzip", "identity", "deflate", "br"}
                     if e in accept_encoding
                 )
             )

setup.py

@@ -85,6 +85,7 @@ setup(
         "tornado>=4.3, <4.5",
         "urwid>=1.3.1, <1.4",
         "watchdog>=0.8.3, <0.9",
+        "brotlipy>=0.3.0, <0.4",
     ],
     extras_require={
         ':sys_platform == "win32"': [

test/netlib/test_encoding.py

@@ -21,6 +21,18 @@ def test_gzip():
         encoding.decode(b"bogus", "gzip")
 
 
+def test_brotli():
+    assert b"string" == encoding.decode(
+        encoding.encode(
+            b"string",
+            "br"
+        ),
+        "br"
+    )
+    with tutils.raises(ValueError):
+        encoding.decode(b"bogus", "br")
+
+
 def test_deflate():
     assert b"string" == encoding.decode(
         encoding.encode(