From 9036c7f7b8b903d27874769589ce4c2fc813d0a5 Mon Sep 17 00:00:00 2001
From: clenk <clenk@mitre.org>
Date: Fri, 9 Jun 2017 12:20:40 -0400
Subject: [PATCH] Add custom properties

---
 stix2/base.py               |  7 ++++++-
 stix2/test/test_identity.py | 31 +++++++++++++++++++++++++++++--
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/stix2/base.py b/stix2/base.py
index c40c9e6..d3cb072 100644
--- a/stix2/base.py
+++ b/stix2/base.py
@@ -90,13 +90,18 @@ class _STIXBase(collections.Mapping):
         self.__now = get_timestamp()
 
         # Detect any keyword arguments not allowed for a specific type
+        custom_props = kwargs.pop('custom_properties', {})
+        if custom_props and not isinstance(custom_props, dict):
+            raise ValueError("'custom_properties' must be a dictionary")
         extra_kwargs = list(set(kwargs) - set(cls._properties))
         if extra_kwargs:
             raise ExtraPropertiesError(cls, extra_kwargs)
 
         # Remove any keyword arguments whose value is None
         setting_kwargs = {}
-        for prop_name, prop_value in kwargs.items():
+        props = kwargs.copy()
+        props.update(custom_props)
+        for prop_name, prop_value in props.items():
             if prop_value is not None:
                 setting_kwargs[prop_name] = prop_value
 
diff --git a/stix2/test/test_identity.py b/stix2/test/test_identity.py
index 3952d25..e4a67e6 100644
--- a/stix2/test/test_identity.py
+++ b/stix2/test/test_identity.py
@@ -19,7 +19,7 @@ EXPECTED = """{
 
 
 def test_identity_example():
-    report = stix2.Identity(
+    identity = stix2.Identity(
         id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
         created="2015-12-21T19:59:11Z",
         modified="2015-12-21T19:59:11Z",
@@ -27,7 +27,7 @@ def test_identity_example():
         identity_class="individual",
     )
 
-    assert str(report) == EXPECTED
+    assert str(identity) == EXPECTED
 
 
 @pytest.mark.parametrize("data", [
@@ -50,4 +50,31 @@ def test_parse_identity(data):
     assert identity.modified == dt.datetime(2015, 12, 21, 19, 59, 11, tzinfo=pytz.utc)
     assert identity.name == "John Smith"
 
+
+def test_identity_custom_property():
+    identity = stix2.Identity(
+        id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
+        created="2015-12-21T19:59:11Z",
+        modified="2015-12-21T19:59:11Z",
+        name="John Smith",
+        identity_class="individual",
+        custom_properties={
+            "foo": "bar",
+        },
+    )
+
+    assert identity.foo == "bar"
+
+
+def test_identity_custom_property_invalid():
+    with pytest.raises(stix2.exceptions.ExtraPropertiesError):
+        stix2.Identity(
+            id="identity--311b2d2d-f010-5473-83ec-1edf84858f4c",
+            created="2015-12-21T19:59:11Z",
+            modified="2015-12-21T19:59:11Z",
+            name="John Smith",
+            identity_class="individual",
+            x_foo="bar",
+        )
+
 # TODO: Add other examples