added remove_custom_stix() utility; tweaked new_version() to make it a little safer

stix2/utils.py

@@ -16,6 +16,9 @@ from .exceptions import (InvalidValueError, RevokeError,
 # timestamps in a single object, the timestamps will vary by a few microseconds.
 NOW = object()
 
+# STIX object properties that cannot be modified
+STIX_UNMOD_PROPERTIES = ["created", "created_by_ref", "id", "type"]
+
 
 class STIXdatetime(dt.datetime):
     def __new__(cls, *args, **kwargs):
@@ -215,7 +218,7 @@ def new_version(data, **kwargs):
     properties_to_change = kwargs.keys()
 
     # Make sure certain properties aren't trying to change
-    for prop in ["created", "created_by_ref", "id", "type"]:
+    for prop in STIX_UNMOD_PROPERTIES:
         if prop in properties_to_change:
             unchangable_properties.append(prop)
     if unchangable_properties:
@@ -227,8 +230,11 @@ def new_version(data, **kwargs):
     elif 'modified' in data:
         old_modified_property = parse_into_datetime(data.get('modified'), precision='millisecond')
         new_modified_property = parse_into_datetime(kwargs['modified'], precision='millisecond')
-        if new_modified_property < old_modified_property:
-            raise InvalidValueError(cls, 'modified', "The new modified datetime cannot be before the current modified datatime.")
+        if new_modified_property <= old_modified_property:
+            raise InvalidValueError(cls, 'modified',
+                                    "The new modified datetime cannot be before the current modified datetime. The new modified time can "
+                                    "also not be equal to the current modified datetime because if a consumer receives two objects that are "
+                                    "different, but have the same id and modified timestamp, it is not defined how the consumer handles the objects.")
     new_obj_inner.update(kwargs)
     # Exclude properties with a value of 'None' in case data is not an instance of a _STIXBase subclass
     return cls(**{k: v for k, v in new_obj_inner.items() if v is not None})
@@ -255,3 +261,47 @@ def get_class_hierarchy_names(obj):
     for cls in obj.__class__.__mro__:
         names.append(cls.__name__)
     return names
+
+
+def remove_custom_stix(stix_obj):
+    """remove any custom STIX objects or properties
+
+    Args:
+        stix_obj (dict OR python-stix obj): a single python-stix object
+                                             or dict of a STIX object
+    """
+
+    if stix_obj["type"].startswith("x_"):
+        # if entire object is custom, discard
+        return None
+
+    custom_props = []
+    for prop in stix_obj.items():
+        if prop[0].startswith("x_"):
+            # for every custom property, record it and set value to None
+            # (so we can pass it to new_version() and it will be dropped)
+            custom_props.append((prop[0], None))
+
+    if custom_props:
+        # obtain set of object properties that can be transferred
+        # to a new object version. This is 1)custom props with their
+        # values set to None, and 2)any properties left that are not
+        # unmodifiable STIX properties or the "modified" property
+
+        # set of properties that are not supplied to new_version()
+        # to be used for updating properties. This includes unmodifiable
+        # properties (properties that new_version() just re-uses from the
+        # existing STIX object) and the "modified" property. We dont supply the
+        # "modified" property so that new_version() creates a new datetime
+        # value for this property
+        non_supplied_props = STIX_UNMOD_PROPERTIES + ["modified"]
+
+        props = [(prop, stix_obj[prop]) for prop in stix_obj if prop not in non_supplied_props]
+
+        # add to set the custom properties we want to get rid of (with their value=None)
+        props.extend(custom_props)
+
+        return new_version(stix_obj, **(dict(props)))
+
+    else:
+        return stix_obj