diff --git a/stix2/base.py b/stix2/base.py index 9fe1617..3874721 100644 --- a/stix2/base.py +++ b/stix2/base.py @@ -8,10 +8,9 @@ import simplejson as json import six from .exceptions import ( - AtLeastOnePropertyError, CustomContentError, DependentPropertiesError, - ExtraPropertiesError, ImmutableError, InvalidObjRefError, - InvalidValueError, MissingPropertiesError, - MutuallyExclusivePropertiesError, + AtLeastOnePropertyError, DependentPropertiesError, ExtraPropertiesError, + ImmutableError, InvalidObjRefError, InvalidValueError, + MissingPropertiesError, MutuallyExclusivePropertiesError, ) from .markings.utils import validate from .utils import NOW, find_property_index, format_datetime, get_timestamp @@ -93,14 +92,6 @@ class _STIXBase(collections.Mapping): # No point in wrapping InvalidValueError in another # InvalidValueError... so let those propagate. raise - except CustomContentError as exc: - if not self.__allow_custom: - six.raise_from( - InvalidValueError( - self.__class__, prop_name, reason=str(exc), - ), - exc, - ) except Exception as exc: six.raise_from( InvalidValueError( diff --git a/stix2/properties.py b/stix2/properties.py index b9a5aff..1095156 100644 --- a/stix2/properties.py +++ b/stix2/properties.py @@ -568,7 +568,10 @@ class ExtensionsProperty(DictionaryProperty): else: raise ValueError("Cannot determine extension type.") else: - raise CustomContentError("Can't parse unknown extension type: {}".format(key)) + if self.allow_custom: + dictified[key] = subvalue + else: + raise CustomContentError("Can't parse unknown extension type: {}".format(key)) return dictified diff --git a/stix2/test/v20/test_custom.py b/stix2/test/v20/test_custom.py index 3cfc23f..6d127f2 100644 --- a/stix2/test/v20/test_custom.py +++ b/stix2/test/v20/test_custom.py @@ -1,6 +1,7 @@ import pytest import stix2 +import stix2.v20 from ...exceptions import InvalidValueError from .constants import FAKE_TIME, IDENTITY_ID, MARKING_DEFINITION_ID @@ -885,6 +886,49 @@ def test_parse_observable_with_custom_extension(): assert parsed.extensions['x-new-ext'].property2 == 12 +def test_custom_and_spec_extension_mix(): + """ + Try to make sure that when allow_custom=True, encountering a custom + extension doesn't result in a completely uncleaned extensions property. + """ + + file_obs = stix2.v20.File( + name="my_file.dat", + extensions={ + "x-custom1": { + "a": 1, + "b": 2, + }, + "ntfs-ext": { + "sid": "S-1-whatever", + }, + "x-custom2": { + "z": 99.9, + "y": False, + }, + "raster-image-ext": { + "image_height": 1024, + "image_width": 768, + "bits_per_pixel": 32, + }, + }, + allow_custom=True, + ) + + assert file_obs.extensions["x-custom1"] == {"a": 1, "b": 2} + assert file_obs.extensions["x-custom2"] == {"y": False, "z": 99.9} + assert file_obs.extensions["ntfs-ext"].sid == "S-1-whatever" + assert file_obs.extensions["raster-image-ext"].image_height == 1024 + + # Both of these should have been converted to objects, not left as dicts. + assert isinstance( + file_obs.extensions["raster-image-ext"], stix2.v20.RasterImageExt, + ) + assert isinstance( + file_obs.extensions["ntfs-ext"], stix2.v20.NTFSExt, + ) + + @pytest.mark.parametrize( "data", [ # URL is not in EXT_MAP diff --git a/stix2/test/v21/test_custom.py b/stix2/test/v21/test_custom.py index a5c9244..47d484a 100644 --- a/stix2/test/v21/test_custom.py +++ b/stix2/test/v21/test_custom.py @@ -2,6 +2,7 @@ import pytest import stix2 import stix2.base +import stix2.v21 from ...exceptions import InvalidValueError from .constants import FAKE_TIME, IDENTITY_ID, MARKING_DEFINITION_ID @@ -889,6 +890,49 @@ def test_parse_observable_with_custom_extension(): assert parsed.extensions['x-new-ext'].property2 == 12 +def test_custom_and_spec_extension_mix(): + """ + Try to make sure that when allow_custom=True, encountering a custom + extension doesn't result in a completely uncleaned extensions property. + """ + + file_obs = stix2.v21.File( + name="my_file.dat", + extensions={ + "x-custom1": { + "a": 1, + "b": 2, + }, + "ntfs-ext": { + "sid": "S-1-whatever", + }, + "x-custom2": { + "z": 99.9, + "y": False, + }, + "raster-image-ext": { + "image_height": 1024, + "image_width": 768, + "bits_per_pixel": 32, + }, + }, + allow_custom=True, + ) + + assert file_obs.extensions["x-custom1"] == {"a": 1, "b": 2} + assert file_obs.extensions["x-custom2"] == {"y": False, "z": 99.9} + assert file_obs.extensions["ntfs-ext"].sid == "S-1-whatever" + assert file_obs.extensions["raster-image-ext"].image_height == 1024 + + # Both of these should have been converted to objects, not left as dicts. + assert isinstance( + file_obs.extensions["raster-image-ext"], stix2.v21.RasterImageExt, + ) + assert isinstance( + file_obs.extensions["ntfs-ext"], stix2.v21.NTFSExt, + ) + + @pytest.mark.parametrize( "data", [ # URL is not in EXT_MAP