Merge pull request #18 from yeti-platform/entitysearch_tag

Add support for tags in search_entities and update tests

Author: tomchop

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"

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,
             },

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(

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(