Add gif parser and tests

2 files changed, 34 insertions, 1 deletions

diff --git a/mitmproxy/contentviews/image/image_parser.py b/mitmproxy/contentviews/image/image_parser.py
index 0af58a88..b104d105 100644
--- a/mitmproxy/contentviews/image/image_parser.py
+++ b/mitmproxy/contentviews/image/image_parser.py
@@ -4,6 +4,7 @@ import typing
 from kaitaistruct import KaitaiStream
 
 from mitmproxy.contrib.kaitaistruct import png
+from mitmproxy.contrib.kaitaistruct import gif
 
 Metadata = typing.List[typing.Tuple[str, str]]
 
@@ -28,3 +29,29 @@ def parse_png(data: bytes) -> Metadata:
         elif chunk.type == 'zTXt':
             parts.append((chunk.body.keyword, chunk.body.text_datastream.decode('iso8859-1')))
     return parts
+
+
+def parse_gif(data: bytes) -> Metadata:
+    img = gif.Gif(KaitaiStream(io.BytesIO(data)))
+    parts = [
+        ('Format', 'Compuserve GIF')
+    ]
+    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':
+            ext_blocks.append(block)
+    comment_blocks = []
+    for block in ext_blocks:
+        if block.body.label._name_ == 'comment':
+            comment_blocks.append(block)
+    for block in comment_blocks:
+        entries = block.body.body.entries
+        for entry in entries:
+            comment = entry.bytes
+            if comment is not b'':
+                parts.append(('comment', str(comment)))
+    return parts
diff --git a/mitmproxy/contentviews/image/view.py b/mitmproxy/contentviews/image/view.py
index 08a70795..9caf9a6c 100644
--- a/mitmproxy/contentviews/image/view.py
+++ b/mitmproxy/contentviews/image/view.py
@@ -22,11 +22,17 @@ class ViewImage(base.View):
     ]
 
     def __call__(self, data, **metadata):
-        if imghdr.what('', h=data) == 'png':
+        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
+        elif image_type == 'gif':
+            f = "GIF"
+            parts = image_parser.parse_gif(data)
+            fmt = base.format_dict(multidict.MultiDict(parts))
+            return "%s image" % f, fmt
         try:
             img = Image.open(io.BytesIO(data))
         except IOError: