fix Elasticsearch lint

salt/sensoroni/files/analyzers/elasticsearch/elasticsearch.py

@@ -30,10 +30,10 @@ def checkConfigRequirements(conf):
     if not conf['base_url']:
         sys.exit(126)
     if not conf['timestamp_field_name']:
-        sys.exit(126)   
+        sys.exit(126)
     if not conf['cert_path']:
-        sys.exit(126)    
-    return True   
+        sys.exit(126)
+    return True
 
 
 def buildReq(conf, input):
@@ -41,7 +41,7 @@ def buildReq(conf, input):
     # based off of user configurable values
     num_results = conf['num_results']
 
-    if conf['map'] != None:
+    if conf['map'] is not None:
         mappings = conf['map']
     else:
         mappings = dict()
@@ -61,7 +61,7 @@ def buildReq(conf, input):
             "bool": {
                 "must": [{
                     "wildcard": {
-                        type: input['value'],                        
+                        type: input['value'],
                     },
                 }
                 ],

salt/sensoroni/files/analyzers/elasticsearch/elasticsearch_test.py

@@ -3,11 +3,12 @@ import sys
 from unittest.mock import patch, MagicMock
 import unittest
 import elasticsearch
-import helpers
 import json
 from datetime import datetime, timedelta
 
-class TestElasticSearchMethods(unittest.TestCase):   
+
+class TestElasticSearchMethods(unittest.TestCase):
+
     '''Test that the analyzer main method work as expect when not given enough input'''
     def test_main_missing_input(self):
         with patch('sys.exit', new=MagicMock()) as sysmock:
@@ -19,137 +20,74 @@ class TestElasticSearchMethods(unittest.TestCase):
 
     '''Test that analyzer main method work as expect when all required input is given'''
     def test_main_success(self):
-            conf = {"base_url":"test", "auth_user":"test", "auth_pwd":"test", "num_results":10,"api_key":"test","index":"test","time_delta_minutes": 14400,"timestamp_field_name":"test", "map":{}, "cert_path":""}
-            with patch('elasticsearch.helpers.loadConfig', new=MagicMock(return_value=conf))as mock_yaml:
-                with patch('sys.stdout', new=StringIO()) as mock_cmd:
-                    with patch('elasticsearch.analyze', new=MagicMock(return_value={'foo': 'bar'})) as mock:
-                            sys.argv = ["cmd", "conf"]
-                            elasticsearch.main()
-                            expected = '{"foo": "bar"}\n'
-                            self.assertEqual(mock_cmd.getvalue(), expected)
-                            mock.assert_called_once()
-                            mock_yaml.assert_called_once()
-                            
+        conf = {"base_url": "test", "auth_user": "test", "auth_pwd": "test", "api_key": "test", "index": "test", "time_delta_minutes": 14400, "map": {}, "cert_path": ""}
+        with patch('elasticsearch.helpers.loadConfig', new=MagicMock(return_value=conf))as mock_yaml:
+            with patch('sys.stdout', new=StringIO()) as mock_cmd:
+                with patch('elasticsearch.analyze', new=MagicMock(return_value={'foo': 'bar'})) as mock:
+                    sys.argv = ["cmd", "conf"]
+                    elasticsearch.main()
+                    expected = '{"foo": "bar"}\n'
+                    self.assertEqual(mock_cmd.getvalue(), expected)
+                    mock.assert_called_once()
+                    mock_yaml.assert_called_once()
+
     '''Test that checks for empty and none values in configurables'''
     def test_checkConfigRequirements(self):
-        conf = {"base_url":"", "auth_user":"", "auth_pwd":"", "num_results":None,"api_key":"","index":"","time_delta_minutes": None,"timestamp_field_name":"", "map":{}, "cert_path":""}
+        conf = {"base_url": "", "auth_user": "", "auth_pwd": "", "num_results": None, "api_key": "", "index": "", "time_delta_minutes": None, "timestamp_field_name": "", "map": {}, "cert_path": ""}
         with self.assertRaises(SystemExit) as cm:
             elasticsearch.checkConfigRequirements(conf)
         self.assertEqual(cm.exception.code, 126)
-
-
     '''Test that checks buildReq method, by comparing a mock buildReq result with an expectedQuery, used a mock object to simulate an expectedQuery
         since Elasticsearch buildReq uses values in the config'''
     def test_buildReq(self):
         numberOfResults = 1
         observableType = "hash"
         expectedQuery = {
-        "from": 0,
-        "size": numberOfResults,
-        "query": {
-            "bool": {
-                "must": [{
-                    "wildcard": {
-                        observableType: observableType,
-                    },
-                }
-                ],
-                "filter": {
-                    "range": {
-                        "@timestamp": {
-                            "gte": ('2023-11-29T14:23:45'),
-                            "lte": ('2023-11-29T14:23:45')
+            "from": 0,
+            "size": numberOfResults,
+            "query": {
+                "bool": {
+                    "must": [{
+                        "wildcard": {
+                            observableType: observableType,
+                        },
+                    }
+                    ],
+                    "filter": {
+                        "range": {
+                            "@timestamp": {
+                                "gte": ('2023-11-29T14:23:45'),
+                                "lte": ('2023-11-29T14:23:45')
+                            }
                         }
                     }
                 }
             }
         }
-    }      
         with patch('elasticsearch.buildReq', new=MagicMock(return_value=expectedQuery)) as mock:
-            response = elasticsearch.buildReq(observableType,numberOfResults)     
+            response = elasticsearch.buildReq(observableType, numberOfResults)
             self.assertEqual(json.dumps(response), json.dumps(expectedQuery))
             mock.assert_called_once()
 
     def test_wrongbuildReq(self):
-            result={'map':'123','artifactType':'hash','timestamp_field_name':'abc', 'time_delta_minutes':14400, 'num_results':10,'value':'0' }
-            cur_time = datetime.now()
-            start_time = cur_time - timedelta(minutes=result['time_delta_minutes'])
-            query=elasticsearch.buildReq(result, result)
-            comparequery=json.dumps({
-                "from": 0,
-                "size":10,
-                "query": {
-                    "bool":{
-                        "must": [{
-                            "wildcard": {
-                                'hash': result['value'],
-                            },
-                        }
-                        ],
-                        "filter":{
-                            "range":{
-                                result['timestamp_field_name']:{
-                                    "gte": start_time.strftime('%Y-%m-%dT%H:%M:%S'),
-                                    "lte": cur_time.strftime('%Y-%m-%dT%H:%M:%S')
-                                }
-                            }
-                        }
-                    }
-                }
-                
-            })
-            self.assertEqual(query, comparequery )
-            
-    def test_rightbuildReq(self):
-            result={'map':{'hash':'testingHash'},'artifactType':'hash','timestamp_field_name':'abc', 'time_delta_minutes':14400, 'num_results':10,'value':'0'}
-            cur_time = datetime.now()
-            start_time = cur_time - timedelta(minutes=result['time_delta_minutes'])
-            query=elasticsearch.buildReq(result, result)
-            comparequery=json.dumps({
-                "from": 0,
-                "size": 10,
-                "query": {
-                    "bool":{
-                        "must":[{
-                                "wildcard": {
-                                    result['map'][result['artifactType']]: result['value'],
-                                },
-                            }
-                        ]
-                        ,
-                        "filter":{
-                            "range":{
-                                result['timestamp_field_name']:{
-                                    "gte": start_time.strftime('%Y-%m-%dT%H:%M:%S'),
-                                    "lte": cur_time.strftime('%Y-%m-%dT%H:%M:%S')
-                                }
-                            }
-                        }
-                    }
-                }
-            })
-            self.assertEqual(query, comparequery )
-
-    def test_rightbuildReq100result(self):
-        result={'map':{'hash':'testingHash'},'artifactType':'hash','timestamp_field_name':'abc', 'time_delta_minutes':14400, 'num_results':100,'value':'0'}
+        result = {'map': '123', 'artifactType': 'hash', 'timestamp_field_name': 'abc', 'time_delta_minutes': 14400, 'num_results': 10, 'value': '0'}
         cur_time = datetime.now()
         start_time = cur_time - timedelta(minutes=result['time_delta_minutes'])
-        query=elasticsearch.buildReq(result, result)
-        comparequery=json.dumps({
+        query = elasticsearch.buildReq(result, result)
+        comparequery = json.dumps({
             "from": 0,
-            "size": 100,
+            "size": 10,
             "query": {
-                "bool":{
-                    "must":[{
-                            "wildcard": {
-                                result['map'][result['artifactType']]: result['value'],
-                            },
-                        }
-                    ]
-                    ,
-                    "filter":{
-                        "range":{
-                            result['timestamp_field_name']:{
+                "bool": {
+                    "must": [{
+                        "wildcard": {
+                            'hash': result['value'],
+                        },
+                    }
+                    ],
+                    "filter": {
+                        "range": {
+                            result['timestamp_field_name']: {
                                 "gte": start_time.strftime('%Y-%m-%dT%H:%M:%S'),
                                 "lte": cur_time.strftime('%Y-%m-%dT%H:%M:%S')
                             }
@@ -158,15 +96,71 @@ class TestElasticSearchMethods(unittest.TestCase):
                 }
             }
         })
-        self.assertEqual(query, comparequery )
+        self.assertEqual(query, comparequery)
 
+    def test_rightbuildReq(self):
+        result = {'map': {'hash': 'testingHash'}, 'artifactType': 'hash', 'timestamp_field_name': 'abc', 'time_delta_minutes': 14400, 'num_results': 10, 'value': '0'}
+        cur_time = datetime.now()
+        start_time = cur_time - timedelta(minutes=result['time_delta_minutes'])
+        query = elasticsearch.buildReq(result, result)
+        comparequery = json.dumps({
+            "from": 0,
+            "size": 10,
+            "query": {
+                "bool": {
+                    "must": [{
+                            "wildcard": {
+                                result['map'][result['artifactType']]: result['value'],
+                            },
+                    }],
+                    "filter": {
+                        "range": {
+                            result['timestamp_field_name']: {
+                                "gte": start_time.strftime('%Y-%m-%dT%H:%M:%S'),
+                                "lte": cur_time.strftime('%Y-%m-%dT%H:%M:%S')
+                            }
+                        }
+                    }
+                }
+            }
+        })
+        self.assertEqual(query, comparequery)
+
+    def test_rightbuildReq100result(self):
+        result = {'map': {'hash': 'testingHash'}, 'artifactType': 'hash', 'timestamp_field_name': 'abc', 'time_delta_minutes': 14400, 'num_results': 100, 'value': '0'}
+        cur_time = datetime.now()
+        start_time = cur_time - timedelta(minutes=result['time_delta_minutes'])
+        query = elasticsearch.buildReq(result, result)
+        comparequery = json.dumps({
+            "from": 0,
+            "size": 100,
+            "query": {
+                "bool": {
+                    "must": [{
+                            "wildcard": {
+                                result['map'][result['artifactType']]: result['value'],
+                            },
+                    }],
+                    "filter": {
+                        "range": {
+                            result['timestamp_field_name']: {
+                                "gte": start_time.strftime('%Y-%m-%dT%H:%M:%S'),
+                                "lte": cur_time.strftime('%Y-%m-%dT%H:%M:%S')
+                            }
+                        }
+                    }
+                }
+            }
+        })
+        self.assertEqual(query, comparequery)
 
     '''Test that checks sendReq method to expect a response from a requests.post'''
     def test_sendReq(self):
-        conf = {"base_url":"test", "auth_user":"test", "auth_pwd":"test", "api_key":"test","index":"test", "cert_path":""}
+        conf = {"base_url": "test", "auth_user": "test", "auth_pwd": "test", "api_key": "test", "index": "test", "cert_path": ""}
         with patch('requests.post', new=MagicMock(return_value=MagicMock())) as mock:
             response = elasticsearch.sendReq(conf, 'example_query')
-            self.assertIsNotNone(response)    
+            self.assertIsNotNone(response)
+            mock.assert_called_once
 
     '''Test that checks prepareResults method, by comparing a mock prepareResults return_value with an expectedResult'''
     def test_prepareResults(self):
@@ -179,16 +173,17 @@ class TestElasticSearchMethods(unittest.TestCase):
             response = elasticsearch.prepareResults(raw)
             self.assertEqual(expectedResult, response)
             mock.assert_called_once()
-
     '''Test that checks analyze method, simulated sendReq and prepareResults with 2 mock objects and variables sendReqOutput and prepareResultOutput,
-            input created for analyze method call and then we compared results['summary'] with 'Documents returned: 5' '''    
+            input created for analyze method call and then we compared results['summary'] with 'Documents returned: 5' '''
     def test_analyze(self):
         sendReqOutput = {'_id': "0", "hash": "123"}
-        input = '{"artifactType":"hash", "value":"123"}'
-        prepareResultOutput = {'response': {'_id': "0", "hash": "123"},'summary': "Documents returned: 5", 'status': 'info'}
-        conf = {"base_url":"test", "auth_user":"test", "auth_pwd":"test", "num_results":10,"api_key":"test","index":"test","time_delta_minutes": 14400,"timestamp_field_name":"test", "map":{}, "cert_path":"test"}
+        input = '{"artifactType": "hash", "value": "123"}'
+        prepareResultOutput = {'response': {'_id': "0", "hash": "123"}, 'summary': "Documents returned: 5", 'status': 'info'}
+        conf = {"base_url": "test", "auth_user": "test", "auth_pwd": "test", "num_results": 10, "api_key": "test", "index": "test",
+                "time_delta_minutes": 14400, "timestamp_field_name": "test", "map": {}, "cert_path": "test"}
         with patch('elasticsearch.sendReq', new=MagicMock(return_value=sendReqOutput)) as mock:
             with patch('elasticsearch.prepareResults', new=MagicMock(return_value=prepareResultOutput)) as mock2:
                 results = elasticsearch.analyze(conf, input)
                 self.assertEqual(results["summary"], "Documents returned: 5")
-                mock.assert_called_once()
+                mock.assert_called_once()
+                mock2.assert_called_once()