Merge pull request #12 from yeti-platform/bloom

Add bloom search endpoint

Author: tomchop

tests/api.py

@@ -62,6 +62,19 @@ def test_search_observables(self, mock_post):
             json={"query": {"value": "test_value"}, "count": 0},
         )
 
+    @patch("yeti.api.requests.Session.post")
+    def test_search_bloom(self, mock_post):
+        mock_response = MagicMock()
+        mock_response.content = b'[{"value": "test.com", "hits": ["filter1"]}]'
+        mock_post.return_value = mock_response
+
+        result = self.api.search_bloom(["test.com"])
+        self.assertEqual(result, [{"value": "test.com", "hits": ["filter1"]}])
+        mock_post.assert_called_with(
+            "http://fake-url/api/v2/bloom/search",
+            json={"values": ["test.com"]},
+        )
+
     @patch("yeti.api.requests.Session.post")
     def test_new_entity(self, mock_post):
         mock_response = MagicMock()

yeti/api.py

@@ -212,6 +212,23 @@ def search_observables(self, value: str) -> list[YetiObject]:
         )
         return json.loads(response)["observables"]
 
+    def search_bloom(self, values: list[str]) -> list[dict[str, Any]]:
+        """Searches for a list of observable values in Yeti's bloom filters.
+
+        Args:
+          values: The list of observable values to search for.
+
+        Returns:
+          A list of dicts representing hits, e.g.
+
+            {"value": "example.com", hits:["filter1"]}
+        """
+        params = {"values": values}
+        response = self.do_request(
+            "POST", f"{self._url_root}/api/v2/bloom/search", json_data=params
+        )
+        return json.loads(response)
+
     def new_entity(
         self, entity: dict[str, Any], tags: list[str] | None = None
     ) -> YetiObject: