slackbot.py

slackbot.py

__all__ = [
    'SlackBot',
    'RecBot',
]

import os
import time
import re
import subprocess
from functools import wraps

from slackclient import SlackClient
import log_helper


def command_handler(func):
    """
    Decorator which runs preprocessing on command argument
    - removes regexp from command text
    - replaces aliases
    - posts return value of function to channel
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        regexp = kwargs.get('regexp', '')
        command = kwargs.get('command', '')
        channel = kwargs.get('channel', '')
        self = args[0]
        if command and regexp:
            command = self._preprocess_command(command, regexp)
            kwargs['command'] = command

        message = func(*args, **kwargs)
        self._post(text=message, channel=channel)
        return message

    return wrapper


# noinspection PyUnusedLocal
class SlackBot:

    def __init__(self, token: str='', logger=None):

        self.logger = logger or log_helper.configure_logger('slackbot', verbose=False)

        self._token = token or os.environ.get('SLACK_BOT_TOKEN')
        self._read_websocket_delay_sec = 1  # 1 second delay between reading from firehose
        self._slack_client = SlackClient(self._token)

        self.rules = [
            (r'^!echo', self.echo),
        ]
        self.aliases = [
            (r'\bll\b', 'ls -lah'),
        ]

    def listen(self):
        try:
            self._connect()
            self.logger.info("StarterBot connected and running!")
            while True:
                data = self._slack_client.rtm_read()
                for command, channel in self._parse_slack_output(data):
                    self._handle_command(command, channel)
                time.sleep(self._read_websocket_delay_sec)
        except KeyboardInterrupt:
            self.logger.info('Graceful exit initiated')

    @command_handler
    def echo(self, command: str, channel: str, regexp: str) -> str:
        return command

    def _connect(self):
        connection = self._slack_client.rtm_connect()
        if not connection:
            raise IOError('Connection failed.')

    def _handle_command(self, command: str, channel: str):
        for pattern, action in self.rules:
            if re.match(pattern, command):
                action(command=command, channel=channel, regexp=pattern)

    @staticmethod
    def _do_read_message(message: str) -> bool:
        return len(message) and message[0] == '!'

    def _post(self, text: str, channel: str):
        self._slack_client.api_call("chat.postMessage", channel=channel, text=text, as_user=True)

    def _preprocess_command(self, command: str, regexp: str) -> str:
        command = re.sub(regexp, '', command)
        for pattern, alias in self.aliases:
            command = re.sub(pattern, alias, command)
        return command

    def _parse_slack_output(self, slack_rtm_output: list) -> tuple:
        """
        The Slack Real Time Messaging API is an events firehose.
        this parsing function returns None unless a message is
        directed at the Bot, based on its ID.
        """
        output_list = slack_rtm_output
        if output_list and len(output_list) > 0:
            for output in output_list:

                event_type = output.get('type', '')
                if event_type != 'message':
                    continue

                message = output.get('text', '')
                if not self._do_read_message(message):
                    continue

                # return text after the @ mention, whitespace removed
                message = message.strip().lower()
                channel = output['channel']
                self.logger.debug(f'{message} {channel}')
                yield message, channel


# noinspection PyUnusedLocal
class RecBot(SlackBot):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self._max_output_chars = 2000
        self.rules.extend([
            (r'^!run', self.run),
        ])
        self.aliases.extend([
            (r'\bll\b', 'ls -lah'),
            (r'show my log', 'cat /tmp/slackbot.log'),
        ])
     
    @command_handler
    def run(self, command: str, channel: str, regexp: str) -> str:
        try:
            output = subprocess.check_output(command, stderr=subprocess.STDOUT, shell=True)
            output = output.decode('utf8').strip()

            if len(output) > self._max_output_chars:
                output = output[:self._max_output_chars]
                output += '\n...'

            message = f'result: \n```{output}```'

        except subprocess.CalledProcessError as e:
            output = e.output.decode('utf8').strip()
            message = f'Failed to run command. \n ```{output}```'

        return message


if __name__ == "__main__":
    bot = RecBot()
    bot.listen()