telegram_baskit_bot.py

#!/usr/bin/env python
from threading import Thread
from io import StringIO
import telepot, re, sys, os, requests, subprocess

class TelegramBaskitBot(Thread):
    def __init__(self, token, adminuser=None, adminid=None, address=None, port=None, sudo=True, baskit=None):
        # Populate the required fields
        self.token = token
        self.adminuser = adminuser
        self.admin_id = adminid
        self.use_sudo = sudo
        self.sudo_user = 'root'
        self.remsg = re.compile(r'\!baskit (.*)')

        # Check to see if an address has been set.  If it has, then lets use it.
        # If one hasn't, then we will lookup the public IP that this host has
        # and use that.
        if address:
            self.address = address
        else:
            self.address = requests.get('http://ipinfo.io/ip').content.strip('\n')

        # The default minecraft port is 25565.  If one hasn't been specified at
        # instantiation, then we will go ahead and use that port number.
        if port:
            self.port = port
        else:
            self.port = 25565

        # Now to set the baskit path.  We will need this to know where the
        # baskit binary lives.
        if baskit:
            self.baskit_path = baskit
        else:
            self.baskit_path = '/usr/local/bin/baskit'

        # Now lets go ahead and create the bot that we will use.
        self.bot = telepot.Bot(self.token)
        self.bot.setWebhook()

        # Set the username that we have gathered from the bot.
        self.username = ('@{0}').format(bot.getMe()['username'])

    def command(self, command):
        '''
        Generic Shell command method
        '''
        result = subprocess.run(command, stdout=subprocess.PIPE)
        return result.stdout

    def baskit(self, command):
        '''
        Sends a command to the baskit interpretor
        '''
        # The base command.  As we're just accepting a string for the command
        # we will pass it on as such.  It shouldn't cause an issue.
        cmd = [self.baskit_path, command]

        # If sudo is specified, then prepend the sudo command.
        if self.use_sudo:
            cmd = ['sudo', '-u', self.sudo_user] + cmd
        return self.comand(cmd)

    def bot_send(self, channel, message):
        '''
        Sends a message through the bot to the specified channel.
        '''
        self.bot.sendMessage(channel, message)

    def run(self):
        '''
        The run method.  This will be run by using the start method.
        '''
        def myid(message):
            return ('{0}:{1}').format(
                message['chat']['id'], 
                message['chat']['username']
            )

        def handler(message):
            '''
            Generic message handler.  Here is where we parse and feed data back
            and forth from the bot.
            '''
            command = self.remsg.findall(message['text'])

            # We want to make sure that the command isn't empty.
            if len(command) > 0:
                # Well looky there, we got something, lets overload the command
                # variable with the first response we got.  We dont want to look
                # any anything else that may be returned by our greedy regex.
                command = command[0]

                # If the user is an admin, we will let them run whatever commands
                # they would like to run on the host.
                if message['chat']['username'] == self.adminuser:
                    if command == 'sysreboot':
                        self.bot_send(self.admin_id, ('Restarting {0}.').format(self.address))
                        self.command(['sudo', 'reboot'])
                    elif command == 'myid':
                        self.bot_send(message['chat']['id'], myid(message))
                    elif command not in ['myip', 'help']:
                        self.bot_send(self.admin_id, self.baskit(command))
                else:
                    # These are the restricted commands that a non-admin can run.
                    if command == 'players':
                        self.bot_send(message['chat']['id'], self.baskit('players'))
                    elif command == 'myid':
                        self.bot_send(message['chat']['id'], myid(message))

                if command == 'help':
                    msg = [
                        '!baskit myid - Returns the chat id and username in ID:USERNAME format',
                        '!baskit players - Returns the list of players'
                    ]
                    if message['chat']['username'] == self.adminuser:
                        msg.append('!baskit sysreboot - Reboots the host')
                        msg.append('!baskit COMMAND - Will pipe COMMAND to baskit')
                        self.bot_send(self.admin_id, msg.join('\n'))
                    else:
                        self.bot_send(message['chat']['id'], msg.join('\n'))

        # Lastly lets go ahead and run the loop.
        self.bot.message_loop(handler)


if __name__ == '__main__':

    # first thing we need to do is hunt down the baskit.conf file.  We will be
    # using a new section in the config file for our purposes.
    # Example:
    # [TelegramBot]
    # token = TOKEN
    # adminid = TELEGRAMID
    # adminuser = TELEGRAMUSERNAME
    import configparser
    config = configparser.ConfigParser()
    config.read('/etc/baskit.conf', 'baskit.conf')

    bot = TelegramBaskitBot(config.get('TelegramBot', 'token'))
    if config.has_option('TelegramBot', 'adminid'):
        bot.admin_id = config.get('TelegramBot', 'adminid')
    if config.has_option('TelegramBot', 'adminuser'):
        bot.adminuser = config.get('TelegramBot', 'adminuser')

    # Start thew thread and join it to the main thread.
    bot.start()
    bot.join()