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.
pull/1/head
Michael Chisholm 2021-06-25 22:20:02 -04:00
parent 70718063d3
commit c7b4840232
3 changed files with 63 additions and 15 deletions

View File

@ -14,7 +14,7 @@ import stix2.hashes
from .base import _STIXBase from .base import _STIXBase
from .exceptions import CustomContentError, DictionaryKeyError, STIXError from .exceptions import CustomContentError, DictionaryKeyError, STIXError
from .parsing import parse, parse_observable from .parsing import parse, parse_observable
from .registry import STIX2_OBJ_MAPS from .registry import class_for_type
from .utils import ( from .utils import (
STIXTypeClass, _get_dict, get_class_hierarchy_names, get_type_from_id, STIXTypeClass, _get_dict, get_class_hierarchy_names, get_type_from_id,
is_object, is_stix_type, parse_into_datetime, to_enum, 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") raise ValueError("The extensions property must contain a dictionary")
has_custom = False has_custom = False
extension_type_map = STIX2_OBJ_MAPS[self.spec_version].get('extensions', {})
for key, subvalue in dictified.items(): for key, subvalue in dictified.items():
if key in extension_type_map: cls = class_for_type(key, self.spec_version, "extensions")
cls = extension_type_map[key] if cls:
if type(subvalue) is dict: if isinstance(subvalue, dict):
ext = cls(allow_custom=allow_custom, **subvalue) ext = cls(allow_custom=allow_custom, **subvalue)
elif type(subvalue) is cls: elif isinstance(subvalue, cls):
# If already an instance of an _Extension class, assume it's valid # If already an instance of the registered class, assume
# it's valid
ext = subvalue ext = subvalue
else: 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 has_custom = has_custom or ext.has_custom
@ -805,15 +809,24 @@ class ExtensionsProperty(DictionaryProperty):
dictified[key] = ext dictified[key] = ext
else: 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 has_custom = True
dictified[key] = subvalue
elif key.startswith('extension-definition--'):
_validate_id(key, '2.1', 'extension-definition')
dictified[key] = subvalue
else: else:
raise CustomContentError("Can't parse unknown extension type: {}".format(key)) raise CustomContentError("Can't parse unknown extension type: {}".format(key))
dictified[key] = subvalue
return dictified, has_custom return dictified, has_custom

View File

@ -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( @pytest.mark.parametrize(

View File

@ -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( @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) 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(): def test_register_custom_object():
# Not the way to register custom object. # Not the way to register custom object.
class CustomObject2(object): class CustomObject2(object):