#35 Add unit test for corteutils.Analyzer class

contrib/cortexutils/analyzer.py

@@ -23,7 +23,7 @@ def __init__(self):
         self.data_type = self.get_param('dataType', None, 'Missing dataType field')
         self.tlp = self.get_param('tlp', 2)
 
-        self.check_tlp = self.get_param('config.check_tlp', False)
+        self.enable_check_tlp = self.get_param('config.check_tlp', False)
         self.max_tlp = self.get_param('config.max_tlp', 2)
 
         # Set proxy configuration if available
@@ -33,7 +33,8 @@ def __init__(self):
         self.__set_proxies()
 
         # Finally run check tlp
-        self.__check_tlp()
+        if not (self.__check_tlp()):
+            self.error('TLP is higher than allowed.')
 
     # Not breaking compatibility
     def notSupported(self):
@@ -94,10 +95,9 @@ def __get_param(self, source, name, default=None, message=None):
                 return default
 
     def __check_tlp(self):
-        """Check if tlp is okay or not; reports error if too high."""
+        """Check if tlp is okay or not; returns False if too high."""
 
-        if self.check_tlp and self.tlp > self.max_tlp:
-            self.error('TLP is higher than allowed.')
+        return not (self.enable_check_tlp and self.tlp > self.max_tlp)
 
     def get_data(self):
         """Wrapper for getting data from input dict.

contrib/tests/fixtures/test-minimal-config.json

@@ -1,4 +1,4 @@
 {
     "dataType": "ip",
-    "data": "8.8.8.8"
+    "data": "1.1.1.1"
 }

contrib/tests/fixtures/test-tlp-config.json

@@ -0,0 +1,8 @@
+{
+    "dataType": "ip",
+    "data": "1.1.1.1",
+    "config": {
+        "check_tlp": true,
+        "max_tlp": 2
+    }
+}

contrib/tests/suite.py

@@ -26,21 +26,18 @@ def setUp(self):
     def test_default_config(self):
         self.assertEqual(self.analyzer.data_type, 'ip')
         self.assertEqual(self.analyzer.tlp, 2)
-        self.assertEqual(self.analyzer.check_tlp, False)
+        self.assertEqual(self.analyzer.enable_check_tlp, False)
         self.assertEqual(self.analyzer.max_tlp, 2)
         self.assertEqual(self.analyzer.http_proxy, None)
         self.assertEqual(self.analyzer.https_proxy, None)
 
-        self.assertEqual(self.analyzer.get_param('dataType'), "ip")
-        print('')
-
     def test_artifact_data(self):
-        self.assertEqual(self.analyzer.getData(), "8.8.8.8")
-        self.assertEqual(self.analyzer.get_data(), "8.8.8.8")
+        self.assertEqual(self.analyzer.getData(), "1.1.1.1")
+        self.assertEqual(self.analyzer.get_data(), "1.1.1.1")
 
     def test_params_data(self):
-        self.assertEqual(self.analyzer.getParam('data'), "8.8.8.8")
-        self.assertEqual(self.analyzer.get_param('data'), "8.8.8.8")
+        self.assertEqual(self.analyzer.getParam('data'), "1.1.1.1")
+        self.assertEqual(self.analyzer.get_param('data'), "1.1.1.1")
 
 class TestProxyConfig(unittest.TestCase):
 
@@ -57,5 +54,36 @@ def test_proxy_config(self):
         self.assertEqual(os.environ['http_proxy'], proxy_url)
         self.assertEqual(os.environ['https_proxy'], proxy_url)
 
+class TestTlpConfig(unittest.TestCase):
+
+    def setUp(self):
+        load_test_fixture('fixtures/test-tlp-config.json')
+        self.analyzer = Analyzer()
+
+    def test_check_tlp_disabled(self):
+        self.analyzer.enable_check_tlp = False
+
+        # Using the _Analyzer__check_tlp notation to access managed method
+        # __check_tlp
+        self.assertEqual(self.analyzer._Analyzer__check_tlp(), True)
+
+    def test_check_tlp_ko(self):
+        self.analyzer.enable_check_tlp = True
+        self.analyzer.max_tlp = 1
+        self.analyzer.tlp = 3
+
+        # Using the _Analyzer__check_tlp notation to access managed method
+        # __check_tlp
+        self.assertEqual(self.analyzer._Analyzer__check_tlp(), False)
+
+    def test_check_tlp_ok(self):
+        self.analyzer.enable_check_tlp = True
+        self.analyzer.max_tlp = 3
+        self.analyzer.tlp = 3
+
+        # Using the _Analyzer__check_tlp notation to access managed method
+        # __check_tlp
+        self.assertEqual(self.analyzer._Analyzer__check_tlp(), True)
+
 if __name__ == '__main__':
     unittest.main()