<commit_before>import logging
import asyncio
import discord
from discord.ext import commands
from dog import Cog, checks, utils
from dog.humantime import HumanTime

logger = logging.getLogger(__name__)

class Mod(Cog):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.mute_tasks = {}

    async def __global_check(self, ctx):
        # do not handle guild-specific command disables in
        # dms
        if not isinstance(ctx.channel, discord.abc.GuildChannel):
            return True

        return not await self.bot.command_is_disabled(ctx.guild, ctx.command.name)

    async def on_message(self, message):
        # do not handle invisibility in dms
        if isinstance(message.channel, discord.abc.PrivateChannel):
            return

        if await self.bot.config_is_set(message.guild, 'invisible_announce'):
            if message.author.status is discord.Status.offline:
                reply = ('Hey {0.mention}! You\'re invisible. Stop being '
                         'invisible, please. Thanks.')
                await message.channel.send(reply.format(message.author))

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_messages=True, read_message_history=True)
    @checks.is_moderator()
    async def purge(self, ctx, amount: int):
        """
        Purges messages the last <n> messages.
        """

        # increase one because the command message they sent needs
        # to be purged, too.
        amount = amount + 1
        if amount > 501:
            await ctx.send('That number is too big! 500 is the max.')
            return

        # purge the channel
        messages = await ctx.channel.purge(limit=amount)

        # complete!
        await ctx.send(f'Purge complete. Removed {len(messages)}/{amount}'
                       ' messages.', delete_after=2.5)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def block(self, ctx, someone: discord.Member):
        """
        Blocks someone from reading your channel.

        This simply creates a channel permission overwrite which blocks
        an individual from having Read Messages in this channel.
        """
        if ctx.channel.is_default():
            await ctx.send('Members cannot be blocked from the default channel.')
            return
        existing_overwrite = ctx.channel.overwrites_for(someone)
        existing_overwrite.read_messages = False
        await ctx.channel.set_permissions(someone, overwrite=existing_overwrite)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def unblock(self, ctx, someone: discord.Member):
        """
        Unblocks someone from reading your channel.

        Does the opposite of d?block.
        """
        overwrites = ctx.channel.overwrites_for(someone)
        can_read = overwrites.read_messages
        if can_read is None or can_read:
            # person isn't blocked
            await ctx.send('That person isn\'t (specifically) blocked.')
        elif not can_read:
            # person is blocked
            overwrites.read_messages = None
            if overwrites.is_empty():
                # overwrite is empty now, remove it
                logger.info('overwrite is empty')
                await ctx.channel.set_permissions(someone, overwrite=None)
            else:
                # overwrite still has stuff
                logger.info('overwrite still has things!')
                await ctx.channel.set_permissions(someone, overwrite=overwrites)
            await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def disable(self, ctx, command: str):
        """
        Disables a command in this server.

        In order to disable a group, simply pass in the group's name.
        For example: "d?disable feedback" will disable all feedback commands.
        """
        if self.bot.has_prefix(command):
            await ctx.send('You must leave off the prefix.')
            return
        await self.bot.disable_command(ctx.guild, command)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def enable(self, ctx, command: str):
        """ Enables a command in this server. """
        if self.bot.has_prefix(command):
            await ctx.send('You must leave off the prefix.')
            return
        if not await self.bot.command_is_disabled(ctx.guild, command):
            await ctx.send('That command isn\'t disabled!')
            return
        await self.bot.enable_command(ctx.guild, command)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def disabled(self, ctx):
        """ Shows you disabled commands in this server. """
        keys = await self.bot.redis.keys(f'disabled:{ctx.guild.id}:*')
        disabled = ['d?' + k.decode().split(':')[2] for k in keys]

        if not disabled:
            await ctx.send('There are no disabled commands in this server.')
        else:
            await ctx.send(', '.join(disabled))

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(kick_members=True)
    @checks.bot_perms(kick_members=True)
    async def kick(self, ctx, member: discord.Member):
        """ Kicks someone. """
        try:
            await ctx.guild.kick(member)
        except discord.Forbidden:
            await ctx.send('I can\'t do that.')
        else:
            await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(ban_members=True)
    @checks.bot_perms(ban_members=True)
    async def ban(self, ctx, member: discord.Member, days: int=0):
        """ Bans someone from the server. """
        try:
            await ctx.guild.ban(member, delete_message_days=days)
        except discord.Forbidden:
            await ctx.send('I can\'t do that.')
        else:
            await self.bot.ok(ctx)

    def _embed_field_for(self, member):
        return f'{member.mention} {member.name}#{member.discriminator}'

    def _make_action_embed(self, executor, victim, **kwargs):
        embed = discord.Embed(**kwargs)
        embed.add_field(name='Who did it', value=self._embed_field_for(executor))
        embed.add_field(name='Who was it', value=self._embed_field_for(victim))
        embed.set_footer(text=utils.now())
        return embed

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def unmute(self, ctx, member: discord.Member):
        """ Instantly unmutes someone. """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            await ctx.send('\N{CROSS MARK} I can\'t find the "Muted" role.')
            return

        if mute_role not in member.roles:
            await ctx.send(f'\N{SPEAKER} That person isn\'t muted.')
            return

        if member in self.mute_tasks:
            # cancel the mute task, so they don't get unmuted when the
            # task awakens
            self.mute_tasks[member].cancel()

        try:
            await member.remove_roles(mute_role)
            await ctx.send(f'\N{SPEAKER} Unmuted {member.name}#{member.discriminator}'
                           f' (`{member.id}`).')
        except discord.Forbidden:
            await ctx.send('\N{CROSS MARK} I can\'t do that!')
            return
        except:
            await ctx.send('\N{CROSS MARK} I failed to do that.')
            return

        embed = self._make_action_embed(ctx.author, member,
                                        title='\N{SPEAKER} Member force-unmuted')
        await self.bot.send_modlog(ctx.guild, embed=embed)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_channels=True)
    @checks.is_moderator()
    async def mute_setup(self, ctx):
        """
        Sets up "Muted" permission overrides for all channels.

        If the "Muted" role was not found, it will be created.
        By default, the permission overrides will deny the "Send Messages"
        permission for the people with the "Muted" role. If you want to deny
        the "Read Message" permissions for "Muted" people too, you must set
        the "mutesetup_disallow_read" key. You can do this with
            d?config set mutesetup_disallow_read
        Then, run d?muted_setup. If you already did, you must re-run it.

        Note: If you set that key, you will be prompted to provide a list of
        channels that will be readable (but not writable) of "Muted" status.
        """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            msg = await ctx.send('There is no "Muted" role, I\'ll set it up '
                                 'for you.')
            try:
                mute_role = await ctx.guild.create_role(name='Muted')
            except discord.Forbidden:
                await msg.edit(content='I couldn\'t find the "Muted" role,'
                                       'and I couldn\'t create it either!')
                return

        mute_options = {
            'send_messages': False
        }
        exclusion_list = []

        # beacuse we are disallowing read perms, there may be some channels
        # that people will still need to look at, even when muted.
        # so, prompt the moderator for a list of channels that will still
        # be readable (but not writable) regardless of mute status
        if await self.bot.config_is_set(ctx.guild, 'mutesetup_disallow_read'):
            mute_options['read_messages'] = False
            await ctx.send("""Okay! Please provide a comma separated list of channel
names for me to **exclude from being hidden from users
while they are muted**. For example, you probably don't
want muted users to not be able to see a channel like
#announcements!

If you don't want to exclude any channels, just type "none" and send it.

Example response: "announcements,corkboard,etc"
**Do not mention the channel, or put spaces anywhere.**""")

            msg = await self.bot.wait_for_response(ctx)

            if msg.content != "none":
                # create an exclusion list from the message
                exclusion_list = [discord.utils.get(ctx.guild.channels, name=c)
                                  for c in msg.content.split(',')]

                # if discord.utils.get returned None at some point, then that
                # means that one of the channels could not be found.
                if None in exclusion_list:
                    await ctx.send('Sorry! I didn\'t find one of the channels '
                                   'you listed, or your exclusion list was '
                                   'invalid. Please try the command again.')
                    return

                # format the exclusion list to look nice
                formatted_exclusions = '\n'.join([f'• {c.mention}' for c in
                                                  exclusion_list])

                # say the channels that will be excluded
                await ctx.send('For reference, here is the list of channels I'
                               ' am excluding from being hidden while muted:'
                               f'\n\n{formatted_exclusions}')
            else:
                await ctx.send('Alright, I won\'t exclude any channels.')

        failed = []
        succeeded = 0

        prg = await ctx.send('Doing it...')

        for channel in ctx.guild.channels:
            overwrite = discord.PermissionOverwrite(**mute_options)
            try:
                # if the channel is in the exclusion list, just delete
                # the overwrite
                if channel in exclusion_list:
                    await channel.set_permissions(mute_role, send_messages=False)
                else:
                    await channel.set_permissions(mute_role,
                                                  overwrite=overwrite)
            except discord.Forbidden:
                failed.append(channel.mention)
            else:
                if channel not in exclusion_list:
                    succeeded += 1

        if failed:
            await prg.edit(content=f'All done! I failed to edit **{len(failed)}**'
                                   f' channel(s): {", ".join(failed)}')
        else:
            await prg.edit(content='All done! Everything went OK. I modified '
                           f'{succeeded} channel(s). Note: This server\'s'
                           f' default channel, {ctx.guild.default_channel.mention},'
                           f' can always be read, even if someone is muted.'
                           ' (This is a Discord thing.)')

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(manage_roles=True)
    @checks.bot_perms(manage_roles=True)
    async def vanity(self, ctx, name: str):
        """
        Creates a vanity role.

        A vanity role is defined as a role with no permissions.
        """
        perms = discord.Permissions(permissions=0)
        await ctx.guild.create_role(name=name, permissions=perms)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def mute(self, ctx, member: discord.Member, time: HumanTime):
        """
        Mutes someone for a certain amount of time.

        The "Muted" role must exist on the server in order for this to work.
        The bot can setup the "Muted" role and channel overrides for you
        with the d?mute_setup command.

        d?mute <someone> 5m
            Mutes someone for 5 minutes.
        d?mute <someone> 2h5m
            Mutes someone for 2 hours and 5 minutes.
        d?mute <someone> 5s
            Mutes someone for 5 seconds.
        """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            await ctx.send('\N{CROSS MARK} I can\'t find the "Muted" role.')
            return

        mute = '\N{SPEAKER WITH CANCELLATION STROKE}'
        msg = await ctx.send(f'{mute} Muting {member.name}#{member.discriminator}'
                             f' (`{member.id}`) for {time.seconds} second(s).')

        try:
            await member.add_roles(mute_role)
        except discord.Forbidden:
            await msg.edit(content='\N{CROSS MARK} I can\'t do that!'
                           ' I might be too low on the role hierarchy,'
                           ' or I need permissions.'
                           ' Ensure that the my bot role is placed above'
                           ' the "Muted" role.')
            return
        except:
            await msg.edit(content='\N{CROSS MARK} I failed to do that.')
            return

        async def unmute_task():
            await asyncio.sleep(time.seconds)
            await member.remove_roles(mute_role)
            if await self.bot.config_is_set(ctx.guild, 'unmute_announce'):
                await ctx.send(f'\N{SPEAKER} {member.name}#{member.discriminator}'
                               f' (`{member.id}`) has been unmuted.')

        task = self.bot.loop.create_task(unmute_task())
        self.mute_tasks[member] = task

        embed = self._make_action_embed(ctx.author, member, title=f'{mute} Member muted')
        embed.add_field(name='Duration', value=f'{time.raw} ({time.seconds}s)')
        await self.bot.send_modlog(ctx.guild, embed=embed)

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(manage_guild=True)
    @checks.bot_perms(manage_nicknames=True)
    async def attentionseek(self, ctx, replace_with: str='💩'):
        """
        Changes attention-seeking nicknames.

        This will change the nickname of anybody whose name starts with "!"
        to a name you specify. By default, they are renamed to "💩".

        The renaming of attention-seekers is borrowed from the Discord API
        server.
        """
        attention_seekers = [m for m in ctx.guild.members if
                             m.display_name.startswith('!')]
        succeeded = len(attention_seekers)
        for seeker in attention_seekers:
            try:
                await seeker.edit(nick=replace_with)
            except:
                succeeded -= 1
        failed_count = len(attention_seekers) - succeeded
        await ctx.send(f'Renamed {succeeded} attention seeker(s).'
                       f' Failed to rename {failed_count}.')

def setup(bot):
    bot.add_cog(Mod(bot))
<commit_msg>dog.ext.mod: better cmd help<commit_after>import logging
import asyncio
import discord
from discord.ext import commands
from dog import Cog, checks, utils
from dog.humantime import HumanTime

logger = logging.getLogger(__name__)

class Mod(Cog):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.mute_tasks = {}

    async def __global_check(self, ctx):
        # do not handle guild-specific command disables in
        # dms
        if not isinstance(ctx.channel, discord.abc.GuildChannel):
            return True

        return not await self.bot.command_is_disabled(ctx.guild, ctx.command.name)

    async def on_message(self, message):
        # do not handle invisibility in dms
        if isinstance(message.channel, discord.abc.PrivateChannel):
            return

        if await self.bot.config_is_set(message.guild, 'invisible_announce'):
            if message.author.status is discord.Status.offline:
                reply = ('Hey {0.mention}! You\'re invisible. Stop being '
                         'invisible, please. Thanks.')
                await message.channel.send(reply.format(message.author))

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_messages=True, read_message_history=True)
    @checks.is_moderator()
    async def purge(self, ctx, amount: int):
        """
        Purges messages the last <n> messages.

        This includes the purge command message. The bot will send a message
        upon completion, but will automatically be deleted.
        """

        # increase one because the command message they sent needs
        # to be purged, too.
        amount = amount + 1
        if amount > 501:
            await ctx.send('That number is too big! 500 is the max.')
            return

        # purge the channel
        messages = await ctx.channel.purge(limit=amount)

        # complete!
        await ctx.send(f'Purge complete. Removed {len(messages)}/{amount}'
                       ' messages.', delete_after=2.5)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def block(self, ctx, someone: discord.Member):
        """
        Blocks someone from reading your channel.

        This simply creates a channel permission overwrite which blocks
        an individual from having Read Messages in this channel.

        Any existing permission overwrites for the target is preserved.
        """
        if ctx.channel.is_default():
            await ctx.send('Members cannot be blocked from the default channel.')
            return
        existing_overwrite = ctx.channel.overwrites_for(someone)
        existing_overwrite.read_messages = False
        await ctx.channel.set_permissions(someone, overwrite=existing_overwrite)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def unblock(self, ctx, someone: discord.Member):
        """
        Unblocks someone from reading your channel.

        Does the opposite of d?block.
        """
        overwrites = ctx.channel.overwrites_for(someone)
        can_read = overwrites.read_messages
        if can_read is None or can_read:
            # person isn't blocked
            await ctx.send('That person isn\'t (specifically) blocked.')
        elif not can_read:
            # person is blocked
            overwrites.read_messages = None
            if overwrites.is_empty():
                # overwrite is empty now, remove it
                logger.info('overwrite is empty')
                await ctx.channel.set_permissions(someone, overwrite=None)
            else:
                # overwrite still has stuff
                logger.info('overwrite still has things!')
                await ctx.channel.set_permissions(someone, overwrite=overwrites)
            await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def disable(self, ctx, command: str):
        """
        Disables a command in this server.

        In order to disable a group, simply pass in the group's name.
        For example: "d?disable feedback" will disable all feedback commands.
        """
        if self.bot.has_prefix(command):
            await ctx.send('You must leave off the prefix.')
            return
        await self.bot.disable_command(ctx.guild, command)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def enable(self, ctx, command: str):
        """ Enables a command in this server. """
        if self.bot.has_prefix(command):
            await ctx.send('You must leave off the prefix.')
            return
        if not await self.bot.command_is_disabled(ctx.guild, command):
            await ctx.send('That command isn\'t disabled!')
            return
        await self.bot.enable_command(ctx.guild, command)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def disabled(self, ctx):
        """ Shows disabled commands in this server. """
        keys = await self.bot.redis.keys(f'disabled:{ctx.guild.id}:*')
        disabled = ['d?' + k.decode().split(':')[2] for k in keys]

        if not disabled:
            await ctx.send('There are no disabled commands in this server.')
        else:
            await ctx.send(', '.join(disabled))

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(kick_members=True)
    @checks.bot_perms(kick_members=True)
    async def kick(self, ctx, member: discord.Member):
        """ Kicks someone. """
        try:
            await ctx.guild.kick(member)
        except discord.Forbidden:
            await ctx.send('I can\'t do that.')
        else:
            await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(ban_members=True)
    @checks.bot_perms(ban_members=True)
    async def ban(self, ctx, member: discord.Member, days: int=0):
        """ Bans someone. """
        try:
            await ctx.guild.ban(member, delete_message_days=days)
        except discord.Forbidden:
            await ctx.send('I can\'t do that.')
        else:
            await self.bot.ok(ctx)

    def _embed_field_for(self, member):
        return f'{member.mention} {member.name}#{member.discriminator}'

    def _make_action_embed(self, executor, victim, **kwargs):
        embed = discord.Embed(**kwargs)
        embed.add_field(name='Who did it', value=self._embed_field_for(executor))
        embed.add_field(name='Who was it', value=self._embed_field_for(victim))
        embed.set_footer(text=utils.now())
        return embed

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def unmute(self, ctx, member: discord.Member):
        """ Instantly unmutes someone. """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            await ctx.send('\N{CROSS MARK} I can\'t find the "Muted" role.')
            return

        if mute_role not in member.roles:
            await ctx.send(f'\N{SPEAKER} That person isn\'t muted.')
            return

        if member in self.mute_tasks:
            # cancel the mute task, so they don't get unmuted when the
            # task awakens
            self.mute_tasks[member].cancel()

        try:
            await member.remove_roles(mute_role)
            await ctx.send(f'\N{SPEAKER} Unmuted {member.name}#{member.discriminator}'
                           f' (`{member.id}`).')
        except discord.Forbidden:
            await ctx.send('\N{CROSS MARK} I can\'t do that!')
            return
        except:
            await ctx.send('\N{CROSS MARK} I failed to do that.')
            return

        embed = self._make_action_embed(ctx.author, member,
                                        title='\N{SPEAKER} Member force-unmuted')
        await self.bot.send_modlog(ctx.guild, embed=embed)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_channels=True)
    @checks.is_moderator()
    async def mute_setup(self, ctx):
        """
        Sets up "Muted" permission overrides for all channels.

        If the "Muted" role was not found, it will be created.
        By default, the permission overrides will deny the "Send Messages"
        permission for the people with the "Muted" role. If you want to deny
        the "Read Message" permissions for "Muted" people too, you must set
        the "mutesetup_disallow_read" key. You can do this with
            d?config set mutesetup_disallow_read
        Then, run d?muted_setup. If you already did, you must re-run it.

        Note: If you set that key, you will be prompted to provide a list of
        channels that will be readable (but not writable) of "Muted" status.
        """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            msg = await ctx.send('There is no "Muted" role, I\'ll set it up '
                                 'for you.')
            try:
                mute_role = await ctx.guild.create_role(name='Muted')
            except discord.Forbidden:
                await msg.edit(content='I couldn\'t find the "Muted" role,'
                                       'and I couldn\'t create it either!')
                return

        mute_options = {
            'send_messages': False
        }
        exclusion_list = []

        # beacuse we are disallowing read perms, there may be some channels
        # that people will still need to look at, even when muted.
        # so, prompt the moderator for a list of channels that will still
        # be readable (but not writable) regardless of mute status
        if await self.bot.config_is_set(ctx.guild, 'mutesetup_disallow_read'):
            mute_options['read_messages'] = False
            await ctx.send("""Okay! Please provide a comma separated list of channel
names for me to **exclude from being hidden from users
while they are muted**. For example, you probably don't
want muted users to not be able to see a channel like
#announcements!

If you don't want to exclude any channels, just type "none" and send it.

Example response: "announcements,corkboard,etc"
**Do not mention the channel, or put spaces anywhere.**""")

            msg = await self.bot.wait_for_response(ctx)

            if msg.content != "none":
                # create an exclusion list from the message
                exclusion_list = [discord.utils.get(ctx.guild.channels, name=c)
                                  for c in msg.content.split(',')]

                # if discord.utils.get returned None at some point, then that
                # means that one of the channels could not be found.
                if None in exclusion_list:
                    await ctx.send('Sorry! I didn\'t find one of the channels '
                                   'you listed, or your exclusion list was '
                                   'invalid. Please try the command again.')
                    return

                # format the exclusion list to look nice
                formatted_exclusions = '\n'.join([f'• {c.mention}' for c in
                                                  exclusion_list])

                # say the channels that will be excluded
                await ctx.send('For reference, here is the list of channels I'
                               ' am excluding from being hidden while muted:'
                               f'\n\n{formatted_exclusions}')
            else:
                await ctx.send('Alright, I won\'t exclude any channels.')

        failed = []
        succeeded = 0

        prg = await ctx.send('Doing it...')

        for channel in ctx.guild.channels:
            overwrite = discord.PermissionOverwrite(**mute_options)
            try:
                # if the channel is in the exclusion list, just delete
                # the overwrite
                if channel in exclusion_list:
                    await channel.set_permissions(mute_role, send_messages=False)
                else:
                    await channel.set_permissions(mute_role,
                                                  overwrite=overwrite)
            except discord.Forbidden:
                failed.append(channel.mention)
            else:
                if channel not in exclusion_list:
                    succeeded += 1

        if failed:
            await prg.edit(content=f'All done! I failed to edit **{len(failed)}**'
                                   f' channel(s): {", ".join(failed)}')
        else:
            await prg.edit(content='All done! Everything went OK. I modified '
                           f'{succeeded} channel(s). Note: This server\'s'
                           f' default channel, {ctx.guild.default_channel.mention},'
                           f' can always be read, even if someone is muted.'
                           ' (This is a Discord thing.)')

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(manage_roles=True)
    @checks.bot_perms(manage_roles=True)
    async def vanity(self, ctx, name: str):
        """
        Creates a vanity role.

        A vanity role is defined as a role with no permissions.
        """
        perms = discord.Permissions(permissions=0)
        await ctx.guild.create_role(name=name, permissions=perms)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def mute(self, ctx, member: discord.Member, time: HumanTime):
        """
        Mutes someone for a certain amount of time.

        The "Muted" role must exist on the server in order for this to work.
        The bot can setup the "Muted" role and channel overrides for you
        with the d?mute_setup command.

        d?mute <someone> 5m
            Mutes someone for 5 minutes.
        d?mute <someone> 2h5m
            Mutes someone for 2 hours and 5 minutes.
        d?mute <someone> 5s
            Mutes someone for 5 seconds.
        """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            await ctx.send('\N{CROSS MARK} I can\'t find the "Muted" role.')
            return

        mute = '\N{SPEAKER WITH CANCELLATION STROKE}'
        msg = await ctx.send(f'{mute} Muting {member.name}#{member.discriminator}'
                             f' (`{member.id}`) for {time.seconds} second(s).')

        try:
            await member.add_roles(mute_role)
        except discord.Forbidden:
            await msg.edit(content='\N{CROSS MARK} I can\'t do that!'
                           ' I might be too low on the role hierarchy,'
                           ' or I need permissions.'
                           ' Ensure that the my bot role is placed above'
                           ' the "Muted" role.')
            return
        except:
            await msg.edit(content='\N{CROSS MARK} I failed to do that.')
            return

        async def unmute_task():
            await asyncio.sleep(time.seconds)
            await member.remove_roles(mute_role)
            if await self.bot.config_is_set(ctx.guild, 'unmute_announce'):
                await ctx.send(f'\N{SPEAKER} {member.name}#{member.discriminator}'
                               f' (`{member.id}`) has been unmuted.')

        task = self.bot.loop.create_task(unmute_task())
        self.mute_tasks[member] = task

        embed = self._make_action_embed(ctx.author, member, title=f'{mute} Member muted')
        embed.add_field(name='Duration', value=f'{time.raw} ({time.seconds}s)')
        await self.bot.send_modlog(ctx.guild, embed=embed)

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(manage_guild=True)
    @checks.bot_perms(manage_nicknames=True)
    async def attentionseek(self, ctx, replace_with: str='💩'):
        """
        Changes attention-seeking nicknames.

        This will change the nickname of anybody whose name starts with "!"
        to a name you specify. By default, they are renamed to "💩".

        The renaming of attention-seekers is borrowed from the Discord API
        server.
        """
        attention_seekers = [m for m in ctx.guild.members if
                             m.display_name.startswith('!')]
        succeeded = len(attention_seekers)
        for seeker in attention_seekers:
            try:
                await seeker.edit(nick=replace_with)
            except:
                succeeded -= 1
        failed_count = len(attention_seekers) - succeeded
        await ctx.send(f'Renamed {succeeded} attention seeker(s).'
                       f' Failed to rename {failed_count}.')

def setup(bot):
    bot.add_cog(Mod(bot))
<|endoftext|>


<commit_before>import logging
import asyncio
import discord
from discord.ext import commands
from dog import Cog, checks, utils
from dog.humantime import HumanTime

logger = logging.getLogger(__name__)

class Mod(Cog):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.mute_tasks = {}

    async def __global_check(self, ctx):
        # do not handle guild-specific command disables in
        # dms
        if not isinstance(ctx.channel, discord.abc.GuildChannel):
            return True

        return not await self.bot.command_is_disabled(ctx.guild, ctx.command.name)

    async def on_message(self, message):
        # do not handle invisibility in dms
        if isinstance(message.channel, discord.abc.PrivateChannel):
            return

        if await self.bot.config_is_set(message.guild, 'invisible_announce'):
            if message.author.status is discord.Status.offline:
                reply = ('Hey {0.mention}! You\'re invisible. Stop being '
                         'invisible, please. Thanks.')
                await message.channel.send(reply.format(message.author))

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_messages=True, read_message_history=True)
    @checks.is_moderator()
    async def purge(self, ctx, amount: int):
        """
        Purges messages the last <n> messages.

        This includes the purge command message. The bot will send a message
        upon completion, but will automatically be deleted.
        """

        # increase one because the command message they sent needs
        # to be purged, too.
        amount = amount + 1
        if amount > 501:
            await ctx.send('That number is too big! 500 is the max.')
            return

        # purge the channel
        messages = await ctx.channel.purge(limit=amount)

        # complete!
        await ctx.send(f'Purge complete. Removed {len(messages)}/{amount}'
                       ' messages.', delete_after=2.5)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def block(self, ctx, someone: discord.Member):
        """
        Blocks someone from reading your channel.

        This simply creates a channel permission overwrite which blocks
        an individual from having Read Messages in this channel.

        Any existing permission overwrites for the target is preserved.
        """
        if ctx.channel.is_default():
            await ctx.send('Members cannot be blocked from the default channel.')
            return
        existing_overwrite = ctx.channel.overwrites_for(someone)
        existing_overwrite.read_messages = False
        await ctx.channel.set_permissions(someone, overwrite=existing_overwrite)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def unblock(self, ctx, someone: discord.Member):
        """
        Unblocks someone from reading your channel.

        Does the opposite of d?block.
        """
        overwrites = ctx.channel.overwrites_for(someone)
        can_read = overwrites.read_messages
        if can_read is None or can_read:
            # person isn't blocked
            await ctx.send('That person isn\'t (specifically) blocked.')
        elif not can_read:
            # person is blocked
            overwrites.read_messages = None
            if overwrites.is_empty():
                # overwrite is empty now, remove it
                logger.info('overwrite is empty')
                await ctx.channel.set_permissions(someone, overwrite=None)
            else:
                # overwrite still has stuff
                logger.info('overwrite still has things!')
                await ctx.channel.set_permissions(someone, overwrite=overwrites)
            await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def disable(self, ctx, command: str):
        """
        Disables a command in this server.

        In order to disable a group, simply pass in the group's name.
        For example: "d?disable feedback" will disable all feedback commands.
        """
        if self.bot.has_prefix(command):
            await ctx.send('You must leave off the prefix.')
            return
        await self.bot.disable_command(ctx.guild, command)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def enable(self, ctx, command: str):
        """ Enables a command in this server. """
        if self.bot.has_prefix(command):
            await ctx.send('You must leave off the prefix.')
            return
        if not await self.bot.command_is_disabled(ctx.guild, command):
            await ctx.send('That command isn\'t disabled!')
            return
        await self.bot.enable_command(ctx.guild, command)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def disabled(self, ctx):
        """ Shows disabled commands in this server. """
        keys = await self.bot.redis.keys(f'disabled:{ctx.guild.id}:*')
        disabled = ['d?' + k.decode().split(':')[2] for k in keys]

        if not disabled:
            await ctx.send('There are no disabled commands in this server.')
        else:
            await ctx.send(', '.join(disabled))

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(kick_members=True)
    @checks.bot_perms(kick_members=True)
    async def kick(self, ctx, member: discord.Member):
        """ Kicks someone. """
        try:
            await ctx.guild.kick(member)
        except discord.Forbidden:
            await ctx.send('I can\'t do that.')
        else:
            await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(ban_members=True)
    @checks.bot_perms(ban_members=True)
    async def ban(self, ctx, member: discord.Member, days: int=0):
        """ Bans someone. """
        try:
            await ctx.guild.ban(member, delete_message_days=days)
        except discord.Forbidden:
            await ctx.send('I can\'t do that.')
        else:
            await self.bot.ok(ctx)

    def _embed_field_for(self, member):
        return f'{member.mention} {member.name}#{member.discriminator}'

    def _make_action_embed(self, executor, victim, **kwargs):
        embed = discord.Embed(**kwargs)
        embed.add_field(name='Who did it', value=self._embed_field_for(executor))
        embed.add_field(name='Who was it', value=self._embed_field_for(victim))
        embed.set_footer(text=utils.now())
        return embed

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def unmute(self, ctx, member: discord.Member):
        """ Instantly unmutes someone. """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            await ctx.send('\N{CROSS MARK} I can\'t find the "Muted" role.')
            return

        if mute_role not in member.roles:
            await ctx.send(f'\N{SPEAKER} That person isn\'t muted.')
            return

        if member in self.mute_tasks:
            # cancel the mute task, so they don't get unmuted when the
            # task awakens
            self.mute_tasks[member].cancel()

        try:
            await member.remove_roles(mute_role)
            await ctx.send(f'\N{SPEAKER} Unmuted {member.name}#{member.discriminator}'
                           f' (`{member.id}`).')
        except discord.Forbidden:
            await ctx.send('\N{CROSS MARK} I can\'t do that!')
            return
        except:
            await ctx.send('\N{CROSS MARK} I failed to do that.')
            return

        embed = self._make_action_embed(ctx.author, member,
                                        title='\N{SPEAKER} Member force-unmuted')
        await self.bot.send_modlog(ctx.guild, embed=embed)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_channels=True)
    @checks.is_moderator()
    async def mute_setup(self, ctx):
        """
        Sets up "Muted" permission overrides for all channels.

        If the "Muted" role was not found, it will be created.
        By default, the permission overrides will deny the "Send Messages"
        permission for the people with the "Muted" role. If you want to deny
        the "Read Message" permissions for "Muted" people too, you must set
        the "mutesetup_disallow_read" key. You can do this with
            d?config set mutesetup_disallow_read
        Then, run d?muted_setup. If you already did, you must re-run it.

        Note: If you set that key, you will be prompted to provide a list of
        channels that will be readable (but not writable) of "Muted" status.
        """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            msg = await ctx.send('There is no "Muted" role, I\'ll set it up '
                                 'for you.')
            try:
                mute_role = await ctx.guild.create_role(name='Muted')
            except discord.Forbidden:
                await msg.edit(content='I couldn\'t find the "Muted" role,'
                                       'and I couldn\'t create it either!')
                return

        mute_options = {
            'send_messages': False
        }
        exclusion_list = []

        # beacuse we are disallowing read perms, there may be some channels
        # that people will still need to look at, even when muted.
        # so, prompt the moderator for a list of channels that will still
        # be readable (but not writable) regardless of mute status
        if await self.bot.config_is_set(ctx.guild, 'mutesetup_disallow_read'):
            mute_options['read_messages'] = False
            await ctx.send("""Okay! Please provide a comma separated list of channel
names for me to **exclude from being hidden from users
while they are muted**. For example, you probably don't
want muted users to not be able to see a channel like
#announcements!

If you don't want to exclude any channels, just type "none" and send it.

Example response: "announcements,corkboard,etc"
**Do not mention the channel, or put spaces anywhere.**""")

            msg = await self.bot.wait_for_response(ctx)

            if msg.content != "none":
                # create an exclusion list from the message
                exclusion_list = [discord.utils.get(ctx.guild.channels, name=c)
                                  for c in msg.content.split(',')]

                # if discord.utils.get returned None at some point, then that
                # means that one of the channels could not be found.
                if None in exclusion_list:
                    await ctx.send('Sorry! I didn\'t find one of the channels '
                                   'you listed, or your exclusion list was '
                                   'invalid. Please try the command again.')
                    return

                # format the exclusion list to look nice
                formatted_exclusions = '\n'.join([f'• {c.mention}' for c in
                                                  exclusion_list])

                # say the channels that will be excluded
                await ctx.send('For reference, here is the list of channels I'
                               ' am excluding from being hidden while muted:'
                               f'\n\n{formatted_exclusions}')
            else:
                await ctx.send('Alright, I won\'t exclude any channels.')

        failed = []
        succeeded = 0

        prg = await ctx.send('Doing it...')

        for channel in ctx.guild.channels:
            overwrite = discord.PermissionOverwrite(**mute_options)
            try:
                # if the channel is in the exclusion list, just delete
                # the overwrite
                if channel in exclusion_list:
                    await channel.set_permissions(mute_role, send_messages=False)
                else:
                    await channel.set_permissions(mute_role,
                                                  overwrite=overwrite)
            except discord.Forbidden:
                failed.append(channel.mention)
            else:
                if channel not in exclusion_list:
                    succeeded += 1

        if failed:
            await prg.edit(content=f'All done! I failed to edit **{len(failed)}**'
                                   f' channel(s): {", ".join(failed)}')
        else:
            await prg.edit(content='All done! Everything went OK. I modified '
                           f'{succeeded} channel(s). Note: This server\'s'
                           f' default channel, {ctx.guild.default_channel.mention},'
                           f' can always be read, even if someone is muted.'
                           ' (This is a Discord thing.)')

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(manage_roles=True)
    @checks.bot_perms(manage_roles=True)
    async def vanity(self, ctx, *, name: str):
        """
        Creates a vanity role.

        A vanity role is defined as a role with no permissions.
        """
        perms = discord.Permissions(permissions=0)
        await ctx.guild.create_role(name=name, permissions=perms)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def mute(self, ctx, member: discord.Member, time: HumanTime):
        """
        Mutes someone for a certain amount of time.

        The "Muted" role must exist on the server in order for this to work.
        The bot can setup the "Muted" role and channel overrides for you
        with the d?mute_setup command.

        d?mute <someone> 5m
            Mutes someone for 5 minutes.
        d?mute <someone> 2h5m
            Mutes someone for 2 hours and 5 minutes.
        d?mute <someone> 5s
            Mutes someone for 5 seconds.
        """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            await ctx.send('\N{CROSS MARK} I can\'t find the "Muted" role.')
            return

        mute = '\N{SPEAKER WITH CANCELLATION STROKE}'
        msg = await ctx.send(f'{mute} Muting {member.name}#{member.discriminator}'
                             f' (`{member.id}`) for {time.seconds} second(s).')

        try:
            await member.add_roles(mute_role)
        except discord.Forbidden:
            await msg.edit(content='\N{CROSS MARK} I can\'t do that!'
                           ' I might be too low on the role hierarchy,'
                           ' or I need permissions.'
                           ' Ensure that the my bot role is placed above'
                           ' the "Muted" role.')
            return
        except:
            await msg.edit(content='\N{CROSS MARK} I failed to do that.')
            return

        async def unmute_task():
            await asyncio.sleep(time.seconds)
            await member.remove_roles(mute_role)
            if await self.bot.config_is_set(ctx.guild, 'unmute_announce'):
                await ctx.send(f'\N{SPEAKER} {member.name}#{member.discriminator}'
                               f' (`{member.id}`) has been unmuted.')

        task = self.bot.loop.create_task(unmute_task())
        self.mute_tasks[member] = task

        embed = self._make_action_embed(ctx.author, member, title=f'{mute} Member muted')
        embed.add_field(name='Duration', value=f'{time.raw} ({time.seconds}s)')
        await self.bot.send_modlog(ctx.guild, embed=embed)

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(manage_guild=True)
    @checks.bot_perms(manage_nicknames=True)
    async def attentionseek(self, ctx, replace_with: str='💩'):
        """
        Changes attention-seeking nicknames.

        This will change the nickname of anybody whose name starts with "!"
        to a name you specify. By default, they are renamed to "💩".

        The renaming of attention-seekers is borrowed from the Discord API
        server.
        """
        attention_seekers = [m for m in ctx.guild.members if
                             m.display_name.startswith('!')]
        succeeded = len(attention_seekers)
        for seeker in attention_seekers:
            try:
                await seeker.edit(nick=replace_with)
            except:
                succeeded -= 1
        failed_count = len(attention_seekers) - succeeded
        await ctx.send(f'Renamed {succeeded} attention seeker(s).'
                       f' Failed to rename {failed_count}.')

def setup(bot):
    bot.add_cog(Mod(bot))
<commit_msg>mod: allow users to be plonked from dog<commit_after>import logging
import asyncio
import discord
from discord.ext import commands
from dog import Cog, checks, utils
from dog.humantime import HumanTime

logger = logging.getLogger(__name__)

class Mod(Cog):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.mute_tasks = {}

    async def __global_check(self, ctx):
        # do not handle guild-specific command disables in
        # dms
        if not isinstance(ctx.channel, discord.abc.GuildChannel):
            return True

        banned_names = [
            'Can\'t Use Dog', 'can\'t use dog', 'Dog Plonk', 'dog plonk',
            'Banned from Dog', 'banned from dog'
        ]

        for banned_name in banned_names:
            if discord.utils.get(ctx.author.roles, name=banned_name):
                return False

        return not await self.bot.command_is_disabled(ctx.guild, ctx.command.name)

    async def on_message(self, message):
        # do not handle invisibility in dms
        if isinstance(message.channel, discord.abc.PrivateChannel):
            return

        if await self.bot.config_is_set(message.guild, 'invisible_announce'):
            if message.author.status is discord.Status.offline:
                reply = ('Hey {0.mention}! You\'re invisible. Stop being '
                         'invisible, please. Thanks.')
                await message.channel.send(reply.format(message.author))

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_messages=True, read_message_history=True)
    @checks.is_moderator()
    async def purge(self, ctx, amount: int):
        """
        Purges messages the last <n> messages.

        This includes the purge command message. The bot will send a message
        upon completion, but will automatically be deleted.
        """

        # increase one because the command message they sent needs
        # to be purged, too.
        amount = amount + 1
        if amount > 501:
            await ctx.send('That number is too big! 500 is the max.')
            return

        # purge the channel
        messages = await ctx.channel.purge(limit=amount)

        # complete!
        await ctx.send(f'Purge complete. Removed {len(messages)}/{amount}'
                       ' messages.', delete_after=2.5)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def block(self, ctx, someone: discord.Member):
        """
        Blocks someone from reading your channel.

        This simply creates a channel permission overwrite which blocks
        an individual from having Read Messages in this channel.

        Any existing permission overwrites for the target is preserved.
        """
        if ctx.channel.is_default():
            await ctx.send('Members cannot be blocked from the default channel.')
            return
        existing_overwrite = ctx.channel.overwrites_for(someone)
        existing_overwrite.read_messages = False
        await ctx.channel.set_permissions(someone, overwrite=existing_overwrite)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def unblock(self, ctx, someone: discord.Member):
        """
        Unblocks someone from reading your channel.

        Does the opposite of d?block.
        """
        overwrites = ctx.channel.overwrites_for(someone)
        can_read = overwrites.read_messages
        if can_read is None or can_read:
            # person isn't blocked
            await ctx.send('That person isn\'t (specifically) blocked.')
        elif not can_read:
            # person is blocked
            overwrites.read_messages = None
            if overwrites.is_empty():
                # overwrite is empty now, remove it
                logger.info('overwrite is empty')
                await ctx.channel.set_permissions(someone, overwrite=None)
            else:
                # overwrite still has stuff
                logger.info('overwrite still has things!')
                await ctx.channel.set_permissions(someone, overwrite=overwrites)
            await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def disable(self, ctx, command: str):
        """
        Disables a command in this server.

        In order to disable a group, simply pass in the group's name.
        For example: "d?disable feedback" will disable all feedback commands.
        """
        if self.bot.has_prefix(command):
            await ctx.send('You must leave off the prefix.')
            return
        await self.bot.disable_command(ctx.guild, command)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def enable(self, ctx, command: str):
        """ Enables a command in this server. """
        if self.bot.has_prefix(command):
            await ctx.send('You must leave off the prefix.')
            return
        if not await self.bot.command_is_disabled(ctx.guild, command):
            await ctx.send('That command isn\'t disabled!')
            return
        await self.bot.enable_command(ctx.guild, command)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.is_moderator()
    async def disabled(self, ctx):
        """ Shows disabled commands in this server. """
        keys = await self.bot.redis.keys(f'disabled:{ctx.guild.id}:*')
        disabled = ['d?' + k.decode().split(':')[2] for k in keys]

        if not disabled:
            await ctx.send('There are no disabled commands in this server.')
        else:
            await ctx.send(', '.join(disabled))

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(kick_members=True)
    @checks.bot_perms(kick_members=True)
    async def kick(self, ctx, member: discord.Member):
        """ Kicks someone. """
        try:
            await ctx.guild.kick(member)
        except discord.Forbidden:
            await ctx.send('I can\'t do that.')
        else:
            await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(ban_members=True)
    @checks.bot_perms(ban_members=True)
    async def ban(self, ctx, member: discord.Member, days: int=0):
        """ Bans someone. """
        try:
            await ctx.guild.ban(member, delete_message_days=days)
        except discord.Forbidden:
            await ctx.send('I can\'t do that.')
        else:
            await self.bot.ok(ctx)

    def _embed_field_for(self, member):
        return f'{member.mention} {member.name}#{member.discriminator}'

    def _make_action_embed(self, executor, victim, **kwargs):
        embed = discord.Embed(**kwargs)
        embed.add_field(name='Who did it', value=self._embed_field_for(executor))
        embed.add_field(name='Who was it', value=self._embed_field_for(victim))
        embed.set_footer(text=utils.now())
        return embed

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def unmute(self, ctx, member: discord.Member):
        """ Instantly unmutes someone. """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            await ctx.send('\N{CROSS MARK} I can\'t find the "Muted" role.')
            return

        if mute_role not in member.roles:
            await ctx.send(f'\N{SPEAKER} That person isn\'t muted.')
            return

        if member in self.mute_tasks:
            # cancel the mute task, so they don't get unmuted when the
            # task awakens
            self.mute_tasks[member].cancel()

        try:
            await member.remove_roles(mute_role)
            await ctx.send(f'\N{SPEAKER} Unmuted {member.name}#{member.discriminator}'
                           f' (`{member.id}`).')
        except discord.Forbidden:
            await ctx.send('\N{CROSS MARK} I can\'t do that!')
            return
        except:
            await ctx.send('\N{CROSS MARK} I failed to do that.')
            return

        embed = self._make_action_embed(ctx.author, member,
                                        title='\N{SPEAKER} Member force-unmuted')
        await self.bot.send_modlog(ctx.guild, embed=embed)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_channels=True)
    @checks.is_moderator()
    async def mute_setup(self, ctx):
        """
        Sets up "Muted" permission overrides for all channels.

        If the "Muted" role was not found, it will be created.
        By default, the permission overrides will deny the "Send Messages"
        permission for the people with the "Muted" role. If you want to deny
        the "Read Message" permissions for "Muted" people too, you must set
        the "mutesetup_disallow_read" key. You can do this with
            d?config set mutesetup_disallow_read
        Then, run d?muted_setup. If you already did, you must re-run it.

        Note: If you set that key, you will be prompted to provide a list of
        channels that will be readable (but not writable) of "Muted" status.
        """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            msg = await ctx.send('There is no "Muted" role, I\'ll set it up '
                                 'for you.')
            try:
                mute_role = await ctx.guild.create_role(name='Muted')
            except discord.Forbidden:
                await msg.edit(content='I couldn\'t find the "Muted" role,'
                                       'and I couldn\'t create it either!')
                return

        mute_options = {
            'send_messages': False
        }
        exclusion_list = []

        # beacuse we are disallowing read perms, there may be some channels
        # that people will still need to look at, even when muted.
        # so, prompt the moderator for a list of channels that will still
        # be readable (but not writable) regardless of mute status
        if await self.bot.config_is_set(ctx.guild, 'mutesetup_disallow_read'):
            mute_options['read_messages'] = False
            await ctx.send("""Okay! Please provide a comma separated list of channel
names for me to **exclude from being hidden from users
while they are muted**. For example, you probably don't
want muted users to not be able to see a channel like
#announcements!

If you don't want to exclude any channels, just type "none" and send it.

Example response: "announcements,corkboard,etc"
**Do not mention the channel, or put spaces anywhere.**""")

            msg = await self.bot.wait_for_response(ctx)

            if msg.content != "none":
                # create an exclusion list from the message
                exclusion_list = [discord.utils.get(ctx.guild.channels, name=c)
                                  for c in msg.content.split(',')]

                # if discord.utils.get returned None at some point, then that
                # means that one of the channels could not be found.
                if None in exclusion_list:
                    await ctx.send('Sorry! I didn\'t find one of the channels '
                                   'you listed, or your exclusion list was '
                                   'invalid. Please try the command again.')
                    return

                # format the exclusion list to look nice
                formatted_exclusions = '\n'.join([f'• {c.mention}' for c in
                                                  exclusion_list])

                # say the channels that will be excluded
                await ctx.send('For reference, here is the list of channels I'
                               ' am excluding from being hidden while muted:'
                               f'\n\n{formatted_exclusions}')
            else:
                await ctx.send('Alright, I won\'t exclude any channels.')

        failed = []
        succeeded = 0

        prg = await ctx.send('Doing it...')

        for channel in ctx.guild.channels:
            overwrite = discord.PermissionOverwrite(**mute_options)
            try:
                # if the channel is in the exclusion list, just delete
                # the overwrite
                if channel in exclusion_list:
                    await channel.set_permissions(mute_role, send_messages=False)
                else:
                    await channel.set_permissions(mute_role,
                                                  overwrite=overwrite)
            except discord.Forbidden:
                failed.append(channel.mention)
            else:
                if channel not in exclusion_list:
                    succeeded += 1

        if failed:
            await prg.edit(content=f'All done! I failed to edit **{len(failed)}**'
                                   f' channel(s): {", ".join(failed)}')
        else:
            await prg.edit(content='All done! Everything went OK. I modified '
                           f'{succeeded} channel(s). Note: This server\'s'
                           f' default channel, {ctx.guild.default_channel.mention},'
                           f' can always be read, even if someone is muted.'
                           ' (This is a Discord thing.)')

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(manage_roles=True)
    @checks.bot_perms(manage_roles=True)
    async def vanity(self, ctx, *, name: str):
        """
        Creates a vanity role.

        A vanity role is defined as a role with no permissions.
        """
        perms = discord.Permissions(permissions=0)
        await ctx.guild.create_role(name=name, permissions=perms)
        await self.bot.ok(ctx)

    @commands.command()
    @commands.guild_only()
    @checks.bot_perms(manage_roles=True)
    @checks.is_moderator()
    async def mute(self, ctx, member: discord.Member, time: HumanTime):
        """
        Mutes someone for a certain amount of time.

        The "Muted" role must exist on the server in order for this to work.
        The bot can setup the "Muted" role and channel overrides for you
        with the d?mute_setup command.

        d?mute <someone> 5m
            Mutes someone for 5 minutes.
        d?mute <someone> 2h5m
            Mutes someone for 2 hours and 5 minutes.
        d?mute <someone> 5s
            Mutes someone for 5 seconds.
        """
        mute_role = discord.utils.get(ctx.guild.roles, name='Muted')

        if not mute_role:
            await ctx.send('\N{CROSS MARK} I can\'t find the "Muted" role.')
            return

        mute = '\N{SPEAKER WITH CANCELLATION STROKE}'
        msg = await ctx.send(f'{mute} Muting {member.name}#{member.discriminator}'
                             f' (`{member.id}`) for {time.seconds} second(s).')

        try:
            await member.add_roles(mute_role)
        except discord.Forbidden:
            await msg.edit(content='\N{CROSS MARK} I can\'t do that!'
                           ' I might be too low on the role hierarchy,'
                           ' or I need permissions.'
                           ' Ensure that the my bot role is placed above'
                           ' the "Muted" role.')
            return
        except:
            await msg.edit(content='\N{CROSS MARK} I failed to do that.')
            return

        async def unmute_task():
            await asyncio.sleep(time.seconds)
            await member.remove_roles(mute_role)
            if await self.bot.config_is_set(ctx.guild, 'unmute_announce'):
                await ctx.send(f'\N{SPEAKER} {member.name}#{member.discriminator}'
                               f' (`{member.id}`) has been unmuted.')

        task = self.bot.loop.create_task(unmute_task())
        self.mute_tasks[member] = task

        embed = self._make_action_embed(ctx.author, member, title=f'{mute} Member muted')
        embed.add_field(name='Duration', value=f'{time.raw} ({time.seconds}s)')
        await self.bot.send_modlog(ctx.guild, embed=embed)

    @commands.command()
    @commands.guild_only()
    @commands.has_permissions(manage_guild=True)
    @checks.bot_perms(manage_nicknames=True)
    async def attentionseek(self, ctx, replace_with: str='💩'):
        """
        Changes attention-seeking nicknames.

        This will change the nickname of anybody whose name starts with "!"
        to a name you specify. By default, they are renamed to "💩".

        The renaming of attention-seekers is borrowed from the Discord API
        server.
        """
        attention_seekers = [m for m in ctx.guild.members if
                             m.display_name.startswith('!')]
        succeeded = len(attention_seekers)
        for seeker in attention_seekers:
            try:
                await seeker.edit(nick=replace_with)
            except:
                succeeded -= 1
        failed_count = len(attention_seekers) - succeeded
        await ctx.send(f'Renamed {succeeded} attention seeker(s).'
                       f' Failed to rename {failed_count}.')

def setup(bot):
    bot.add_cog(Mod(bot))
<|endoftext|>


<commit_before>#!/usr/bin/env python
import operator
import time

import lasagne
import numpy as np
import theano
import theano.tensor as T
from sklearn.metrics import confusion_matrix

import fileio
import perceptron

batchsize = 10

def iterate_minibatches(inputs, targets, batchsize, shuffle=False):
    assert len(inputs) == len(targets)
    if shuffle:
        indices = np.arange(len(inputs))
        np.random.shuffle(indices)
    for start_idx in range(0, len(inputs) - batchsize + 1, batchsize):
        if shuffle:
            excerpt = indices[start_idx:start_idx + batchsize]
        else:
            excerpt = slice(start_idx, start_idx + batchsize)
        yield inputs[excerpt], targets[excerpt]

fr = fileio.FaceReader("../data/training.csv", "../data/training.pkl.gz", fast_nrows=110)
fr.Read()

X, Y = fr.GetData()

train_X = X[1:100]
train_Y = Y[1:100]

test_X = X[100:110]
test_Y = Y[100:110]

mlp = perceptron.ConvolutionalMLP(
	(batchsize, 1, 96, 96), # input shape
	0.2, # input drop-rate
	[96*96, 96*96], # hidden_layer_widths,
	[0.5, 0.5], # hidden_drop_rate,
	lasagne.nonlinearities.rectify, # hidden_layer_nonlinearity,
	30) # output_width
print(mlp)
mlp.BuildNetwork()

# Finally, launch the training loop.
print("Starting training...")
# We iterate over epochs:
num_epochs=3
for epoch in range(num_epochs):
    # In each epoch, we do a full pass over the training data:
    train_err = 0
    train_batches = 0
    start_time = time.time()
    for batch in iterate_minibatches(train_X, train_Y, batchsize, shuffle=False):
        inputs, targets = batch
        train_err += mlp.train_fn(inputs, targets)
        train_batches += 1

    # And a full pass over the validation data:
    val_err = 0
    val_batches = 0
    for batch in iterate_minibatches(test_X, test_Y, batchsize, shuffle=False):
        inputs, targets = batch
        err = mlp.val_fn(inputs, targets)
        val_err += err
        val_batches += 1

    # Then we print the results for this epoch:
    print("Epoch {} of {} took {:.3f}s".format(
        epoch + 1, num_epochs, time.time() - start_time))
    print("  training loss:\t\t{:.6f}".format(train_err / train_batches))
    print("  validation loss:\t\t{:.6f}".format(val_err / val_batches))

# After training, we compute and print the test error:
test_err = 0
test_batches = 0
for batch in iterate_minibatches(data.test_X, data.test_Y, batchsize, shuffle=False):
    inputs, targets = batch
    err = mlp.val_fn(inputs, targets)
    test_err += err
    test_batches += 1
print("Final results:")
print("  test loss:\t\t\t{:.6f}".format(test_err / test_batches))

# Predict the first 500 (batchsize)
# y_pred = mlp.Predict(data.test_X[0:batchsize])

import code
code.interact(local=locals())<commit_msg>fix bug in main<commit_after>#!/usr/bin/env python
import operator
import time

import lasagne
import numpy as np
import theano
import theano.tensor as T
from sklearn.metrics import confusion_matrix

import fileio
import perceptron

batchsize = 10

def iterate_minibatches(inputs, targets, batchsize, shuffle=False):
    assert len(inputs) == len(targets)
    if shuffle:
        indices = np.arange(len(inputs))
        np.random.shuffle(indices)
    for start_idx in range(0, len(inputs) - batchsize + 1, batchsize):
        if shuffle:
            excerpt = indices[start_idx:start_idx + batchsize]
        else:
            excerpt = slice(start_idx, start_idx + batchsize)
        yield inputs[excerpt], targets[excerpt]

fr = fileio.FaceReader("../data/training.csv", "../data/training.pkl.gz", fast_nrows=110)
fr.Read()

X, Y = fr.GetData()

train_X = X[1:100]
train_Y = Y[1:100]

test_X = X[100:110]
test_Y = Y[100:110]

mlp = perceptron.ConvolutionalMLP(
	(batchsize, 1, 96, 96), # input shape
	0.2, # input drop-rate
	[96*96, 96*96], # hidden_layer_widths,
	[0.5, 0.5], # hidden_drop_rate,
	lasagne.nonlinearities.rectify, # hidden_layer_nonlinearity,
	30) # output_width
print(mlp)
mlp.BuildNetwork()

# Finally, launch the training loop.
print("Starting training...")
# We iterate over epochs:
num_epochs=3
for epoch in range(num_epochs):
    # In each epoch, we do a full pass over the training data:
    train_err = 0
    train_batches = 0
    start_time = time.time()
    for batch in iterate_minibatches(train_X, train_Y, batchsize, shuffle=False):
        inputs, targets = batch
        train_err += mlp.train_fn(inputs, targets)
        train_batches += 1

    # And a full pass over the validation data:
    val_err = 0
    val_batches = 0
    for batch in iterate_minibatches(test_X, test_Y, batchsize, shuffle=False):
        inputs, targets = batch
        err = mlp.val_fn(inputs, targets)
        val_err += err
        val_batches += 1

    # Then we print the results for this epoch:
    print("Epoch {} of {} took {:.3f}s".format(
        epoch + 1, num_epochs, time.time() - start_time))
    print("  training loss:\t\t{:.6f}".format(train_err / train_batches))
    print("  validation loss:\t\t{:.6f}".format(val_err / val_batches))

# After training, we compute and print the test error:
test_err = 0
test_batches = 0
for batch in iterate_minibatches(test_X, test_Y, batchsize, shuffle=False):
    inputs, targets = batch
    err = mlp.val_fn(inputs, targets)
    test_err += err
    test_batches += 1
print("Final results:")
print("  test loss:\t\t\t{:.6f}".format(test_err / test_batches))

# Predict the first 500 (batchsize)
# y_pred = mlp.Predict(data.test_X[0:batchsize])

import code
code.interact(local=locals())<|endoftext|>


<commit_before># @license
# Copyright 2016 Google Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import posixpath

static_content_filenames = set(['main.bundle.js', 'chunk_worker.bundle.js', 'styles.css', 'index.html'])

mime_type_map = {
    '.css': 'text/css',
    '.js': 'application/javascript',
    '.html': 'text/html',
    '.map': 'application/json'
}

def guess_mime_type_from_path(path):
    return mime_type_map.get(posixpath.splitext(path)[1], 'application/octet-stream')

class StaticContentSource(object):
    def get(self, name):
        if name == '':
            name = 'index.html'
        return self.get_content(name), guess_mime_type_from_path(name)

    def get_content(self, name):
        raise NotImplementedError

class PkgResourcesContentSource(StaticContentSource):
    def get_content(self, name):
        import pkg_resources
        if name not in static_content_filenames:
            raise ValueError('Invalid static resource name: %r' % name)
        if pkg_resources.resource_exists(__name__, name):
            return pkg_resources.resource_string(__name__, name)
        raise ValueError(
            'Static resources not built.  Run: "python setup.py bundle_client" or use an alternative static content source.'
        )

class HttpSource(StaticContentSource):
    def __init__(self, url):
        self.url = url

    def get_content(self, name):
        import requests
        full_url = posixpath.join(self.url, name)
        r = requests.get(full_url)
        if r.status_code >= 200 and r.status_code < 300:
            return r.content
        raise ValueError('Failed to retrieve %r: %s' % (full_url, r.reason))

class FileSource(StaticContentSource):
    def __init__(self, path=None, file_open=open):
        if path is None:
            path = os.path.join(os.path.dirname(__file__), '../../../dist/dev-python')
        self.file_path = path
        self.file_open = file_open

    def get_content(self, name):
        full_path = os.path.join(self.file_path, name)
        try:
            with self.file_open(full_path, 'rb') as f:
                return f.read()
        except Exception as e:
            raise ValueError('Failed to read local path %r: %s' % (full_path, e))

dist_dev_static_content_source = FileSource()

def get_default_static_content_source():
    return PkgResourcesContentSource()

def get_static_content_source(source=None, url=None, path=None, file_open=None):
    if source is not None:
        return source
    elif url is not None:
        return HttpSource(url)
    elif path is not None:
        return FileSource(path=path, file_open=file_open)
    else:
        return get_default_static_content_source()
<commit_msg>python(FileSource): properly handle default value for file_open<commit_after># @license
# Copyright 2016 Google Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import posixpath

static_content_filenames = set(['main.bundle.js', 'chunk_worker.bundle.js', 'styles.css', 'index.html'])

mime_type_map = {
    '.css': 'text/css',
    '.js': 'application/javascript',
    '.html': 'text/html',
    '.map': 'application/json'
}

def guess_mime_type_from_path(path):
    return mime_type_map.get(posixpath.splitext(path)[1], 'application/octet-stream')

class StaticContentSource(object):
    def get(self, name):
        if name == '':
            name = 'index.html'
        return self.get_content(name), guess_mime_type_from_path(name)

    def get_content(self, name):
        raise NotImplementedError

class PkgResourcesContentSource(StaticContentSource):
    def get_content(self, name):
        import pkg_resources
        if name not in static_content_filenames:
            raise ValueError('Invalid static resource name: %r' % name)
        if pkg_resources.resource_exists(__name__, name):
            return pkg_resources.resource_string(__name__, name)
        raise ValueError(
            'Static resources not built.  Run: "python setup.py bundle_client" or use an alternative static content source.'
        )

class HttpSource(StaticContentSource):
    def __init__(self, url):
        self.url = url

    def get_content(self, name):
        import requests
        full_url = posixpath.join(self.url, name)
        r = requests.get(full_url)
        if r.status_code >= 200 and r.status_code < 300:
            return r.content
        raise ValueError('Failed to retrieve %r: %s' % (full_url, r.reason))

class FileSource(StaticContentSource):
    def __init__(self, path=None, file_open=None):
        if path is None:
            path = os.path.join(os.path.dirname(__file__), '../../../dist/dev-python')
        self.file_path = path
        self.file_open = file_open or open

    def get_content(self, name):
        full_path = os.path.join(self.file_path, name)
        try:
            with self.file_open(full_path, 'rb') as f:
                return f.read()
        except Exception as e:
            raise ValueError('Failed to read local path %r: %s' % (full_path, e))

dist_dev_static_content_source = FileSource()

def get_default_static_content_source():
    return PkgResourcesContentSource()

def get_static_content_source(source=None, url=None, path=None, file_open=None):
    if source is not None:
        return source
    elif url is not None:
        return HttpSource(url)
    elif path is not None:
        return FileSource(path=path, file_open=file_open)
    else:
        return get_default_static_content_source()
<|endoftext|>


<commit_before>from bitarray import bitarray
import struct
import random
import hashlib

# Number of simultaneous requests made to "prime the pump" after handshake
SIM_REQUESTS = 20

class Peer(object):
    # Can't initialize without a dictionary. Handshake
    # takes place using socket before peer init
    def __init__(self, sock, reactor, torrent):
        self.sock = sock
        self.sock.setblocking(True)
        self.reactor = reactor
        self.torrent = torrent
        self.valid_indices = []
        self.bitfield = None
        self.max_size = 16 * 1024
        self.states = {'reading_length': 0, 'reading_id': 1,
                       'reading_message': 2}
        self.save_state = {'state': self.states['reading_length'],
                           'length': 0, 'message_id': None,
                           'message': '', 'remainder': ''}
        self.message_codes = ['choke', 'unchoke', 'interested',
                              'not interested', 'have', 'bitfield', 'request',
                              'piece', 'cancel', 'port']
        self.ischoking = True
        self.isinterested = False

    def fileno(self):
        return self.sock.fileno()

    def getpeername(self):
        return self.sock.getpeername()

    def read(self):
        bytes = self.sock.recv(self.max_size)
        if len(bytes) == 0:
            print 'Got 0 bytes from fileno {}.'.format(self.fileno())
            self.torrent.kill_peer(self)
        self.process_input(bytes)

    def process_input(self, bytes):
        while bytes:
            if self.save_state['state'] == self.states['reading_length']:
                bytes = self.get_message_length(bytes)
            elif self.save_state['state'] == self.states['reading_id']:
                bytes = self.get_message_id(bytes)
            elif self.save_state['state'] == self.states['reading_message']:
                bytes = self.get_message(bytes)

    def get_message_length(self, instr):

            # If we already have a partial message, start with that
            if self.save_state['remainder']:
                print 'We have a remainder at the top of get_message_length'
                instr = self.save_state['remainder'] + instr
                self.save_state['remainder'] = ''

            # If we have four bytes we can at least read the length
            if len(instr) >= 4:

                # Need 0 index because struct.unpack returns tuple
                # save_state['length'] is based on what the peer *says*, not
                # on the length of the actual message
                self.save_state['length'] = struct.unpack('!i', instr[0:4])[0]
                if self.save_state['length'] == 0:
                    self.keep_alive()
                    self.save_state['state'] = self.states['reading_length']
                    return instr[4:]
                else:
                    self.save_state['state'] = self.states['reading_id']
                    return instr[4:]

            # Less than four bytes and we save + wait for next read
            # Increeedibly unlikely to happen
            else:
                self.save_state['remainder'] = instr
                return ''

    def get_message_id(self, instr):
        self.save_state['message_id'] = struct.unpack('b', instr[0])[0]
        self.save_state['state'] = self.states['reading_message']
        return instr[1:]

    def get_message(self, instr):
        # Since one byte is getting used up for the message_id
        length_after_id = self.save_state['length'] - 1
        if length_after_id == 0:
            self.save_state['state'] = self.states['reading_length']
            self.save_state['message_id'] = None
            self.save_state['message'] = ''
            return instr

        if self.save_state['remainder']:
            instr = self.save_state['remainder'] + instr

        # If we have more than what we need we act on the full message and
        # return the rest
        if len(instr) >= length_after_id:

            self.save_state['message'] = instr[:length_after_id]

            # If we hit handle_message we know that we have a FULL MESSAGE
            # All the stateful stuff can go in the garbage
            self.handle_message()
            self.reset_state()
            return instr[length_after_id:]

        # Otherwise we stash what we have and keep things the way they are
        else:
            self.save_state['remainder'] = instr
            return None

    def reset_state(self):
        self.save_state['state'] = self.states['reading_length']
        self.save_state['length'] = 0
        self.save_state['message_id'] = None
        self.save_state['message'] = ''
        self.save_state['remainder'] = ''

    # This is only getting called when I have a complete message
    def handle_message(self):
        if self.save_state['message_id'] == 0:
            self.pchoke()
        elif self.save_state['message_id'] == 1:
            self.punchoke()
        elif self.save_state['message_id'] == 2:
            self.pinterested()
        elif self.save_state['message_id'] == 3:
            self.pnotinterested()
        elif self.save_state['message_id'] == 4:
            self.phave()
        elif self.save_state['message_id'] == 5:
            self.pbitfield()
        elif self.save_state['message_id'] == 6:
            self.prequest()
        elif self.save_state['message_id'] == 7:
            self.ppiece(self.save_state['message'])
        elif self.save_state['message_id'] == 8:
            self.pcancel()
        elif self.save_state['message_id'] == 9:
            pass

    def pchoke(self):
        print 'choke'
        self.ischoking = True

    def punchoke(self):
        print 'unchoke'
        self.ischoking = False

    def pinterested(self):
        print 'pinterested'

    def pnotinterested(self):
        print 'pnotinterested'

    def phave(self):
        index = struct.unpack('>i', self.save_state['message'])[0]
        self.bitfield[index] = True

    def pbitfield(self):
        self.bitfield = bitarray()
        self.bitfield.frombytes(self.save_state['message'])
        self.interested()
        self.unchoke()
        self.mad_requests()

    def prequest(self):
        print 'prequest'

    def ppiece(self, content):
        '''
        Process a piece that we've received from a peer, writing it out to
        one or more files
        '''
        piece_index, block_begin = struct.unpack('!ii', content[0:8])
        piece_dict = {'kind': 'piece', 'peer': self.sock.getpeername(),
                      'piece_index': piece_index}
        self.torrent.switchboard.try_vis_handoff(piece_dict)
        block = content[8:]
        if hashlib.sha1(block).digest() == (self.torrent.torrent_dict['info']
                                            ['pieces']
                                            [20 * piece_index:20 * piece_index
                                             + 20]):
            print 'hash matches'
            print ('writing piece {}. Length is '
                   '{}').format(repr(block)[:10] + '...', len(block))
            byte_index = piece_index * self.torrent.piece_length
            self.torrent.switchboard.write(byte_index, block)
            self.torrent.switchboard.mark_off(piece_index)
            print self.torrent.switchboard.bitfield
            if self.torrent.switchboard.complete:
                print '\nDownload complete\n'
                self.reactor.is_running = False
        else:
            raise Exception("hash of piece doesn't"
                            "match hash in torrent_dict")
        self.request_piece()

    def pcancel(self):
        print 'pcancel'

    def read_timeout(self):
        self.mad_requests()

    def mad_requests(self):
        for i in xrange(SIM_REQUESTS):
            self.request_piece()
            print ('Just made {}th consecutive request'
                   'to {}'.format(i, self.fileno()))

    def request_piece(self):
        if not self.valid_indices:
            # We want a list of all indices where:
            #   - We're interested in the piece
            #     (i.e., it's in torrent.outfile.bitfield)
            #   - The peer has the piece (it's available)
            for i in range(self.torrent.num_pieces):
                if (self.torrent.switchboard.bitfield[i] is True
                        and self.bitfield[i] is True):
                    self.valid_indices.append(i)
            if not self.valid_indices:
                return
            else:
                random.shuffle(self.valid_indices)
        self.request(self.valid_indices.pop())

    def interested(self):
        packet = ''.join(struct.pack('!ib', 1, 2))
        self.sock.send(packet)

    def unchoke(self):
        packet = struct.pack('!ib', 1, 1)
        self.sock.send(packet)

    def keep_alive(self):
        print 'inside keep_alive'

    def write(self):
        pass

    def request(self, want_piece):
        if want_piece == self.torrent.num_pieces - 1:
            piece_length = self.torrent.last_piece_length
        else:
            piece_length = self.torrent.piece_length
        packet = ''.join(struct.pack('!ibiii', 13, 6, want_piece, 0,
                         piece_length))
        bytes = self.sock.send(packet)
        request_dict = {'kind': 'request',
                        'peer': self.sock.getpeername(),
                        'piece': want_piece}
        self.torrent.switchboard.try_vis_handoff(request_dict)
        print 'next request:', request_dict
        if bytes != len(packet):
            raise Exception('couldnt send request')
<commit_msg>Added activate message to peer init<commit_after>from bitarray import bitarray
import struct
import random
import hashlib
import pudb

# Number of simultaneous requests made to "prime the pump" after handshake
SIM_REQUESTS = 20

class Peer(object):
    # Can't initialize without a dictionary. Handshake
    # takes place using socket before peer init
    def __init__(self, sock, reactor, torrent):
        self.sock = sock
        self.sock.setblocking(True)
        self.reactor = reactor
        self.torrent = torrent
        self.valid_indices = []
        self.bitfield = None
        self.max_size = 16 * 1024
        self.states = {'reading_length': 0, 'reading_id': 1,
                       'reading_message': 2}
        self.save_state = {'state': self.states['reading_length'],
                           'length': 0, 'message_id': None,
                           'message': '', 'remainder': ''}
        self.message_codes = ['choke', 'unchoke', 'interested',
                              'not interested', 'have', 'bitfield', 'request',
                              'piece', 'cancel', 'port']
        self.ischoking = True
        self.isinterested = False
        activate_dict = {'kind': 'activate', 'address': self.getpeername()}
        # pudb.set_trace()
        self.torrent.switchboard.try_vis_handoff(activate_dict)

    def fileno(self):
        return self.sock.fileno()

    def getpeername(self):
        return self.sock.getpeername()

    def read(self):
        bytes = self.sock.recv(self.max_size)
        if len(bytes) == 0:
            print 'Got 0 bytes from fileno {}.'.format(self.fileno())
            self.torrent.kill_peer(self)
        self.process_input(bytes)

    def process_input(self, bytes):
        while bytes:
            if self.save_state['state'] == self.states['reading_length']:
                bytes = self.get_message_length(bytes)
            elif self.save_state['state'] == self.states['reading_id']:
                bytes = self.get_message_id(bytes)
            elif self.save_state['state'] == self.states['reading_message']:
                bytes = self.get_message(bytes)

    def get_message_length(self, instr):

            # If we already have a partial message, start with that
            if self.save_state['remainder']:
                print 'We have a remainder at the top of get_message_length'
                instr = self.save_state['remainder'] + instr
                self.save_state['remainder'] = ''

            # If we have four bytes we can at least read the length
            if len(instr) >= 4:

                # Need 0 index because struct.unpack returns tuple
                # save_state['length'] is based on what the peer *says*, not
                # on the length of the actual message
                self.save_state['length'] = struct.unpack('!i', instr[0:4])[0]
                if self.save_state['length'] == 0:
                    self.keep_alive()
                    self.save_state['state'] = self.states['reading_length']
                    return instr[4:]
                else:
                    self.save_state['state'] = self.states['reading_id']
                    return instr[4:]

            # Less than four bytes and we save + wait for next read
            # Increeedibly unlikely to happen
            else:
                self.save_state['remainder'] = instr
                return ''

    def get_message_id(self, instr):
        self.save_state['message_id'] = struct.unpack('b', instr[0])[0]
        self.save_state['state'] = self.states['reading_message']
        return instr[1:]

    def get_message(self, instr):
        # Since one byte is getting used up for the message_id
        length_after_id = self.save_state['length'] - 1
        if length_after_id == 0:
            self.save_state['state'] = self.states['reading_length']
            self.save_state['message_id'] = None
            self.save_state['message'] = ''
            return instr

        if self.save_state['remainder']:
            instr = self.save_state['remainder'] + instr

        # If we have more than what we need we act on the full message and
        # return the rest
        if len(instr) >= length_after_id:

            self.save_state['message'] = instr[:length_after_id]

            # If we hit handle_message we know that we have a FULL MESSAGE
            # All the stateful stuff can go in the garbage
            self.handle_message()
            self.reset_state()
            return instr[length_after_id:]

        # Otherwise we stash what we have and keep things the way they are
        else:
            self.save_state['remainder'] = instr
            return None

    def reset_state(self):
        self.save_state['state'] = self.states['reading_length']
        self.save_state['length'] = 0
        self.save_state['message_id'] = None
        self.save_state['message'] = ''
        self.save_state['remainder'] = ''

    # This is only getting called when I have a complete message
    def handle_message(self):
        if self.save_state['message_id'] == 0:
            self.pchoke()
        elif self.save_state['message_id'] == 1:
            self.punchoke()
        elif self.save_state['message_id'] == 2:
            self.pinterested()
        elif self.save_state['message_id'] == 3:
            self.pnotinterested()
        elif self.save_state['message_id'] == 4:
            self.phave()
        elif self.save_state['message_id'] == 5:
            self.pbitfield()
        elif self.save_state['message_id'] == 6:
            self.prequest()
        elif self.save_state['message_id'] == 7:
            self.ppiece(self.save_state['message'])
        elif self.save_state['message_id'] == 8:
            self.pcancel()
        elif self.save_state['message_id'] == 9:
            pass

    def pchoke(self):
        print 'choke'
        self.ischoking = True

    def punchoke(self):
        print 'unchoke'
        self.ischoking = False

    def pinterested(self):
        print 'pinterested'

    def pnotinterested(self):
        print 'pnotinterested'

    def phave(self):
        index = struct.unpack('>i', self.save_state['message'])[0]
        self.bitfield[index] = True

    def pbitfield(self):
        self.bitfield = bitarray()
        self.bitfield.frombytes(self.save_state['message'])
        self.interested()
        self.unchoke()
        self.mad_requests()

    def prequest(self):
        print 'prequest'

    def ppiece(self, content):
        '''
        Process a piece that we've received from a peer, writing it out to
        one or more files
        '''
        piece_index, block_begin = struct.unpack('!ii', content[0:8])
        piece_dict = {'kind': 'piece', 'peer': self.sock.getpeername(),
                      'piece_index': piece_index}
        self.torrent.switchboard.try_vis_handoff(piece_dict)
        block = content[8:]
        if hashlib.sha1(block).digest() == (self.torrent.torrent_dict['info']
                                            ['pieces']
                                            [20 * piece_index:20 * piece_index
                                             + 20]):
            print 'hash matches'
            print ('writing piece {}. Length is '
                   '{}').format(repr(block)[:10] + '...', len(block))
            byte_index = piece_index * self.torrent.piece_length
            self.torrent.switchboard.write(byte_index, block)
            self.torrent.switchboard.mark_off(piece_index)
            print self.torrent.switchboard.bitfield
            if self.torrent.switchboard.complete:
                print '\nDownload complete\n'
                self.reactor.is_running = False
        else:
            raise Exception("hash of piece doesn't"
                            "match hash in torrent_dict")
        self.request_piece()

    def pcancel(self):
        print 'pcancel'

    def read_timeout(self):
        self.mad_requests()

    def mad_requests(self):
        for i in xrange(SIM_REQUESTS):
            self.request_piece()
            print ('Just made {}th consecutive request'
                   'to {}'.format(i, self.fileno()))

    def request_piece(self):
        if not self.valid_indices:
            # We want a list of all indices where:
            #   - We're interested in the piece
            #     (i.e., it's in torrent.outfile.bitfield)
            #   - The peer has the piece (it's available)
            for i in range(self.torrent.num_pieces):
                if (self.torrent.switchboard.bitfield[i] is True
                        and self.bitfield[i] is True):
                    self.valid_indices.append(i)
            if not self.valid_indices:
                return
            else:
                random.shuffle(self.valid_indices)
        self.request(self.valid_indices.pop())

    def interested(self):
        packet = ''.join(struct.pack('!ib', 1, 2))
        self.sock.send(packet)

    def unchoke(self):
        packet = struct.pack('!ib', 1, 1)
        self.sock.send(packet)

    def keep_alive(self):
        print 'inside keep_alive'

    def write(self):
        pass

    def request(self, want_piece):
        if want_piece == self.torrent.num_pieces - 1:
            piece_length = self.torrent.last_piece_length
        else:
            piece_length = self.torrent.piece_length
        packet = ''.join(struct.pack('!ibiii', 13, 6, want_piece, 0,
                         piece_length))
        bytes = self.sock.send(packet)
        request_dict = {'kind': 'request',
                        'peer': self.sock.getpeername(),
                        'piece': want_piece}
        self.torrent.switchboard.try_vis_handoff(request_dict)
        print 'next request:', request_dict
        if bytes != len(packet):
            raise Exception('couldnt send request')
<|endoftext|>