Add created to object factory, clean up code
parent
6ddad22810
commit
4dba41afc8
|
@ -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)
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue