diff --git a/stix2/common.py b/stix2/common.py index a2e6918..d7994c6 100644 --- a/stix2/common.py +++ b/stix2/common.py @@ -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([ diff --git a/stix2/markings/__init__.py b/stix2/markings/__init__.py index 4f72e4c..41c761d 100644 --- a/stix2/markings/__init__.py +++ b/stix2/markings/__init__.py @@ -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 diff --git a/stix2/sdo.py b/stix2/sdo.py index 77c781a..ea4bd13 100644 --- a/stix2/sdo.py +++ b/stix2/sdo.py @@ -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([ diff --git a/stix2/sro.py b/stix2/sro.py index af483bc..0a8048d 100644 --- a/stix2/sro.py +++ b/stix2/sro.py @@ -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([ diff --git a/stix2/test/test_markings.py b/stix2/test/test_markings.py index 0c6069a..456bf92 100644 --- a/stix2/test/test_markings.py +++ b/stix2/test/test_markings.py @@ -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