Merge pull request #1402 from mhils/console-fixes

minor fixes

4 files changed, 27 insertions, 23 deletions

diff --git a/mitmproxy/console/flowview.py b/mitmproxy/console/flowview.py
index 789066fc..11c72151 100644
--- a/mitmproxy/console/flowview.py
+++ b/mitmproxy/console/flowview.py
@@ -6,6 +6,7 @@ import sys
 import traceback
 
 import urwid
+from typing import Optional, Union  # noqa
 
 from mitmproxy import contentviews
 from mitmproxy import controller
@@ -105,7 +106,8 @@ footer = [
 class FlowViewHeader(urwid.WidgetWrap):
 
     def __init__(self, master, f):
-        self.master, self.flow = master, f
+        self.master = master  # type: "mitmproxy.console.master.ConsoleMaster"
+        self.flow = f  # type: models.HTTPFlow
         self._w = common.format_flow(
             f,
             False,
@@ -530,13 +532,6 @@ class FlowView(tabs.Tabs):
         )
         signals.flow_change.send(self, flow = self.flow)
 
-    def delete_body(self, t):
-        if self.tab_offset == TAB_REQ:
-            self.flow.request.content = None
-        else:
-            self.flow.response.content = None
-        signals.flow_change.send(self, flow = self.flow)
-
     def keypress(self, size, key):
         key = super(self.__class__, self).keypress(size, key)
 
@@ -545,6 +540,8 @@ class FlowView(tabs.Tabs):
             return
 
         key = common.shortcuts(key)
+
+        conn = None  # type: Optional[Union[models.HTTPRequest, models.HTTPResponse]]
         if self.tab_offset == TAB_REQ:
             conn = self.flow.request
         elif self.tab_offset == TAB_RESP:
@@ -691,15 +688,8 @@ class FlowView(tabs.Tabs):
                     args = (scope, self.flow, common.copy_to_clipboard_or_prompt)
                 )
             elif key == "x":
-                signals.status_prompt_onekey.send(
-                    prompt = "Delete body",
-                    keys = (
-                        ("completely", "c"),
-                        ("mark as missing", "m"),
-                    ),
-                    callback = self.delete_body
-                )
-                key = None
+                conn.content = None
+                signals.flow_change.send(self, flow=self.flow)
             elif key == "v":
                 if conn.raw_content:
                     t = conn.headers.get("content-type")
@@ -713,7 +703,9 @@ class FlowView(tabs.Tabs):
                 self.flow.backup()
                 e = conn.headers.get("content-encoding", "identity")
                 if e != "identity":
-                    if not conn.decode():
+                    try:
+                        conn.decode()
+                    except ValueError:
                         signals.status_message.send(
                             message = "Could not decode - invalid data?"
                         )
diff --git a/mitmproxy/contentviews.py b/mitmproxy/contentviews.py
index afdaad7f..e155bc01 100644
--- a/mitmproxy/contentviews.py
+++ b/mitmproxy/contentviews.py
@@ -20,6 +20,8 @@ import logging
 import subprocess
 import sys
 
+from typing import Mapping  # noqa
+
 import html2text
 import lxml.etree
 import lxml.html
@@ -76,6 +78,7 @@ def pretty_json(s):
 
 
 def format_dict(d):
+    # type: (Mapping[Union[str,bytes], Union[str,bytes]]) -> Generator[Tuple[Union[str,bytes], Union[str,bytes]]]
     """
     Helper function that transforms the given dictionary into a list of
         ("key",   key  )
@@ -85,7 +88,7 @@ def format_dict(d):
     max_key_len = max(len(k) for k in d.keys())
     max_key_len = min(max_key_len, KEY_MAX)
     for key, value in d.items():
-        key += ":"
+        key += b":" if isinstance(key, bytes) else u":"
         key = key.ljust(max_key_len + 2)
         yield [
             ("header", key),
@@ -106,12 +109,16 @@ class View(object):
     prompt = ()
     content_types = []
 
-    def __call__(self, data, **metadata):
+    def __call__(
+            self,
+            data,  # type: bytes
+            **metadata
+    ):
         """
         Transform raw data into human-readable output.
 
         Args:
-            data: the data to decode/format as bytes.
+            data: the data to decode/format.
             metadata: optional keyword-only arguments for metadata. Implementations must not
                 rely on a given argument being present.
 
@@ -278,6 +285,10 @@ class ViewURLEncoded(View):
     content_types = ["application/x-www-form-urlencoded"]
 
     def __call__(self, data, **metadata):
+        try:
+            data = data.decode("ascii", "strict")
+        except ValueError:
+            return None
         d = url.decode(data)
         return "URLEncoded form", format_dict(multidict.MultiDict(d))
 
diff --git a/netlib/http/url.py b/netlib/http/url.py
index 2fc6e7ee..1c8c007a 100644
--- a/netlib/http/url.py
+++ b/netlib/http/url.py
@@ -82,6 +82,7 @@ def unparse(scheme, host, port, path=""):
 
 
 def encode(s):
+    # type: (six.text_type, bytes) -> str
     """
         Takes a list of (key, value) tuples and returns a urlencoded string.
     """
diff --git a/test/mitmproxy/test_contentview.py b/test/mitmproxy/test_contentview.py
index 2db9ab40..aad53b37 100644
--- a/test/mitmproxy/test_contentview.py
+++ b/test/mitmproxy/test_contentview.py
@@ -59,10 +59,10 @@ class TestContentView:
         assert f[0] == "Query"
 
     def test_view_urlencoded(self):
-        d = url.encode([("one", "two"), ("three", "four")])
+        d = url.encode([("one", "two"), ("three", "four")]).encode()
         v = cv.ViewURLEncoded()
         assert v(d)
-        d = url.encode([("adsfa", "")])
+        d = url.encode([("adsfa", "")]).encode()
         v = cv.ViewURLEncoded()
         assert v(d)