Allow the setting of khashmir ID's suffix.
[quix0rs-apt-p2p.git] / apt_p2p_Khashmir / khash.py
1
2 """Functions to deal with hashes (node IDs and keys).
3
4 @var HASH_LENGTH: the length of the hash to use in bytes
5 """
6
7 from sha import sha
8 from os import urandom
9
10 from twisted.trial import unittest
11
12 HASH_LENGTH = 20
13
14 def intify(hstr):
15     """Convert a hash (big-endian) to a long python integer."""
16     assert len(hstr) == HASH_LENGTH
17     return long(hstr.encode('hex'), 16)
18
19 def stringify(num):
20     """Convert a long python integer to a hash."""
21     str = hex(num)[2:]
22     if str[-1] == 'L':
23         str = str[:-1]
24     if len(str) % 2 != 0:
25         str = '0' + str
26     str = str.decode('hex')
27     return (HASH_LENGTH - len(str)) *'\x00' + str
28     
29 def distance(a, b):
30     """Calculate the distance between two hashes expressed as strings."""
31     return intify(a) ^ intify(b)
32
33 def newID(suffix = ''):
34     """Get a new pseudorandom globally unique hash string.
35     
36     @param suffix: an optional string to end the ID with
37     """
38     assert len(suffix) < 20, 'The suffix must be shorter than the SHA1 hash'
39     h = sha()
40     h.update(urandom(HASH_LENGTH))
41     if suffix:
42         return h.digest()[:-len(suffix)] + suffix
43     else:
44         return h.digest()
45
46 def newIDInRange(min, max):
47     """Get a new pseudorandom globally unique hash string in the range."""
48     return stringify(randRange(min,max))
49     
50 def randRange(min, max):
51     """Get a new pseudorandom globally unique hash number in the range."""
52     return min + intify(newID()) % (max - min)
53     
54 def newTID():
55     """Get a new pseudorandom transaction ID number."""
56     return randRange(-2**30, 2**30)
57
58 class TestNewID(unittest.TestCase):
59     """Test the newID function."""
60     def testLength(self):
61         self.failUnlessEqual(len(newID()), HASH_LENGTH)
62     def testHundreds(self):
63         for x in xrange(100):
64             self.testLength()
65     def testSuffix(self):
66         id = newID('T012')
67         self.failUnlessEqual(len(id), HASH_LENGTH)
68         self.failUnlessEqual(id[-4:], 'T012')
69
70 class TestIntify(unittest.TestCase):
71     """Test the intify function."""
72     known = [('\0' * HASH_LENGTH, 0),
73             ('\xff' * HASH_LENGTH, 2L**(HASH_LENGTH*8) - 1),
74             ]
75     def testKnown(self):
76         for str, value in self.known: 
77             self.failUnlessEqual(intify(str),  value)
78     def testEndianessOnce(self):
79         h = newID()
80         while h[-1] == '\xff':
81             h = newID()
82         k = h[:-1] + chr(ord(h[-1]) + 1)
83         self.failUnlessEqual(intify(k) - intify(h), 1)
84     def testEndianessLots(self):
85         for x in xrange(100):
86             self.testEndianessOnce()
87
88 class TestDisantance(unittest.TestCase):
89     """Test the distance function."""
90     known = [
91             (("\0" * HASH_LENGTH, "\xff" * HASH_LENGTH), 2L**(HASH_LENGTH*8) -1),
92             ((sha("foo").digest(), sha("foo").digest()), 0),
93             ((sha("bar").digest(), sha("bar").digest()), 0)
94             ]
95     def testKnown(self):
96         for pair, dist in self.known:
97             self.failUnlessEqual(distance(pair[0], pair[1]), dist)
98     def testCommutitive(self):
99         for i in xrange(100):
100             x, y, z = newID(), newID(), newID()
101             self.failUnlessEqual(distance(x,y) ^ distance(y, z), distance(x, z))
102         
103 class TestRandRange(unittest.TestCase):
104     """Test the randRange function."""
105     def testOnce(self):
106         a = intify(newID())
107         b = intify(newID())
108         if a < b:
109             c = randRange(a, b)
110             self.failUnlessEqual(a <= c < b, True, "output out of range %d  %d  %d" % (b, c, a))
111         else:
112             c = randRange(b, a)
113             self.failUnlessEqual(b <= c < a, True, "output out of range %d  %d  %d" % (b, c, a))
114
115     def testOneHundredTimes(self):
116         for i in xrange(100):
117             self.testOnce()