ddos-mitigator/geoip-lookup.py

97 lines
2.7 KiB
Python
Raw Normal View History

#!/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