diff --git a/docs/guide/custom.ipynb b/docs/guide/custom.ipynb index bc19749..e651084 100644 --- a/docs/guide/custom.ipynb +++ b/docs/guide/custom.ipynb @@ -1155,10 +1155,10 @@ "source": [ "from stix2 import File, CustomExtension\n", "\n", - "@CustomExtension(File, 'x-new-ext', {\n", - " 'property1': properties.StringProperty(required=True),\n", - " 'property2': properties.IntegerProperty(),\n", - "})\n", + "@CustomExtension(File, 'x-new-ext', [\n", + " ('property1', properties.StringProperty(required=True)),\n", + " ('property2', properties.IntegerProperty()),\n", + "])\n", "class NewExtension():\n", " pass\n", "\n", diff --git a/stix2/test/test_custom.py b/stix2/test/test_custom.py index cc8b32b..a14503f 100644 --- a/stix2/test/test_custom.py +++ b/stix2/test/test_custom.py @@ -435,10 +435,10 @@ def test_observed_data_with_custom_observable_object(): assert ob_data.objects['0'].property1 == 'something' -@stix2.observables.CustomExtension(stix2.DomainName, 'x-new-ext', { - 'property1': stix2.properties.StringProperty(required=True), - 'property2': stix2.properties.IntegerProperty(), -}) +@stix2.observables.CustomExtension(stix2.DomainName, 'x-new-ext', [ + ('property1', stix2.properties.StringProperty(required=True)), + ('property2', stix2.properties.IntegerProperty()), +]) class NewExtension(): def __init__(self, property2=None, **kwargs): if property2 and property2 < 10: @@ -485,9 +485,9 @@ def test_custom_extension_invalid_observable(): class Foo(object): pass with pytest.raises(ValueError) as excinfo: - @stix2.observables.CustomExtension(Foo, 'x-new-ext', { - 'property1': stix2.properties.StringProperty(required=True), - }) + @stix2.observables.CustomExtension(Foo, 'x-new-ext', [ + ('property1', stix2.properties.StringProperty(required=True)), + ]) class FooExtension(): pass # pragma: no cover assert str(excinfo.value) == "'observable' must be a valid Observable class!" @@ -495,9 +495,9 @@ def test_custom_extension_invalid_observable(): class Bar(stix2.observables._Observable): pass with pytest.raises(ValueError) as excinfo: - @stix2.observables.CustomExtension(Bar, 'x-new-ext', { - 'property1': stix2.properties.StringProperty(required=True), - }) + @stix2.observables.CustomExtension(Bar, 'x-new-ext', [ + ('property1', stix2.properties.StringProperty(required=True)), + ]) class BarExtension(): pass assert "Unknown observable type" in str(excinfo.value) @@ -506,9 +506,9 @@ def test_custom_extension_invalid_observable(): class Baz(stix2.observables._Observable): _type = 'Baz' with pytest.raises(ValueError) as excinfo: - @stix2.observables.CustomExtension(Baz, 'x-new-ext', { - 'property1': stix2.properties.StringProperty(required=True), - }) + @stix2.observables.CustomExtension(Baz, 'x-new-ext', [ + ('property1', stix2.properties.StringProperty(required=True)), + ]) class BazExtension(): pass assert "Unknown observable type" in str(excinfo.value) @@ -520,21 +520,29 @@ def test_custom_extension_no_properties(): @stix2.observables.CustomExtension(stix2.DomainName, 'x-new-ext2', None) class BarExtension(): pass - assert "'properties' must be a dict!" in str(excinfo.value) + assert "Must supply a list, containing tuples." in str(excinfo.value) def test_custom_extension_empty_properties(): + with pytest.raises(ValueError) as excinfo: + @stix2.observables.CustomExtension(stix2.DomainName, 'x-new-ext2', []) + class BarExtension(): + pass + assert "Must supply a list, containing tuples." in str(excinfo.value) + + +def test_custom_extension_dict_properties(): with pytest.raises(ValueError) as excinfo: @stix2.observables.CustomExtension(stix2.DomainName, 'x-new-ext2', {}) class BarExtension(): pass - assert "'properties' must be a dict!" in str(excinfo.value) + assert "Must supply a list, containing tuples." in str(excinfo.value) def test_custom_extension_no_init_1(): - @stix2.observables.CustomExtension(stix2.DomainName, 'x-new-extension', { - 'property1': stix2.properties.StringProperty(required=True), - }) + @stix2.observables.CustomExtension(stix2.DomainName, 'x-new-extension', [ + ('property1', stix2.properties.StringProperty(required=True)), + ]) class NewExt(): pass @@ -543,9 +551,9 @@ def test_custom_extension_no_init_1(): def test_custom_extension_no_init_2(): - @stix2.observables.CustomExtension(stix2.DomainName, 'x-new-ext2', { - 'property1': stix2.properties.StringProperty(required=True), - }) + @stix2.observables.CustomExtension(stix2.DomainName, 'x-new-ext2', [ + ('property1', stix2.properties.StringProperty(required=True)), + ]) class NewExt2(object): pass @@ -594,3 +602,8 @@ def test_register_custom_object(): stix2._register_type(CustomObject2) # Note that we will always check against newest OBJ_MAP. assert (CustomObject2._type, CustomObject2) in stix2.OBJ_MAP.items() + + +def test_extension_property_location(): + assert 'extensions' in stix2.v20.observables.OBJ_MAP_OBSERVABLE['x-new-observable']._properties + assert 'extensions' not in stix2.v20.observables.EXT_MAP['domain-name']['x-new-ext']._properties diff --git a/stix2/v20/observables.py b/stix2/v20/observables.py index 83600b0..39a8f19 100644 --- a/stix2/v20/observables.py +++ b/stix2/v20/observables.py @@ -979,6 +979,9 @@ def CustomObservable(type='x-custom-observable', properties=None): "is not a ListProperty containing ObjectReferenceProperty." % prop_name) _properties.update(properties) + _properties.update([ + ('extensions', ExtensionsProperty(enclosing_type=_type)), + ]) def __init__(self, **kwargs): _Observable.__init__(self, **kwargs) @@ -1030,12 +1033,10 @@ def CustomExtension(observable=None, type='x-custom-observable', properties=None class _Custom(cls, _Extension): _type = type - _properties = { - 'extensions': ExtensionsProperty(enclosing_type=_type), - } + _properties = OrderedDict() - if not isinstance(properties, dict) or not properties: - raise ValueError("'properties' must be a dict!") + if not properties or not isinstance(properties, list): + raise ValueError("Must supply a list, containing tuples. For example, [('property1', IntegerProperty())]") _properties.update(properties)