Fix formatting/whitespace

salt/sensoroni/files/analyzers/urlscan/urlscan_test.py

@@ -0,0 +1,121 @@
+from io import StringIO
+import sys
+from unittest.mock import patch, MagicMock, PropertyMock, call
+from urlscan import urlscan
+import unittest
+
+
+class TestUrlScanMethods(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"]
+                urlscan.main()
+                self.assertEqual(mock_stderr.getvalue(), "usage: cmd [-h] [-c CONFIG_FILE] 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('urlscan.urlscan.analyze', new=MagicMock(return_value=output)) as mock:
+                sys.argv = ["cmd", "input"]
+                urlscan.main()
+                expected = '{"foo": "bar"}\n'
+                self.assertEqual(mock_stdout.getvalue(), expected)
+                mock.assert_called_once()
+
+    def test_checkConfigRequirements_notEnabled(self):
+        conf = {"not_a_key": "abcd12345"}
+        with self.assertRaises(SystemExit) as cm:
+            urlscan.checkConfigRequirements(conf)
+        self.assertEqual(cm.exception.code, 126)
+
+    def test_checkConfigRequirements_noApikey(self):
+        conf = {"enabled": True, "not_a_key": "abcd12345"}
+        with self.assertRaises(SystemExit) as cm:
+            urlscan.checkConfigRequirements(conf)
+        self.assertEqual(cm.exception.code, 126)
+
+    def test_checkConfigRequirements_Exist(self):
+        conf = {"enabled": True, "api_key": "abcd12345"}
+        config_exists = urlscan.checkConfigRequirements(conf)
+        self.assertTrue(config_exists)
+
+    def test_buildReq(self):
+        conf = {'base_url': 'https://myurl/api/v1/', 'api_key': 'abcd12345', 'visibility': 'public'}
+        artifact_type = "url"
+        artifact_value = "https://abc.com"
+        result = urlscan.buildReq(conf, artifact_type, artifact_value)
+        self.assertEqual("https://myurl/api/v1/scan/", result[0])
+        self.assertEqual({'API-Key': 'abcd12345'}, result[1])
+
+    def test_sendReq(self):
+        with patch('requests.request', new=MagicMock(return_value=MagicMock())) as mock:
+            headers = {"API-Key": "abcd1234"}
+            data = {"url": "https://urlscan.io", "visibility": "public"}
+            response = urlscan.sendReq("https://myurl", headers=headers, data=data)
+            mock.assert_called_once_with("POST", url="https://myurl", headers={"API-Key": "abcd1234"}, data={"url": "https://urlscan.io", "visibility": "public"})
+            self.assertIsNotNone(response)
+
+    def test_getReport_noRetry(self):
+        output_report = MagicMock()
+        type(output_report).status_code = PropertyMock(return_value=404)
+        output_report_body = {"requests": "body"}
+        output_report.json.return_value = output_report_body
+        with patch('requests.request', new=MagicMock(return_value=output_report)) as mock:
+            result = urlscan.getReport({'timeout': 0}, "https://abc.com/report")
+            self.assertEqual(404, result.status_code)
+            mock.assert_called_once()
+
+    def test_getReport_withRetry(self):
+        output_report = MagicMock()
+        type(output_report).status_code = PropertyMock(return_value=404)
+        output_report_body = {"requests": "body"}
+        output_report.json.return_value = output_report_body
+        with patch('requests.request', new=MagicMock(return_value=output_report)) as mock:
+            result = urlscan.getReport({'timeout': 3}, "https://abc.com/report")
+            self.assertEqual(404, result.status_code)
+            mock.assert_has_calls([call('GET', 'https://abc.com/report'), call('GET', 'https://abc.com/report')])
+
+    def test_prepareResults_sus(self):
+        raw = {"requests": [{"request": {"requestId": "1"}}], "verdicts": {"overall": {"score": 50, "malicious": False, "hasVerdicts": False}}}
+        results = urlscan.prepareResults(raw)
+        self.assertEqual(results["response"], raw)
+        self.assertEqual(results["summary"], "suspicious")
+        self.assertEqual(results["status"], "caution")
+
+    def test_prepareResults_mal(self):
+        raw = {"requests": [{"request": {"requestId": "2"}}], "verdicts": {"overall": {"score": 100, "malicious": True, "hasVerdicts": False}}}
+        results = urlscan.prepareResults(raw)
+        self.assertEqual(results["response"], raw)
+        self.assertEqual(results["summary"], "malicious")
+        self.assertEqual(results["status"], "threat")
+
+    def test_prepareResults_info(self):
+        raw = {"requests": [{"request": {"requestId": "3"}}], "verdicts": {"overall": {"score": 0, "malicious": False, "hasVerdicts": False}}}
+        results = urlscan.prepareResults(raw)
+        self.assertEqual(results["response"], raw)
+        self.assertEqual(results["summary"], "Scan complete.")
+        self.assertEqual(results["status"], "info")
+
+    def test_prepareResults_error(self):
+        raw = {}
+        results = urlscan.prepareResults(raw)
+        self.assertEqual(results["response"], raw)
+        self.assertEqual(results["summary"], "internal_failure")
+        self.assertEqual(results["status"], "caution")
+
+    def test_analyze(self):
+        output_req = "https://myurl/report"
+        output_report = MagicMock()
+        output_report_body = {"requests": [{"request": {"requestId": "3"}}], "verdicts": {"overall": {"score": 0, "malicious": False, "hasVerdicts": False}}}
+        output_report.json.return_value = output_report_body
+        artifactInput = '{"value":"https://abc.com","artifactType":"url"}'
+        conf = {'enabled': True, 'base_url': 'https://myurl/api/v1/', 'api_key': 'abcd12345', 'visibility': 'public'}
+        with patch('urlscan.urlscan.sendReq', new=MagicMock(return_value=output_req)) as mock_req:
+            with patch('urlscan.urlscan.getReport', new=MagicMock(return_value=output_report)) as mock_report:
+                results = urlscan.analyze(conf, artifactInput)
+                self.assertEqual(results["summary"], "Scan complete.")
+                mock_req.assert_called_once()
+                mock_report.assert_called_once()