2a83929d537d5e64c99875648e6ddb153f7ec14c
[~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         return HashBlacklistContent(self.hashobj.copy(), self.blacklist,
84                                     self.maxlen)
85
86 class DecompressedHash(object):
87     """Apply a decompression function before the hash. This class provides the
88     hashlib interface (update, hexdigest, copy) excluding digest and name."""
89     def __init__(self, decompressor, hashobj):
90         """
91         @param decompressor: a decompression object like bz2.BZ2Decompressor or
92             lzma.LZMADecompressor. It has to provide methods decompress and
93             copy as well as an unused_data attribute. It may provide a flush
94             method.
95         @param hashobj: a hashlib-like obj providing methods update, hexdigest
96             and copy
97         """
98         self.decompressor = decompressor
99         self.hashobj = hashobj
100
101     def update(self, data):
102         self.hashobj.update(self.decompressor.decompress(data))
103
104     def hexdigest(self):
105         if not hasattr(self.decompressor, "flush"):
106             if self.decompressor.unused_data:
107                 raise ValueError("decompressor did not consume all data")
108             return self.hashobj.hexdigest()
109         tmpdecomp = self.decompressor.copy()
110         data = tmpdecomp.flush()
111         if tmpdecomp.unused_data:
112             raise ValueError("decompressor did not consume all data")
113         tmphash = self.hashobj.copy()
114         tmphash.update(data)
115         return tmphash.hexdigest()
116
117     def copy(self):
118         return DecompressedHash(self.decompressor.copy(), self.hashobj.copy())
119
120 class SuppressingHash(object):
121     """A hash that silences exceptions from the update and hexdigest methods of
122     a hashlib-like object. If an exception has occurred, hexdigest always
123     returns None."""
124     def __init__(self, hashobj, exceptions=()):
125         """
126         @param hashobj: a hashlib-like object providing methods update, copy
127             and hexdigest. If a name attribute is present, it is mirrored as
128             well.
129         @type exceptions: tuple
130         @param exceptions: exception classes to be suppressed
131         """
132         self.hashobj = hashobj
133         self.exceptions = exceptions
134         if hasattr(hashobj, "name"):
135             self.name = hashobj.name
136
137     def update(self, data):
138         if self.hashobj:
139             try:
140                 self.hashobj.update(data)
141             except self.exceptions:
142                 self.hashobj = None
143
144     def hexdigest(self):
145         if self.hashobj:
146             try:
147                 return self.hashobj.hexdigest()
148             except self.exceptions:
149                 self.hashobj = None
150         return None
151
152     def copy(self):
153         if self.hashobj:
154             return SuppressingHash(self.hashobj.copy(), self.exceptions)
155         return SuppressingHash(None, self.exceptions)
156
157 def hash_file(hashobj, filelike, blocksize=65536):
158     """Feed the entire contents from the given filelike to the given hashobj.
159     @param hashobj: hashlib-like object providing an update method
160     @param filelike: file-like object providing read(size)
161     """
162     data = filelike.read(blocksize)
163     while data:
164         hashobj.update(data)
165         data = filelike.read(blocksize)
166     return hashobj
167
168 class HashedStream(object):
169     """A file-like object, that supports sequential reading and hashes the
170     contents on the fly."""
171     def __init__(self, filelike, hashobj):
172         """
173         @param filelike: a file-like object, that must support the read method
174         @param hashobj: a hashlib-like object providing update and hexdigest
175         """
176         self.filelike = filelike
177         self.hashobj = hashobj
178
179     def read(self, length):
180         data = self.filelike.read(length)
181         self.hashobj.update(data)
182         return data
183
184     def hexdigest(self):
185         return self.hashobj.hexdigest()
186
187     def validate(self, hexdigest):
188         """Soak up any remaining input and validate the read data using the
189         given hexdigest.
190         @raises ValueError: when the hash does not match
191         """
192         while self.read(65536):
193             pass
194         if self.hexdigest() != hexdigest:
195             raise ValueError("hash sum mismatch")