diff --git a/geoip-lookup.py b/geoip-lookup.py index ccf4fc5..28ff122 100755 --- a/geoip-lookup.py +++ b/geoip-lookup.py @@ -26,13 +26,12 @@ class LookupException(Exception): """ pass - -def get_county_code(ipaddress, dbfile): +def get_geoip_result(ipaddress, dbfile): """ - Determine the country code that the given ipaddress comes from. + Perform a lookup of the given ipaddress in the database. :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) + :return: The lookup result """ if not ipaddress: raise LookupException("No address given") @@ -43,18 +42,17 @@ def get_county_code(ipaddress, dbfile): try: reader = geoip2.database.Reader(dbfile) dbtype = reader.metadata().database_type - country = None + result = 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 + 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': - - if not country: + else: raise LookupException("Unsupported DB type: " + dbtype) - return country.iso_code + return result except FileNotFoundError as e: raise LookupException(e.args) except maxminddb.errors.InvalidDatabaseError as e: @@ -67,6 +65,30 @@ def get_county_code(ipaddress, dbfile): if reader: reader.close() +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) + """ + result = get_geoip_result(ipaddress, dbfile) + + return result.country.iso_code + +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 "," + """ + result = get_geoip_result(ipaddress, dbfile) + country = result.country.iso_code + continent = result.continent.code + network = result.traits.network + + return "%s,%s,%s" % (country, continent, network) def parse_command_line(argv): """ @@ -78,26 +100,34 @@ def parse_command_line(argv): 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 + + 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 ip.v4.add.ress" + :param argv: Format: "-f /path/to/database.mmdb [--detail|-d] ip.v4.add.ress" :return: """ try: - (dbfile, ipaddress) = parse_command_line(argv) - code = get_county_code(ipaddress, dbfile) + (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 detail: + print("Unknown,,") + else: + print("Unknown") if __name__ == '__main__': try: