From b636a320df0ebf7ae4744ad4ec7a7488768d347b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Fri, 21 Sep 2018 13:02:15 -0700
Subject: [PATCH] new: toggle warning list, add test case

---
 pymisp/aping.py                 | 12 ++++++++++++
 tests/testlive_comprehensive.py | 24 ++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/pymisp/aping.py b/pymisp/aping.py
index 790153a..81f3342 100644
--- a/pymisp/aping.py
+++ b/pymisp/aping.py
@@ -35,6 +35,18 @@ class ExpandedPyMISP(PyMISP):
             to_return['OR'] = or_parameters
         return to_return
 
+    def toggle_warninglist(self, warninglist_id: int, force_enable: bool=None):
+        '''Toggle (enable/disable) the status of a warninglist by ID.
+        :param warninglist_id: ID of the WarningList
+        :param force_enable: Force the warning list in the enabled state (does nothing is already enabled)
+        '''
+        query = {'id': warninglist_id}
+        if force_enable is not None:
+            query['enabled'] = force_enable
+        url = urljoin(self.root_url, '/warninglists/toggleEnable')
+        response = self._prepare_request('POST', url, json.dumps(query))
+        return self._check_response(response)
+
     def make_timestamp(self, value: DateTypes):
         if isinstance(value, datetime):
             return datetime.timestamp()
diff --git a/tests/testlive_comprehensive.py b/tests/testlive_comprehensive.py
index e9b8fc3..205fc6e 100644
--- a/tests/testlive_comprehensive.py
+++ b/tests/testlive_comprehensive.py
@@ -580,6 +580,30 @@ class TestComprehensive(unittest.TestCase):
             # events = self.user_misp_connector.search(value='stuff', searchall=True, pythonify=True)
             # self.assertEqual(len(events), 1)
 
+            # warninglist
+            # FIXME: the warning lists ID aren't deterministic
+            response = self.admin_misp_connector.toggle_warninglist('17', force_enable=True)  # enable ipv4 DNS.
+            self.assertDictEqual(response, {'saved': True, 'success': '1 warninglist(s) enabled'})
+            second.add_attribute('ip-src', '9.9.9.9')
+            second = self.user_misp_connector.update_event(second)
+
+            events = self.user_misp_connector.search(eventid=second.id, pythonify=True)
+            self.assertEqual(len(events), 1)
+            self.assertEqual(events[0].id, second.id)
+            self.assertEqual(len(events[0].attributes), 3)
+
+            events = self.user_misp_connector.search(eventid=second.id, enforce_warninglist=False, pythonify=True)
+            self.assertEqual(len(events), 1)
+            self.assertEqual(events[0].id, second.id)
+            self.assertEqual(len(events[0].attributes), 3)
+
+            events = self.user_misp_connector.search(eventid=second.id, enforce_warninglist=True, pythonify=True)
+            self.assertEqual(len(events), 1)
+            self.assertEqual(events[0].id, second.id)
+            self.assertEqual(len(events[0].attributes), 2)
+            response = self.admin_misp_connector.toggle_warninglist('17')  # disable ipv4 DNS.
+            self.assertDictEqual(response, {'saved': True, 'success': '1 warninglist(s) disabled'})
+
             time.sleep(1)
             # attachments
             with open('tests/testlive_comprehensive.py', 'rb') as f: