Fix as per feedback and add more tests

4 files changed, 20 insertions, 32 deletions

diff --git a/mitmproxy/contentviews/__init__.py b/mitmproxy/contentviews/__init__.py
index b28f9d8f..357172e3 100644
--- a/mitmproxy/contentviews/__init__.py
+++ b/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():
diff --git a/mitmproxy/contentviews/image/__init__.py b/mitmproxy/contentviews/image/__init__.py
index e69de29b..d2de66d0 100644
--- a/mitmproxy/contentviews/image/__init__.py
+++ b/mitmproxy/contentviews/image/__init__.py
@@ -0,0 +1 @@
+from .view import ViewImage
diff --git a/mitmproxy/contentviews/image/image_parser.py b/mitmproxy/contentviews/image/image_parser.py
index 8e5762c9..e523def5 100644
--- a/mitmproxy/contentviews/image/image_parser.py
+++ b/mitmproxy/contentviews/image/image_parser.py
@@ -1,34 +1,22 @@
+import typing
+
 from kaitaistruct import KaitaiStream
 
 from . import png
 
-def get_png(data):
+Metadata = typing.List[typing.Tuple[str, str]]
+
+def parse_png(data: bytes) -> Metadata:
     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]
+    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':
-            gamma = chunk.gamma_int / 100000
-            parts['gamma'] = gamma
+            parts.append(tuple(['gamma', str(chunk.body.gamma_int / 100000)]))
         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
-
-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
+            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
diff --git a/mitmproxy/contentviews/image/pillow.py b/mitmproxy/contentviews/image/view.py
index c48d665a..4d13c917 100644
--- a/mitmproxy/contentviews/image/pillow.py
+++ b/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: