Add the possibility to return continent code and network in CIDR notation in addition to the country code. Needs error handling.

This commit is contained in:
Manuel Friedli 2021-07-14 03:37:50 +02:00
parent 9b4a55dfb2
commit f7f833e289

View file

@ -38,7 +38,7 @@ def get_county_code(ipaddress, dbfile):
raise LookupException("No address given") raise LookupException("No address given")
if not dbfile: if not dbfile:
raise LookupException("No db file given") raise LookupException("No db file given")
reader = None reader = None
try: try:
reader = geoip2.database.Reader(dbfile) reader = geoip2.database.Reader(dbfile)
@ -53,7 +53,7 @@ def get_county_code(ipaddress, dbfile):
if not country: if not country:
raise LookupException("Unsupported DB type: " + dbtype) raise LookupException("Unsupported DB type: " + dbtype)
return country.iso_code return country.iso_code
except FileNotFoundError as e: except FileNotFoundError as e:
raise LookupException(e.args) raise LookupException(e.args)
@ -67,6 +67,18 @@ def get_county_code(ipaddress, dbfile):
if reader: if reader:
reader.close() reader.close()
def get_details(ipaddress, dbfile):
reader = None
try:
reader = geoip2.database.Reader(dbfile)
res = reader.city(ipaddress)
country = res.country.iso_code
continent = res.continent.code
network = res.traits.network
return "%s|%s|%s" % (country, continent, network)
finally:
if reader:
reader.close()
def parse_command_line(argv): def parse_command_line(argv):
""" """
@ -78,10 +90,11 @@ def parse_command_line(argv):
dbfile = None dbfile = None
parser = argparse.ArgumentParser(description='Get the country code from an IP address') 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('-f', dest='dbfile', required=True, help="Path to the GeoIP2 database file")
parser.add_argument('--details', 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") parser.add_argument('address', help="The IP address to check")
args = parser.parse_args() args = parser.parse_args()
return args.dbfile, args.address return args.dbfile, args.address, args.detail
def main(argv): def main(argv):
""" """
@ -91,8 +104,12 @@ def main(argv):
:return: :return:
""" """
try: try:
(dbfile, ipaddress) = parse_command_line(argv) (dbfile, ipaddress, detail) = parse_command_line(argv)
code = get_county_code(ipaddress, dbfile) if detail:
code = get_details(ipaddress, dbfile)
else:
code = get_county_code(ipaddress, dbfile)
print(code) print(code)
except LookupException as e: except LookupException as e:
print(e.args, file=sys.stderr) print(e.args, file=sys.stderr)