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_missing_input(self): with patch('sys.stdout', new=StringIO()) as mock_cmd: sys.argv = ["cmd"] malwarebazaar.main() self.assertEqual(mock_cmd.getvalue(), 'ERROR: Input is not in proper JSON format\n') 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: 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(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(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( 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'}]} 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(input) results2 = malwarebazaar.analyze(input2) results3 = malwarebazaar.analyze(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'}]} 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(input) results2 = malwarebazaar.analyze(input2) results3 = malwarebazaar.analyze(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: response = malwarebazaar.sendReq( {'baseUrl': 'https://www.randurl.xyz'}, 'example_data') self.assertIsNotNone(response) mock.assert_called_once()