From f7f833e28987fcf31012345fe80c78406456236a Mon Sep 17 00:00:00 2001
From: Manuel Friedli <manuel@fritteli.ch>
Date: Wed, 14 Jul 2021 03:37:50 +0200
Subject: [PATCH] Add the possibility to return continent code and network in
 CIDR notation in addition to the country code. Needs error handling.

---
 geoip-lookup.py | 27 ++++++++++++++++++++++-----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/geoip-lookup.py b/geoip-lookup.py
index ccf4fc5..33b1378 100755
--- a/geoip-lookup.py
+++ b/geoip-lookup.py
@@ -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)