Fix handling of custom extensions: make sure when
allow_custom=True that you never get a half-cleaned property value.master
parent
fb834d83f6
commit
c212c7c678
|
@ -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(
|
||||||
|
|
|
@ -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
|
||||||
|
|
|
@ -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
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue