Refactoring.

This commit is contained in:
Manuel Friedli 2021-07-15 22:02:57 +02:00
parent 4147d03607
commit 75b8e2fb6c

View file

@ -26,51 +26,12 @@ class LookupException(Exception):
""" """
pass 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 ipaddress: The IP address to look up
:param dbfile: The path to the GeoIP2/GeoLite2 database file (Country or City database) :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")
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: if not ipaddress:
raise LookupException("No address given") raise LookupException("No address given")
@ -91,10 +52,7 @@ def get_details(ipaddress, dbfile):
else: else:
raise LookupException("Unsupported DB type: " + dbtype) raise LookupException("Unsupported DB type: " + dbtype)
country = result.country.iso_code return result
continent = result.continent.code
network = result.traits.network
return "%s,%s,%s" % (country, continent, network)
except FileNotFoundError as e: except FileNotFoundError as e:
raise LookupException(e.args) raise LookupException(e.args)
except maxminddb.errors.InvalidDatabaseError as e: except maxminddb.errors.InvalidDatabaseError as e:
@ -107,6 +65,31 @@ def get_details(ipaddress, dbfile):
if reader: if reader:
reader.close() 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): 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 Parse the command line. First, the database file must be specified ("-f /path/to/db/file.mmdb"), then the IP address
@ -141,7 +124,10 @@ def main(argv):
print(code) print(code)
except LookupException as e: except LookupException as e:
print(e.args, file=sys.stderr) print(e.args, file=sys.stderr)
print("Unknown") if detail:
print("Unknown,,")
else:
print("Unknown")
if __name__ == '__main__': if __name__ == '__main__':
try: try: