Minor changes to custom object decorators. Add new marking decorator.

stix2/__init__.py

@@ -4,8 +4,9 @@
 
 from . import exceptions
 from .common import (TLP_AMBER, TLP_GREEN, TLP_RED, TLP_WHITE,
-                     ExternalReference, GranularMarking, KillChainPhase,
+                     CustomMarking, ExternalReference, GranularMarking,
-                     MarkingDefinition, StatementMarking, TLPMarking)
+                     KillChainPhase, MarkingDefinition, StatementMarking,
+                     TLPMarking)
 from .core import Bundle, _register_type, parse
 from .environment import ObjectFactory
 from .observables import (URL, AlternateDataStream, ArchiveExt, Artifact,

stix2/common.py

@@ -109,17 +109,54 @@ class MarkingDefinition(_STIXBase):
         super(MarkingDefinition, self).__init__(**kwargs)
 
 
-def register_marking(new_marking):
-    """Register a custom STIX Marking Definition type.
-    """
-    OBJ_MAP_MARKING[new_marking._type] = new_marking
-
-
 OBJ_MAP_MARKING = {
     'tlp': TLPMarking,
     'statement': StatementMarking,
 }
 
+
+def _register_marking(cls):
+    """Register a custom STIX Marking Definition type.
+    """
+    OBJ_MAP_MARKING[cls._type] = cls
+    return cls
+
+
+def CustomMarking(type='x-custom-marking', properties=None):
+    """
+    Custom STIX Marking decorator.
+
+    Examples:
+
+        @CustomMarking('x-custom-marking', [
+            ('property1', StringProperty(required=True)),
+            ('property2', IntegerProperty()),
+        ])
+        class MyNewMarkingObjectType():
+            pass
+
+    """
+    def custom_builder(cls):
+
+        class _Custom(cls, _STIXBase):
+            _type = type
+            _properties = OrderedDict()
+
+            if not properties or not isinstance(properties, list):
+                raise ValueError("Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]")
+
+            _properties.update(properties)
+
+            def __init__(self, **kwargs):
+                _STIXBase.__init__(self, **kwargs)
+                cls.__init__(self, **kwargs)
+
+        _register_marking(_Custom)
+        return _Custom
+
+    return custom_builder
+
+
 TLP_WHITE = MarkingDefinition(
     id="marking-definition--613f2e26-407d-48c7-9eca-b8e91df99dc9",
     created="2017-01-20T00:00:00.000Z",

stix2/observables.py

@@ -804,7 +804,7 @@ def CustomObservable(type='x-custom-observable', properties=None):
                 ('type', TypeProperty(_type)),
             ])
 
-            if not properties:
+            if not properties or not isinstance(properties, list):
                 raise ValueError("Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]")
 
             _properties.update(properties)

stix2/sdo.py

@@ -330,13 +330,10 @@ def CustomObject(type='x-custom-type', properties=None):
                 ('modified', TimestampProperty(default=lambda: NOW, precision='millisecond')),
             ])
 
-            if not properties:
+            if not properties or not isinstance(properties, list):
                 raise ValueError("Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]")
 
-            normal_properties = [x for x in properties if not x[0].startswith("x_")]
+            _properties.update([x for x in properties if not x[0].startswith("x_")])
-            custom_properties = [x for x in properties if x[0].startswith("x_")]
-
-            _properties.update(normal_properties)
 
             # This is to follow the general properties structure.
             _properties.update([
@@ -348,7 +345,7 @@ def CustomObject(type='x-custom-type', properties=None):
             ])
 
             # Put all custom properties at the bottom, sorted alphabetically.
-            _properties.update(sorted(custom_properties, key=lambda x: x[0]))
+            _properties.update(sorted([x for x in properties if x[0].startswith("x_")], key=lambda x: x[0]))
 
             def __init__(self, **kwargs):
                 _STIXBase.__init__(self, **kwargs)