diff --git a/mitmproxy/contentviews/__init__.py b/mitmproxy/contentviews/__init__.py index b28f9d8f3..357172e3e 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 e69de29bb..d2de66d02 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 8e5762c99..e523def56 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): - 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 diff --git a/mitmproxy/contentviews/image/pillow.py b/mitmproxy/contentviews/image/view.py similarity index 92% rename from mitmproxy/contentviews/image/pillow.py rename to mitmproxy/contentviews/image/view.py index c48d665a7..4d13c917e 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: diff --git a/setup.py b/setup.py index 85eece95a..f2b5db3e8 100644 --- a/setup.py +++ b/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", diff --git a/test/mitmproxy/contentviews/test_image.py b/test/mitmproxy/contentviews/test_image.py index 000b9da55..9e7e28f56 100644 --- a/test/mitmproxy/contentviews/test_image.py +++ b/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", diff --git a/test/mitmproxy/contentviews/test_image_parser.py b/test/mitmproxy/contentviews/test_image_parser.py index d41163927..2544e3f10 100644 --- a/test/mitmproxy/contentviews/test_image_parser.py +++ b/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 diff --git a/test/mitmproxy/data/png_parser/aspect.png b/test/mitmproxy/data/png_parser/aspect.png new file mode 100644 index 000000000..17c01913f Binary files /dev/null and b/test/mitmproxy/data/png_parser/aspect.png differ diff --git a/test/mitmproxy/data/png_parser/ct0n0g04.png b/test/mitmproxy/data/png_parser/ct0n0g04.png new file mode 100644 index 000000000..40d1e062f Binary files /dev/null and b/test/mitmproxy/data/png_parser/ct0n0g04.png differ diff --git a/test/mitmproxy/data/png_parser/ct1n0g04.png b/test/mitmproxy/data/png_parser/ct1n0g04.png new file mode 100644 index 000000000..3ba110aa7 Binary files /dev/null and b/test/mitmproxy/data/png_parser/ct1n0g04.png differ diff --git a/test/mitmproxy/data/png_parser/cten0g04.png b/test/mitmproxy/data/png_parser/cten0g04.png new file mode 100644 index 000000000..a6a56faf2 Binary files /dev/null and b/test/mitmproxy/data/png_parser/cten0g04.png differ diff --git a/test/mitmproxy/data/png_parser/g07n0g16.png b/test/mitmproxy/data/png_parser/g07n0g16.png new file mode 100644 index 000000000..d6a47c2d5 Binary files /dev/null and b/test/mitmproxy/data/png_parser/g07n0g16.png differ