Add to list properties like `external_references`, but include option to
replace insteadstix2.1
parent
e233a424fb
commit
2e45cacd52
|
@ -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)
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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"
|
||||
|
|
Loading…
Reference in New Issue