Prevent BitTorrent nodes from infiltrating the DHT.
[quix0rs-apt-p2p.git] / apt_p2p_Khashmir / util.py
1
2 """Some utitlity functions for use in apt-p2p's khashmir DHT."""
3
4 from twisted.trial import unittest
5
6 from khash import HASH_LENGTH
7
8 def bucket_stats(l):
9     """Given a list of khashmir instances, finds min, max, and average number of nodes in tables."""
10     max = avg = 0
11     min = None
12     def count(buckets):
13         c = 0
14         for bucket in buckets:
15             c = c + bucket.len()
16         return c
17     for node in l:
18         c = count(node.table.buckets)
19         if min == None:
20             min = c
21         elif c < min:
22             min = c
23         if c > max:
24             max = c
25         avg = avg + c
26     avg = avg / len(l)
27     return {'min':min, 'max':max, 'avg':avg}
28
29 def uncompact(s):
30     """Extract the contact info from a compact node representation.
31     
32     @type s: C{string}
33     @param s: the compact representation
34     @rtype: C{dictionary}
35     @return: the node ID, IP address and port to contact the node on
36     @raise ValueError: if the compact representation doesn't exist
37     """
38     if (len(s) != HASH_LENGTH+6):
39         raise ValueError, "not compact node info: %r" % (s, )
40     id = s[:HASH_LENGTH]
41     host = '.'.join([str(ord(i)) for i in s[HASH_LENGTH:(HASH_LENGTH+4)]])
42     port = (ord(s[HASH_LENGTH+4]) << 8) | ord(s[HASH_LENGTH+5])
43     return {'id': id, 'host': host, 'port': port}
44
45 def compact(id, host, port):
46     """Create a compact representation of node contact info.
47     
48     @type id: C{string}
49     @param id: the node ID
50     @type host: C{string}
51     @param host: the IP address of the node
52     @type port: C{int}
53     @param port: the port number to contact the node on
54     @rtype: C{string}
55     @return: the compact representation
56     @raise ValueError: if the compact representation doesn't exist
57     """
58     
59     s = id + ''.join([chr(int(i)) for i in host.split('.')]) + \
60           chr((port & 0xFF00) >> 8) + chr(port & 0xFF)
61     if len(s) != 26:
62         raise ValueError
63     return s
64
65 def byte_format(s):
66     """Format a byte size for reading by the user.
67     
68     @type s: C{long}
69     @param s: the number of bytes
70     @rtype: C{string}
71     @return: the formatted size with appropriate units
72     """
73     if (s < 1):
74         r = str(int(s*1000.0)/1000.0) + 'B'
75     elif (s < 10):
76         r = str(int(s*100.0)/100.0) + 'B'
77     elif (s < 102):
78         r = str(int(s*10.0)/10.0) + 'B'
79     elif (s < 1024):
80         r = str(int(s)) + 'B'
81     elif (s < 10485):
82         r = str(int((s/1024.0)*100.0)/100.0) + 'KiB'
83     elif (s < 104857):
84         r = str(int((s/1024.0)*10.0)/10.0) + 'KiB'
85     elif (s < 1048576):
86         r = str(int(s/1024)) + 'KiB'
87     elif (s < 10737418L):
88         r = str(int((s/1048576.0)*100.0)/100.0) + 'MiB'
89     elif (s < 107374182L):
90         r = str(int((s/1048576.0)*10.0)/10.0) + 'MiB'
91     elif (s < 1073741824L):
92         r = str(int(s/1048576)) + 'MiB'
93     elif (s < 1099511627776L):
94         r = str(int((s/1073741824.0)*100.0)/100.0) + 'GiB'
95     else:
96         r = str(int((s/1099511627776.0)*100.0)/100.0) + 'TiB'
97     return(r)
98
99 class TestUtil(unittest.TestCase):
100     """Tests for the utilities."""
101     
102     timeout = 5
103     myid = '\xca\xec\xb8\x0c\x00\xe7\x07\xf8~])\x8f\x9d\xe5_B\xff\x1a\xc4!'
104     host = '165.234.1.34'
105     port = 61234
106
107     def test_compact(self):
108         d = uncompact(compact(self.myid, self.host, self.port))
109         self.failUnlessEqual(d['id'], self.myid)
110         self.failUnlessEqual(d['host'], self.host)
111         self.failUnlessEqual(d['port'], self.port)
112