class ChatBot(object):
"""
Generic base class for chatbots. Subclasses must provide
a reply method and a messages property generator
"""
def __init__(self, redis_host, redis_port):
self.db = JobsDB(redis_host, redis_port)
self.builtins = { 'help': self._help,
'jobs': self._jobs,
'logs': self._logs,
'cancel': self._cancel,
'confirm': self._confirm,
'workers': self._workers }
self.message_queue = []
self.executor = ThreadPoolExecutor(max_workers=10)
self.executor.submit(self._message_worker)
@abc.abstractmethod
def reply(self, text, channel):
raise NotImplementedError
@abc.abstractproperty
def messages(self):
"""
Generator yielding message tuples
in the form (text,user,channel)
"""
raise NotImplementedError
def _message_worker(self):
for msg in self.messages:
try:
self._process_msg(*msg)
except Exception as e:
log.error(e)
def _process_msg(self, text, user, channel):
"""
"""
command, args = self._parse_command(text)
if command in self.builtins:
self.reply(self.builtins[command](args), channel)
self.reply('EOF', channel)
else:
ok, result = self.db.create_job(command, args=args, initiator=user)
if not ok:
self.reply(result, channel)
self.reply('EOF', channel)
return
job_id = result
log.info('Created job %s' % job_id)
job = self.db.get_job(job_id)
if job['status'] == 'pending':
self.reply('%s needs confirmation' % str(job_id), channel)
self.reply('EOF', channel)
if job['chatbot_stream'] != 'False':
self.executor.submit(self._output_handler, job_id, channel)
@staticmethod
def _parse_command(text):
"""
Parse message text; return command and arguments
"""
words = text.split(' ')
cmd = words.pop(0)
args = ' '.join(words)
return cmd, args
def _output_handler(self, job_id, channel, stream=True):
"""
Worker to send the output of a given job_id to a given channel
params:
- stream(bool): Toggle streaming output as it comes in
vs posting when a job finishes. Default False.
"""
active = False
prefix = '[%s]' % job_id[-8:]
log.debug('output handler spawned for job %s' % job_id)
# sleep on jobs awaiting confirmation
while not active:
job = self.db.get_job(job_id)
if job['status'] == 'canceled':
return
if job['status'] != 'pending':
active = True
else:
sleep(1)
if stream:
for line in self.db.get_log(job_id, timestamp=False):
self.reply(prefix + line, channel)
else:
msg = ''
for line in self.db.get_log(job_id, timestamp=False):
msg += prefix + line + '\n'
self.reply(msg, channel)
self.reply('EOF', channel)
######
# Builtin command methods
######
def _confirm(self, arg):
""" Confirm a pending job """
job = self.db.get_job(arg)
if not job:
return 'no such job id'
if job['status'] != 'pending':
return 'job not awaiting confirm'
self.db.update_job(arg, 'status', 'ready')
def _cancel(self, arg):
""" Cancel a pending job """
job = self.db.get_job(arg)
if not job:
return 'no such job id'
ok, result = self.db.cancel_job(arg)
if not ok:
return result
return 'job %s canceled' % job['id']
def _workers(self, arg):
""" Show registered job workers """
workers = self.db.get_workers()
if not workers:
return 'no registered workers'
else:
return [('%s(%s)' % (w['name'], w['host'])) for w in workers]
def _jobs(self, arg):
""" Show jobs by currrent state """
subcommands = ['pending', 'running', 'completed', 'all']
if arg not in subcommands:
return 'argument must be one of %s' % ','.join(subcommands)
jobs = self.db.get_jobs(status=arg)
if not jobs:
return 'no matching jobs found'
formatted = []
for j in jobs:
created = format_time(j['created'])
formatted.append('%s %s(%s) %s' %
(created, j['id'], j['name'], j['status']))
return formatted
def _logs(self, args):
""" Show logs for a given job id """
if not self.db.get_job(args):
return 'no matching jobs found'
return self.db.get_stored_log(args)
def _help(self, args):
""" Show this help dialog """
builtin_cmds = ['Builtin commands:']
for cmd, func in sorted(self.builtins.items()):
builtin_cmds.append(' [ %s ]%s' % (cmd, func.__doc__))
action_cmds = ['Action commands:']
for cmd in sorted([ a['name'] for a in self.db.get_actions() ]):
action_cmds.append(' [ %s ]' % cmd)
return builtin_cmds + action_cmds
replace('\n', '').replace(' ', '')
class JobWorker(object):
"""
Multivac worker process. Spawns jobs, streams job stdout/stderr,
and creates actions and groups in redis from config file.
"""
def __init__(self, redis_host, redis_port, config_path):
self.pids = {} # dict of job_id:subprocess object
self.db = JobsDB(redis_host, redis_port)
self.config_path = config_path
self.read_config(self.config_path)
self.name = self._get_name()
self.executor = ThreadPoolExecutor(max_workers=10)
self.run()
def run(self):
print('Starting Multivac Job Worker %s' % self.name)
while True:
self.db.register_worker(self.name, socket.getfqdn())
# spawn ready jobs
for job in self.db.get_jobs(status='ready'):
self.executor.submit(self._job_worker, job)
# collect ended processes
pids = deepcopy(self.pids)
for job_id, pid in pids.items():
if not self._is_running(pid):
self.db.cleanup_job(job_id)
del self.pids[job_id]
print('completed job %s' % job['id'])
# re-read config if modified
if os.stat(self.config_path).st_mtime != self.config_mtime:
log.warn('re-reading modified config %s' % self.config_path)
self.read_config(self.config_path)
# cancel pending jobs exceeding timeout
now = time()
for job in self.db.get_jobs(status='pending'):
if (now - int(job['created'])) > pending_job_timeout:
print('canceling unconfirmed job %s' % job['id'])
self.db.cancel_job(job['id'])
sleep(1)
def read_config(self, path):
with open(path, 'r') as of:
config = yaml.load(of.read())
self.config_mtime = os.stat(path).st_mtime
if 'groups' in config:
self._read_groups(config['groups'])
self._read_actions(config['actions'])
def _read_groups(self, groups):
self.db.purge_groups()
for group,members in groups.items():
self.db.add_group(group, members)
log.info('loaded group %s' % (group))
def _read_actions(self, actions):
self.db.purge_actions()
for a in actions:
action = { 'confirm_required': False,
'chatbot_stream': True,
'allow_groups': 'all' }
action.update(a)
if isinstance(action['allow_groups'], list):
action['allow_groups'] = ','.join(action['allow_groups'])
self.db.add_action(action)
log.info('loaded action %s' % (action['name']))
def _get_name(self):
"""
Randomly generate a unique name for this worker
"""
name = names.get_first_name()
if name in self.db.get_workers():
self._get_name()
else:
return name
def _is_running(self, pid):
try:
os.kill(pid, 0)
except OSError:
return False
return True
def _job_worker(self, job):
print('running job %s' % job['id'])
self.db.update_job(job['id'], 'status', 'running')
if job['args']:
cmdline = shlex.split(job['cmd'] + ' ' + job['args'])
else:
cmdline = job['cmd']
proc = subprocess.Popen(
cmdline,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
self.pids[job['id']] = proc.pid
self.executor.submit(self._log_worker,
job['id'],
proc.stdout,
proc.stderr)
proc.wait()
def _log_worker(self, job_id, stdout, stderr):
log.debug('Log handler started for job %s' % job_id)
while True:
output = self._read(stdout)
error = self._read(stderr)
if output:
output = self._sanitize(output)
self.db.append_job_log(job_id, output)
log.debug('%s-STDOUT: %s' % (job_id, output))
if error:
error = self._sanitize(error)
self.db.append_job_log(job_id, error)
log.debug('%s-STDOUT: %s' % (job_id, error))
# exit when job has been collected
if job_id not in self.pids:
log.debug('Log handler stopped for job %s' % job_id)
return
def _sanitize(self, line):
line = line.decode('utf-8')
line = line.replace('\n', '')
return line
def _read(self, pipe):
"""
Non-blocking method for reading fd
"""
fd = pipe.fileno()
fl = fcntl.fcntl(fd, fcntl.F_GETFL)
fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
try:
return pipe.read()
except:
return ""