Fix: Removal Logic Bug in Annotation Mixins (#3323)

* Fixing a bug in the set_status function that allowed for multiple entries in that list.
* Also fixed similar issues with the `remove_label` and `remove_comment` functions.
* Added a warning for now to surface instances where the problem still occurs.

---------

Co-authored-by: Janosch <99879757+jkppr@users.noreply.github.com>

Author: jaegeral

timesketch/models/annotations.py

@@ -18,6 +18,7 @@
 from __future__ import unicode_literals
 
 import json
+import logging
 import six
 
 from sqlalchemy import Column
@@ -32,6 +33,8 @@
 from timesketch.models import BaseModel
 from timesketch.models import db_session
 
+logger = logging.getLogger("timesketch.models.annotations")
+
 
 class BaseAnnotation(object):
     """Base class with common attributes."""
@@ -140,9 +143,20 @@ def remove_label(self, label):
         Args:
             label: Name of the label.
         """
-        for label_obj in self.labels:
-            if label_obj.label.lower() != label.lower():
-                continue
+        labels_to_remove = [
+            label_obj
+            for label_obj in self.labels
+            if label_obj.label.lower() == label.lower()
+        ]
+
+        if not labels_to_remove:
+            logger.warning(
+                "Attempted to remove non-existent label: %s from object: %s",
+                str(label),
+                str(type(self).__name__),
+            )
+
+        for label_obj in labels_to_remove:
             self.labels.remove(label_obj)
         db_session.add(self)
         db_session.commit()
@@ -244,15 +258,25 @@ def remove_comment(self, comment_id):
 
         Args:
             comment_id: Id of the comment.
+
+        Returns:
+            True if the comment was removed, False otherwise.
         """
-        for comment_obj in self.comments:
-            if comment_obj.id == int(comment_id):
-                self.comments.remove(comment_obj)
-                db_session.add(self)
-                db_session.commit()
-                return True
 
-        return False
+        comments_to_remove = [
+            comment_obj
+            for comment_obj in self.comments
+            if comment_obj.id == int(comment_id)
+        ]
+        if not comments_to_remove:
+            logger.debug("Comment to delete not found")
+            return False  # Comment not found
+        for comment_obj in comments_to_remove:
+            logger.debug("Removing comment")
+            self.comments.remove(comment_obj)
+        db_session.add(self)
+        db_session.commit()
+        return True
 
     def get_comment(self, comment_id):
         """Retrieves a comment.
@@ -329,8 +353,7 @@ def set_status(self, status):
         Args:
             status: Name of the status
         """
-        for _status in self.status:
-            self.status.remove(_status)
+        self.status = []  # replace the list with an empty list.
         self.status.append(self.Status(user=None, status=status))
         db_session.add(self)
         db_session.commit()
@@ -339,11 +362,30 @@ def set_status(self, status):
     def get_status(self):
         """Get the current status.
 
+        Only one status should be in the database at a time.
+
+        Raises:
+            RuntimeError: If more than one status is available.
+
         Returns:
             The status as a string
         """
         if not self.status:
             self.status.append(self.Status(user=None, status="new"))
+        if len(self.status) > 1:
+            self_id = self.id if hasattr(self, "id") else None
+            # TODO: Change from warning to raising an exception once we ensured
+            # it won't affect the deployment.
+            # raise RuntimeError(
+            # "More than one status available for object [%s] with ID: [%s]",
+            #     str(type(self).__name__),
+            #     str(self_id),
+            # )
+            logging.warning(
+                "More than one status available for object [%s] with ID: [%s]",
+                str(type(self).__name__),
+                str(self_id),
+            )
         return self.status[0]
 
 

timesketch/tsctl.py

@@ -612,9 +612,9 @@ def sketch_info(sketch_id):
                 "user_id",
             ],
         ]
-        for status in sketch.status:
+        for _status in sketch.status:
             status_table.append(
-                [status.id, status.status, status.created_at, status.user_id]
+                [_status.id, _status.status, _status.created_at, _status.user_id]
             )
         print("Status:")
         print_table(status_table)
@@ -671,6 +671,22 @@ def timeline_status(timeline_id, action, status):
             ]
         )
         print_table(table_data)
+
+        status_table = [
+            [
+                "id",
+                "status",
+                "created_at",
+                "user_id",
+            ],
+        ]
+        for _status in timeline.status:
+            status_table.append(
+                [_status.id, _status.status, _status.created_at, _status.user_id]
+            )
+        print("Status:")
+        print_table(status_table)
+
     elif action == "set":
         timeline = Timeline.query.filter_by(id=timeline_id).first()
         if not timeline: