new: Add get_object to ExpandedPyMISP.

Fix #372
MISP · Apr 4, 2019 · 99b2052 · 99b2052
1 parent 1de4c9d
commit 99b2052
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 6 deletions.
diff --git a/pymisp/aping.py b/pymisp/aping.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
-from .exceptions import MISPServerError, NewEventError, UpdateEventError, UpdateAttributeError, PyMISPNotImplementedYet
+from .exceptions import MISPServerError, PyMISPUnexpectedResponse, PyMISPNotImplementedYet
 from .api import PyMISP, everything_broken
 from .mispevent import MISPEvent, MISPAttribute, MISPSighting, MISPLog, MISPObject
 from typing import TypeVar, Optional, Tuple, List, Dict, Union
@@ -100,7 +100,7 @@ def get_event(self, event_id: int):
     def add_object(self, event_id: int, misp_object: MISPObject):
         created_object = super().add_object(event_id, misp_object)
         if isinstance(created_object, str):
-            raise NewEventError(f'Unexpected response from server: {created_object}')
+            raise PyMISPUnexpectedResponse(f'Unexpected response from server: {created_object}')
         elif 'errors' in created_object:
             return created_object
         o = MISPObject(misp_object.name)
@@ -110,17 +110,31 @@ def add_object(self, event_id: int, misp_object: MISPObject):
     def update_object(self, misp_object: MISPObject):
         updated_object = super().edit_object(misp_object)
         if isinstance(updated_object, str):
-            raise NewEventError(f'Unexpected response from server: {updated_object}')
+            raise PyMISPUnexpectedResponse(f'Unexpected response from server: {updated_object}')
         elif 'errors' in updated_object:
             return updated_object
         o = MISPObject(misp_object.name)
         o.from_dict(**updated_object)
         return o
 
+    def get_object(self, object_id: int):
+        """Get an object
+
+        :param obj_id: Object id to get
+        """
+        misp_object = super().get_object(object_id)
+        if isinstance(misp_object, str):
+            raise PyMISPUnexpectedResponse(f'Unexpected response from server: {misp_object}')
+        elif 'errors' in misp_object:
+            return misp_object
+        o = MISPObject(misp_object['Object']['name'])
+        o.from_dict(**misp_object)
+        return o
+
     def add_event(self, event: MISPEvent):
         created_event = super().add_event(event)
         if isinstance(created_event, str):
-            raise NewEventError(f'Unexpected response from server: {created_event}')
+            raise PyMISPUnexpectedResponse(f'Unexpected response from server: {created_event}')
         elif 'errors' in created_event:
             return created_event
         e = MISPEvent()
@@ -130,7 +144,7 @@ def add_event(self, event: MISPEvent):
     def update_event(self, event: MISPEvent):
         updated_event = super().update_event(event.uuid, event)
         if isinstance(updated_event, str):
-            raise UpdateEventError(f'Unexpected response from server: {updated_event}')
+            raise PyMISPUnexpectedResponse(f'Unexpected response from server: {updated_event}')
         elif 'errors' in updated_event:
             return updated_event
         e = MISPEvent()
@@ -140,7 +154,7 @@ def update_event(self, event: MISPEvent):
     def update_attribute(self, attribute: MISPAttribute):
         updated_attribute = super().update_attribute(attribute.uuid, attribute)
         if isinstance(updated_attribute, str):
-            raise UpdateAttributeError(f'Unexpected response from server: {updated_attribute}')
+            raise PyMISPUnexpectedResponse(f'Unexpected response from server: {updated_attribute}')
         elif 'errors' in updated_attribute:
             return updated_attribute
         a = MISPAttribute()

diff --git a/tests/testlive_comprehensive.py b/tests/testlive_comprehensive.py
@@ -911,10 +911,16 @@ def test_update_object(self):
         ip_dom.add_attribute('ip', value='8.8.8.8')
         first.add_object(ip_dom)
         try:
+            # Update with full event
             first = self.user_misp_connector.add_event(first)
             first.objects[0].add_attribute('ip', value='8.9.9.8')
             first = self.user_misp_connector.update_event(first)
             self.assertEqual(first.objects[0].attributes[2].value, '8.9.9.8')
+            # Update object only
+            misp_object = self.user_misp_connector.get_object(first.objects[0].id)
+            misp_object.attributes[2].value = '8.9.9.9'
+            misp_object = self.user_misp_connector.update_object(misp_object)
+            self.assertEqual(misp_object.attributes[2].value, '8.9.9.9')
         finally:
             # Delete event
             self.admin_misp_connector.delete_event(first.id)