From b2db32a2c7054fcf33bb6a243b9dfdc3475edde0 Mon Sep 17 00:00:00 2001 From: Wes Lambert Date: Thu, 21 Apr 2022 17:33:24 +0000 Subject: [PATCH] Add function/test for non-existent VT api_key --- salt/sensoroni/files/analyzers/virustotal/virustotal.py | 9 +++++++++ .../files/analyzers/virustotal/virustotal_test.py | 6 ++++++ 2 files changed, 15 insertions(+) diff --git a/salt/sensoroni/files/analyzers/virustotal/virustotal.py b/salt/sensoroni/files/analyzers/virustotal/virustotal.py index 040669079..f06c5d14f 100644 --- a/salt/sensoroni/files/analyzers/virustotal/virustotal.py +++ b/salt/sensoroni/files/analyzers/virustotal/virustotal.py @@ -3,6 +3,14 @@ import requests import argparse import helpers import os +import sys + + +def checkConfigRequirements(conf): + if "api_key" not in conf or len(conf['api_key']) == 0: + sys.exit(126) + else: + return True def buildHeaders(conf): @@ -58,6 +66,7 @@ def prepareResults(raw): def analyze(conf, input): + checkConfigRequirements(conf) meta = helpers.loadMetadata(__file__) data = helpers.parseArtifact(input) helpers.checkSupportedType(meta, data["artifactType"]) diff --git a/salt/sensoroni/files/analyzers/virustotal/virustotal_test.py b/salt/sensoroni/files/analyzers/virustotal/virustotal_test.py index 50bbe3cb6..103e137ce 100644 --- a/salt/sensoroni/files/analyzers/virustotal/virustotal_test.py +++ b/salt/sensoroni/files/analyzers/virustotal/virustotal_test.py @@ -25,6 +25,12 @@ class TestVirusTotalMethods(unittest.TestCase): self.assertEqual(mock_stdout.getvalue(), expected) mock.assert_called_once() + def checkConfigRequirements(self): + conf = {"not_a_key": "abcd12345"} + with self.assertRaises(SystemExit) as cm: + virustotal.checkConfigRequirements(conf) + self.assertEqual(cm.exception.code, 126) + def test_buildHeaders(self): result = virustotal.buildHeaders({"api_key": "xyz"}) self.assertEqual("xyz", result["x-apikey"])