Add a shortcut for marking functions

So marking functions like `add_markings()` can be directly called as methods on
SDO/SRO classes. Note that they return new versions of the objects
because objects are immutable.
stix2.1
Chris Lenk 2017-10-03 10:28:58 -04:00
parent 58db629de8
commit 2bdf83401a
5 changed files with 43 additions and 17 deletions

View File

@ -3,6 +3,7 @@
from collections import OrderedDict from collections import OrderedDict
from .base import _STIXBase from .base import _STIXBase
from .markings import MarkingsMixin
from .properties import (HashesProperty, IDProperty, ListProperty, Property, from .properties import (HashesProperty, IDProperty, ListProperty, Property,
ReferenceProperty, SelectorProperty, StringProperty, ReferenceProperty, SelectorProperty, StringProperty,
TimestampProperty, TypeProperty) TimestampProperty, TypeProperty)
@ -76,7 +77,7 @@ class MarkingProperty(Property):
raise ValueError("must be a Statement, TLP Marking or a registered marking.") raise ValueError("must be a Statement, TLP Marking or a registered marking.")
class MarkingDefinition(_STIXBase): class MarkingDefinition(_STIXBase, MarkingsMixin):
_type = 'marking-definition' _type = 'marking-definition'
_properties = OrderedDict() _properties = OrderedDict()
_properties.update([ _properties.update([

View File

@ -212,3 +212,16 @@ def is_marked(obj, marking=None, selectors=None, inherited=False, descendants=Fa
result = result or object_markings.is_marked(obj, object_marks) result = result or object_markings.is_marked(obj, object_marks)
return result return result
class MarkingsMixin():
pass
# Note that all of these methods will return a new object because of immutability
MarkingsMixin.get_markings = get_markings
MarkingsMixin.set_markings = set_markings
MarkingsMixin.remove_markings = remove_markings
MarkingsMixin.add_markings = add_markings
MarkingsMixin.clear_markings = clear_markings
MarkingsMixin.is_marked = is_marked

View File

@ -6,6 +6,7 @@ import stix2
from .base import _STIXBase from .base import _STIXBase
from .common import ExternalReference, GranularMarking, KillChainPhase from .common import ExternalReference, GranularMarking, KillChainPhase
from .markings import MarkingsMixin
from .observables import ObservableProperty from .observables import ObservableProperty
from .properties import (BooleanProperty, IDProperty, IntegerProperty, from .properties import (BooleanProperty, IDProperty, IntegerProperty,
ListProperty, PatternProperty, ReferenceProperty, ListProperty, PatternProperty, ReferenceProperty,
@ -13,7 +14,7 @@ from .properties import (BooleanProperty, IDProperty, IntegerProperty,
from .utils import NOW from .utils import NOW
class AttackPattern(_STIXBase): class AttackPattern(_STIXBase, MarkingsMixin):
_type = 'attack-pattern' _type = 'attack-pattern'
_properties = OrderedDict() _properties = OrderedDict()
@ -34,7 +35,7 @@ class AttackPattern(_STIXBase):
]) ])
class Campaign(_STIXBase): class Campaign(_STIXBase, MarkingsMixin):
_type = 'campaign' _type = 'campaign'
_properties = OrderedDict() _properties = OrderedDict()
@ -58,7 +59,7 @@ class Campaign(_STIXBase):
]) ])
class CourseOfAction(_STIXBase): class CourseOfAction(_STIXBase, MarkingsMixin):
_type = 'course-of-action' _type = 'course-of-action'
_properties = OrderedDict() _properties = OrderedDict()
@ -78,7 +79,7 @@ class CourseOfAction(_STIXBase):
]) ])
class Identity(_STIXBase): class Identity(_STIXBase, MarkingsMixin):
_type = 'identity' _type = 'identity'
_properties = OrderedDict() _properties = OrderedDict()
@ -101,7 +102,7 @@ class Identity(_STIXBase):
]) ])
class Indicator(_STIXBase): class Indicator(_STIXBase, MarkingsMixin):
_type = 'indicator' _type = 'indicator'
_properties = OrderedDict() _properties = OrderedDict()
@ -125,7 +126,7 @@ class Indicator(_STIXBase):
]) ])
class IntrusionSet(_STIXBase): class IntrusionSet(_STIXBase, MarkingsMixin):
_type = 'intrusion-set' _type = 'intrusion-set'
_properties = OrderedDict() _properties = OrderedDict()
@ -152,7 +153,7 @@ class IntrusionSet(_STIXBase):
]) ])
class Malware(_STIXBase): class Malware(_STIXBase, MarkingsMixin):
_type = 'malware' _type = 'malware'
_properties = OrderedDict() _properties = OrderedDict()
@ -173,7 +174,7 @@ class Malware(_STIXBase):
]) ])
class ObservedData(_STIXBase): class ObservedData(_STIXBase, MarkingsMixin):
_type = 'observed-data' _type = 'observed-data'
_properties = OrderedDict() _properties = OrderedDict()
@ -195,7 +196,7 @@ class ObservedData(_STIXBase):
]) ])
class Report(_STIXBase): class Report(_STIXBase, MarkingsMixin):
_type = 'report' _type = 'report'
_properties = OrderedDict() _properties = OrderedDict()
@ -217,7 +218,7 @@ class Report(_STIXBase):
]) ])
class ThreatActor(_STIXBase): class ThreatActor(_STIXBase, MarkingsMixin):
_type = 'threat-actor' _type = 'threat-actor'
_properties = OrderedDict() _properties = OrderedDict()
@ -245,7 +246,7 @@ class ThreatActor(_STIXBase):
]) ])
class Tool(_STIXBase): class Tool(_STIXBase, MarkingsMixin):
_type = 'tool' _type = 'tool'
_properties = OrderedDict() _properties = OrderedDict()
@ -267,7 +268,7 @@ class Tool(_STIXBase):
]) ])
class Vulnerability(_STIXBase): class Vulnerability(_STIXBase, MarkingsMixin):
_type = 'vulnerability' _type = 'vulnerability'
_properties = OrderedDict() _properties = OrderedDict()
@ -316,7 +317,7 @@ def CustomObject(type='x-custom-type', properties=None):
def custom_builder(cls): def custom_builder(cls):
class _Custom(cls, _STIXBase): class _Custom(cls, _STIXBase, MarkingsMixin):
_type = type _type = type
_properties = OrderedDict() _properties = OrderedDict()
_properties.update([ _properties.update([

View File

@ -4,13 +4,14 @@ from collections import OrderedDict
from .base import _STIXBase from .base import _STIXBase
from .common import ExternalReference, GranularMarking from .common import ExternalReference, GranularMarking
from .markings import MarkingsMixin
from .properties import (BooleanProperty, IDProperty, IntegerProperty, from .properties import (BooleanProperty, IDProperty, IntegerProperty,
ListProperty, ReferenceProperty, StringProperty, ListProperty, ReferenceProperty, StringProperty,
TimestampProperty, TypeProperty) TimestampProperty, TypeProperty)
from .utils import NOW from .utils import NOW
class Relationship(_STIXBase): class Relationship(_STIXBase, MarkingsMixin):
_type = 'relationship' _type = 'relationship'
_properties = OrderedDict() _properties = OrderedDict()
@ -45,7 +46,7 @@ class Relationship(_STIXBase):
super(Relationship, self).__init__(**kwargs) super(Relationship, self).__init__(**kwargs)
class Sighting(_STIXBase): class Sighting(_STIXBase, MarkingsMixin):
_type = 'sighting' _type = 'sighting'
_properties = OrderedDict() _properties = OrderedDict()
_properties.update([ _properties.update([

View File

@ -241,4 +241,14 @@ def test_marking_wrong_type_construction():
assert str(excinfo.value) == "Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]" assert str(excinfo.value) == "Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]"
# TODO: Add other examples def test_campaign_add_markings():
campaign = stix2.Campaign(
id="campaign--8e2e2d2b-17d4-4cbf-938f-98ee46b3cd3f",
created_by_ref="identity--f431f809-377b-45e0-aa1c-6a4751cae5ff",
created="2016-04-06T20:03:00Z",
modified="2016-04-06T20:03:00Z",
name="Green Group Attacks Against Finance",
description="Campaign by Green Group against a series of targets in the financial services sector.",
)
campaign = campaign.add_markings(TLP_WHITE)
assert campaign.object_marking_refs[0] == TLP_WHITE.id