hptoad/hptoad.py

371 rivejä
13 KiB
Python
Executable File

#!/usr/bin/env python3
import asyncio
import functools
import importlib.util
import logging
import os
import re
import signal
import sys
import types
import slixmpp
opts = {
"muc": "room@conference.example.com",
"nick": "botname",
"jid": "botname@example.com",
"resource": "resource",
"password": "password",
"connect": "xmpp.example.org:5222",
}
class HptoadPlugin:
name = None
def __init__(self, name):
spec = importlib.util.find_spec("plugins.%s" % name)
module = types.ModuleType(spec.name)
spec.loader.exec_module(module)
self.name = "plugins.%s" % name
self._obj = module.Plugin()
# Calls a function in the plugin, returns a dict with three variables.
# If the called function returned something else, all extras will be
# discarded.
# "handled" must be set to True if the plugin did something with the given
# information, and if the call did nothing, then False.
# "reply" is a string with the answer and "error" is a string with an error
# information or an empty string.
@asyncio.coroutine
def _call(self, cb_name, *args, **kwargs):
ret = {}
try:
if hasattr(self._obj, cb_name):
func = getattr(self._obj, cb_name)
if asyncio.iscoroutinefunction(func):
ret = yield from func(*args, **kwargs)
else:
ret = func(*args, **kwargs)
if not ret or type(ret) != dict:
ret = {}
except Exception as e:
ret = {"error": "%s: %s" % (type(e).__name__, str(e))}
return {"handled": bool(ret.get("handled", False)),
"reply": str(ret.get("reply", "")),
"error": str(ret.get("error", ""))}
@asyncio.coroutine
def call_initiate(self):
result = yield from self._call("initiate")
return result
@asyncio.coroutine
def call_question(self, body, nick, from_id, is_admin):
result = yield from self._call("question", body, nick,
from_id, is_admin)
return result
@asyncio.coroutine
def call_command(self, command, body, nick, from_id, is_admin):
result = yield from self._call("command", command, body, nick,
from_id, is_admin)
return result
@asyncio.coroutine
def call_chat_message(self, body, nick, from_id, is_admin):
result = yield from self._call("chat_message", body, nick,
from_id, is_admin)
return result
class Hptoad:
plugins = {}
def __init__(self, opts, timeout=5.0):
self.client = slixmpp.ClientXMPP("%s/%s" % (opts["jid"],
opts["resource"]),
opts["password"])
self.client.register_plugin("xep_0199") # XMPP Ping.
self.client.register_plugin("xep_0045") # XMPP MUC.
self.muc_obj = self.client.plugin["xep_0045"]
self.muc_is_joined = False
self.timeout = timeout
self.logger = logging.getLogger(self.__class__.__name__)
self.logger.addHandler(logging.NullHandler())
self.logger.setLevel(logging.DEBUG)
self.jid = opts["jid"]
self.connect_host = opts["connect"]
self.muc = opts["muc"]
self.pure_bot_nick = opts["nick"]
self.bot_nick = self.pure_bot_nick
def register_handlers(self):
self.client.add_event_handler("failed_all_auth",
self.on_failed_all_auth)
self.client.add_event_handler("session_start", self.on_session_start)
self.client.add_event_handler("got_online", self.on_got_online)
self.client.add_event_handler("disconnected", self.on_disconnected)
self.client.add_event_handler("message", self.on_message)
self.client.add_event_handler("muc::%s::presence" % self.muc,
self.on_muc_presence)
def connect(self):
# Reset the nick.
self.bot_nick = self.pure_bot_nick
if self.connect_host:
connect = self.connect_host.split(":", 1)
if len(connect) != 2 or not connect[1].isdigit():
self.logger.critical("Conn: Connection server format is " +
"invalid, should be example.org:5222")
sys.exit(1)
else:
connect = ()
self.client.connect(connect)
@asyncio.coroutine
def join_muc_loop(self):
while not self.muc_is_joined:
self.muc_obj.join_muc(self.muc, self.bot_nick)
yield from asyncio.sleep(self.timeout)
def join_muc(self):
asyncio.async(self.join_muc_loop())
def log_exception(self, ex):
self.logger.error("%s: %s" % (type(ex).__name__, str(ex)))
def log_message_event(self, event):
self.logger.debug("&{{jabber:client message} %s %s %s %s %s { }}" %
(event["from"], event["id"], event["to"],
event["type"], event["body"]))
def is_muc_admin(self, muc, nick):
if nick not in self.muc_obj.rooms[self.muc]:
return False
affiliation = self.muc_obj.get_jid_property(muc, nick, "affiliation")
return True if affiliation in ("admin", "owner") else False
@asyncio.coroutine
def handle_command(self, command, body, nick, from_id, is_admin):
if command == "megakick": # Megakick.
reply = None
victim = body
if victim:
is_bot_admin = self.is_muc_admin(self.muc, self.bot_nick)
is_victim_admin = self.is_muc_admin(self.muc, victim)
if is_admin and victim != self.bot_nick:
if is_bot_admin and not is_victim_admin and \
victim in self.muc_obj.rooms[self.muc]:
self.muc_obj.set_role(self.muc, victim, "none")
else:
reply = "%s: Can't megakick %s." % (nick, victim)
else:
reply = "%s: GTFO" % nick
else:
reply = "%s: WAT" % nick
if reply:
self.client.send_message(mto=self.muc, mbody=reply,
mtype="groupchat")
else: # Any plugin command.
futures = []
for plugin in self.plugins.values():
future = asyncio.async(plugin.call_command(command, body, nick,
from_id, is_admin))
callback = functools.partial(self.on_plugin_got_result,
nick=nick, from_id=from_id,
is_admin=is_admin)
future.add_done_callback(callback)
futures.append(future)
if futures:
results = yield from asyncio.gather(*futures)
if not [i["handled"] for i in results if i["handled"]]:
self.client.send_message(mto=self.muc,
mbody="%s: WAT" % nick,
mtype="groupchat")
@asyncio.coroutine
def handle_self_message(self, body, nick, from_id):
if body.startswith("!"):
split = body.split(" ", 1)
command = split[0].strip()[1:]
message = split[1] if len(split) > 1 else ""
yield from self.handle_command(command, message, nick,
from_id, True)
elif body and len(body) > 0:
self.client.send_message(mto=self.muc, mbody=body.strip(),
mtype="groupchat")
@asyncio.coroutine
def handle_muc_message(self, body, nick, from_id):
is_admin = self.is_muc_admin(self.muc, nick)
futures = []
# Has to be redone with the current bot nick.
call_regexp = re.compile("^%s[:,]" % re.escape(self.bot_nick))
for plugin in self.plugins.values():
future = asyncio.async(plugin.call_chat_message(body, nick,
from_id, is_admin))
future.add_done_callback(self.on_plugin_got_result)
futures.append(future)
if body.startswith("!"): # Any plugin command.
split = body.split(" ", 1)
command = split[0].strip()[1:]
message = split[1] if len(split) > 1 else ""
yield from self.handle_command(command, message, nick, from_id,
is_admin=is_admin)
elif call_regexp.match(body): # Chat.
message = call_regexp.sub("", body).lstrip()
for plugin in self.plugins.values():
future = asyncio.async(plugin.call_question(message, nick,
from_id, is_admin))
callback = functools.partial(self.on_plugin_got_result,
nick=nick, from_id=from_id,
is_admin=is_admin)
future.add_done_callback(callback)
futures.append(future)
if futures:
yield from asyncio.gather(*futures)
def on_failed_all_auth(self, event):
self.logger.critical("Auth: Could not connect to the server, or " +
"password mismatch!")
sys.exit(1)
def on_session_start(self, event):
self.client.get_roster()
self.client.send_presence(pstatus="is there some food in this world?",
ppriority=12)
def on_got_online(self, event):
self.join_muc()
@asyncio.coroutine
def on_disconnected(self, event):
self.muc_is_joined = False
self.logger.error("Conn: Connection lost, reattempting in %d seconds" %
self.timeout)
yield from asyncio.sleep(self.timeout)
self.connect()
@asyncio.coroutine
def on_message(self, event):
try:
if not event["type"] in ("chat", "normal", "groupchat"):
return
self.log_message_event(event)
body = event["body"]
from_id = event["from"]
if event["type"] == "groupchat":
nick = event["mucnick"]
if nick != self.bot_nick:
yield from self.handle_muc_message(body, nick, from_id)
elif event["from"].bare == self.jid:
# Use resource as a nickname with self messages.
nick = from_id.resource
yield from self.handle_self_message(body, nick, from_id)
except Exception as e:
self.log_exception(e)
@asyncio.coroutine
def on_muc_presence(self, event):
try:
typ = event["muc"]["type"]
from_id = event["from"]
nick = event["muc"]["nick"]
if not typ:
typ = event["type"]
if not nick:
nick = self.muc_obj.get_nick(self.muc, from_id)
if typ == "available":
self.muc_is_joined = True
elif typ == "error":
self.muc_is_joined = False
if event["error"]["code"] == "409":
self.bot_nick = self.bot_nick + "_"
self.join_muc()
elif typ == "unavailable":
if nick == self.bot_nick:
self.muc_is_joined = False
self.bot_nick = self.pure_bot_nick
yield from asyncio.sleep(0.5)
self.join_muc()
except Exception as e:
self.log_exception(e)
def on_plugin_got_result(self, future, nick="", from_id="", is_admin=False):
result = future.result()
if not result:
return
if result["handled"] and result["reply"]:
self.client.send_message(mto=self.muc, mbody=result["reply"],
mtype="groupchat")
if result["error"]:
self.logger.error(result["error"])
if is_admin and from_id:
self.client.send_message(mto=from_id, mbody=result["error"],
mtype="chat")
def import_plugins(self):
plugins = {}
_, _, filenames = next(os.walk("./plugins"), (None, None, []))
for filename in (i for i in filenames if i.endswith(".py")):
try:
plugin = HptoadPlugin(filename[:-3])
plugins[filename[:-3]] = plugin
future = asyncio.async(plugin.call_initiate())
future.add_done_callback(self.on_plugin_got_result)
except Exception as e:
self.log_exception(e)
return plugins
def run(self):
self.plugins = self.import_plugins()
self.register_handlers()
self.connect()
self.client.process(forever=True)
if __name__ == "__main__":
signal.signal(signal.SIGINT, signal.SIG_DFL)
logging.basicConfig(format="%(asctime)s %(message)s",
datefmt="%Y/%m/%d %H:%M:%S")
if os.path.isfile(sys.argv[0]) and os.path.dirname(sys.argv[0]):
os.chdir(os.path.dirname(sys.argv[0]))
hptoad = Hptoad(opts)
while True:
hptoad.run()
logging.error("Unknown: WTF am I doing here?")