#!/usr/bin/python3.8

import getopt
import sys

try:
    import geoip2.database
except ImportError:
    print(
        "Required module geoip2.database not found. On Gentoo Linux, please install dev-python/geoip2 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 dbfile == 'GeoLite2-Country' or dbtype == 'GeoIP2-Country':
            country = reader.country(ipaddress).country
        # ASN is not supported
        # elif dbfile == 'GeoLite2-ASN' or dbtype == 'GeoIP2-ASN':

        if not country:
            raise LookupException("Unsupported DB type: " + dbtype)

        return country.iso_code
    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
    try:
        opts, args = getopt.getopt(argv, "f:")
    except getopt.GetoptError as e:
        raise LookupException("Error parsing command line") from e

    for opt, arg in opts:
        if opt == "-f":
            dbfile = arg
        else:
            raise LookupException("Unknown command line argument")

    if len(args) == 0:
        raise LookupException("No address given on command line")
    return dbfile, args[0]


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 ip.v4.add.ress"
    :return:
    """
    (dbfile, ipaddress) = parse_command_line(argv)
    code = get_county_code(ipaddress, dbfile)
    print(code)


if __name__ == '__main__':
    try:
        main(sys.argv[1:])
    except BaseException as e:
        print("Usage: geoip-lookup.py -f /path/to/geoip2-database.mmdb 192.168.1.1", file=sys.stderr)
        raise e