Merge pull request #469 from TheHive-Project/444-review-procedure-endpoints

#444 - Review procedure endpoints

Author: Kamforka

tests/test_procedure_endpoint.py

@@ -1,10 +1,16 @@
 import pytest
+
 from thehive4py.client import TheHiveApi
 from thehive4py.errors import TheHiveError
 from thehive4py.helpers import now_to_ts
+from thehive4py.query.filters import In
 from thehive4py.types.alert import OutputAlert
 from thehive4py.types.case import OutputCase
-from thehive4py.types.procedure import InputUpdateProcedure, OutputProcedure
+from thehive4py.types.procedure import (
+    InputProcedure,
+    InputUpdateProcedure,
+    OutputProcedure,
+)
 
 
 class TestProcedureEndpoint:
@@ -24,6 +30,42 @@ def test_create_in_alert_and_get(
         fetched_procedure = thehive.procedure.get(procedure_id=created_procedure["_id"])
         assert created_procedure == fetched_procedure
 
+    def test_bulk_create_in_alert_and_find(
+        self, thehive: TheHiveApi, test_alert: OutputAlert
+    ):
+
+        procedures: list[InputProcedure] = [
+            {
+                "occurDate": now_to_ts(),
+                "patternId": "T1059.006",
+                "tactic": "execution",
+                "description": "...",
+            },
+            {
+                "occurDate": now_to_ts(),
+                "patternId": "T1059.007",
+                "tactic": "execution",
+                "description": "...",
+            },
+        ]
+        created_procedures = thehive.procedure.bulk_create_in_alert(
+            alert_id=test_alert["_id"], procedures=procedures
+        )
+
+        fetched_procedures = thehive.procedure.find(
+            filters=In(
+                field="_id",
+                values=[procedure["_id"] for procedure in created_procedures],
+            )
+        )
+        assert sorted(
+            created_procedures,
+            key=lambda procedure: procedure["_id"],
+        ) == sorted(
+            fetched_procedures,
+            key=lambda procedure: procedure["_id"],
+        )
+
     def test_create_in_case_and_get(self, thehive: TheHiveApi, test_case: OutputCase):
         created_procedure = thehive.procedure.create_in_case(
             case_id=test_case["_id"],
@@ -38,6 +80,40 @@ def test_create_in_case_and_get(self, thehive: TheHiveApi, test_case: OutputCase
         fetched_procedure = thehive.procedure.get(procedure_id=created_procedure["_id"])
         assert created_procedure == fetched_procedure
 
+    def test_bulk_create_in_case_and_find(
+        self, thehive: TheHiveApi, test_case: OutputCase
+    ):
+        procedures: list[InputProcedure] = [
+            {
+                "occurDate": now_to_ts(),
+                "patternId": "T1059.006",
+                "tactic": "execution",
+                "description": "...",
+            },
+            {
+                "occurDate": now_to_ts(),
+                "patternId": "T1059.007",
+                "tactic": "execution",
+                "description": "...",
+            },
+        ]
+        created_procedures = thehive.procedure.bulk_create_in_case(
+            case_id=test_case["_id"], procedures=procedures
+        )
+        fetched_procedures = thehive.procedure.find(
+            filters=In(
+                field="_id",
+                values=[procedure["_id"] for procedure in created_procedures],
+            )
+        )
+        assert sorted(
+            created_procedures,
+            key=lambda procedure: procedure["_id"],
+        ) == sorted(
+            fetched_procedures,
+            key=lambda procedure: procedure["_id"],
+        )
+
     def test_delete(self, thehive: TheHiveApi, test_procedure: OutputProcedure):
         procedure_id = test_procedure["_id"]
         thehive.procedure.delete(procedure_id=procedure_id)
@@ -57,3 +133,9 @@ def test_update(self, thehive: TheHiveApi, test_procedure: OutputProcedure):
 
         for key, value in update_fields.items():
             assert updated_procedure.get(key) == value
+
+    def test_bulk_delete(self, thehive: TheHiveApi, test_procedure: OutputProcedure):
+        procedure_id = test_procedure["_id"]
+        thehive.procedure.bulk_delete(procedure_ids=[procedure_id])
+        with pytest.raises(TheHiveError):
+            thehive.procedure.get(procedure_id=procedure_id)

thehive4py/endpoints/procedure.py

@@ -14,48 +14,132 @@
 
 
 class ProcedureEndpoint(EndpointBase):
-    def create_in_alert(
-        self, alert_id: str, procedure: InputProcedure
+    def create_in_case(
+        self, case_id: str, procedure: InputProcedure
     ) -> OutputProcedure:
+        """Create a procedure in a case.
+
+        Args:
+            case_id: The id of the case.
+            procedure: The fields of the procedure to create.
+
+        Returns:
+            The created procedure.
+        """
         return self._session.make_request(
-            "POST", path=f"/api/v1/alert/{alert_id}/procedure", json=procedure
+            "POST", path=f"/api/v1/case/{case_id}/procedure", json=procedure
         )
 
-    def create_in_case(
-        self, case_id: str, procedure: InputProcedure
+    def bulk_create_in_case(
+        self, case_id: str, procedures: List[InputProcedure]
+    ) -> List[OutputProcedure]:
+        """Create several procedures in a case.
+
+        Args:
+            case_id: The id of the case.
+            procedures: The list of procedures to create.
+
+        Returns:
+            The list of created procedures.
+        """
+        return self._session.make_request(
+            "POST",
+            path=f"/api/v1/case/{case_id}/procedures",
+            json={"procedures": procedures},
+        )
+
+    def create_in_alert(
+        self, alert_id: str, procedure: InputProcedure
     ) -> OutputProcedure:
+        """Create a procedure in an alert.
+
+        Args:
+            alert_id: The id of the alert.
+            procedure: The fields of the procedure to create.
+
+        Returns:
+            The created procedure.
+        """
         return self._session.make_request(
-            "POST", path=f"/api/v1/case/{case_id}/procedure", json=procedure
+            "POST", path=f"/api/v1/alert/{alert_id}/procedure", json=procedure
         )
 
-    def get(self, procedure_id: str) -> OutputProcedure:
-        # TODO: temp implementation until a dedicated get endpoint
-        procedures = self._session.make_request(
+    def bulk_create_in_alert(
+        self, alert_id: str, procedures: List[InputProcedure]
+    ) -> List[OutputProcedure]:
+        """Create multiple procedures in an alert.
+
+        Args:
+            alert_id: The id of the alert.
+            procedures: The list of procedures to create.
+
+        Returns:
+            The list of created procedures.
+        """
+        return self._session.make_request(
             "POST",
-            path="/api/v1/query",
-            json={"query": [{"_name": "getProcedure", "idOrName": procedure_id}]},
+            path=f"/api/v1/alert/{alert_id}/procedures",
+            json={"procedures": procedures},
         )
-        try:
-            return procedures[0]
-        except IndexError:
-            raise TheHiveError("404 - Procedure not found")
 
     def delete(self, procedure_id: str) -> None:
+        """Delete a procedure.
+
+        Args:
+            procedure_id: The id of the procedure.
+
+        Returns:
+            N/A
+        """
         return self._session.make_request(
             "DELETE", path=f"/api/v1/procedure/{procedure_id}"
         )
 
     def update(self, procedure_id: str, fields: InputUpdateProcedure) -> None:
+        """Update a procedure.
+
+        Args:
+            procedure_id: The id of the procedure.
+            fields: The fields of the procedure to update.
+
+        Returns:
+            N/A
+        """
         return self._session.make_request(
             "PATCH", path=f"/api/v1/procedure/{procedure_id}", json=fields
         )
 
+    def bulk_delete(self, procedure_ids: List[str]) -> None:
+        """Delete multiple procedures.
+
+        Args:
+            procedure_ids: The list of procedure ids to delete.
+
+        Returns:
+            N/A
+        """
+        return self._session.make_request(
+            "POST",
+            path="/api/v1/procedure/delete/_bulk",
+            json={"ids": procedure_ids},
+        )
+
     def find(
         self,
         filters: Optional[FilterExpr] = None,
         sortby: Optional[SortExpr] = None,
         paginate: Optional[Paginate] = None,
     ) -> List[OutputProcedure]:
+        """Find multiple procedures.
+
+        Args:
+            filters: The filter expressions to apply in the query.
+            sortby: The sort expressions to apply in the query.
+            paginate: The pagination expression to apply in the query.
+
+        Returns:
+            The list of procedures matched by the query or an empty list.
+        """
         query: QueryExpr = [
             {"_name": "listProcedure"},
             *self._build_subquery(filters=filters, sortby=sortby, paginate=paginate),
@@ -67,3 +151,23 @@ def find(
             params={"name": "procedures"},
             json={"query": query},
         )
+
+    def get(self, procedure_id: str) -> OutputProcedure:
+        """Get a procedure by id.
+
+        Args:
+            procedure_id: The id of the procedure.
+
+        Returns:
+            The procedure specified by the id.
+        """
+        # TODO: temp implementation until a dedicated get endpoint
+        procedures = self._session.make_request(
+            "POST",
+            path="/api/v1/query",
+            json={"query": [{"_name": "getProcedure", "idOrName": procedure_id}]},
+        )
+        try:
+            return procedures[0]
+        except IndexError:
+            raise TheHiveError("404 - Procedure not found")

thehive4py/types/procedure.py

@@ -1,34 +1,32 @@
 from typing import TypedDict
 
+from typing_extensions import NotRequired
 
-class InputProcedureRequired(TypedDict):
+
+class InputProcedure(TypedDict):
     occurDate: int
     patternId: str
+    tactic: NotRequired[str]
+    description: NotRequired[str]
 
 
-class InputProcedure(InputProcedureRequired, total=False):
-    tactic: str
-    description: str
-
-
-class OutputProcedureRequired(TypedDict):
+class OutputProcedure(TypedDict):
     _id: str
     _createdAt: int
     _createdBy: str
+    _updatedAt: NotRequired[int]
+    _updatedBy: NotRequired[str]
+    description: NotRequired[str]
     occurDate: int
-    tactic: str
-    tacticLabel: str
+    patternId: NotRequired[str]
+    patternName: NotRequired[str]
+    tactic: NotRequired[str]
+    tacticLabel: NotRequired[str]
     extraData: dict
 
 
-class OutputProcedure(OutputProcedureRequired, total=False):
-    _updatedAt: int
-    _updatedBy: str
-    description: str
-    patternId: str
-    patternName: str
-
-
 class InputUpdateProcedure(TypedDict, total=False):
     description: str
     occurDate: int
+    patternId: str
+    tactic: str