Add search_multiple endpoints (#1258)

Author: tomchop

core/web/apiv2/dfiq.py

@@ -53,6 +53,17 @@ class DFIQSearchRequest(BaseModel):
     page: int = 0
 
 
+class DFIQGetMultipleRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    names: list[str] = []
+    type: dfiq.DFIQType | None = None
+    sorting: list[tuple[str, bool]] = []
+    filter_aliases: list[tuple[str, str]] = []
+    count: int = 50
+    page: int = 0
+
+
 class DFIQSearchResponse(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
@@ -420,3 +431,22 @@ def search(httpreq: Request, request: DFIQSearchRequest) -> DFIQSearchResponse:
         user=httpreq.state.user,
     )
     return DFIQSearchResponse(dfiq=dfiq_objects, total=total)
+
+
+@router.post("/get/multiple")
+def get_multiple(
+    httpreq: Request, request: DFIQGetMultipleRequest
+) -> DFIQSearchResponse:
+    """Gets multiple DFIQ objects by name."""
+    query = {"name__in": request.names}
+    if request.type:
+        query["type"] = request.type
+    dfiq_objects, total = dfiq.DFIQBase.filter(
+        query_args=query,
+        offset=request.page * request.count,
+        count=request.count,
+        sorting=request.sorting,
+        aliases=request.filter_aliases,
+        user=httpreq.state.user,
+    )
+    return DFIQSearchResponse(dfiq=dfiq_objects, total=total)

core/web/apiv2/entities.py

@@ -33,6 +33,17 @@ class EntitySearchRequest(BaseModel):
     page: int = 0
 
 
+class EntityMultipleGetRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    names: list[str] = []
+    type: EntityType | None = None
+    sorting: list[tuple[str, bool]] = []
+    filter_aliases: list[tuple[str, str]] = []
+    count: int = 50
+    page: int = 0
+
+
 class EntitySearchResponse(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
@@ -191,6 +202,26 @@ def search(httpreq: Request, request: EntitySearchRequest) -> EntitySearchRespon
     return response
 
 
+@router.post("/get/multiple")
+def get_multiple(
+    httpreq: Request, request: EntityMultipleGetRequest
+) -> EntitySearchResponse:
+    """Gets multiple entities by name."""
+    query = {"name__in": request.names}
+    if request.type:
+        query["type"] = request.type
+    entities, total = Entity.filter(
+        query_args=query,
+        offset=request.page * request.count,
+        count=request.count,
+        sorting=request.sorting,
+        aliases=request.filter_aliases,
+        links_count=True,
+        user=httpreq.state.user,
+    )
+    return EntitySearchResponse(entities=entities, total=total)
+
+
 @router.post("/tag")
 @rbac.permission_on_ids(roles.Permission.WRITE)
 def tag(httpreq: Request, request: EntityTagRequest) -> EntityTagResponse:

core/web/apiv2/indicators.py

@@ -43,6 +43,17 @@ class IndicatorSearchRequest(BaseModel):
     page: int = 0
 
 
+class IndicatorMultipleGetRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    names: list[str] = []
+    type: IndicatorType | None = None
+    sorting: list[tuple[str, bool]] = []
+    filter_aliases: list[tuple[str, str]] = []
+    count: int = 50
+    page: int = 0
+
+
 class IndicatorSearchResponse(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
@@ -237,6 +248,25 @@ def search(
     return IndicatorSearchResponse(indicators=indicators, total=total)
 
 
+@router.post("/get/multiple")
+def get_multiple(
+    httpreq: Request, request: IndicatorMultipleGetRequest
+) -> IndicatorSearchResponse:
+    """Gets multiple indicators by name."""
+    query = {"name__in": request.names}
+    if request.type:
+        query["type"] = request.type
+    indicators, total = Indicator.filter(
+        query_args=query,
+        offset=request.page * request.count,
+        count=request.count,
+        sorting=request.sorting,
+        aliases=request.filter_aliases,
+        user=httpreq.state.user,
+    )
+    return IndicatorSearchResponse(indicators=indicators, total=total)
+
+
 @router.post("/tag")
 @rbac.permission_on_ids(roles.Permission.WRITE)
 def tag(httpreq: Request, request: IndicatorTagRequest) -> IndicatorTagResponse:

core/web/apiv2/tag.py

@@ -33,6 +33,14 @@ class TagSearchRequest(BaseModel):
     page: int
 
 
+class TagMultipleGetRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    names: list[str] = []
+    count: int
+    page: int
+
+
 class TagSearchResponse(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
@@ -109,6 +117,16 @@ def search(request: TagSearchRequest) -> TagSearchResponse:
     return TagSearchResponse(tags=tags, total=total)
 
 
+@router.post("/get/multiple")
+def get_multiple(request: TagMultipleGetRequest) -> TagSearchResponse:
+    """Gets multiple Tags by name."""
+    request_args = {"name__in": request.names}
+    count = request.count
+    page = request.page
+    tags, total = Tag.filter(request_args, offset=page * count, count=count)
+    return TagSearchResponse(tags=tags, total=total)
+
+
 @router.delete("/{tag_id}")
 def delete(tag_id: str) -> None:
     """Deletes a Tag."""

tests/apiv2/dfiq.py

@@ -678,3 +678,27 @@ def test_to_archive(self):
                 self.assertIn("semi_private_question", content)
                 self.assertIn("public_approach", content)
                 self.assertIn("internal_approach", content)
+
+    def test_get_multiple(self):
+        with open("tests/dfiq_test_data/S1003.yaml", "r") as f:
+            yaml_string = f.read()
+
+        response = client.post(
+            "/api/v2/dfiq/from_yaml",
+            json={
+                "dfiq_yaml": yaml_string,
+                "dfiq_type": dfiq.DFIQType.scenario,
+            },
+        )
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+
+        response = client.post(
+            "/api/v2/dfiq/get/multiple",
+            json={"names": ["scenario1"], "page": 0, "count": 10},
+        )
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        self.assertEqual(len(data["dfiq"]), 1)
+        self.assertEqual(data["total"], 1)
+        self.assertEqual(data["dfiq"][0]["name"], "scenario1")

tests/apiv2/entities.py

@@ -110,6 +110,17 @@ def test_search_entities(self):
         self.assertEqual(len(data["entities"][0]["tags"]), 1, data)
         self.assertIn("ta1", [tag["name"] for tag in data["entities"][0]["tags"]])
 
+    def test_get_multiple_entities(self):
+        response = client.post(
+            "/api/v2/entities/get/multiple",
+            json={"names": ["ta1", "bears", "foo"], "page": 0, "count": 10},
+        )
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        self.assertEqual(len(data["entities"]), 2)
+        names = [e["name"] for e in data["entities"]]
+        self.assertCountEqual(names, ["ta1", "bears"])
+
     def test_search_entities_tagged(self):
         response = client.post(
             "/api/v2/entities/search",

tests/apiv2/indicators.py

@@ -122,6 +122,17 @@ def test_search_indicators(self):
         self.assertEqual(len(data["indicators"][0]["tags"]), 1)
         self.assertIn("hextag", [tag["name"] for tag in data["indicators"][0]["tags"]])
 
+    def test_get_multiple_indicators(self):
+        response = client.post(
+            "/api/v2/indicators/get/multiple",
+            json={"names": ["hex", "localhost", "foo"], "page": 0, "count": 10},
+        )
+        self.assertEqual(response.status_code, 200)
+        data = response.json()
+        self.assertEqual(len(data["indicators"]), 2)
+        names = [indicator["name"] for indicator in data["indicators"]]
+        self.assertCountEqual(names, ["hex", "localhost"])
+
     def test_search_indicators_tagged(self):
         response = client.post(
             "/api/v2/indicators/search",

tests/apiv2/tags.py

@@ -73,6 +73,22 @@ def test_tag_search(self):
         data = response.json()
         self.assertEqual(len(data), 2)
 
+    def test_tag_get_multiple(self):
+        response = client.post("/api/v2/tags/", json={"name": "tag2-test"})
+        response = client.post("/api/v2/tags/", json={"name": "tag3-test"})
+        self.assertEqual(response.status_code, 200)
+
+        response = client.post(
+            "/api/v2/tags/get/multiple",
+            json={"names": ["tag1", "tag2-test"], "page": 0, "count": 10},
+        )
+        self.assertEqual(response.status_code, 200)
+        data = response.json()
+        self.assertEqual(len(data["tags"]), 2)
+        self.assertEqual(data["total"], 2)
+        self.assertEqual(data["tags"][0]["name"], "tag1")
+        self.assertEqual(data["tags"][1]["name"], "tag2-test")
+
     def test_tag_delete(self):
         response = client.delete(f"/api/v2/tags/{self.tag.id}")
         self.assertEqual(response.status_code, 200)