Fix: calling save() doesn't overwrite all the object's contents (fixes #947) (#954)

Author: tomchop

core/database_arango.py

@@ -204,12 +204,19 @@ def _update(self, document_json):
         newdoc["id"] = newdoc.pop("_key")
         return newdoc
 
-    def save(self: TYetiObject) -> TYetiObject:
+    def save(
+            self: TYetiObject,
+            exclude_overwrite: list[str] = ['created', 'tags', 'context']
+        ) -> TYetiObject:
         """Inserts or updates a Yeti object into the database.
 
         We need to pass the JSON representation of the object to the database
         because it may contain fields that are not JSON serializable by arango.
 
+        Args:
+          exclude_overwrite: Exclude overwriting these fields if observable
+            already exists in the database.
+
         Returns:
           The created Yeti object.
         """
@@ -219,8 +226,12 @@ def save(self: TYetiObject) -> TYetiObject:
         else:
             result = self._insert(self.model_dump_json())
             if not result:
-                result = self._update(self.model_dump_json(exclude={"created"}))
-        return self.__class__(**result)
+                result = self._update(self.model_dump_json(exclude=exclude_overwrite))
+        yeti_object = self.__class__(**result)
+        #TODO: Override this if we decide to implement YetiTagModel
+        if hasattr(self, 'tags'):
+            yeti_object.get_tags()
+        return yeti_object
 
     @classmethod
     def list(cls: Type[TYetiObject]) -> Iterable[TYetiObject]:

tests/schemas/observable.py

@@ -23,6 +23,29 @@ def test_observable_create(self) -> None:
         self.assertIsNotNone(result.id)
         self.assertEqual(result.value, "toto.com")
 
+    def test_observable_update(self) -> None:
+        """Tests that calling save() on an observable treats it as PATCH."""
+        result = registry_key.RegistryKey(
+            key="Microsoft\\Windows\\CurrentVersion\\Run",
+            value="persist",
+            data=b"cmd.exe",
+            hive=registry_key.RegistryHive.HKEY_LOCAL_MACHINE_Software).save()
+        result.tag(['tag1'])
+        result.add_context(source='source1', context={'some': 'info'})
+        self.assertEqual(list(result.tags.keys()), ['tag1'])
+        self.assertEqual(
+            result.context[0], {'source': 'source1', 'some': 'info'})
+        result = registry_key.RegistryKey(
+            key="Microsoft\\Windows\\CurrentVersion\\RunOnce",
+            value="persist",
+            data=b"other.exe",
+            hive=registry_key.RegistryHive.HKEY_LOCAL_MACHINE_Software).save()
+        self.assertEqual(result.key, "Microsoft\\Windows\\CurrentVersion\\RunOnce")
+        self.assertEqual(result.data, b"other.exe")
+        self.assertEqual(list(result.tags.keys()), ['tag1'])
+        self.assertEqual(
+            result.context[0], {'source': 'source1', 'some': 'info'})
+
     def test_create_generic_observable(self):
         result = generic_observable.GenericObservable(value="Some_String").save()
         self.assertIsNotNone(result.id)