12:49aec9763d9e
Anton Shestakov <av6@dwimlabs.net>, Sat, 14 May 2016 18:03:07 +0800
docs: describe basic functionality, prerequisites and config

next change 13:d4079b5ce7a0
previous change 11:1be0d90ed46b

rollbot.py

Permissions: -rw-r--r--

Other formats: Feeds:
#!/usr/bin/env python
import logging
import random
import re
import sys
from argparse import ArgumentParser, FileType, SUPPRESS
import yaml
import sleekxmpp
if sys.version_info < (3, 0):
reload(sys)
sys.setdefaultencoding('utf8')
else:
raw_input = input
class Command(object):
name = 'command'
def __init__(self, rollbot):
self.rollbot = rollbot
def usage(self):
return '(no help available, try experimenting)'
class Roll(Command):
name = 'roll'
regex = re.compile('^(?P<times>\d+)?d(?P<sides>\d+)$', re.IGNORECASE)
def usage(self):
return '%s <X>D<Y> means roll Y-sided die X times, if X is 1 it can be omitted' % self.name
def respond(self, words, message):
if len(words) == 2:
match = self.regex.match(words[1])
if match is not None:
times = match.group('times')
if times is None:
times = 1
sides = match.group('sides')
rolls, total = self.roll(int(times), int(sides))
return self.format(rolls, total)
return '%s usage: %s' % (self.name, self.usage())
def roll(self, times, sides):
results = [random.randint(1, sides) for i in range(times)]
total = sum(results)
return results, total
def format(self, rolls, total):
if len(rolls) == 1:
return '%d' % total
else:
return '(%s) = %d' % (' + '.join('%d' % r for r in rolls), total)
class Help(Command):
name = 'help'
def usage(self):
return '%s or %s <command>' % (self.name, self.name)
def respond(self, words, message):
if len(words) == 2:
cmdname = words[1]
if cmdname in self.rollbot.commands:
command = self.rollbot.commands[cmdname]
return '%s usage: %s' % (command.name, command.usage())
else:
return 'no such command: %s' % cmdname
elif len(words) == 1:
cmds = ' '.join(sorted(self.rollbot.commands.keys()))
return 'available commands: %s' % cmds
else:
return '%s usage: %s' % (self.name, self.usage())
class RollBot(sleekxmpp.ClientXMPP):
def __init__(self, jid, password, nick, prefix):
sleekxmpp.ClientXMPP.__init__(self, jid, password)
self.nick = nick
self.prefix = prefix
self.commands = {}
self.add_command(Roll)
self.add_command(Help)
self.add_event_handler('session_start', self.start)
self.add_event_handler('message', self.message)
self.add_event_handler('groupchat_message', self.muc_message)
self.add_event_handler('groupchat_invite', self.muc_invite)
self.add_event_handler('groupchat_direct_invite', self.muc_direct_invite)
def add_command(self, commandcls):
self.commands[commandcls.name] = commandcls(self)
def start(self, event):
self.get_roster()
self.send_presence()
def message(self, msg):
if msg['type'] not in ('chat', 'normal'):
return
mbody = msg['body'].lstrip()
if mbody.startswith(self.prefix):
mbody = mbody[len(self.prefix):]
mbody = mbody.lstrip(' ')
rbody = self.respond(mbody)
if rbody:
msg.reply(rbody)
msg.send()
def muc_message(self, msg):
if msg['mucnick'] == self.nick:
return
mbody = msg['body'].lstrip()
if mbody.startswith(self.nick):
mbody = mbody[len(self.nick):]
elif mbody.startswith(self.prefix):
mbody = mbody[len(self.prefix):]
else:
return
mbody = mbody.lstrip(':, ')
rbody = self.respond(mbody)
if rbody:
rbody = '%s: %s' % (msg['mucnick'], rbody)
msg.reply(rbody)
msg.send()
def muc_invite(self, invite):
room = invite['from']
password = '' # https://github.com/fritzy/SleekXMPP/issues/409
self.plugin['xep_0045'].joinMUC(room, self.nick, password=password, wait=True)
def muc_direct_invite(self, invite):
room = invite['groupchat_invite']['jid']
password = invite['groupchat_invite']['password']
self.plugin['xep_0045'].joinMUC(room, self.nick, password=password, wait=True)
def respond(self, message, prefix=None):
response = ''
words = message.split()
if words:
cmdname = words[0].lower()
if cmdname in self.commands:
command = self.commands[cmdname]
response = command.respond(words, message)
return response
def lookup(key, args, config, default=None):
if hasattr(args, key):
return getattr(args, key)
return config.get(key, default)
def main():
parser = ArgumentParser(argument_default=SUPPRESS)
parser.add_argument('-c', '--config', help='configuration file (YAML)', type=FileType('r'))
group = parser.add_argument_group('main configuration', 'options that can be specified in the configuration file')
group.add_argument('-j', '--jid', help='JID to use')
group.add_argument('-p', '--password', help='password to use')
group.add_argument('-n', '--nick', help='MUC nickname (default: rollbot)')
group.add_argument('--prefix', help='command prefix (default: !)')
group = parser.add_mutually_exclusive_group()
group.add_argument('-q', '--quiet', help='set logging to ERROR', action='store_const', dest='loglevel', const=logging.ERROR, default=logging.INFO)
group.add_argument('-d', '--debug', help='set logging to DEBUG', action='store_const', dest='loglevel', const=logging.DEBUG, default=logging.INFO)
args = parser.parse_args()
logging.basicConfig(level=args.loglevel, format='%(asctime)s %(levelname)-8s %(message)s')
if hasattr(args, 'config'):
config = yaml.safe_load(args.config)
else:
config = {}
jid = lookup('jid', args, config)
password = lookup('password', args, config)
nick = lookup('nick', args, config, default='rollbot')
prefix = lookup('prefix', args, config, default='!')
rollbot = RollBot(jid, password, nick, prefix)
rollbot.auto_authorize = True
rollbot.auto_subscribe = True
rollbot.register_plugin('xep_0030') # Service Discovery
rollbot.register_plugin('xep_0045') # Multi-User Chat
rollbot.register_plugin('xep_0092') # Software Version
rollbot.register_plugin('xep_0199') # XMPP Ping
rollbot.register_plugin('xep_0202') # Entity Time
rollbot.register_plugin('xep_0249') # Direct MUC Invitations
if rollbot.connect():
logging.info('RollBot connected')
rollbot.process(block=True)
logging.info('RollBot disconnected')
else:
logging.fatal("RollBot couldn't connect")
if __name__ == '__main__':
main()