Add created to object factory, clean up code

stix2.1
clenk 2017-07-13 09:45:43 -04:00
parent 6ddad22810
commit 4dba41afc8
2 changed files with 54 additions and 33 deletions

View File

@ -1,21 +1,27 @@
class ObjectFactory(object): class ObjectFactory(object):
def __init__(self, created_by=None, object_markings=None, def __init__(self, created_by_ref=None, created=None,
granular_markings=None, external_references=None): external_references=None, object_marking_refs=None,
self.created_by = created_by granular_markings=None):
self.object_markings = object_markings self.created_by_ref = created_by_ref
self.granular_markings = granular_markings self.created = created
self.external_references = external_references self.external_references = external_references
self.object_marking_refs = object_marking_refs
self.granular_markings = granular_markings
def create(self, cls, **kwargs): def create(self, cls, **kwargs):
if self.created_by is not None: if 'created_by_ref' not in kwargs and self.created_by_ref is not None:
kwargs['created_by_ref'] = self.created_by kwargs['created_by_ref'] = self.created_by_ref
if self.object_markings is not None: if 'created' not in kwargs and self.created is not None:
kwargs['object_marking_refs'] = self.object_markings kwargs['created'] = self.created
if self.granular_markings is not None: if 'modified' not in kwargs:
kwargs['granular_markings'] = self.granular_markings kwargs['modified'] = self.created
if self.external_references is not None: if 'external_references' not in kwargs and self.external_references is not None:
kwargs['external_references'] = self.external_references kwargs['external_references'] = self.external_references
if 'object_marking_refs' not in kwargs and self.object_marking_refs is not None:
kwargs['object_marking_refs'] = self.object_marking_refs
if 'granular_markings' not in kwargs and self.granular_markings is not None:
kwargs['granular_markings'] = self.granular_markings
return cls(**kwargs) return cls(**kwargs)

View File

@ -1,42 +1,34 @@
import stix2 import stix2
from .constants import IDENTITY_ID, IDENTITY_KWARGS, INDICATOR_KWARGS from .constants import (FAKE_TIME, IDENTITY_ID, IDENTITY_KWARGS,
INDICATOR_KWARGS)
def test_object_factory_created_by_ref_str(): def test_object_factory_created_by_ref_str():
factory = stix2.ObjectFactory(created_by=IDENTITY_ID) factory = stix2.ObjectFactory(created_by_ref=IDENTITY_ID)
ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS) ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
assert ind.created_by_ref == IDENTITY_ID assert ind.created_by_ref == IDENTITY_ID
def test_object_factory_created_by_ref_obj(): def test_object_factory_created_by_ref_obj():
id_obj = stix2.Identity(id=IDENTITY_ID, **IDENTITY_KWARGS) id_obj = stix2.Identity(id=IDENTITY_ID, **IDENTITY_KWARGS)
factory = stix2.ObjectFactory(created_by=id_obj) factory = stix2.ObjectFactory(created_by_ref=id_obj)
ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS) ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
assert ind.created_by_ref == IDENTITY_ID assert ind.created_by_ref == IDENTITY_ID
def test_object_factory_obj_markings(): def test_object_factory_override_default():
stmt_marking = stix2.StatementMarking("Copyright 2016, Example Corp") factory = stix2.ObjectFactory(created_by_ref=IDENTITY_ID)
mark_def = stix2.MarkingDefinition(definition_type="statement", new_id = "identity--983b3172-44fe-4a80-8091-eb8098841fe8"
definition=stmt_marking) ind = factory.create(stix2.Indicator, created_by_ref=new_id, **INDICATOR_KWARGS)
factory = stix2.ObjectFactory(object_markings=[mark_def, stix2.TLP_AMBER]) assert ind.created_by_ref == new_id
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(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(): def test_object_factory_created():
marking = stix2.GranularMarking(marking_ref=stix2.TLP_AMBER, factory = stix2.ObjectFactory(created=FAKE_TIME)
selectors="created_by_ref")
factory = stix2.ObjectFactory(created_by=IDENTITY_ID,
granular_markings=marking)
ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS) ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
assert "created_by_ref" in ind.granular_markings[0].selectors assert ind.created == FAKE_TIME
assert ind.modified == FAKE_TIME
def test_object_factory_external_resource(): def test_object_factory_external_resource():
@ -46,3 +38,26 @@ def test_object_factory_external_resource():
ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS) ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
assert ind.external_references[0].source_name == "ACME Threat Intel" assert ind.external_references[0].source_name == "ACME Threat Intel"
assert ind.external_references[0].description == "Threat report" assert ind.external_references[0].description == "Threat report"
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(object_marking_refs=[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(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_granular_markings():
marking = stix2.GranularMarking(marking_ref=stix2.TLP_AMBER,
selectors="created_by_ref")
factory = stix2.ObjectFactory(created_by_ref=IDENTITY_ID,
granular_markings=marking)
ind = factory.create(stix2.Indicator, **INDICATOR_KWARGS)
assert "created_by_ref" in ind.granular_markings[0].selectors