1#!/usr/bin/python
2
3# Audio Tools, a module and set of tools for manipulating audio data
4# Copyright (C) 2007-2014  Brian Langenberger
5
6# This program is free software; you can redistribute it and/or modify
7# it under the terms of the GNU General Public License as published by
8# the Free Software Foundation; either version 2 of the License, or
9# (at your option) any later version.
10
11# This program is distributed in the hope that it will be useful,
12# but WITHOUT ANY WARRANTY; without even the implied warranty of
13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14# GNU General Public License for more details.
15
16# You should have received a copy of the GNU General Public License
17# along with this program; if not, write to the Free Software
18# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
19
20from audiotools.bitstream import BitstreamReader, format_size
21from audiotools import InvalidImage
22
23
24def image_metrics(file_data):
25    """returns an ImageMetrics subclass from a string of file data
26
27    raises InvalidImage if there is an error parsing the file
28    or its type is unknown"""
29
30    if file_data[0:3] == b"\xff\xd8\xff":
31        return __JPEG__.parse(file_data)
32    elif file_data[0:8] == b'\x89\x50\x4E\x47\x0D\x0A\x1A\x0A':
33        return __PNG__.parse(file_data)
34    elif file_data[0:3] == b'GIF':
35        return __GIF__.parse(file_data)
36    elif file_data[0:2] == b'BM':
37        return __BMP__.parse(file_data)
38    elif (file_data[0:2] == b'II') or (file_data[0:2] == b'MM'):
39        return __TIFF__.parse(file_data)
40    else:
41        from audiotools.text import ERR_IMAGE_UNKNOWN_TYPE
42        raise InvalidImage(ERR_IMAGE_UNKNOWN_TYPE)
43
44
45#######################
46# JPEG
47#######################
48
49
50class ImageMetrics:
51    """a container for image data"""
52
53    def __init__(self, width, height, bits_per_pixel, color_count, mime_type):
54        """fields are as follows:
55
56        width          - image width as an integer number of pixels
57        height         - image height as an integer number of pixels
58        bits_per_pixel - the number of bits per pixel as an integer
59        color_count    - for palette-based images, the total number of colors
60        mime_type      - the image's MIME type, as a unicode string
61
62        all of the ImageMetrics subclasses implement these fields
63        in addition, they all implement a parse() classmethod
64        used to parse binary string data and return something
65        imageMetrics compatible
66        """
67
68        self.width = width
69        self.height = height
70        self.bits_per_pixel = bits_per_pixel
71        self.color_count = color_count
72        self.mime_type = mime_type
73
74    def __repr__(self):
75        return "ImageMetrics(%s,%s,%s,%s,%s)" % \
76               (repr(self.width),
77                repr(self.height),
78                repr(self.bits_per_pixel),
79                repr(self.color_count),
80                repr(self.mime_type))
81
82    @classmethod
83    def parse(cls, file_data):
84        raise NotImplementedError()
85
86
87class InvalidJPEG(InvalidImage):
88    """raised if a JPEG cannot be parsed correctly"""
89
90    pass
91
92
93class __JPEG__(ImageMetrics):
94    def __init__(self, width, height, bits_per_pixel):
95        ImageMetrics.__init__(self, width, height, bits_per_pixel,
96                              0, u'image/jpeg')
97
98    @classmethod
99    def parse(cls, file_data):
100        def segments(reader):
101            if reader.read(8) != 0xFF:
102                from audiotools.text import ERR_IMAGE_INVALID_JPEG_MARKER
103                raise InvalidJPEG(ERR_IMAGE_INVALID_JPEG_MARKER)
104            segment_type = reader.read(8)
105
106            while segment_type != 0xDA:
107                if segment_type not in {0xD8, 0xD9}:
108                    yield (segment_type, reader.substream(reader.read(16) - 2))
109                else:
110                    yield (segment_type, None)
111
112                if reader.read(8) != 0xFF:
113                    from audiotools.text import ERR_IMAGE_INVALID_JPEG_MARKER
114                    raise InvalidJPEG(ERR_IMAGE_INVALID_JPEG_MARKER)
115                segment_type = reader.read(8)
116
117        try:
118            for (segment_type,
119                 segment_data) in segments(BitstreamReader(file_data, False)):
120                if (segment_type in {0xC0, 0xC1, 0xC2, 0xC3,
121                                     0xC5, 0XC5, 0xC6, 0xC7,
122                                     0xC9, 0xCA, 0xCB, 0xCD,
123                                     0xCE, 0xCF}):  # start of frame
124                    (data_precision,
125                     image_height,
126                     image_width,
127                     components) = segment_data.parse("8u 16u 16u 8u")
128                    return __JPEG__(width=image_width,
129                                    height=image_height,
130                                    bits_per_pixel=data_precision * components)
131        except IOError:
132            from audiotools.text import ERR_IMAGE_IOERROR_JPEG
133            raise InvalidJPEG(ERR_IMAGE_IOERROR_JPEG)
134
135#######################
136# PNG
137#######################
138
139
140class InvalidPNG(InvalidImage):
141    """raised if a PNG cannot be parsed correctly"""
142
143    pass
144
145
146class __PNG__(ImageMetrics):
147    def __init__(self, width, height, bits_per_pixel, color_count):
148        ImageMetrics.__init__(self, width, height, bits_per_pixel, color_count,
149                              u'image/png')
150
151    @classmethod
152    def parse(cls, file_data):
153        def chunks(reader):
154            if reader.read_bytes(8) != b'\x89\x50\x4E\x47\x0D\x0A\x1A\x0A':
155                from audiotools.text import ERR_IMAGE_INVALID_PNG
156                raise InvalidPNG(ERR_IMAGE_INVALID_PNG)
157            (chunk_length, chunk_type) = reader.parse("32u 4b")
158            while chunk_type != b'IEND':
159                yield (chunk_type,
160                       chunk_length,
161                       reader.substream(chunk_length))
162                chunk_crc = reader.read(32)
163                (chunk_length, chunk_type) = reader.parse("32u 4b")
164
165        ihdr = None
166        plte_length = 0
167
168        try:
169            for (chunk_type,
170                 chunk_length,
171                 chunk_data) in chunks(BitstreamReader(file_data, False)):
172                if chunk_type == b'IHDR':
173                    ihdr = chunk_data
174                elif chunk_type == b'PLTE':
175                    plte_length = chunk_length
176
177            if ihdr is None:
178                from audiotools.text import ERR_IMAGE_INVALID_PNG
179                raise InvalidPNG(ERR_IMAGE_INVALID_PNG)
180
181            (width,
182             height,
183             bit_depth,
184             color_type,
185             compression_method,
186             filter_method,
187             interlace_method) = ihdr.parse("32u 32u 8u 8u 8u 8u 8u")
188        except IOError:
189            from audiotools.text import ERR_IMAGE_IOERROR_PNG
190            raise InvalidPNG(ERR_IMAGE_IOERROR_PNG)
191
192        if color_type == 0:    # grayscale
193            return cls(width=width,
194                       height=height,
195                       bits_per_pixel=bit_depth,
196                       color_count=0)
197        elif color_type == 2:  # RGB
198            return cls(width=width,
199                       height=height,
200                       bits_per_pixel=bit_depth * 3,
201                       color_count=0)
202        elif color_type == 3:  # palette
203            if (plte_length % 3) != 0:
204                from audiotools.text import ERR_IMAGE_INVALID_PLTE
205                raise InvalidPNG(ERR_IMAGE_INVALID_PLTE)
206            else:
207                return cls(width=width,
208                           height=height,
209                           bits_per_pixel=8,
210                           color_count=plte_length // 3)
211        elif color_type == 4:  # grayscale + alpha
212            return cls(width=width,
213                       height=height,
214                       bits_per_pixel=bit_depth * 2,
215                       color_count=0)
216        elif color_type == 6:  # RGB + alpha
217            return cls(width=width,
218                       height=height,
219                       bits_per_pixel=bit_depth * 4,
220                       color_count=0)
221
222#######################
223# BMP
224#######################
225
226
227class InvalidBMP(InvalidImage):
228    """raised if a BMP cannot be parsed correctly"""
229
230    pass
231
232
233class __BMP__(ImageMetrics):
234    def __init__(self, width, height, bits_per_pixel, color_count):
235        ImageMetrics.__init__(self, width, height, bits_per_pixel, color_count,
236                              u'image/x-ms-bmp')
237
238    @classmethod
239    def parse(cls, file_data):
240        try:
241            (magic_number,
242             file_size,
243             data_offset,
244             header_size,
245             width,
246             height,
247             color_planes,
248             bits_per_pixel,
249             compression_method,
250             image_size,
251             horizontal_resolution,
252             vertical_resolution,
253             colors_used,
254             important_colors_used) = BitstreamReader(file_data, True).parse(
255                "2b 32u 16p 16p 32u " +
256                "32u 32u 32u 16u 16u 32u 32u 32u 32u 32u 32u")
257        except IOError:
258            from audiotools.text import ERR_IMAGE_IOERROR_BMP
259            raise InvalidBMP(ERR_IMAGE_IOERROR_BMP)
260
261        if magic_number != b'BM':
262            from audiotools.text import ERR_IMAGE_INVALID_BMP
263            raise InvalidBMP(ERR_IMAGE_INVALID_BMP)
264        else:
265            return cls(width=width,
266                       height=height,
267                       bits_per_pixel=bits_per_pixel,
268                       color_count=colors_used)
269
270
271#######################
272# GIF
273#######################
274
275
276class InvalidGIF(InvalidImage):
277    """raised if a GIF cannot be parsed correctly"""
278
279    pass
280
281
282class __GIF__(ImageMetrics):
283    def __init__(self, width, height, color_count):
284        ImageMetrics.__init__(self, width, height, 8, color_count,
285                              u'image/gif')
286
287    @classmethod
288    def parse(cls, file_data):
289        try:
290            (gif,
291             version,
292             width,
293             height,
294             color_table_size) = BitstreamReader(file_data, True).parse(
295                "3b 3b 16u 16u 3u 5p")
296        except IOError:
297            from audiotools.text import ERR_IMAGE_IOERROR_GIF
298            raise InvalidGIF(ERR_IMAGE_IOERROR_GIF)
299
300        if gif != b'GIF':
301            from audiotools.text import ERR_IMAGE_INVALID_GIF
302            raise InvalidGIF(ERR_IMAGE_INVALID_GIF)
303        else:
304            return cls(width=width,
305                       height=height,
306                       color_count=2 ** (color_table_size + 1))
307
308
309#######################
310# TIFF
311#######################
312
313
314class InvalidTIFF(InvalidImage):
315    """raised if a TIFF cannot be parsed correctly"""
316
317    pass
318
319
320class __TIFF__(ImageMetrics):
321    def __init__(self, width, height, bits_per_pixel, color_count):
322        ImageMetrics.__init__(self, width, height,
323                              bits_per_pixel, color_count,
324                              u'image/tiff')
325
326    @classmethod
327    def parse(cls, file_data):
328        from io import BytesIO
329
330        def tags(file, order):
331            while True:
332                reader = BitstreamReader(file, order)
333                # read all the tags in an IFD
334                tag_count = reader.read(16)
335                sub_reader = reader.substream(tag_count * 12)
336                next_ifd = reader.read(32)
337
338                for i in range(tag_count):
339                    (tag_code,
340                     tag_datatype,
341                     tag_value_count) = sub_reader.parse("16u 16u 32u")
342                    if tag_datatype == 1:    # BYTE type
343                        tag_struct = "8u" * tag_value_count
344                    elif tag_datatype == 3:  # SHORT type
345                        tag_struct = "16u" * tag_value_count
346                    elif tag_datatype == 4:  # LONG type
347                        tag_struct = "32u" * tag_value_count
348                    else:                      # all other types
349                        tag_struct = "4b"
350                    if format_size(tag_struct) <= 32:
351                        yield (tag_code, sub_reader.parse(tag_struct))
352                        sub_reader.skip(32 - format_size(tag_struct))
353                    else:
354                        offset = sub_reader.read(32)
355                        file.seek(offset, 0)
356                        yield (tag_code,
357                               BitstreamReader(file, order).parse(tag_struct))
358
359                if next_ifd != 0:
360                    file.seek(next_ifd, 0)
361                else:
362                    break
363
364        file = BytesIO(file_data)
365        try:
366            byte_order = file.read(2)
367            if byte_order == b'II':
368                order = 1
369            elif byte_order == b'MM':
370                order = 0
371            else:
372                from audiotools.text import ERR_IMAGE_INVALID_TIFF
373                raise InvalidTIFF(ERR_IMAGE_INVALID_TIFF)
374            reader = BitstreamReader(file, order)
375            if reader.read(16) != 42:
376                from audiotools.text import ERR_IMAGE_INVALID_TIFF
377                raise InvalidTIFF(ERR_IMAGE_INVALID_TIFF)
378
379            initial_ifd = reader.read(32)
380            file.seek(initial_ifd, 0)
381
382            width = 0
383            height = 0
384            bits_per_pixel = 0
385            color_count = 0
386            for (tag_id, tag_values) in tags(file, order):
387                if tag_id == 0x0100:
388                    width = tag_values[0]
389                elif tag_id == 0x0101:
390                    height = tag_values[0]
391                elif tag_id == 0x0102:
392                    bits_per_pixel = sum(tag_values)
393                elif tag_id == 0x0140:
394                    color_count = len(tag_values) // 3
395        except IOError:
396            from audiotools.text import ERR_IMAGE_IOERROR_TIFF
397            raise InvalidTIFF(ERR_IMAGE_IOERROR_TIFF)
398
399        return cls(width=width,
400                   height=height,
401                   bits_per_pixel=bits_per_pixel,
402                   color_count=color_count)
403