securityonion/salt/sensoroni/files/analyzers/malwarehashregistry/malwarehashregistry_test.py

from io import StringIO
import sys
from unittest.mock import patch, MagicMock
from malwarehashregistry import malwarehashregistry
import unittest


class TestMalwareHashRegistryMethods(unittest.TestCase):

    def test_main_missing_input(self):
        with patch('sys.exit', new=MagicMock()) as sysmock:
            with patch('sys.stderr', new=StringIO()) as mock_stderr:
                sys.argv = ["cmd"]
                malwarehashregistry.main()
                self.assertEqual(mock_stderr.getvalue(), "usage: cmd [-h] artifact\ncmd: error: the following arguments are required: artifact\n")
                sysmock.assert_called_once_with(2)

    def test_main_success(self):
        output = {"foo": "bar"}
        with patch('sys.stdout', new=StringIO()) as mock_stdout:
            with patch('malwarehashregistry.malwarehashregistry.analyze', new=MagicMock(return_value=output)) as mock:
                sys.argv = ["cmd", "input"]
                malwarehashregistry.main()
                expected = '{"foo": "bar"}\n'
                self.assertEqual(mock_stdout.getvalue(), expected)
                mock.assert_called_once()

    def test_sendReq(self):
        output = "84af04b8e69682782607a0c5796ca56999eda6b3 1563161433 35"
        hash = "abcd1234"
        server = "hash.cymru.com"
        flags = 0
        options = {"whoishost": server}
        with patch('whois.NICClient.whois_lookup', new=MagicMock(return_value=output)) as mock:
            response = malwarehashregistry.sendReq(hash)
            mock.assert_called_once_with(options, hash, flags)
            self.assertIsNotNone(response)
            self.assertEqual(response["hash"], "84af04b8e69682782607a0c5796ca56999eda6b3")
            self.assertRegex(response["last_seen"], r'2019-..-07 ..:..:33')  # host running this test won't always use UTC
            self.assertEqual(response["av_detection_percentage"], 35)

    def test_sendReqNoData(self):
        output = "84af04b8e69682782607a0c5796ca5696b3 NO_DATA"
        hash = "abcd1234"
        server = "hash.cymru.com"
        flags = 0
        options = {"whoishost": server}
        with patch('whois.NICClient.whois_lookup', new=MagicMock(return_value=output)) as mock:
            response = malwarehashregistry.sendReq(hash)
            mock.assert_called_once_with(options, hash, flags)
            self.assertIsNotNone(response)
            self.assertEqual(response, {"hash": "84af04b8e69682782607a0c5796ca5696b3", "last_seen": "NO_DATA", "av_detection_percentage": 0})

    def test_prepareResults_none(self):
        raw = {"hash": "14af04b8e69682782607a0c5796ca56999eda6b3", "last_seen": "NO_DATA", "av_detection_percentage": 0}
        results = malwarehashregistry.prepareResults(raw)
        self.assertEqual(results["response"], raw)
        self.assertEqual(results["summary"], "no_results")
        self.assertEqual(results["status"], "ok")

    def test_prepareResults_harmless(self):
        raw = {"hash": "14af04b8e69682782607a0c5796ca56999eda6b3", "last_seen": "123456", "av_detection_percentage": 0}
        results = malwarehashregistry.prepareResults(raw)
        self.assertEqual(results["response"], raw)
        self.assertEqual(results["summary"], "harmless")
        self.assertEqual(results["status"], "ok")

    def test_prepareResults_sus(self):
        raw = {"hash": "14af04b8e69682782607a0c5796ca56999eda6b3", "last_seen": "123456", "av_detection_percentage": 1}
        results = malwarehashregistry.prepareResults(raw)
        self.assertEqual(results["response"], raw)
        self.assertEqual(results["summary"], "suspicious")
        self.assertEqual(results["status"], "caution")

    def test_prepareResults_mal(self):
        raw = {"hash": "14af04b8e69682782607a0c5796ca56999eda6b3", "last_seen": "123456", "av_detection_percentage": 51}
        results = malwarehashregistry.prepareResults(raw)
        self.assertEqual(results["response"], raw)
        self.assertEqual(results["summary"], "malicious")
        self.assertEqual(results["status"], "threat")

    def test_prepareResults_error(self):
        raw = {}
        results = malwarehashregistry.prepareResults(raw)
        self.assertEqual(results["response"], raw)
        self.assertEqual(results["summary"], "internal_failure")
        self.assertEqual(results["status"], "caution")

    def test_analyze(self):
        output = {"hash": "14af04b8e69682782607a0c5796ca56999eda6b3", "last_seen": "NO_DATA", "av_detection_percentage": 0}
        artifactInput = '{"value": "14af04b8e69682782607a0c5796ca56999eda6b3", "artifactType": "hash"}'
        with patch('malwarehashregistry.malwarehashregistry.sendReq', new=MagicMock(return_value=output)) as mock:
            results = malwarehashregistry.analyze(artifactInput)
            self.assertEqual(results["summary"], "no_results")
            mock.assert_called_once()