Files
securityonion/salt/sensoroni/files/analyzers/malwarebazaar/malwarebazaar_test.py
2025-08-20 20:59:23 -05:00

255 lines
12 KiB
Python

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)