Repository: aau-network-security/richkit

Issue:
is_outdated function fails 
- in `test_util.py` file is_outdated function is failing, needs to be fixed.

Discussion:
mrturkmen06: It could be nice to merge this changes before merging others #118 , #115, @kidmose 

Patch:
diff --git a/richkit/lookup/__init__.py b/richkit/lookup/__init__.py
index 6f246d7..b5b87b5 100644
--- a/richkit/lookup/__init__.py
+++ b/richkit/lookup/__init__.py
@@ -34,3 +34,13 @@ def registered_country(ip_address):
     :param ip_address: IP Address (string)
     """
     return geo.get_registered_country(ip_address)
+
+def maxmindb_licence_key(license_key):
+    """
+    Return license key for MaxMind DB
+    Retrieve license key for usage of MaxMindDb
+
+    If it is not present print warning
+    """
+
+    return geo.get_license_key(license_key)
\ No newline at end of file
diff --git a/richkit/lookup/geo.py b/richkit/lookup/geo.py
index 56af766..7fab6b4 100644
--- a/richkit/lookup/geo.py
+++ b/richkit/lookup/geo.py
@@ -2,6 +2,24 @@ from richkit.lookup.util import MaxMindDB
 import os
 
 
+def get_license_key(license_key='MAXMIND_LICENSE_KEY'):
+    """
+    @param license_key: Name of environment variable
+    @return: license of MaxMindDB from environent variables as string
+    @return: in case of error, returns Exception, more specifically KeyError
+    """
+    try:
+        maxmind_db_license = os.environ[license_key]
+        return maxmind_db_license
+    except Exception:
+        print("\nWARNING: No MAXMIND LICENSE KEY Found in environment variables")
+        print("\nUsage of lookup module might be affected due to no MaxMind DB License".strip())
+        print("\nMore info ? Check here: https://github.com/aau-network-security/richkit/wiki/Retrieve-and-configure"
+              "-licence-key".strip())
+        print("Proceeding anyway...")
+        return 'NOLICENSEKEYFOUND'
+
+
 def get_country(ip_address):
     """
     Return the country code of a given IP address
@@ -9,6 +27,7 @@ def get_country(ip_address):
     :param ip_address: IP Address (string)
 
     """
+
     try:
         country_code_db = MaxMindDB((
             "https://download.maxmind.com/app/geoip_download?"
@@ -16,7 +35,7 @@ def get_country(ip_address):
             "license_key={license_key}&"
             "suffix=tar.gz"
         ).format(
-            license_key=os.environ['MAXMIND_LICENSE_KEY'],
+            license_key=get_license_key(),
         ), "cc"
         )
         result = country_code_db.get_data(ip_address)
@@ -40,7 +59,7 @@ def get_registered_country(ip_address):
             "license_key={license_key}&"
             "suffix=tar.gz"
         ).format(
-            license_key=os.environ['MAXMIND_LICENSE_KEY'],
+            license_key=get_license_key(),
         ), "cc"
         )
         result = country_code_db.get_data(ip_address)
@@ -64,7 +83,7 @@ def get_asn(ip_address):
             "license_key={license_key}&"
             "suffix=tar.gz"
         ).format(
-            license_key=os.environ['MAXMIND_LICENSE_KEY'],
+            license_key=get_license_key(),
         ), "asn"
         )
         result = country_code_db.get_data(ip_address)
diff --git a/richkit/lookup/util.py b/richkit/lookup/util.py
index dbd83db..caa6285 100644
--- a/richkit/lookup/util.py
+++ b/richkit/lookup/util.py
@@ -45,8 +45,10 @@ class MaxMindDB:
         self.path_db = maxmind_directory
         if MaxMindDB.get_db_path(self) is None:
             MaxMindDB.get_db(self)
+        #  weeks = 1 because the database is updated once a week.
+        #  if it is downloaded more than one week ago, it will be removed and updated
 
-        if self.is_outdated():
+        if self.get_age() > timedelta(weeks=1):
             os.remove(self.get_db_path())
             MaxMindDB.get_db(self)
 
@@ -122,5 +124,3 @@ class MaxMindDB:
         )
         return delta
 
-    def is_outdated(self):
-        return self.get_age() > timedelta(weeks=1)


Test patch:
diff --git a/richkit/test/lookup/test_geo.py b/richkit/test/lookup/test_geo.py
index d3a8163..40a05dd 100644
--- a/richkit/test/lookup/test_geo.py
+++ b/richkit/test/lookup/test_geo.py
@@ -1,3 +1,5 @@
+import os
+
 from richkit import lookup
 import unittest
 
@@ -15,3 +17,10 @@ class LookupTestCase(unittest.TestCase):
     def test_registered_country(self):
         registered_country = lookup.registered_country("8.8.8.8")
         assert registered_country == 'US'
+
+    def test_maxmindb_licence_key(self):
+        test_license_key = os.environ["TEST_LICENSE_KEY"] = "LICENSEKEY"
+        license_key = lookup.maxmindb_licence_key("TEST_LICENSE_KEY")
+        non_existing_license_key = lookup.maxmindb_licence_key("NON-EXISTING")
+        self.assertTrue(license_key, test_license_key)
+        self.assertIs(non_existing_license_key, 'NOLICENSEKEYFOUND')
diff --git a/richkit/test/lookup/test_util.py b/richkit/test/lookup/test_util.py
index 7b0a27e..f93929e 100644
--- a/richkit/test/lookup/test_util.py
+++ b/richkit/test/lookup/test_util.py
@@ -1,3 +1,6 @@
+import time
+from datetime import datetime
+
 from richkit.lookup import util
 from richkit.lookup.util import MaxMindDB
 import os
@@ -101,8 +104,11 @@ class MaxMindDBTestCase(unittest.TestCase):
         s.get_db()
         # Check if file is present
         p = s.get_db_path()
+        s_age = s.get_age()  # get_age() function is tested over here
         self.assertIsNotNone(p, "get_db did not a path to the db")
         self.assertTrue(Path(p).exists())
+        self.assertTrue(s_age.microseconds)
+
 
     def test_extracted_db(self):
         s = StubMaxMindDB()
@@ -110,6 +116,3 @@ class MaxMindDBTestCase(unittest.TestCase):
         with self.assertRaises(Exception):
             s.unpack()
 
-    def test_is_outdated(self):
-        obj = MaxMindDB(MaxMindDB.MASTERURL, "cc")
-        self.assertFalse(obj.is_outdated())