diff --git a/pymisp/mispevent.py b/pymisp/mispevent.py index 0c7dad6..c660ed4 100644 --- a/pymisp/mispevent.py +++ b/pymisp/mispevent.py @@ -1076,6 +1076,9 @@ class MISPObject(AbstractMISP): def add_reference(self, referenced_uuid, relationship_type, comment=None, **kwargs): """Add a link (uuid) to an other object""" + if isinstance(referenced_uuid, AbstractMISP): + # Allow to pass an object or an attribute instead of its UUID + referenced_uuid = referenced_uuid.uuid if kwargs.get('object_uuid'): # Load existing object object_uuid = kwargs.pop('object_uuid') diff --git a/tests/test_mispevent.py b/tests/test_mispevent.py index 9344eba..854fca7 100644 --- a/tests/test_mispevent.py +++ b/tests/test_mispevent.py @@ -82,7 +82,7 @@ class TestMISPEvent(unittest.TestCase): del a.uuid self.mispevent.objects[0].uuid = 'a' self.mispevent.objects[1].uuid = 'b' - self.mispevent.objects[0].add_reference('b', 'baz', comment='foo') + self.mispevent.objects[0].add_reference(self.mispevent.objects[1], 'baz', comment='foo') self.assertEqual(self.mispevent.objects[0].references[0].relationship_type, 'baz') with open('tests/mispevent_testfiles/event_obj_attr_tag.json', 'r') as f: ref_json = json.load(f)