From 951885a5dd2f1dd72a67390caa1a07f10f24c8c2 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Tue, 2 Aug 2016 20:36:19 -0700 Subject: [PATCH] simplify contentview logic --- mitmproxy/builtins/dumper.py | 36 +++-------- mitmproxy/console/flowview.py | 35 ++-------- mitmproxy/contentviews.py | 81 +++++++++++++++-------- netlib/http/__init__.py | 1 + test/mitmproxy/builtins/test_dumper.py | 9 ++- test/mitmproxy/test_contentview.py | 90 +++++++++++++++----------- 6 files changed, 127 insertions(+), 125 deletions(-) diff --git a/mitmproxy/builtins/dumper.py b/mitmproxy/builtins/dumper.py index 59f9349d7..699d4678d 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 d0e6bb114..5c72be091 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 dacef36db..f95efb99d 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 af95f4d09..02a37dd37 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 6287fe861..1c7173e03 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 66cad47bd..f0afdc0bc 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()