Use python-debian for parsing RFC 822 files (untested).
[quix0rs-apt-p2p.git] / apt_dht / Hash.py
1
2 from binascii import b2a_hex, a2b_hex
3 import sys
4
5 from twisted.internet import threads, defer
6 from twisted.trial import unittest
7
8 class HashError(ValueError):
9     """An error has occurred while hashing a file."""
10     
11 class HashObject:
12     """Manages hashes and hashing for a file."""
13     
14     """The priority ordering of hashes, and how to extract them."""
15     ORDER = [ {'name': 'sha1', 
16                    'AptPkgRecord': 'SHA1Hash', 
17                    'AptSrcRecord': False, 
18                    'AptIndexRecord': 'SHA1',
19                    'old_module': 'sha',
20                    'hashlib_func': 'sha1',
21                    },
22               {'name': 'sha256',
23                    'AptPkgRecord': 'SHA256Hash', 
24                    'AptSrcRecord': False, 
25                    'AptIndexRecord': 'SHA256',
26                    'hashlib_func': 'sha256',
27                    },
28               {'name': 'md5',
29                    'AptPkgRecord': 'MD5Hash', 
30                    'AptSrcRecord': True, 
31                    'AptIndexRecord': 'MD5SUM',
32                    'old_module': 'md5',
33                    'hashlib_func': 'md5',
34                    },
35             ]
36     
37     def __init__(self):
38         self.hashTypeNum = 0    # Use the first if nothing else matters
39         self.expHash = None
40         self.expHex = None
41         self.expSize = None
42         self.expNormHash = None
43         self.fileHasher = None
44         self.fileHash = None
45         self.fileHex = None
46         self.fileNormHash = None
47         self.done = True
48         self.result = 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)]
54                     break
55         
56     def _norm_hash(self, hashString, bits=None, bytes=None):
57         if bits is not None:
58             bytes = (bits - 1) // 8 + 1
59         else:
60             if bytes is None:
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]
66         return hashString
67
68     #### Methods for returning the expected hash
69     def expected(self):
70         """Get the expected hash."""
71         return self.expHash
72     
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)
77         return self.expHex
78     
79     def normexpected(self, bits=None, bytes=None):
80         """Normalize the binary hash for the given length.
81         
82         You must specify one of bits or bytes.
83         """
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
87
88     #### Methods for hashing data
89     def new(self, force = False):
90         """Generate a new hashing object suitable for hashing a file.
91         
92         @param force: set to True to force creating a new hasher even if
93             the hash has been verified already
94         """
95         if self.result is None or force == True:
96             self.result = None
97             self.size = 0
98             self.done = False
99             if sys.version_info < (2, 5):
100                 mod = __import__(self.ORDER[self.hashTypeNum]['old_module'], globals(), locals(), [])
101                 self.fileHasher = mod.new()
102             else:
103                 import hashlib
104                 func = getattr(hashlib, self.ORDER[self.hashTypeNum]['hashlib_func'])
105                 self.fileHasher = func()
106
107     def update(self, data):
108         """Add more data to the file hasher."""
109         if self.result is None:
110             if self.done:
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)
116         
117     def digest(self):
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()
123             self.done = True
124         return self.fileHash
125
126     def hexdigest(self):
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())
130         return self.fileHex
131         
132     def norm(self, bits=None, bytes=None):
133         """Normalize the binary hash for the given length.
134         
135         You must specify one of bits or bytes.
136         """
137         if self.fileNormHash is None:
138             self.fileNormHash = self._norm_hash(self.digest(), bits, bytes)
139         return self.fileNormHash
140
141     def verify(self):
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)
145         return self.result
146     
147     def hashInThread(self, file):
148         """Hashes a file in a separate thread, callback with the result."""
149         file.restat(False)
150         if not file.exists():
151             df = defer.Deferred()
152             df.errback(HashError("file not found"))
153             return df
154         
155         df = threads.deferToThread(self._hashInThread, file)
156         return df
157     
158     def _hashInThread(self, file):
159         """Hashes a file, returning itself as the result."""
160         f = file.open()
161         self.new(force = True)
162         data = f.read(4096)
163         while data:
164             self.update(data)
165             data = f.read(4096)
166         self.digest()
167         return self
168
169     #### Methods for setting the expected hash
170     def set(self, hashType, hashHex, size):
171         """Initialize the hash object.
172         
173         @param hashType: must be one of the dictionaries from L{ORDER}
174         """
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)
179         
180     def setFromIndexRecord(self, record):
181         """Set the hash from the cache of index file records.
182         
183         @type record: C{dictionary}
184         @param record: keys are hash types, values are tuples of (hash, size)
185         """
186         for hashType in self.ORDER:
187             result = record.get(hashType['AptIndexRecord'], None)
188             if result:
189                 self.set(hashType, result[0], result[1])
190                 return True
191         return False
192
193     def setFromPkgRecord(self, record, size):
194         """Set the hash from Apt's binary packages cache.
195         
196         @param record: whatever is returned by apt_pkg.GetPkgRecords()
197         """
198         for hashType in self.ORDER:
199             hashHex = getattr(record, hashType['AptPkgRecord'], None)
200             if hashHex:
201                 self.set(hashType, hashHex, size)
202                 return True
203         return False
204     
205     def setFromSrcRecord(self, record):
206         """Set the hash from Apt's source package records cache.
207         
208         Currently very simple since Apt only tracks MD5 hashes of source files.
209         
210         @type record: (C{string}, C{int}, C{string})
211         @param record: the hash, size and path of the source file
212         """
213         for hashType in self.ORDER:
214             if hashType['AptSrcRecord']:
215                 self.set(hashType, record[0], record[1])
216                 return True
217         return False
218
219 class TestHashObject(unittest.TestCase):
220     """Unit tests for the hash objects."""
221     
222     timeout = 5
223     if sys.version_info < (2, 4):
224         skip = "skippingme"
225     
226     def test_normalize(self):
227         h = HashObject()
228         h.set(h.ORDER[0], b2a_hex('12345678901234567890'), '0')
229         self.failUnless(h.normexpected(bits = 160) == '12345678901234567890')
230         h = HashObject()
231         h.set(h.ORDER[0], b2a_hex('12345678901234567'), '0')
232         self.failUnless(h.normexpected(bits = 160) == '12345678901234567\000\000\000')
233         h = HashObject()
234         h.set(h.ORDER[0], b2a_hex('1234567890123456789012345'), '0')
235         self.failUnless(h.normexpected(bytes = 20) == '12345678901234567890')
236         h = HashObject()
237         h.set(h.ORDER[0], b2a_hex('1234567890123456789'), '0')
238         self.failUnless(h.normexpected(bytes = 20) == '1234567890123456789\000')
239         h = HashObject()
240         h.set(h.ORDER[0], b2a_hex('123456789012345678901'), '0')
241         self.failUnless(h.normexpected(bits = 160) == '12345678901234567890')
242
243     def test_failure(self):
244         h = HashObject()
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')
250     
251     def test_sha1(self):
252         h = HashObject()
253         found = False
254         for hashType in h.ORDER:
255             if hashType['name'] == 'sha1':
256                 found = True
257                 break
258         self.failUnless(found == True)
259         h.set(hashType, 'c722df87e1acaa64b27aac4e174077afc3623540', '19')
260         h.new()
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)
265         
266     def test_md5(self):
267         h = HashObject()
268         found = False
269         for hashType in h.ORDER:
270             if hashType['name'] == 'md5':
271                 found = True
272                 break
273         self.failUnless(found == True)
274         h.set(hashType, '2a586bcd1befc5082c872dcd96a01403', '19')
275         h.new()
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)
280         
281     def test_sha256(self):
282         h = HashObject()
283         found = False
284         for hashType in h.ORDER:
285             if hashType['name'] == 'sha256':
286                 found = True
287                 break
288         self.failUnless(found == True)
289         h.set(hashType, '55b971f64d9772f733de03f23db39224f51a455cc5ad4c2db9d5740d2ab259a7', '19')
290         h.new()
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)
295
296     if sys.version_info < (2, 5):
297         test_sha256.skip = "SHA256 hashes are not supported by Python until version 2.5"