Merge pull request #90 from oasis-open/created-by

Add created_by() function to Environment

stix2/environment.py

@@ -152,3 +152,22 @@ class Environment(object):
     def parse(self, *args, **kwargs):
         return _parse(*args, **kwargs)
     parse.__doc__ = _parse.__doc__
+
+    def creator_of(self, obj):
+        """Retrieve the Identity refered to by the object's `created_by_ref`.
+
+        Args:
+            obj: The STIX object whose `created_by_ref` property will be looked
+                up.
+
+        Returns:
+            The STIX object's creator, or
+            None, if the object contains no `created_by_ref` property or the
+                object's creator cannot be found.
+
+        """
+        creator_id = obj.get('created_by_ref', '')
+        if creator_id:
+            return self.get(creator_id)
+        else:
+            return None

stix2/sources/__init__.py

@@ -266,6 +266,8 @@ class CompositeDataSource(DataSource):
         # remove duplicate versions
         if len(all_data) > 0:
             all_data = deduplicate(all_data)
+        else:
+            return None
 
         # reduce to most recent version
         stix_obj = sorted(all_data, key=lambda k: k['modified'], reverse=True)[0]

stix2/sources/memory.py

@@ -215,10 +215,13 @@ class MemorySource(DataSource):
 
         all_data = self.query(query=query, _composite_filters=_composite_filters)
 
+        if all_data:
             # reduce to most recent version
             stix_obj = sorted(all_data, key=lambda k: k['modified'])[0]
 
             return stix_obj
+        else:
+            return None
 
     def all_versions(self, stix_id, _composite_filters=None):
         """retrieve STIX objects from in-memory dict via STIX ID, all versions of it

stix2/test/test_environment.py

@@ -184,3 +184,35 @@ def test_parse_malware():
     assert mal.modified == FAKE_TIME
     assert mal.labels == ['ransomware']
     assert mal.name == "Cryptolocker"
+
+
+def test_created_by():
+    identity = stix2.Identity(**IDENTITY_KWARGS)
+    factory = stix2.ObjectFactory(created_by_ref=identity.id)
+    env = stix2.Environment(store=stix2.MemoryStore(), factory=factory)
+    env.add(identity)
+
+    ind = env.create(stix2.Indicator, **INDICATOR_KWARGS)
+    creator = env.creator_of(ind)
+    assert creator is identity
+
+
+def test_created_by_no_datasource():
+    identity = stix2.Identity(**IDENTITY_KWARGS)
+    factory = stix2.ObjectFactory(created_by_ref=identity.id)
+    env = stix2.Environment(factory=factory)
+
+    ind = env.create(stix2.Indicator, **INDICATOR_KWARGS)
+    with pytest.raises(AttributeError) as excinfo:
+        env.creator_of(ind)
+    assert 'Environment has no data source' in str(excinfo.value)
+
+
+def test_created_by_not_found():
+    identity = stix2.Identity(**IDENTITY_KWARGS)
+    factory = stix2.ObjectFactory(created_by_ref=identity.id)
+    env = stix2.Environment(store=stix2.MemoryStore(), factory=factory)
+
+    ind = env.create(stix2.Indicator, **INDICATOR_KWARGS)
+    creator = env.creator_of(ind)
+    assert creator is None