diff --git a/pymisp/abstract.py b/pymisp/abstract.py index af4a724..47b9cf7 100644 --- a/pymisp/abstract.py +++ b/pymisp/abstract.py @@ -179,7 +179,7 @@ class AbstractMISP(MutableMapping, MISPFileCache, metaclass=ABCMeta): """Load a JSON string""" self.from_dict(**loads(json_string)) - def to_dict(self) -> Dict: + def to_dict(self, json_format: bool = False) -> Dict: """Dump the class to a dictionary. This method automatically removes the timestamp recursively in every object that has been edited is order to let MISP update the event accordingly.""" @@ -192,6 +192,16 @@ class AbstractMISP(MutableMapping, MISPFileCache, metaclass=ABCMeta): continue elif isinstance(val, str): val = val.strip() + elif json_format: + if isinstance(val, AbstractMISP): + val = val.to_json(True) + elif isinstance(val, (datetime, date)): + val = val.isoformat() + elif isinstance(val, Enum): + val = val.value + elif isinstance(val, UUID): + val = str(val) + if attribute == 'timestamp': if not self.__force_timestamps and is_edited: # In order to be accepted by MISP, the timestamp of an object @@ -201,7 +211,7 @@ class AbstractMISP(MutableMapping, MISPFileCache, metaclass=ABCMeta): continue else: val = self._datetime_to_timestamp(val) - if (attribute in ['first_seen', 'last_seen', 'datetime'] + if (attribute in ('first_seen', 'last_seen', 'datetime') and isinstance(val, datetime) and not val.tzinfo): # Need to make sure the timezone is set. Otherwise, it will be processed as UTC on the server diff --git a/pymisp/mispevent.py b/pymisp/mispevent.py index 65e4c94..a45dff2 100644 --- a/pymisp/mispevent.py +++ b/pymisp/mispevent.py @@ -528,8 +528,8 @@ class MISPAttribute(AbstractMISP): super().from_dict(**kwargs) - def to_dict(self) -> Dict: - to_return = super().to_dict() + def to_dict(self, json_format: bool = False) -> Dict: + to_return = super().to_dict(json_format) if self.data: to_return['data'] = base64.b64encode(self.data.getvalue()).decode() return to_return @@ -967,10 +967,10 @@ class MISPObject(AbstractMISP): to_return.append(a) return to_return - def to_dict(self, strict: bool = False) -> Dict: + def to_dict(self, json_format: bool = False, strict: bool = False) -> Dict: if strict or self._strict and self._known_template: self._validate() - return super(MISPObject, self).to_dict() + return super(MISPObject, self).to_dict(json_format) def to_json(self, sort_keys: bool = False, indent: Optional[int] = None, strict: bool = False): if strict or self._strict and self._known_template: @@ -1730,8 +1730,8 @@ class MISPEvent(AbstractMISP): super(MISPEvent, self).from_dict(**kwargs) - def to_dict(self) -> Dict: - to_return = super().to_dict() + def to_dict(self, json_format: bool = False) -> Dict: + to_return = super().to_dict(json_format) if to_return.get('date'): if isinstance(self.date, datetime): diff --git a/tests/test_mispevent.py b/tests/test_mispevent.py index 251d9ee..7d43b37 100644 --- a/tests/test_mispevent.py +++ b/tests/test_mispevent.py @@ -79,6 +79,15 @@ class TestMISPEvent(unittest.TestCase): ref_json = json.load(f) self.assertEqual(self.mispevent.to_json(sort_keys=True, indent=2), json.dumps(ref_json, sort_keys=True, indent=2)) + def test_to_dict_json_format(self): + misp_event = MISPEvent() + av_signature_object = MISPObject("av-signature") + av_signature_object.add_attribute("signature", "EICAR") + av_signature_object.add_attribute("software", "ClamAv") + misp_event.add_object(av_signature_object) + + self.assertEqual(json.loads(misp_event.to_json()), misp_event.to_dict(json_format=True)) + def test_object_tag(self): self.mispevent.add_object(name='file', strict=True) a = self.mispevent.objects[0].add_attribute('filename', value='')