From 1ba064734be640824e4e7699aef7d56b0631cb6b Mon Sep 17 00:00:00 2001
From: Greg Back <gback@mitre.org>
Date: Thu, 2 Feb 2017 09:50:35 -0600
Subject: [PATCH] Special handling for timestamp fields.

If a type has more than one timestamp field that should be automatically generated,
we want them to all be same, not vary by milliseconds.
---
 stix2/__init__.py | 83 +++++++++++++++++------------------------------
 1 file changed, 30 insertions(+), 53 deletions(-)

diff --git a/stix2/__init__.py b/stix2/__init__.py
index ca7739f..c9628c2 100644
--- a/stix2/__init__.py
+++ b/stix2/__init__.py
@@ -5,6 +5,12 @@ import uuid
 
 import pytz
 
+
+# Sentinel value for fields that should be set to the current time.
+# We can't use the standard 'default' approach, since if there are multiple
+# timestamps in a single object, the timestamps will vary by a few microseconds.
+NOW = object()
+
 DEFAULT_ERROR = "{type} must have {field}='{expected}'."
 COMMON_PROPERTIES = {
     'type': {
@@ -17,8 +23,12 @@ COMMON_PROPERTIES = {
         'expected': (lambda x: x._type + "--"),
         'error_msg': "{type} {field} values must begin with '{expected}'."
     },
-    'created': {},
-    'modified': {},
+    'created': {
+        'default': NOW,
+    },
+    'modified': {
+        'default': NOW,
+    },
 }
 
 
@@ -52,10 +62,18 @@ class _STIXBase(collections.Mapping):
     def _make_id(cls):
         return cls._type + "--" + str(uuid.uuid4())
 
-    @classmethod
-    def _check_kwargs(cls, **kwargs):
+    def __init__(self, **kwargs):
+        cls = self.__class__
         class_name = cls.__name__
 
+        # Use the same timestamp for any auto-generated datetimes
+        now = datetime.datetime.now(tz=pytz.UTC)
+
+        # Detect any keyword arguments not allowed for a specific type
+        extra_kwargs = list(set(kwargs) - set(cls._properties))
+        if extra_kwargs:
+            raise TypeError("unexpected keyword arguments: " + str(extra_kwargs))
+
         for prop_name, prop_metadata in cls._properties.items():
             if prop_name not in kwargs:
                 if prop_metadata.get('required'):
@@ -63,7 +81,11 @@ class _STIXBase(collections.Mapping):
                     raise ValueError(msg.format(type=class_name,
                                                 field=prop_name))
                 if prop_metadata.get('default'):
-                    kwargs[prop_name] = prop_metadata['default'](cls)
+                    default = prop_metadata['default']
+                    if default == NOW:
+                        kwargs[prop_name] = now
+                    else:
+                        kwargs[prop_name] = default(cls)
                 elif prop_metadata.get('fixed'):
                     kwargs[prop_name] = prop_metadata['fixed']
 
@@ -85,18 +107,6 @@ class _STIXBase(collections.Mapping):
                     )
                     raise ValueError(msg)
 
-        return kwargs
-
-    def __init__(self, **kwargs):
-        # TODO: move all of this back into init, once we check the right things
-        # in the right order, or move after the unexpected check.
-        kwargs = self._check_kwargs(**kwargs)
-
-        # Detect any keyword arguments not allowed for a specific type
-        extra_kwargs = list(set(kwargs) - set(self.__class__._properties))
-        if extra_kwargs:
-            raise TypeError("unexpected keyword arguments: " + str(extra_kwargs))
-
         self._inner = kwargs
 
     def __getitem__(self, key):
@@ -155,7 +165,9 @@ class Indicator(_STIXBase):
         'pattern': {
             'required': True,
         },
-        'valid_from': {},
+        'valid_from': {
+            'default': NOW,
+        },
     })
 
     def __init__(self, **kwargs):
@@ -171,18 +183,6 @@ class Indicator(_STIXBase):
         # - valid_until
         # - kill_chain_phases
 
-        # TODO: do we care about the performance penalty of creating this
-        # if we won't need it?
-        now = datetime.datetime.now(tz=pytz.UTC)
-
-        # TODO: remove once we check all the fields in the right order
-        kwargs = self._check_kwargs(**kwargs)
-
-        kwargs.update({
-            'created': kwargs.get('created', now),
-            'modified': kwargs.get('modified', now),
-            'valid_from': kwargs.get('valid_from', now),
-        })
         super(Indicator, self).__init__(**kwargs)
 
 
@@ -210,17 +210,6 @@ class Malware(_STIXBase):
         # - description
         # - kill_chain_phases
 
-        # TODO: do we care about the performance penalty of creating this
-        # if we won't need it?
-        now = datetime.datetime.now(tz=pytz.UTC)
-
-        # TODO: remove once we check all the fields in the right order
-        kwargs = self._check_kwargs(**kwargs)
-
-        kwargs.update({
-            'created': kwargs.get('created', now),
-            'modified': kwargs.get('modified', now),
-        })
         super(Malware, self).__init__(**kwargs)
 
 
@@ -260,13 +249,6 @@ 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 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):
@@ -275,9 +257,4 @@ class Relationship(_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),
-        })
-
         super(Relationship, self).__init__(**kwargs)