minecraft_manager/bot/discord.py

import asyncio
import logging
import traceback
import threading
from enum import Enum

import discord
from discord.ext import commands
from django.conf import settings

logger = logging.getLogger(__name__)

discord_loop = None
discord_bot = None
description = '''
A Discord bot connected to an MCM instance.
'''


class DiscordStatus(Enum):
    STOPPED = 0
    STARTING = 1
    STARTED = 2
    STOPPING = 3

    def __str__(self):
        return self.name.title()

    def is_running(self):
        return self != DiscordStatus.STOPPED


discord_status = DiscordStatus.STOPPED


class Discord(commands.Bot):
    discord_game = 'MCM'
    prefix = getattr(settings, 'DISCORD_BOT_PREFIX', '!')
    auth_roles = getattr(settings, 'DISCORD_BOT_ROLES', [])
    superuser_roles = getattr(settings, 'DISCORD_SUPERUSER_ROLES', [])
    error_users = getattr(settings, 'DISCORD_ERROR_USERS', [])

    def __init__(self, token, loop):
        super().__init__(command_prefix=self.prefix, description=description, case_insensitive=True, help_command=None, activity=discord.Game(name=self.discord_game), loop=loop)
        self.token = token
        self.load_extension("minecraft_manager.bot.commands")

    async def on_ready(self):
        global discord_status
        discord_status = DiscordStatus.STARTED
        print('Logged in as')
        print(self.user.name)
        print(self.user.id)
        print(discord.__version__)
        print('Voice Loaded: {0}'.format(discord.opus.is_loaded()))
        print('OAuth URL: https://discordapp.com/api/oauth2/authorize?client_id={0}&permissions=0&scope=bot'.format(self.user.id))
        print('------')
        print('Logged in as {0} ({1}) with discord.py v{2}'.format(self.user.name, self.user.id, discord.__version__))

    async def on_disconnect(self):
        global discord_status
        discord_status = DiscordStatus.STOPPED

    async def discord_message(self, dest, message):
        if isinstance(message, discord.Embed):
            for idx, field in enumerate(message.fields):
                if not field.value:
                    message.set_field_at(idx, name=field.name, value="N/A")
            await dest.send(embed=message)
        else:
            await dest.send(message)

    async def on_command_error(self, context, exception):
        if not isinstance(exception, commands.CommandInvokeError):
            return
        if hasattr(exception, "original"):
            error = ''.join(traceback.format_tb(exception.original.__traceback__))
        else:
            error = exception
        logger.error(error)
        for user_id in self.error_users:
            user = self.get_user(user_id)
            if user:
                await self.discord_message(user, '```python\n{}```'.format(error))

    def run_bot(self):
        global discord_loop
        try:
            discord_loop.run_until_complete(self.start(self.token))
        except KeyboardInterrupt:
            logger.info("Bot received keyboard interrupt")
        except Exception as e:
            print(e)
            logger.info('Bot encountered the following unhandled exception %s', e)


def start():
    global discord_loop, discord_bot, discord_status
    if discord_status != DiscordStatus.STOPPED:
        return
    token = getattr(settings, 'DISCORD_BOT_TOKEN', None)
    discord_loop = asyncio.new_event_loop()
    discord_bot = Discord(token, discord_loop)
    thread = threading.Thread(target=discord_bot.run_bot)
    thread.start()
    discord_status = DiscordStatus.STARTING


def stop():
    global discord_loop, discord_bot, discord_status
    if discord_status == DiscordStatus.STARTED:
        discord_loop.create_task(discord_bot.close())
        discord_status = DiscordStatus.STOPPING
        discord_loop = None
        discord_bot = None


def restart():
    def _restart():
        stop()
        while discord_status.is_running():
            pass
        start()
    if discord_status != DiscordStatus.STARTED:
        return
    thread = threading.Thread(target=_restart)
    thread.start()


def status():
    return discord_status.is_running()


def display():
    return str(discord_status)


class DestType(Enum):
    CHANNEL = 1
    USER = 2


def send(dest_type: DestType, dest_id: int, message: str = "", embed: discord.Embed = None):
    async def _send():
        if dest_type == DestType.CHANNEL:
            dest = discord_bot.get_channel(dest_id)
        elif dest_type == DestType.USER:
            dest = discord_bot.get_user(dest_id)
        else:
            return
        if message is not None:
            await dest.send(message)
        if embed is not None:
            for idx, field in enumerate(embed.fields):
                if not field.value:
                    embed.set_field_at(idx, name=field.name, value="N/A")
            await dest.send(embed=embed)
    global discord_loop, discord_bot
    if discord_loop:
        discord_loop.create_task(_send())