class AntiSpam: def __init__(self): self.whitelist = (list(SUDO_USERS) or []) + [OWNER_ID] Duration.CUSTOM = 15 # Custom duration, 15 seconds self.sec_limit = RequestRate(6, Duration.CUSTOM) # 6 / Per 15 Seconds self.min_limit = RequestRate(20, Duration.MINUTE) # 20 / Per minute self.hour_limit = RequestRate(100, Duration.HOUR) # 100 / Per hour self.daily_limit = RequestRate(1000, Duration.DAY) # 1000 / Per day self.limiter = Limiter( self.sec_limit, self.min_limit, self.hour_limit, self.daily_limit, bucket_class=MemoryListBucket) def check_user(self, user): """ Return True if user is to be ignored else False """ if user in self.whitelist: return False try: self.limiter.try_acquire(user) return False except BucketFullException: return True
class AntiSpam: def __init__(self): self.whitelist = (DEV_USERS or []) + (DRAGONS or []) + ( WOLVES or []) + (DEMONS or []) + (TIGERS or []) #Values are HIGHLY experimental, its recommended you pay attention to our commits as we will be adjusting the values over time with what suits best. Duration.CUSTOM = 15 # Custom duration, 15 seconds self.sec_limit = RequestRate(6, Duration.CUSTOM) # 6 / Per 15 Seconds self.min_limit = RequestRate(20, Duration.MINUTE) # 20 / Per minute self.hour_limit = RequestRate(100, Duration.HOUR) # 100 / Per hour self.daily_limit = RequestRate(1000, Duration.DAY) # 1000 / Per day self.limiter = Limiter(self.sec_limit, self.min_limit, self.hour_limit, self.daily_limit, bucket_class=MemoryListBucket) def check_user(self, user): """ Return True if user is to be ignored else False """ if user in self.whitelist: return False try: self.limiter.try_acquire(user) return False except BucketFullException: return True
def limit(user): rate = RequestRate(2, Duration.DAY) limiter = Limiter(rate, bucket_class=MemoryListBucket) identity = user try: limiter.try_acquire(identity) return True except BucketFullException as err: return False
def test_simple_01(time_function): """Single-rate Limiter with RedisBucket""" rate = RequestRate(3, 5 * Duration.SECOND) limiter = Limiter( rate, bucket_class=RedisBucket, # Separate buckets used to distinct values from previous run, # as time_function return value has different int part. bucket_kwargs={ "redis_pool": pool, "bucket_name": str(time_function) }, time_function=time_function, ) item = "vutran_list" with pytest.raises(BucketFullException): for _ in range(4): limiter.try_acquire(item) sleep(6) limiter.try_acquire(item) vol = limiter.get_current_volume(item) assert vol == 1 limiter.try_acquire(item) limiter.try_acquire(item) with pytest.raises(BucketFullException): limiter.try_acquire(item)
def test_simple_01(): """ Single-rate Limiter with RedisBucket """ rate = RequestRate(3, 5 * Duration.SECOND) limiter = Limiter( rate, bucket_class=RedisBucket, bucket_kwargs={ "redis_pool": pool, "bucket_name": "test-bucket-1" }, ) item = 'vutran_list' with pytest.raises(BucketFullException): for _ in range(4): limiter.try_acquire(item) sleep(6) limiter.try_acquire(item) vol = limiter.get_current_volume(item) assert vol == 1 limiter.try_acquire(item) limiter.try_acquire(item) with pytest.raises(BucketFullException): limiter.try_acquire(item)
def test_remaining_time(time_function): """The remaining_time metadata returned from a BucketFullException should take into account the time elapsed during limited calls (including values less than 1 second). """ limiter2 = Limiter(RequestRate(5, Duration.SECOND), time_function=time_function) for _ in range(5): limiter2.try_acquire("item") sleep(0.1) try: limiter2.try_acquire("item") except BucketFullException as err: delay_time = err.meta_info["remaining_time"] assert 0.8 < delay_time < 0.9
def test_sleep(time_function): """Make requests at a rate of 6 requests per 5 seconds (avg. 1.2 requests per second). If each request takes ~0.5 seconds, then the bucket should be full after 6 requests (3 seconds). Run 15 requests, and expect a total of 2 delays required to stay within the rate limit. """ rate = RequestRate(6, 5 * Duration.SECOND) limiter = Limiter(rate, time_function=time_function) track_sleep = Mock(side_effect=sleep) # run time.sleep() and track the number of calls start = time() for i in range(15): try: limiter.try_acquire("test") print(f"[{time() - start:07.4f}] Pushed: {i+1} items") sleep(0.5) # Simulated request rate except BucketFullException as err: print(err.meta_info) track_sleep(err.meta_info["remaining_time"]) print(f"Elapsed: {time() - start:07.4f} seconds") assert track_sleep.call_count == 2
def test_flushing(): """Multi-rates Limiter with RedisBucket""" rate_1 = RequestRate(5, 5 * Duration.SECOND) limiter = Limiter( rate_1, bucket_class=RedisBucket, bucket_kwargs={ "redis_pool": pool, "bucket_name": "Flushing-Bucket", }, ) item = "redis-test-item" for _ in range(3): limiter.try_acquire(item) size = limiter.get_current_volume(item) assert size == 3 assert limiter.flush_all() == 1 size = limiter.get_current_volume(item) assert size == 0
def test_acquire(time_function): # Separate buckets used to distinct values from previous run, # as time_function return value has different int part. bucket_kwargs = { "bucket_name": str(time_function), "redis_pool": redis_connection, } limiter = Limiter( rate, bucket_class=SCRedisBucket, bucket_kwargs=bucket_kwargs, time_function=time_function, ) try: for _ in range(5): limiter.try_acquire("item-id") sleep(2) except BucketFullException as err: print(err.meta_info) assert round(err.meta_info["remaining_time"]) == 6
def test_simple_01(): """Single-rate Limiter""" with pytest.raises(InvalidParams): # No rates provided Limiter() with pytest.raises(InvalidParams): rate_1 = RequestRate(3, 5 * Duration.SECOND) rate_2 = RequestRate(4, 5 * Duration.SECOND) Limiter(rate_1, rate_2) rate = RequestRate(3, 5 * Duration.SECOND) with pytest.raises(ImmutableClassProperty): rate.limit = 10 with pytest.raises(ImmutableClassProperty): rate.interval = 10 limiter = Limiter(rate) item = "vutran" has_raised = False try: for _ in range(4): limiter.try_acquire(item) sleep(1) except BucketFullException as err: has_raised = True print(err) assert str(err) assert isinstance(err.meta_info, dict) assert 1.9 < err.meta_info["remaining_time"] < 2.0 assert has_raised sleep(6) limiter.try_acquire(item) vol = limiter.get_current_volume(item) assert vol == 1 limiter.try_acquire(item) limiter.try_acquire(item) with pytest.raises(BucketFullException): limiter.try_acquire(item)
def test_simple_03(): """Single-rate Limiter with MemoryListBucket""" rate = RequestRate(3, 5 * Duration.SECOND) limiter = Limiter(rate, bucket_class=MemoryListBucket) item = "vutran_list" with pytest.raises(BucketFullException): for _ in range(4): limiter.try_acquire(item) sleep(6) limiter.try_acquire(item) vol = limiter.get_current_volume(item) assert vol == 1 limiter.try_acquire(item) limiter.try_acquire(item) with pytest.raises(BucketFullException): limiter.try_acquire(item)
def test_simple_01(): """ Single-rate Limiter """ rate = RequestRate(3, 5 * Duration.SECOND) limiter = Limiter(rate) item = 'vutran' with pytest.raises(BucketFullException): for _ in range(4): limiter.try_acquire(item) sleep(6) limiter.try_acquire(item) vol = limiter.get_current_volume(item) assert vol == 1 limiter.try_acquire(item) limiter.try_acquire(item) with pytest.raises(BucketFullException): limiter.try_acquire(item)
def test_simple_02(): """Multi-rates Limiter""" rate_1 = RequestRate(5, 5 * Duration.SECOND) rate_2 = RequestRate(7, 9 * Duration.SECOND) limiter2 = Limiter(rate_1, rate_2) item = "tranvu" err = None with pytest.raises(BucketFullException) as err: # Try add 6 items within 5 seconds # Exceed Rate-1 for _ in range(6): limiter2.try_acquire(item) print(err.value.meta_info) assert limiter2.get_current_volume(item) == 5 sleep(6) # Still shorter than Rate-2 interval, so all items must be kept limiter2.try_acquire(item) limiter2.try_acquire(item) assert limiter2.get_current_volume(item) == 7 with pytest.raises(BucketFullException) as err: # Exceed Rate-2 limiter2.try_acquire(item) print(err.value.meta_info) sleep(6) # 12 seconds passed limiter2.try_acquire(item) # Only items within last 9 seconds kept, plus the new one assert limiter2.get_current_volume(item) == 3 # print('Bucket Rate-1:', limiter2.get_filled_slots(rate_1, item)) # print('Bucket Rate-2:', limiter2.get_filled_slots(rate_2, item)) # Within the nearest 5 second interval # Rate-1 has only 1 item, so we can add 4 more limiter2.try_acquire(item) limiter2.try_acquire(item) limiter2.try_acquire(item) limiter2.try_acquire(item) with pytest.raises(BucketFullException): # Exceed Rate-1 again limiter2.try_acquire(item) # Withint the nearest 9 second-interval, we have 7 items assert limiter2.get_current_volume(item) == 7 # Fast forward to 6 more seconds # Bucket Rate-1 is refreshed and empty by now # Bucket Rate-2 has now only 5 items sleep(6) # print('Bucket Rate-1:', limiter2.get_filled_slots(rate_1, item)) # print('Bucket Rate-2:', limiter2.get_filled_slots(rate_2, item)) limiter2.try_acquire(item) limiter2.try_acquire(item) with pytest.raises(BucketFullException): # Exceed Rate-2 again limiter2.try_acquire(item) assert limiter2.flush_all() == 1 assert limiter2.get_current_volume(item) == 0
def test_simple_02(time_function): """Multi-rates Limiter with RedisBucket""" rate_1 = RequestRate(5, 5 * Duration.SECOND) rate_2 = RequestRate(7, 9 * Duration.SECOND) limiter4 = Limiter( rate_1, rate_2, bucket_class=RedisBucket, bucket_kwargs={ "redis_pool": pool, # Separate buckets used to distinct values from previous run, # as time_function return value has different int part. "bucket_name": str(time_function), }, time_function=time_function, ) item = "redis-test-item" with pytest.raises(BucketFullException): # Try add 6 items within 5 seconds # Exceed Rate-1 for _ in range(6): limiter4.try_acquire(item) assert limiter4.get_current_volume(item) == 5 sleep(6.5) # Still shorter than Rate-2 interval, so all items must be kept limiter4.try_acquire(item) # print('Bucket Rate-1:', limiter4.get_filled_slots(rate_1, item)) # print('Bucket Rate-2:', limiter4.get_filled_slots(rate_2, item)) limiter4.try_acquire(item) assert limiter4.get_current_volume(item) == 7 with pytest.raises(BucketFullException): # Exceed Rate-2 limiter4.try_acquire(item) sleep(6) # 12 seconds passed limiter4.try_acquire(item) # Only items within last 9 seconds kept, plus the new one assert limiter4.get_current_volume(item) == 3 # print('Bucket Rate-1:', limiter4.get_filled_slots(rate_1, item)) # print('Bucket Rate-2:', limiter4.get_filled_slots(rate_2, item)) # Within the nearest 5 second interval # Rate-1 has only 1 item, so we can add 4 more limiter4.try_acquire(item) limiter4.try_acquire(item) limiter4.try_acquire(item) limiter4.try_acquire(item) with pytest.raises(BucketFullException): # Exceed Rate-1 again limiter4.try_acquire(item) # Withint the nearest 9 second-interval, we have 7 items assert limiter4.get_current_volume(item) == 7 # Fast forward to 6 more seconds # Bucket Rate-1 is refreshed and empty by now # Bucket Rate-2 has now only 5 items sleep(6) # print('Bucket Rate-1:', limiter4.get_filled_slots(rate_1, item)) # print('Bucket Rate-2:', limiter4.get_filled_slots(rate_2, item)) limiter4.try_acquire(item) limiter4.try_acquire(item) with pytest.raises(BucketFullException): # Exceed Rate-2 again limiter4.try_acquire(item)
class ModPlus(commands.Cog): """Ultimate Moderation Cog for RedBot""" def __init__(self, bot): self.bot = bot self.config = Config.get_conf(self, 8818154, force_registration=True) # PyRateLimit.init(redis_host="localhost", redis_port=6379) hourly_rate5 = RequestRate(5, Duration.HOUR) hourly_rate3 = RequestRate(3, Duration.HOUR) self.kicklimiter = Limiter(hourly_rate5) self.banlimiter = Limiter(hourly_rate3) # self.kicklimit = PyRateLimit() # self.kicklimit.create(3600, 5) # self.banlimit = PyRateLimit() # self.banlimit.create(3600, 3) default_global = { 'notifs': { 'kick': [], 'ban': [], 'mute': [], 'jail': [], 'channelperms': [], 'editchannel': [], 'deletemessages': [], 'ratelimit': [], 'adminrole': [], 'bot':[], 'warn':[] }, 'notifchannels' : { 'kick': [], 'ban': [], 'mute': [], 'jail': [], 'channelperms': [], 'editchannel': [], 'deletemessages': [], 'ratelimit': [], 'adminrole': [], 'bot':[], 'warn':[] } } default_guild = { 'perms': { 'kick': [], 'ban': [], 'mute': [], 'jail': [], 'channelperms': [], 'editchannel': [], 'deletemessages': [], 'warn': [] }, 'roles': { 'warning1': None, 'warning2': None, 'warning3+': None, 'jailed': None, 'muted': None } } self.config.register_guild(**default_guild) self.config.register_global(**default_global) self.permkeys = [ 'kick', 'ban', 'mute', 'jail', 'channelperms', 'editchannel', 'deletemessages', 'warn' ] self.notifkeys = [ 'kick', 'ban', 'mute', 'jail', 'channelperms', 'editchannel', 'deletemessages', 'ratelimit', 'adminrole', 'bot', 'warn' ] self.rolekeys = [ 'warning1', 'warning2', 'warning3+', 'jailed', 'muted' ] # Notifications Part @commands.group(aliases=['notifs', 'notif']) # CHANNEL @checks.mod() async def adminnotifications(self, ctx): """Configure what notifications to get""" pass @adminnotifications.group(name='channel') async def notifschannel(self, ctx): """Configure a channel to recieve notifications""" pass @adminnotifications.command(name='info') async def notifsinfo(self, ctx): """Get information about notification system""" await ctx.send(NOTIF_SYS_INFO) @adminnotifications.command(name='add') async def notifsadd(self, ctx, notifkey: str, user: discord.Member = None): """Get notified about something""" if user is None: user = ctx.author notifkey = notifkey.strip().lower() if notifkey not in self.notifkeys: return await ctx.send(ERROR_MESSAGES['NOTIF_UNRECOGNIZED']) data = await self.config.notifs() if user.id in data[notifkey]: return await ctx.send(f"{user.display_name} is already getting notified about {notifkey}") data[notifkey].append(user.id) await self.config.notifs.set(data) return await ctx.send(f"{user.display_name} will now be notified on {notifkey}") @adminnotifications.command(name='remove') async def notifsremove(self, ctx, notifkey: str, user: discord.Member = None): """Stop getting notified about something""" if user is None: user = ctx.author notifkey = notifkey.strip().lower() if notifkey not in self.notifkeys: return await ctx.send(ERROR_MESSAGES['NOTIF_UNRECOGNIZED']) data = await self.config.notifs() if user.id not in data[notifkey]: return await ctx.send(f"{user.display_name} isn't currently getting notified about {notifkey}") data[notifkey].remove(user.id) await self.config.notifs.set(data) return await ctx.send(f"{user.display_name} will now stop being notified about {notifkey}") # SHOW NOTIFICATIONS @adminnotifications.command(name='list') async def notifslist(self, ctx, user: discord.Member = None): """Show which notifications you / a user has enabled""" if user is None: user = ctx.author data = await self.config.notifs() notifs = [] for notif in data: if user.id in data[notif]: notifs.append(notif) await ctx.send(f'{user.display_name} is getting notified for the following: ' + ', '.join(notifs)) # Channel notifications @notifschannel.command(name='add') async def channelnotifsadd(self, ctx, notifkey: str, channel: discord.TextChannel): """Get notified about something (channel)""" notifkey = notifkey.strip().lower() if notifkey not in self.notifkeys: return await ctx.send(ERROR_MESSAGES['NOTIF_UNRECOGNIZED']) data = await self.config.notifchannels() channeldata = [channel.guild.id, channel.id] if channeldata in data[notifkey]: return await ctx.send(f"{channel.name} is already getting notified about {notifkey}") data[notifkey].append(channeldata) await self.config.notifchannels.set(data) return await ctx.send(f"{channel.name} will now be notified on {notifkey}") @notifschannel.command(name='remove') async def channelnotifsremove(self, ctx, notifkey: str, channel: discord.TextChannel): """Stop getting notified about something (channel)""" notifkey = notifkey.strip().lower() if notifkey not in self.notifkeys: return await ctx.send(ERROR_MESSAGES['NOTIF_UNRECOGNIZED']) data = await self.config.notifchannels() channeldata = [channel.guild.id, channel.id] if channeldata not in data[notifkey]: return await ctx.send(f"{channel.name} isn't currently getting notified about {notifkey}") data[notifkey].remove(channeldata) await self.config.notifchannels.set(data) return await ctx.send(f"{channel.name} will now stop being notified about {notifkey}") @notifschannel.command(name='list') async def channelnotifslist(self, ctx, channel: discord.TextChannel): """Show which notifications a channel has enabled""" data = await self.config.notifchannels() channeldata = [channel.guild.id, channel.id] notifs = [] for notif in data: if channeldata in data[notif]: notifs.append(notif) await ctx.send(f'{channel.name} is getting notified for the following: ' + ', '.join(notifs)) # NOTIFY FUNCTION async def notify(self, notifkey, payload): data = await self.config.all() for userid in data['notifs'][notifkey]: user: discord.User = await self.bot.fetch_user(userid) try: await user.send(payload) except Exception: pass for channel in data['notifchannels'][notifkey]: guild: discord.guild = self.bot.get_guild(channel[0]) if guild is not None: txtchannel = guild.get_channel(channel[1]) try: await txtchannel.send(payload, allowed_mentions=discord.AllowedMentions.all()) except Exception: pass # Admin Logging @commands.Cog.listener(name='on_guild_role_update') async def role_add_admin(self, old: discord.Role, new: discord.Role): if new.permissions.administrator and not old.permissions.administrator: await self.notify('adminrole', f'@everyone Role {new.mention}({new.id}) was updated to contain administrator permission. \n IN: {old.guild.name}({old.guild.id})') @commands.Cog.listener(name='on_member_join') async def join_bot(self, member: discord.Member): """Detect if new joining member is a bot""" if member.bot: await self.notify('bot', f'@everyone Role Bot {member.mention}({member.id}) was added. \n IN: {member.guild.name}({member.guild.id})') @commands.Cog.listener(name='on_member_update') async def member_admin(self, old: discord.Member, new: discord.Member): new_roles = [] for role in new.roles: if role not in old.roles: new_roles.append(role) for role in new_roles: if role.permissions.administrator: await self.notify('adminrole', f'@everyone Member{new.mention}({new.id}) was updated to contain administrator permission. \n IN: {old.guild.name}({old.guild.id})') async def rate_limit_exceeded(self, user: discord.Member, type): """Called to removed all moderation roles when a mod has hit ratelimit""" allmodroles = [] data = await self.config.guild(user.guild).perms() for perm in data: for role in data[perm]: if role not in allmodroles: allmodroles.append(role) rm_mention = [] broken = [] issue = False for role in user.roles: if role.id in allmodroles: try: await user.remove_roles(role, reason='Rate limit exceeded.') rm_mention.append(role.mention) rm_mention.append('(' + str(role.id) + ')') except Exception: issue = True broken.append(role.mention) broken.append('(' + str(role.id) + ')') if issue: await self.notify('ratelimit', "Removing roles in the ratelimit below ended in error. The user has a role above the bot. The following roles could not be removed: " + ', '.join(broken)) payload = f"@everyone {type} ratelimit has been exceeded by {user.mention} ({user.display_name}, {user.id}). The following roles with power have been removed: " + ', '.join(rm_mention) await self.notify('ratelimit', payload) async def action_check(self, ctx, permkey): if await self.bot.is_admin(ctx.author) or await self.bot.is_owner(ctx.author) or ctx.author.guild_permissions.administrator: # Admin auto-bypass return True data = await self.config.guild(ctx.guild).all() canrun = False for role in ctx.author.roles: if role.id in data['perms'][permkey]: canrun = True break if not canrun: return False if permkey == 'kick': try: self.kicklimiter.try_acquire(str(ctx.author.id)) except BucketFullException: await self.rate_limit_exceeded(ctx.author, 'kick') return False elif permkey == 'ban': try: self.banlimiter.try_acquire(str(ctx.author.id)) except BucketFullException: await self.rate_limit_exceeded(ctx.author, 'kick') return False return True @commands.group() @checks.admin() async def modpset(self, ctx): """Configure Mod Plus""" pass @modpset.group(aliases=['perms', 'perm']) async def permissions(self, ctx): """Configure Role Permissions""" pass @permissions.command(name='info') async def permsinfo(self, ctx): """Get info about perms system""" await ctx.send(PERM_SYS_INFO) @permissions.command(name='add') async def permsadd(self, ctx, role: discord.Role, *, permkey: str): """Add Perms to a Role""" permkey = permkey.strip().lower() if permkey not in self.permkeys: return await ctx.send(ERROR_MESSAGES['PERM_UNRECOGNIZED']) data = await self.config.guild(ctx.guild).get_raw("perms", permkey) if role.id in data: return await ctx.send(f"{role.name} already has the {permkey} permission") data.append(role.id) await self.config.guild(ctx.guild).set_raw("perms", permkey, value=data) return await ctx.send(f"{role.name} was sucessfully given the {permkey} permission") @permissions.command(name='remove') async def permsremove(self, ctx, role: discord.Role, *, permkey: str): """Revoke Perms from a Role""" permkey = permkey.strip().lower() if permkey not in self.permkeys: return await ctx.send(ERROR_MESSAGES['PERM_UNRECOGNIZED']) data = await self.config.guild(ctx.guild).get_raw("perms", permkey) if role.id not in data: return await ctx.send(f"{role.name} doesn't have the {permkey} permission") data.remove(role.id) await self.config.guild(ctx.guild).set_raw("perms", permkey, value=data) return await ctx.send(f"{role.name} has sucessfully been revoked the {permkey} permission") @permissions.group(name='list') async def permslist(self, ctx): """List Permissions""" pass @permslist.command(name='perm', aliases=['perms', 'permission']) async def list_perm_by_perm(self, ctx, *, permkey): """List Roles which have the given perm""" permkey = permkey.strip().lower() if permkey not in self.permkeys: return await ctx.send(ERROR_MESSAGES['PERM_UNRECOGNIZED']) data = await self.config.guild(ctx.guild).get_raw("perms", permkey) rolenames = [] for roleid in data: role = ctx.guild.get_role(roleid) if role is None: continue rolenames.append(role.mention) output = f"Roles that have {permkey} permission: " + ', '.join(rolenames) if len(rolenames) == 0: output = f"No Roles have the {permkey} permission." await ctx.send(output) @permslist.command(name='role') async def list_perms_by_role(self, ctx, role: discord.Role): """List which permissions a role has""" data = await self.config.guild(ctx.guild).perms() perms = [] for permkey in data: if role.id in data[permkey]: perms.append(permkey) if len(perms) == 0: await ctx.send(f"{role.name} has no permissions.") else: output = f"{role.name} has the following for permisssions: " + ', '.join(perms) await ctx.send(output)