From c7b48402329d10b25522c95fa7b717ee55b81fff Mon Sep 17 00:00:00 2001 From: Michael Chisholm Date: Fri, 25 Jun 2021 22:20:02 -0400 Subject: [PATCH] Code style and behavioral improvements in ExtensionsProperty: - Changed an error message "Cannot determine extension type." At that point in the code, we in fact have a registered extension type and the class for it, so it didn't seem to make any sense. I changed it to say that an extension of that type couldn't be created from type of value given. Because this is really about typing (in both the old and new code), I also changed the exception class to TypeError. - Changed customization behavior: unregistered "extension-definition--" style extensions are no longer considered custom. Their mere presence can no longer be considered a customization, since it doesn't fit the criteria, and in fact the extension mechanism is supposed to supplant the old customization system anyway. Unregistered extensions of other types are still considered custom. - Fixed up unit tests according to the exception message change. - Added a unit test for unregistered "extension-definition--" style extensions. --- stix2/properties.py | 39 +++++++++++++++++++++++------------ stix2/test/v20/test_custom.py | 2 +- stix2/test/v21/test_custom.py | 37 ++++++++++++++++++++++++++++++++- 3 files changed, 63 insertions(+), 15 deletions(-) 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):