Document the DHT's main khashmir module.

[quix0rs-apt-p2p.git] / apt_dht / util.py

diff --git a/apt_dht/util.py b/apt_dht/util.py
index 84449006336bad591d449fd1c29ea40fff7c5ae7..1f571e56572392d5cdca9679f7a9160dc6f2dba7 100644 (file)
--- a/apt_dht/util.py
+++ b/apt_dht/util.py
@@ -1,19 +1,33 @@
-## Copyright 2002-2003 Andrew Loewenstern, All Rights Reserved
-# see LICENSE.txt for license information
+
+"""Some utitlity functions for use in the apt-dht program.
+
+@var isLocal: a compiled regular expression suitable for testing if an
+    IP address is from a known local or private range
+"""
 
 import os, re
 
 from twisted.python import log
+from twisted.trial import unittest
 
 isLocal = re.compile('^(192\.168\.[0-9]{1,3}\.[0-9]{1,3})|'+
                      '(10\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})|'+
                      '(172\.0?([1][6-9])|([2][0-9])|([3][0-1])\.[0-9]{1,3}\.[0-9]{1,3})|'+
                      '(127\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})$')
 
-def findMyIPAddr(addrs, intended_port):
+def findMyIPAddr(addrs, intended_port, local_ok = False):
+    """Find the best IP address to use from a list of possibilities.
+    
+    @param addrs: the list of possible IP addresses
+    @param intended_port: the port that was supposed to be used
+    @param local_ok: whether known local/private IP ranges are allowed
+        (defaults to False)
+    @return: the preferred IP address, or None if one couldn't be found
+    """
     log.msg("got addrs: %r" % (addrs,))
     my_addr = None
     
+    # Try to find an address using the ifconfig function
     try:
         ifconfig = os.popen("/sbin/ifconfig |/bin/grep inet|"+
                             "/usr/bin/awk '{print $2}' | "+
@@ -21,23 +35,24 @@ def findMyIPAddr(addrs, intended_port):
     except:
         ifconfig = []
 
-    # Get counts for all the non-local addresses returned
+    # Get counts for all the non-local addresses returned from ifconfig
     addr_count = {}
     for addr in ifconfig:
-        if not isLocal.match(addr):
+        if local_ok or not isLocal.match(addr):
             addr_count.setdefault(addr, 0)
             addr_count[addr] += 1
     
+    # If only one was found, use it as a starting point
     local_addrs = addr_count.keys()    
     if len(local_addrs) == 1:
         my_addr = local_addrs[0]
         log.msg('Found remote address from ifconfig: %r' % (my_addr,))
     
-    # Get counts for all the non-local addresses returned
+    # Get counts for all the non-local addresses returned from the DHT
     addr_count = {}
     port_count = {}
     for addr in addrs:
-        if not isLocal.match(addr[0]):
+        if local_ok or not isLocal.match(addr[0]):
             addr_count.setdefault(addr[0], 0)
             addr_count[addr[0]] += 1
             port_count.setdefault(addr[1], 0)
@@ -62,7 +77,8 @@ def findMyIPAddr(addrs, intended_port):
             popular_count = port_count[port]
         elif port_count[port] == popular_count:
             popular_port.append(port)
-            
+
+    # Check to make sure the port isn't being changed
     port = intended_port
     if len(port_count.keys()) > 1:
         log.msg('Problem, multiple ports have been found: %r' % (port_count,))
@@ -73,6 +89,7 @@ def findMyIPAddr(addrs, intended_port):
     else:
         log.msg('Port was not found')
 
+    # If one is popular, use that address
     if len(popular_addr) == 1:
         log.msg('Found popular address: %r' % (popular_addr[0],))
         if my_addr and my_addr != popular_addr[0]:
@@ -91,6 +108,7 @@ def findMyIPAddr(addrs, intended_port):
     return my_addr
 
 def ipAddrFromChicken():
+    """Retrieve a possible IP address from the ipchecken website."""
     import urllib
     ip_search = re.compile('\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}')
     try:
@@ -101,3 +119,49 @@ def ipAddrFromChicken():
          return current_ip
     except Exception:
          return []
+
+def uncompact(s):
+    """Extract the contact info from a compact peer representation.
+    
+    @type s: C{string}
+    @param s: the compact representation
+    @rtype: (C{string}, C{int})
+    @return: the IP address and port number to contact the peer on
+    @raise ValueError: if the compact representation doesn't exist
+    """
+    if (len(s) != 6):
+        raise ValueError
+    ip = '.'.join([str(ord(i)) for i in s[0:4]])
+    port = (ord(s[4]) << 8) | ord(s[5])
+    return (ip, port)
+
+def compact(ip, port):
+    """Create a compact representation of peer contact info.
+    
+    @type ip: C{string}
+    @param ip: the IP address of the peer
+    @type port: C{int}
+    @param port: the port number to contact the peer on
+    @rtype: C{string}
+    @return: the compact representation
+    @raise ValueError: if the compact representation doesn't exist
+    """
+    
+    s = ''.join([chr(int(i)) for i in ip.split('.')]) + \
+          chr((port & 0xFF00) >> 8) + chr(port & 0xFF)
+    if len(s) != 6:
+        raise ValueError
+    return s
+
+class TestUtil(unittest.TestCase):
+    """Tests for the utilities."""
+    
+    timeout = 5
+    ip = '165.234.1.34'
+    port = 61234
+
+    def test_compact(self):
+        """Make sure compacting is reversed correctly by uncompacting."""
+        d = uncompact(compact(self.ip, self.port))
+        self.failUnlessEqual(d[0], self.ip)
+        self.failUnlessEqual(d[1], self.port)