Switch to local GeoIP instead of remote service

contrib/map/backend/warden-map.py

@@ -6,17 +6,18 @@
 # Copyright (C) 2016 Cesnet z.s.p.o
 # Use of this source is governed by a 3-clause BSD-style license, see LICENSE file.
 
+import GeoIP
+import codecs
+import time
 
 def getLastEvents(events, client, key, cert, cacert, secret):
 
-  try:
-    ses = Session()
-    req = Request('POST', 'https://warden-hub.cesnet.cz/warden3/getEvents?client='+ client + ('&secret='+ secret if secret else "")+'&count=' + events)
-    pre = req.prepare()
-    res = ses.send(pre, cert = (cert, key), verify=cacert)
-  except requests.exceptions.RequestException as error:
-    print error
-    sys.exit(1)
+  ses = Session()
+  #req = Request('POST', 'https://warden-hub.cesnet.cz/warden3/getEvents?client='+ client + ('&secret='+ secret if secret else "")+'&count=' + events + '&id=0')
+  req = Request('POST', 'https://warden-hub.cesnet.cz/warden3/getEvents?client='+ client + ('&secret='+ secret if secret else "")+'&count=' + events)
+  #req = Request('POST', 'https://warden-hub.cesnet.cz/warden3/getEvents?nocat=Other&client='+ client + ('&secret='+ secret if secret else "")+'&count=' + events)
+  pre = req.prepare()
+  res = ses.send(pre, cert = (cert, key), verify=cacert)
 
   data = res.json()
   i = 0
@@ -29,7 +30,7 @@ def getLastEvents(events, client, key, cert, cacert, secret):
           if (key == 'origin') or (key == 'destination'):
             event[key] = {}
             if 'IP4' in p[value][0]:
-              event[key]['ip'] = p[value][0]['IP4']
+              event[key]['ip'] = p[value][0]['IP4'][0]
             else:
               event[key] = {}
           elif (key == 'event'):
@@ -49,21 +50,19 @@ def getLastEvents(events, client, key, cert, cacert, secret):
 
   return eventsList
 
-def getGeolocation(ip):
+def getGeolocation(ip, db):
 
-  try:
-    response = requests.get('http://freegeoip.net/json/' + str(ip[0]))
-  except requests.exceptions.RequestException as error:
-    print error
-    sys.exit(1)
+  data = db.record_by_addr(ip)
 
-  try:
-    json_data = json.loads(response.text)
-  except ValueError as error:
-    print error
-    sys.exit(1)
-
-  return {'latitude': json_data['latitude'], 'longitude': json_data['longitude'], 'country_name': json_data['country_name'], 'city': json_data['city']}
+  if not data:
+    return {}
+  else:
+    return {
+      'latitude': data['latitude'], 
+      'longitude': data['longitude'], 
+      'country_name': unicode(data['country_name'], "utf-8") if data['country_name'] else None, 
+      'city': unicode(data['city'], "utf-8") if data['city'] else None
+    }
 
 def main(args):
 
@@ -79,13 +78,16 @@ def main(args):
   else:
     path = 'warden-map.json'
 
+  db = GeoIP.open("GeoLiteCity.dat", GeoIP.GEOIP_MEMORY_CACHE)
+  db.set_charset(GeoIP.GEOIP_CHARSET_UTF8)
+
   wardenEvents = getLastEvents(events, client, key, cert, cacert, secret)
 
   for p in wardenEvents:
     for target in {'origin', 'destination'}:
       geoData = {}
       if 'ip' in p[target]:
-        geoData = getGeolocation(p[target]['ip'])
+        geoData = getGeolocation(p[target]['ip'], db)
         for value in {'latitude', 'longitude', 'country_name', 'city'}:
           if value in geoData:
             if not geoData[value]:
@@ -102,12 +104,10 @@ def main(args):
         p[target]['latitude'] = 49.743
         p[target]['longitude'] = 15.338
 
-  try:
-    with open(path, 'w') as outfile:
-      json.dump(wardenEvents, outfile)
-  except IOError:
-    print "Error: File does not appear to exist."
-    sys.exit(1)
+  wardenEvents.append(int(time.time()));
+
+  with codecs.open(path, 'w', encoding="utf-8") as outfile:
+    json.dump(wardenEvents, outfile, ensure_ascii = False)
 
   return 0