diff --git a/.gitignore b/.gitignore index 6f30344f..818c7ac5 100644 --- a/.gitignore +++ b/.gitignore @@ -10,4 +10,4 @@ results cron/ipblocklist.txt alerts/generic_alerts /.project -.data +/data diff --git a/cron/update_geolite_db.py b/cron/update_geolite_db.py index e5ae14a1..b3178bcb 100755 --- a/cron/update_geolite_db.py +++ b/cron/update_geolite_db.py @@ -17,6 +17,7 @@ import tempfile import tarfile sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../lib')) +from geo_ip import GeoIP from utilities.logger import logger, initLogger @@ -37,10 +38,14 @@ def fetch_db_data(db_download_location): def save_db_data(save_path, db_data): - temp_save_path = save_path + ".bak" + 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) + logger.debug("Testing temp geolite db file") + geo_ip = GeoIP(temp_save_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) diff --git a/lib/geo_ip.py b/lib/geo_ip.py index b52d9c81..e66c38be 100644 --- a/lib/geo_ip.py +++ b/lib/geo_ip.py @@ -3,8 +3,9 @@ import geoip2.database class GeoIP(object): - def __init__(self): - db_location = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../data/GeoLite2-City.mmdb") + def __init__(self, db_location=None): + if db_location is None: + db_location = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../data/GeoLite2-City.mmdb") try: self.db = geoip2.database.Reader(db_location) except IOError: