diff --git a/pymisp/mispevent.py b/pymisp/mispevent.py index e843d6f..68263b7 100644 --- a/pymisp/mispevent.py +++ b/pymisp/mispevent.py @@ -1016,6 +1016,12 @@ class MISPObject(AbstractMISP): else: self._known_template = False + if 'distribution' in kwargs and kwargs['distribution'] is not None: + self.distribution = kwargs.pop('distribution') + self.distribution = int(self.distribution) + if self.distribution not in [0, 1, 2, 3, 4, 5]: + raise NewAttributeError('{} is invalid, the distribution has to be in 0, 1, 2, 3, 4, 5'.format(self.distribution)) + if kwargs.get('timestamp'): if sys.version_info >= (3, 3): self.timestamp = datetime.datetime.fromtimestamp(int(kwargs.pop('timestamp')), datetime.timezone.utc) diff --git a/tests/testlive_comprehensive.py b/tests/testlive_comprehensive.py index b33976f..6b9d460 100644 --- a/tests/testlive_comprehensive.py +++ b/tests/testlive_comprehensive.py @@ -3,7 +3,7 @@ import unittest -from pymisp import ExpandedPyMISP, MISPEvent, MISPOrganisation, MISPUser, Distribution, ThreatLevel, Analysis +from pymisp import ExpandedPyMISP, MISPEvent, MISPOrganisation, MISPUser, Distribution, ThreatLevel, Analysis, MISPObject from datetime import datetime, timedelta, date from io import BytesIO @@ -14,7 +14,7 @@ try: except ImportError as e: print(e) url = 'http://localhost:8080' - key = 'BSip0zVadeFDeolkX2g7MHx8mrlr0uE04hh6CQj0' + key = 'LBelWqKY9SQyG0huZzAMqiEBl6FODxpgRRXMsZFu' from uuid import uuid4 @@ -439,6 +439,49 @@ class TestComprehensive(unittest.TestCase): self.admin_misp_connector.delete_event(first.id) self.admin_misp_connector.delete_event(second.id) + def test_default_distribution(self): + '''The default distributions on the VM are This community only for the events and Inherit from event for attr/obj)''' + first = self.create_simple_event() + del first.distribution + o = first.add_object(name='file') + o.add_attribute('filename', value='foo.exe') + try: + # Event create + first = self.user_misp_connector.add_event(first) + self.assertEqual(first.distribution, Distribution.this_community_only.value) + self.assertEqual(first.attributes[0].distribution, Distribution.inherit.value) + self.assertEqual(first.objects[0].distribution, Distribution.inherit.value) + self.assertEqual(first.objects[0].attributes[0].distribution, Distribution.inherit.value) + # Event edit + first.add_attribute('ip-dst', '12.54.76.43') + o = first.add_object(name='file') + o.add_attribute('filename', value='foo2.exe') + first = self.user_misp_connector.update_event(first) + self.assertEqual(first.attributes[1].distribution, Distribution.inherit.value) + self.assertEqual(first.objects[1].distribution, Distribution.inherit.value) + self.assertEqual(first.objects[1].attributes[0].distribution, Distribution.inherit.value) + # Attribute create + attribute = self.user_misp_connector.add_named_attribute(first, 'comment', 'bar') + # FIXME: Add helper that returns a list of MISPAttribute + self.assertEqual(attribute[0]['Attribute']['distribution'], str(Distribution.inherit.value)) + # Object - add + o = MISPObject('file') + o.add_attribute('filename', value='blah.exe') + new_obj = self.user_misp_connector.add_object(first.id, o.template_uuid, o) + # FIXME: Add helper that returns a MISPObject + self.assertEqual(new_obj['Object']['distribution'], str(Distribution.inherit.value)) + self.assertEqual(new_obj['Object']['Attribute'][0]['distribution'], str(Distribution.inherit.value)) + # Object - edit + clean_obj = MISPObject(**new_obj['Object']) + clean_obj.from_dict(**new_obj['Object']) + clean_obj.add_attribute('filename', value='blah.exe') + new_obj = self.user_misp_connector.edit_object(clean_obj) + for a in new_obj['Object']['Attribute']: + self.assertEqual(a['distribution'], str(Distribution.inherit.value)) + finally: + # Delete event + self.admin_misp_connector.delete_event(first.id) + def test_simple_event(self): '''Search a bunch of parameters: * Value not existing