Update iptogeo.py class

This commit is contained in:
Gregory Soutade 2016-04-12 07:56:10 +02:00
parent 9bb188b6fc
commit e9c2509006

View File

@ -1,6 +1,25 @@
#!/usr/bin/env python #!/usr/bin/env python
# -*- coding: utf-8 -*- # -*- 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 socket
import struct import struct
@ -27,26 +46,44 @@ class IPToGeo(object):
5 : 'Unsupported IP version', 5 : 'Unsupported IP version',
6 : 'IP not found'} 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_addr = remote_addr
self._remote_port = remote_port self._remote_port = remote_port
self._timeout = timeout self._timeout = timeout
self._family = family
self._create_socket() self._create_socket()
def _create_socket(self): 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: if not self._timeout is None:
self._socket.settimeout(self._timeout) self._socket.settimeout(self._timeout)
self._socket.connect((self._remote_addr, self._remote_port)) 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 = ''
packet += struct.pack('<IBBBBI', IPToGeo.MAGIC, IPToGeo.VERSION, IPToGeo.REQ, packet += struct.pack('<IBBBBI', IPToGeo.MAGIC, IPToGeo.VERSION, IPToGeo.REQ,
0, #err 0, #err
IPToGeo.IPV4, # ip type ip_type, # ip type
0) # flags 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('<III', 0, 0, 0) # ipv6
packet += struct.pack('<I', 0) # country code packet += struct.pack('<I', 0) # country code
@ -58,18 +95,35 @@ class IPToGeo(object):
if magic != IPToGeo.MAGIC: if magic != IPToGeo.MAGIC:
raise IPToGeoException('Invalid magic %08x' % (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: if err != 0:
raise IPToGeoException(IPToGeo.ERRORS[err]) raise IPToGeoException(IPToGeo.ERRORS[err])
(cc0, cc1, cc2, cc3) = struct.unpack_from('BBBB', packet, 7*4) (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): def ip_to_geo(self, ip):
splitted_ip = [int(a) for a in ip.split('.')] 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: try:
self._socket.send(packet) self._socket.send(packet)
except IOError, e: except IOError, e: