mirror of
https://github.com/Security-Onion-Solutions/securityonion.git
synced 2025-12-06 17:22:49 +01:00
73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
from io import StringIO
|
|
import sys
|
|
from unittest.mock import patch, MagicMock
|
|
from urlhaus import urlhaus
|
|
import unittest
|
|
|
|
|
|
class TestUrlhausMethods(unittest.TestCase):
|
|
|
|
def test_main_missing_input(self):
|
|
with patch('sys.stdout', new=StringIO()) as mock_stdout:
|
|
sys.argv = ["cmd"]
|
|
urlhaus.main()
|
|
self.assertEqual(mock_stdout.getvalue(), "ERROR: Missing input JSON\n")
|
|
|
|
def test_main_success(self):
|
|
output = {"foo": "bar"}
|
|
with patch('sys.stdout', new=StringIO()) as mock_stdout:
|
|
with patch('urlhaus.urlhaus.analyze', new=MagicMock(return_value=output)) as mock:
|
|
sys.argv = ["cmd", "input"]
|
|
urlhaus.main()
|
|
expected = '{"foo": "bar"}\n'
|
|
self.assertEqual(mock_stdout.getvalue(), expected)
|
|
mock.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:
|
|
meta = {"baseUrl": "myurl"}
|
|
response = urlhaus.sendReq(meta, "mypayload")
|
|
mock.assert_called_once_with("POST", "myurl", data="mypayload")
|
|
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"}
|
|
artifactInput = '{"value":"foo","artifactType":"url"}'
|
|
with patch('urlhaus.urlhaus.sendReq', new=MagicMock(return_value=output)) as mock:
|
|
results = urlhaus.analyze(artifactInput)
|
|
self.assertEqual(results["summary"], "malware_download")
|
|
mock.assert_called_once()
|