Feedback

5 files changed, 28 insertions, 12 deletions

diff --git a/mitmproxy/contentviews/image/__init__.py b/mitmproxy/contentviews/image/__init__.py
index d2de66d0..0d0f06e0 100644
--- a/mitmproxy/contentviews/image/__init__.py
+++ b/mitmproxy/contentviews/image/__init__.py
@@ -1 +1 @@
-from .view import ViewImage
+from .view import ViewImage  # noqa
diff --git a/mitmproxy/contentviews/image/image_parser.py b/mitmproxy/contentviews/image/image_parser.py
index 11d66c61..0af58a88 100644
--- a/mitmproxy/contentviews/image/image_parser.py
+++ b/mitmproxy/contentviews/image/image_parser.py
@@ -1,3 +1,4 @@
+import io
 import typing
 
 from kaitaistruct import KaitaiStream
@@ -6,8 +7,9 @@ from mitmproxy.contrib.kaitaistruct import png
 
 Metadata = typing.List[typing.Tuple[str, str]]
 
+
 def parse_png(data: bytes) -> Metadata:
-    img = png.Png(KaitaiStream(data))
+    img = png.Png(KaitaiStream(io.BytesIO(data)))
     parts = [
         ('Format', 'Portable network graphics')
     ]
diff --git a/mitmproxy/contentviews/image/view.py b/mitmproxy/contentviews/image/view.py
index 4d13c917..08a70795 100644
--- a/mitmproxy/contentviews/image/view.py
+++ b/mitmproxy/contentviews/image/view.py
@@ -1,4 +1,5 @@
-import io, imghdr
+import io
+import imghdr
 
 from PIL import ExifTags
 from PIL import Image
@@ -7,7 +8,6 @@ from mitmproxy.types import multidict
 from . import image_parser
 
 from mitmproxy.contentviews import base
-from kaitaistruct import KaitaiStream
 
 
 class ViewImage(base.View):
@@ -24,7 +24,7 @@ class ViewImage(base.View):
     def __call__(self, data, **metadata):
         if imghdr.what('', h=data) == 'png':
             f = "PNG"
-            parts = image_parser.parse_png(io.BytesIO(data))
+            parts = image_parser.parse_png(data)
             fmt = base.format_dict(multidict.MultiDict(parts))
             return "%s image" % f, fmt
         try:
diff --git a/setup.py b/setup.py
index f2b5db3e..406f9e6a 100644
--- a/setup.py
+++ b/setup.py
@@ -71,7 +71,7 @@ setup(
         "html2text>=2016.1.8, <=2016.9.19",
         "hyperframe>=4.0.1, <5",
         "jsbeautifier>=1.6.3, <1.7",
-        "kaitaistruct>=0.5",
+        "kaitaistruct>=0.5, <0.6",
         "Pillow>=3.2, <4.1",
         "passlib>=1.6.5, <1.8",
         "pyasn1>=0.1.9, <0.2",
diff --git a/test/mitmproxy/contentviews/test_image_parser.py b/test/mitmproxy/contentviews/test_image_parser.py
index ba215b93..62a07f56 100644
--- a/test/mitmproxy/contentviews/test_image_parser.py
+++ b/test/mitmproxy/contentviews/test_image_parser.py
@@ -1,10 +1,9 @@
-import io
-
 import pytest
 
 from mitmproxy.contentviews.image import image_parser
 from mitmproxy.test import tutils
 
+
 @pytest.mark.parametrize("filename, metadata", {
     # no textual data
     "mitmproxy/data/png_parser/ct0n0g04.png": [
@@ -20,7 +19,12 @@ from mitmproxy.test import tutils
         ('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.')
+        ('Description', 'A compilation of a set of images created to test the\n'
+         'various 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.')
     ],
     # with compressed textual data
     "mitmproxy/data/png_parser/ctzn0g04.png": [
@@ -30,7 +34,12 @@ from mitmproxy.test import tutils
         ('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.')
+        ('Description', 'A compilation of a set of images created to test the\n'
+         'various 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.')
     ],
     # UTF-8 international text - english
     "mitmproxy/data/png_parser/cten0g04.png": [
@@ -40,7 +49,12 @@ from mitmproxy.test import tutils
         ('Title', 'PngSuite'),
         ('Author', 'Willem van Schaik (willem@schaik.com)'),
         ('Copyright', 'Copyright Willem van Schaik, Canada 2011'),
-        ('Description', 'A compilation of a set of images created to test the various color-types of the PNG format. Included are black&white, color, paletted, with alpha channel, with transparency formats. All bit-depths allowed according to the spec are present.'), ('Software', 'Created on a NeXTstation color using "pnmtopng".'), ('Disclaimer', 'Freeware.')
+        ('Description', 'A compilation of a set of images created to test the '
+         'various color-types of the PNG format. Included are black&white, color,'
+         ' paletted, with alpha channel, with transparency formats. All bit-depths'
+         ' allowed according to the spec are present.'),
+        ('Software', 'Created on a NeXTstation color using "pnmtopng".'),
+        ('Disclaimer', 'Freeware.')
     ],
     # check gamma value
     "mitmproxy/data/png_parser/g07n0g16.png": [
@@ -59,4 +73,4 @@ from mitmproxy.test import tutils
 }.items())
 def test_parse_png(filename, metadata):
     with open(tutils.test_data.path(filename), "rb") as f:
-         assert metadata == image_parser.parse_png(io.BytesIO(f.read()))
+        assert metadata == image_parser.parse_png(f.read())