Merge pull request #286 from chisholm/custom_content_error_fix

Fix handling of custom extensions: uncleaned property values

stix2/base.py

@@ -8,10 +8,9 @@ import simplejson as json
 import six
 
 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 .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(

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
 
 

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

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