iptogeo/tests/iptogeo.py

#!/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

class IPToGeoException(Exception):
    pass

class IPToGeo(object):

    MAGIC = 0x179E08EF
    VERSION = 1
    REQ = 1
    RESP = 0
    IPV4 = 4
    IPV6 = 16

    IP_NOT_FOUND = 6

    PACKET_SIZE = 32
    
    ERRORS = {1 : 'Bad magic',
              2 : 'Bad version',
              3 : 'Bad request field' ,
              4 : 'Bad IP version',
              5 : 'Unsupported IP version',
              6 : 'IP not found'}
    
    def __init__(self, remote_addr='127.0.0.1', remote_port=53333, timeout=None):
        self._remote_addr = remote_addr
        self._remote_port = remote_port
        self._timeout = timeout
        
        self._create_socket()

    def _create_socket(self):
        self._socket = socket.socket(socket.AF_INET, 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):
        packet = ''
        packet += struct.pack('<IBBBBI', IPToGeo.MAGIC, IPToGeo.VERSION, IPToGeo.REQ,
                              0, #err
                              IPToGeo.IPV4, # ip type
                              0) # flags
        packet += struct.pack('<BBBB', ip[0], ip[1], ip[2], ip[3]) # ipv4
        packet += struct.pack('<III', 0, 0, 0) # ipv6
        packet += struct.pack('<I', 0) # country code

        return packet

    def _check_request(self, packet):
        (magic, version, req, err, ip_type, flags, ipv4, ipv6b, ipv6c, ipv6d) = struct.unpack_from('<IBBBBIIIII', packet, 0)

        if magic != IPToGeo.MAGIC:
            raise IPToGeoException('Invalid magic %08x' % (magic))
        
        if err == IPToGeo.IP_NOT_FOUND: return (ipv4, 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))
        
    def ip_to_geo(self, ip):
        splitted_ip = [int(a) for a in ip.split('.')]

        packet = self._create_request(splitted_ip)
        try:
            self._socket.send(packet)
        except IOError, e:
            # Give another chance (we may have been disconnected due to timeout)
            self._create_socket()
            self._socket.send(packet)
        packet = self._socket.recv(IPToGeo.PACKET_SIZE)
        if not packet:
            raise IPToGeoException('Error, empty packet')
        (ip, country_code) = self._check_request(packet)
        if country_code:
            # convert to string
            country_code = '%c%c' % (country_code[0], country_code[1])
        return (ip, country_code)

    def close(self):
        self._socket.close()