From 79475586d86ec6741d651db0dff6a556b992064e Mon Sep 17 00:00:00 2001
From: Emmanuelle Vargas-Gonzalez <emmanuelle@mitre.org>
Date: Fri, 3 Nov 2017 14:17:36 -0400
Subject: [PATCH] Revert object_properties() to #85 fix. Update tests
 accordingly

---
 stix2/base.py                   | 16 +++++++--------
 stix2/test/test_data_sources.py | 35 +++++++++++++++++++++++++--------
 stix2/test/test_filesystem.py   | 10 +++++-----
 stix2/test/test_identity.py     | 11 +++++++++++
 4 files changed, 51 insertions(+), 21 deletions(-)

diff --git a/stix2/base.py b/stix2/base.py
index d7c7553..76b07b8 100644
--- a/stix2/base.py
+++ b/stix2/base.py
@@ -40,7 +40,14 @@ class _STIXBase(collections.Mapping):
     """Base class for STIX object types"""
 
     def object_properties(self):
-        return list(self._properties.keys())
+        props = set(self._properties.keys())
+        custom_props = list(set(self._inner.keys()) - props)
+        custom_props.sort()
+
+        all_properties = list(self._properties.keys())
+        all_properties.extend(custom_props)  # Any custom properties to the bottom
+
+        return all_properties
 
     def _check_property(self, prop_name, prop, kwargs):
         if prop_name not in kwargs:
@@ -102,13 +109,6 @@ class _STIXBase(collections.Mapping):
             extra_kwargs = list(set(kwargs) - set(cls._properties))
             if extra_kwargs:
                 raise ExtraPropertiesError(cls, extra_kwargs)
-        else:
-            from .properties import CustomProperty
-
-            # The custom properties will get added to the bottom.
-            # Matched with a CustomProperty.
-            extra_kwargs = list(set(kwargs) - set(cls._properties))
-            self._properties.update([(x, CustomProperty()) for x in extra_kwargs])
 
         # Remove any keyword arguments whose value is None
         setting_kwargs = {}
diff --git a/stix2/test/test_data_sources.py b/stix2/test/test_data_sources.py
index d798a0c..514ba40 100644
--- a/stix2/test/test_data_sources.py
+++ b/stix2/test/test_data_sources.py
@@ -498,29 +498,48 @@ def test_composite_datasource_operations():
                    objects=STIX_OBJS1,
                    spec_version="2.0",
                    type="bundle")
-    cds = CompositeDataSource()
-    ds1 = MemorySource(stix_data=BUNDLE1)
-    ds2 = MemorySource(stix_data=STIX_OBJS2)
+    cds1 = CompositeDataSource()
+    ds1_1 = MemorySource(stix_data=BUNDLE1)
+    ds1_2 = MemorySource(stix_data=STIX_OBJS2)
 
-    cds.add_data_sources([ds1, ds2])
+    cds2 = CompositeDataSource()
+    ds2_1 = MemorySource(stix_data=BUNDLE1)
+    ds2_2 = MemorySource(stix_data=STIX_OBJS2)
 
-    indicators = cds.all_versions("indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f")
+    cds1.add_data_sources([ds1_1, ds1_2])
+    cds2.add_data_sources([ds2_1, ds2_2])
+
+    indicators = cds1.all_versions("indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f")
 
     # In STIX_OBJS2 changed the 'modified' property to a later time...
     assert len(indicators) == 2
 
-    indicator = cds.get("indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f")
+    cds1.add_data_sources([cds2])
+
+    indicator = cds1.get("indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f")
 
     assert indicator["id"] == "indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f"
     assert indicator["modified"] == "2017-01-31T13:49:53.935Z"
     assert indicator["type"] == "indicator"
 
-    query = [
+    query1 = [
         Filter("type", "=", "indicator")
     ]
 
-    results = cds.query(query)
+    query2 = [
+        Filter("valid_from", "=", "2017-01-27T13:49:53.935382Z")
+    ]
+
+    cds1.filters.update(query2)
+
+    results = cds1.query(query1)
 
     # STIX_OBJS2 has indicator with later time, one with different id, one with
     # original time in STIX_OBJS1
     assert len(results) == 3
+
+    indicator = cds1.get("indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f")
+
+    assert indicator["id"] == "indicator--d81f86b9-975b-bc0b-775e-810c5ad45a4f"
+    assert indicator["modified"] == "2017-01-31T13:49:53.935Z"
+    assert indicator["type"] == "indicator"
diff --git a/stix2/test/test_filesystem.py b/stix2/test/test_filesystem.py
index 1610615..85f6966 100644
--- a/stix2/test/test_filesystem.py
+++ b/stix2/test/test_filesystem.py
@@ -340,7 +340,7 @@ def test_filesystem_object_with_custom_property(fs_store):
 
     fs_store.add(camp, True)
 
-    camp_r = fs_store.get(camp.id)
+    camp_r = fs_store.get(camp.id, allow_custom=True)
     assert camp_r.id == camp.id
     assert camp_r.x_empire == camp.x_empire
 
@@ -352,9 +352,9 @@ def test_filesystem_object_with_custom_property_in_bundle(fs_store):
                     allow_custom=True)
 
     bundle = Bundle(camp, allow_custom=True)
-    fs_store.add(bundle, True)
+    fs_store.add(bundle, allow_custom=True)
 
-    camp_r = fs_store.get(camp.id)
+    camp_r = fs_store.get(camp.id, allow_custom=True)
     assert camp_r.id == camp.id
     assert camp_r.x_empire == camp.x_empire
 
@@ -367,9 +367,9 @@ def test_filesystem_custom_object(fs_store):
         pass
 
     newobj = NewObj(property1='something')
-    fs_store.add(newobj, True)
+    fs_store.add(newobj, allow_custom=True)
 
-    newobj_r = fs_store.get(newobj.id)
+    newobj_r = fs_store.get(newobj.id, allow_custom=True)
     assert newobj_r.id == newobj.id
     assert newobj_r.property1 == 'something'
 
diff --git a/stix2/test/test_identity.py b/stix2/test/test_identity.py
index a9415fe..8e3dd42 100644
--- a/stix2/test/test_identity.py
+++ b/stix2/test/test_identity.py
@@ -62,4 +62,15 @@ def test_parse_no_type():
             "identity_class": "individual"
         }""")
 
+
+def test_identity_with_custom():
+    identity = stix2.Identity(
+        name="John Smith",
+        identity_class="individual",
+        custom_properties={'x_foo': 'bar'}
+    )
+
+    assert identity.x_foo == "bar"
+    assert "x_foo" in identity.object_properties()
+
 # TODO: Add other examples