diff --git a/pyproject.toml b/pyproject.toml index 9e51daa..7061792 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "yeti-python" -version = "2.0.8" +version = "2.1.0" description = "Python bindings for the Yeti API" authors = ["tomchop"] license = "Apache" diff --git a/tests/api.py b/tests/api.py index bd8835d..d28e199 100644 --- a/tests/api.py +++ b/tests/api.py @@ -66,13 +66,17 @@ def test_search_entities(self, mock_post): mock_post.return_value = mock_response result = self.api.search_entities( - name="test_entity", description="test_description" + name="test_entity", description="test_description", tags=["tag1"] ) self.assertEqual(result, [{"name": "test_entity"}]) mock_post.assert_called_with( "http://fake-url/api/v2/entities/search", json={ - "query": {"name": "test_entity", "description": "test_description"}, + "query": { + "name": "test_entity", + "description": "test_description", + "tags": ["tag1"], + }, "count": 100, "page": 0, }, diff --git a/tests/e2e.py b/tests/e2e.py index 6cd5eb6..45c66e3 100644 --- a/tests/e2e.py +++ b/tests/e2e.py @@ -60,6 +60,23 @@ def test_search_entities(self): self.assertEqual(result[0]["name"], "testSearch") self.assertEqual(result[0]["tags"][0]["name"], "testtag") + def test_search_entities_with_tags(self): + self.api.auth_api_key(os.getenv("YETI_API_KEY")) + self.api.new_entity( + { + "name": "testSearchWithTags", + "type": "malware", + "description": "test", + }, + tags=["testtag1", "testtag2"], + ) + time.sleep(5) + result = self.api.search_entities( + name="testSear", description="tes", tags=["testtag1"] + ) + self.assertEqual(len(result), 1, result) + self.assertEqual(result[0]["name"], "testSearchWithTags") + def test_get_multiple_entities(self): self.api.auth_api_key(os.getenv("YETI_API_KEY")) self.api.new_entity( diff --git a/yeti/api.py b/yeti/api.py index 389a053..abf75f8 100644 --- a/yeti/api.py +++ b/yeti/api.py @@ -224,6 +224,8 @@ def search_indicators( pattern: The pattern of the indicator to search for. description: The description of the indicator to search for. (substring match) tags: The tags of the indicator to search for. + count: The number of results to return (default is 100, which means all). + page: The page of results to return (default is 0, which means the first page). Returns: The response from the API; a list of dicts representing indicators. @@ -299,6 +301,7 @@ def search_entities( name: str | None = None, entity_type: str | None = None, description: str | None = None, + tags: list[str] | None = None, count: int = 100, page: int = 0, ) -> list[YetiObject]: @@ -310,6 +313,7 @@ def search_entities( name: The name of the entity to search for (substring match). entity_type: The type of the entity to search for. description: The description of the entity to search for. (substring match) + tags: The tags of the entity to search for. count: The number of results to return (default is 100, which means all). page: The page of results to return (default is 0, which means the first page). @@ -326,6 +330,8 @@ def search_entities( query["type"] = entity_type if description: query["description"] = description + if tags: + query["tags"] = tags params = {"query": query, "count": count, "page": page} response = self.do_request(