c91fb64381136d1c3e8d320cdfaf86232787eab3
[~helmut/debian-dedup.git] / dedup / hashing.py
1 import itertools
2 try:
3     from itertools import imap as map
4 except ImportError:
5     pass # in python3 map is already imap
6
7 class HashBlacklist(object):
8     """Turn a hashlib-like object into a hash that returns None for some
9     blacklisted hashes instead of the real hash value.
10
11     We only work with hexdigests here, so diget() disappears. The methods
12     copy and update as well as the name attribute keep working as expected.
13     """
14     def __init__(self, hashobj, blacklist=()):
15         """
16         @param hashobj: a hashlib-like object
17         @param blacklist: an object providing __contains__.
18             hexdigest values which are contained in the blacklist
19             are turned into None values
20         """
21         self.hashobj = hashobj
22         self.blacklist = blacklist
23         self.update = self.hashobj.update
24
25     @property
26     def name(self):
27         return self.hashobj.name
28
29     def hexdigest(self):
30         digest = self.hashobj.hexdigest()
31         if digest in self.blacklist:
32             return None
33         return digest
34
35     def copy(self):
36         return HashBlacklist(self.hashobj.copy(), self.blacklist)
37
38 class HashBlacklistContent(object):
39     """Turn a hashlib-like object into a hash that returns None for some
40     blacklisted content instead of the real hash value. Unlike HashBlacklist,
41     not the output of the hash is considered, but its input."""
42
43     def __init__(self, hashobj, blacklist=(), maxlen=None):
44         """
45         @param hashobj: a hashlib-like object
46         @param blacklist: an object providing __contains__.
47             hash inputs which are contained in the blacklist
48             are turned into None values
49         @param maxlen: the maximum length of a blacklisted input.
50             Defaults to max(map(len, blacklist)), so if it is absent,
51             the blacklist must support iteration.
52         """
53         self.hashobj = hashobj
54         self.blacklist = blacklist
55         if maxlen is None:
56             # the chain avoids passing the empty sequence to max
57             maxlen = max(itertools.chain((0,), map(len, blacklist)))
58         self.maxlen = maxlen
59         self.stored = b""
60
61     @property
62     def name(self):
63         return self.hashobj.name
64
65     def update(self, data):
66         if self.stored is not None:
67             self.stored += data
68             if len(self.stored) > self.maxlen:
69                 self.stored = None
70         self.hashobj.update(data)
71
72     def digest(self):
73         if self.stored is not None and self.stored in self.blacklist:
74             return None
75         return self.hashobj.digest()
76
77     def hexdigest(self):
78         if self.stored is not None and self.stored in self.blacklist:
79             return None
80         return self.hashobj.hexdigest()
81
82     def copy(self):
83         new = HashBlacklistContent(self.hashobj.copy(), self.blacklist,
84                                    self.maxlen)
85         new.stored = self.stored
86         return new
87
88 class DecompressedHash(object):
89     """Apply a decompression function before the hash. This class provides the
90     hashlib interface (update, hexdigest, copy) excluding digest and name."""
91     def __init__(self, decompressor, hashobj):
92         """
93         @param decompressor: a decompression object like bz2.BZ2Decompressor or
94             lzma.LZMADecompressor. It has to provide methods decompress and
95             copy as well as an unused_data attribute. It may provide a flush
96             method.
97         @param hashobj: a hashlib-like obj providing methods update, hexdigest
98             and copy
99         """
100         self.decompressor = decompressor
101         self.hashobj = hashobj
102
103     def update(self, data):
104         self.hashobj.update(self.decompressor.decompress(data))
105
106     def hexdigest(self):
107         if not hasattr(self.decompressor, "flush"):
108             if self.decompressor.unused_data:
109                 raise ValueError("decompressor did not consume all data")
110             return self.hashobj.hexdigest()
111         tmpdecomp = self.decompressor.copy()
112         data = tmpdecomp.flush()
113         if tmpdecomp.unused_data:
114             raise ValueError("decompressor did not consume all data")
115         tmphash = self.hashobj.copy()
116         tmphash.update(data)
117         return tmphash.hexdigest()
118
119     def copy(self):
120         return DecompressedHash(self.decompressor.copy(), self.hashobj.copy())
121
122 class SuppressingHash(object):
123     """A hash that silences exceptions from the update and hexdigest methods of
124     a hashlib-like object. If an exception has occurred, hexdigest always
125     returns None."""
126     def __init__(self, hashobj, exceptions=()):
127         """
128         @param hashobj: a hashlib-like object providing methods update, copy
129             and hexdigest. If a name attribute is present, it is mirrored as
130             well.
131         @type exceptions: tuple
132         @param exceptions: exception classes to be suppressed
133         """
134         self.hashobj = hashobj
135         self.exceptions = exceptions
136         if hasattr(hashobj, "name"):
137             self.name = hashobj.name
138
139     def update(self, data):
140         if self.hashobj:
141             try:
142                 self.hashobj.update(data)
143             except self.exceptions:
144                 self.hashobj = None
145
146     def hexdigest(self):
147         if self.hashobj:
148             try:
149                 return self.hashobj.hexdigest()
150             except self.exceptions:
151                 self.hashobj = None
152         return None
153
154     def copy(self):
155         if self.hashobj:
156             return SuppressingHash(self.hashobj.copy(), self.exceptions)
157         return SuppressingHash(None, self.exceptions)
158
159 def hash_file(hashobj, filelike, blocksize=65536):
160     """Feed the entire contents from the given filelike to the given hashobj.
161     @param hashobj: hashlib-like object providing an update method
162     @param filelike: file-like object providing read(size)
163     """
164     data = filelike.read(blocksize)
165     while data:
166         hashobj.update(data)
167         data = filelike.read(blocksize)
168     return hashobj
169
170 class HashedStream(object):
171     """A file-like object, that supports sequential reading and hashes the
172     contents on the fly."""
173     def __init__(self, filelike, hashobj):
174         """
175         @param filelike: a file-like object, that must support the read method
176         @param hashobj: a hashlib-like object providing update and hexdigest
177         """
178         self.filelike = filelike
179         self.hashobj = hashobj
180
181     def read(self, length):
182         data = self.filelike.read(length)
183         self.hashobj.update(data)
184         return data
185
186     def hexdigest(self):
187         return self.hashobj.hexdigest()
188
189     def validate(self, hexdigest):
190         """Soak up any remaining input and validate the read data using the
191         given hexdigest.
192         @raises ValueError: when the hash does not match
193         """
194         while self.read(65536):
195             pass
196         if self.hexdigest() != hexdigest:
197             raise ValueError("hash sum mismatch")