AttackProtector: add handling of the exempt capability

AttackProtector/config.py

@@ -76,6 +76,9 @@ AttackProtector = conf.registerPlugin('AttackProtector')
 # conf.registerGlobalValue(AttackProtector, 'someConfigVariableName',
 #     registry.Boolean(False, """Help for someConfigVariableName."""))
 
+conf.registerGlobalValue(AttackProtector, 'exempt',
+    registry.String('nopunish', _("""If a user has this capability, he won't be
+        punished by AttackProtector""")))
 conf.registerGlobalValue(AttackProtector, 'delay',
     registry.Integer(10, _("""Determines how long (in seconds) the plugin will
     wait before being enabled. A too low value makes the bot believe that

AttackProtector/plugin.py

@@ -32,6 +32,7 @@ import re
 import time
 
 import supybot.utils as utils
+import supybot.ircdb as ircdb
 from supybot.commands import *
 import supybot.plugins as plugins
 import supybot.ircmsgs as ircmsgs
@@ -160,6 +161,15 @@ class AttackProtector(callbacks.Plugin):
         prefix = lastItem.prefix
         nick = prefix.split('!')[0]
         kind = lastItem.kind
+
+        try:
+            ircdb.users.getUser(msg.prefix) # May raise KeyError
+            capability = self.registryValue('exempt')
+            if capability:
+                if ircdb.checkCapability(msg.prefix, capability):
+                    return
+        except KeyError:
+            pass
         punishment = self.registryValue('%s.punishment' % kind, channel)
         reason = _('%s flood detected') % kind
         if punishment == 'kick':

AttackProtector/test.py

@@ -30,9 +30,10 @@
 
 import time
 from supybot.test import *
+import supybot.ircdb as ircdb
 
 class AttackProtectorTestCase(ChannelPluginTestCase):
-    plugins = ('AttackProtector', 'Utilities')
+    plugins = ('AttackProtector', 'Utilities', 'User')
     config = {'supybot.plugins.AttackProtector.join.detection': '5p2',
               'supybot.plugins.AttackProtector.part.punishment':
               'command echo hi !'}
@@ -169,6 +170,29 @@ class AttackProtectorTestCase(ChannelPluginTestCase):
         self.failIf(self._getIfAnswerIsThisKick('message') == False,
                     'No reaction to both notice and privmsg flood.')
 
+    #################################
+    # Test trusted users
+    def testDoesNotPunishTrustedUsers(self):
+        feedMsg = PluginTestCase.feedMsg
+        feedMsg(self, 'register toto foobarbaz')
+        feedMsg(self, 'identify toto foobarbaz')
+        [self.irc.takeMsg() for x in 'xx']
+        self.assertNotError('eval '
+                '''"ircdb.users.getUser(1).capabilities.add('nopunish')"''')
+        self.assertResponse('capabilities', '[nopunish]')
+
+        try:
+            for i in range(1, 5):
+                msg = ircmsgs.join(self.channel, prefix=self.prefix)
+                self.irc.feedMsg(msg)
+            self.failIf(self._getIfAnswerIsThisBan(), 'Punishes trusted user')
+        finally:
+            feedMsg(self, 'hostmask remove toto %s' % self.prefix)
+            feedMsg(self, 'unidentify') # Otherwise, other tests would fail
+            [self.irc.takeMsg() for x in 'xx']
+            self.assertNotRegexp('whoami', 'toto')
+            self.assertError('capabilities')
+
     #################################
     # Global tests
     def testCleanCollection(self):