From 2c03fb96c25a3fd09b1243476cb8f6544e2ef875 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Fri, 17 Aug 2018 15:09:17 +0200
Subject: [PATCH] new: [test] Attribute modification

---
 pymisp/mispevent.py             |  2 +
 tests/testlive_comprehensive.py | 83 +++++++++++++++++++++------------
 2 files changed, 56 insertions(+), 29 deletions(-)

diff --git a/pymisp/mispevent.py b/pymisp/mispevent.py
index a8092f1..8cb342e 100644
--- a/pymisp/mispevent.py
+++ b/pymisp/mispevent.py
@@ -163,6 +163,8 @@ class MISPAttribute(AbstractMISP):
         return misp_shadow_attribute
 
     def from_dict(self, **kwargs):
+        if kwargs.get('Attribute'):
+            kwargs = kwargs.get('Attribute')
         if kwargs.get('type') and kwargs.get('category'):
             if kwargs['type'] not in self.__category_type_mapping[kwargs['category']]:
                 if self.__strict:
diff --git a/tests/testlive_comprehensive.py b/tests/testlive_comprehensive.py
index 8bfe8ed..3dcd8bf 100644
--- a/tests/testlive_comprehensive.py
+++ b/tests/testlive_comprehensive.py
@@ -3,7 +3,7 @@
 
 import unittest
 
-from pymisp import ExpandedPyMISP, MISPEvent, MISPOrganisation, MISPUser, Distribution, ThreatLevel, Analysis
+from pymisp import ExpandedPyMISP, MISPEvent, MISPOrganisation, MISPUser, Distribution, ThreatLevel, Analysis, MISPAttribute
 from datetime import datetime, timedelta
 
 import time
@@ -85,37 +85,41 @@ class TestComprehensive(unittest.TestCase):
 
     def test_search_value_attribute(self):
         me = self.create_event_org_only()
-        # Create event
-        created_event = self.admin_misp_connector.add_event(me)
-        c_me = MISPEvent()
-        c_me.load(created_event)
-        # Search as admin
-        response = self.admin_misp_connector.search(controller='attributes', value=me.attributes[0].value)
-        self.assertEqual(len(response), 1)
-        # Connect as user
-        user_misp_connector = ExpandedPyMISP(url, self.test_usr.authkey)
-        # Search as user
-        response = user_misp_connector.search(controller='attributes', value=me.attributes[0].value)
-        self.assertEqual(response, [])
-        # Delete event
-        self.admin_misp_connector.delete_event(c_me.id)
+        try:
+            # Create event
+            created_event = self.admin_misp_connector.add_event(me)
+            c_me = MISPEvent()
+            c_me.load(created_event)
+            # Search as admin
+            response = self.admin_misp_connector.search(controller='attributes', value=me.attributes[0].value)
+            self.assertEqual(len(response), 1)
+            # Connect as user
+            user_misp_connector = ExpandedPyMISP(url, self.test_usr.authkey)
+            # Search as user
+            response = user_misp_connector.search(controller='attributes', value=me.attributes[0].value)
+            self.assertEqual(response, [])
+        finally:
+            # Delete event
+            self.admin_misp_connector.delete_event(c_me.id)
 
     def test_search_tag_event(self):
         me = self.create_event_with_tags()
-        # Create event
-        created_event = self.admin_misp_connector.add_event(me)
-        c_me = MISPEvent()
-        c_me.load(created_event)
-        # Search as admin
-        response = self.admin_misp_connector.search(tags='tlp:white___test')
-        self.assertEqual(len(response), 1)
-        # Connect as user
-        user_misp_connector = ExpandedPyMISP(url, self.test_usr.authkey)
-        # Search as user
-        response = user_misp_connector.search(value='tlp:white___test')
-        self.assertEqual(response, [])
-        # Delete event
-        self.admin_misp_connector.delete_event(c_me.id)
+        try:
+            # Create event
+            created_event = self.admin_misp_connector.add_event(me)
+            c_me = MISPEvent()
+            c_me.load(created_event)
+            # Search as admin
+            response = self.admin_misp_connector.search(tags='tlp:white___test')
+            self.assertEqual(len(response), 1)
+            # Connect as user
+            user_misp_connector = ExpandedPyMISP(url, self.test_usr.authkey)
+            # Search as user
+            response = user_misp_connector.search(value='tlp:white___test')
+            self.assertEqual(response, [])
+        finally:
+            # Delete event
+            self.admin_misp_connector.delete_event(c_me.id)
 
     @unittest.skip("currently failing")
     def test_search_tag_event_fancy(self):
@@ -301,6 +305,27 @@ class TestComprehensive(unittest.TestCase):
             # Delete event
             self.admin_misp_connector.delete_event(first_to_delete.id)
 
+    def test_edit_attribute(self):
+        first = self.create_event_org_only()
+        user_misp_connector = ExpandedPyMISP(url, self.test_usr.authkey, debug=False)
+        try:
+            first.attributes[0].comment = 'This is the original comment'
+            first_created_event = user_misp_connector.add_event(first)
+            first_to_delete = MISPEvent()
+            first_to_delete.load(first_created_event)
+            first_to_delete.attributes[0].comment = 'This is the modified comment'
+            response = user_misp_connector.update_attribute(first_to_delete.attributes[0].id, first_to_delete.attributes[0])
+            tmp_attr = MISPAttribute()
+            tmp_attr.from_dict(**response)
+            self.assertEqual(tmp_attr.comment, 'This is the modified comment')
+            response = user_misp_connector.change_comment(first_to_delete.attributes[0].uuid, 'This is the modified comment, again')
+            tmp_attr = MISPAttribute()
+            tmp_attr.from_dict(**response)
+            self.assertEqual(tmp_attr.comment, 'This is the modified comment, again')
+        finally:
+            # Delete event
+            self.admin_misp_connector.delete_event(first_to_delete.id)
+
     @unittest.skip("currently failing")
     def test_search_tag_attribute(self):
         me = self.create_event_with_tags()