From d4e92dd81388484afb2480f89ffcfd68b76fb422 Mon Sep 17 00:00:00 2001
From: clenk <clenk@mitre.org>
Date: Tue, 13 Jun 2017 10:26:43 -0400
Subject: [PATCH] Allow adding validation to custom object types

_Custom.__init__() is a solution to recursion errors arising from using
a decorator. See https://stackoverflow.com/q/14739809/1013457
---
 stix2/sdo.py              | 33 +++++++++++++++++++++++++++++++--
 stix2/test/test_custom.py |  4 ++--
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/stix2/sdo.py b/stix2/sdo.py
index b38e5a8..905a637 100644
--- a/stix2/sdo.py
+++ b/stix2/sdo.py
@@ -195,11 +195,35 @@ class Vulnerability(_STIXBase):
 
 
 def CustomObject(type='x-custom-type', properties={}):
-    """Custom STIX Object type decorator"""
+    """Custom STIX Object type decorator
+
+    Example 1:
+
+    @CustomObject('type-name', {
+        'property1': StringProperty(required=True),
+        'property2': IntegerProperty(),
+    })
+    class MyNewObjectType():
+        pass
+
+    Supply an __init__() function to add any special validations to the custom
+    type. Don't call super().__init() though - doing so will cause an error.
+
+    Example 2:
+
+    @CustomObject('type-name', {
+        'property1': StringProperty(required=True),
+        'property2': IntegerProperty(),
+    })
+    class MyNewObjectType():
+        def __init__(self, property2=None, **kwargs):
+            if property2 and property2 < 10:
+                raise ValueError("'property2' is too small.")
+    """
 
     def custom_builder(cls):
 
-        class _Custom(_STIXBase):
+        class _Custom(cls, _STIXBase):
             _type = type
             _properties = COMMON_PROPERTIES.copy()
             _properties.update({
@@ -207,6 +231,11 @@ def CustomObject(type='x-custom-type', properties={}):
                 'type': TypeProperty(_type),
             })
             _properties.update(properties)
+
+            def __init__(self, **kwargs):
+                _STIXBase.__init__(self, **kwargs)
+                cls.__init__(self, **kwargs)
+
         stix2._register_type(_Custom)
         return _Custom
 
diff --git a/stix2/test/test_custom.py b/stix2/test/test_custom.py
index 2b1b5c5..30c26c6 100644
--- a/stix2/test/test_custom.py
+++ b/stix2/test/test_custom.py
@@ -78,7 +78,7 @@ def test_parse_identity_custom_property(data):
 })
 class NewType():
     def __init__(self, property2=None, **kwargs):
-        if property2 < 10:
+        if property2 and property2 < 10:
             raise ValueError("'property2' is too small.")
 
 
@@ -90,7 +90,7 @@ def test_custom_object_type():
         NewType(property2=42)
 
     with pytest.raises(ValueError):
-        NewType(property2=4)
+        NewType(property1='something', property2=4)
 
 
 def test_parse_custom_object_type():