From 2e45cacd52b22dc27a5df76498c633c8b8301a8b Mon Sep 17 00:00:00 2001
From: clenk <clenk@mitre.org>
Date: Mon, 17 Jul 2017 14:56:13 -0400
Subject: [PATCH] Add to list properties like `external_references`, but
 include option to replace instead

---
 stix2/environment.py           | 16 +++++++++++++++-
 stix2/properties.py            | 11 ++++-------
 stix2/test/test_environment.py | 24 ++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/stix2/environment.py b/stix2/environment.py
index a80cf47..89ce3d4 100644
--- a/stix2/environment.py
+++ b/stix2/environment.py
@@ -2,7 +2,8 @@
 class ObjectFactory(object):
 
     def __init__(self, created_by_ref=None, created=None,
-                 external_references=None, object_marking_refs=None):
+                 external_references=None, object_marking_refs=None,
+                 list_append=True):
 
         self._defaults = {}
         if created_by_ref:
@@ -16,12 +17,25 @@ class ObjectFactory(object):
             self._defaults['external_references'] = external_references
         if object_marking_refs:
             self._defaults['object_marking_refs'] = object_marking_refs
+        self._list_append = list_append
+        self._list_properties = ['external_references', 'object_marking_refs']
 
     def create(self, cls, **kwargs):
         # Use self.defaults as the base, but update with any explicit args
         # provided by the user.
         properties = dict(**self._defaults)
         if kwargs:
+            if self._list_append:
+                # Append provided items to list properties instead of replacing them
+                for list_prop in set(self._list_properties).intersection(kwargs.keys(), properties.keys()):
+                    kwarg_prop = kwargs.pop(list_prop)
+                    if kwarg_prop is None:
+                        del properties[list_prop]
+                        continue
+                    if not isinstance(properties[list_prop], list):
+                        properties[list_prop] = [properties[list_prop]]
+                    properties[list_prop].append(kwarg_prop)
+
             properties.update(**kwargs)
 
         return cls(**properties)
diff --git a/stix2/properties.py b/stix2/properties.py
index 686d1c3..71e4bd9 100644
--- a/stix2/properties.py
+++ b/stix2/properties.py
@@ -5,7 +5,7 @@ import inspect
 import re
 import uuid
 
-from six import text_type
+from six import string_types, text_type
 
 from .base import _Observable, _STIXBase
 from .exceptions import DictionaryKeyError
@@ -101,12 +101,9 @@ class ListProperty(Property):
             iter(value)
         except TypeError:
             raise ValueError("must be an iterable.")
-        try:
-            if isinstance(value, _STIXBase) or isinstance(value, basestring):
-                value = [value]
-        except NameError:
-            if isinstance(value, str):
-                value = [value]
+
+        if isinstance(value, (_STIXBase, string_types)):
+            value = [value]
 
         result = []
         for item in value:
diff --git a/stix2/test/test_environment.py b/stix2/test/test_environment.py
index 549a854..9e5ab98 100644
--- a/stix2/test/test_environment.py
+++ b/stix2/test/test_environment.py
@@ -39,6 +39,9 @@ def test_object_factory_external_resource():
     assert ind.external_references[0].source_name == "ACME Threat Intel"
     assert ind.external_references[0].description == "Threat report"
 
+    ind2 = factory.create(stix2.Indicator, external_references=None, **INDICATOR_KWARGS)
+    assert 'external_references' not in ind2
+
 
 def test_object_factory_obj_markings():
     stmt_marking = stix2.StatementMarking("Copyright 2016, Example Corp")
@@ -52,3 +55,24 @@ def test_object_factory_obj_markings():
     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_list_append():
+    ext_ref = stix2.ExternalReference(source_name="ACME Threat Intel",
+                                      description="Threat report from ACME")
+    ext_ref2 = stix2.ExternalReference(source_name="Yet Another Threat Report",
+                                       description="Threat report from YATR")
+    factory = stix2.ObjectFactory(external_references=ext_ref)
+    ind = factory.create(stix2.Indicator, external_references=ext_ref2, **INDICATOR_KWARGS)
+    assert ind.external_references[1].source_name == "Yet Another Threat Report"
+
+
+def test_object_factory_list_replace():
+    ext_ref = stix2.ExternalReference(source_name="ACME Threat Intel",
+                                      description="Threat report from ACME")
+    ext_ref2 = stix2.ExternalReference(source_name="Yet Another Threat Report",
+                                       description="Threat report from YATR")
+    factory = stix2.ObjectFactory(external_references=ext_ref, list_append=False)
+    ind = factory.create(stix2.Indicator, external_references=ext_ref2, **INDICATOR_KWARGS)
+    assert len(ind.external_references) == 1
+    assert ind.external_references[0].source_name == "Yet Another Threat Report"