2 """Hash and store hash information for a file.
4 @var PIECE_SIZE: the piece size to use for hashing pieces of files
8 from binascii import b2a_hex, a2b_hex
11 from twisted.internet import threads, defer
12 from twisted.trial import unittest
16 class HashError(ValueError):
17 """An error has occurred while hashing a file."""
20 """Manages hashes and hashing for a file.
22 @ivar ORDER: the priority ordering of hashes, and how to extract them
26 ORDER = [ {'name': 'sha1',
28 'AptPkgRecord': 'SHA1Hash',
29 'AptSrcRecord': False,
30 'AptIndexRecord': 'SHA1',
32 'hashlib_func': 'sha1',
36 'AptPkgRecord': 'SHA256Hash',
37 'AptSrcRecord': False,
38 'AptIndexRecord': 'SHA256',
39 'hashlib_func': 'sha256',
43 'AptPkgRecord': 'MD5Hash',
45 'AptIndexRecord': 'MD5SUM',
47 'hashlib_func': 'md5',
51 def __init__(self, digest = None, size = None, pieces = ''):
52 """Initialize the hash object."""
53 self.hashTypeNum = 0 # Use the first if nothing else matters
54 if sys.version_info < (2, 5):
55 # sha256 is not available in python before 2.5, remove it
56 for hashType in self.ORDER:
57 if hashType['name'] == 'sha256':
58 del self.ORDER[self.ORDER.index(hashType)]
64 self.expNormHash = None
65 self.fileHasher = None
66 self.pieceHasher = None
67 self.fileHash = digest
68 self.pieceHash = [pieces[x:x+20] for x in xrange(0, len(pieces), 20)]
71 self.fileNormHash = None
76 def new(self, force = False):
77 """Generate a new hashing object suitable for hashing a file.
79 @param force: set to True to force creating a new object even if
80 the hash has been verified already
82 if self.result is None or force:
85 self.fileHasher = self._new()
86 if self.ORDER[self.hashTypeNum]['name'] == 'sha1':
87 self.pieceHasher = None
89 self.pieceHasher = self._newSHA1()
95 self.fileNormHash = None
98 """Create a new hashing object according to the hash type."""
99 if sys.version_info < (2, 5):
100 mod = __import__(self.ORDER[self.hashTypeNum]['old_module'], globals(), locals(), [])
104 func = getattr(hashlib, self.ORDER[self.hashTypeNum]['hashlib_func'])
108 """Create a new SHA1 hashing object."""
109 if sys.version_info < (2, 5):
114 return hashlib.sha1()
116 def update(self, data):
117 """Add more data to the file hasher."""
118 if self.result is None:
120 raise HashError, "Already done, you can't add more data after calling digest() or verify()"
121 if self.fileHasher is None:
122 raise HashError, "file hasher not initialized"
124 if not self.pieceHasher and self.size + len(data) > PIECE_SIZE:
125 # Hash up to the piece size
126 self.fileHasher.update(data[:(PIECE_SIZE - self.size)])
127 data = data[(PIECE_SIZE - self.size):]
128 self.size = PIECE_SIZE
131 # Save the first piece digest and initialize a new piece hasher
132 self.pieceHash.append(self.fileHasher.digest())
133 self.pieceHasher = self._newSHA1()
136 # Loop in case the data contains multiple pieces
137 while self.pieceSize + len(data) > PIECE_SIZE:
138 # Save the piece hash and start a new one
139 self.pieceHasher.update(data[:(PIECE_SIZE - self.pieceSize)])
140 self.pieceHash.append(self.pieceHasher.digest())
141 self.pieceHasher = self._newSHA1()
143 # Don't forget to hash the data normally
144 self.fileHasher.update(data[:(PIECE_SIZE - self.pieceSize)])
145 data = data[(PIECE_SIZE - self.pieceSize):]
146 self.size += PIECE_SIZE - self.pieceSize
149 # Hash any remaining data
150 self.pieceHasher.update(data)
151 self.pieceSize += len(data)
153 self.fileHasher.update(data)
154 self.size += len(data)
156 def hashInThread(self, file):
157 """Hashes a file in a separate thread, returning a deferred that will callback with the result."""
159 if not file.exists():
160 return defer.fail(HashError("file not found"))
162 df = threads.deferToThread(self._hashInThread, file)
165 def _hashInThread(self, file):
166 """Hashes a file, returning itself as the result."""
168 self.new(force = True)
176 #{ Checking hashes of data
177 def pieceDigests(self):
178 """Get the piece hashes of the added file data."""
180 return self.pieceHash
183 """Get the hash of the added file data."""
184 if self.fileHash is None:
185 if self.fileHasher is None:
186 raise HashError, "you must hash some data first"
187 self.fileHash = self.fileHasher.digest()
190 # Save the last piece hash
192 self.pieceHash.append(self.pieceHasher.digest())
196 """Get the hash of the added file data in hex format."""
197 if self.fileHex is None:
198 self.fileHex = b2a_hex(self.digest())
202 """Verify that the added file data hash matches the expected hash."""
203 if self.result is None and self.fileHash is not None and self.expHash is not None:
204 self.result = (self.fileHash == self.expHash and self.size == self.expSize)
209 """Get the expected hash."""
212 def hexexpected(self):
213 """Get the expected hash in hex format."""
214 if self.expHex is None and self.expHash is not None:
215 self.expHex = b2a_hex(self.expHash)
218 #{ Setting the expected hash
219 def set(self, hashType, hashHex, size):
220 """Initialize the hash object.
222 @param hashType: must be one of the dictionaries from L{ORDER}
224 self.hashTypeNum = self.ORDER.index(hashType) # error if not found
225 self.expHex = hashHex
226 self.expSize = int(size)
227 self.expHash = a2b_hex(self.expHex)
229 def setFromIndexRecord(self, record):
230 """Set the hash from the cache of index file records.
232 @type record: C{dictionary}
233 @param record: keys are hash types, values are tuples of (hash, size)
235 for hashType in self.ORDER:
236 result = record.get(hashType['AptIndexRecord'], None)
238 self.set(hashType, result[0], result[1])
242 def setFromPkgRecord(self, record, size):
243 """Set the hash from Apt's binary packages cache.
245 @param record: whatever is returned by apt_pkg.GetPkgRecords()
247 for hashType in self.ORDER:
248 hashHex = getattr(record, hashType['AptPkgRecord'], None)
250 self.set(hashType, hashHex, size)
254 def setFromSrcRecord(self, record):
255 """Set the hash from Apt's source package records cache.
257 Currently very simple since Apt only tracks MD5 hashes of source files.
259 @type record: (C{string}, C{int}, C{string})
260 @param record: the hash, size and path of the source file
262 for hashType in self.ORDER:
263 if hashType['AptSrcRecord']:
264 self.set(hashType, record[0], record[1])
268 class TestHashObject(unittest.TestCase):
269 """Unit tests for the hash objects."""
272 if sys.version_info < (2, 4):
275 def test_failure(self):
276 """Tests that the hash object fails when treated badly."""
278 h.set(h.ORDER[0], b2a_hex('12345678901234567890'), '0')
279 self.failUnlessRaises(HashError, h.digest)
280 self.failUnlessRaises(HashError, h.hexdigest)
281 self.failUnlessRaises(HashError, h.update, 'gfgf')
283 def test_pieces(self):
284 """Tests updating of pieces a little at a time."""
287 for i in xrange(120*1024):
288 h.update('1234567890')
289 pieces = h.pieceDigests()
290 self.failUnless(h.digest() == '1(j\xd2q\x0b\n\x91\xd2\x13\x90\x15\xa3E\xcc\xb0\x8d.\xc3\xc5')
291 self.failUnless(len(pieces) == 3)
292 self.failUnless(pieces[0] == ',G \xd8\xbbPl\xf1\xa3\xa0\x0cW\n\xe6\xe6a\xc9\x95/\xe5')
293 self.failUnless(pieces[1] == '\xf6V\xeb/\xa8\xad[\x07Z\xf9\x87\xa4\xf5w\xdf\xe1|\x00\x8e\x93')
294 self.failUnless(pieces[2] == 'M[\xbf\xee\xaa+\x19\xbaV\xf699\r\x17o\xcb\x8e\xcfP\x19')
296 def test_pieces_at_once(self):
297 """Tests the updating of multiple pieces at once."""
300 h.update('1234567890'*120*1024)
301 self.failUnless(h.digest() == '1(j\xd2q\x0b\n\x91\xd2\x13\x90\x15\xa3E\xcc\xb0\x8d.\xc3\xc5')
302 pieces = h.pieceDigests()
303 self.failUnless(len(pieces) == 3)
304 self.failUnless(pieces[0] == ',G \xd8\xbbPl\xf1\xa3\xa0\x0cW\n\xe6\xe6a\xc9\x95/\xe5')
305 self.failUnless(pieces[1] == '\xf6V\xeb/\xa8\xad[\x07Z\xf9\x87\xa4\xf5w\xdf\xe1|\x00\x8e\x93')
306 self.failUnless(pieces[2] == 'M[\xbf\xee\xaa+\x19\xbaV\xf699\r\x17o\xcb\x8e\xcfP\x19')
308 def test_pieces_boundaries(self):
309 """Tests the updating exactly to piece boundaries."""
312 h.update('1234567890'*52428)
314 assert h.size % PIECE_SIZE == 0
316 h.update('1234567890'*52428)
318 assert h.size % PIECE_SIZE == 0
320 h.update('1234567890'*18022)
321 assert h.size == 10*120*1024
322 pieces = h.pieceDigests()
323 self.failUnless(h.digest() == '1(j\xd2q\x0b\n\x91\xd2\x13\x90\x15\xa3E\xcc\xb0\x8d.\xc3\xc5')
324 self.failUnless(len(pieces) == 3)
325 self.failUnless(pieces[0] == ',G \xd8\xbbPl\xf1\xa3\xa0\x0cW\n\xe6\xe6a\xc9\x95/\xe5')
326 self.failUnless(pieces[1] == '\xf6V\xeb/\xa8\xad[\x07Z\xf9\x87\xa4\xf5w\xdf\xe1|\x00\x8e\x93')
327 self.failUnless(pieces[2] == 'M[\xbf\xee\xaa+\x19\xbaV\xf699\r\x17o\xcb\x8e\xcfP\x19')
329 def test_pieces_other_hashes(self):
330 """Tests updating of pieces a little at a time."""
332 for hashType in h.ORDER:
333 if hashType['name'] != 'sha1':
334 h.hashTypeNum = h.ORDER.index(hashType)
336 assert h.ORDER[h.hashTypeNum]['name'] != 'sha1'
338 for i in xrange(120*1024):
339 h.update('1234567890')
340 pieces = h.pieceDigests()
341 self.failUnless(len(pieces) == 3)
342 self.failUnless(pieces[0] == ',G \xd8\xbbPl\xf1\xa3\xa0\x0cW\n\xe6\xe6a\xc9\x95/\xe5')
343 self.failUnless(pieces[1] == '\xf6V\xeb/\xa8\xad[\x07Z\xf9\x87\xa4\xf5w\xdf\xe1|\x00\x8e\x93')
344 self.failUnless(pieces[2] == 'M[\xbf\xee\xaa+\x19\xbaV\xf699\r\x17o\xcb\x8e\xcfP\x19')
347 """Test hashing using the SHA1 hash."""
350 for hashType in h.ORDER:
351 if hashType['name'] == 'sha1':
354 self.failUnless(found == True)
355 h.set(hashType, '3bba0a5d97b7946ad2632002bf9caefe2cb18e00', '19')
357 h.update('apt-p2p is the best')
358 self.failUnless(h.hexdigest() == '3bba0a5d97b7946ad2632002bf9caefe2cb18e00')
359 self.failUnlessRaises(HashError, h.update, 'gfgf')
360 self.failUnless(h.verify() == True)
363 """Test hashing using the MD5 hash."""
366 for hashType in h.ORDER:
367 if hashType['name'] == 'md5':
370 self.failUnless(found == True)
371 h.set(hashType, '6b5abdd30d7ed80edd229f9071d8c23c', '19')
373 h.update('apt-p2p is the best')
374 self.failUnless(h.hexdigest() == '6b5abdd30d7ed80edd229f9071d8c23c')
375 self.failUnlessRaises(HashError, h.update, 'gfgf')
376 self.failUnless(h.verify() == True)
378 def test_sha256(self):
379 """Test hashing using the SHA256 hash."""
382 for hashType in h.ORDER:
383 if hashType['name'] == 'sha256':
386 self.failUnless(found == True)
387 h.set(hashType, '47f2238a30a0340faa2bf01a9bdc42ba77b07b411cda1e24cd8d7b5c4b7d82a7', '19')
389 h.update('apt-p2p is the best')
390 self.failUnless(h.hexdigest() == '47f2238a30a0340faa2bf01a9bdc42ba77b07b411cda1e24cd8d7b5c4b7d82a7')
391 self.failUnlessRaises(HashError, h.update, 'gfgf')
392 self.failUnless(h.verify() == True)
394 if sys.version_info < (2, 5):
395 test_sha256.skip = "SHA256 hashes are not supported by Python until version 2.5"