Skip to content

Instantly share code, notes, and snippets.

@dmitryhd
Last active November 7, 2019 20:11
Show Gist options
  • Save dmitryhd/704a8f463ff68ff55831ea6d3fb94b36 to your computer and use it in GitHub Desktop.
Save dmitryhd/704a8f463ff68ff55831ea6d3fb94b36 to your computer and use it in GitHub Desktop.
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment