Merge pull request #225 from emmanvg/memory-datastore-fix

Memory datastore fix

stix2/datastore/memory.py

@@ -44,14 +44,14 @@ def _add(store, stix_data=None, allow_custom=True, version=None):
         # Map ID directly to the object, if it is a marking.  Otherwise,
         # map to a family, so we can track multiple versions.
         if _is_marking(stix_obj):
-            store._data[stix_obj.id] = stix_obj
+            store._data[stix_obj["id"]] = stix_obj
 
         else:
-            if stix_obj.id in store._data:
-                obj_family = store._data[stix_obj.id]
+            if stix_obj["id"] in store._data:
+                obj_family = store._data[stix_obj["id"]]
             else:
                 obj_family = _ObjectFamily()
-                store._data[stix_obj.id] = obj_family
+                store._data[stix_obj["id"]] = obj_family
 
             obj_family.add(stix_obj)
 
@@ -64,8 +64,8 @@ def _is_marking(obj_or_id):
     :return: True if a marking definition, False otherwise.
     """
 
-    if isinstance(obj_or_id, _STIXBase):
-        id_ = obj_or_id.id
+    if isinstance(obj_or_id, (_STIXBase, dict)):
+        id_ = obj_or_id["id"]
     else:
         id_ = obj_or_id
 
@@ -84,14 +84,14 @@ class _ObjectFamily(object):
         self.latest_version = None
 
     def add(self, obj):
-        self.all_versions[obj.modified] = obj
+        self.all_versions[obj["modified"]] = obj
         if self.latest_version is None or \
-                obj.modified > self.latest_version.modified:
+                obj["modified"] > self.latest_version["modified"]:
             self.latest_version = obj
 
     def __str__(self):
         return "<<{}; latest={}>>".format(self.all_versions,
-                                          self.latest_version.modified)
+                                          self.latest_version["modified"])
 
     def __repr__(self):
         return str(self)