From 675a29dbfba52328550973a05f9d012cb302c784 Mon Sep 17 00:00:00 2001
From: Greg Back <gback@mitre.org>
Date: Thu, 2 Feb 2017 08:53:46 -0600
Subject: [PATCH] Add support for required fields with no default values.

---
 stix2/__init__.py | 66 +++++++++++++++++++++++------------------------
 1 file changed, 33 insertions(+), 33 deletions(-)

diff --git a/stix2/__init__.py b/stix2/__init__.py
index 7458f1e..ca7739f 100644
--- a/stix2/__init__.py
+++ b/stix2/__init__.py
@@ -58,6 +58,10 @@ class _STIXBase(collections.Mapping):
 
         for prop_name, prop_metadata in cls._properties.items():
             if prop_name not in kwargs:
+                if prop_metadata.get('required'):
+                    msg = "Missing required field for {type}: '{field}'."
+                    raise ValueError(msg.format(type=class_name,
+                                                field=prop_name))
                 if prop_metadata.get('default'):
                     kwargs[prop_name] = prop_metadata['default'](cls)
                 elif prop_metadata.get('fixed'):
@@ -145,8 +149,12 @@ class Indicator(_STIXBase):
     _type = 'indicator'
     _properties = COMMON_PROPERTIES.copy()
     _properties.update({
-        'labels': {},
-        'pattern': {},
+        'labels': {
+            'required': True,
+        },
+        'pattern': {
+            'required': True,
+        },
         'valid_from': {},
     })
 
@@ -170,12 +178,6 @@ class Indicator(_STIXBase):
         # TODO: remove once we check all the fields in the right order
         kwargs = self._check_kwargs(**kwargs)
 
-        if not kwargs.get('labels'):
-            raise ValueError("Missing required field for Indicator: 'labels'.")
-
-        if not kwargs.get('pattern'):
-            raise ValueError("Missing required field for Indicator: 'pattern'.")
-
         kwargs.update({
             'created': kwargs.get('created', now),
             'modified': kwargs.get('modified', now),
@@ -189,8 +191,12 @@ class Malware(_STIXBase):
     _type = 'malware'
     _properties = COMMON_PROPERTIES.copy()
     _properties.update({
-        'labels': {},
-        'name': {},
+        'labels': {
+            'required': True,
+        },
+        'name': {
+            'required': True,
+        },
     })
 
     def __init__(self, **kwargs):
@@ -211,12 +217,6 @@ class Malware(_STIXBase):
         # TODO: remove once we check all the fields in the right order
         kwargs = self._check_kwargs(**kwargs)
 
-        if not kwargs.get('labels'):
-            raise ValueError("Missing required field for Malware: 'labels'.")
-
-        if not kwargs.get('name'):
-            raise ValueError("Missing required field for Malware: 'name'.")
-
         kwargs.update({
             'created': kwargs.get('created', now),
             'modified': kwargs.get('modified', now),
@@ -229,9 +229,15 @@ class Relationship(_STIXBase):
     _type = 'relationship'
     _properties = COMMON_PROPERTIES.copy()
     _properties.update({
-        'relationship_type': {},
-        'source_ref': {},
-        'target_ref': {},
+        'relationship_type': {
+            'required': True,
+        },
+        'source_ref': {
+            'required': True,
+        },
+        'target_ref': {
+            'required': True,
+        },
     })
 
     # Explicitly define the first three kwargs to make readable Relationship declarations.
@@ -246,9 +252,7 @@ class Relationship(_STIXBase):
 
         # - description
 
-        # TODO: remove once we check all the fields in the right order
-        kwargs = self._check_kwargs(**kwargs)
-
+        # Allow (source_ref, relationship_type, target_ref) as positional args.
         if source_ref and not kwargs.get('source_ref'):
             kwargs['source_ref'] = source_ref
         if relationship_type and not kwargs.get('relationship_type'):
@@ -256,28 +260,24 @@ class Relationship(_STIXBase):
         if target_ref and not kwargs.get('target_ref'):
             kwargs['target_ref'] = target_ref
 
+        # TODO: remove once we check all the fields in the right order
+        kwargs = self._check_kwargs(**kwargs)
+
         # TODO: do we care about the performance penalty of creating this
         # if we won't need it?
         now = datetime.datetime.now(tz=pytz.UTC)
 
-        if not kwargs.get('relationship_type'):
-            raise ValueError("Missing required field for Relationship: 'relationship_type'.")
-
-        if not kwargs.get('source_ref'):
-            raise ValueError("Missing required field for Relationship: 'source_ref'.")
-        elif isinstance(kwargs['source_ref'], _STIXBase):
+        # If actual STIX objects (vs. just the IDs) are passed in, extract the
+        # ID values to use in the Relationship object.
+        if kwargs.get('source_ref') and isinstance(kwargs['source_ref'], _STIXBase):
             kwargs['source_ref'] = kwargs['source_ref'].id
 
-        if not kwargs.get('target_ref'):
-            raise ValueError("Missing required field for Relationship: 'target_ref'.")
-        elif isinstance(kwargs['target_ref'], _STIXBase):
+        if kwargs.get('target_ref') and isinstance(kwargs['target_ref'], _STIXBase):
             kwargs['target_ref'] = kwargs['target_ref'].id
 
         kwargs.update({
             'created': kwargs.get('created', now),
             'modified': kwargs.get('modified', now),
-            'relationship_type': kwargs['relationship_type'],
-            'target_ref': kwargs['target_ref'],
         })
 
         super(Relationship, self).__init__(**kwargs)