simplify contentview logic

6 files changed, 127 insertions, 125 deletions

diff --git a/mitmproxy/builtins/dumper.py b/mitmproxy/builtins/dumper.py
index 59f9349d..699d4678 100644
--- a/mitmproxy/builtins/dumper.py
+++ b/mitmproxy/builtins/dumper.py
@@ -63,30 +63,12 @@ class Dumper(object):
             )
             self.echo(headers, ident=4)
         if self.flow_detail >= 3:
-            try:
-                content = message.content
-            except ValueError:
-                content = message.get_content(strict=False)
-
-            if content is None:
-                self.echo("(content missing)", ident=4)
-            elif content:
-                self.echo("")
-
-                try:
-                    _, lines = contentviews.get_content_view(
-                        contentviews.get("Auto"),
-                        content,
-                        headers=getattr(message, "headers", None)
-                    )
-                except exceptions.ContentViewException:
-                    s = "Content viewer failed: \n" + traceback.format_exc()
-                    ctx.log.debug(s)
-                    _, lines = contentviews.get_content_view(
-                        contentviews.get("Raw"),
-                        content,
-                        headers=getattr(message, "headers", None)
-                    )
+                _, lines, error = contentviews.get_message_content_view(
+                    contentviews.get("Auto"),
+                    message
+                )
+                if error:
+                    ctx.log.debug(error)
 
                 styles = dict(
                     highlight=dict(bold=True),
@@ -105,13 +87,13 @@ class Dumper(object):
                 else:
                     lines_to_echo = lines
 
-                lines_to_echo = list(lines_to_echo)
-
                 content = u"\r\n".join(
                     u"".join(colorful(line)) for line in lines_to_echo
                 )
+                if content:
+                    self.echo("")
+                    self.echo(content)
 
-                self.echo(content)
                 if next(lines, None):
                     self.echo("(cut off)", ident=4, dim=True)
 
diff --git a/mitmproxy/console/flowview.py b/mitmproxy/console/flowview.py
index d0e6bb11..5c72be09 100644
--- a/mitmproxy/console/flowview.py
+++ b/mitmproxy/console/flowview.py
@@ -206,36 +206,11 @@ class FlowView(tabs.Tabs):
             )
 
     def _get_content_view(self, message, viewmode, max_lines, _):
-
-        try:
-            content = message.content
-            if content != message.raw_content:
-                enc = "[decoded {}]".format(
-                    message.headers.get("content-encoding")
-                )
-            else:
-                enc = None
-        except ValueError:
-            content = message.raw_content
-            enc = "[cannot decode]"
-        try:
-            query = None
-            if isinstance(message, models.HTTPRequest):
-                query = message.query
-            description, lines = contentviews.get_content_view(
-                viewmode, content, headers=message.headers, query=query
-            )
-        except exceptions.ContentViewException:
-            s = "Content viewer failed: \n" + traceback.format_exc()
-            signals.add_log(s, "error")
-            description, lines = contentviews.get_content_view(
-                contentviews.get("Raw"), content, headers=message.headers
-            )
-            description = description.replace("Raw", "Couldn't parse: falling back to Raw")
-
-        if enc:
-            description = " ".join([enc, description])
-
+        description, lines, error = contentviews.get_message_content_view(
+            viewmode, message
+        )
+        if error:
+            signals.add_log(error, "error")
         # Give hint that you have to tab for the response.
         if description == "No content" and isinstance(message, models.HTTPRequest):
             description = "No request content (press tab to view response)"
diff --git a/mitmproxy/contentviews.py b/mitmproxy/contentviews.py
index dacef36d..f95efb99 100644
--- a/mitmproxy/contentviews.py
+++ b/mitmproxy/contentviews.py
@@ -14,31 +14,27 @@ requests, the query parameters are passed as the ``query`` keyword argument.
 """
 from __future__ import absolute_import, print_function, division
 
+import cssutils
 import datetime
+import html2text
+import jsbeautifier
 import json
 import logging
-import subprocess
-import sys
-
-from typing import Mapping  # noqa
-
-import html2text
 import lxml.etree
 import lxml.html
 import six
+import subprocess
+import traceback
 from PIL import ExifTags
 from PIL import Image
-from six import BytesIO
-
-import cssutils
-import jsbeautifier
-
 from mitmproxy import exceptions
 from mitmproxy.contrib.wbxml import ASCommandResponse
 from netlib import http
 from netlib import multidict
-from netlib.http import url
 from netlib import strutils
+from netlib.http import url
+from six import BytesIO
+from typing import Mapping  # noqa
 
 try:
     import pyamf
@@ -612,6 +608,39 @@ def safe_to_print(lines, encoding="utf8"):
         yield clean_line
 
 
+def get_message_content_view(viewmode, message):
+    """
+    Like get_content_view, but also handles message encoding.
+    """
+    try:
+        content = message.content
+    except ValueError:
+        content = message.raw_content
+        enc = "[cannot decode]"
+    else:
+        if isinstance(message, http.Message) and content != message.raw_content:
+            enc = "[decoded {}]".format(
+                message.headers.get("content-encoding")
+            )
+        else:
+            enc = None
+
+    if content is None:
+        return "", iter([[("error", "content missing")]]), None
+
+    query = message.query if isinstance(message, http.Request) else None
+    headers = message.headers if isinstance(message, http.Message) else None
+
+    description, lines, error = get_content_view(
+        viewmode, content, headers=headers, query=query
+    )
+
+    if enc:
+        description = "{} {}".format(enc, description)
+
+    return description, lines, error
+
+
 def get_content_view(viewmode, data, **metadata):
     """
         Args:
@@ -619,24 +648,24 @@ def get_content_view(viewmode, data, **metadata):
             data, **metadata: arguments passed to View instance.
 
         Returns:
-            A (description, content generator) tuple.
+            A (description, content generator, error) tuple.
+            If the content view raised an exception generating the view,
+            the exception is returned in error and the flow is formatted in raw mode.
             In contrast to calling the views directly, text is always safe-to-print unicode.
-
-        Raises:
-            ContentViewException, if the content view threw an error.
     """
     try:
         ret = viewmode(data, **metadata)
+        if ret is None:
+            ret = "Couldn't parse: falling back to Raw", get("Raw")(data, **metadata)[1]
+        desc, content = ret
+        error = None
     # Third-party viewers can fail in unexpected ways...
-    except Exception as e:
-        six.reraise(
-            exceptions.ContentViewException,
-            exceptions.ContentViewException(str(e)),
-            sys.exc_info()[2]
-        )
-    if not ret:
+    except Exception:
         desc = "Couldn't parse: falling back to Raw"
         _, content = get("Raw")(data, **metadata)
-    else:
-        desc, content = ret
-    return desc, safe_to_print(content)
+        error = "{} Content viewer failed: \n{}".format(
+            getattr(viewmode, "name"),
+            traceback.format_exc()
+        )
+
+    return desc, safe_to_print(content), error
diff --git a/netlib/http/__init__.py b/netlib/http/__init__.py
index af95f4d0..02a37dd3 100644
--- a/netlib/http/__init__.py
+++ b/netlib/http/__init__.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import, print_function, division
 from netlib.http.request import Request
 from netlib.http.response import Response
+from netlib.http.message import Message
 from netlib.http.headers import Headers, parse_content_type
 from netlib.http.message import decoded
 from netlib.http import http1, http2, status_codes, multipart
diff --git a/test/mitmproxy/builtins/test_dumper.py b/test/mitmproxy/builtins/test_dumper.py
index 6287fe86..1c7173e0 100644
--- a/test/mitmproxy/builtins/test_dumper.py
+++ b/test/mitmproxy/builtins/test_dumper.py
@@ -15,7 +15,7 @@ class TestDumper(mastertest.MasterTest):
         d = dumper.Dumper()
         sio = StringIO()
 
-        updated = set(["tfile", "flow_detail"])
+        updated = {"tfile", "flow_detail"}
         d.configure(dump.Options(tfile = sio, flow_detail = 0), updated)
         d.response(tutils.tflow())
         assert not sio.getvalue()
@@ -66,10 +66,9 @@ class TestDumper(mastertest.MasterTest):
 
 
 class TestContentView(mastertest.MasterTest):
-    @mock.patch("mitmproxy.contentviews.get_content_view")
-    def test_contentview(self, get_content_view):
-        se = exceptions.ContentViewException(""), ("x", iter([]))
-        get_content_view.side_effect = se
+    @mock.patch("mitmproxy.contentviews.ViewAuto.__call__")
+    def test_contentview(self, view_auto):
+        view_auto.side_effect = exceptions.ContentViewException("")
 
         s = state.State()
         sio = StringIO()
diff --git a/test/mitmproxy/test_contentview.py b/test/mitmproxy/test_contentview.py
index 66cad47b..f0afdc0b 100644
--- a/test/mitmproxy/test_contentview.py
+++ b/test/mitmproxy/test_contentview.py
@@ -1,3 +1,4 @@
+import mock
 from mitmproxy.exceptions import ContentViewException
 from netlib.http import Headers
 from netlib.http import url
@@ -5,6 +6,7 @@ from netlib import multidict
 
 import mitmproxy.contentviews as cv
 from . import tutils
+import netlib.tutils
 
 try:
     import pyamf
@@ -180,43 +182,6 @@ Larry
         assert f[0] == "Query"
         assert [x for x in f[1]] == [[("header", "foo: "), ("text", "bar")]]
 
-    def test_get_content_view(self):
-        r = cv.get_content_view(
-            cv.get("Raw"),
-            b"[1, 2, 3]",
-            headers=Headers(content_type="application/json")
-        )
-        assert "Raw" in r[0]
-
-        r = cv.get_content_view(
-            cv.get("Auto"),
-            b"[1, 2, 3]",
-            headers=Headers(content_type="application/json")
-        )
-        assert r[0] == "JSON"
-
-        r = cv.get_content_view(
-            cv.get("Auto"),
-            b"[1, 2",
-            headers=Headers(content_type="application/json")
-        )
-        assert "Raw" in r[0]
-
-        r = cv.get_content_view(
-            cv.get("Auto"),
-            b"[1, 2, 3]",
-            headers=Headers(content_type="application/vnd.api+json")
-        )
-        assert r[0] == "JSON"
-
-        tutils.raises(
-            ContentViewException,
-            cv.get_content_view,
-            cv.get("AMF"),
-            b"[1, 2",
-            headers=Headers()
-        )
-
     def test_add_cv(self):
         class TestContentView(cv.View):
             name = "test"
@@ -233,6 +198,57 @@ Larry
         )
 
 
+def test_get_content_view():
+    desc, lines, err = cv.get_content_view(
+        cv.get("Raw"),
+        b"[1, 2, 3]",
+    )
+    assert "Raw" in desc
+    assert list(lines)
+    assert not err
+
+    desc, lines, err = cv.get_content_view(
+        cv.get("Auto"),
+        b"[1, 2, 3]",
+        headers=Headers(content_type="application/json")
+    )
+    assert desc == "JSON"
+
+    desc, lines, err = cv.get_content_view(
+        cv.get("JSON"),
+        b"[1, 2",
+    )
+    assert "Couldn't parse" in desc
+
+    with mock.patch("mitmproxy.contentviews.ViewAuto.__call__") as view_auto:
+        view_auto.side_effect = ValueError
+
+        desc, lines, err = cv.get_content_view(
+            cv.get("JSON"),
+            b"[1, 2",
+        )
+        assert err
+        assert "Couldn't parse" in desc
+
+
+def test_get_message_content_view():
+    r = netlib.tutils.treq()
+    desc, lines, err = cv.get_message_content_view(cv.get("Raw"), r)
+    assert desc == "Raw"
+
+    r.encode("gzip")
+    desc, lines, err = cv.get_message_content_view(cv.get("Raw"), r)
+    assert desc == "[decoded gzip] Raw"
+
+    r.headers["content-encoding"] = "deflate"
+    desc, lines, err = cv.get_message_content_view(cv.get("Raw"), r)
+    assert desc == "[cannot decode] Raw"
+
+    r.content = None
+    desc, lines, err = cv.get_message_content_view(cv.get("Raw"), r)
+    assert list(lines) == [[("error", "content missing")]]
+
+
 if pyamf:
     def test_view_amf_request():
         v = cv.ViewAMF()