diff --git a/bot/mozdefbot.py b/bot/mozdefbot.py index 33077548..0c6bc5d0 100755 --- a/bot/mozdefbot.py +++ b/bot/mozdefbot.py @@ -124,7 +124,8 @@ def isIP(ip): def ipLocation(ip): location = "" try: - geoip = GeoIP() + geoip_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../data/GeoLite2-City.mmdb") + geoip = GeoIP(geoip_data_dir) geoDict = geoip.lookup_ip(ip) if geoDict is not None: if 'error' in geoDict: diff --git a/mozdef_util/mozdef_util/geo_ip.py b/mozdef_util/mozdef_util/geo_ip.py index 96d0a03b..e1a8eda8 100644 --- a/mozdef_util/mozdef_util/geo_ip.py +++ b/mozdef_util/mozdef_util/geo_ip.py @@ -1,11 +1,8 @@ -import os import geoip2.database class GeoIP(object): - 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") + def __init__(self, db_location): try: self.db = geoip2.database.Reader(db_location) except IOError: diff --git a/mozdef_util/setup.py b/mozdef_util/setup.py index 8728bf71..8677f1f3 100644 --- a/mozdef_util/setup.py +++ b/mozdef_util/setup.py @@ -56,6 +56,6 @@ setup( test_suite='tests', tests_require=[], url='https://github.com/mozilla/MozDef/tree/master/lib', - version='1.0.1', + version='1.0.2', zip_safe=False, ) diff --git a/mq/plugins/geoip.py b/mq/plugins/geoip.py index 001f8a98..34db7494 100644 --- a/mq/plugins/geoip.py +++ b/mq/plugins/geoip.py @@ -27,7 +27,8 @@ class message(object): ''' self.registration = ['sourceipaddress', 'destinationipaddress'] self.priority = 20 - self.geoip = GeoIP() + geoip_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../data/GeoLite2-City.mmdb") + self.geoip = GeoIP(geoip_data_dir) def ipLocation(self, ip): location = dict() diff --git a/requirements.txt b/requirements.txt index e9d00ab9..e0a12e74 100644 --- a/requirements.txt +++ b/requirements.txt @@ -33,7 +33,7 @@ KitnIRC==0.2.6 kombu==4.1.0 meld3==1.0.2 mozdef-client==1.0.11 -mozdef-util==1.0.1 +mozdef-util==1.0.2 MySQL-python==1.2.5 netaddr==0.7.1 nose==1.3.7 diff --git a/tests/lib/test_geo_ip.py b/tests/lib/test_geo_ip.py index 098c0678..a04ee9c0 100644 --- a/tests/lib/test_geo_ip.py +++ b/tests/lib/test_geo_ip.py @@ -1,14 +1,10 @@ -import os -import sys from mozdef_util.geo_ip import GeoIP class TestGeoIPLookup(object): - def setup(self): - self.geo_ip = GeoIP() - # Unfortunately since the db file is not present by default # we verify the error def test_without_db_file(self): - geo_dict = self.geo_ip.lookup_ip('129.21.1.40') + geo_ip = GeoIP("nonexistent_db") + geo_dict = geo_ip.lookup_ip('129.21.1.40') assert geo_dict['error'] == 'No Geolite DB Found!'