Fix handling of custom extensions: make sure when

allow_custom=True that you never get a half-cleaned property
value.
master
Michael Chisholm 2019-08-26 17:10:54 -04:00
parent fb834d83f6
commit c212c7c678
4 changed files with 95 additions and 13 deletions

View File

@ -8,10 +8,9 @@ import simplejson as json
import six import six
from .exceptions import ( from .exceptions import (
AtLeastOnePropertyError, CustomContentError, DependentPropertiesError, AtLeastOnePropertyError, DependentPropertiesError, ExtraPropertiesError,
ExtraPropertiesError, ImmutableError, InvalidObjRefError, ImmutableError, InvalidObjRefError, InvalidValueError,
InvalidValueError, MissingPropertiesError, MissingPropertiesError, MutuallyExclusivePropertiesError,
MutuallyExclusivePropertiesError,
) )
from .markings.utils import validate from .markings.utils import validate
from .utils import NOW, find_property_index, format_datetime, get_timestamp 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 # No point in wrapping InvalidValueError in another
# InvalidValueError... so let those propagate. # InvalidValueError... so let those propagate.
raise 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: except Exception as exc:
six.raise_from( six.raise_from(
InvalidValueError( InvalidValueError(

View File

@ -567,6 +567,9 @@ class ExtensionsProperty(DictionaryProperty):
dictified[key] = subvalue dictified[key] = subvalue
else: else:
raise ValueError("Cannot determine extension type.") raise ValueError("Cannot determine extension type.")
else:
if self.allow_custom:
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))
return dictified return dictified

View File

@ -1,6 +1,7 @@
import pytest import pytest
import stix2 import stix2
import stix2.v20
from ...exceptions import InvalidValueError from ...exceptions import InvalidValueError
from .constants import FAKE_TIME, IDENTITY_ID, MARKING_DEFINITION_ID 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 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 partially-cleaned 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( @pytest.mark.parametrize(
"data", [ "data", [
# URL is not in EXT_MAP # URL is not in EXT_MAP

View File

@ -2,6 +2,7 @@ import pytest
import stix2 import stix2
import stix2.base import stix2.base
import stix2.v21
from ...exceptions import InvalidValueError from ...exceptions import InvalidValueError
from .constants import FAKE_TIME, IDENTITY_ID, MARKING_DEFINITION_ID 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 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 partially-cleaned 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( @pytest.mark.parametrize(
"data", [ "data", [
# URL is not in EXT_MAP # URL is not in EXT_MAP