From 78c9f4f605ea421924a459a748d55fd569db6472 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= Date: Mon, 25 Nov 2019 16:35:56 +0100 Subject: [PATCH] chg: Few more improvements on the feed export --- pymisp/abstract.py | 4 ++-- pymisp/mispevent.py | 19 ++++++++++--------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/pymisp/abstract.py b/pymisp/abstract.py index 7ea7ced..3e7f850 100644 --- a/pymisp/abstract.py +++ b/pymisp/abstract.py @@ -290,7 +290,7 @@ class AbstractMISP(MutableMapping, MISPFileCache): if getattr(self, field, None) is not None: if field in ['timestamp', 'publish_timestamp']: to_return[field] = self._datetime_to_timestamp(getattr(self, field)) - elif field == 'date': + elif isinstance(getattr(self, field), (datetime.datetime, datetime.date)): to_return[field] = getattr(self, field).isoformat() else: to_return[field] = getattr(self, field) @@ -354,7 +354,7 @@ class AbstractMISP(MutableMapping, MISPFileCache): def _datetime_to_timestamp(self, d): """Convert a datetime.datetime object to a timestamp (int)""" - if isinstance(d, (int, str)) or (sys.version_info < (3, 0) and isinstance(d, unicode)): + if isinstance(d, (int, float, str)) or (sys.version_info < (3, 0) and isinstance(d, unicode)): # Assume we already have a timestamp return int(d) if sys.version_info >= (3, 3): diff --git a/pymisp/mispevent.py b/pymisp/mispevent.py index dda64fb..b16b068 100644 --- a/pymisp/mispevent.py +++ b/pymisp/mispevent.py @@ -142,10 +142,7 @@ class MISPAttribute(AbstractMISP): h.update(to_encode.encode("utf-8")) return [h.hexdigest()] - def _to_feed(self, valid_distributions): - if (hasattr(self, 'distribution') and self.distribution is not None - and self.distribution not in valid_distributions): - return False + def _to_feed(self): to_return = super(MISPAttribute, self)._to_feed() if self.data: to_return['data'] = base64.b64encode(self.data.getvalue()).decode() @@ -530,7 +527,9 @@ class MISPEvent(AbstractMISP): if not hasattr(self, r): raise PyMISPError('The field {} is required to generate the event feed output.') - if hasattr(self, 'distribution') and int(self.distribution) not in valid_distributions: + if (hasattr(self, 'distribution') + and self.distribution is not None + and int(self.distribution) not in valid_distributions): return self._set_default() @@ -547,22 +546,24 @@ class MISPEvent(AbstractMISP): for attribute in self.attributes: if (valid_distributions and attribute.get('distribution') is not None and attribute.distribution not in valid_distributions): continue - to_return['Attribute'].append(attribute._to_feed(valid_distributions)) + to_return['Attribute'].append(attribute._to_feed()) if with_meta: to_return['_hashes'] += attribute.hash_values('md5') if self.objects: - to_return['Object']['Attribute'] = [] + to_return['Object'] = [] for obj in self.objects: if (valid_distributions and obj.get('distribution') is not None and obj.distribution not in valid_distributions): continue - to_return['Object'] = obj._to_feed() + obj_to_attach = obj._to_feed() + obj_to_attach['Attribute'] = [] for attribute in obj.attributes: if (valid_distributions and attribute.get('distribution') is not None and attribute.distribution not in valid_distributions): continue - to_return['Object']['Attribute'].append(attribute._to_feed(valid_distributions)) + obj_to_attach['Attribute'].append(attribute._to_feed()) if with_meta: to_return['_hashes'] += attribute.hash_values('md5') + to_return['Object'].append(obj_to_attach) return to_return