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/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([

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

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([

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([

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