new: Allow to pass an object template to MISPObject.__init__

MISPObject part of #6670
pull/667/head
Raphaël Vinot 2020-12-01 14:00:43 +01:00
parent bbc68f4f2d
commit 0d86a4339f
3 changed files with 71 additions and 24 deletions

View File

@ -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

View File

@ -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",

View File

@ -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