2 from binascii import b2a_hex, a2b_hex
5 from twisted.trial import unittest
8 """Manages hashes and hashing for a file."""
10 """The priority ordering of hashes, and how to extract them."""
11 ORDER = [ {'name': 'sha1',
12 'AptPkgRecord': 'SHA1Hash',
13 'AptSrcRecord': False,
14 'AptIndexRecord': 'SHA1',
16 'hashlib_func': 'sha1',
19 'AptPkgRecord': 'SHA256Hash',
20 'AptSrcRecord': False,
21 'AptIndexRecord': 'SHA256',
22 'hashlib_func': 'sha256',
25 'AptPkgRecord': 'MD5Hash',
27 'AptIndexRecord': 'MD5Sum',
29 'hashlib_func': 'md5',
34 self.hashTypeNum = 0 # Use the first if nothing else matters
38 self.expNormHash = None
39 self.fileHasher = None
42 self.fileNormHash = None
45 if sys.version_info < (2, 5):
46 # sha256 is not available in python before 2.5, remove it
47 for hashType in self.ORDER:
48 if hashType['name'] == 'sha256':
49 del self.ORDER[self.ORDER.index(hashType)]
52 def _norm_hash(self, hashString, bits=None, bytes=None):
54 bytes = (bits - 1) // 8 + 1
56 assert bytes is not None, "you must specify one of bits or bytes"
57 if len(hashString) < bytes:
58 hashString = hashString + '\000'*(bytes - len(hashString))
59 elif len(hashString) > bytes:
60 hashString = hashString[:bytes]
63 #### Methods for returning the expected hash
65 """Get the expected hash."""
68 def hexexpected(self):
69 """Get the expected hash in hex format."""
70 if self.expHex is None and self.expHash is not None:
71 self.expHex = b2a_hex(self.expHash)
74 def normexpected(self, bits=None, bytes=None):
75 """Normalize the binary hash for the given length.
77 You must specify one of bits or bytes.
79 if self.expNormHash is None and self.expHash is not None:
80 self.expNormHash = self._norm_hash(self.expHash, bits, bytes)
81 return self.expNormHash
83 #### Methods for hashing data
84 def new(self, force = False):
85 """Generate a new hashing object suitable for hashing a file.
87 @param force: set to True to force creating a new hasher even if
88 the hash has been verified already
90 if self.result is None or force == True:
94 if sys.version_info < (2, 5):
95 mod = __import__(self.ORDER[self.hashTypeNum]['old_module'], globals(), locals(), [])
96 self.fileHasher = mod.new()
99 func = getattr(hashlib, self.ORDER[self.hashTypeNum]['hashlib_func'])
100 self.fileHasher = func()
102 def update(self, data):
103 """Add more data to the file hasher."""
104 if self.result is None:
105 assert self.done == False, "Already done, you can't add more data after calling digest() or verify()"
106 assert self.fileHasher is not None, "file hasher not initialized"
107 self.fileHasher.update(data)
108 self.size += len(data)
111 """Get the hash of the added file data."""
112 if self.fileHash is None:
113 assert self.fileHasher is not None, "you must hash some data first"
114 self.fileHash = self.fileHasher.digest()
119 """Get the hash of the added file data in hex format."""
120 if self.fileHex is None:
121 self.fileHex = b2a_hex(self.digest())
124 def norm(self, bits=None, bytes=None):
125 """Normalize the binary hash for the given length.
127 You must specify one of bits or bytes.
129 if self.fileNormHash is None:
130 self.fileNormHash = self._norm_hash(self.digest(), bits, bytes)
131 return self.fileNormHash
134 """Verify that the added file data hash matches the expected hash."""
135 if self.result is None and self.fileHash is not None and self.expHash is not None:
136 self.result = (self.fileHash == self.expHash and self.size == self.expSize)
139 #### Methods for setting the expected hash
140 def set(self, hashType, hashHex, size):
141 """Initialize the hash object.
143 @param hashType: must be one of the dictionaries from L{ORDER}
145 self.hashTypeNum = self.ORDER.index(hashType) # error if not found
146 self.expHex = hashHex
147 self.expSize = int(size)
148 self.expHash = a2b_hex(self.expHex)
150 def setFromIndexRecord(self, record):
151 """Set the hash from the cache of index file records.
153 @type record: C{dictionary}
154 @param record: keys are hash types, values are tuples of (hash, size)
156 for hashType in self.ORDER:
157 result = record.get(hashType['AptIndexRecord'], None)
159 self.set(hashType, result[0], result[1])
163 def setFromPkgRecord(self, record, size):
164 """Set the hash from Apt's binary packages cache.
166 @param record: whatever is returned by apt_pkg.GetPkgRecords()
168 for hashType in self.ORDER:
169 hashHex = getattr(record, hashType['AptPkgRecord'], None)
171 self.set(hashType, hashHex, size)
175 def setFromSrcRecord(self, record):
176 """Set the hash from Apt's source package records cache.
178 Currently very simple since Apt only tracks MD5 hashes of source files.
180 @type record: (C{string}, C{int}, C{string})
181 @param record: the hash, size and path of the source file
183 for hashType in self.ORDER:
184 if hashType['AptSrcRecord']:
185 self.set(hashType, record[0], record[1])
189 class TestHashObject(unittest.TestCase):
190 """Unit tests for the hash objects."""
193 if sys.version_info < (2, 4):
196 def test_normalize(self):
198 h.set(h.ORDER[0], b2a_hex('12345678901234567890'), '0')
199 self.failUnless(h.normexpected(bits = 160) == '12345678901234567890')
201 h.set(h.ORDER[0], b2a_hex('12345678901234567'), '0')
202 self.failUnless(h.normexpected(bits = 160) == '12345678901234567\000\000\000')
204 h.set(h.ORDER[0], b2a_hex('1234567890123456789012345'), '0')
205 self.failUnless(h.normexpected(bytes = 20) == '12345678901234567890')
207 h.set(h.ORDER[0], b2a_hex('1234567890123456789'), '0')
208 self.failUnless(h.normexpected(bytes = 20) == '1234567890123456789\000')
210 h.set(h.ORDER[0], b2a_hex('123456789012345678901'), '0')
211 self.failUnless(h.normexpected(bits = 160) == '12345678901234567890')
213 def test_failure(self):
215 h.set(h.ORDER[0], b2a_hex('12345678901234567890'), '0')
216 self.failUnlessRaises(AssertionError, h.normexpected)
217 self.failUnlessRaises(AssertionError, h.digest)
218 self.failUnlessRaises(AssertionError, h.hexdigest)
219 self.failUnlessRaises(AssertionError, h.update, 'gfgf')
224 for hashType in h.ORDER:
225 if hashType['name'] == 'sha1':
228 self.failUnless(found == True)
229 h.set(hashType, 'c722df87e1acaa64b27aac4e174077afc3623540', '19')
231 h.update('apt-dht is the best')
232 self.failUnless(h.hexdigest() == 'c722df87e1acaa64b27aac4e174077afc3623540')
233 self.failUnlessRaises(AssertionError, h.update, 'gfgf')
234 self.failUnless(h.verify() == True)
239 for hashType in h.ORDER:
240 if hashType['name'] == 'md5':
243 self.failUnless(found == True)
244 h.set(hashType, '2a586bcd1befc5082c872dcd96a01403', '19')
246 h.update('apt-dht is the best')
247 self.failUnless(h.hexdigest() == '2a586bcd1befc5082c872dcd96a01403')
248 self.failUnlessRaises(AssertionError, h.update, 'gfgf')
249 self.failUnless(h.verify() == True)
251 def test_sha256(self):
254 for hashType in h.ORDER:
255 if hashType['name'] == 'sha256':
258 self.failUnless(found == True)
259 h.set(hashType, '55b971f64d9772f733de03f23db39224f51a455cc5ad4c2db9d5740d2ab259a7', '19')
261 h.update('apt-dht is the best')
262 self.failUnless(h.hexdigest() == '55b971f64d9772f733de03f23db39224f51a455cc5ad4c2db9d5740d2ab259a7')
263 self.failUnlessRaises(AssertionError, h.update, 'gfgf')
264 self.failUnless(h.verify() == True)
266 if sys.version_info < (2, 5):
267 test_sha256.skip = "SHA256 hashes are not supported by Python until version 2.5"