From c4bfcf3e1bb0674af8d616d1a9c0598aa9c5ad3e Mon Sep 17 00:00:00 2001 From: "Brian S. Stephan" Date: Tue, 12 May 2015 20:45:18 -0500 Subject: [PATCH] add BotAdmin, IrcPlugin models a lot of stuff in here around support for loading plugins from arbitrary files. plugins have a basic amount of initialization and then hook into the core IRC event system it makes sense to have modules respond to regexes, so there's some handler stuff for that --- it was the most popular way to do stuff in the old version of the bot we need to check that people trying to load plugins are admins, so there's some stuff for that, too the expectation is that many features from here are happen in plugins, rather than modifying the core bot --- dr_botzo/ircbot/admin.py | 4 +- dr_botzo/ircbot/bot.py | 296 +++++++++++++++++- dr_botzo/ircbot/lib.py | 42 +++ .../ircbot/management/commands/runircbot.py | 2 + .../migrations/0002_botadmin_ircplugin.py | 29 ++ .../migrations/0003_auto_20150512_1934.py | 19 ++ dr_botzo/ircbot/models.py | 25 ++ 7 files changed, 407 insertions(+), 10 deletions(-) create mode 100644 dr_botzo/ircbot/lib.py create mode 100644 dr_botzo/ircbot/migrations/0002_botadmin_ircplugin.py create mode 100644 dr_botzo/ircbot/migrations/0003_auto_20150512_1934.py diff --git a/dr_botzo/ircbot/admin.py b/dr_botzo/ircbot/admin.py index bce6aaf..4738d93 100644 --- a/dr_botzo/ircbot/admin.py +++ b/dr_botzo/ircbot/admin.py @@ -2,7 +2,9 @@ from django.contrib import admin -from ircbot.models import IrcChannel +from ircbot.models import BotAdmin, IrcChannel, IrcPlugin +admin.site.register(BotAdmin) admin.site.register(IrcChannel) +admin.site.register(IrcPlugin) diff --git a/dr_botzo/ircbot/bot.py b/dr_botzo/ircbot/bot.py index b1920d1..1b1a687 100644 --- a/dr_botzo/ircbot/bot.py +++ b/dr_botzo/ircbot/bot.py @@ -1,6 +1,9 @@ """Provide the base IRC client bot which other code can latch onto.""" +import bisect +import collections import logging +import re import ssl import sys @@ -11,19 +14,128 @@ from irc.connection import Factory from irc.dict import IRCDict import irc.modes -from ircbot.models import IrcChannel +import ircbot.lib as ircbotlib +from ircbot.models import IrcChannel, IrcPlugin log = logging.getLogger('ircbot.bot') +class PrioritizedRegexHandler(collections.namedtuple('Base', ('priority', 'regex', 'callback'))): + def __lt__(self, other): + "when sorting prioritized handlers, only use the priority" + return self.priority < other.priority + + +class DrReactor(irc.client.Reactor): + """Customize the basic IRC library's Reactor with more features.""" + + def __do_nothing(*args, **kwargs): + pass + + def __init__(self, on_connect=__do_nothing, on_disconnect=__do_nothing, on_schedule=__do_nothing): + """Initialize our custom stuff.""" + + super(DrReactor, self).__init__(on_connect=on_connect, on_disconnect=on_disconnect, + on_schedule=on_schedule) + self.regex_handlers = {} + + def add_global_regex_handler(self, event, regex, handler, priority=0): + """Adds a global handler function for a specific event type and regex. + + Arguments: + + event --- Event type (a string). + + handler -- Callback function taking connection and event + parameters. + + priority --- A number (the lower the number, the higher priority). + + The handler function is called whenever the specified event is + triggered in any of the connections and the regex matches. + See documentation for the Event class. + + The handler functions are called in priority order (lowest + number is highest priority). If a handler function returns + "NO MORE", no more handlers will be called. + + This is basically an extension of add_global_handler(), either may + work, though it turns out most modules probably want this one. + """ + + handler = PrioritizedRegexHandler(priority, regex, handler) + with self.mutex: + log.debug(u"in add_global_regex_handler") + event_regex_handlers = self.regex_handlers.setdefault(event, []) + bisect.insort(event_regex_handlers, handler) + + def remove_global_regex_handler(self, event, handler): + """Removes a global regex handler function. + + Arguments: + + event -- Event type (a string). + handler -- Callback function. + + Returns 1 on success, otherwise 0. + """ + + with self.mutex: + if not event in self.regex_handlers: + return 0 + for h in self.regex_handlers[event]: + if handler == h.callback: + self.regex_handlers[event].remove(h) + return 1 + + def _handle_event(self, connection, event): + """Handle an Event event incoming on ServerConnection connection. + + Also supports regex handlers. + """ + + log.debug(u"in DrReactor._handle_event") + + with self.mutex: + # doing regex version first as it has the potential to be more specific + log.debug(u"checking regex handlers for %s", event.type) + matching_handlers = sorted( + self.regex_handlers.get("all_events", []) + + self.regex_handlers.get(event.type, []) + ) + log.debug(u"got %d", len(matching_handlers)) + for handler in matching_handlers: + log.debug(u"checking %s vs. %s", handler, event.arguments) + for line in event.arguments: + match = re.search(handler.regex, line) + if match: + log.debug(u"match!") + result = handler.callback(connection, event, match) + if result == "NO MORE": + return + + matching_handlers = sorted( + self.handlers.get("all_events", []) + + self.handlers.get(event.type, []) + ) + for handler in matching_handlers: + result = handler.callback(connection, event) + if result == "NO MORE": + return + + + class IRCBot(irc.client.SimpleIRCClient): """A single-server IRC bot class.""" + reactor_class = DrReactor + def __init__(self, reconnection_interval=60): super(IRCBot, self).__init__() self.channels = IRCDict() + self.plugins = [] # set up the server list self.server_list = settings.IRCBOT_SERVER_LIST @@ -38,8 +150,16 @@ class IRCBot(irc.client.SimpleIRCClient): self._realname = settings.IRCBOT_REALNAME # handlers - for i in ['disconnect', 'join', 'kick', 'mode', 'namreply', 'nick', 'part', 'quit', 'endofmotd']: + for i in ['disconnect', 'join', 'kick', 'mode', 'namreply', 'nick', 'part', 'quit', 'welcome']: self.connection.add_global_handler(i, getattr(self, '_on_' + i), -20) + self.connection.reactor.add_global_regex_handler('pubmsg', r'^!load\s+([\S]+)$', + getattr(self, 'handle_load'), -20) + self.connection.reactor.add_global_regex_handler('privmsg', r'^!load\s+([\S]+)$', + getattr(self, 'handle_load'), -20) + self.connection.reactor.add_global_regex_handler('pubmsg', r'^!unload\s+([\S]+)$', + getattr(self, 'handle_unload'), -20) + self.connection.reactor.add_global_regex_handler('privmsg', r'^!unload\s+([\S]+)$', + getattr(self, 'handle_unload'), -20) def _connected_checker(self): if not self.connection.is_connected(): @@ -97,13 +217,6 @@ class IRCBot(irc.client.SimpleIRCClient): # Mode on self... XXX pass - def _on_endofmotd(self, c, e): - """Join autojoin channels when the MOTD is over.""" - - for chan in IrcChannel.objects.filter(autojoin=True): - log.info(u"autojoining %s", chan.name) - self.connection.join(chan) - def _on_namreply(self, c, e): """Get the list of names in a channel. @@ -155,6 +268,34 @@ class IRCBot(irc.client.SimpleIRCClient): if ch.has_user(nick): ch.remove_user(nick) + def _on_welcome(self, connection, event): + """Initialize/run a bunch of on-connection stuff. + + Set the nickmask that the ircd tells us is us, autojoin channels, etc. + + Args: + connection source connection + event incoming event + + """ + + what = event.arguments[0] + + log.debug("welcome: %s", what) + + for chan in IrcChannel.objects.filter(autojoin=True): + log.info(u"autojoining %s", chan.name) + self.connection.join(chan) + + for plugin in IrcPlugin.objects.filter(autoload=True): + log.info(u"autoloading %s", plugin.path) + self._load_plugin(connection, event, plugin.path, feedback=False) + + match = re.search(r'(\S+!\S+@\S+)', what) + if match: + self.nickmask = match.group(1) + log.debug("setting nickmask: %s", self.nickmask) + def die(self, msg="Bye, cruel world!"): """Let the bot die. @@ -216,12 +357,149 @@ class IRCBot(irc.client.SimpleIRCClient): def on_dccchat(self, c, e): pass + def handle_load(self, connection, event, match): + """Handle IRC requests to load a plugin.""" + + log.debug(u"is admin?: %s", str(ircbotlib.is_admin(event.source))) + if ircbotlib.is_admin(event.source): + plugin_path = match.group(1) + log.debug(u"calling _load_plugin on %s", plugin_path) + self._load_plugin(connection, event, plugin_path) + + def _load_plugin(self, connection, event, plugin_path, feedback=True): + """Load an IRC plugin. + + The general assumption here is that a plugin's init loads its hooks and handlers. + """ + + log.debug(u"trying to load %s", plugin_path) + + dest = None + if feedback: + if irc.client.is_channel(event.target): + dest = event.target + else: + dest = irc.client.NickMask(event.source).nick + + for path, plugin in self.plugins: + if plugin_path == path: + if feedback: + self.privmsg(dest, "Plugin '{0:s}' is already loaded.".format(plugin_path)) + return + + # not loaded, let's get to work + try: + __import__(plugin_path) + module = sys.modules[plugin_path] + plugin = module.plugin(self, connection, event) + plugin.start() + self.plugins.append((plugin_path, plugin)) + + # give it a model + plugin_model, c = IrcPlugin.objects.get_or_create(path=plugin_path) + + if feedback: + self.privmsg(dest, "Plugin '{0:s}' loaded.".format(plugin_path)) + except ImportError as e: + log.error("Error loading '{0:s}'".format(plugin_path)) + log.exception(e) + if feedback: + self.privmsg(dest, "Plugin '{0:s}' could not be loaded.".format(plugin_path)) + + def handle_unload(self, connection, event, match): + """Handle IRC requests to unload a plugin.""" + + log.debug(u"is admin?: %s", str(ircbotlib.is_admin(event.source))) + if ircbotlib.is_admin(event.source): + plugin_path = match.group(1) + log.debug(u"calling _unload_plugin on %s", plugin_path) + self._unload_plugin(connection, event, plugin_path) + + def _unload_plugin(self, connection, event, plugin_path): + """Attempt to unload and del a module if it's loaded.""" + + log.debug(u"trying to unload %s", plugin_path) + + if irc.client.is_channel(event.target): + dest = event.target + else: + dest = irc.client.NickMask(event.source).nick + + for path, plugin in self.plugins: + if plugin_path == path: + self.plugins.remove((path, plugin)) + plugin.stop() + del plugin + del sys.modules[plugin_path] + + self.privmsg(dest, "Plugin '{0:s}' unloaded.".format(plugin_path)) + return + + # guess it was never loaded + self.privmsg(dest, "Plugin '{0:s}' is not loaded.".format(plugin_path)) + + def privmsg(self, target, text): + """Send a PRIVMSG command. + + Args: + target the destination nick/channel + text the message to send + + """ + + if not target: + return + + log.debug("OUTGOING PRIVMSG: t[%s] m[%s]", target, text) + + splitter = "..." + + # split messages that are too long. Max length is 512. + # TODO: this does not properly handle when the nickmask has been + # masked by the ircd + # is the above still the case? + space = 512 - len('\r\n') - len(' PRIVMSG ') - len(target) - len(' :') - len(self.nickmask) - len(' :') + splitspace = space - (len(splitter) + 1) + + if len(text) > space: + times = 1 + + while len(text) > splitspace: + splitpos = text.rfind(' ', 0, splitspace) + splittext = text[0:splitpos] + ' ' + splitter + text = splitter + ' ' + text[splitpos+1:] + self.connection.send_raw("PRIVMSG {0:s} :{1:s}".format(target, splittext)) + + times = times + 1 + if times >= 4: + # this is stupidly long, abort + return + + # done splitting + self.connection.send_raw("PRIVMSG {0:s} :{1:s}".format(target, text)) + else: + self.connection.send_raw("PRIVMSG {0:s} :{1:s}".format(target, text)) + def start(self): """Start the bot.""" self._connect() super(IRCBot, self).start() + def sigint_handler(self, signal, frame): + """Cleanly shutdown on SIGINT.""" + + log.debug(u"shutting down") + for path, plugin in self.plugins: + log.debug(u"trying to shut down %s", path) + self.plugins.remove((path, plugin)) + plugin.stop() + del plugin + del sys.modules[path] + + self.disconnect("Shutting down...") + sys.exit() + class Channel(object): """A class for keeping information about an IRC channel.""" diff --git a/dr_botzo/ircbot/lib.py b/dr_botzo/ircbot/lib.py new file mode 100644 index 0000000..c8f4f2a --- /dev/null +++ b/dr_botzo/ircbot/lib.py @@ -0,0 +1,42 @@ +"""Library and convenience methods for the IRC bot and plugins.""" + +import logging + +from ircbot.models import BotAdmin + + +log = logging.getLogger('ircbot.lib') + + +class Plugin(object): + + """Plugin base class.""" + + def __init__(self, bot, connection, event): + """Initialization stuff here --- global handlers, configs from database, so on.""" + + self.bot = bot + self.connection = connection + self.event = event + + log.info(u"initialized %s", self.__class__.__name__) + + def start(self): + """Initialization stuff here --- global handlers, configs from database, so on.""" + + log.info(u"started %s", self.__class__.__name__) + + def stop(self): + """Teardown stuff here --- unregister handlers, for example.""" + + log.info(u"stopped %s", self.__class__.__name__) + + +def is_admin(source): + """Check if the provided event source is a bot admin.""" + + if source in BotAdmin.objects.values_list('nickmask', flat=True): + log.debug(u"in is_admin; True") + return True + log.debug(u"in is_admin; False") + return False diff --git a/dr_botzo/ircbot/management/commands/runircbot.py b/dr_botzo/ircbot/management/commands/runircbot.py index a7a4bb7..a09b53c 100644 --- a/dr_botzo/ircbot/management/commands/runircbot.py +++ b/dr_botzo/ircbot/management/commands/runircbot.py @@ -1,6 +1,7 @@ """Start the IRC bot via Django management command.""" import logging +import signal from django.core.management import BaseCommand @@ -22,4 +23,5 @@ class Command(BaseCommand): """Start the IRC bot and spin forever.""" irc = IRCBot() + signal.signal(signal.SIGINT, irc.sigint_handler) irc.start() diff --git a/dr_botzo/ircbot/migrations/0002_botadmin_ircplugin.py b/dr_botzo/ircbot/migrations/0002_botadmin_ircplugin.py new file mode 100644 index 0000000..019b0e7 --- /dev/null +++ b/dr_botzo/ircbot/migrations/0002_botadmin_ircplugin.py @@ -0,0 +1,29 @@ +# -*- coding: utf-8 -*- +from __future__ import unicode_literals + +from django.db import models, migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('ircbot', '0001_initial'), + ] + + operations = [ + migrations.CreateModel( + name='BotAdmin', + fields=[ + ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)), + ('nickmask', models.CharField(unique=True, max_length=200)), + ], + ), + migrations.CreateModel( + name='IrcPlugin', + fields=[ + ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)), + ('path', models.CharField(unique=True, max_length=200)), + ('autojoin', models.BooleanField(default=False)), + ], + ), + ] diff --git a/dr_botzo/ircbot/migrations/0003_auto_20150512_1934.py b/dr_botzo/ircbot/migrations/0003_auto_20150512_1934.py new file mode 100644 index 0000000..bca5846 --- /dev/null +++ b/dr_botzo/ircbot/migrations/0003_auto_20150512_1934.py @@ -0,0 +1,19 @@ +# -*- coding: utf-8 -*- +from __future__ import unicode_literals + +from django.db import models, migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('ircbot', '0002_botadmin_ircplugin'), + ] + + operations = [ + migrations.RenameField( + model_name='ircplugin', + old_name='autojoin', + new_name='autoload', + ), + ] diff --git a/dr_botzo/ircbot/models.py b/dr_botzo/ircbot/models.py index 9c84bb0..46e98e6 100644 --- a/dr_botzo/ircbot/models.py +++ b/dr_botzo/ircbot/models.py @@ -8,6 +8,18 @@ from django.db import models log = logging.getLogger('ircbot.models') +class BotAdmin(models.Model): + + """Configure admins, which can do things through the bot that others can't.""" + + nickmask = models.CharField(max_length=200, unique=True) + + def __unicode__(self): + """String representation.""" + + return u"{0:s}".format(self.nickmask) + + class IrcChannel(models.Model): """Track channel settings.""" @@ -19,3 +31,16 @@ class IrcChannel(models.Model): """String representation.""" return u"{0:s}".format(self.name) + + +class IrcPlugin(models.Model): + + """Represent an IRC plugin and its loading settings.""" + + path = models.CharField(max_length=200, unique=True) + autoload = models.BooleanField(default=False) + + def __unicode__(self): + """String representation.""" + + return u"{0:s}".format(self.path)