Merge pull request #180 from oasis-open/179-allow-custom

Fix 2 more allow custom bugs

stix2/base.py

@@ -135,10 +135,12 @@ class _STIXBase(collections.Mapping):
         custom_props = kwargs.pop('custom_properties', {})
         if custom_props and not isinstance(custom_props, dict):
             raise ValueError("'custom_properties' must be a dictionary")
-        if not allow_custom:
-            extra_kwargs = list(set(kwargs) - set(cls._properties))
+        if not self.__allow_custom:
+            extra_kwargs = list(set(kwargs) - set(self._properties))
             if extra_kwargs:
                 raise ExtraPropertiesError(cls, extra_kwargs)
+        if custom_props:
+            self.__allow_custom = True
 
         # Remove any keyword arguments whose value is None
         setting_kwargs = {}
@@ -149,17 +151,17 @@ class _STIXBase(collections.Mapping):
                 setting_kwargs[prop_name] = prop_value
 
         # Detect any missing required properties
-        required_properties = set(get_required_properties(cls._properties))
+        required_properties = set(get_required_properties(self._properties))
         missing_kwargs = required_properties - set(setting_kwargs)
         if missing_kwargs:
             raise MissingPropertiesError(cls, missing_kwargs)
 
-        for prop_name, prop_metadata in cls._properties.items():
+        for prop_name, prop_metadata in self._properties.items():
             self._check_property(prop_name, prop_metadata, setting_kwargs)
 
         # Cache defaulted optional properties for serialization
         defaulted = []
-        for name, prop in cls._properties.items():
+        for name, prop in self._properties.items():
             try:
                 if (not prop.required and not hasattr(prop, '_fixed_value') and
                         prop.default() == setting_kwargs[name]):

stix2/core.py

@@ -63,7 +63,7 @@ class Bundle(_STIXBase):
                 kwargs['objects'] = list(args) + kwargs.get('objects', [])
 
         self.__allow_custom = kwargs.get('allow_custom', False)
-        self._properties['objects'].allow_custom = kwargs.get('allow_custom', False)
+        self._properties['objects'].contained.allow_custom = kwargs.get('allow_custom', False)
 
         super(Bundle, self).__init__(**kwargs)
 

stix2/test/test_custom.py

@@ -24,6 +24,20 @@ def test_identity_custom_property():
         )
     assert str(excinfo.value) == "'custom_properties' must be a dictionary"
 
+    with pytest.raises(stix2.exceptions.ExtraPropertiesError) as excinfo:
+        stix2.Identity(
+            id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
+            created="2015-12-21T19:59:11Z",
+            modified="2015-12-21T19:59:11Z",
+            name="John Smith",
+            identity_class="individual",
+            custom_properties={
+                "foo": "bar",
+            },
+            foo="bar",
+            )
+    assert "Unexpected properties for Identity" in str(excinfo.value)
+
     identity = stix2.Identity(
         id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
         created="2015-12-21T19:59:11Z",
@@ -34,7 +48,6 @@ def test_identity_custom_property():
             "foo": "bar",
         },
     )
-
     assert identity.foo == "bar"
 
 
@@ -88,13 +101,61 @@ def test_parse_identity_custom_property(data):
     assert identity.foo == "bar"
 
 
-def test_custom_property_in_bundled_object():
+def test_custom_property_object_in_bundled_object():
     bundle = stix2.Bundle(IDENTITY_CUSTOM_PROP, allow_custom=True)
 
     assert bundle.objects[0].x_foo == "bar"
     assert '"x_foo": "bar"' in str(bundle)
 
 
+def test_custom_properties_object_in_bundled_object():
+    obj = stix2.Identity(
+        name="John Smith",
+        identity_class="individual",
+        custom_properties={
+            "x_foo": "bar",
+        }
+    )
+    bundle = stix2.Bundle(obj, allow_custom=True)
+
+    assert bundle.objects[0].x_foo == "bar"
+    assert '"x_foo": "bar"' in str(bundle)
+
+
+def test_custom_property_dict_in_bundled_object():
+    custom_identity = {
+        'type': 'identity',
+        'id': 'identity--311b2d2d-f010-5473-83ec-1edf84858f4c',
+        'created': '2015-12-21T19:59:11Z',
+        'name': 'John Smith',
+        'identity_class': 'individual',
+        'x_foo': 'bar',
+    }
+    with pytest.raises(stix2.exceptions.ExtraPropertiesError):
+        bundle = stix2.Bundle(custom_identity)
+
+    bundle = stix2.Bundle(custom_identity, allow_custom=True)
+    assert bundle.objects[0].x_foo == "bar"
+    assert '"x_foo": "bar"' in str(bundle)
+
+
+def test_custom_properties_dict_in_bundled_object():
+    custom_identity = {
+        'type': 'identity',
+        'id': 'identity--311b2d2d-f010-5473-83ec-1edf84858f4c',
+        'created': '2015-12-21T19:59:11Z',
+        'name': 'John Smith',
+        'identity_class': 'individual',
+        'custom_properties': {
+            'x_foo': 'bar',
+        },
+    }
+    bundle = stix2.Bundle(custom_identity)
+
+    assert bundle.objects[0].x_foo == "bar"
+    assert '"x_foo": "bar"' in str(bundle)
+
+
 def test_custom_property_in_observed_data():
     artifact = stix2.File(
         allow_custom=True,