chg: Allow to pass an AbstractMISP to add_reference

Fix #379
pull/383/head
Raphaël Vinot 2019-04-16 11:41:40 +02:00
parent e545a7cbc1
commit b67d2e024c
2 changed files with 4 additions and 1 deletions

View File

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

View File

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