start of automated property checking.

stix2/__init__.py

@@ -6,6 +6,18 @@ import uuid
 import pytz
 
 
+COMMON_PROPERTIES = {
+    'type': {
+        'default': (lambda x: x._type),
+        'validate': (lambda x, val: val == x._type),
+        'error_msg': "{type} must have {field}='{expected}'.",
+    },
+    'id': {},
+    'created': {},
+    'modified': {},
+}
+
+
 def format_datetime(dt):
     # TODO: how to handle naive datetime
 
@@ -29,13 +41,18 @@ class _STIXBase(collections.Mapping):
     def _check_kwargs(cls, **kwargs):
         class_name = cls.__name__
 
-        # Ensure that, if provided, the 'type' kwarg is correct.
+        for prop_name, prop_metadata in cls._properties.items():
-        required_type = cls._type
+            if prop_metadata.get('default') and prop_name not in kwargs:
-        if not kwargs.get('type'):
+                kwargs[prop_name] = prop_metadata['default'](cls)
-            kwargs['type'] = required_type
+
-        if kwargs['type'] != required_type:
+            if prop_metadata.get('validate'):
-            msg = "{0} must have type='{1}'."
+                if not prop_metadata['validate'](cls, kwargs[prop_name]):
-            raise ValueError(msg.format(class_name, required_type))
+                    msg = prop_metadata['error_msg'].format(
+                        type=class_name,
+                        field=prop_name,
+                        expected=prop_metadata.get('default')(cls),
+                    )
+                    raise ValueError(msg)
 
         id_prefix = cls._type + "--"
         if not kwargs.get('id'):
@@ -85,12 +102,16 @@ class _STIXBase(collections.Mapping):
 class Bundle(_STIXBase):
 
     _type = 'bundle'
-    _properties = [
+    _properties = {
-        'type',
+        'type': {
-        'id',
+            'default': (lambda x: x._type),
-        'spec_version',
+            'validate': (lambda x, val: val == x._type),
-        'objects',
+            'error_msg': "{type} must have {field}='{expected}'.",
-    ]
+        },
+        'id': {},
+        'spec_version': {},
+        'objects': {},
+    }
 
     def __init__(self, **kwargs):
         # TODO: remove once we check all the fields in the right order
@@ -119,15 +140,12 @@ class Bundle(_STIXBase):
 class Indicator(_STIXBase):
 
     _type = 'indicator'
-    _properties = [
+    _properties = COMMON_PROPERTIES.copy()
-        'type',
+    _properties.update({
-        'id',
+        'labels': {},
-        'created',
+        'pattern': {},
-        'modified',
+        'valid_from': {},
-        'labels',
+    })
-        'pattern',
-        'valid_from',
-    ]
 
     def __init__(self, **kwargs):
         # TODO:
@@ -177,14 +195,11 @@ class Indicator(_STIXBase):
 class Malware(_STIXBase):
 
     _type = 'malware'
-    _properties = [
+    _properties = COMMON_PROPERTIES.copy()
-        'type',
+    _properties.update({
-        'id',
+        'labels': {},
-        'created',
+        'name': {},
-        'modified',
+    })
-        'labels',
-        'name',
-    ]
 
     def __init__(self, **kwargs):
         # TODO:
@@ -230,15 +245,12 @@ class Malware(_STIXBase):
 class Relationship(_STIXBase):
 
     _type = 'relationship'
-    _properties = [
+    _properties = COMMON_PROPERTIES.copy()
-        'type',
+    _properties.update({
-        'id',
+        'relationship_type': {},
-        'created',
+        'source_ref': {},
-        'modified',
+        'target_ref': {},
-        'relationship_type',
+    })
-        'source_ref',
-        'target_ref',
-    ]
 
     # Explicitly define the first three kwargs to make readable Relationship declarations.
     def __init__(self, source_ref=None, relationship_type=None, target_ref=None,