Merge branch 'master' of https://github.com/tiwilliam/pygeoip

Conflicts:
	pygeoip/__init__.py
	tests/config.py
	tests/test_pygeoip.py

pygeoip/__init__.py

@@ -187,7 +187,9 @@ class GeoIP(GeoIPBase):
         @rtype: int
         """
         offset = 0
-        for depth in range(31, -1, -1):
+        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
                 endIndex = startIndex + (2 * self._recordLength)
@@ -202,7 +204,6 @@ 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]
@@ -379,7 +380,8 @@ class GeoIP(GeoIPBase):
         if not ipnum:
             raise ValueError("Invalid IP address: %s" % addr)
 
-        if self._databaseType != const.COUNTRY_EDITION:
+        COUNTY_EDITIONS = (const.COUNTRY_EDITION, const.COUNTRY_EDITION_V6)
+        if self._databaseType not in COUNTY_EDITIONS:
             message = 'Invalid database type, expected Country'
             raise GeoIPError(message)
 
@@ -396,8 +398,17 @@ class GeoIP(GeoIPBase):
         @rtype: str
         """
         try:
-            if self._databaseType == const.COUNTRY_EDITION:
-                return const.COUNTRY_CODES[self.id_by_addr(addr)]
+            COUNTRY_EDITIONS = (const.COUNTRY_EDITION, const.COUNTRY_EDITION_V6)
+            if self._databaseType in COUNTRY_EDITIONS:
+                ipv = 6 if addr.find(':') >= 0 else 4
+                if ipv == 4 and self._databaseType != const.COUNTRY_EDITION:
+                    raise ValueError('Invalid database type; expected IPv6 address')
+                if ipv == 6 and self._databaseType != const.COUNTRY_EDITION_V6:
+                    raise ValueError('Invalid database type; expected IPv4 address')
+
+                country_id = self.id_by_addr(addr)
+
+                return const.COUNTRY_CODES[country_id]
             elif self._databaseType in const.REGION_CITY_EDITIONS:
                 return self.region_by_addr(addr)['country_code']
 
@@ -416,9 +427,24 @@ class GeoIP(GeoIPBase):
         @return: 2-letter country code
         @rtype: str
         """
+        if self._databaseType == const.COUNTRY_EDITION_V6:
+            return self._country_code_by_name_v6(hostname)
+        else:
+            return self._country_code_by_name_v4(hostname)
+
+    def _country_code_by_name_v4(self, hostname):
         addr = socket.gethostbyname(hostname)
         return self.country_code_by_addr(addr)
 
+    def _country_code_by_name_v6(self, hostname):
+        try:
+            response = socket.getaddrinfo(hostname, 0, socket.AF_INET6)
+        except socket.gaierror:
+            return ''
+        family, socktype, proto, canonname, sockaddr = response[0]
+        address, port, flow, scope = sockaddr
+        return self.country_code_by_addr(address)
+
     def country_name_by_addr(self, addr):
         """
         Returns full country name for specified IP address.

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))
+    max_index = len(ipnum) - 1
+    return sum(ipnum[max_index - i] << (i * 32) for i in range(len(ipnum)))

tests/config.py

@@ -3,6 +3,7 @@ import os.path as path
 DATA_DIR = path.join(path.dirname(path.realpath(__file__)), 'data')
 
 COUNTRY_DB_PATH = path.join(DATA_DIR, 'GeoIP.dat')
+COUNTRY_V6_DB_PATH = path.join(DATA_DIR, 'GeoIPv6.dat')
 REGION_DB_PATH = path.join(DATA_DIR, 'GeoIPRegion.dat')
 CITY_DB_PATH = path.join(DATA_DIR, 'GeoLiteCity.dat')
 ORG_DB_PATH = path.join(DATA_DIR, 'GeoIPOrg.dat')