2 from binascii import b2a_hex, a2b_hex
5 from twisted.internet import threads, defer
6 from twisted.trial import unittest
8 class HashError(ValueError):
9 """An error has occurred while hashing a file."""
12 """Manages hashes and hashing for a file."""
14 """The priority ordering of hashes, and how to extract them."""
15 ORDER = [ {'name': 'sha1',
16 'AptPkgRecord': 'SHA1Hash',
17 'AptSrcRecord': False,
18 'AptIndexRecord': 'SHA1',
20 'hashlib_func': 'sha1',
23 'AptPkgRecord': 'SHA256Hash',
24 'AptSrcRecord': False,
25 'AptIndexRecord': 'SHA256',
26 'hashlib_func': 'sha256',
29 'AptPkgRecord': 'MD5Hash',
31 'AptIndexRecord': 'MD5Sum',
33 'hashlib_func': 'md5',
38 self.hashTypeNum = 0 # Use the first if nothing else matters
42 self.expNormHash = None
43 self.fileHasher = None
46 self.fileNormHash = None
49 if sys.version_info < (2, 5):
50 # sha256 is not available in python before 2.5, remove it
51 for hashType in self.ORDER:
52 if hashType['name'] == 'sha256':
53 del self.ORDER[self.ORDER.index(hashType)]
56 def _norm_hash(self, hashString, bits=None, bytes=None):
58 bytes = (bits - 1) // 8 + 1
61 raise HashError, "you must specify one of bits or bytes"
62 if len(hashString) < bytes:
63 hashString = hashString + '\000'*(bytes - len(hashString))
64 elif len(hashString) > bytes:
65 hashString = hashString[:bytes]
68 #### Methods for returning the expected hash
70 """Get the expected hash."""
73 def hexexpected(self):
74 """Get the expected hash in hex format."""
75 if self.expHex is None and self.expHash is not None:
76 self.expHex = b2a_hex(self.expHash)
79 def normexpected(self, bits=None, bytes=None):
80 """Normalize the binary hash for the given length.
82 You must specify one of bits or bytes.
84 if self.expNormHash is None and self.expHash is not None:
85 self.expNormHash = self._norm_hash(self.expHash, bits, bytes)
86 return self.expNormHash
88 #### Methods for hashing data
89 def new(self, force = False):
90 """Generate a new hashing object suitable for hashing a file.
92 @param force: set to True to force creating a new hasher even if
93 the hash has been verified already
95 if self.result is None or force == True:
99 if sys.version_info < (2, 5):
100 mod = __import__(self.ORDER[self.hashTypeNum]['old_module'], globals(), locals(), [])
101 self.fileHasher = mod.new()
104 func = getattr(hashlib, self.ORDER[self.hashTypeNum]['hashlib_func'])
105 self.fileHasher = func()
107 def update(self, data):
108 """Add more data to the file hasher."""
109 if self.result is None:
111 raise HashError, "Already done, you can't add more data after calling digest() or verify()"
112 if self.fileHasher is None:
113 raise HashError, "file hasher not initialized"
114 self.fileHasher.update(data)
115 self.size += len(data)
118 """Get the hash of the added file data."""
119 if self.fileHash is None:
120 if self.fileHasher is None:
121 raise HashError, "you must hash some data first"
122 self.fileHash = self.fileHasher.digest()
127 """Get the hash of the added file data in hex format."""
128 if self.fileHex is None:
129 self.fileHex = b2a_hex(self.digest())
132 def norm(self, bits=None, bytes=None):
133 """Normalize the binary hash for the given length.
135 You must specify one of bits or bytes.
137 if self.fileNormHash is None:
138 self.fileNormHash = self._norm_hash(self.digest(), bits, bytes)
139 return self.fileNormHash
142 """Verify that the added file data hash matches the expected hash."""
143 if self.result is None and self.fileHash is not None and self.expHash is not None:
144 self.result = (self.fileHash == self.expHash and self.size == self.expSize)
147 def hashInThread(self, file):
148 """Hashes a file in a separate thread, callback with the result."""
150 if not file.exists():
151 df = defer.Deferred()
152 df.errback(HashError("file not found"))
155 df = threads.deferToThread(self._hashInThread, file)
158 def _hashInThread(self, file):
159 """Hashes a file, returning itself as the result."""
161 self.new(force = True)
169 #### Methods for setting the expected hash
170 def set(self, hashType, hashHex, size):
171 """Initialize the hash object.
173 @param hashType: must be one of the dictionaries from L{ORDER}
175 self.hashTypeNum = self.ORDER.index(hashType) # error if not found
176 self.expHex = hashHex
177 self.expSize = int(size)
178 self.expHash = a2b_hex(self.expHex)
180 def setFromIndexRecord(self, record):
181 """Set the hash from the cache of index file records.
183 @type record: C{dictionary}
184 @param record: keys are hash types, values are tuples of (hash, size)
186 for hashType in self.ORDER:
187 result = record.get(hashType['AptIndexRecord'], None)
189 self.set(hashType, result[0], result[1])
193 def setFromPkgRecord(self, record, size):
194 """Set the hash from Apt's binary packages cache.
196 @param record: whatever is returned by apt_pkg.GetPkgRecords()
198 for hashType in self.ORDER:
199 hashHex = getattr(record, hashType['AptPkgRecord'], None)
201 self.set(hashType, hashHex, size)
205 def setFromSrcRecord(self, record):
206 """Set the hash from Apt's source package records cache.
208 Currently very simple since Apt only tracks MD5 hashes of source files.
210 @type record: (C{string}, C{int}, C{string})
211 @param record: the hash, size and path of the source file
213 for hashType in self.ORDER:
214 if hashType['AptSrcRecord']:
215 self.set(hashType, record[0], record[1])
219 class TestHashObject(unittest.TestCase):
220 """Unit tests for the hash objects."""
223 if sys.version_info < (2, 4):
226 def test_normalize(self):
228 h.set(h.ORDER[0], b2a_hex('12345678901234567890'), '0')
229 self.failUnless(h.normexpected(bits = 160) == '12345678901234567890')
231 h.set(h.ORDER[0], b2a_hex('12345678901234567'), '0')
232 self.failUnless(h.normexpected(bits = 160) == '12345678901234567\000\000\000')
234 h.set(h.ORDER[0], b2a_hex('1234567890123456789012345'), '0')
235 self.failUnless(h.normexpected(bytes = 20) == '12345678901234567890')
237 h.set(h.ORDER[0], b2a_hex('1234567890123456789'), '0')
238 self.failUnless(h.normexpected(bytes = 20) == '1234567890123456789\000')
240 h.set(h.ORDER[0], b2a_hex('123456789012345678901'), '0')
241 self.failUnless(h.normexpected(bits = 160) == '12345678901234567890')
243 def test_failure(self):
245 h.set(h.ORDER[0], b2a_hex('12345678901234567890'), '0')
246 self.failUnlessRaises(HashError, h.normexpected)
247 self.failUnlessRaises(HashError, h.digest)
248 self.failUnlessRaises(HashError, h.hexdigest)
249 self.failUnlessRaises(HashError, h.update, 'gfgf')
254 for hashType in h.ORDER:
255 if hashType['name'] == 'sha1':
258 self.failUnless(found == True)
259 h.set(hashType, 'c722df87e1acaa64b27aac4e174077afc3623540', '19')
261 h.update('apt-dht is the best')
262 self.failUnless(h.hexdigest() == 'c722df87e1acaa64b27aac4e174077afc3623540')
263 self.failUnlessRaises(HashError, h.update, 'gfgf')
264 self.failUnless(h.verify() == True)
269 for hashType in h.ORDER:
270 if hashType['name'] == 'md5':
273 self.failUnless(found == True)
274 h.set(hashType, '2a586bcd1befc5082c872dcd96a01403', '19')
276 h.update('apt-dht is the best')
277 self.failUnless(h.hexdigest() == '2a586bcd1befc5082c872dcd96a01403')
278 self.failUnlessRaises(HashError, h.update, 'gfgf')
279 self.failUnless(h.verify() == True)
281 def test_sha256(self):
284 for hashType in h.ORDER:
285 if hashType['name'] == 'sha256':
288 self.failUnless(found == True)
289 h.set(hashType, '55b971f64d9772f733de03f23db39224f51a455cc5ad4c2db9d5740d2ab259a7', '19')
291 h.update('apt-dht is the best')
292 self.failUnless(h.hexdigest() == '55b971f64d9772f733de03f23db39224f51a455cc5ad4c2db9d5740d2ab259a7')
293 self.failUnlessRaises(HashError, h.update, 'gfgf')
294 self.failUnless(h.verify() == True)
296 if sys.version_info < (2, 5):
297 test_sha256.skip = "SHA256 hashes are not supported by Python until version 2.5"