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

-"""
-Misc. utility functions. It is part of the pygeoip package.
-
-@author: Jennifer Ennis <zaylea at gmail dot com>
-
-@license:
-Copyright(C) 2004 MaxMind LLC
-
-This program is free software: you can redistribute it and/or modify
-it under the terms of the GNU Lesser General Public License as published by
-the Free Software Foundation, either version 3 of the License, or
-(at your option) any later version.
-
-This program is distributed in the hope that it will be useful,
-but WITHOUT ANY WARRANTY; without even the implied warranty of
-MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-GNU General Public License for more details.
-
-You should have received a copy of the GNU Lesser General Public License
-along with this program.  If not, see <http://www.gnu.org/licenses/lgpl.txt>.
-"""
-
-import six
-
-def ip2long(ip):
-    """
-    Convert a IPv4 address into a 32-bit integer.
-    
-    @param ip: quad-dotted IPv4 address
-    @type ip: str
-    @return: network byte order 32-bit integer
-    @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])
-    else:
-        ip_long = long(ip_array[0]) * 16777216 + long(ip_array[1]) * 65536 + long(ip_array[2]) * 256 + long(ip_array[3])
-    return ip_long  
-
+"""
+Misc. utility functions. It is part of the pygeoip package.
+
+@author: Jennifer Ennis <zaylea at gmail dot com>
+
+@license:
+Copyright(C) 2004 MaxMind LLC
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Lesser General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU Lesser General Public License
+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
+    @type ip: str
+    @return: network byte order 32-bit integer
+    @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])
+    else:
+        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)))