ddos-mitigator/geoip-lookup.py

#!/usr/bin/python

import argparse
import sys

try:
    import geoip2.database
    import geoip2.errors
except ImportError:
    print(
        "Required modules geoip2.database and geoip2.errors not found. On Gentoo Linux, please install dev-python/geoip2 from the 'fritteli' overlay.",
        file=sys.stderr)
    exit(1)

try:
    import maxminddb.errors
except ImportError:
    print(
        "Required module maxminddb.errors not found. On Gentoo Linux, please install dev-python/maxminddb from the 'fritteli' overlay.",
        file=sys.stderr)
    exit(1)

class LookupException(Exception):
    """
    General Exception class that is raised if anything goes wrong.
    """
    pass

def get_county_code(ipaddress, dbfile):
    """
    Determine the country code that the given ipaddress comes from.
    :param ipaddress: The IP address to look up
    :param dbfile: The path to the GeoIP2/GeoLite2 database file (Country or City database)
    :return: The ISO country code (2 letters)
    """
    if not ipaddress:
        raise LookupException("No address given")
    if not dbfile:
        raise LookupException("No db file given")

    reader = None
    try:
        reader = geoip2.database.Reader(dbfile)
        dbtype = reader.metadata().database_type
        country = None
        if dbtype == 'GeoLite2-City' or dbtype == 'GeoIP2-City':
            country = reader.city(ipaddress).country
        elif dbtype == 'GeoLite2-Country' or dbtype == 'GeoIP2-Country':
            country = reader.country(ipaddress).country
        # ASN is not supported
        # elif dbfile == 'GeoLite2-ASN' or dbtype == 'GeoIP2-ASN':
        else:
            raise LookupException("Unsupported DB type: " + dbtype)

        return country.iso_code
    except FileNotFoundError as e:
        raise LookupException(e.args)
    except maxminddb.errors.InvalidDatabaseError as e:
        raise LookupException(e.args)
    except geoip2.errors.AddressNotFoundError as e:
        raise LookupException(e.args)
    except ValueError as e:
        raise LookupException(e.args)
    finally:
        if reader:
            reader.close()

def get_details(ipaddress, dbfile):
    """
    Determine the country code, continent code and the network that the given ipaddress comes from.
    :param ipaddress: The IP address to look up
    :param dbfile: The path to the GeoIP2/GeoLite2 database file (Country or City database)
    :return: A string consisting of the ISO country code (2 letters), ISO continent code (2 letters) and network (CIDR notation), concatenated by ","
    """
    if not ipaddress:
        raise LookupException("No address given")
    if not dbfile:
        raise LookupException("No db file given")

    reader = None
    try:
        reader = geoip2.database.Reader(dbfile)
        dbtype = reader.metadata().database_type
        result = None
        if dbtype == 'GeoLite2-City' or dbtype == 'GeoIP2-City':
            result = reader.city(ipaddress)
        elif dbtype == 'GeoLite2-Country' or dbtype == 'GeoIP2-Country':
            result = reader.country(ipaddress)
        # ASN is not supported
        # elif dbfile == 'GeoLite2-ASN' or dbtype == 'GeoIP2-ASN':
        else:
            raise LookupException("Unsupported DB type: " + dbtype)

        country = result.country.iso_code
        continent = result.continent.code
        network = result.traits.network
        return "%s,%s,%s" % (country, continent, network)
    except FileNotFoundError as e:
        raise LookupException(e.args)
    except maxminddb.errors.InvalidDatabaseError as e:
        raise LookupException(e.args)
    except geoip2.errors.AddressNotFoundError as e:
        raise LookupException(e.args)
    except ValueError as e:
        raise LookupException(e.args)
    finally:
        if reader:
            reader.close()

def parse_command_line(argv):
    """
    Parse the command line. First, the database file must be specified ("-f /path/to/db/file.mmdb"), then the IP address
    to look up
    :param argv:
    :return:
    """
    dbfile = None
    parser = argparse.ArgumentParser(description='Get the country code from an IP address')
    parser.add_argument('-f', dest='dbfile', required=True, help="Path to the GeoIP2 database file")
    parser.add_argument('-d', dest='detail', action='store_const', const=True, default=False, help="Verbose output: Print continent code and network along with country code")
    parser.add_argument('--detail', dest='detail', action='store_const', const=True, default=False, help="Verbose output: Print continent code and network along with country code")
    parser.add_argument('address', help="The IP address to check")
    args = parser.parse_args()

    return args.dbfile, args.address, args.detail

def main(argv):
    """
    Read the database file and the IP address from the command line and print the corresponding ISO country code on
    stdout.
    :param argv: Format: "-f /path/to/database.mmdb [--detail|-d] ip.v4.add.ress"
    :return:
    """
    try:
        (dbfile, ipaddress, detail) = parse_command_line(argv)
        if detail:
            code = get_details(ipaddress, dbfile)
        else:
            code = get_county_code(ipaddress, dbfile)

        print(code)
    except LookupException as e:
        print(e.args, file=sys.stderr)
        print("Unknown")

if __name__ == '__main__':
    try:
        main(sys.argv[1:])
    except BaseException as e:
        print("Unknown")
        raise e