mirror of
https://github.com/Security-Onion-Solutions/securityonion.git
synced 2025-12-08 10:12:53 +01:00
82 lines
3.4 KiB
Python
82 lines
3.4 KiB
Python
from io import StringIO
|
|
import sys
|
|
from unittest.mock import patch, MagicMock
|
|
import unittest
|
|
from urlhaus import urlhaus
|
|
|
|
|
|
class TestUrlhausMethods(unittest.TestCase):
|
|
|
|
def test_main_success(self):
|
|
output = {"foo": "bar"}
|
|
config = {"api_key": "test_key"}
|
|
with patch('sys.stdout', new=StringIO()) as mock_stdout:
|
|
with patch('urlhaus.urlhaus.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"]
|
|
urlhaus.main()
|
|
expected = '{"foo": "bar"}\n'
|
|
self.assertEqual(mock_stdout.getvalue(), expected)
|
|
mock_analyze.assert_called_once()
|
|
mock_config.assert_called_once()
|
|
|
|
def test_buildReq(self):
|
|
result = urlhaus.buildReq("test")
|
|
self.assertEqual("test", result["url"])
|
|
|
|
def test_sendReq(self):
|
|
with patch('requests.request', new=MagicMock(return_value=MagicMock())) as mock:
|
|
conf = {"api_key": "test_key"}
|
|
meta = {"baseUrl": "myurl"}
|
|
response = urlhaus.sendReq(conf, meta, "mypayload")
|
|
mock.assert_called_once_with("POST", "myurl", data="mypayload", headers={"Auth-Key": "test_key"})
|
|
self.assertIsNotNone(response)
|
|
|
|
def test_prepareResults_none(self):
|
|
raw = {"query_status": "no_results"}
|
|
results = urlhaus.prepareResults(raw)
|
|
self.assertEqual(results["response"], raw)
|
|
self.assertEqual(results["summary"], "no_results")
|
|
self.assertEqual(results["status"], "ok")
|
|
|
|
def test_prepareResults_invalidUrl(self):
|
|
raw = {"query_status": "invalid_url"}
|
|
results = urlhaus.prepareResults(raw)
|
|
self.assertEqual(results["response"], raw)
|
|
self.assertEqual(results["summary"], "invalid_url")
|
|
self.assertEqual(results["status"], "caution")
|
|
|
|
def test_prepareResults_threat(self):
|
|
raw = {"query_status": "invalid_url"} # This is overrided in this scenario
|
|
raw["threat"] = "bad_actor"
|
|
results = urlhaus.prepareResults(raw)
|
|
self.assertEqual(results["response"], raw)
|
|
self.assertEqual(results["summary"], "bad_actor")
|
|
self.assertEqual(results["status"], "threat")
|
|
|
|
def test_prepareResults_error(self):
|
|
raw = {}
|
|
results = urlhaus.prepareResults(raw)
|
|
self.assertEqual(results["response"], raw)
|
|
self.assertEqual(results["summary"], "internal_failure")
|
|
self.assertEqual(results["status"], "caution")
|
|
|
|
def test_analyze(self):
|
|
output = {"threat": "malware_download"}
|
|
config = {"api_key": "test_key"}
|
|
artifactInput = '{"value":"foo","artifactType":"url"}'
|
|
with patch('urlhaus.urlhaus.sendReq', new=MagicMock(return_value=output)) as mock:
|
|
results = urlhaus.analyze(config, artifactInput)
|
|
self.assertEqual(results["summary"], "malware_download")
|
|
mock.assert_called_once()
|
|
|
|
def test_checkConfigRequirements_valid(self):
|
|
config = {"api_key": "test_key"}
|
|
self.assertTrue(urlhaus.checkConfigRequirements(config))
|
|
|
|
def test_checkConfigRequirements_missing_key(self):
|
|
config = {}
|
|
with self.assertRaises(SystemExit) as cm:
|
|
urlhaus.checkConfigRequirements(config)
|
|
self.assertEqual(cm.exception.code, 126)
|