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

Try some refactoring. Needs testing.

This commit is contained in:
Manuel Friedli 2021-07-15 21:51:14 +02:00
parent f7f833e289
commit 0062ae0b2b
1 changed files with 37 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
geoip-lookup.py
View File

@ -26,7 +26,6 @@ class LookupException(Exception):
"""
pass
def get_county_code(ipaddress, dbfile):
"""
Determine the country code that the given ipaddress comes from.
@ -38,7 +37,7 @@ def get_county_code(ipaddress, dbfile):
raise LookupException("No address given")
if not dbfile:
raise LookupException("No db file given")
reader = None
try:
reader = geoip2.database.Reader(dbfile)
@ -50,10 +49,9 @@ def get_county_code(ipaddress, dbfile):
country = reader.country(ipaddress).country
# 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
except FileNotFoundError as e:
raise LookupException(e.args)
@ -68,14 +66,38 @@ def get_county_code(ipaddress, dbfile):
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:
raise LookupException("No address given")
if not dbfile:
raise LookupException("No db file given")
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)
dbtype = reader.metadata().database_type
result = None
country = None
continent = None
network = None
if dbtype == 'GeoLite2-City' or dbtype == 'GeoIP2-City':
result = reader.city(ipaddress)
elif dbfile == 'GeoLite2-Country' or dbtype == 'GeoIP2-Country':
result = reader.country(ipaddress)
# ASN is not supported
# elif dbfile == 'GeoLite2-ASN' or dbtype == 'GeoIP2-ASN':
else:
raise LookupException("Unsupported DB type: " + dbtype)
country = result.country.iso_code
continent = result.continent.code
network = result.traits.network
return "%s,%s,%s" % (country, continent, network)
finally:
if reader:
reader.close()
@ -90,17 +112,18 @@ 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('--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('-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, 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:
@ -109,13 +132,12 @@ def main(argv):
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 __name__ == '__main__':
try:
main(sys.argv[1:])