XMPP bot written in Python
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

hptoad.py 13KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370
  1. #!/usr/bin/env python3
  2. import asyncio
  3. import functools
  4. import importlib.util
  5. import logging
  6. import os
  7. import re
  8. import signal
  9. import sys
  10. import types
  11. import slixmpp
  12. opts = {
  13. "muc": "room@conference.example.com",
  14. "nick": "botname",
  15. "jid": "botname@example.com",
  16. "resource": "resource",
  17. "password": "password",
  18. "connect": "xmpp.example.org:5222",
  19. }
  20. class HptoadPlugin:
  21. name = None
  22. def __init__(self, name):
  23. spec = importlib.util.find_spec("plugins.%s" % name)
  24. module = types.ModuleType(spec.name)
  25. spec.loader.exec_module(module)
  26. self.name = "plugins.%s" % name
  27. self._obj = module.Plugin()
  28. # Calls a function in the plugin, returns a dict with three variables.
  29. # If the called function returned something else, all extras will be
  30. # discarded.
  31. # "handled" must be set to True if the plugin did something with the given
  32. # information, and if the call did nothing, then False.
  33. # "reply" is a string with the answer and "error" is a string with an error
  34. # information or an empty string.
  35. @asyncio.coroutine
  36. def _call(self, cb_name, *args, **kwargs):
  37. ret = {}
  38. try:
  39. if hasattr(self._obj, cb_name):
  40. func = getattr(self._obj, cb_name)
  41. if asyncio.iscoroutinefunction(func):
  42. ret = yield from func(*args, **kwargs)
  43. else:
  44. ret = func(*args, **kwargs)
  45. if not ret or type(ret) != dict:
  46. ret = {}
  47. except Exception as e:
  48. ret = {"error": "%s: %s" % (type(e).__name__, str(e))}
  49. return {"handled": bool(ret.get("handled", False)),
  50. "reply": str(ret.get("reply", "")),
  51. "error": str(ret.get("error", ""))}
  52. @asyncio.coroutine
  53. def call_initiate(self):
  54. result = yield from self._call("initiate")
  55. return result
  56. @asyncio.coroutine
  57. def call_question(self, body, nick, from_id, is_admin):
  58. result = yield from self._call("question", body, nick,
  59. from_id, is_admin)
  60. return result
  61. @asyncio.coroutine
  62. def call_command(self, command, body, nick, from_id, is_admin):
  63. result = yield from self._call("command", command, body, nick,
  64. from_id, is_admin)
  65. return result
  66. @asyncio.coroutine
  67. def call_chat_message(self, body, nick, from_id, is_admin):
  68. result = yield from self._call("chat_message", body, nick,
  69. from_id, is_admin)
  70. return result
  71. class Hptoad:
  72. plugins = {}
  73. def __init__(self, opts, timeout=5.0):
  74. self.client = slixmpp.ClientXMPP("%s/%s" % (opts["jid"],
  75. opts["resource"]),
  76. opts["password"])
  77. self.client.register_plugin("xep_0199") # XMPP Ping.
  78. self.client.register_plugin("xep_0045") # XMPP MUC.
  79. self.muc_obj = self.client.plugin["xep_0045"]
  80. self.muc_is_joined = False
  81. self.timeout = timeout
  82. self.logger = logging.getLogger(self.__class__.__name__)
  83. self.logger.addHandler(logging.NullHandler())
  84. self.logger.setLevel(logging.DEBUG)
  85. self.jid = opts["jid"]
  86. self.connect_host = opts["connect"]
  87. self.muc = opts["muc"]
  88. self.pure_bot_nick = opts["nick"]
  89. self.bot_nick = self.pure_bot_nick
  90. def register_handlers(self):
  91. self.client.add_event_handler("failed_all_auth",
  92. self.on_failed_all_auth)
  93. self.client.add_event_handler("session_start", self.on_session_start)
  94. self.client.add_event_handler("got_online", self.on_got_online)
  95. self.client.add_event_handler("disconnected", self.on_disconnected)
  96. self.client.add_event_handler("message", self.on_message)
  97. self.client.add_event_handler("muc::%s::presence" % self.muc,
  98. self.on_muc_presence)
  99. def connect(self):
  100. # Reset the nick.
  101. self.bot_nick = self.pure_bot_nick
  102. if self.connect_host:
  103. connect = self.connect_host.split(":", 1)
  104. if len(connect) != 2 or not connect[1].isdigit():
  105. self.logger.critical("Conn: Connection server format is " +
  106. "invalid, should be example.org:5222")
  107. sys.exit(1)
  108. else:
  109. connect = ()
  110. self.client.connect(connect)
  111. @asyncio.coroutine
  112. def join_muc_loop(self):
  113. while not self.muc_is_joined:
  114. self.muc_obj.join_muc(self.muc, self.bot_nick)
  115. yield from asyncio.sleep(self.timeout)
  116. def join_muc(self):
  117. asyncio.async(self.join_muc_loop())
  118. def log_exception(self, ex):
  119. self.logger.error("%s: %s" % (type(ex).__name__, str(ex)))
  120. def log_message_event(self, event):
  121. self.logger.debug("&{{jabber:client message} %s %s %s %s %s { }}" %
  122. (event["from"], event["id"], event["to"],
  123. event["type"], event["body"]))
  124. def is_muc_admin(self, muc, nick):
  125. if nick not in self.muc_obj.rooms[self.muc]:
  126. return False
  127. affiliation = self.muc_obj.get_jid_property(muc, nick, "affiliation")
  128. return True if affiliation in ("admin", "owner") else False
  129. @asyncio.coroutine
  130. def handle_command(self, command, body, nick, from_id, is_admin):
  131. if command == "megakick": # Megakick.
  132. reply = None
  133. victim = body
  134. if victim:
  135. is_bot_admin = self.is_muc_admin(self.muc, self.bot_nick)
  136. is_victim_admin = self.is_muc_admin(self.muc, victim)
  137. if is_admin and victim != self.bot_nick:
  138. if is_bot_admin and not is_victim_admin and \
  139. victim in self.muc_obj.rooms[self.muc]:
  140. self.muc_obj.set_role(self.muc, victim, "none")
  141. else:
  142. reply = "%s: Can't megakick %s." % (nick, victim)
  143. else:
  144. reply = "%s: GTFO" % nick
  145. else:
  146. reply = "%s: WAT" % nick
  147. if reply:
  148. self.client.send_message(mto=self.muc, mbody=reply,
  149. mtype="groupchat")
  150. else: # Any plugin command.
  151. futures = []
  152. for plugin in self.plugins.values():
  153. future = asyncio.async(plugin.call_command(command, body, nick,
  154. from_id, is_admin))
  155. callback = functools.partial(self.on_plugin_got_result,
  156. nick=nick, from_id=from_id,
  157. is_admin=is_admin)
  158. future.add_done_callback(callback)
  159. futures.append(future)
  160. if futures:
  161. results = yield from asyncio.gather(*futures)
  162. if not [i["handled"] for i in results if i["handled"]]:
  163. self.client.send_message(mto=self.muc,
  164. mbody="%s: WAT" % nick,
  165. mtype="groupchat")
  166. @asyncio.coroutine
  167. def handle_self_message(self, body, nick, from_id):
  168. if body.startswith("!"):
  169. split = body.split(" ", 1)
  170. command = split[0].strip()[1:]
  171. message = split[1] if len(split) > 1 else ""
  172. yield from self.handle_command(command, message, nick,
  173. from_id, True)
  174. elif body and len(body) > 0:
  175. self.client.send_message(mto=self.muc, mbody=body.strip(),
  176. mtype="groupchat")
  177. @asyncio.coroutine
  178. def handle_muc_message(self, body, nick, from_id):
  179. is_admin = self.is_muc_admin(self.muc, nick)
  180. futures = []
  181. # Has to be redone with the current bot nick.
  182. call_regexp = re.compile("^%s[:,]" % re.escape(self.bot_nick))
  183. for plugin in self.plugins.values():
  184. future = asyncio.async(plugin.call_chat_message(body, nick,
  185. from_id, is_admin))
  186. future.add_done_callback(self.on_plugin_got_result)
  187. futures.append(future)
  188. if body.startswith("!"): # Any plugin command.
  189. split = body.split(" ", 1)
  190. command = split[0].strip()[1:]
  191. message = split[1] if len(split) > 1 else ""
  192. yield from self.handle_command(command, message, nick, from_id,
  193. is_admin=is_admin)
  194. elif call_regexp.match(body): # Chat.
  195. message = call_regexp.sub("", body).lstrip()
  196. for plugin in self.plugins.values():
  197. future = asyncio.async(plugin.call_question(message, nick,
  198. from_id, is_admin))
  199. callback = functools.partial(self.on_plugin_got_result,
  200. nick=nick, from_id=from_id,
  201. is_admin=is_admin)
  202. future.add_done_callback(callback)
  203. futures.append(future)
  204. if futures:
  205. yield from asyncio.gather(*futures)
  206. def on_failed_all_auth(self, event):
  207. self.logger.critical("Auth: Could not connect to the server, or " +
  208. "password mismatch!")
  209. sys.exit(1)
  210. def on_session_start(self, event):
  211. self.client.get_roster()
  212. self.client.send_presence(pstatus="is there some food in this world?",
  213. ppriority=12)
  214. def on_got_online(self, event):
  215. self.join_muc()
  216. @asyncio.coroutine
  217. def on_disconnected(self, event):
  218. self.muc_is_joined = False
  219. self.logger.error("Conn: Connection lost, reattempting in %d seconds" %
  220. self.timeout)
  221. yield from asyncio.sleep(self.timeout)
  222. self.connect()
  223. @asyncio.coroutine
  224. def on_message(self, event):
  225. try:
  226. if not event["type"] in ("chat", "normal", "groupchat"):
  227. return
  228. self.log_message_event(event)
  229. body = event["body"]
  230. from_id = event["from"]
  231. if event["type"] == "groupchat":
  232. nick = event["mucnick"]
  233. if nick != self.bot_nick:
  234. yield from self.handle_muc_message(body, nick, from_id)
  235. elif event["from"].bare == self.jid:
  236. # Use resource as a nickname with self messages.
  237. nick = from_id.resource
  238. yield from self.handle_self_message(body, nick, from_id)
  239. except Exception as e:
  240. self.log_exception(e)
  241. @asyncio.coroutine
  242. def on_muc_presence(self, event):
  243. try:
  244. typ = event["muc"]["type"]
  245. from_id = event["from"]
  246. nick = event["muc"]["nick"]
  247. if not typ:
  248. typ = event["type"]
  249. if not nick:
  250. nick = self.muc_obj.get_nick(self.muc, from_id)
  251. if typ == "available":
  252. self.muc_is_joined = True
  253. elif typ == "error":
  254. self.muc_is_joined = False
  255. if event["error"]["code"] == "409":
  256. self.bot_nick = self.bot_nick + "_"
  257. self.join_muc()
  258. elif typ == "unavailable":
  259. if nick == self.bot_nick:
  260. self.muc_is_joined = False
  261. self.bot_nick = self.pure_bot_nick
  262. yield from asyncio.sleep(0.5)
  263. self.join_muc()
  264. except Exception as e:
  265. self.log_exception(e)
  266. def on_plugin_got_result(self, future, nick="", from_id="", is_admin=False):
  267. result = future.result()
  268. if not result:
  269. return
  270. if result["handled"] and result["reply"]:
  271. self.client.send_message(mto=self.muc, mbody=result["reply"],
  272. mtype="groupchat")
  273. if result["error"]:
  274. self.logger.error(result["error"])
  275. if is_admin and from_id:
  276. self.client.send_message(mto=from_id, mbody=result["error"],
  277. mtype="chat")
  278. def import_plugins(self):
  279. plugins = {}
  280. _, _, filenames = next(os.walk("./plugins"), (None, None, []))
  281. for filename in (i for i in filenames if i.endswith(".py")):
  282. try:
  283. plugin = HptoadPlugin(filename[:-3])
  284. plugins[filename[:-3]] = plugin
  285. future = asyncio.async(plugin.call_initiate())
  286. future.add_done_callback(self.on_plugin_got_result)
  287. except Exception as e:
  288. self.log_exception(e)
  289. return plugins
  290. def run(self):
  291. self.plugins = self.import_plugins()
  292. self.register_handlers()
  293. self.connect()
  294. self.client.process(forever=True)
  295. if __name__ == "__main__":
  296. signal.signal(signal.SIGINT, signal.SIG_DFL)
  297. logging.basicConfig(format="%(asctime)s %(message)s",
  298. datefmt="%Y/%m/%d %H:%M:%S")
  299. if os.path.isfile(sys.argv[0]) and os.path.dirname(sys.argv[0]):
  300. os.chdir(os.path.dirname(sys.argv[0]))
  301. hptoad = Hptoad(opts)
  302. while True:
  303. hptoad.run()
  304. logging.error("Unknown: WTF am I doing here?")