From 6792ec40587bde8dbd7fac67c35038afc126e80b Mon Sep 17 00:00:00 2001
From: Angelo Agatino Nicolosi <angelonicolosi@yanchware.com>
Date: Sat, 30 Jul 2016 14:43:53 +0200
Subject: Integrated encode/decoder for brotli

---
 mitmproxy/console/flowview.py |  2 ++
 netlib/encoding.py            | 15 +++++++++++++--
 netlib/http/message.py        |  2 +-
 netlib/http/request.py        |  2 +-
 setup.py                      |  1 +
 test/netlib/test_encoding.py  | 12 ++++++++++++
 6 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/mitmproxy/console/flowview.py b/mitmproxy/console/flowview.py
index c354563f..4cde955a 100644
--- a/mitmproxy/console/flowview.py
+++ b/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)
diff --git a/netlib/encoding.py b/netlib/encoding.py
index da282194..9c8acff7 100644
--- a/netlib/encoding.py
+++ b/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"]
diff --git a/netlib/http/message.py b/netlib/http/message.py
index be35b8d1..ce92bab1 100644
--- a/netlib/http/message.py
+++ b/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.
 
diff --git a/netlib/http/request.py b/netlib/http/request.py
index 061217a3..d59fead4 100644
--- a/netlib/http/request.py
+++ b/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
                 )
             )
diff --git a/setup.py b/setup.py
index a1580d8b..3915e421 100644
--- a/setup.py
+++ b/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"': [
diff --git a/test/netlib/test_encoding.py b/test/netlib/test_encoding.py
index a5e81379..08e69ec5 100644
--- a/test/netlib/test_encoding.py
+++ b/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(
-- 
cgit v1.2.3