Fix as per feedback and add more tests

mitmproxy/contentviews/__init__.py

@@ -23,9 +23,8 @@ from mitmproxy.net import http
 from mitmproxy.utils import strutils
 from . import (
     auto, raw, hex, json, xml_html, html_outline, wbxml, javascript, css,
-    urlencoded, multipart, query, protobuf
+    urlencoded, multipart, image, query, protobuf
 )
-from .image import pillow
 from .base import View, VIEW_CUTOFF, KEY_MAX, format_text, format_dict
 
 views = []  # type: List[View]
@@ -171,7 +170,7 @@ add(javascript.ViewJavaScript())
 add(css.ViewCSS())
 add(urlencoded.ViewURLEncoded())
 add(multipart.ViewMultipart())
-add(pillow.ViewImage())
+add(image.ViewImage())
 add(query.ViewQuery())
 
 if protobuf.ViewProtobuf.is_available():

mitmproxy/contentviews/image/__init__.py

@@ -0,0 +1 @@
+from .view import ViewImage

mitmproxy/contentviews/image/image_parser.py

@@ -1,34 +1,22 @@
+import typing
+
 from kaitaistruct import KaitaiStream
 
 from . import png
 
-def get_png(data):
-    img = png.Png(KaitaiStream(data))
-    parts = {'format': 'Portable network graphics'}
-    f = 'PNG'
-    width = img.ihdr.width
-    height = img.ihdr.height
-    parts["width"] = width
-    parts["height"] = height
-    for i in range(0, len(img.chunks)):
-        chunk = img.chunks[i]
-        if chunk.type == 'gAMA':
-            gamma = chunk.gamma_int / 100000
-            parts['gamma'] = gamma
-        elif chunk.type == 'pHYs':
-            aspectx = chunk.pixels_per_unit_x
-            aspecty = chunk.pixels_per_unit_y
-            parts["aspectx"] = aspectx
-            parts["aspecty"] = aspecty
-    return f, parts
+Metadata = typing.List[typing.Tuple[str, str]]
 
-def format_contentviews(parts):
-    ret = []
-    ret.append(tuple(['Format', parts["format"]]))
-    if "width" in parts:
-        ret.append(tuple(['Size', str(parts["width"]) + " x " + str(parts["height"]) + " px"]))
-    if "aspectx" in parts:
-        ret.append(tuple(['aspect', '(' + str(parts["aspectx"]) + ', ' + str(parts["aspecty"]) + ')']))
-    if "gamma" in parts:
-        ret.append(tuple(['gamma', str(parts["gamma"])]))
-    return ret
+def parse_png(data: bytes) -> Metadata:
+    img = png.Png(KaitaiStream(data))
+    parts = [tuple(['Format', 'Portable network graphics'])]
+    parts.append(tuple(['Size', str(img.ihdr.width) + " x " + str(img.ihdr.height) + " px"]))
+    for chunk in img.chunks:
+        if chunk.type == 'gAMA':
+            parts.append(tuple(['gamma', str(chunk.body.gamma_int / 100000)]))
+        elif chunk.type == 'pHYs':
+            aspectx = chunk.body.pixels_per_unit_x
+            aspecty = chunk.body.pixels_per_unit_y
+            parts.append(tuple(['aspect', str(aspectx) + " x " + str(aspecty)]))
+        elif chunk.type == 'tEXt':
+            parts.append(tuple([chunk.body.keyword, chunk.body.text]))
+    return parts

mitmproxy/contentviews/image/view.py

@@ -23,8 +23,8 @@ class ViewImage(base.View):
 
     def __call__(self, data, **metadata):
         if imghdr.what('', h=data) == 'png':
-            f, parts = image_parser.get_png(io.BytesIO(data))
-            parts = image_parser.format_contentviews(parts)
+            f = "PNG"
+            parts = image_parser.parse_png(io.BytesIO(data))
             fmt = base.format_dict(multidict.MultiDict(parts))
             return "%s image" % f, fmt
         try:

setup.py

@@ -71,6 +71,7 @@ setup(
         "html2text>=2016.1.8, <=2016.9.19",
         "hyperframe>=4.0.1, <5",
         "jsbeautifier>=1.6.3, <1.7",
+        "kaitaistruct>=0.5",
         "Pillow>=3.2, <4.1",
         "passlib>=1.6.5, <1.8",
         "pyasn1>=0.1.9, <0.2",

test/mitmproxy/contentviews/test_image.py

@@ -1,10 +1,10 @@
-from mitmproxy.contentviews.image import pillow
+from mitmproxy.contentviews import image
 from mitmproxy.test import tutils
 from . import full_eval
 
 
 def test_view_image():
-    v = full_eval(pillow.ViewImage())
+    v = full_eval(image.ViewImage())
     for img in [
         "mitmproxy/data/image.png",
         "mitmproxy/data/image.gif",

test/mitmproxy/contentviews/test_image_parser.py

@@ -3,12 +3,54 @@ import io
 from mitmproxy.contentviews.image import image_parser
 from mitmproxy.test import tutils
 
-def test_png_parser():
-    img = "mitmproxy/data/image.png"
-    with open(tutils.test_data.path(img), "rb") as f:
-        fmt, parts = image_parser.get_png(io.BytesIO(f.read()))
-        assert fmt == "PNG"
-        assert parts
-        assert parts["width"] == 174
-        assert parts["height"] == 174
-        assert parts["format"] == "Portable network graphics"
+class TestPngParser:
+    def test_png_parser(self):
+        img = "mitmproxy/data/image.png"
+        with open(tutils.test_data.path(img), "rb") as f:
+            parts = image_parser.parse_png(io.BytesIO(f.read()))
+            assert parts
+            assert tuple(['Size', '174 x 174 px']) in parts
+            assert tuple(["Format", "Portable network graphics"]) in parts
+
+    def test_textual_data(self):
+        img = "mitmproxy/data/png_parser/ct1n0g04.png"
+        with open(tutils.test_data.path(img), "rb") as f:
+            parts = image_parser.parse_png(io.BytesIO(f.read()))
+            assert parts
+            expected = [
+                ('Title', 'PngSuite'),
+                ('Author', 'Willem A.J. van Schaik\n(willem@schaik.com)'),
+                ('Copyright', 'Copyright Willem van Schaik, Singapore 1995-96'),
+                ('Description', 'A compilation of a set of images created to test the\nvarious color-types of the PNG format. Included are\nblack&white, color, paletted, with alpha channel, with\ntransparency formats. All bit-depths allowed according\nto the spec are present.'),
+                ('Software', 'Created on a NeXTstation color using "pnmtopng".'),
+                ('Disclaimer', 'Freeware.')
+            ]
+            for data in expected:
+                assert data in parts
+
+    def test_no_textual_data(self):
+        img = "mitmproxy/data/png_parser/ct0n0g04.png"
+        with open(tutils.test_data.path(img), "rb") as f:
+            parts = image_parser.parse_png(io.BytesIO(f.read()))
+            assert parts
+            metadata = [
+                ('Format', 'Portable network graphics'),
+                ('Size', '32 x 32 px'),
+                ('gamma', '1.0')
+                ]
+            parts = [data for data in parts if data not in metadata]
+            assert not parts
+
+    def test_gamma(self):
+        img = "mitmproxy/data/png_parser/g07n0g16.png"
+        with open(tutils.test_data.path(img), "rb") as f:
+            parts = image_parser.parse_png(io.BytesIO(f.read()))
+            assert parts
+            assert ('gamma', '0.7') in parts
+
+    def test_gamma(self):
+        img = "mitmproxy/data/png_parser/aspect.png"
+        with open(tutils.test_data.path(img), "rb") as f:
+            parts = image_parser.parse_png(io.BytesIO(f.read()))
+            assert parts
+            assert ('aspect', '72 x 72') in parts