mitmdump: colorize output, add content views

3 files changed, 192 insertions, 81 deletions

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 17b47dd2..dd44dc69 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -1,14 +1,16 @@
 from __future__ import absolute_import, print_function
-import json
 import sys
 import os
+import traceback
 
-from netlib.http.semantics import CONTENT_MISSING
-import netlib.utils
+import click
 
-from . import flow, filt, utils
-from .protocol import http
 
+from netlib.http.semantics import CONTENT_MISSING
+import netlib.utils
+from . import flow, filt, contentview
+from .exceptions import ContentViewException
+from .models import HTTPRequest
 
 class DumpError(Exception):
     pass
@@ -55,24 +57,6 @@ class Options(object):
                 setattr(self, i, None)
 
 
-def str_response(resp):
-    r = "%s %s" % (resp.code, resp.msg)
-    if resp.is_replay:
-        r = "[replay] " + r
-    return r
-
-
-def str_request(f, showhost):
-    if f.client_conn:
-        c = f.client_conn.address.host
-    else:
-        c = "[replay]"
-    r = "%s %s %s" % (c, f.request.method, f.request.pretty_url(showhost))
-    if f.request.stickycookie:
-        r = "[stickycookie] " + r
-    return r
-
-
 class DumpMaster(flow.FlowMaster):
     def __init__(self, server, options, outfile=sys.stdout):
         flow.FlowMaster.__init__(self, server, flow.State())
@@ -163,73 +147,162 @@ class DumpMaster(flow.FlowMaster):
     def add_event(self, e, level="info"):
         needed = dict(error=0, info=1, debug=2).get(level, 1)
         if self.o.verbosity >= needed:
-            print(e, file=self.outfile)
-            self.outfile.flush()
+            self.echo(
+                e,
+                fg="red" if level == "error" else None,
+                dim=(level == "debug")
+            )
 
     @staticmethod
-    def indent(n, t):
-        l = str(t).strip().splitlines()
+    def indent(n, text):
+        l = str(text).strip().splitlines()
         pad = " " * n
         return "\n".join(pad + i for i in l)
 
-    def _print_message(self, message):
+    def echo(self, text, indent=None, **style):
+        if indent:
+            text = self.indent(indent, text)
+        click.secho(text, file=self.outfile, **style)
+
+    def _echo_message(self, message):
         if self.o.flow_detail >= 2:
-            print(self.indent(4, str(message.headers)), file=self.outfile)
+            headers = "\r\n".join(
+                "{}: {}".format(
+                    click.style(k, fg="blue", bold=True),
+                    click.style(v, fg="blue"))
+                    for k, v in message.headers.fields
+            )
+            self.echo(headers, indent=4)
         if self.o.flow_detail >= 3:
             if message.content == CONTENT_MISSING:
-                print(self.indent(4, "(content missing)"), file=self.outfile)
+                self.echo("(content missing)", indent=4)
             elif message.content:
-                print("", file=self.outfile)
-                content = message.get_decoded_content()
-                if not utils.isBin(content):
-                    try:
-                        jsn = json.loads(content)
-                        print(
-                            self.indent(
-                                4,
-                                json.dumps(
-                                    jsn,
-                                    indent=2)),
-                            file=self.outfile)
-                    except ValueError:
-                        print(self.indent(4, content), file=self.outfile)
-                else:
-                    d = netlib.utils.hexdump(content)
-                    d = "\n".join("%s\t%s %s" % i for i in d)
-                    print(self.indent(4, d), file=self.outfile)
+                self.echo("")
+                cutoff = sys.maxsize if self.o.flow_detail >= 4 else contentview.VIEW_CUTOFF
+                try:
+                    type, lines = contentview.get_content_view(
+                        contentview.get("Auto"),
+                        message.headers,
+                        message.body,
+                        cutoff,
+                        isinstance(message, HTTPRequest)
+                    )
+                except ContentViewException:
+                    s = "Content viewer failed: \n" + traceback.format_exc()
+                    self.add_event(s, "debug")
+                    type, lines = contentview.get_content_view(
+                        contentview.get("Raw"),
+                        message.headers,
+                        message.body,
+                        cutoff,
+                        isinstance(message, HTTPRequest)
+                    )
+
+                styles = dict(
+                    highlight=dict(bold=True),
+                    offset=dict(fg="blue"),
+                    header=dict(fg="green", bold=True),
+                    text=dict(fg="green")
+                )
+
+                def colorful(line):
+                    yield "    "  # we can already indent here
+                    for (style, text) in line:
+                        yield click.style(text, **styles.get(style, {}))
+
+                content = "\r\n".join(
+                    "".join(colorful(line)) for line in lines
+                )
+                self.echo(content)
+
         if self.o.flow_detail >= 2:
-            print("", file=self.outfile)
+            self.echo("")
 
-    def _process_flow(self, f):
-        self.state.delete_flow(f)
-        if self.filt and not f.match(self.filt):
-            return
+    def _echo_request_line(self, flow):
+        if flow.request.stickycookie:
+            stickycookie = click.style("[stickycookie] ", fg="yellow", bold=True)
+        else:
+            stickycookie = ""
+
+        if flow.client_conn:
+            client = click.style(flow.client_conn.address.host, bold=True)
+        else:
+            client = click.style("[replay]", fg="yellow", bold=True)
+
+        method = flow.request.method
+        method_color=dict(
+            GET="green",
+            DELETE="red"
+        ).get(method.upper(), "magenta")
+        method = click.style(method, fg=method_color, bold=True)
+        url = click.style(flow.request.pretty_url(self.showhost), bold=True)
+
+        line = "{stickycookie}{client} {method} {url}".format(
+            stickycookie=stickycookie,
+            client=client,
+            method=method,
+            url=url
+        )
+        self.echo(line)
+
+    def _echo_response_line(self, flow):
+        if flow.response.is_replay:
+            replay = click.style("[replay] ", fg="yellow", bold=True)
+        else:
+            replay = ""
+
+        code = flow.response.status_code
+        code_color = None
+        if 200 <= code < 300:
+            code_color = "green"
+        elif 300 <= code < 400:
+            code_color = "magenta"
+        elif 400 <= code < 600:
+            code_color = "red"
+        code = click.style(str(code), fg=code_color, bold=True, blink=(code == 418))
+        msg = click.style(flow.response.msg, fg=code_color, bold=True)
+
+        if flow.response.content == CONTENT_MISSING:
+            size = "(content missing)"
+        else:
+            size = netlib.utils.pretty_size(len(flow.response.content))
+        size = click.style(size, bold=True)
+
+        arrows = click.style("<<", bold=True)
 
+        line = "{replay} {arrows} {code} {msg} {size}".format(
+            replay=replay,
+            arrows=arrows,
+            code=code,
+            msg=msg,
+            size=size
+        )
+        self.echo(line)
+
+    def echo_flow(self, f):
         if self.o.flow_detail == 0:
             return
 
         if f.request:
-            print(str_request(f, self.showhost), file=self.outfile)
-            self._print_message(f.request)
+            self._echo_request_line(f)
+            self._echo_message(f.request)
 
         if f.response:
-            if f.response.content == CONTENT_MISSING:
-                sz = "(content missing)"
-            else:
-                sz = netlib.utils.pretty_size(len(f.response.content))
-            print(
-                " << %s %s" %
-                (str_response(
-                    f.response),
-                    sz),
-                file=self.outfile)
-            self._print_message(f.response)
+            self._echo_response_line(f)
+            self._echo_message(f.response)
 
         if f.error:
-            print(" << {}".format(f.error.msg), file=self.outfile)
+            self.echo(" << {}".format(f.error.msg), bold=True, fg="red")
 
         self.outfile.flush()
 
+    def _process_flow(self, f):
+        self.state.delete_flow(f)
+        if self.filt and not f.match(self.filt):
+            return
+
+        self.echo_flow(f)
+
     def handle_request(self, f):
         flow.FlowMaster.handle_request(self, f)
         if f:
diff --git a/setup.py b/setup.py
index 896d0248..50d4604f 100644
--- a/setup.py
+++ b/setup.py
@@ -23,15 +23,17 @@ deps = {
     "html2text>=2015.4.14",
     "construct>=2.5.2",
     "six>=1.9.0",
+    "lxml>=3.3.6",
+    "Pillow>=2.3.0",
 }
 # A script -> additional dependencies dict.
 scripts = {
     "mitmproxy": {
         "urwid>=1.3",
-        "lxml>=3.3.6",
-        "Pillow>=2.3.0",
     },
-    "mitmdump": set(),
+    "mitmdump": {
+        "click>=5.1",
+    },
     "mitmweb": set()
 }
 # Developer dependencies
diff --git a/test/test_dump.py b/test/test_dump.py
index c76f555f..0fc4cd4d 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -1,5 +1,7 @@
 import os
 from cStringIO import StringIO
+from libmproxy.contentview import ViewAuto
+from libmproxy.exceptions import ContentViewException
 from libmproxy.models import HTTPResponse
 
 import netlib.tutils
@@ -12,17 +14,51 @@ import mock
 
 
 def test_strfuncs():
-    t = HTTPResponse.wrap(netlib.tutils.tresp())
-    t.is_replay = True
-    dump.str_response(t)
-
-    f = tutils.tflow()
-    f.client_conn = None
-    f.request.stickycookie = True
-    assert "stickycookie" in dump.str_request(f, False)
-    assert "stickycookie" in dump.str_request(f, True)
-    assert "replay" in dump.str_request(f, False)
-    assert "replay" in dump.str_request(f, True)
+    o = dump.Options()
+    m = dump.DumpMaster(None, o)
+
+    m.outfile = StringIO()
+    m.o.flow_detail = 0
+    m.echo_flow(tutils.tflow())
+    assert not m.outfile.getvalue()
+
+    m.o.flow_detail = 4
+    m.echo_flow(tutils.tflow())
+    assert m.outfile.getvalue()
+
+    m.outfile = StringIO()
+    m.echo_flow(tutils.tflow(resp=True))
+    assert "<<" in m.outfile.getvalue()
+
+    m.outfile = StringIO()
+    m.echo_flow(tutils.tflow(err=True))
+    assert "<<" in m.outfile.getvalue()
+
+    flow = tutils.tflow()
+    flow.request = netlib.tutils.treq()
+    flow.request.stickycookie = True
+    flow.client_conn = mock.MagicMock()
+    flow.client_conn.address.host = "foo"
+    flow.response = netlib.tutils.tresp(content=CONTENT_MISSING)
+    flow.response.is_replay = True
+    flow.response.code = 300
+    m.echo_flow(flow)
+
+
+    flow = tutils.tflow(resp=netlib.tutils.tresp("{"))
+    flow.response.headers["content-type"] = "application/json"
+    flow.response.code = 400
+    m.echo_flow(flow)
+
+
+@mock.patch("libmproxy.contentview.get_content_view")
+def test_contentview(get_content_view):
+    get_content_view.side_effect = ContentViewException(""), ("x", [])
+
+    o = dump.Options(flow_detail=4, verbosity=3)
+    m = dump.DumpMaster(None, o, StringIO())
+    m.echo_flow(tutils.tflow())
+    assert "Content viewer failed" in m.outfile.getvalue()
 
 
 class TestDumpMaster: