From c9b5d240730cc3ca4c0dfb30d11298472f44e53e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Wed, 11 Dec 2019 23:12:14 +0100
Subject: [PATCH] fix: Add missing fields to event & attribute for the feed
 output

---
 pymisp/abstract.py  | 11 +++++++++++
 pymisp/mispevent.py | 34 +++++++++++++++++++++++++++++-----
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/pymisp/abstract.py b/pymisp/abstract.py
index 5081a7c..c7c3fa2 100644
--- a/pymisp/abstract.py
+++ b/pymisp/abstract.py
@@ -290,6 +290,8 @@ class AbstractMISP(MutableMapping, MISPFileCache):
     def _to_feed(self):
         if not hasattr(self, '_fields_for_feed'):
             raise PyMISPError('Unable to export in the feed format, _fields_for_feed is missing.')
+        if hasattr(self, '_set_default') and callable(self._set_default):
+            self._set_default()
         to_return = {}
         for field in self._fields_for_feed:
             if getattr(self, field, None) is not None:
@@ -299,6 +301,11 @@ class AbstractMISP(MutableMapping, MISPFileCache):
                     to_return[field] = getattr(self, field).isoformat()
                 else:
                     to_return[field] = getattr(self, field)
+            else:
+                if field == 'data':
+                    # data in attribute is special
+                    continue
+                raise PyMISPError('The field {} is required in {} when generating a feed.'.format(field, self.__class__.__name__))
         return to_return
 
     def to_json(self, sort_keys=False, indent=None):
@@ -423,6 +430,10 @@ class MISPTag(AbstractMISP):
             kwargs = kwargs.get('Tag')
         super(MISPTag, self).from_dict(**kwargs)
 
+    def _set_default(self):
+        if not hasattr(self, 'colour'):
+            self.colour = '#ffffff'
+
     def _to_feed(self):
         if hasattr(self, 'exportable') and not self.exportable:
             return False
diff --git a/pymisp/mispevent.py b/pymisp/mispevent.py
index 81d0c2e..16f66ea 100644
--- a/pymisp/mispevent.py
+++ b/pymisp/mispevent.py
@@ -104,7 +104,7 @@ def make_bool(value):
 
 class MISPAttribute(AbstractMISP):
     _fields_for_feed = {'uuid', 'value', 'category', 'type', 'comment', 'data',
-                        'timestamp', 'to_ids', 'object_relation', 'disable_correlation'}
+                        'timestamp', 'to_ids', 'disable_correlation'}
 
     def __init__(self, describe_types=None, strict=False):
         """Represents an Attribute
@@ -142,6 +142,12 @@ class MISPAttribute(AbstractMISP):
             h.update(to_encode.encode("utf-8"))
             return [h.hexdigest()]
 
+    def _set_default(self):
+        if not hasattr(self, 'comment'):
+            self.comment = ''
+        if not hasattr(self, 'timestamp'):
+            self.timestamp = datetime.datetime.timestamp(datetime.datetime.now())
+
     def _to_feed(self):
         to_return = super(MISPAttribute, self)._to_feed()
         if self.data:
@@ -473,15 +479,19 @@ class MISPEvent(AbstractMISP):
         self.ShadowAttribute = []
 
     def _set_default(self):
-        """There are a few keys that could be set by default"""
+        """There are a few keys that could, or need to be set by default for the feed generator"""
         if not hasattr(self, 'published'):
             self.published = True
         if not hasattr(self, 'uuid'):
             self.uuid = str(uuid.uuid4())
+        if not hasattr(self, 'extends_uuid'):
+            self.extends_uuid = ''
         if not hasattr(self, 'date'):
             self.set_date(datetime.date.today())
         if not hasattr(self, 'timestamp'):
             self.timestamp = datetime.datetime.timestamp(datetime.datetime.now())
+        if not hasattr(self, 'publish_timestamp'):
+            self.publish_timestamp = datetime.datetime.timestamp(datetime.datetime.now())
         if not hasattr(self, 'analysis'):
             # analysis: 0 means initial, 1 ongoing, 2 completed
             self.analysis = 2
@@ -534,8 +544,6 @@ class MISPEvent(AbstractMISP):
                 and int(self.distribution) not in valid_distributions):
             return
 
-        self._set_default()
-
         to_return = super(MISPEvent, self)._to_feed()
         if with_meta:
             to_return['_hashes'] = []
@@ -567,7 +575,7 @@ class MISPEvent(AbstractMISP):
                         to_return['_hashes'] += attribute.hash_values('md5')
                 to_return['Object'].append(obj_to_attach)
 
-        return to_return
+        return {'Event': to_return}
 
     @property
     def known_types(self):
@@ -1001,6 +1009,13 @@ class MISPObjectReference(AbstractMISP):
 
     def __init__(self):
         super(MISPObjectReference, self).__init__()
+        self.uuid = str(uuid.uuid4())
+
+    def _set_default(self):
+        if not hasattr(self, 'comment'):
+            self.comment = ''
+        if not hasattr(self, 'timestamp'):
+            self.timestamp = datetime.datetime.timestamp(datetime.datetime.now())
 
     def from_dict(self, **kwargs):
         if 'ObjectReference' in kwargs:
@@ -1202,6 +1217,9 @@ class MISPSighting(AbstractMISP):
 
 class MISPObjectAttribute(MISPAttribute):
 
+    _fields_for_feed = {'uuid', 'object_relation', 'value', 'category', 'type',
+                        'comment', 'data', 'timestamp', 'to_ids', 'disable_correlation'}
+
     def __init__(self, definition):
         super(MISPObjectAttribute, self).__init__()
         self._definition = definition
@@ -1353,6 +1371,12 @@ class MISPObject(AbstractMISP):
         self.template_version = self._definition['version']
         return True
 
+    def _set_default(self):
+        if not hasattr(self, 'comment'):
+            self.comment = ''
+        if not hasattr(self, 'timestamp'):
+            self.timestamp = datetime.datetime.timestamp(datetime.datetime.now())
+
     def _to_feed(self):
         to_return = super(MISPObject, self)._to_feed()
         if self.references: