Add granular markings to object factory

stix2.1
clenk 2017-07-12 15:22:50 -04:00
parent 659ad3df89
commit 593f16662a
2 changed files with 19 additions and 10 deletions

View File

@ -1,17 +1,17 @@
class ObjectFactory(object):
def __init__(self, created_by=None, markings=None):
def __init__(self, created_by=None, object_markings=None, granular_markings=None):
self.created_by = created_by
self.markings = markings
def create(self, cls=None, **kwargs):
if cls is None:
raise ValueError('No STIX object class provided')
self.object_markings = object_markings
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.markings is not None:
kwargs['object_marking_refs'] = self.markings
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
return cls(**kwargs)

View File

@ -20,11 +20,20 @@ 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(markings=[mark_def, stix2.TLP_AMBER])
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(markings=stix2.TLP_RED)
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_granular_markings():
marking = stix2.GranularMarking(marking_ref=stix2.TLP_AMBER,
selectors="created_by_ref")
factory = stix2.ObjectFactory(created_by=IDENTITY_ID,
granular_markings=marking)
ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
assert "created_by_ref" in ind.granular_markings[0].selectors