Merge pull request #74 from oasis-open/custom-props-in-bundled-objs

Allow objects with custom properties in bundles

stix2/core.py

@@ -15,6 +15,10 @@ from .utils import get_dict
 
 class STIXObjectProperty(Property):
 
+    def __init__(self, allow_custom=False):
+        self.allow_custom = allow_custom
+        super(STIXObjectProperty, self).__init__()
+
     def clean(self, value):
         try:
             dictified = get_dict(value)
@@ -25,7 +29,10 @@ class STIXObjectProperty(Property):
         if 'type' in dictified and dictified['type'] == 'bundle':
             raise ValueError('This property may not contain a Bundle object')
 
-        parsed_obj = parse(dictified)
+        if self.allow_custom:
+            parsed_obj = parse(dictified, allow_custom=True)
+        else:
+            parsed_obj = parse(dictified)
         return parsed_obj
 
 
@@ -48,6 +55,10 @@ class Bundle(_STIXBase):
             else:
                 kwargs['objects'] = list(args) + kwargs.get('objects', [])
 
+        allow_custom = kwargs.get('allow_custom', False)
+        if allow_custom:
+            self._properties['objects'] = ListProperty(STIXObjectProperty(True))
+
         super(Bundle, self).__init__(**kwargs)
 
 

stix2/properties.py

@@ -124,7 +124,11 @@ class ListProperty(Property):
                 obj_type = self.contained.type
             elif type(self.contained).__name__ is 'STIXObjectProperty':
                 # ^ this way of checking doesn't require a circular import
-                obj_type = type(valid)
+                # valid is already an instance of a python-stix2 class; no need
+                # to turn it into a dictionary and then pass it to the class
+                # constructor again
+                result.append(valid)
+                continue
             else:
                 obj_type = self.contained
 

stix2/test/test_custom.py

@@ -81,6 +81,18 @@ def test_parse_identity_custom_property(data):
     assert identity.foo == "bar"
 
 
+def test_custom_property_in_bundled_object():
+    identity = stix2.Identity(
+        name="John Smith",
+        identity_class="individual",
+        x_foo="bar",
+        allow_custom=True,
+    )
+    bundle = stix2.Bundle(identity, allow_custom=True)
+
+    assert bundle.objects[0].x_foo == "bar"
+
+
 @stix2.sdo.CustomObject('x-new-type', [
     ('property1', stix2.properties.StringProperty(required=True)),
     ('property2', stix2.properties.IntegerProperty()),