Extend the geoip fetcher to also download the ASN database. Changes the configuration file syntax slightly. (#1562)

cron/update_geolite_db.conf

@@ -1,5 +1,6 @@
 [options]
-db_location = /opt/mozdef/envs/mozdef/data/GeoLite2-City.mmdb
-db_download_location = https://updates.maxmind.com/geoip/databases/GeoLite2-City/update
+db_file = GeoLite2-City.mmdb
+asn_db_file = GeoLite2-ASN.mmdb
+db_store_location = /opt/mozdef/envs/mozdef/data
 account_id = <insert account number>
 license_key = <insert license key>

cron/update_geolite_db.py

@@ -6,25 +6,27 @@
 # Copyright (c) 2017 Mozilla Corporation
 
 import sys
-import os
-from configlib import getConfig, OptionParser
-
 import requests
 import tempfile
 import gzip
 
+from configlib import getConfig, OptionParser
+from tempfile import mkstemp
+from os import close, fsync, path, rename
+
 from mozdef_util.geo_ip import GeoIP
 from mozdef_util.utilities.logger import logger, initLogger
 
 
-def fetch_db_data(db_download_location):
+def fetch_db_data(db_file):
+    db_download_location = 'https://updates.maxmind.com/geoip/databases/' + db_file[:-5] + '/update'
     logger.debug('Fetching db data from ' + db_download_location)
     auth_creds = (options.account_id, options.license_key)
     response = requests.get(db_download_location, auth=auth_creds)
     if not response.ok:
         raise Exception("Received bad response from maxmind server: {0}".format(response.text))
     db_raw_data = response.content
-    with tempfile.NamedTemporaryFile(mode='wb') as temp:
+    with tempfile.NamedTemporaryFile(mode='wb', prefix=db_file + '.zip.', suffix='.tmp', dir=options.db_store_location) as temp:
         logger.debug('Writing compressed gzip to temp file: ' + temp.name)
         temp.write(db_raw_data)
         temp.flush()
@@ -34,24 +36,31 @@ def fetch_db_data(db_download_location):
         return data
 
 
-def save_db_data(save_path, db_data):
-    temp_save_path = save_path + ".tmp"
-    logger.debug("Saving db data to " + temp_save_path)
-    with open(temp_save_path, "wb+") as text_file:
-        text_file.write(db_data)
+def save_db_data(db_file, db_data):
+    save_path = path.join(options.db_store_location, db_file)
+    fd, temp_path = mkstemp(suffix='.tmp', prefix=db_file, dir=options.db_store_location)
+    with open(temp_path, 'wb') as temp:
+        logger.debug("Saving db data to " + temp_path)
+        temp.write(db_data)
+        fsync(temp.fileno())
+        temp.flush()
         logger.debug("Testing temp geolite db file")
-    geo_ip = GeoIP(temp_save_path)
+        geo_ip = GeoIP(temp_path)
         # Do a generic lookup to verify we don't get any errors (malformed data)
         geo_ip.lookup_ip('8.8.8.8')
         logger.debug("Moving temp file to " + save_path)
-    os.rename(temp_save_path, save_path)
+    close(fd)
+    rename(temp_path, save_path)
 
 
 def main():
     logger.debug('Starting')
-    logger.debug(options)
-    db_data = fetch_db_data(options.db_download_location)
-    save_db_data(options.db_location, db_data)
+
+    db_data = fetch_db_data(options.db_file)
+    asn_db_data = fetch_db_data(options.asn_db_file)
+
+    save_db_data(options.db_file, db_data)
+    save_db_data(options.asn_db_file, asn_db_data)
 
 
 def initConfig():
@@ -60,9 +69,9 @@ def initConfig():
     options.sysloghostname = getConfig('sysloghostname', 'localhost', options.configfile)
     options.syslogport = getConfig('syslogport', 514, options.configfile)
 
-    options.db_download_location = getConfig('db_download_location', '', options.configfile)
-    options.db_location = getConfig('db_location', '', options.configfile)
-
+    options.db_store_location = getConfig('db_store_location', '', options.configfile)
+    options.db_file = getConfig('db_file', '', options.configfile)
+    options.asn_db_file = getConfig('asn_db_file', '', options.configfile)
     options.account_id = getConfig('account_id', '', options.configfile)
     options.license_key = getConfig('license_key', '', options.configfile)