Update iptogeo.py class

plugins/post_analysis/iptogeo.py

@@ -1,6 +1,25 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+#
+# Copyright 2016 Grégory Soutadé
+#
+# This file is part of iptogeo.
+#
+# iptogeo is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# iptogeo 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 General Public License
+# along with iptogeo.  If not, see <http://www.gnu.org/licenses/>.
+#
+
 import socket
 import struct
 
@@ -27,26 +46,44 @@ class IPToGeo(object):
               5 : 'Unsupported IP version',
               6 : 'IP not found'}
     
-    def __init__(self, remote_addr='127.0.0.1', remote_port=53333, timeout=None):
+    def __init__(self, remote_addr='127.0.0.1', remote_port=53333, timeout=None, family=socket.AF_INET):
         self._remote_addr = remote_addr
         self._remote_port = remote_port
         self._timeout = timeout
+        self._family = family
 
         self._create_socket()
 
     def _create_socket(self):
-        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self._socket = socket.socket(self._family, socket.SOCK_STREAM)
         if not self._timeout is None:
             self._socket.settimeout(self._timeout)
         self._socket.connect((self._remote_addr, self._remote_port))
 
-    def _create_request(self, ip):
+    def _extend_ipv6(self, ipv6):
+        tmp = ''
+        for s in ipv6.split(':'):
+            if not s: break
+            while len(s) != 4:
+                s = '0' + s
+            tmp += s
+        while len(tmp) < 16*2:
+            tmp += '0'
+        res = ''
+        for i in range(0, 15*2, 2):
+            res += tmp[i] + tmp[i+1] + ':'
+        res += tmp[30] + tmp[31]
+
+        return res
+
+    def _create_request(self, ip, ip_type):
         packet = ''
         packet += struct.pack('<IBBBBI', IPToGeo.MAGIC, IPToGeo.VERSION, IPToGeo.REQ,
                               0, #err
-                              IPToGeo.IPV4, # ip type
+                              ip_type, # ip type
                               0) # flags
-        packet += struct.pack('<BBBB', ip[0], ip[1], ip[2], ip[3]) # ipv4
+        for i in ip:
+            packet += struct.pack('<B', i) # ipv4
         packet += struct.pack('<III', 0, 0, 0) # ipv6
         packet += struct.pack('<I', 0) # country code
 
@@ -58,18 +95,35 @@ class IPToGeo(object):
         if magic != IPToGeo.MAGIC:
             raise IPToGeoException('Invalid magic %08x' % (magic))
         
-        if err == IPToGeo.IP_NOT_FOUND: return (ipv4, None) # IP not found
+        ip_res = '%08x' % (ipv4)
+        if ip_type == IPToGeo.IPV6:
+            ip_res += '%08x' % (ipv6b)
+            ip_res += '%08x' % (ipv6c)
+            ip_res += '%08x' % (ipv6d)
+
+        if err == IPToGeo.IP_NOT_FOUND: return (ip_res, None) # IP not found
         if err != 0:
             raise IPToGeoException(IPToGeo.ERRORS[err])
 
         (cc0, cc1, cc2, cc3) = struct.unpack_from('BBBB', packet, 7*4)
 
-        return (ipv4, '%c%c%c%c' % (cc0, cc1, cc2, cc3))
+        return (ip_res, '%c%c%c%c' % (cc0, cc1, cc2, cc3))
         
     def ip_to_geo(self, ip):
+        ip_type = IPToGeo.IPV4
+        if ip.find('.') >= 0:
             splitted_ip = [int(a) for a in ip.split('.')]
+            if len(splitted_ip) != 4:
+                raise Exception('Bad IP %s' % (ip))
+        elif ip.find(':') >= 0:
+            splitted_ip = [int(a, 16) for a in self._extend_ipv6(ip).split(':')]
+            if len(splitted_ip) != 16:
+                raise Exception('Bad IP %s' % (ip))
+            ip_type = IPToGeo.IPV6
+        else:
+            raise Exception('Bad IP %s' % (ip))
 
-        packet = self._create_request(splitted_ip)
+        packet = self._create_request(splitted_ip, ip_type)
         try:
             self._socket.send(packet)
         except IOError, e: