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:
parent
9b4a55dfb2
commit
f7f833e289
1 changed files with 22 additions and 5 deletions
|
@ -38,7 +38,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)
|
||||
|
@ -53,7 +53,7 @@ def get_county_code(ipaddress, dbfile):
|
|||
|
||||
if not country:
|
||||
raise LookupException("Unsupported DB type: " + dbtype)
|
||||
|
||||
|
||||
return country.iso_code
|
||||
except FileNotFoundError as e:
|
||||
raise LookupException(e.args)
|
||||
|
@ -67,6 +67,18 @@ def get_county_code(ipaddress, dbfile):
|
|||
if reader:
|
||||
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):
|
||||
"""
|
||||
|
@ -78,10 +90,11 @@ 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('address', help="The IP address to check")
|
||||
args = parser.parse_args()
|
||||
|
||||
return args.dbfile, args.address
|
||||
return args.dbfile, args.address, args.detail
|
||||
|
||||
def main(argv):
|
||||
"""
|
||||
|
@ -91,8 +104,12 @@ def main(argv):
|
|||
:return:
|
||||
"""
|
||||
try:
|
||||
(dbfile, ipaddress) = parse_command_line(argv)
|
||||
code = get_county_code(ipaddress, dbfile)
|
||||
(dbfile, ipaddress, detail) = parse_command_line(argv)
|
||||
if detail:
|
||||
code = get_details(ipaddress, dbfile)
|
||||
else:
|
||||
code = get_county_code(ipaddress, dbfile)
|
||||
|
||||
print(code)
|
||||
except LookupException as e:
|
||||
print(e.args, file=sys.stderr)
|
||||
|
|
Loading…
Reference in a new issue