Add granular markings to object factory
parent
659ad3df89
commit
593f16662a
|
@ -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)
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue