Last active
November 7, 2019 20:11
-
-
Save dmitryhd/704a8f463ff68ff55831ea6d3fb94b36 to your computer and use it in GitHub Desktop.
slackbot.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
__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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment