From da0f6ef7d241649e86d8356d80c45a9def22fe02 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Tue, 25 Sep 2018 09:32:17 -0400
Subject: [PATCH] new: Add test for warninglists

---
 pymisp/aping.py                 | 14 ++++++++++--
 tests/testlive_comprehensive.py | 39 ++++++++++++++++-----------------
 2 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/pymisp/aping.py b/pymisp/aping.py
index 7478e8d..641b192 100644
--- a/pymisp/aping.py
+++ b/pymisp/aping.py
@@ -35,12 +35,22 @@ class ExpandedPyMISP(PyMISP):
             to_return['OR'] = or_parameters
         return to_return
 
-    def toggle_warninglist(self, warninglist_id: int, force_enable: bool=None):
+    def toggle_warninglist(self, warninglist_id: List[int]=None, warninglist_name: List[str]=None, 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 warninglist_id is None and warninglist_name is None:
+            raise Exception('Either warninglist_id or warninglist_name is required.')
+        query = {}
+        if warninglist_id is not None:
+            if not isinstance(warninglist_id, list):
+                warninglist_id = [warninglist_id]
+            query['id'] = warninglist_id
+        if warninglist_name is not None:
+            if not isinstance(warninglist_name, list):
+                warninglist_name = [warninglist_name]
+            query['name'] = warninglist_name
         if force_enable is not None:
             query['enabled'] = force_enable
         url = urljoin(self.root_url, '/warninglists/toggleEnable')
diff --git a/tests/testlive_comprehensive.py b/tests/testlive_comprehensive.py
index 96bd46b..594ce39 100644
--- a/tests/testlive_comprehensive.py
+++ b/tests/testlive_comprehensive.py
@@ -580,29 +580,28 @@ class TestComprehensive(unittest.TestCase):
             # self.assertEqual(len(events), 1)
 
             # warninglist
-            # FIXME: the warning lists ID aren't deterministic
-            if local:
-                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)
+            response = self.admin_misp_connector.toggle_warninglist(warninglist_name='%dns resolv%', force_enable=True)  # enable ipv4 DNS.
+            # response = self.admin_misp_connector.toggle_warninglist(warninglist_id=[17], force_enable=True)  # enable ipv4 DNS.
+            self.assertDictEqual(response, {'saved': True, 'success': '3 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, 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=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'})
+            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(warninglist_name='%dns resolv%')  # disable ipv4 DNS.
+            self.assertDictEqual(response, {'saved': True, 'success': '3 warninglist(s) toggled'})
 
             time.sleep(1)
             # attachments