From 0d86a4339f13f9d6634ec1523a25ecae226f8340 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= Date: Tue, 1 Dec 2020 14:00:43 +0100 Subject: [PATCH] new: Allow to pass an object template to MISPObject.__init__ MISPObject part of #6670 --- pymisp/mispevent.py | 51 +++++++++++++------ .../mispevent_testfiles/misp_custom_obj.json | 1 - tests/test_mispevent.py | 43 +++++++++++++--- 3 files changed, 71 insertions(+), 24 deletions(-) diff --git a/pymisp/mispevent.py b/pymisp/mispevent.py index 3366faa..f5f65c0 100644 --- a/pymisp/mispevent.py +++ b/pymisp/mispevent.py @@ -640,8 +640,15 @@ class MISPObject(AbstractMISP): self.name: str = name self._known_template: bool = False self.id: int + self._definition: Optional[Dict] - self._set_template(kwargs.get('misp_objects_path_custom')) + misp_objects_template_custom = kwargs.pop('misp_objects_template_custom', None) + misp_objects_path_custom = kwargs.pop('misp_objects_path_custom', None) + if misp_objects_template_custom: + self._set_template(misp_objects_template_custom=misp_objects_template_custom) + else: + # Fall back to default path if None + self._set_template(misp_objects_path_custom=misp_objects_path_custom) self.uuid: str = str(uuid.uuid4()) self.first_seen: datetime @@ -679,14 +686,19 @@ class MISPObject(AbstractMISP): self.standalone = standalone def _load_template_path(self, template_path: Union[Path, str]) -> bool: - self._definition: Optional[Dict] = self._load_json(template_path) - if not self._definition: + template = self._load_json(template_path) + if not template: + self._definition = None return False + self._load_template(template) + return True + + def _load_template(self, template: Dict) -> None: + self._definition = template setattr(self, 'meta-category', self._definition['meta-category']) self.template_uuid = self._definition['uuid'] self.description = self._definition['description'] self.template_version = self._definition['version'] - return True def _set_default(self): if not hasattr(self, 'comment'): @@ -715,16 +727,21 @@ class MISPObject(AbstractMISP): self.name = object_name self._set_template(misp_objects_path_custom) - def _set_template(self, misp_objects_path_custom: Optional[Union[Path, str]] = None): - if misp_objects_path_custom: - # If misp_objects_path_custom is given, and an object with the given name exists, use that. - if isinstance(misp_objects_path_custom, str): - self.misp_objects_path = Path(misp_objects_path_custom) - else: - self.misp_objects_path = misp_objects_path_custom + def _set_template(self, misp_objects_path_custom: Optional[Union[Path, str]] = None, misp_objects_template_custom: Optional[Dict] = None): + if misp_objects_template_custom: + # A complete template was given to the constructor + self._load_template(misp_objects_template_custom) + self._known_template = True + else: + if misp_objects_path_custom: + # If misp_objects_path_custom is given, and an object with the given name exists, use that. + if isinstance(misp_objects_path_custom, str): + self.misp_objects_path = Path(misp_objects_path_custom) + else: + self.misp_objects_path = misp_objects_path_custom - # Try to get the template - self._known_template = self._load_template_path(self.misp_objects_path / self.name / 'definition.json') + # Try to get the template + self._known_template = self._load_template_path(self.misp_objects_path / self.name / 'definition.json') if not self._known_template and self._strict: raise UnknownMISPObjectTemplate('{} is unknown in the MISP object directory.'.format(self.name)) @@ -789,6 +806,10 @@ class MISPObject(AbstractMISP): else: self._known_template = False + # depending on how the object is initialized, we may have a few keys to pop + kwargs.pop('misp_objects_template_custom', None) + kwargs.pop('misp_objects_path_custom', None) + if 'distribution' in kwargs and kwargs['distribution'] is not None: self.distribution = kwargs.pop('distribution') self.distribution = int(self.distribution) @@ -903,9 +924,7 @@ class MISPObject(AbstractMISP): else: attribute = MISPObjectAttribute({}) # Overwrite the parameters of self._default_attributes_parameters with the ones of value - attribute.from_dict(object_relation=object_relation, **dict(self._default_attributes_parameters, **value)) - # FIXME New syntax python3 only, keep for later. - # attribute.from_dict(object_relation=object_relation, **{**self._default_attributes_parameters, **value}) + attribute.from_dict(object_relation=object_relation, **{**self._default_attributes_parameters, **value}) self.__fast_attribute_access[object_relation].append(attribute) self.Attribute.append(attribute) self.edited = True diff --git a/tests/mispevent_testfiles/misp_custom_obj.json b/tests/mispevent_testfiles/misp_custom_obj.json index 6cb6ff2..a58ae5a 100644 --- a/tests/mispevent_testfiles/misp_custom_obj.json +++ b/tests/mispevent_testfiles/misp_custom_obj.json @@ -22,7 +22,6 @@ "description": "TestTemplate.", "distribution": "5", "meta-category": "file", - "misp_objects_path_custom": "tests/mispevent_testfiles", "name": "test_object_template", "sharing_group_id": "0", "template_uuid": "4ec55cc6-9e49-4c64-b794-03c25c1a6589", diff --git a/tests/test_mispevent.py b/tests/test_mispevent.py index 9025d90..251d9ee 100644 --- a/tests/test_mispevent.py +++ b/tests/test_mispevent.py @@ -3,7 +3,6 @@ import unittest import json -import sys from io import BytesIO import glob import hashlib @@ -285,17 +284,47 @@ class TestMISPEvent(unittest.TestCase): misp_obj = self.mispevent.get_object_by_id(1556) self.assertEqual(misp_obj.uuid, '5a3cd604-e11c-4de5-bbbf-c170950d210f') - def test_userdefined_object(self): + def test_userdefined_object_custom_template(self): + self.init_event() + with open('tests/mispevent_testfiles/test_object_template/definition.json') as f: + template = json.load(f) + self.mispevent.add_object(name='test_object_template', strict=True, + misp_objects_template_custom=template) + with self.assertRaises(InvalidMISPObject) as e: + # Fail on required + self.mispevent.to_json(sort_keys=True, indent=2) + self.assertEqual(e.exception.message, '{\'member3\'} are required.') + + a = self.mispevent.objects[0].add_attribute('member3', value='foo') + del a.uuid + with self.assertRaises(InvalidMISPObject) as e: + # Fail on requiredOneOf + self.mispevent.to_json(sort_keys=True, indent=2) + self.assertEqual(e.exception.message, 'At least one of the following attributes is required: member1, member2') + + a = self.mispevent.objects[0].add_attribute('member1', value='bar') + del a.uuid + a = self.mispevent.objects[0].add_attribute('member1', value='baz') + del a.uuid + with self.assertRaises(InvalidMISPObject) as e: + # member1 is not a multiple + self.mispevent.to_json(sort_keys=True, indent=2) + self.assertEqual(e.exception.message, 'Multiple occurrences of member1 is not allowed') + + self.mispevent.objects[0].attributes = self.mispevent.objects[0].attributes[:2] + self.mispevent.objects[0].uuid = 'a' + with open('tests/mispevent_testfiles/misp_custom_obj.json', 'r') as f: + ref_json = json.load(f) + del self.mispevent.uuid + self.assertEqual(self.mispevent.to_json(sort_keys=True, indent=2), json.dumps(ref_json, sort_keys=True, indent=2)) + + def test_userdefined_object_custom_dir(self): self.init_event() self.mispevent.add_object(name='test_object_template', strict=True, misp_objects_path_custom='tests/mispevent_testfiles') with self.assertRaises(InvalidMISPObject) as e: # Fail on required self.mispevent.to_json(sort_keys=True, indent=2) - if sys.version_info >= (3, ): - self.assertEqual(e.exception.message, '{\'member3\'} are required.') - else: - # Python2 bullshit - self.assertEqual(e.exception.message, 'set([u\'member3\']) are required.') + self.assertEqual(e.exception.message, '{\'member3\'} are required.') a = self.mispevent.objects[0].add_attribute('member3', value='foo') del a.uuid