Add created to object factory, clean up code

stix2/environment.py

@@ -1,21 +1,27 @@
 
 class ObjectFactory(object):
 
-    def __init__(self, created_by=None, object_markings=None,
-                 granular_markings=None, external_references=None):
-        self.created_by = created_by
-        self.object_markings = object_markings
-        self.granular_markings = granular_markings
+    def __init__(self, created_by_ref=None, created=None,
+                 external_references=None, object_marking_refs=None,
+                 granular_markings=None):
+        self.created_by_ref = created_by_ref
+        self.created = created
         self.external_references = external_references
+        self.object_marking_refs = object_marking_refs
+        self.granular_markings = granular_markings
 
     def create(self, cls, **kwargs):
-        if self.created_by is not None:
-            kwargs['created_by_ref'] = self.created_by
-        if self.object_markings is not None:
-            kwargs['object_marking_refs'] = self.object_markings
-        if self.granular_markings is not None:
-            kwargs['granular_markings'] = self.granular_markings
-        if self.external_references is not None:
+        if 'created_by_ref' not in kwargs and self.created_by_ref is not None:
+            kwargs['created_by_ref'] = self.created_by_ref
+        if 'created' not in kwargs and self.created is not None:
+            kwargs['created'] = self.created
+            if 'modified' not in kwargs:
+                kwargs['modified'] = self.created
+        if 'external_references' not in kwargs and self.external_references is not None:
             kwargs['external_references'] = self.external_references
+        if 'object_marking_refs' not in kwargs and self.object_marking_refs is not None:
+            kwargs['object_marking_refs'] = self.object_marking_refs
+        if 'granular_markings' not in kwargs and self.granular_markings is not None:
+            kwargs['granular_markings'] = self.granular_markings
 
         return cls(**kwargs)

stix2/test/test_environment.py

@@ -1,42 +1,34 @@
 import stix2
 
-from .constants import IDENTITY_ID, IDENTITY_KWARGS, INDICATOR_KWARGS
+from .constants import (FAKE_TIME, IDENTITY_ID, IDENTITY_KWARGS,
+                        INDICATOR_KWARGS)
 
 
 def test_object_factory_created_by_ref_str():
-    factory = stix2.ObjectFactory(created_by=IDENTITY_ID)
+    factory = stix2.ObjectFactory(created_by_ref=IDENTITY_ID)
     ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
     assert ind.created_by_ref == IDENTITY_ID
 
 
 def test_object_factory_created_by_ref_obj():
     id_obj = stix2.Identity(id=IDENTITY_ID, **IDENTITY_KWARGS)
-    factory = stix2.ObjectFactory(created_by=id_obj)
+    factory = stix2.ObjectFactory(created_by_ref=id_obj)
     ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
     assert ind.created_by_ref == IDENTITY_ID
 
 
-def test_object_factory_obj_markings():
-    stmt_marking = stix2.StatementMarking("Copyright 2016, Example Corp")
-    mark_def = stix2.MarkingDefinition(definition_type="statement",
-                                       definition=stmt_marking)
-    factory = stix2.ObjectFactory(object_markings=[mark_def, stix2.TLP_AMBER])
-    ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
-    assert mark_def.id in ind.object_marking_refs
-    assert stix2.TLP_AMBER.id in ind.object_marking_refs
-
-    factory = stix2.ObjectFactory(object_markings=stix2.TLP_RED)
-    ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
-    assert stix2.TLP_RED.id in ind.object_marking_refs
+def test_object_factory_override_default():
+    factory = stix2.ObjectFactory(created_by_ref=IDENTITY_ID)
+    new_id = "identity--983b3172-44fe-4a80-8091-eb8098841fe8"
+    ind = factory.create(stix2.Indicator, created_by_ref=new_id, **INDICATOR_KWARGS)
+    assert ind.created_by_ref == new_id
 
 
-def test_object_factory_granular_markings():
-    marking = stix2.GranularMarking(marking_ref=stix2.TLP_AMBER,
-                                    selectors="created_by_ref")
-    factory = stix2.ObjectFactory(created_by=IDENTITY_ID,
-                                  granular_markings=marking)
+def test_object_factory_created():
+    factory = stix2.ObjectFactory(created=FAKE_TIME)
     ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
-    assert "created_by_ref" in ind.granular_markings[0].selectors
+    assert ind.created == FAKE_TIME
+    assert ind.modified == FAKE_TIME
 
 
 def test_object_factory_external_resource():
@@ -46,3 +38,26 @@ def test_object_factory_external_resource():
     ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
     assert ind.external_references[0].source_name == "ACME Threat Intel"
     assert ind.external_references[0].description == "Threat report"
+
+
+def test_object_factory_obj_markings():
+    stmt_marking = stix2.StatementMarking("Copyright 2016, Example Corp")
+    mark_def = stix2.MarkingDefinition(definition_type="statement",
+                                       definition=stmt_marking)
+    factory = stix2.ObjectFactory(object_marking_refs=[mark_def, stix2.TLP_AMBER])
+    ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
+    assert mark_def.id in ind.object_marking_refs
+    assert stix2.TLP_AMBER.id in ind.object_marking_refs
+
+    factory = stix2.ObjectFactory(object_marking_refs=stix2.TLP_RED)
+    ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
+    assert stix2.TLP_RED.id in ind.object_marking_refs
+
+
+def test_object_factory_granular_markings():
+    marking = stix2.GranularMarking(marking_ref=stix2.TLP_AMBER,
+                                    selectors="created_by_ref")
+    factory = stix2.ObjectFactory(created_by_ref=IDENTITY_ID,
+                                  granular_markings=marking)
+    ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
+    assert "created_by_ref" in ind.granular_markings[0].selectors