refactored main into start_rtmbot.py, refactored so that RtmBot class is more self sufficient and can be imported and used more easily, removed unused objects from main

This commit is contained in:
Thomas Zakrajsek 2015-08-02 18:29:11 -04:00
parent c6a8b30d2d
commit 45e9f7780c
3 changed files with 84 additions and 65 deletions

3
rtmbot/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from core import *
site_config = {}

View file

@ -14,21 +14,42 @@ from argparse import ArgumentParser
from slackclient import SlackClient
def dbg(debug_string):
if debug:
logging.info(debug_string)
class RtmBot(object):
def __init__(self, token):
def __init__(self, config):
#set the config object
self.config = config
global site_config
site_config = self.config
#set slack token
self.token = config.get('SLACK_TOKEN')
#set working directory for loading plugins or other files
working_directory = os.path.dirname(sys.argv[0])
self.directory = self.config.get('BASE_PATH', working_directory)
if not self.directory.startswith('/'):
path = '{}/{}'.format(os.getcwd(), self.directory)
self.directory = os.path.abspath(path)
#establish logging
log_file = config.get('LOGFILE', 'rtmbot.log')
logging.basicConfig(filename=log_file,
level=logging.INFO,
format='%(asctime)s %(message)s')
logging.info(self.directory)
self.debug = self.config.has_key('DEBUG')
#initialize stateful fields
self.last_ping = 0
self.token = token
self.bot_plugins = []
self.slack_client = None
def _dbg(self, debug_string):
if self.debug:
logging.info(debug_string)
def connect(self):
"""Convenience method that creates Server instance"""
self.slack_client = SlackClient(self.token)
self.slack_client.rtm_connect()
def start(self):
def _start(self):
self.connect()
self.load_plugins()
while True:
@ -38,19 +59,30 @@ class RtmBot(object):
self.output()
self.autoping()
time.sleep(.1)
def start(self):
if self.config.has_key('DAEMON'):
if self.config.get('DAEMON'):
import daemon
with daemon.DaemonContext():
self._start()
self._start()
def autoping(self):
#hardcode the interval to 3 seconds
now = int(time.time())
if now > self.last_ping + 3:
self.slack_client.server.ping()
self.last_ping = now
def input(self, data):
if "type" in data:
function_name = "process_" + data["type"]
dbg("got {}".format(function_name))
self._dbg("got {}".format(function_name))
for plugin in self.bot_plugins:
plugin.register_jobs()
plugin.do(function_name, data)
def output(self):
for plugin in self.bot_plugins:
limiter = False
@ -63,18 +95,23 @@ class RtmBot(object):
message = output[1].encode('ascii','ignore')
channel.send_message("{}".format(message))
limiter = True
def crons(self):
for plugin in self.bot_plugins:
plugin.do_jobs()
def load_plugins(self):
for plugin in glob.glob(directory+'/plugins/*'):
for plugin in glob.glob(self.directory+'/plugins/*'):
sys.path.insert(0, plugin)
sys.path.insert(0, directory+'/plugins/')
for plugin in glob.glob(directory+'/plugins/*.py') + glob.glob(directory+'/plugins/*/*.py'):
sys.path.insert(0, self.directory+'/plugins/')
for plugin in glob.glob(self.directory+'/plugins/*.py') + glob.glob(self.directory+'/plugins/*/*.py'):
logging.info(plugin)
name = plugin.split('/')[-1][:-3]
# try:
self.bot_plugins.append(Plugin(name))
if name in self.config:
logging.info("config found for: " + name)
plugin_config = self.config.get(name)
self.bot_plugins.append(Plugin(name, plugin_config))
# except:
# print "error loading plugin %s" % name
@ -83,13 +120,12 @@ class Plugin(object):
self.name = name
self.jobs = []
self.module = __import__(name)
self.module.config = plugin_config
self.register_jobs()
self.outputs = []
if name in config:
logging.info("config found for: " + name)
self.module.config = config[name]
if 'setup' in dir(self.module):
self.module.setup()
def register_jobs(self):
if 'crontable' in dir(self.module):
for interval, function in self.module.crontable:
@ -98,6 +134,7 @@ class Plugin(object):
self.module.crontable = []
else:
self.module.crontable = []
def do(self, function_name, data):
if function_name in dir(self.module):
#this makes the plugin fail with stack trace in debug mode
@ -105,17 +142,19 @@ class Plugin(object):
try:
eval("self.module."+function_name)(data)
except:
dbg("problem in module {} {}".format(function_name, data))
self._dbg("problem in module {} {}".format(function_name, data))
else:
eval("self.module."+function_name)(data)
if "catch_all" in dir(self.module):
try:
self.module.catch_all(data)
except:
dbg("problem in catch all")
self._dbg("problem in catch all")
def do_jobs(self):
for job in self.jobs:
job.check()
def do_output(self):
output = []
while True:
@ -134,17 +173,20 @@ class Job(object):
self.function = function
self.interval = interval
self.lastrun = 0
def __str__(self):
return "{} {} {}".format(self.function, self.interval, self.lastrun)
def __repr__(self):
return self.__str__()
def check(self):
if self.lastrun + self.interval < time.time():
if not debug:
try:
self.function()
except:
dbg("problem")
self._dbg("problem")
else:
self.function()
self.lastrun = time.time()
@ -152,50 +194,3 @@ class Job(object):
class UnknownChannel(Exception):
pass
def main_loop():
if "LOGFILE" in config:
logging.basicConfig(filename=config["LOGFILE"], level=logging.INFO, format='%(asctime)s %(message)s')
logging.info(directory)
try:
bot.start()
except KeyboardInterrupt:
sys.exit(0)
except:
logging.exception('OOPS')
def parse_args():
parser = ArgumentParser()
parser.add_argument(
'-c',
'--config',
help='Full path to config file.',
metavar='path'
)
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
directory = os.path.dirname(sys.argv[0])
if not directory.startswith('/'):
directory = os.path.abspath("{}/{}".format(os.getcwd(),
directory
))
config = yaml.load(file(args.config or 'rtmbot.conf', 'r'))
debug = config["DEBUG"]
bot = RtmBot(config["SLACK_TOKEN"])
site_plugins = []
files_currently_downloading = []
job_hash = {}
if config.has_key("DAEMON"):
if config["DAEMON"]:
import daemon
with daemon.DaemonContext():
main_loop()
main_loop()

21
start_rtmbot.py Executable file
View file

@ -0,0 +1,21 @@
#!/usr/bin/env python
from argparse import ArgumentParser
import yaml
from rtmbot import RtmBot
def parse_args():
parser = ArgumentParser()
parser.add_argument(
'-c',
'--config',
help='Full path to config file.',
metavar='path'
)
return parser.parse_args()
# load args with config path
args = parse_args()
config = yaml.load(file(args.config or 'rtmbot.conf', 'r'))
bot = RtmBot(config)
bot.start()