linux
/
ddos-mitigator
1
0
Fork 0
Code Issues Pull Requests Releases 5 Wiki Activity

Merge pull request 'feature/handle-python-exceptions' (#5) from feature/handle-python-exceptions into master 

Reviewed-on: #5
This commit is contained in:
Manuel Friedli 2020-11-24 23:10:37 +01:00
parent a1d5821cc7 31f0901134
commit 22bc52d665
1 changed files with 32 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
geoip-lookup.py
View File

@ -1,16 +1,24 @@
#!/usr/bin/python3.8 #!/usr/bin/python3.8
import getopt import argparse
import sys import sys
try: try:
import geoip2.database import geoip2.database
import geoip2.errors
except ImportError: except ImportError:
print( print(
"Required module geoip2.database not found. On Gentoo Linux, please install dev-python/geoip2 from the 'fritteli' overlay.", "Required modules geoip2.database and geoip2.errors not found. On Gentoo Linux, please install dev-python/geoip2 from the 'fritteli' overlay.",
file=sys.stderr) file=sys.stderr)
exit(1) 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): class LookupException(Exception):
""" """
@ -47,6 +55,14 @@ def get_county_code(ipaddress, dbfile):
raise LookupException("Unsupported DB type: " + dbtype) raise LookupException("Unsupported DB type: " + dbtype)
return country.iso_code 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: finally:
if reader: if reader:
reader.close() reader.close()
@ -60,21 +76,12 @@ def parse_command_line(argv):
:return: :return:
""" """
dbfile = None dbfile = None
try: parser = argparse.ArgumentParser(description='Get the country code from an IP address')
opts, args = getopt.getopt(argv, "f:") parser.add_argument('-f', dest='dbfile', required=True, help="Path to the GeoIP2 database file")
except getopt.GetoptError as e: parser.add_argument('address', help="The IP address to check")
raise LookupException("Error parsing command line") from e args = parser.parse_args()
for opt, arg in opts: return args.dbfile, args.address
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): def main(argv):
""" """
@ -83,14 +90,18 @@ def main(argv):
:param argv: Format: "-f /path/to/database.mmdb ip.v4.add.ress" :param argv: Format: "-f /path/to/database.mmdb ip.v4.add.ress"
:return: :return:
""" """
(dbfile, ipaddress) = parse_command_line(argv) try:
code = get_county_code(ipaddress, dbfile) (dbfile, ipaddress) = parse_command_line(argv)
print(code) code = get_county_code(ipaddress, dbfile)
print(code)
except LookupException as e:
print(e.args, file=sys.stderr)
print("Unknown")
if __name__ == '__main__': if __name__ == '__main__':
try: try:
main(sys.argv[1:]) main(sys.argv[1:])
except BaseException as e: except BaseException as e:
print("Usage: geoip-lookup.py -f /path/to/geoip2-database.mmdb 192.168.1.1", file=sys.stderr) print("Unknown")
raise e raise e