From 7d6f8d922b695b2899f0d73c6b61373bfe7d4cf2 Mon Sep 17 00:00:00 2001 From: Elijah Gibson <82835625+semphorin@users.noreply.github.com> Date: Mon, 18 Dec 2023 15:57:41 -0500 Subject: [PATCH] Update malwarebazaar_test.py Flake8 linting --- .../malwarebazaar/malwarebazaar_test.py | 147 +++++++++++------- 1 file changed, 95 insertions(+), 52 deletions(-) diff --git a/salt/sensoroni/files/analyzers/malwarebazaar/malwarebazaar_test.py b/salt/sensoroni/files/analyzers/malwarebazaar/malwarebazaar_test.py index 79903142d..b9f3112f2 100644 --- a/salt/sensoroni/files/analyzers/malwarebazaar/malwarebazaar_test.py +++ b/salt/sensoroni/files/analyzers/malwarebazaar/malwarebazaar_test.py @@ -15,61 +15,68 @@ class TestMalwarebazaarMethods(unittest.TestCase): def test_main_success(self): with patch('sys.stdout', new=StringIO()) as mock_cmd: - with patch('malwarebazaar.malwarebazaar.analyze', new=MagicMock(return_value={'test': 'val'})) as mock: + with patch('malwarebazaar.malwarebazaar.analyze', + new=MagicMock(return_value={'test': 'val'})) as mock: sys.argv = ["cmd", "input"] malwarebazaar.main() expected = '{"test": "val"}\n' self.assertEqual(mock_cmd.getvalue(), expected) mock.assert_called_once() - + def isInJson_string_found_in_dict(self): test_string = "helo" input_json = { - "value":"test", - "test":"value", - "arr":["Foo", "Bar", "Hello"], - "dict1":{"helo":"val", "key":"val"} - } - self.assertEqual(malwarebazaar.isInJson(test_string), True) - + "value": "test", + "test": "value", + "arr": ["Foo", "Bar", "Hello"], + "dict1": {"helo": "val", "key": "val"} + } + self.assertEqual(malwarebazaar.isInJson(input_json, test_string), True) + def isInJson_string_found_in_arr(self): test_string = "helo" input_json = { - "value":"test", - "test":"value", - "arr":["Foo", "Bar", "helo"], - "dict1":{"Hello":"val", "key":"val"} - } - self.assertEqual(malwarebazaar.isInJson(test_string), True) - + "value": "test", + "test": "value", + "arr": ["Foo", "Bar", "helo"], + "dict1": {"Hello": "val", "key": "val"} + } + self.assertEqual(malwarebazaar.isInJson(input_json, test_string), True) + def isInJson_string_not_found(self): test_string = "ValNotInJSON" input_json = { - "value":"test", - "test":"value", - "arr":["Foo", "Bar", "helo"], - "dict1":{"Hello":"val", "key":"val"} - } - self.assertEqual(malwarebazaar.isInJson(test_string), False) + "value": "test", + "test": "value", + "arr": ["Foo", "Bar", "helo"], + "dict1": {"Hello": "val", "key": "val"} + } + self.assertEqual(malwarebazaar.isInJson( + input_json, test_string), False) def test_analyze(self): - """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 'no result' """ - sendReqOutput = {'threat': 'no_result', "query_status": "ok", 'data': [{'sha256_hash': 'notavalidhash'}]} + """simulated sendReq and prepareResults with 2 mock objects + and variables sendReqOutput and prep_res_sim, + input created for analyze method call + and then we compared results['summary'] with 'no result' """ + sendReqOutput = {'threat': 'no_result', "query_status": "ok", + 'data': [{'sha256_hash': 'notavalidhash'}]} input = '{"artifactType": "hash", "value": "1234"}' input2 = '{"artifactType": "tlsh", "value": "1234"}' input3 = '{"artifactType": "gimphash", "value": "1234"}' - prepareResultOutput = {'response': '', - 'summary': 'no result', 'status': 'info'} + prep_res_sim = {'response': '', + 'summary': 'no result', 'status': 'info'} - with patch('malwarebazaar.malwarebazaar.sendReq', new=MagicMock(return_value=sendReqOutput)) as mock: - with patch('malwarebazaar.malwarebazaar.prepareResults', new=MagicMock(return_value=prepareResultOutput)) as mock2: + with patch('malwarebazaar.malwarebazaar.sendReq', + new=MagicMock(return_value=sendReqOutput)) as mock: + with patch('malwarebazaar.malwarebazaar.prepareResults', + new=MagicMock(return_value=prep_res_sim)) as mock2: results = malwarebazaar.analyze(input) results2 = malwarebazaar.analyze(input2) results3 = malwarebazaar.analyze(input3) - self.assertEqual(results["summary"], prepareResultOutput['summary']) - self.assertEqual(results2["summary"], prepareResultOutput['summary']) - self.assertEqual(results3["summary"], prepareResultOutput['summary']) + self.assertEqual(results["summary"], prep_res_sim['summary']) + self.assertEqual(results2["summary"], prep_res_sim['summary']) + self.assertEqual(results3["summary"], prep_res_sim['summary']) self.assertEqual(results["status"], "info") self.assertEqual(results2["status"], "info") self.assertEqual(results3["status"], "info") @@ -77,23 +84,28 @@ class TestMalwarebazaarMethods(unittest.TestCase): mock.assert_called() def test_analyze_result(self): - """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 'no result' """ - sendReqOutput = {'threat': 'threat', "query_status": "notok", 'data': [{'sha256_hash': 'validhash'}]} + """simulated sendReq and prepareResults with 2 mock objects + and variables sendReqOutput and prep_res_sim, + input created for analyze method call + and then we compared results['summary'] with 'no result' """ + sendReqOutput = {'threat': 'threat', "query_status": "notok", 'data': [ + {'sha256_hash': 'validhash'}]} input = '{"artifactType": "hash", "value": "1234"}' input2 = '{"artifactType": "tlsh", "value": "1234"}' input3 = '{"artifactType": "gimphash", "value": "1234"}' - prepareResultOutput = {'response': '', - 'summary': 'Bad', 'status': 'threat'} + prep_res_sim = {'response': '', + 'summary': 'Bad', 'status': 'threat'} - with patch('malwarebazaar.malwarebazaar.sendReq', new=MagicMock(return_value=sendReqOutput)) as mock: - with patch('malwarebazaar.malwarebazaar.prepareResults', new=MagicMock(return_value=prepareResultOutput)) as mock2: + with patch('malwarebazaar.malwarebazaar.sendReq', + new=MagicMock(return_value=sendReqOutput)) as mock: + with patch('malwarebazaar.malwarebazaar.prepareResults', + new=MagicMock(return_value=prep_res_sim)) as mock2: results = malwarebazaar.analyze(input) results2 = malwarebazaar.analyze(input2) results3 = malwarebazaar.analyze(input3) - self.assertEqual(results["summary"], prepareResultOutput['summary']) - self.assertEqual(results2["summary"], prepareResultOutput['summary']) - self.assertEqual(results3["summary"], prepareResultOutput['summary']) + self.assertEqual(results["summary"], prep_res_sim['summary']) + self.assertEqual(results2["summary"], prep_res_sim['summary']) + self.assertEqual(results3["summary"], prep_res_sim['summary']) self.assertEqual(results["status"], "threat") self.assertEqual(results2["status"], "threat") self.assertEqual(results3["status"], "threat") @@ -110,47 +122,76 @@ class TestMalwarebazaarMethods(unittest.TestCase): def test_prepareResults_empty(self): # raw is empty raw = {} - expected = {'response': raw, 'status': 'caution', 'summary': 'internal_failure'} + expected = {'response': raw, 'status': 'caution', + 'summary': 'internal_failure'} results = malwarebazaar.prepareResults(raw) self.assertEqual(results, expected) def test_prepareResults_threat(self): - raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', 'vendor_intel': {'ReversingLabs': {'status': 'MALICIOUS'}}, 'signature': 'abcd1234', 'tags': ['tag1']}]} + raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', + 'vendor_intel': + {'ReversingLabs': + {'status': + 'MALICIOUS'}}, + 'signature': 'abcd1234', + 'tags': ['tag1']}]} expected = {'response': raw, 'status': 'threat', 'summary': 'abcd1234'} results = malwarebazaar.prepareResults(raw) self.assertEqual(results, expected) def test_prepareResults_caution(self): # raw is empty - raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', 'vendor_intel': {'Triage': {'score': '6'}}, 'signature': 'abcd1234', 'tags': ['tag1']}]} - expected = {'response': raw, 'status': 'caution', 'summary': 'abcd1234'} + raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', + 'vendor_intel': + {'Triage': {'score': '6'}}, + 'signature': 'abcd1234', + 'tags': ['tag1']}]} + expected = {'response': raw, + 'status': 'caution', 'summary': 'abcd1234'} results = malwarebazaar.prepareResults(raw) self.assertEqual(results, expected) def test_prepareResults_info(self): # raw is empty - raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', 'vendor_intel': {'Triage': {'score': '3'}}, 'signature': 'abcd1234', 'tags': ['tag1']}]} + raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', + 'vendor_intel': + {'Triage': {'score': '3'}}, + 'signature': 'abcd1234', + 'tags': ['tag1']}]} expected = {'response': raw, 'status': 'info', 'summary': 'abcd1234'} results = malwarebazaar.prepareResults(raw) self.assertEqual(results, expected) def test_prepareResults_ok(self): # raw is empty - raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', 'vendor_intel': {'Triage': {'score': '1'}}, 'signature': 'abcd1234', 'tags': ['tag1']}]} + raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', + 'vendor_intel': + {'Triage': {'score': '1'}}, + 'signature': 'abcd1234', + 'tags': ['tag1']}]} expected = {'response': raw, 'status': 'ok', 'summary': 'abcd1234'} results = malwarebazaar.prepareResults(raw) self.assertEqual(results, expected) def test_prepareResults_ok_tags(self): # raw is empty - raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', 'vendor_intel': {'Triage': {'score': '1'}}, 'tags': ['tag1']}]} + raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', + 'vendor_intel': + {'Triage': {'score': '1'}}, + 'tags': ['tag1']}]} expected = {'response': raw, 'status': 'ok', 'summary': 'tag1'} results = malwarebazaar.prepareResults(raw) self.assertEqual(results, expected) def test_prepareResults_ok_yomi(self): # raw is empty - raw = {'query_status': 'ok', 'data': [{'sha256_hash': 'validhash', 'vendor_intel': {'YOROI_YOMI': {'detection': 'detection1', 'summary': '0.1'}}}]} + raw = {'query_status': 'ok', + 'data': [{'sha256_hash': 'validhash', + 'vendor_intel': + {'YOROI_YOMI': + {'detection': + 'detection1', + 'summary': '0.1'}}}]} expected = {'response': raw, 'status': 'ok', 'summary': 'detection1'} results = malwarebazaar.prepareResults(raw) self.assertEqual(results, expected) @@ -170,9 +211,11 @@ class TestMalwarebazaarMethods(unittest.TestCase): self.assertEqual( result, {'query': 'get_tlsh', 'tlsh': ''}) - # simulate API response and makes sure sendReq gives a response, we are just checking if sendReq gives back anything + # simulate API response and makes sure sendReq gives a response, + # we are just checking if sendReq gives back anything def test_sendReq(self): - with patch('requests.post', new=MagicMock(return_value=MagicMock())) as mock: + with patch('requests.post', + new=MagicMock(return_value=MagicMock())) as mock: response = malwarebazaar.sendReq( {'baseUrl': 'https://www.randurl.xyz'}, 'example_data') self.assertIsNotNone(response)