Track user channel modes by user object, not nickname

4 files changed, 24 insertions, 9 deletions

diff --git a/IRCChannel.py b/IRCChannel.py
index 3ecfb110..2cc40654 100644
--- a/IRCChannel.py
+++ b/IRCChannel.py
@@ -28,20 +28,35 @@ class Channel(object):
         self.users.add(user)
     def remove_user(self, user):
         self.users.remove(user)
+        for mode in self.modes[:]:
+            if mode in self.sever.mode_prefixes and user in self.modes[mode]:
+                self.modes[mode].discard(user)
+                if not len(self.modes[mode]):
+                    del self.modes[mode]
     def has_user(self, user):
         return user in self.users
     def add_mode(self, mode, args=None):
         if not mode in self.modes:
             self.modes[mode] = set([])
         if args:
-            self.modes[mode].add(args.lower())
+            if mode in self.server.mode_prefixes:
+                user = self.server.get_user(args)
+                if user:
+                    self.modes[mode].add(user)
+            else:
+                self.modes[mode].add(args.lower())
         self.bot.events.on("mode").on("channel").call(
             channel=self, mode=mode, args=args, remove=False)
     def remove_mode(self, mode, args=None):
         if not args:
             del self.modes[mode]
         else:
-            self.modes[mode].discard(args.lower())
+            if mode in self.server.mode_prefixes:
+                user = self.server.get_user(args)
+                if user:
+                    self.modes[mode].discard(user)
+            else:
+                self.modes[mode].discard(args.lower())
             if not len(self.modes[mode]):
                 del self.modes[mode]
         self.bot.events.on("mode").on("channel").call(
@@ -68,17 +83,17 @@ class Channel(object):
     def send_ban(self, hostmask):
         self.server.send_mode(self.name, "+b", hostmask)
 
-    def mode_or_above(self, nickname, mode):
+    def mode_or_above(self, user, mode):
         mode_orders = list(self.server.mode_prefixes.values())
         mode_index = mode_orders.index(mode)
         for mode in mode_orders[:mode_index+1]:
-            if nickname.lower() in self.modes.get(mode, []):
+            if user in self.modes.get(mode, []):
                 return True
         return False
 
-    def get_user_status(self, nickname):
+    def get_user_status(self, user):
         modes = ""
         for mode in self.server.mode_prefixes.values():
-            if nickname.lower() in self.modes.get(mode, []):
+            if user in self.modes.get(mode, []):
                 modes += mode
         return modes
diff --git a/modules/channel_op.py b/modules/channel_op.py
index 03d97426..1093ad10 100644
--- a/modules/channel_op.py
+++ b/modules/channel_op.py
@@ -90,7 +90,7 @@ class Module(object):
         if len(set(nicknames) & set(event["message_split"])) >= event["channel"].get_setting(
                 "highlight-spam-threshold", 10):
             protection_enabled = event["channel"].get_setting("highlight-spam-protection", False)
-            has_mode = event["channel"].mode_or_above(event["user"].nickname, "v")
+            has_mode = event["channel"].mode_or_above(event["user"], "v")
             should_ban = event["channel"].get_setting("highlight-spam-ban", False)
             if protection_enabled and not has_mode:
                 if should_ban:
diff --git a/modules/check_mode.py b/modules/check_mode.py
index 97da52d7..a7a8e788 100644
--- a/modules/check_mode.py
+++ b/modules/check_mode.py
@@ -10,6 +10,6 @@ class Module(object):
         if event["is_channel"] and event["hook"].kwargs.get(
                 "require_mode"):
             required_mode = event["hook"].kwargs.get("require_mode")[0]
-            if not event["target"].mode_or_above(event["user"].nickname,
+            if not event["target"].mode_or_above(event["user"],
                     required_mode):
                 return "You do not have permission to do this"
diff --git a/modules/set.py b/modules/set.py
index ddc2a758..9b0678a9 100644
--- a/modules/set.py
+++ b/modules/set.py
@@ -50,7 +50,7 @@ class Module(object):
         self._set(self.settings, event, event["user"])
 
     def channel_set(self, event):
-        if event["target"].mode_or_above(event["user"].nickname,
+        if event["target"].mode_or_above(event["user"],
                 "o"):
             self._set(self.channel_settings, event, event["target"])
         else: