fix: handle sigma errors and default to offline tests

Author: Neo23x0

.github/workflows/integration-tests.yml

@@ -26,7 +26,7 @@ jobs:
 
       - name: Run live integration smoke tests
         run: |
-          python -m pytest -v \
+          python -m pytest -m integration -v \
             tests/test_basic.py::test_status \
             tests/test_basic.py::test_demo_rules_json \
             tests/test_cli_and_sigma.py::test_sigma_zip_updates_retrieved_rule_count_live_demo

Makefile

@@ -1,4 +1,4 @@
-.PHONY: help prepare-dev test lint release release-test
+.PHONY: help prepare-dev test test-integration lint release release-test
 
 VENV_NAME?=venv
 VENV_ACTIVATE=. $(VENV_NAME)/bin/activate
@@ -8,7 +8,9 @@ help:
 	@echo "make prepare-dev"
 	@echo "       prepare development environment, use only once"
 	@echo "make test"
-	@echo "       run tests"
+	@echo "       run offline tests"
+	@echo "make test-integration"
+	@echo "       run live integration tests"
 	@echo "make lint"
 	@echo "       run pylint and mypy"
 
@@ -28,6 +30,9 @@ $(VENV_NAME)/bin/activate: setup.py
 test: venv
 	${PYTHON} -m pytest -v
 
+test-integration: venv
+	${PYTHON} -m pytest -m integration -v
+
 release:
 	rm -rf ./dist/*
 	${PYTHON} -m pip install --upgrade setuptools wheel
@@ -50,4 +55,3 @@ lint: venv
 
 
 
-

pytest.ini

@@ -1,3 +1,4 @@
 [pytest]
+addopts = -m "not integration"
 markers =
     integration: tests that call the live Valhalla service

tests/test_cli_and_sigma.py

@@ -7,7 +7,7 @@
 
 import valhallaAPI.valhalla as valhalla_module
 import valhallaAPI.valhalla_cli as valhalla_cli
-from valhallaAPI.valhalla import ValhallaAPI
+from valhallaAPI.valhalla import ValhallaAPI, ApiError
 from valhallaAPI.version import __version__
 
 DEMO_KEY = ValhallaAPI.DEMO_KEY
@@ -75,6 +75,44 @@ def fake_post(url, data=None, proxies=None, headers=None):
     ]
 
 
+def test_sigma_json_error_response_does_not_keyerror(monkeypatch):
+    def fake_post(url, data=None, proxies=None, headers=None):
+        assert url.endswith("/getsigma")
+        return MockResponse(
+            {
+                "status": "error",
+                "message": "demo failure",
+            }
+        )
+
+    monkeypatch.setattr(valhalla_module.requests, "post", fake_post)
+    v = ValhallaAPI(api_key=DEMO_KEY)
+
+    response = v.get_sigma_rules_json(search="suspicious", private_only=True)
+
+    assert response["status"] == "error"
+    assert v.last_retrieved_rules_count == 0
+
+
+def test_sigma_zip_raises_api_error_on_error_response(monkeypatch):
+    def fake_post(url, data=None, proxies=None, headers=None):
+        assert url.endswith("/getsigma")
+        return MockResponse(
+            {
+                "status": "error",
+                "message": "demo failure",
+            }
+        )
+
+    monkeypatch.setattr(valhalla_module.requests, "post", fake_post)
+    v = ValhallaAPI(api_key=DEMO_KEY)
+
+    with pytest.raises(ApiError) as exc:
+        v.get_sigma_rules_zip(search="suspicious", private_only=True)
+
+    assert exc.value.message == "demo failure"
+
+
 @pytest.mark.integration
 def test_sigma_zip_updates_retrieved_rule_count_live_demo():
     v = ValhallaAPI(api_key=DEMO_KEY)

valhallaAPI/valhalla.py

@@ -313,12 +313,11 @@ def get_sigma_rules_json(self, search="", private_only=False):
         # Load JSON
         rules_response = json.loads(r.text)
 
-        if search:
-            rules_response['rules'] = filter_search(rules_response['rules'], query=search)
-        if private_only:
-            rules_response['rules'] = filter_privateonly(rules_response['rules'])
-
         if 'rules' in rules_response:
+            if search:
+                rules_response['rules'] = filter_search(rules_response['rules'], query=search)
+            if private_only:
+                rules_response['rules'] = filter_privateonly(rules_response['rules'])
             self.last_retrieved_rules_count = len(rules_response['rules'])
         else:
             self.last_retrieved_rules_count = 0
@@ -332,6 +331,12 @@ def get_sigma_rules_zip(self, search="", private_only=False):
         :return:
         """
         rules_response = self.get_sigma_rules_json(search, private_only)
+
+        if 'status' in rules_response:
+            if rules_response['status'] == "error":
+                raise ApiError(rules_response['message'])
+        if 'rules' not in rules_response:
+            raise ApiError("Unexpected response from Valhalla API")
 
         zip_buffer = io.BytesIO()
         with zipfile.ZipFile(file=zip_buffer, mode='w') as zip_file: