From a79354cd525413417278bd7ca37138c604a11b3d Mon Sep 17 00:00:00 2001
From: James Lu <james@overdrivenetworks.com>
Date: Sun, 22 Dec 2019 23:52:02 -0800
Subject: [PATCH] apply_modes: fix removing multiple ban modes in one command

---
 classes.py                    |  9 +++++----
 test/protocol_test_fixture.py | 16 ++++++++++++++++
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/classes.py b/classes.py
index 4d4c19e..5f9e875 100644
--- a/classes.py
+++ b/classes.py
@@ -1106,10 +1106,11 @@ class PyLinkNetworkCoreWithUtils(PyLinkNetworkCore):
 
                 existing = mapping.get(real_mode[0])
                 arg = real_mode[1]
-                # When removing a mode: check all existing modes with the same character
-                # If no args were needed on removal, remove all modes with that letter
-                # If an arg was given, remove all modes matching the arg (IRC case insensitive)
-                if existing is not None:
+                # Mode requires argument for removal (case insensitive)
+                if real_mode[0] in (supported_modes['*A'] + supported_modes['*B']):
+                    modelist.discard((real_mode[0], self.to_lower(arg)))
+                # Mode does not require argument for removal - remove all modes entries with the same character
+                else:
                     while existing:
                         oldvalue = existing.pop()
                         if arg is None or self.to_lower(arg) == self.to_lower(oldvalue):
diff --git a/test/protocol_test_fixture.py b/test/protocol_test_fixture.py
index 99ae013..fb3b05b 100644
--- a/test/protocol_test_fixture.py
+++ b/test/protocol_test_fixture.py
@@ -575,6 +575,13 @@ class BaseProtocolTest(unittest.TestCase):
         self.p.apply_modes('#', [('+n', None), ('-i', None)])
         self.assertEqual(c.modes, {('n', None)})
 
+        c = self.p.channels['#Magic'] = Channel(self.p, name='#Magic')
+        self.p.apply_modes('#Magic', [('+m', None), ('+n', None), ('+i', None)])
+        self.assertEqual(c.modes, {('m', None), ('n', None), ('i', None)}, "Modes should be added")
+
+        self.p.apply_modes('#Magic', [('-i', None), ('-n', None)])
+        self.assertEqual(c.modes, {('m', None)}, "Modes should be removed")
+
     def test_apply_modes_channel_remove_nonexistent(self):
         c = self.p.channels['#abc'] = Channel(self.p, name='#abc')
         self.p.apply_modes('#abc', [('+t', None)])
@@ -634,6 +641,15 @@ class BaseProtocolTest(unittest.TestCase):
         self.p.apply_modes('#Magic', [('-b', '*!*@BEST.HOST')])
         self.assertFalse(c.modes, "Ban on *!*@best.host should be removed (different case)")
 
+    def test_apply_modes_channel_ban_multiple(self):
+        c = self.p.channels['#Magic'] = Channel(self.p, name='#Magic')
+        self.p.apply_modes('#Magic', [('+b', '*!*@test.host'), ('+b', '*!*@best.host'), ('+b', '*!*@guest.host')])
+        self.assertEqual(c.modes, {('b', '*!*@test.host'), ('b', '*!*@best.host'), ('b', '*!*@guest.host')},
+                         "Bans should be added")
+
+        self.p.apply_modes('#Magic', [('-b', '*!*@best.host'), ('-b', '*!*@guest.host'), ('-b', '*!*@test.host')])
+        self.assertEqual(c.modes, set(), "Bans should be removed")
+
     def test_apply_modes_channel_mode_cycle(self):
         c = self.p.channels['#Magic'] = Channel(self.p, name='#Magic')
         self.p.apply_modes('#Magic', [('+b', '*!*@example.net'), ('-b', '*!*@example.net')])