Add custom Cyber Observables

stix2/__init__.py

@@ -114,12 +114,11 @@ def parse(data, allow_custom=False):
     try:
         obj_class = OBJ_MAP[obj['type']]
     except KeyError:
-        # TODO handle custom objects
-        raise exceptions.ParseError("Can't parse unknown object type '%s'!" % obj['type'])
+        raise exceptions.ParseError("Can't parse unknown object type '%s'! For custom types, use the CustomObject decorator." % obj['type'])
     return obj_class(allow_custom=allow_custom, **obj)
 
 
-def parse_observable(data, _valid_refs, allow_custom=False):
+def parse_observable(data, _valid_refs=[], allow_custom=False):
     """Deserialize a string or file-like object into a STIX Cyber Observable object.
 
     Args:
@@ -139,8 +138,7 @@ def parse_observable(data, _valid_refs, allow_custom=False):
     try:
         obj_class = OBJ_MAP_OBSERVABLE[obj['type']]
     except KeyError:
-        # TODO handle custom observable objects
-        raise exceptions.ParseError("Can't parse unknown object type '%s'!" % obj['type'])
+        raise exceptions.ParseError("Can't parse unknown object type '%s'! For custom observables, use the CustomObservable decorator." % obj['type'])
 
     if 'extensions' in obj and obj['type'] in EXT_MAP:
         for name, ext in obj['extensions'].items():
@@ -157,3 +155,10 @@ def _register_type(new_type):
     """
 
     OBJ_MAP[new_type._type] = new_type
+
+
+def _register_observable(new_observable):
+    """Register a custom STIX Cyber Observable type.
+    """
+
+    OBJ_MAP_OBSERVABLE[new_observable._type] = new_observable

stix2/observables.py

@@ -5,6 +5,8 @@ embedded in Email Message objects, inherit from _STIXBase instead of Observable
 and do not have a '_type' attribute.
 """
 
+import stix2
+
 from .base import _Extension, _Observable, _STIXBase
 from .exceptions import AtLeastOnePropertyError, DependentPropertiesError
 from .properties import (BinaryProperty, BooleanProperty, DictionaryProperty,
@@ -586,3 +588,27 @@ class X509Certificate(_Observable):
         'subject_public_key_exponent': IntegerProperty(),
         'x509_v3_extensions': EmbeddedObjectProperty(type=X509V3ExtenstionsType),
     }
+
+
+def CustomObservable(type='x-custom-observable', properties={}):
+    """Custom STIX Cyber Observable type decorator
+
+    """
+
+    def custom_builder(cls):
+
+        class _Custom(cls, _Observable):
+            _type = type
+            _properties = {
+                'type': TypeProperty(_type),
+            }
+            _properties.update(properties)
+
+            def __init__(self, **kwargs):
+                _Observable.__init__(self, **kwargs)
+                cls.__init__(self, **kwargs)
+
+        stix2._register_observable(_Custom)
+        return _Custom
+
+    return custom_builder

stix2/sdo.py

@@ -199,7 +199,7 @@ def CustomObject(type='x-custom-type', properties={}):
 
     Example 1:
 
-    @CustomObject('type-name', {
+    @CustomObject('x-type-name', {
         'property1': StringProperty(required=True),
         'property2': IntegerProperty(),
     })
@@ -211,7 +211,7 @@ def CustomObject(type='x-custom-type', properties={}):
 
     Example 2:
 
-    @CustomObject('type-name', {
+    @CustomObject('x-type-name', {
         'property1': StringProperty(required=True),
         'property2': IntegerProperty(),
     })

stix2/test/test_custom.py

@@ -72,7 +72,7 @@ def test_parse_identity_custom_property(data):
     assert identity.foo == "bar"
 
 
-@stix2.sdo.CustomObject('new-type', {
+@stix2.sdo.CustomObject('x-new-type', {
     'property1': stix2.properties.StringProperty(required=True),
     'property2': stix2.properties.IntegerProperty(),
 })
@@ -95,10 +95,74 @@ def test_custom_object_type():
 
 def test_parse_custom_object_type():
     nt_string = """{
-        "type": "new-type",
+        "type": "x-new-type",
         "created": "2015-12-21T19:59:11Z",
         "property1": "something"
     }"""
 
     nt = stix2.parse(nt_string)
     assert nt.property1 == 'something'
+
+
+@stix2.observables.CustomObservable('x-new-observable', {
+    'property1': stix2.properties.StringProperty(required=True),
+    'property2': stix2.properties.IntegerProperty(),
+})
+class NewObservable():
+    def __init__(self, property2=None, **kwargs):
+        if property2 and property2 < 10:
+            raise ValueError("'property2' is too small.")
+
+
+def test_custom_observable_object():
+    no = NewObservable(property1='something')
+    assert no.property1 == 'something'
+
+    with pytest.raises(stix2.exceptions.MissingPropertiesError):
+        NewObservable(property2=42)
+
+    with pytest.raises(ValueError):
+        NewObservable(property1='something', property2=4)
+
+
+def test_parse_custom_observable_object():
+    nt_string = """{
+        "type": "x-new-observable",
+        "property1": "something"
+    }"""
+
+    nt = stix2.parse_observable(nt_string)
+    assert nt.property1 == 'something'
+
+
+def test_observable_custom_property():
+    with pytest.raises(ValueError):
+        NewObservable(
+            property1='something',
+            custom_properties="foobar",
+        )
+
+    no = NewObservable(
+        property1='something',
+        custom_properties={
+            "foo": "bar",
+        },
+    )
+    assert no.foo == "bar"
+
+
+def test_observable_custom_property_invalid():
+    with pytest.raises(stix2.exceptions.ExtraPropertiesError):
+        NewObservable(
+            property1='something',
+            x_foo="bar",
+        )
+
+
+def test_observable_custom_property_allowed():
+    no = NewObservable(
+        property1='something',
+        x_foo="bar",
+        allow_custom=True,
+    )
+    assert no.x_foo == "bar"