Repository: oasis-open/cti-python-stix2

Issue:
Calling add_markings on a existing STIX object with custom properties causes an exception
Here is the stack trace:

Traceback (most recent call last):
 ```
 File "/Users/rpiazza/projects/capec-attack/capec2stix/capec2stix.py", line 197, in <module>
    main()
  File "/Users/rpiazza/projects/capec-attack/capec2stix/capec2stix.py", line 189, in main
    tlos = convert2stix(attack_pattern_catalog, marking_def.id, ident.id)
  File "/Users/rpiazza/projects/capec-attack/capec2stix/capec2stix.py", line 170, in convert2stix
    return convert_attack_pattern(attack_pattern_catalog.Attack_Patterns.Attack_Pattern[0],  marking_def_id, ident_id)
  File "/Users/rpiazza/projects/capec-attack/capec2stix/capec2stix.py", line 158, in convert_attack_pattern
    ap.add_markings(marking_def_id)
  File "/Users/rpiazza/py-envs/python3.5/lib/python3.5/site-packages/stix2/markings/__init__.py", line 143, in add_markings
    return object_markings.add_markings(obj, marking)
  File "/Users/rpiazza/py-envs/python3.5/lib/python3.5/site-packages/stix2/markings/object_markings.py", line 40, in add_markings
    return new_version(obj, object_marking_refs=list(object_markings))
  File "/Users/rpiazza/py-envs/python3.5/lib/python3.5/site-packages/stix2/utils.py", line 234, in new_version
    return cls(**{k: v for k, v in new_obj_inner.items() if v is not None})
  File "/Users/rpiazza/py-envs/python3.5/lib/python3.5/site-packages/stix2/base.py", line 111, in __init__
    raise ExtraPropertiesError(cls, extra_kwargs)
stix2.exceptions.ExtraPropertiesError: Unexpected properties for AttackPattern: (x_capec_abstraction, x_capec_consequences, x_capec_example_instances, x_capec_likelihood_of_attack, x_capec_prerequisites, x_capec_skills_required, x_capec_typical_severity, x_resources_required).
```

Discussion:

Patch:
diff --git a/stix2/markings/granular_markings.py b/stix2/markings/granular_markings.py
index be5d258..7c227d9 100644
--- a/stix2/markings/granular_markings.py
+++ b/stix2/markings/granular_markings.py
@@ -116,9 +116,9 @@ def remove_markings(obj, marking, selectors):
     granular_markings = utils.compress_markings(granular_markings)
 
     if granular_markings:
-        return new_version(obj, granular_markings=granular_markings)
+        return new_version(obj, granular_markings=granular_markings, allow_custom=True)
     else:
-        return new_version(obj, granular_markings=None)
+        return new_version(obj, granular_markings=None, allow_custom=True)
 
 
 def add_markings(obj, marking, selectors):
@@ -152,7 +152,7 @@ def add_markings(obj, marking, selectors):
 
     granular_marking = utils.expand_markings(granular_marking)
     granular_marking = utils.compress_markings(granular_marking)
-    return new_version(obj, granular_markings=granular_marking)
+    return new_version(obj, granular_markings=granular_marking, allow_custom=True)
 
 
 def clear_markings(obj, selectors):
@@ -207,9 +207,9 @@ def clear_markings(obj, selectors):
     granular_markings = utils.compress_markings(granular_markings)
 
     if granular_markings:
-        return new_version(obj, granular_markings=granular_markings)
+        return new_version(obj, granular_markings=granular_markings, allow_custom=True)
     else:
-        return new_version(obj, granular_markings=None)
+        return new_version(obj, granular_markings=None, allow_custom=True)
 
 
 def is_marked(obj, marking=None, selectors=None, inherited=False, descendants=False):
diff --git a/stix2/markings/object_markings.py b/stix2/markings/object_markings.py
index c0375c3..a169fe3 100644
--- a/stix2/markings/object_markings.py
+++ b/stix2/markings/object_markings.py
@@ -37,7 +37,7 @@ def add_markings(obj, marking):
 
     object_markings = set(obj.get("object_marking_refs", []) + marking)
 
-    return new_version(obj, object_marking_refs=list(object_markings))
+    return new_version(obj, object_marking_refs=list(object_markings), allow_custom=True)
 
 
 def remove_markings(obj, marking):
@@ -69,9 +69,9 @@ def remove_markings(obj, marking):
 
     new_markings = [x for x in object_markings if x not in marking]
     if new_markings:
-        return new_version(obj, object_marking_refs=new_markings)
+        return new_version(obj, object_marking_refs=new_markings, allow_custom=True)
     else:
-        return new_version(obj, object_marking_refs=None)
+        return new_version(obj, object_marking_refs=None, allow_custom=True)
 
 
 def set_markings(obj, marking):
@@ -103,7 +103,7 @@ def clear_markings(obj):
         A new version of the given SDO or SRO with object_marking_refs cleared.
 
     """
-    return new_version(obj, object_marking_refs=None)
+    return new_version(obj, object_marking_refs=None, allow_custom=True)
 
 
 def is_marked(obj, marking=None):
diff --git a/stix2/utils.py b/stix2/utils.py
index 73337d0..37ff166 100644
--- a/stix2/utils.py
+++ b/stix2/utils.py
@@ -251,7 +251,7 @@ def revoke(data):
 
     if data.get("revoked"):
         raise RevokeError("revoke")
-    return new_version(data, revoked=True)
+    return new_version(data, revoked=True, allow_custom=True)
 
 
 def get_class_hierarchy_names(obj):


Test patch:
diff --git a/stix2/test/test_custom.py b/stix2/test/test_custom.py
index 7c1832b..76ad61b 100644
--- a/stix2/test/test_custom.py
+++ b/stix2/test/test_custom.py
@@ -2,7 +2,14 @@ import pytest
 
 import stix2
 
-from .constants import FAKE_TIME
+from .constants import FAKE_TIME, MARKING_DEFINITION_ID
+
+IDENTITY_CUSTOM_PROP = stix2.Identity(
+    name="John Smith",
+    identity_class="individual",
+    x_foo="bar",
+    allow_custom=True,
+)
 
 
 def test_identity_custom_property():
@@ -82,18 +89,38 @@ def test_parse_identity_custom_property(data):
 
 
 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)
+    bundle = stix2.Bundle(IDENTITY_CUSTOM_PROP, allow_custom=True)
 
     assert bundle.objects[0].x_foo == "bar"
     assert '"x_foo": "bar"' in str(bundle)
 
 
+def test_identity_custom_property_revoke():
+    identity = IDENTITY_CUSTOM_PROP.revoke()
+    assert identity.x_foo == "bar"
+
+
+def test_identity_custom_property_edit_markings():
+    marking_obj = stix2.MarkingDefinition(
+        id=MARKING_DEFINITION_ID,
+        definition_type="statement",
+        definition=stix2.StatementMarking(statement="Copyright 2016, Example Corp")
+    )
+    marking_obj2 = stix2.MarkingDefinition(
+        id=MARKING_DEFINITION_ID,
+        definition_type="statement",
+        definition=stix2.StatementMarking(statement="Another one")
+    )
+
+    # None of the following should throw exceptions
+    identity = IDENTITY_CUSTOM_PROP.add_markings(marking_obj)
+    identity2 = identity.add_markings(marking_obj2, ['x_foo'])
+    identity2.remove_markings(marking_obj.id)
+    identity2.remove_markings(marking_obj2.id, ['x_foo'])
+    identity2.clear_markings()
+    identity2.clear_markings('x_foo')
+
+
 def test_custom_marking_no_init_1():
     @stix2.CustomMarking('x-new-obj', [
         ('property1', stix2.properties.StringProperty(required=True)),