from io import StringIO import sys from unittest.mock import patch, MagicMock from malwarebazaar import malwarebazaar import unittest class TestMalwarebazaarMethods(unittest.TestCase): def test_main_success(self): output = {"test": "val"} config = {"api_key": "test_key"} with patch('sys.stdout', new=StringIO()) as mock_stdout: with patch('malwarebazaar.malwarebazaar.analyze', new=MagicMock(return_value=output)) as mock_analyze: with patch('helpers.loadConfig', new=MagicMock(return_value=config)) as mock_config: sys.argv = ["cmd", "input"] malwarebazaar.main() expected = '{"test": "val"}\n' self.assertEqual(mock_stdout.getvalue(), expected) mock_analyze.assert_called_once() mock_config.assert_called_once() def test_isInJson_tail_greater_than_max_depth(self): max_depth = 1000 tail = 2000 test_string = "helo" input_json = { "value": "test", "test": "value", "arr": ["Foo", "Bar", "Hello"], "dict1": {"key1": "val", "key2": "helo"} } self.assertEqual(malwarebazaar.isInJson(input_json, test_string, max_depth, tail), False) def test_isInJson_string_found_in_dict(self): test_string = "helo" input_json = { "value": "test", "test": "value", "arr": ["Foo", "Bar", "Hello"], "dict1": {"key1": "val", "key2": "helo"} } self.assertEqual(malwarebazaar.isInJson(input_json, test_string), True) def test_isInJson_dict_in_list(self): max_depth = 1000 tail = 1 test_string = "helo" input_json = { "key1": "test", "key2": "value", "key3": ["Foo", "Bar", "Hello"], "nested_list": [{"key1": "val", "key2": "helo"}] } self.assertEqual(malwarebazaar.isInJson(input_json, test_string, max_depth, tail), True) def test_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(input_json, test_string), True) def test_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(input_json, test_string), False) def test_analyze(self): """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'}]} config = {"api_key": "test_key"} input = '{"artifactType": "hash", "value": "1234"}' input2 = '{"artifactType": "tlsh", "value": "1234"}' input3 = '{"artifactType": "gimphash", "value": "1234"}' 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=prep_res_sim)) as mock2: results = malwarebazaar.analyze(config, input) results2 = malwarebazaar.analyze(config, input2) results3 = malwarebazaar.analyze(config, input3) 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") mock2.assert_called() mock.assert_called() def test_analyze_result(self): """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'}]} config = {"api_key": "test_key"} input = '{"artifactType": "hash", "value": "1234"}' input2 = '{"artifactType": "tlsh", "value": "1234"}' input3 = '{"artifactType": "gimphash", "value": "1234"}' 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=prep_res_sim)) as mock2: results = malwarebazaar.analyze(config, input) results2 = malwarebazaar.analyze(config, input2) results3 = malwarebazaar.analyze(config, input3) 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") mock2.assert_called() mock.assert_called() def test_prepareResults_illegal_search_term(self): # illegal search term raw = {'query_status': 'illegal_search_term'} expected = {'response': raw, 'status': 'info', 'summary': 'no result'} results = malwarebazaar.prepareResults(raw) self.assertEqual(results, expected) def test_prepareResults_empty(self): # raw is empty raw = {} 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']}]} 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'} 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']}]} 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']}]} 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']}]} 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'}}}]} expected = {'response': raw, 'status': 'ok', 'summary': 'detection1'} results = malwarebazaar.prepareResults(raw) self.assertEqual(results, expected) def test_buildReqGimqhash(self): result = malwarebazaar.buildReq('gimphash', '') self.assertEqual( result, {'query': 'get_gimphash', 'gimphash': ''}) def test_buildReqHash(self): result = malwarebazaar.buildReq('hash', '') self.assertEqual( result, {'query': 'get_info', 'hash': ''}) def test_buildReqtlshhash(self): result = malwarebazaar.buildReq('tlsh', '') 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 def test_sendReq(self): with patch('requests.post', new=MagicMock(return_value=MagicMock())) as mock: conf = {"api_key": "test_key"} response = malwarebazaar.sendReq( conf, {'baseUrl': 'https://www.randurl.xyz'}, 'example_data') self.assertIsNotNone(response) mock.assert_called_once() def test_checkConfigRequirements_valid(self): config = {"api_key": "test_key"} self.assertTrue(malwarebazaar.checkConfigRequirements(config)) def test_checkConfigRequirements_missing_key(self): config = {} with self.assertRaises(SystemExit) as cm: malwarebazaar.checkConfigRequirements(config) self.assertEqual(cm.exception.code, 126)