added caching to lookup function

Author: GalaxP

geolite/README.md

@@ -31,8 +31,8 @@ Below are the fields that will be set for each database type:
   * uint32 `country_geoname_id`
   * uint32 `is_in_european_union`
   * string `city_name` 
-  * float `latitude`
-  * float `longitude`
+  * float  `latitude`
+  * float  `longitude`
   * uint32 `accuracy_radius`
 
 * `asn`
@@ -61,6 +61,10 @@ module takes the following parameters:
 
   * Specify the type of GeoLite database. The default value is `country`.
 
+* `-c` `--cache` number
+
+  * Specify the number of lookup calls that will be cached. Set to `0` to disable caching. The default value is `128`.
+
 
 ## Example
 The following command :

geolite/geolite.py

@@ -36,6 +36,7 @@
 
 
 import argparse
+from functools import lru_cache
 import pytrap
 import geoip2.database
 
@@ -58,6 +59,10 @@
                     help="type of GeoLite database",
                     choices=['country', 'city', 'asn'],
                     default="country")
+parser.add_argument('-c', "--cache",
+                    type=int,
+                    help="number of cached lookups. If set to 0, caching is disabled.",
+                    default=128)
 
 # parse command line arguments
 args = parser.parse_args()
@@ -88,9 +93,19 @@
 trap.setRequiredFmt(0, fmttype, fmtspec)
 rec = pytrap.UnirecTemplate(fmtspec)
 
-
 # open GeoLite2 database reader
 with geoip2.database.Reader(args.db) as reader:
+    @lru_cache(maxsize=args.cache)
+    def lookup_in_database(lookup_ip):
+        #Function to look up the IP address in the GeoLite2 database.
+        if args.type == 'city':
+            return reader.city(lookup_ip)
+        elif args.type == 'country':
+            return reader.country(lookup_ip)
+        elif args.type == 'asn':
+            return reader.asn(lookup_ip)
+    
+
     # main loop
     while True:
         try:
@@ -114,12 +129,7 @@
                 ip = rec.get(data, field)
                 output.ip = ip
                 try:
-                    if args.type == 'city':
-                        geolocation = reader.city(str(ip))
-                    elif args.type == 'country':
-                        geolocation = reader.country(str(ip))
-                    elif args.type == 'asn':
-                        geolocation = reader.asn(str(ip))
+                    geolocation = lookup_in_database(str(ip))
                 except:
                     continue
                 if geolocation is None: