MISP
/
PyMISP
mirror of https://github.com/MISP/PyMISP
2
0
Fork 0
Code Issues Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pymisp/mispevent.py
View File

@ -1076,6 +1076,9 @@ class MISPObject(AbstractMISP):
def add_reference(self, referenced_uuid, relationship_type, comment=None, **kwargs): def add_reference(self, referenced_uuid, relationship_type, comment=None, **kwargs):
"""Add a link (uuid) to an other object""" """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'): if kwargs.get('object_uuid'):
# Load existing object # Load existing object
object_uuid = kwargs.pop('object_uuid') object_uuid = kwargs.pop('object_uuid')

2
tests/test_mispevent.py
View File

@ -82,7 +82,7 @@ class TestMISPEvent(unittest.TestCase):
del a.uuid del a.uuid
self.mispevent.objects[0].uuid = 'a' self.mispevent.objects[0].uuid = 'a'
self.mispevent.objects[1].uuid = 'b' 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') self.assertEqual(self.mispevent.objects[0].references[0].relationship_type, 'baz')
with open('tests/mispevent_testfiles/event_obj_attr_tag.json', 'r') as f: with open('tests/mispevent_testfiles/event_obj_attr_tag.json', 'r') as f:
ref_json = json.load(f) ref_json = json.load(f)