Merge pull request #9 from yeti-platform/overlays

Add function to get Yara bundles with overlays

Author: tomchop

README.md

@@ -1,2 +1,5 @@
 # yeti-python
+
 Python client for the Yeti v2 API
+
+[![Unit tests](https://github.com/yeti-platform/yeti-python/actions/workflows/unittests.yml/badge.svg)](https://github.com/yeti-platform/yeti-python/actions/workflows/unittests.yml)

tests/api.py

@@ -295,6 +295,30 @@ def test_error_message(self, mock_post):
         self.assertEqual(str(raised.exception), "error_message")
         self.assertEqual(raised.exception.status_code, 400)
 
+    @patch("yeti.api.requests.Session.post")
+    def test_get_yara_bundle_with_overlays(self, mock_post):
+        # Mock the YARA bundle response
+        mock_response = MagicMock()
+        mock_response.content = b'{"bundle": "bundlestring"}'
+        mock_post.return_value = mock_response
+
+        # Call the method with overlays
+        result = self.api.get_yara_bundle_with_overlays(
+            overlays=["overlay1", "overlay2"]
+        )
+
+        # Check the result
+        self.assertEqual(result, {"bundle": "bundlestring"})
+        mock_post.assert_called_with(
+            "http://fake-url/api/v2/indicators/yara/bundle",
+            json={
+                "ids": [],
+                "tags": [],
+                "exclude_tags": [],
+                "overlays": ["overlay1", "overlay2"],
+            },
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

yeti/api.py

@@ -274,6 +274,45 @@ def patch_indicator(
         )
         return json.loads(response)
 
+    def get_yara_bundle_with_overlays(
+        self,
+        ids: list[str] | None = None,
+        tags: list[str] | None = None,
+        exclude_tags: list[str] | None = None,
+        overlays: list[str] | None = None,
+    ) -> str:
+        """Gets a Yara bundle with overlays.
+
+        Args:
+            ids: The list of IDs to include in the bundle.
+            tags: Include Yara rules with this tag in the bundle.
+            exclude_tags: Remove Yara rules with this tag from the bundle.
+            overlays: The list of overlays to include in the bundle.
+        """
+        if ids is None:
+            ids = []
+        if tags is None:
+            tags = []
+        if exclude_tags is None:
+            exclude_tags = []
+        if overlays is None:
+            overlays = []
+
+        params = {
+            "ids": ids,
+            "tags": tags,
+            "exclude_tags": exclude_tags,
+            "overlays": overlays,
+        }
+
+        result = self.do_request(
+            "POST",
+            f"{self._url_root}/api/v2/indicators/yara/bundle",
+            json_data=params,
+        )
+
+        return json.loads(result)
+
     def search_dfiq(self, name: str, dfiq_type: str | None = None) -> list[YetiObject]:
         """Searches for a DFIQ in Yeti.