Repository: oasis-open/cti-python-stix2

Issue:
Adding SCO to the MemoryStore results in error
Below an example of how to trigger the problem:

``` python
from stix2.v21 import observables
from stix2 import MemoryStore

mem_store = MemoryStore()

file_observable = observables.File(
    name="example.exe",
    size=68 * 1000,
    magic_number_hex="50000000",
    hashes={
        "SHA-256": "841a8921140aba50671ebb0770fecc4ee308c4952cfeff8de154ab14eeef4649"
    }
)

mem_store.add(file_observable)
```

I obtain the following traceback:

```
Traceback (most recent call last):
  File "C:/Users/user-1/example_project/mem_store.py", line 98, in <module>
    example()
  File "C:/Users/user-1/example_project/mem_store.py", line 94, in example
    mem_store.add(file_observable)
  File "C:\Users\user-1\venv\lib\site-packages\stix2\datastore\__init__.py", line 216, in add
    return self.sink.add(*args, **kwargs)
  File "C:\Users\user-1\venv\lib\site-packages\stix2\datastore\memory.py", line 185, in add
    _add(self, stix_data, self.allow_custom, version)
  File "C:\Users\user-1\venv\lib\site-packages\stix2\datastore\memory.py", line 62, in _add
    obj_family.add(stix_obj)
  File "C:\Users\user-1\venv\lib\site-packages\stix2\datastore\memory.py", line 77, in add
    self.all_versions[obj["modified"]] = obj
  File "C:\Users\user-1\venv\lib\site-packages\stix2\base.py", line 195, in __getitem__
    return self._inner[key]
KeyError: 'modified'
```


Discussion:

Patch:
diff --git a/stix2/datastore/memory.py b/stix2/datastore/memory.py
index 52bf4c8..52da168 100644
--- a/stix2/datastore/memory.py
+++ b/stix2/datastore/memory.py
@@ -10,7 +10,6 @@ from stix2.base import _STIXBase
 from stix2.core import parse
 from stix2.datastore import DataSink, DataSource, DataStoreMixin
 from stix2.datastore.filters import FilterSet, apply_common_filters
-from stix2.utils import is_marking
 
 
 def _add(store, stix_data, allow_custom=True, version=None):
@@ -47,12 +46,10 @@ def _add(store, stix_data, allow_custom=True, version=None):
         else:
             stix_obj = parse(stix_data, allow_custom, version)
 
-        # 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
-
-        else:
+        # Map ID to a _ObjectFamily if the object is versioned, so we can track
+        # multiple versions.  Otherwise, map directly to the object.  All
+        # versioned objects should have a "modified" property.
+        if "modified" in stix_obj:
             if stix_obj["id"] in store._data:
                 obj_family = store._data[stix_obj["id"]]
             else:
@@ -61,6 +58,9 @@ def _add(store, stix_data, allow_custom=True, version=None):
 
             obj_family.add(stix_obj)
 
+        else:
+            store._data[stix_obj["id"]] = stix_obj
+
 
 class _ObjectFamily(object):
     """
@@ -267,12 +267,12 @@ class MemorySource(DataSource):
         """
         stix_obj = None
 
-        if is_marking(stix_id):
-            stix_obj = self._data.get(stix_id)
-        else:
-            object_family = self._data.get(stix_id)
-            if object_family:
-                stix_obj = object_family.latest_version
+        mapped_value = self._data.get(stix_id)
+        if mapped_value:
+            if isinstance(mapped_value, _ObjectFamily):
+                stix_obj = mapped_value.latest_version
+            else:
+                stix_obj = mapped_value
 
         if stix_obj:
             all_filters = list(
@@ -300,17 +300,13 @@ class MemorySource(DataSource):
 
         """
         results = []
-        stix_objs_to_filter = None
-        if is_marking(stix_id):
-            stix_obj = self._data.get(stix_id)
-            if stix_obj:
-                stix_objs_to_filter = [stix_obj]
-        else:
-            object_family = self._data.get(stix_id)
-            if object_family:
-                stix_objs_to_filter = object_family.all_versions.values()
+        mapped_value = self._data.get(stix_id)
+        if mapped_value:
+            if isinstance(mapped_value, _ObjectFamily):
+                stix_objs_to_filter = mapped_value.all_versions.values()
+            else:
+                stix_objs_to_filter = [mapped_value]
 
-        if stix_objs_to_filter:
             all_filters = list(
                 itertools.chain(
                     _composite_filters or [],


Test patch:
diff --git a/stix2/test/v20/test_datastore_memory.py b/stix2/test/v20/test_datastore_memory.py
index fba96dd..7852746 100644
--- a/stix2/test/v20/test_datastore_memory.py
+++ b/stix2/test/v20/test_datastore_memory.py
@@ -423,3 +423,24 @@ def test_object_family_internal_components(mem_source):
 
     assert "latest=2017-01-27 13:49:53.936000+00:00>>" in str_representation
     assert "latest=2017-01-27 13:49:53.936000+00:00>>" in repr_representation
+
+
+def test_unversioned_objects(mem_store):
+    marking = {
+        "type": "marking-definition",
+        "id": "marking-definition--48e83cde-e902-4404-85b3-6e81f75ccb62",
+        "created": "1988-01-02T16:44:04.000Z",
+        "definition_type": "statement",
+        "definition": {
+            "statement": "Copyright (C) ACME Corp.",
+        },
+    }
+
+    mem_store.add(marking)
+
+    obj = mem_store.get(marking["id"])
+    assert obj["id"] == marking["id"]
+
+    objs = mem_store.all_versions(marking["id"])
+    assert len(objs) == 1
+    assert objs[0]["id"] == marking["id"]
diff --git a/stix2/test/v21/test_datastore_memory.py b/stix2/test/v21/test_datastore_memory.py
index 4f63a06..e07943c 100644
--- a/stix2/test/v21/test_datastore_memory.py
+++ b/stix2/test/v21/test_datastore_memory.py
@@ -438,3 +438,38 @@ def test_object_family_internal_components(mem_source):
 
     assert "latest=2017-01-27 13:49:53.936000+00:00>>" in str_representation
     assert "latest=2017-01-27 13:49:53.936000+00:00>>" in repr_representation
+
+
+def test_unversioned_objects(mem_store):
+    marking = {
+        "type": "marking-definition",
+        "spec_version": "2.1",
+        "id": "marking-definition--48e83cde-e902-4404-85b3-6e81f75ccb62",
+        "created": "1988-01-02T16:44:04.000Z",
+        "definition_type": "statement",
+        "definition": {
+            "statement": "Copyright (C) ACME Corp.",
+        },
+    }
+
+    file_sco = {
+        "type": "file",
+        "id": "file--bbd59c0c-1aa4-44f1-96de-80b8325372c7",
+        "name": "cats.png",
+    }
+
+    mem_store.add([marking, file_sco])
+
+    obj = mem_store.get(marking["id"])
+    assert obj["id"] == marking["id"]
+
+    obj = mem_store.get(file_sco["id"])
+    assert obj["id"] == file_sco["id"]
+
+    objs = mem_store.all_versions(marking["id"])
+    assert len(objs) == 1
+    assert objs[0]["id"] == marking["id"]
+
+    objs = mem_store.all_versions(file_sco["id"])
+    assert len(objs) == 1
+    assert objs[0]["id"] == file_sco["id"]