Merge pull request #1997 from s4chin/add-gif-parser

Add gif parser and tests

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

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:

mitmproxy/contrib/kaitaistruct/gif.py

@@ -0,0 +1,247 @@
+# This is a generated file! Please edit source .ksy file and use kaitai-struct-compiler to rebuild
+# The source was png.ksy from here - https://github.com/kaitai-io/kaitai_struct_formats/blob/562154250bea0081fed4e232751b934bc270a0c7/image/gif.ksy
+
+import array
+import struct
+import zlib
+from enum import Enum
+
+from kaitaistruct import KaitaiStruct, KaitaiStream, BytesIO
+
+
+class Gif(KaitaiStruct):
+
+    class BlockType(Enum):
+        extension = 33
+        local_image_descriptor = 44
+        end_of_file = 59
+
+    class ExtensionLabel(Enum):
+        graphic_control = 249
+        comment = 254
+        application = 255
+    def __init__(self, _io, _parent=None, _root=None):
+        self._io = _io
+        self._parent = _parent
+        self._root = _root if _root else self
+        self.header = self._root.Header(self._io, self, self._root)
+        self.logical_screen_descriptor = self._root.LogicalScreenDescriptor(self._io, self, self._root)
+        if self.logical_screen_descriptor.has_color_table:
+            self._raw_global_color_table = self._io.read_bytes((self.logical_screen_descriptor.color_table_size * 3))
+            io = KaitaiStream(BytesIO(self._raw_global_color_table))
+            self.global_color_table = self._root.ColorTable(io, self, self._root)
+
+        self.blocks = []
+        while not self._io.is_eof():
+            self.blocks.append(self._root.Block(self._io, self, self._root))
+
+
+    class ImageData(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.lzw_min_code_size = self._io.read_u1()
+            self.subblocks = self._root.Subblocks(self._io, self, self._root)
+
+
+    class ColorTableEntry(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.red = self._io.read_u1()
+            self.green = self._io.read_u1()
+            self.blue = self._io.read_u1()
+
+
+    class LogicalScreenDescriptor(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.screen_width = self._io.read_u2le()
+            self.screen_height = self._io.read_u2le()
+            self.flags = self._io.read_u1()
+            self.bg_color_index = self._io.read_u1()
+            self.pixel_aspect_ratio = self._io.read_u1()
+
+        @property
+        def has_color_table(self):
+            if hasattr(self, '_m_has_color_table'):
+                return self._m_has_color_table
+
+            self._m_has_color_table = (self.flags & 128) != 0
+            return self._m_has_color_table
+
+        @property
+        def color_table_size(self):
+            if hasattr(self, '_m_color_table_size'):
+                return self._m_color_table_size
+
+            self._m_color_table_size = (2 << (self.flags & 7))
+            return self._m_color_table_size
+
+
+    class LocalImageDescriptor(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.left = self._io.read_u2le()
+            self.top = self._io.read_u2le()
+            self.width = self._io.read_u2le()
+            self.height = self._io.read_u2le()
+            self.flags = self._io.read_u1()
+            if self.has_color_table:
+                self._raw_local_color_table = self._io.read_bytes((self.color_table_size * 3))
+                io = KaitaiStream(BytesIO(self._raw_local_color_table))
+                self.local_color_table = self._root.ColorTable(io, self, self._root)
+
+            self.image_data = self._root.ImageData(self._io, self, self._root)
+
+        @property
+        def has_color_table(self):
+            if hasattr(self, '_m_has_color_table'):
+                return self._m_has_color_table
+
+            self._m_has_color_table = (self.flags & 128) != 0
+            return self._m_has_color_table
+
+        @property
+        def has_interlace(self):
+            if hasattr(self, '_m_has_interlace'):
+                return self._m_has_interlace
+
+            self._m_has_interlace = (self.flags & 64) != 0
+            return self._m_has_interlace
+
+        @property
+        def has_sorted_color_table(self):
+            if hasattr(self, '_m_has_sorted_color_table'):
+                return self._m_has_sorted_color_table
+
+            self._m_has_sorted_color_table = (self.flags & 32) != 0
+            return self._m_has_sorted_color_table
+
+        @property
+        def color_table_size(self):
+            if hasattr(self, '_m_color_table_size'):
+                return self._m_color_table_size
+
+            self._m_color_table_size = (2 << (self.flags & 7))
+            return self._m_color_table_size
+
+
+    class Block(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.block_type = self._root.BlockType(self._io.read_u1())
+            _on = self.block_type
+            if _on == self._root.BlockType.extension:
+                self.body = self._root.Extension(self._io, self, self._root)
+            elif _on == self._root.BlockType.local_image_descriptor:
+                self.body = self._root.LocalImageDescriptor(self._io, self, self._root)
+
+
+    class ColorTable(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.entries = []
+            while not self._io.is_eof():
+                self.entries.append(self._root.ColorTableEntry(self._io, self, self._root))
+
+
+
+    class Header(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.magic = self._io.ensure_fixed_contents(3, struct.pack('3b', 71, 73, 70))
+            self.version = self._io.read_bytes(3)
+
+
+    class ExtGraphicControl(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.block_size = self._io.ensure_fixed_contents(1, struct.pack('1b', 4))
+            self.flags = self._io.read_u1()
+            self.delay_time = self._io.read_u2le()
+            self.transparent_idx = self._io.read_u1()
+            self.terminator = self._io.ensure_fixed_contents(1, struct.pack('1b', 0))
+
+        @property
+        def transparent_color_flag(self):
+            if hasattr(self, '_m_transparent_color_flag'):
+                return self._m_transparent_color_flag
+
+            self._m_transparent_color_flag = (self.flags & 1) != 0
+            return self._m_transparent_color_flag
+
+        @property
+        def user_input_flag(self):
+            if hasattr(self, '_m_user_input_flag'):
+                return self._m_user_input_flag
+
+            self._m_user_input_flag = (self.flags & 2) != 0
+            return self._m_user_input_flag
+
+
+    class Subblock(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.num_bytes = self._io.read_u1()
+            self.bytes = self._io.read_bytes(self.num_bytes)
+
+
+    class ExtApplication(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.application_id = self._root.Subblock(self._io, self, self._root)
+            self.subblocks = []
+            while True:
+                _ = self._root.Subblock(self._io, self, self._root)
+                self.subblocks.append(_)
+                if _.num_bytes == 0:
+                    break
+
+
+    class Subblocks(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.entries = []
+            while True:
+                _ = self._root.Subblock(self._io, self, self._root)
+                self.entries.append(_)
+                if _.num_bytes == 0:
+                    break
+
+
+    class Extension(KaitaiStruct):
+        def __init__(self, _io, _parent=None, _root=None):
+            self._io = _io
+            self._parent = _parent
+            self._root = _root if _root else self
+            self.label = self._root.ExtensionLabel(self._io.read_u1())
+            _on = self.label
+            if _on == self._root.ExtensionLabel.application:
+                self.body = self._root.ExtApplication(self._io, self, self._root)
+            elif _on == self._root.ExtensionLabel.comment:
+                self.body = self._root.Subblocks(self._io, self, self._root)
+            elif _on == self._root.ExtensionLabel.graphic_control:
+                self.body = self._root.ExtGraphicControl(self._io, self, self._root)
+            else:
+                self.body = self._root.Subblocks(self._io, self, self._root)

test/mitmproxy/contentviews/test_image_parser.py

@@ -6,13 +6,13 @@ from mitmproxy.test import tutils
 
 @pytest.mark.parametrize("filename, metadata", {
     # no textual data
-    "mitmproxy/data/png_parser/ct0n0g04.png": [
+    "mitmproxy/data/image_parser/ct0n0g04.png": [
         ('Format', 'Portable network graphics'),
         ('Size', '32 x 32 px'),
         ('gamma', '1.0')
     ],
     # with textual data
-    "mitmproxy/data/png_parser/ct1n0g04.png": [
+    "mitmproxy/data/image_parser/ct1n0g04.png": [
         ('Format', 'Portable network graphics'),
         ('Size', '32 x 32 px'),
         ('gamma', '1.0'),
@@ -27,7 +27,7 @@ from mitmproxy.test import tutils
         ('Disclaimer', 'Freeware.')
     ],
     # with compressed textual data
-    "mitmproxy/data/png_parser/ctzn0g04.png": [
+    "mitmproxy/data/image_parser/ctzn0g04.png": [
         ('Format', 'Portable network graphics'),
         ('Size', '32 x 32 px'),
         ('gamma', '1.0'),
@@ -42,7 +42,7 @@ from mitmproxy.test import tutils
         ('Disclaimer', 'Freeware.')
     ],
     # UTF-8 international text - english
-    "mitmproxy/data/png_parser/cten0g04.png": [
+    "mitmproxy/data/image_parser/cten0g04.png": [
         ('Format', 'Portable network graphics'),
         ('Size', '32 x 32 px'),
         ('gamma', '1.0'),
@@ -57,13 +57,13 @@ from mitmproxy.test import tutils
         ('Disclaimer', 'Freeware.')
     ],
     # check gamma value
-    "mitmproxy/data/png_parser/g07n0g16.png": [
+    "mitmproxy/data/image_parser/g07n0g16.png": [
         ('Format', 'Portable network graphics'),
         ('Size', '32 x 32 px'),
         ('gamma', '0.7')
     ],
     # check aspect value
-    "mitmproxy/data/png_parser/aspect.png": [
+    "mitmproxy/data/image_parser/aspect.png": [
         ('Format', 'Portable network graphics'),
         ('Size', '1280 x 798 px'),
         ('aspect', '72 x 72'),
@@ -74,3 +74,33 @@ from mitmproxy.test import tutils
 def test_parse_png(filename, metadata):
     with open(tutils.test_data.path(filename), "rb") as f:
         assert metadata == image_parser.parse_png(f.read())
+
+
+@pytest.mark.parametrize("filename, metadata", {
+    # check comment
+    "mitmproxy/data/image_parser/hopper.gif": [
+        ('Format', 'Compuserve GIF'),
+        ('version', 'GIF89a'),
+        ('Size', '128 x 128 px'),
+        ('background', '0'),
+        ('comment', "b'File written by Adobe Photoshop\\xa8 4.0'")
+    ],
+    # check background
+    "mitmproxy/data/image_parser/chi.gif": [
+        ('Format', 'Compuserve GIF'),
+        ('version', 'GIF89a'),
+        ('Size', '320 x 240 px'),
+        ('background', '248'),
+        ('comment', "b'Created with GIMP'")
+    ],
+    # check working with color table
+    "mitmproxy/data/image_parser/iss634.gif": [
+        ('Format', 'Compuserve GIF'),
+        ('version', 'GIF89a'),
+        ('Size', '245 x 245 px'),
+        ('background', '0')
+    ],
+}.items())
+def test_parse_gif(filename, metadata):
+    with open(tutils.test_data.path(filename), 'rb') as f:
+        assert metadata == image_parser.parse_gif(f.read())