image view: add fallback, catch all images but svgs

7 files changed, 49 insertions, 25 deletions

diff --git a/mitmproxy/contentviews/auto.py b/mitmproxy/contentviews/auto.py
index 7b3cbd78..d46a1bd3 100644
--- a/mitmproxy/contentviews/auto.py
+++ b/mitmproxy/contentviews/auto.py
@@ -18,6 +18,8 @@ class ViewAuto(base.View):
                 return contentviews.content_types_map[ct][0](data, **metadata)
             elif strutils.is_xml(data):
                 return contentviews.get("XML/HTML")(data, **metadata)
+            elif ct.startswith("image/"):
+                return contentviews.get("Image")(data, **metadata)
         if metadata.get("query"):
             return contentviews.get("Query")(data, **metadata)
         if data and strutils.is_mostly_bin(data):
diff --git a/mitmproxy/contentviews/image/__init__.py b/mitmproxy/contentviews/image/__init__.py
index 0d0f06e0..33356bd7 100644
--- a/mitmproxy/contentviews/image/__init__.py
+++ b/mitmproxy/contentviews/image/__init__.py
@@ -1 +1,3 @@
-from .view import ViewImage  # noqa
+from .view import ViewImage
+
+__all__ = ["ViewImage"]
diff --git a/mitmproxy/contentviews/image/image_parser.py b/mitmproxy/contentviews/image/image_parser.py
index 1ff3cff7..062fb38e 100644
--- a/mitmproxy/contentviews/image/image_parser.py
+++ b/mitmproxy/contentviews/image/image_parser.py
@@ -13,9 +13,9 @@ Metadata = typing.List[typing.Tuple[str, str]]
 def parse_png(data: bytes) -> Metadata:
     img = png.Png(KaitaiStream(io.BytesIO(data)))
     parts = [
-        ('Format', 'Portable network graphics')
+        ('Format', 'Portable network graphics'),
+        ('Size', "{0} x {1} px".format(img.ihdr.width, img.ihdr.height))
     ]
-    parts.append(('Size', "{0} x {1} px".format(img.ihdr.width, img.ihdr.height)))
     for chunk in img.chunks:
         if chunk.type == 'gAMA':
             parts.append(('gamma', str(chunk.body.gamma_int / 100000)))
@@ -34,13 +34,13 @@ def parse_png(data: bytes) -> Metadata:
 
 def parse_gif(data: bytes) -> Metadata:
     img = gif.Gif(KaitaiStream(io.BytesIO(data)))
+    descriptor = img.logical_screen_descriptor
     parts = [
-        ('Format', 'Compuserve GIF')
+        ('Format', 'Compuserve GIF'),
+        ('Version', "GIF{}".format(img.header.version.decode('ASCII'))),
+        ('Size', "{} x {} px".format(descriptor.screen_width, descriptor.screen_height)),
+        ('background', str(descriptor.bg_color_index))
     ]
-    parts.append(('version', "GIF{0}".format(img.header.version.decode('ASCII'))))
-    descriptor = img.logical_screen_descriptor
-    parts.append(('Size', "{0} x {1} px".format(descriptor.screen_width, descriptor.screen_height)))
-    parts.append(('background', str(descriptor.bg_color_index)))
     ext_blocks = []
     for block in img.blocks:
         if block.block_type.name == 'extension':
diff --git a/mitmproxy/contentviews/image/view.py b/mitmproxy/contentviews/image/view.py
index d77e1ecb..95ee1e43 100644
--- a/mitmproxy/contentviews/image/view.py
+++ b/mitmproxy/contentviews/image/view.py
@@ -8,26 +8,31 @@ from . import image_parser
 class ViewImage(base.View):
     name = "Image"
     prompt = ("image", "i")
+
+    # there is also a fallback in the auto view for image/*.
     content_types = [
         "image/png",
         "image/jpeg",
         "image/gif",
+        "image/vnd.microsoft.icon",
+        "image/x-icon",
+        "image/webp",
     ]
 
     def __call__(self, data, **metadata):
         image_type = imghdr.what('', h=data)
         if image_type == 'png':
-            f = "PNG"
-            parts = image_parser.parse_png(data)
-            fmt = base.format_dict(multidict.MultiDict(parts))
-            return "%s image" % f, fmt
+            image_metadata = image_parser.parse_png(data)
         elif image_type == 'gif':
-            f = "GIF"
-            parts = image_parser.parse_gif(data)
-            fmt = base.format_dict(multidict.MultiDict(parts))
-            return "%s image" % f, fmt
+            image_metadata = image_parser.parse_gif(data)
         elif image_type == 'jpeg':
-            f = "JPEG"
-            parts = image_parser.parse_jpeg(data)
-            fmt = base.format_dict(multidict.MultiDict(parts))
-            return "%s image" % f, fmt
+            image_metadata = image_parser.parse_jpeg(data)
+        else:
+            image_metadata = [
+                ("Image Format", image_type or "unknown")
+            ]
+        if image_type:
+            view_name = "{} Image".format(image_type.upper())
+        else:
+            view_name = "Unknown Image"
+        return view_name, base.format_dict(multidict.MultiDict(image_metadata))
diff --git a/test/mitmproxy/contentviews/image/test_image_parser.py b/test/mitmproxy/contentviews/image/test_image_parser.py
index 3c8bfdf7..3cb44ca6 100644
--- a/test/mitmproxy/contentviews/image/test_image_parser.py
+++ b/test/mitmproxy/contentviews/image/test_image_parser.py
@@ -80,7 +80,7 @@ def test_parse_png(filename, metadata):
     # check comment
     "mitmproxy/data/image_parser/hopper.gif": [
         ('Format', 'Compuserve GIF'),
-        ('version', 'GIF89a'),
+        ('Version', 'GIF89a'),
         ('Size', '128 x 128 px'),
         ('background', '0'),
         ('comment', "b'File written by Adobe Photoshop\\xa8 4.0'")
@@ -88,7 +88,7 @@ def test_parse_png(filename, metadata):
     # check background
     "mitmproxy/data/image_parser/chi.gif": [
         ('Format', 'Compuserve GIF'),
-        ('version', 'GIF89a'),
+        ('Version', 'GIF89a'),
         ('Size', '320 x 240 px'),
         ('background', '248'),
         ('comment', "b'Created with GIMP'")
@@ -96,7 +96,7 @@ def test_parse_png(filename, metadata):
     # check working with color table
     "mitmproxy/data/image_parser/iss634.gif": [
         ('Format', 'Compuserve GIF'),
-        ('version', 'GIF89a'),
+        ('Version', 'GIF89a'),
         ('Size', '245 x 245 px'),
         ('background', '0')
     ],
diff --git a/test/mitmproxy/contentviews/image/test_view.py b/test/mitmproxy/contentviews/image/test_view.py
index 84bcea17..34f655a1 100644
--- a/test/mitmproxy/contentviews/image/test_view.py
+++ b/test/mitmproxy/contentviews/image/test_view.py
@@ -9,8 +9,11 @@ def test_view_image():
         "mitmproxy/data/image.png",
         "mitmproxy/data/image.gif",
         "mitmproxy/data/all.jpeg",
+        # https://bugs.python.org/issue21574
+        # "mitmproxy/data/image.ico",
     ]:
         with open(tutils.test_data.path(img), "rb") as f:
-            assert v(f.read())
+            viewname, lines = v(f.read())
+            assert img.split(".")[-1].upper() in viewname
 
-    assert not v(b"flibble")
+    assert v(b"flibble") == ('Unknown Image', [[('header', 'Image Format: '), ('text', 'unknown')]])
diff --git a/test/mitmproxy/contentviews/test_auto.py b/test/mitmproxy/contentviews/test_auto.py
index a077affa..2ff43139 100644
--- a/test/mitmproxy/contentviews/test_auto.py
+++ b/test/mitmproxy/contentviews/test_auto.py
@@ -30,6 +30,18 @@ def test_view_auto():
     )
     assert f[0].startswith("XML")
 
+    f = v(
+        b"<svg></svg>",
+        headers=http.Headers(content_type="image/svg+xml")
+    )
+    assert f[0].startswith("XML")
+
+    f = v(
+        b"verybinary",
+        headers=http.Headers(content_type="image/new-magic-image-format")
+    )
+    assert f[0] == "Unknown Image"
+
     f = v(b"\xFF" * 30)
     assert f[0] == "Hex"