Merge pull request #451 from TheHive-Project/438-review-task-log-endpoints

#438 - Review task log endpoints

Author: Kamforka

tests/test_task_log_endpoint.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import pytest
+
 from thehive4py.client import TheHiveApi
 from thehive4py.errors import TheHiveError
 from thehive4py.helpers import now_to_ts
@@ -23,6 +24,29 @@ def test_create_and_get(self, thehive: TheHiveApi, test_task: OutputTask):
 
         assert created_log == fetched_log
 
+    def test_create_with_attachment(
+        self, thehive: TheHiveApi, test_task: OutputTask, tmp_path: Path
+    ):
+
+        attachment_paths = [str(tmp_path / f"attachment-{i}.txt") for i in range(2)]
+
+        for path in attachment_paths:
+            with open(path, "w") as attachment_fp:
+                attachment_fp.write(f"content of {path}")
+
+        log_with_attachment = thehive.task_log.create(
+            task_id=test_task["_id"],
+            task_log={
+                "message": "My test log",
+                "includeInTimeline": now_to_ts(),
+                "startDate": now_to_ts(),
+                "attachments": attachment_paths,
+            },
+        )
+
+        assert "attachments" in log_with_attachment
+        assert len(log_with_attachment["attachments"]) == len(attachment_paths)
+
     def test_delete(self, thehive: TheHiveApi, test_task_log: OutputTaskLog):
         thehive.task_log.delete(task_log_id=test_task_log["_id"])
 
@@ -45,7 +69,7 @@ def test_add_and_delete_attachment(
         with open(attachment_path, "w") as attachment_fp:
             attachment_fp.write("some content...")
 
-        thehive.task_log.add_attachments(
+        thehive.task_log.add_attachment(
             task_log_id=test_task_log["_id"], attachment_paths=[attachment_path]
         )
 
@@ -63,3 +87,30 @@ def test_add_and_delete_attachment(
         )
 
         assert attachments == []
+
+    def test_add_and_delete_attachments(
+        self, thehive: TheHiveApi, test_task_log: OutputTaskLog, tmp_path: Path
+    ):
+        attachment_path = str(tmp_path / "my-attachment.txt")
+        with open(attachment_path, "w") as attachment_fp:
+            attachment_fp.write("some content...")
+
+        with pytest.deprecated_call():
+            thehive.task_log.add_attachments(
+                task_log_id=test_task_log["_id"], attachment_paths=[attachment_path]
+            )
+
+        attachments = thehive.task_log.get(task_log_id=test_task_log["_id"]).get(
+            "attachments", []
+        )
+
+        for attachment in attachments:
+            thehive.task_log.delete_attachment(
+                task_log_id=test_task_log["_id"], attachment_id=attachment["_id"]
+            )
+
+        attachments = thehive.task_log.get(task_log_id=test_task_log["_id"]).get(
+            "attachments", []
+        )
+
+        assert attachments == []

thehive4py/endpoints/task_log.py

@@ -1,3 +1,5 @@
+import json as jsonlib
+import warnings
 from typing import List
 
 from thehive4py.endpoints._base import EndpointBase
@@ -6,33 +8,96 @@
 
 
 class TaskLogEndpoint(EndpointBase):
-    def create(self, task_id: str, task_log: InputTaskLog) -> OutputTaskLog:
+    def create(
+        self,
+        task_id: str,
+        task_log: InputTaskLog,
+    ) -> OutputTaskLog:
+        """Create a task log.
+
+        Args:
+            task_id: The id of the task to create the log in.
+            task_log: The body of the task log.
+
+        Returns:
+            The created task_log.
+        """
+        if "attachments" in task_log:
+            files: List[tuple] = [
+                ("attachments", self._fileinfo_from_filepath(attachment_path))
+                for attachment_path in task_log["attachments"]
+            ]
+            files.append(("_json", jsonlib.dumps(task_log)))
+            kwargs: dict = {"files": files}
+        else:
+            kwargs = {"json": task_log}
         return self._session.make_request(
-            "POST", path=f"/api/v1/task/{task_id}/log", json=task_log
+            "POST", path=f"/api/v1/task/{task_id}/log", **kwargs
         )
 
-    def get(self, task_log_id: str) -> OutputTaskLog:
-        # TODO: temp implementation until a dedicated get endpoint
-        logs = self._session.make_request(
-            "POST",
-            path="/api/v1/query",
-            json={"query": [{"_name": "getLog", "idOrName": task_log_id}]},
-        )
+    def delete(self, task_log_id: str) -> None:
+        """Delete a task log.
 
-        try:
-            return logs[0]
-        except IndexError:
-            raise TheHiveError("404 - Task Log Not Found")
+        Args:
+            task_log_id: The id of the task log.
 
-    def delete(self, task_log_id: str) -> None:
+        Returns:
+            N/A
+        """
         return self._session.make_request("DELETE", path=f"/api/v1/log/{task_log_id}")
 
     def update(self, task_log_id: str, fields: InputUpdateTaskLog) -> None:
+        """Update a task log.
+
+        Args:
+            task_log_id: The id of the task log.
+            fields: The fields of the task log to update.
+
+        Returns:
+            N/A
+        """
         return self._session.make_request(
             "PATCH", path=f"/api/v1/log/{task_log_id}", json=fields
         )
 
+    def add_attachment(self, task_log_id: str, attachment_paths: List[str]) -> None:
+        """Add attachments to a task log.
+
+        Args:
+            task_log_id: The id of the task log.
+            attachment_paths: List of paths to the attachments to create.
+
+        Returns:
+            The created task log attachments.
+        """
+        files = [
+            ("attachments", self._fileinfo_from_filepath(attachment_path))
+            for attachment_path in attachment_paths
+        ]
+        return self._session.make_request(
+            "POST", f"/api/v1/log/{task_log_id}/attachments", files=files
+        )
+
     def add_attachments(self, task_log_id: str, attachment_paths: List[str]) -> None:
+        """Add attachments to a task log.
+
+        !!! warning
+            Deprecated: use [task_log.add_attachment]
+            [thehive4py.endpoints.task_log.TaskLogEndpoint.add_attachment]
+            instead
+
+        Args:
+            task_log_id: The id of the task log.
+            attachment_paths: List of paths to the attachments to create.
+
+        Returns:
+            The created task log attachments.
+        """
+        warnings.warn(
+            message=("Deprecated: use the task_log.add_attachment method instead"),
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
         files = [
             ("attachments", self._fileinfo_from_filepath(attachment_path))
             for attachment_path in attachment_paths
@@ -42,6 +107,36 @@ def add_attachments(self, task_log_id: str, attachment_paths: List[str]) -> None
         )
 
     def delete_attachment(self, task_log_id: str, attachment_id: str) -> None:
+        """Delete a task log attachment.
+
+        Args:
+            task_log_id: The id of the task log.
+            attachment_id: The id of the task log attachment.
+
+        Returns:
+            N/A
+        """
         return self._session.make_request(
             "DELETE", f"/api/v1/log/{task_log_id}/attachments/{attachment_id}"
         )
+
+    def get(self, task_log_id: str) -> OutputTaskLog:
+        """Get a task log by id.
+
+        Args:
+            task_log_id: The id of the task log.
+
+        Returns:
+            The task log specified by the id.
+        """
+        # TODO: temp implementation until a dedicated get endpoint [if ever ;)]
+        logs = self._session.make_request(
+            "POST",
+            path="/api/v1/query",
+            json={"query": [{"_name": "getLog", "idOrName": task_log_id}]},
+        )
+
+        try:
+            return logs[0]
+        except IndexError:
+            raise TheHiveError("404 - Task Log Not Found")

thehive4py/types/task_log.py

@@ -1,5 +1,7 @@
 from typing import List, TypedDict
 
+from thehive4py.types.attachment import OutputAttachment
+
 
 class InputTaskLogRequired(TypedDict):
     message: str
@@ -8,6 +10,7 @@ class InputTaskLogRequired(TypedDict):
 class InputTaskLog(InputTaskLogRequired, total=False):
     startDate: int
     includeInTimeline: int
+    attachments: List[str]
 
 
 class OutputTaskLogRequired(TypedDict):
@@ -24,7 +27,7 @@ class OutputTaskLogRequired(TypedDict):
 class OutputTaskLog(OutputTaskLogRequired, total=False):
     _updatedBy: str
     _updatedAt: int
-    attachments: List[dict]  # TODO: typehint
+    attachments: List[OutputAttachment]
     includeInTimeline: int