memed/bot/ext/rsudo.py

import logging
import asyncio

import discord

from .common import Cog

log = logging.getLogger(__name__)


async def shell(command: str):
    process = await asyncio.create_subprocess_shell(
        command,
        stderr=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
    )

    out, err = map(lambda s: s.decode('utf-8'), await process.communicate())
    return f'{out}{err}'.strip()


class Rsudo(Cog):
    def __init__(self, bot):
        super().__init__(bot)
        self.command_channel = None

    async def create_request(self, message):
        if not self.command_channel:
            self.command_channel = self.bot.get_channel(
                self.bot.config.command_channel)

        if not self.command_channel:
            log.warning('command channel not found')
            return

        # parse it, follows format command,uid
        uid = message.split(',')[-1]
        command = ','.join(message.split(',')[:-1])

        log.info(f'[rsudo] {uid!r} {command!r}')
        e = discord.Embed(title=f'rsudo from uid {uid}')
        e.add_field(name='command', value=f'`{command}`')
        m = await self.command_channel.send(embed=e)

        await m.add_reaction('\N{WHITE HEAVY CHECK MARK}')
        await m.add_reaction('\N{CROSS MARK}')

        emotes = ['\N{WHITE HEAVY CHECK MARK}',
                  '\N{CROSS MARK}']

        def check(reaction, user):
            return user and \
                reaction.emoji in emotes and \
                any(x.id == self.bot.config.admin_role for x in user.roles)

        reaction, user = await self.bot.wait_for('reaction_add', check=check)

        if reaction.emoji == emotes[0]:
            # do
            await self.command_channel.send('executing')
            out = await shell(command)
            fmt = f'```\n{out}\n```'
            await self.command_channel.send(fmt)
        else:
            # dont do
            await self.command_channel.send('denied with success')


def setup(bot):
    bot.add_cog(Rsudo(bot))