diff --git a/stix2/__init__.py b/stix2/__init__.py index fa015f4..f9f685b 100644 --- a/stix2/__init__.py +++ b/stix2/__init__.py @@ -150,3 +150,10 @@ def parse_observable(data, _valid_refs, allow_custom=False): obj['extensions'][name] = ext_class(allow_custom=allow_custom, **obj['extensions'][name]) return obj_class(allow_custom=allow_custom, **obj) + + +def _register_type(new_type): + """Register a custom STIX Object type. + """ + + OBJ_MAP[new_type._type] = new_type diff --git a/stix2/sdo.py b/stix2/sdo.py index 73bb34d..b38e5a8 100644 --- a/stix2/sdo.py +++ b/stix2/sdo.py @@ -1,5 +1,7 @@ """STIX 2.0 Domain Objects""" +import stix2 + from .base import _STIXBase from .common import COMMON_PROPERTIES from .other import KillChainPhase @@ -190,3 +192,22 @@ class Vulnerability(_STIXBase): 'name': StringProperty(required=True), 'description': StringProperty(), }) + + +def CustomObject(type='x-custom-type', properties={}): + """Custom STIX Object type decorator""" + + def custom_builder(cls): + + class _Custom(_STIXBase): + _type = type + _properties = COMMON_PROPERTIES.copy() + _properties.update({ + 'id': IDProperty(_type), + 'type': TypeProperty(_type), + }) + _properties.update(properties) + stix2._register_type(_Custom) + return _Custom + + return custom_builder diff --git a/stix2/test/test_custom.py b/stix2/test/test_custom.py new file mode 100644 index 0000000..2b1b5c5 --- /dev/null +++ b/stix2/test/test_custom.py @@ -0,0 +1,104 @@ +import pytest + +import stix2 + + +def test_identity_custom_property(): + with pytest.raises(ValueError): + stix2.Identity( + id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c", + created="2015-12-21T19:59:11Z", + modified="2015-12-21T19:59:11Z", + name="John Smith", + identity_class="individual", + custom_properties="foobar", + ) + + identity = stix2.Identity( + id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c", + created="2015-12-21T19:59:11Z", + modified="2015-12-21T19:59:11Z", + name="John Smith", + identity_class="individual", + custom_properties={ + "foo": "bar", + }, + ) + + assert identity.foo == "bar" + + +def test_identity_custom_property_invalid(): + with pytest.raises(stix2.exceptions.ExtraPropertiesError): + stix2.Identity( + id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c", + created="2015-12-21T19:59:11Z", + modified="2015-12-21T19:59:11Z", + name="John Smith", + identity_class="individual", + x_foo="bar", + ) + + +def test_identity_custom_property_allowed(): + identity = stix2.Identity( + id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c", + created="2015-12-21T19:59:11Z", + modified="2015-12-21T19:59:11Z", + name="John Smith", + identity_class="individual", + x_foo="bar", + allow_custom=True, + ) + assert identity.x_foo == "bar" + + +@pytest.mark.parametrize("data", [ + """{ + "type": "identity", + "id": "identity--311b2d2d-f010-5473-83ec-1edf84858f4c", + "created": "2015-12-21T19:59:11Z", + "modified": "2015-12-21T19:59:11Z", + "name": "John Smith", + "identity_class": "individual", + "foo": "bar" + }""", +]) +def test_parse_identity_custom_property(data): + with pytest.raises(stix2.exceptions.ExtraPropertiesError): + identity = stix2.parse(data) + + identity = stix2.parse(data, allow_custom=True) + assert identity.foo == "bar" + + +@stix2.sdo.CustomObject('new-type', { + 'property1': stix2.properties.StringProperty(required=True), + 'property2': stix2.properties.IntegerProperty(), +}) +class NewType(): + def __init__(self, property2=None, **kwargs): + if property2 < 10: + raise ValueError("'property2' is too small.") + + +def test_custom_object_type(): + nt = NewType(property1='something') + assert nt.property1 == 'something' + + with pytest.raises(stix2.exceptions.MissingPropertiesError): + NewType(property2=42) + + with pytest.raises(ValueError): + NewType(property2=4) + + +def test_parse_custom_object_type(): + nt_string = """{ + "type": "new-type", + "created": "2015-12-21T19:59:11Z", + "property1": "something" + }""" + + nt = stix2.parse(nt_string) + assert nt.property1 == 'something' diff --git a/stix2/test/test_identity.py b/stix2/test/test_identity.py index edc3400..b2c166c 100644 --- a/stix2/test/test_identity.py +++ b/stix2/test/test_identity.py @@ -51,75 +51,6 @@ def test_parse_identity(data): assert identity.name == "John Smith" -def test_identity_custom_property(): - with pytest.raises(ValueError): - stix2.Identity( - id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c", - created="2015-12-21T19:59:11Z", - modified="2015-12-21T19:59:11Z", - name="John Smith", - identity_class="individual", - custom_properties="foobar", - ) - - identity = stix2.Identity( - id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c", - created="2015-12-21T19:59:11Z", - modified="2015-12-21T19:59:11Z", - name="John Smith", - identity_class="individual", - custom_properties={ - "foo": "bar", - }, - ) - - assert identity.foo == "bar" - - -def test_identity_custom_property_invalid(): - with pytest.raises(stix2.exceptions.ExtraPropertiesError): - stix2.Identity( - id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c", - created="2015-12-21T19:59:11Z", - modified="2015-12-21T19:59:11Z", - name="John Smith", - identity_class="individual", - x_foo="bar", - ) - - -def test_identity_custom_property_allowed(): - identity = stix2.Identity( - id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c", - created="2015-12-21T19:59:11Z", - modified="2015-12-21T19:59:11Z", - name="John Smith", - identity_class="individual", - x_foo="bar", - allow_custom=True, - ) - assert identity.x_foo == "bar" - - -@pytest.mark.parametrize("data", [ - """{ - "type": "identity", - "id": "identity--311b2d2d-f010-5473-83ec-1edf84858f4c", - "created": "2015-12-21T19:59:11Z", - "modified": "2015-12-21T19:59:11Z", - "name": "John Smith", - "identity_class": "individual", - "foo": "bar" - }""", -]) -def test_parse_identity_custom_property(data): - with pytest.raises(stix2.exceptions.ExtraPropertiesError): - identity = stix2.parse(data) - - identity = stix2.parse(data, allow_custom=True) - assert identity.foo == "bar" - - def test_parse_no_type(): with pytest.raises(stix2.exceptions.ParseError): stix2.parse("""