From 8d0ee2924c3513b3b2bb640d166a7400e8f93b83 Mon Sep 17 00:00:00 2001
From: Thomas Chopitea <tomchop@gmail.com>
Date: Fri, 11 Apr 2025 16:54:15 +0000
Subject: [PATCH 1/2] Add function to get Yara bundles with overlays

---
 tests/api.py | 24 ++++++++++++++++++++++++
 yeti/api.py  | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/tests/api.py b/tests/api.py
index f8912c7..9202362 100644
--- a/tests/api.py
+++ b/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()
diff --git a/yeti/api.py b/yeti/api.py
index aee183f..5625b24 100644
--- a/yeti/api.py
+++ b/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.
 

From 33c8332ed374808f264b6b95babe45322d483afd Mon Sep 17 00:00:00 2001
From: Thomas Chopitea <tomchop@gmail.com>
Date: Fri, 11 Apr 2025 16:55:51 +0000
Subject: [PATCH 2/2] Add fancy badge

---
 README.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/README.md b/README.md
index df9a6c5..cee9629 100644
--- a/README.md
+++ b/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)