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}') r = self._prepare_request('HEAD', f'attributes/view/{attribute_id}')
return self._check_head_response(r) 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 """Add an attribute to an existing MISP event
:param event: event to extend :param event: event to extend
:param attribute: attribute to add. NOTE MISP 2.4.113+: you can pass a list of attributes. :param attribute: attribute or (MISP version 2.4.113+) list of attributes to add.
In that case, the pythonified response is the following: {'attributes': [MISPAttribute], 'errors': {errors by attributes}} 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 :param pythonify: Returns a PyMISP Object instead of the plain json output
""" """
event_id = get_uuid_or_id_from_abstract_misp(event) event_id = get_uuid_or_id_from_abstract_misp(event)
@ -582,6 +583,13 @@ class PyMISP:
if 'errors' in new_attribute: if 'errors' in new_attribute:
to_return['errors'] = new_attribute['errors'] to_return['errors'] = new_attribute['errors']
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']: for new_attr in new_attribute['Attribute']:
a = MISPAttribute() a = MISPAttribute()
a.from_dict(**new_attr) a.from_dict(**new_attr)

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.') self.assertEqual(similar_error['errors'][1]['errors']['value'][0], 'A similar attribute already exists for this event.')
# Test add multiple attributes at once # 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 = MISPAttribute()
attr1.value = '1.2.3.4' attr1.value = '1.2.3.4'
attr1.type = 'ip-dst' attr1.type = 'ip-dst'
@ -1761,7 +1768,7 @@ class TestComprehensive(unittest.TestCase):
# Test attribute*S* # Test attribute*S*
attributes = self.admin_misp_connector.attributes() attributes = self.admin_misp_connector.attributes()
self.assertEqual(len(attributes), 6) self.assertEqual(len(attributes), 7)
# attributes = self.user_misp_connector.attributes() # attributes = self.user_misp_connector.attributes()
# self.assertEqual(len(attributes), 5) # self.assertEqual(len(attributes), 5)
# Test event*S* # Test event*S*