webapp.py: fuse two sql queries in get_details
[~helmut/debian-dedup.git] / dedup / image.py
1 import io
2 import struct
3
4 import PIL.Image
5
6 class ImageHash:
7     """A hash on the contents of an image data type supported by PIL. This
8     disregards mode, depth and meta information. Note that due to limitations
9     in PIL and the image format (interlacing) the full contents are stored and
10     decoded in hexdigest."""
11     maxsize = 1024 * 1024 * 32
12     # max memory usage is about 5 * maxpixels in bytes
13     maxpixels = 1024 * 1024 * 32
14
15     def __init__(self, hashobj):
16         """
17         @param hashobj: a hashlib-like object
18         """
19         self.hashobj = hashobj
20         self.imagedetected = False
21         self.content = io.BytesIO()
22
23     def detect(self):
24         raise NotImplementedError
25
26     def update(self, data):
27         self.content.write(data)
28         if self.content.tell() > self.maxsize:
29             raise ValueError("maximum image size exceeded")
30         if not self.imagedetected:
31             self.imagedetected = self.detect()
32
33     def copy(self):
34         new = self.__class__(self.hashobj.copy())
35         new.imagedetected = self.imagedetected
36         new.content = io.BytesIO(self.content.getvalue())
37         return new
38
39     def hexdigest(self):
40         if not self.imagedetected:
41             raise ValueError("not a image")
42         hashobj = self.hashobj.copy()
43         pos = self.content.tell()
44         try:
45             self.content.seek(0)
46             try:
47                 img = PIL.Image.open(self.content)
48             except IOError:
49                 raise ValueError("broken header")
50             width, height = img.size
51             pack = lambda elem: struct.pack("BBBB", *elem)
52             # special casing easy modes reduces memory usage
53             if img.mode == "L":
54                 pack = lambda elem: struct.pack("BBBB", elem, elem, elem, 255)
55             elif img.mode == "RGB":
56                 pack = lambda elem: struct.pack("BBBB", *(elem + (255,)))
57             elif img.mode != "RGBA":
58                 try:
59                     img = img.convert("RGBA")
60                 except (SyntaxError, IndexError, IOError):
61                     # crazy stuff from PIL
62                     raise ValueError("error reading image")
63             try:
64                 for elem in img.getdata():
65                     hashobj.update(pack(elem))
66             except (SyntaxError, IndexError, IOError): # crazy stuff from PIL
67                 raise ValueError("error reading image")
68         finally:
69             self.content.seek(pos)
70         return "%s%8.8x%8.8x" % (hashobj.hexdigest(), width, height)
71
72     @property
73     def name(self):
74         return self.name_prefix + self.hashobj.name
75
76
77 class PNGHash(ImageHash):
78     """A hash on the contents of a PNG image."""
79     name_prefix = "png_"
80
81     def detect(self):
82         if self.content.tell() < 33: # header + IHDR
83             return False
84         curvalue = self.content.getvalue()
85         if curvalue.startswith(b"\x89PNG\r\n\x1a\n\0\0\0\x0dIHDR"):
86             width, height = struct.unpack(">II", curvalue[16:24])
87             if width * height > self.maxpixels:
88                 raise ValueError("maximum image pixels exceeded")
89             return True
90         raise ValueError("not a png image")
91
92 class GIFHash(ImageHash):
93     """A hash on the contents of the first frame of a GIF image."""
94     name_prefix = "gif_"
95
96     def detect(self):
97         if self.content.tell() < 10: # magic + logical dimension
98             return False
99         curvalue = self.content.getvalue()
100         if curvalue.startswith((b"GIF87a", b"GIF89a")):
101             width, height = struct.unpack("<HH", curvalue[6:10])
102             if width * height > self.maxpixels:
103                 raise ValueError("maximum image pixels exceeded")
104             return True
105         raise ValueError("not a png image")