From 0e658255a8d0bfa670a0fec8ab8588ec6b51dfd0 Mon Sep 17 00:00:00 2001
From: clenk <clenk@mitre.org>
Date: Wed, 30 Aug 2017 15:33:28 -0400
Subject: [PATCH] Move check to when custom observable is defined

Catch properties named "_ref" or "_refs" that aren't
ObjectReferenceProperty when the custom observable is defined, not when
it is instantiated.
---
 stix2/base.py             | 16 ++++++----------
 stix2/observables.py      | 10 ++++++++++
 stix2/test/test_custom.py | 35 +++++++++++++++++++++++++----------
 3 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/stix2/base.py b/stix2/base.py
index 7de193b..185ff35 100644
--- a/stix2/base.py
+++ b/stix2/base.py
@@ -205,18 +205,14 @@ class _Observable(_STIXBase):
         try:
             allowed_types = prop.contained.valid_types
         except AttributeError:
-            try:
-                allowed_types = prop.valid_types
-            except AttributeError:
-                raise ValueError("'%s' is named like an object reference property but "
-                                 "is not an ObjectReferenceProperty or a ListProperty "
-                                 "containing ObjectReferenceProperty." % prop_name)
+            allowed_types = prop.valid_types
+
+        try:
+            ref_type = self._STIXBase__valid_refs[ref]
+        except TypeError:
+            raise ValueError("'%s' must be created with _valid_refs as a dict, not a list." % self.__class__.__name__)
 
         if allowed_types:
-            try:
-                ref_type = self._STIXBase__valid_refs[ref]
-            except TypeError:
-                raise ValueError("'%s' must be created with _valid_refs as a dict, not a list." % self.__class__.__name__)
             if ref_type not in allowed_types:
                 raise InvalidObjRefError(self.__class__, prop_name, "object reference '%s' is of an invalid type '%s'" % (ref, ref_type))
 
diff --git a/stix2/observables.py b/stix2/observables.py
index 50eb6d1..4c29785 100644
--- a/stix2/observables.py
+++ b/stix2/observables.py
@@ -766,6 +766,16 @@ def CustomObservable(type='x-custom-observable', properties={}):
             _properties = {
                 'type': TypeProperty(_type),
             }
+            # Check properties ending in "_ref/s" are ObjectReferenceProperties
+            for prop_name, prop in properties.items():
+                if prop_name.endswith('_ref') and prop.__class__.__name__ != 'ObjectReferenceProperty':
+                    raise ValueError("'%s' is named like an object reference property but "
+                                     "is not an ObjectReferenceProperty." % prop_name)
+                elif (prop_name.endswith('_refs') and (prop.__class__.__name__ != 'ListProperty'
+                                                       or prop.contained.__class__.__name__ != 'ObjectReferenceProperty')):
+                    raise ValueError("'%s' is named like an object reference property but "
+                                     "is not a ListProperty containing ObjectReferenceProperty." % prop_name)
+
             _properties.update(properties)
 
             def __init__(self, **kwargs):
diff --git a/stix2/test/test_custom.py b/stix2/test/test_custom.py
index d2bb16f..313a0a8 100644
--- a/stix2/test/test_custom.py
+++ b/stix2/test/test_custom.py
@@ -153,20 +153,35 @@ def test_custom_observable_object():
 
 
 def test_custom_observable_object_invalid_ref_property():
-    @stix2.observables.CustomObservable('x-new-obs', {
-        'property1': stix2.properties.StringProperty(required=True),
-        'property_ref': stix2.properties.StringProperty(),
-    })
-    class NewObs():
-        pass
-
     with pytest.raises(ValueError) as excinfo:
-        NewObs(_valid_refs={'1': 'file'},
-               property1='something',
-               property_ref='1')
+        @stix2.observables.CustomObservable('x-new-obs', {
+            'property_ref': stix2.properties.StringProperty(),
+        })
+        class NewObs():
+            pass
     assert "is named like an object reference property but is not an ObjectReferenceProperty" in str(excinfo.value)
 
 
+def test_custom_observable_object_invalid_refs_property():
+    with pytest.raises(ValueError) as excinfo:
+        @stix2.observables.CustomObservable('x-new-obs', {
+            'property_refs': stix2.properties.StringProperty(),
+        })
+        class NewObs():
+            pass
+    assert "is named like an object reference property but is not a ListProperty containing ObjectReferenceProperty" in str(excinfo.value)
+
+
+def test_custom_observable_object_invalid_refs_list_property():
+    with pytest.raises(ValueError) as excinfo:
+        @stix2.observables.CustomObservable('x-new-obs', {
+            'property_refs': stix2.properties.ListProperty(stix2.properties.StringProperty),
+        })
+        class NewObs():
+            pass
+    assert "is named like an object reference property but is not a ListProperty containing ObjectReferenceProperty" in str(excinfo.value)
+
+
 def test_custom_observable_object_invalid_valid_refs():
     @stix2.observables.CustomObservable('x-new-obs', {
         'property1': stix2.properties.StringProperty(required=True),