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 .base import _STIXBase
from .markings import MarkingsMixin
from .properties import (HashesProperty, IDProperty, ListProperty, Property,
ReferenceProperty, SelectorProperty, StringProperty,
TimestampProperty, TypeProperty)
@ -76,7 +77,7 @@ class MarkingProperty(Property):
raise ValueError("must be a Statement, TLP Marking or a registered marking.")
class MarkingDefinition(_STIXBase):
class MarkingDefinition(_STIXBase, MarkingsMixin):
_type = 'marking-definition'
_properties = OrderedDict()
_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)
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 .common import ExternalReference, GranularMarking, KillChainPhase
from .markings import MarkingsMixin
from .observables import ObservableProperty
from .properties import (BooleanProperty, IDProperty, IntegerProperty,
ListProperty, PatternProperty, ReferenceProperty,
@ -13,7 +14,7 @@ from .properties import (BooleanProperty, IDProperty, IntegerProperty,
from .utils import NOW
class AttackPattern(_STIXBase):
class AttackPattern(_STIXBase, MarkingsMixin):
_type = 'attack-pattern'
_properties = OrderedDict()
@ -34,7 +35,7 @@ class AttackPattern(_STIXBase):
])
class Campaign(_STIXBase):
class Campaign(_STIXBase, MarkingsMixin):
_type = 'campaign'
_properties = OrderedDict()
@ -58,7 +59,7 @@ class Campaign(_STIXBase):
])
class CourseOfAction(_STIXBase):
class CourseOfAction(_STIXBase, MarkingsMixin):
_type = 'course-of-action'
_properties = OrderedDict()
@ -78,7 +79,7 @@ class CourseOfAction(_STIXBase):
])
class Identity(_STIXBase):
class Identity(_STIXBase, MarkingsMixin):
_type = 'identity'
_properties = OrderedDict()
@ -101,7 +102,7 @@ class Identity(_STIXBase):
])
class Indicator(_STIXBase):
class Indicator(_STIXBase, MarkingsMixin):
_type = 'indicator'
_properties = OrderedDict()
@ -125,7 +126,7 @@ class Indicator(_STIXBase):
])
class IntrusionSet(_STIXBase):
class IntrusionSet(_STIXBase, MarkingsMixin):
_type = 'intrusion-set'
_properties = OrderedDict()
@ -152,7 +153,7 @@ class IntrusionSet(_STIXBase):
])
class Malware(_STIXBase):
class Malware(_STIXBase, MarkingsMixin):
_type = 'malware'
_properties = OrderedDict()
@ -173,7 +174,7 @@ class Malware(_STIXBase):
])
class ObservedData(_STIXBase):
class ObservedData(_STIXBase, MarkingsMixin):
_type = 'observed-data'
_properties = OrderedDict()
@ -195,7 +196,7 @@ class ObservedData(_STIXBase):
])
class Report(_STIXBase):
class Report(_STIXBase, MarkingsMixin):
_type = 'report'
_properties = OrderedDict()
@ -217,7 +218,7 @@ class Report(_STIXBase):
])
class ThreatActor(_STIXBase):
class ThreatActor(_STIXBase, MarkingsMixin):
_type = 'threat-actor'
_properties = OrderedDict()
@ -245,7 +246,7 @@ class ThreatActor(_STIXBase):
])
class Tool(_STIXBase):
class Tool(_STIXBase, MarkingsMixin):
_type = 'tool'
_properties = OrderedDict()
@ -267,7 +268,7 @@ class Tool(_STIXBase):
])
class Vulnerability(_STIXBase):
class Vulnerability(_STIXBase, MarkingsMixin):
_type = 'vulnerability'
_properties = OrderedDict()
@ -316,7 +317,7 @@ def CustomObject(type='x-custom-type', properties=None):
def custom_builder(cls):
class _Custom(cls, _STIXBase):
class _Custom(cls, _STIXBase, MarkingsMixin):
_type = type
_properties = OrderedDict()
_properties.update([

View File

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