Add to list properties like `external_references`, but include option to

replace instead
stix2.1
clenk 2017-07-17 14:56:13 -04:00
parent e233a424fb
commit 2e45cacd52
3 changed files with 43 additions and 8 deletions

View File

@ -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)

View File

@ -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:

View File

@ -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"