From 2b6023e7bb98aa8276882d00df55255bac016ef2 Mon Sep 17 00:00:00 2001
From: Chris Lenk <clenk@mitre.org>
Date: Mon, 9 Oct 2017 17:33:12 -0400
Subject: [PATCH] Allow objects with custom properties in bundles

---
 stix2/core.py             | 13 ++++++++++++-
 stix2/properties.py       |  6 +++++-
 stix2/test/test_custom.py | 12 ++++++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/stix2/core.py b/stix2/core.py
index 3eaabb0..8ee11f5 100644
--- a/stix2/core.py
+++ b/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)
 
 
diff --git a/stix2/properties.py b/stix2/properties.py
index afe994f..ca7f04c 100644
--- a/stix2/properties.py
+++ b/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
 
diff --git a/stix2/test/test_custom.py b/stix2/test/test_custom.py
index ff432c1..48529b9 100644
--- a/stix2/test/test_custom.py
+++ b/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()),