diff options
author | Sachin Kelkar <sachinkel19@gmail.com> | 2017-02-06 17:57:59 +0530 |
---|---|---|
committer | Sachin Kelkar <sachinkel19@gmail.com> | 2017-02-06 17:57:59 +0530 |
commit | 4ef5de2cce6bc794b760102f2a5b82dd28d078c0 (patch) | |
tree | 276522d35d2bd93d5484ca85faa3a654592d6d08 /mitmproxy | |
parent | 2316c0fb74efefd76970c685b8f77f45834f1490 (diff) | |
download | mitmproxy-4ef5de2cce6bc794b760102f2a5b82dd28d078c0.tar.gz mitmproxy-4ef5de2cce6bc794b760102f2a5b82dd28d078c0.tar.bz2 mitmproxy-4ef5de2cce6bc794b760102f2a5b82dd28d078c0.zip |
Add gif parser and tests
Diffstat (limited to 'mitmproxy')
-rw-r--r-- | mitmproxy/contentviews/image/image_parser.py | 27 | ||||
-rw-r--r-- | mitmproxy/contentviews/image/view.py | 8 | ||||
-rw-r--r-- | mitmproxy/contrib/kaitaistruct/gif.py | 241 |
3 files changed, 275 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: diff --git a/mitmproxy/contrib/kaitaistruct/gif.py b/mitmproxy/contrib/kaitaistruct/gif.py new file mode 100644 index 00000000..0a68b815 --- /dev/null +++ b/mitmproxy/contrib/kaitaistruct/gif.py @@ -0,0 +1,241 @@ +# This is a generated file! Please edit source .ksy file and use kaitai-struct-compiler to rebuild + +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.GlobalColorTable(io, self, self._root) + + self.blocks = [] + while not self._io.is_eof(): + self.blocks.append(self._root.Block(self._io, self, self._root)) + + + class GlobalColorTable(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 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() + 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 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) |