diff --git a/stix2/properties.py b/stix2/properties.py index 0f68f6a..e63206a 100644 --- a/stix2/properties.py +++ b/stix2/properties.py @@ -14,7 +14,7 @@ import stix2.hashes from .base import _STIXBase from .exceptions import CustomContentError, DictionaryKeyError, STIXError from .parsing import parse, parse_observable -from .registry import STIX2_OBJ_MAPS +from .registry import class_for_type from .utils import ( STIXTypeClass, _get_dict, get_class_hierarchy_names, get_type_from_id, is_object, is_stix_type, parse_into_datetime, to_enum, @@ -781,17 +781,21 @@ class ExtensionsProperty(DictionaryProperty): raise ValueError("The extensions property must contain a dictionary") has_custom = False - extension_type_map = STIX2_OBJ_MAPS[self.spec_version].get('extensions', {}) for key, subvalue in dictified.items(): - if key in extension_type_map: - cls = extension_type_map[key] - if type(subvalue) is dict: + cls = class_for_type(key, self.spec_version, "extensions") + if cls: + if isinstance(subvalue, dict): ext = cls(allow_custom=allow_custom, **subvalue) - elif type(subvalue) is cls: - # If already an instance of an _Extension class, assume it's valid + elif isinstance(subvalue, cls): + # If already an instance of the registered class, assume + # it's valid ext = subvalue else: - raise ValueError("Cannot determine extension type.") + raise TypeError( + "Can't create extension '{}' from {}.".format( + key, type(subvalue) + ) + ) has_custom = has_custom or ext.has_custom @@ -805,15 +809,24 @@ class ExtensionsProperty(DictionaryProperty): dictified[key] = ext else: - if allow_custom: + # If an unregistered "extension-definition--" style extension, + # we don't know what's supposed to be in it, so we can't + # determine whether there's anything custom. So, assume there + # are no customizations. If it's a different type of extension, + # non-registration implies customization (since all spec-defined + # extensions should be pre-registered with the library). + + if key.startswith('extension-definition--'): + _validate_id( + key, self.spec_version, 'extension-definition--' + ) + elif allow_custom: has_custom = True - dictified[key] = subvalue - elif key.startswith('extension-definition--'): - _validate_id(key, '2.1', 'extension-definition') - dictified[key] = subvalue else: raise CustomContentError("Can't parse unknown extension type: {}".format(key)) + dictified[key] = subvalue + return dictified, has_custom diff --git a/stix2/test/v20/test_custom.py b/stix2/test/v20/test_custom.py index fd45ff4..4647167 100644 --- a/stix2/test/v20/test_custom.py +++ b/stix2/test/v20/test_custom.py @@ -795,7 +795,7 @@ def test_custom_extension_wrong_observable_type(): }, ) - assert 'Cannot determine extension type' in excinfo.value.reason + assert "Can't create extension 'ntfs-ext' from" in excinfo.value.reason @pytest.mark.parametrize( diff --git a/stix2/test/v21/test_custom.py b/stix2/test/v21/test_custom.py index 4f60129..ff0464b 100644 --- a/stix2/test/v21/test_custom.py +++ b/stix2/test/v21/test_custom.py @@ -984,7 +984,7 @@ def test_custom_extension_wrong_observable_type(): }, ) - assert 'Cannot determine extension type' in excinfo.value.reason + assert "Can't create extension 'ntfs-ext'" in excinfo.value.reason @pytest.mark.parametrize( @@ -1228,6 +1228,41 @@ def test_parse_observable_with_unregistered_custom_extension(data): assert not isinstance(parsed_ob['extensions']['x-foobar-ext'], stix2.base._STIXBase) +def test_unregistered_new_style_extension(): + + f_dict = { + "type": "file", + "name": "foo.txt", + "extensions": { + "extension-definition--31adb724-a9a4-44b6-8ec2-fd4b181c9507": { + "extension-type": "property-extension", + "a": 1, + "b": True + } + } + } + + f = stix2.parse(f_dict, allow_custom=False) + + assert f.extensions[ + "extension-definition--31adb724-a9a4-44b6-8ec2-fd4b181c9507" + ]["a"] == 1 + assert f.extensions[ + "extension-definition--31adb724-a9a4-44b6-8ec2-fd4b181c9507" + ]["b"] + assert not f.has_custom + + f = stix2.parse(f_dict, allow_custom=True) + + assert f.extensions[ + "extension-definition--31adb724-a9a4-44b6-8ec2-fd4b181c9507" + ]["a"] == 1 + assert f.extensions[ + "extension-definition--31adb724-a9a4-44b6-8ec2-fd4b181c9507" + ]["b"] + assert not f.has_custom + + def test_register_custom_object(): # Not the way to register custom object. class CustomObject2(object):