Initial commit for IPv6 support

pygeoip/__init__.py

@@ -176,7 +176,6 @@ class GeoIP(GeoIPBase):
         """
 
         ipnum = ip2long(addr)
-
         if not ipnum:
             raise ValueError("Invalid IP address: %s" % addr)
 
@@ -196,10 +195,11 @@ class GeoIP(GeoIPBase):
         @return: offset of start of record
         @rtype: int
         """
-        offset = 0
 
-        for depth in range(31, -1, -1):
+        offset = 0
+        seek_depth = 127 if len(str(ipnum)) > 10 else 31
 
+        for depth in range(seek_depth, -1, -1):
             if self._flags & const.MEMORY_CACHE:
                 startIndex = 2 * self._recordLength * offset
                 length = 2 * self._recordLength
@@ -214,22 +214,15 @@ class GeoIP(GeoIPBase):
             for i in range(2):
                 for j in range(self._recordLength):
                     x[i] += ord(buf[self._recordLength * i + j]) << (j * 8)
-
             if ipnum & (1 << depth):
-
                 if x[1] >= self._databaseSegments:
                     return x[1]
-
                 offset = x[1]
-
             else:
-
                 if x[0] >= self._databaseSegments:
                     return x[0]
-
                 offset = x[0]
 
-
         raise Exception('Error traversing database - perhaps it is corrupt?')
 
     def _get_org(self, ipnum):
@@ -436,7 +429,6 @@ class GeoIP(GeoIPBase):
             else:
                 raise GeoIPError('Invalid database type; country_* methods expect '\
                                  'Country, City, or Region database')
-
         except ValueError:
             raise GeoIPError('*_by_addr methods only accept IP addresses. Use *_by_name for hostnames. (Address: %s)' % addr)
 

pygeoip/util.py

@@ -21,9 +21,25 @@ along with this program.  If not, see <http://www.gnu.org/licenses/lgpl.txt>.
 """
 
 import six
+import struct
+import socket
+from array import array
+
 
 def ip2long(ip):
     """
+    Wrapper function for IPv4 and IPv6 converters
+    @param ip: IPv4 or IPv6 address
+    @type ip: str
+    """
+    if ip.find(':') >= 0:
+        return ip2long_v6(ip)
+    else:
+        return ip2long_v4(ip)
+
+
+def ip2long_v4(ip):
+    """
     Convert a IPv4 address into a 32-bit integer.
 
     @param ip: quad-dotted IPv4 address
@@ -32,7 +48,6 @@ def ip2long(ip):
     @rtype: int
     """
     ip_array = ip.split('.')
-    
     if six.PY3:
         # int and long are unified in py3
         ip_long = int(ip_array[0]) * 16777216 + int(ip_array[1]) * 65536 + int(ip_array[2]) * 256 + int(ip_array[3])
@@ -40,3 +55,17 @@ def ip2long(ip):
         ip_long = long(ip_array[0]) * 16777216 + long(ip_array[1]) * 65536 + long(ip_array[2]) * 256 + long(ip_array[3])
     return ip_long
 
+
+def ip2long_v6(ip):
+    """
+    Convert a IPv6 address into long.
+
+    @param ip: IPv6 address
+    @type ip: str
+    @return: network byte order long
+    @rtype: long
+    """
+    ipbyte = socket.inet_pton(socket.AF_INET6, ip)
+    ipnum = array('L', struct.unpack('!4L', ipbyte))
+    ipnum.reverse()
+    return sum(ipnum[i] << (i * 32) for i in range(len(ipnum)))