chg: Improve add_attribute with a list

Related: #655
pull/665/head
Raphaël Vinot 2020-11-24 13:50:09 +01:00
parent 872005d0eb
commit 35860b49bd
2 changed files with 23 additions and 8 deletions

View File

@ -563,12 +563,13 @@ class PyMISP:
r = self._prepare_request('HEAD', f'attributes/view/{attribute_id}')
return self._check_head_response(r)
def add_attribute(self, event: Union[MISPEvent, int, str, UUID], attribute: MISPAttribute, pythonify: bool = False) -> Union[Dict, MISPAttribute, MISPShadowAttribute]:
def add_attribute(self, event: Union[MISPEvent, int, str, UUID], attribute: Union[MISPAttribute, list], pythonify: bool = False) -> Union[Dict, MISPAttribute, MISPShadowAttribute]:
"""Add an attribute to an existing MISP event
:param event: event to extend
:param attribute: attribute to add. NOTE MISP 2.4.113+: you can pass a list of attributes.
In that case, the pythonified response is the following: {'attributes': [MISPAttribute], 'errors': {errors by attributes}}
:param attribute: attribute or (MISP version 2.4.113+) list of attributes to add.
If a list is passed, the pythonified response is a dict with the following structure:
{'attributes': [MISPAttribute], 'errors': {errors by attributes}}
:param pythonify: Returns a PyMISP Object instead of the plain json output
"""
event_id = get_uuid_or_id_from_abstract_misp(event)
@ -582,10 +583,17 @@ class PyMISP:
if 'errors' in new_attribute:
to_return['errors'] = new_attribute['errors']
for new_attr in new_attribute['Attribute']:
a = MISPAttribute()
a.from_dict(**new_attr)
to_return['attributes'].append(a)
if len(attribute) == 1:
# input list size 1 yields dict, not list of size 1
if 'Attribute' in new_attribute:
a = MISPAttribute()
a.from_dict(**new_attribute['Attribute'])
to_return['attributes'].append(a)
else:
for new_attr in new_attribute['Attribute']:
a = MISPAttribute()
a.from_dict(**new_attr)
to_return['attributes'].append(a)
return to_return
if ('errors' in new_attribute and new_attribute['errors'][0] == 403

View File

@ -1650,6 +1650,13 @@ class TestComprehensive(unittest.TestCase):
self.assertEqual(similar_error['errors'][1]['errors']['value'][0], 'A similar attribute already exists for this event.')
# Test add multiple attributes at once
attr0 = MISPAttribute()
attr0.value = '0.0.0.0'
attr0.type = 'ip-dst'
response = self.user_misp_connector.add_attribute(first, [attr0])
time.sleep(5)
self.assertTrue(isinstance(response['attributes'], list), response['attributes'])
self.assertEqual(response['attributes'][0].value, '0.0.0.0')
attr1 = MISPAttribute()
attr1.value = '1.2.3.4'
attr1.type = 'ip-dst'
@ -1761,7 +1768,7 @@ class TestComprehensive(unittest.TestCase):
# Test attribute*S*
attributes = self.admin_misp_connector.attributes()
self.assertEqual(len(attributes), 6)
self.assertEqual(len(attributes), 7)
# attributes = self.user_misp_connector.attributes()
# self.assertEqual(len(attributes), 5)
# Test event*S*