WIP on final version of accepted INFOCOM paper.
[quix0rs-apt-p2p.git] / apt_p2p / util.py
1
2 """Some utitlity functions for use in the apt-p2p program.
3
4 @var isLocal: a compiled regular expression suitable for testing if an
5     IP address is from a known local or private range
6 """
7
8 import os, re
9
10 from twisted.python import log
11 from twisted.trial import unittest
12
13 isLocal = re.compile('^(192\.168\.[0-9]{1,3}\.[0-9]{1,3})|'+
14                      '(10\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})|'+
15                      '(172\.0?1[6-9]\.[0-9]{1,3}\.[0-9]{1,3})|'+
16                      '(172\.0?2[0-9]\.[0-9]{1,3}\.[0-9]{1,3})|'+
17                      '(172\.0?3[0-1]\.[0-9]{1,3}\.[0-9]{1,3})|'+
18                      '(127\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})$')
19
20 def findMyIPAddr(addrs, intended_port, local_ok = False):
21     """Find the best IP address to use from a list of possibilities.
22     
23     @param addrs: the list of possible IP addresses
24     @param intended_port: the port that was supposed to be used
25     @param local_ok: whether known local/private IP ranges are allowed
26         (defaults to False)
27     @return: the preferred IP address, or None if one couldn't be found
28     """
29     log.msg("got addrs: %r" % (addrs,))
30     my_addr = None
31     
32     # Try to find an address using the ifconfig function
33     try:
34         ifconfig = os.popen("/sbin/ifconfig |/bin/grep inet|"+
35                             "/usr/bin/awk '{print $2}' | "+
36                             "sed -e s/.*://", "r").read().strip().split('\n')
37     except:
38         ifconfig = []
39
40     # Get counts for all the non-local addresses returned from ifconfig
41     addr_count = {}
42     for addr in ifconfig:
43         if local_ok or not isLocal.match(addr):
44             addr_count.setdefault(addr, 0)
45             addr_count[addr] += 1
46     
47     # If only one was found, use it as a starting point
48     local_addrs = addr_count.keys()    
49     if len(local_addrs) == 1:
50         my_addr = local_addrs[0]
51         log.msg('Found remote address from ifconfig: %r' % (my_addr,))
52     
53     # Get counts for all the non-local addresses returned from the DHT
54     addr_count = {}
55     port_count = {}
56     for addr in addrs:
57         if local_ok or not isLocal.match(addr[0]):
58             addr_count.setdefault(addr[0], 0)
59             addr_count[addr[0]] += 1
60             port_count.setdefault(addr[1], 0)
61             port_count[addr[1]] += 1
62     
63     # Find the most popular address
64     popular_addr = []
65     popular_count = 0
66     for addr in addr_count:
67         if addr_count[addr] > popular_count:
68             popular_addr = [addr]
69             popular_count = addr_count[addr]
70         elif addr_count[addr] == popular_count:
71             popular_addr.append(addr)
72     
73     # Find the most popular port
74     popular_port = []
75     popular_count = 0
76     for port in port_count:
77         if port_count[port] > popular_count:
78             popular_port = [port]
79             popular_count = port_count[port]
80         elif port_count[port] == popular_count:
81             popular_port.append(port)
82
83     # Check to make sure the port isn't being changed
84     port = intended_port
85     if len(port_count.keys()) > 1:
86         log.msg('Problem, multiple ports have been found: %r' % (port_count,))
87         if port not in port_count.keys():
88             log.msg('And none of the ports found match the intended one')
89     elif len(port_count.keys()) == 1:
90         port = port_count.keys()[0]
91     else:
92         log.msg('Port was not found')
93
94     # If one is popular, use that address
95     if len(popular_addr) == 1:
96         log.msg('Found popular address: %r' % (popular_addr[0],))
97         if my_addr and my_addr != popular_addr[0]:
98             log.msg('But the popular address does not match: %s != %s' % (popular_addr[0], my_addr))
99         my_addr = popular_addr[0]
100     elif len(popular_addr) > 1:
101         log.msg('Found multiple popular addresses: %r' % (popular_addr,))
102         if my_addr and my_addr not in popular_addr:
103             log.msg('And none of the addresses found match the ifconfig one')
104     else:
105         log.msg('No non-local addresses found: %r' % (popular_addr,))
106         
107     if not my_addr:
108         log.msg("Remote IP Address could not be found for this machine")
109         
110     return my_addr
111
112 def ipAddrFromChicken():
113     """Retrieve a possible IP address from the ipchecken website."""
114     import urllib
115     ip_search = re.compile('\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}')
116     try:
117          f = urllib.urlopen("http://www.ipchicken.com")
118          data = f.read()
119          f.close()
120          current_ip = ip_search.findall(data)
121          return current_ip
122     except:
123          return []
124
125 def uncompact(s):
126     """Extract the contact info from a compact peer representation.
127     
128     @type s: C{string}
129     @param s: the compact representation
130     @rtype: (C{string}, C{int})
131     @return: the IP address and port number to contact the peer on
132     @raise ValueError: if the compact representation doesn't exist
133     """
134     if (len(s) != 6):
135         raise ValueError
136     ip = '.'.join([str(ord(i)) for i in s[0:4]])
137     port = (ord(s[4]) << 8) | ord(s[5])
138     return (ip, port)
139
140 def compact(ip, port):
141     """Create a compact representation of peer contact info.
142     
143     @type ip: C{string}
144     @param ip: the IP address of the peer
145     @type port: C{int}
146     @param port: the port number to contact the peer on
147     @rtype: C{string}
148     @return: the compact representation
149     @raise ValueError: if the compact representation doesn't exist
150     """
151     
152     s = ''.join([chr(int(i)) for i in ip.split('.')]) + \
153           chr((port & 0xFF00) >> 8) + chr(port & 0xFF)
154     if len(s) != 6:
155         raise ValueError
156     return s
157
158 def byte_format(s):
159     """Format a byte size for reading by the user.
160     
161     @type s: C{long}
162     @param s: the number of bytes
163     @rtype: C{string}
164     @return: the formatted size with appropriate units
165     """
166     if (s < 1):
167         r = str(int(s*1000.0)/1000.0) + 'B'
168     elif (s < 10):
169         r = str(int(s*100.0)/100.0) + 'B'
170     elif (s < 102):
171         r = str(int(s*10.0)/10.0) + 'B'
172     elif (s < 1024):
173         r = str(int(s)) + 'B'
174     elif (s < 10485):
175         r = str(int((s/1024.0)*100.0)/100.0) + 'KiB'
176     elif (s < 104857):
177         r = str(int((s/1024.0)*10.0)/10.0) + 'KiB'
178     elif (s < 1048576):
179         r = str(int(s/1024)) + 'KiB'
180     elif (s < 10737418L):
181         r = str(int((s/1048576.0)*100.0)/100.0) + 'MiB'
182     elif (s < 107374182L):
183         r = str(int((s/1048576.0)*10.0)/10.0) + 'MiB'
184     elif (s < 1073741824L):
185         r = str(int(s/1048576)) + 'MiB'
186     elif (s < 1099511627776L):
187         r = str(int((s/1073741824.0)*100.0)/100.0) + 'GiB'
188     else:
189         r = str(int((s/1099511627776.0)*100.0)/100.0) + 'TiB'
190     return(r)
191
192 class TestUtil(unittest.TestCase):
193     """Tests for the utilities."""
194     
195     timeout = 5
196     ip = '165.234.1.34'
197     port = 61234
198
199     def test_compact(self):
200         """Make sure compacting is reversed correctly by uncompacting."""
201         d = uncompact(compact(self.ip, self.port))
202         self.failUnlessEqual(d[0], self.ip)
203         self.failUnlessEqual(d[1], self.port)