diff --git a/tests/api.py b/tests/api.py index 9202362..4eba3d9 100644 --- a/tests/api.py +++ b/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() diff --git a/yeti/api.py b/yeti/api.py index 3a63e47..c41a64b 100644 --- a/yeti/api.py +++ b/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: