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"