jabberbot.py

   1 #!/usr/bin/python
   2 # -*- coding: utf-8 -*-
   3
   4 # JabberBot: A simple jabber/xmpp bot framework
   5 # Copyright (c) 2007-2011 Thomas Perl <thp.io/about>
   6 #
   7 # This program is free software; you can redistribute it and/or modify
   8 # it under the terms of the GNU General Public License as published by
   9 # the Free Software Foundation; either version 3 of the License, or
  10 # (at your option) any later version.
  11 #
  12 # This program is distributed in the hope that it will be useful,
  13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 # GNU General Public License for more details.
  16 #
  17 # You should have received a copy of the GNU General Public License
  18 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  19 #
  20
  21 """A simple framework for creating Jabber/XMPP bots and services"""
  22
  23 import os
  24 import re
  25 import sys
  26
  27 try:
  28     import xmpp
  29 except ImportError:
  30     print >> sys.stderr, 'You need to install xmpppy from http://xmpppy.sf.net/.'
  31     sys.exit(-1)
  32
  33 import time
  34 import inspect
  35 import logging
  36 import traceback
  37
  38 # Will be parsed by setup.py to determine package metadata
  39 __author__ = 'Thomas Perl <m@thp.io>'
  40 __version__ = '0.13'
  41 __website__ = 'http://thp.io/2007/python-jabberbot/'
  42 __license__ = 'GPLv3 or later'
  43
  44 def botcmd(*args, **kwargs):
  45     """Decorator for bot command functions"""
  46
  47     def decorate(func, hidden=False, name=None):
  48         setattr(func, '_jabberbot_command', True)
  49         setattr(func, '_jabberbot_hidden', hidden)
  50         setattr(func, '_jabberbot_command_name', name or func.__name__)
  51         return func
  52
  53     if len(args):
  54         return decorate(args[0], **kwargs)
  55     else:
  56         return lambda func: decorate(func, **kwargs)
  57
  58
  59 class JabberBot(object):
  60     AVAILABLE, AWAY, CHAT, DND, XA, OFFLINE = None, 'away', 'chat', 'dnd', 'xa', 'unavailable'
  61
  62     MSG_AUTHORIZE_ME = 'Hey there. You are not yet on my roster. Authorize my request and I will do the same.'
  63     MSG_NOT_AUTHORIZED = 'You did not authorize my subscription request. Access denied.'
  64
  65     PING_FREQUENCY = 0 # Set to the number of seconds, e.g. 60.
  66     PING_TIMEOUT = 2 # Seconds to wait for a response.
  67
  68     def __init__(self, username, password, res=None, debug=False,
  69             privatedomain=False, acceptownmsgs=False):
  70         """Initializes the jabber bot and sets up commands.
  71
  72         If privatedomain is provided, it should be either
  73         True to only allow subscriptions from the same domain
  74         as the bot or a string that describes the domain for
  75         which subscriptions are accepted (e.g. 'jabber.org').
  76
  77         If acceptownmsgs it set to True, this bot will accept
  78         messages from the same JID that the bot itself has. This
  79         is useful when using JabberBot with a single Jabber account
  80         and multiple instances that want to talk to each other.
  81         """
  82         self.__debug = debug
  83         self.log = logging.getLogger(__name__)
  84         self.__username = username
  85         self.__password = password
  86         self.jid = xmpp.JID(self.__username)
  87         self.res = (res or self.__class__.__name__)
  88         self.conn = None
  89         self.__finished = False
  90         self.__show = None
  91         self.__status = None
  92         self.__seen = {}
  93         self.__threads = {}
  94         self.__lastping = time.time()
  95         self.__privatedomain = privatedomain
  96         self.__acceptownmsgs = acceptownmsgs
  97
  98         self.custom_message_handler = None
  99
 100         self.commands = {}
 101         for name, value in inspect.getmembers(self):
 102             if inspect.ismethod(value) and getattr(value, '_jabberbot_command', False):
 103                 name = getattr(value, '_jabberbot_command_name')
 104                 self.log.debug('Registered command: %s' % name)
 105                 self.commands[name] = value
 106
 107         self.roster = None
 108
 109 ################################
 110
 111     def _send_status(self):
 112         self.conn.send(xmpp.dispatcher.Presence(show=self.__show, status=self.__status))
 113
 114     def __set_status(self, value):
 115         if self.__status != value:
 116             self.__status = value
 117             self._send_status()
 118
 119     def __get_status(self):
 120         return self.__status
 121
 122     status_message = property(fget=__get_status, fset=__set_status)
 123
 124     def __set_show(self, value):
 125         if self.__show != value:
 126             self.__show = value
 127             self._send_status()
 128
 129     def __get_show(self):
 130         return self.__show
 131
 132     status_type = property(fget=__get_show, fset=__set_show)
 133
 134 ################################
 135
 136     def connect(self):
 137         if not self.conn:
 138             if self.__debug:
 139                 conn = xmpp.Client(self.jid.getDomain())
 140             else:
 141                 conn = xmpp.Client(self.jid.getDomain(), debug=[])
 142
 143             conres = conn.connect()
 144             if not conres:
 145                 self.log.error('unable to connect to server %s.' % self.jid.getDomain())
 146                 return None
 147             if conres != 'tls':
 148                 self.log.warning('unable to establish secure connection - TLS failed!')
 149
 150             authres = conn.auth(self.jid.getNode(), self.__password, self.res)
 151             if not authres:
 152                 self.log.error('unable to authorize with server.')
 153                 return None
 154             if authres != 'sasl':
 155                 self.log.warning("unable to perform SASL auth os %s. Old authentication method used!" % self.jid.getDomain())
 156
 157             conn.sendInitPresence()
 158             self.conn = conn
 159             self.roster = self.conn.Roster.getRoster()
 160             self.log.info('*** roster ***')
 161             for contact in self.roster.getItems():
 162                 self.log.info('  %s' % contact)
 163             self.log.info('*** roster ***')
 164             self.conn.RegisterHandler('message', self.callback_message)
 165             self.conn.RegisterHandler('presence', self.callback_presence)
 166
 167         return self.conn
 168
 169     def join_room(self, room, username=None):
 170         """Join the specified multi-user chat room"""
 171         if username is None:
 172             username = self.__username.split('@')[0]
 173         my_room_JID = '/'.join((room, username))
 174         self.connect().send(xmpp.Presence(to=my_room_JID))
 175
 176     def quit(self):
 177         """Stop serving messages and exit.
 178
 179         I find it is handy for development to run the
 180         jabberbot in a 'while true' loop in the shell, so
 181         whenever I make a code change to the bot, I send
 182         the 'reload' command, which I have mapped to call
 183         self.quit(), and my shell script relaunches the
 184         new version.
 185         """
 186         self.__finished = True
 187
 188     def send_message(self, mess):
 189         """Send an XMPP message"""
 190         self.connect().send(mess)
 191
 192     def send_tune(self, song, debug=False):
 193         """Set information about the currently played tune
 194
 195         Song is a dictionary with keys: file, title, artist, album, pos, track,
 196         length, uri. For details see <http://xmpp.org/protocols/tune/>.
 197         """
 198         NS_TUNE = 'http://jabber.org/protocol/tune'
 199         iq = xmpp.Iq(typ='set')
 200         iq.setFrom(self.jid)
 201         iq.pubsub = iq.addChild('pubsub', namespace=xmpp.NS_PUBSUB)
 202         iq.pubsub.publish = iq.pubsub.addChild('publish', attrs={ 'node' : NS_TUNE })
 203         iq.pubsub.publish.item = iq.pubsub.publish.addChild('item', attrs={ 'id' : 'current' })
 204         tune = iq.pubsub.publish.item.addChild('tune')
 205         tune.setNamespace(NS_TUNE)
 206
 207         title = None
 208         if song.has_key('title'):
 209             title = song['title']
 210         elif song.has_key('file'):
 211             title = os.path.splitext(os.path.basename(song['file']))[0]
 212         if title is not None:
 213             tune.addChild('title').addData(title)
 214         if song.has_key('artist'):
 215             tune.addChild('artist').addData(song['artist'])
 216         if song.has_key('album'):
 217             tune.addChild('source').addData(song['album'])
 218         if song.has_key('pos') and song['pos'] > 0:
 219             tune.addChild('track').addData(str(song['pos']))
 220         if song.has_key('time'):
 221             tune.addChild('length').addData(str(song['time']))
 222         if song.has_key('uri'):
 223             tune.addChild('uri').addData(song['uri'])
 224
 225         if debug:
 226             print 'Sending tune:', iq.__str__().encode('utf8')
 227         self.conn.send(iq)
 228
 229     def send(self, user, text, in_reply_to=None, message_type='chat'):
 230         """Sends a simple message to the specified user."""
 231         mess = self.build_message(text)
 232         mess.setTo(user)
 233
 234         if in_reply_to:
 235             mess.setThread(in_reply_to.getThread())
 236             mess.setType(in_reply_to.getType())
 237         else:
 238             mess.setThread(self.__threads.get(user, None))
 239             mess.setType(message_type)
 240
 241         self.send_message(mess)
 242
 243     def send_simple_reply(self, mess, text, private=False):
 244         """Send a simple response to a message"""
 245         self.send_message(self.build_reply(mess, text, private))
 246
 247     def build_reply(self, mess, text=None, private=False):
 248         """Build a message for responding to another message.  Message is NOT sent"""
 249         response = self.build_message(text)
 250         if private:
 251             response.setTo(mess.getFrom())
 252             response.setType('chat')
 253         else:
 254             response.setTo(mess.getFrom().getStripped())
 255             response.setType(mess.getType())
 256         response.setThread(mess.getThread())
 257         return response
 258
 259     def build_message(self, text):
 260         """Builds an xhtml message without attributes."""
 261         text_plain = re.sub(r'<[^>]+>', '', text)
 262         message = xmpp.protocol.Message(body=text_plain)
 263         if text_plain != text:
 264             html = xmpp.Node('html', {'xmlns': 'http://jabber.org/protocol/xhtml-im'})
 265             try:
 266                 html.addChild(node=xmpp.simplexml.XML2Node("<body xmlns='http://www.w3.org/1999/xhtml'>" + text.encode('utf-8') + "</body>"))
 267                 message.addChild(node=html)
 268             except Exception, e:
 269                 # Didn't work, incorrect markup or something.
 270                 # print >> sys.stderr, e, text
 271                 message = xmpp.protocol.Message(body=text_plain)
 272         return message
 273
 274     def get_sender_username(self, mess):
 275         """Extract the sender's user name from a message"""
 276         type = mess.getType()
 277         jid = mess.getFrom()
 278         if type == "groupchat":
 279             username = jid.getResource()
 280         elif type == "chat":
 281             username = jid.getNode()
 282         else:
 283             username = ""
 284         return username
 285
 286     def status_type_changed(self, jid, new_status_type):
 287         """Callback for tracking status types (available, away, offline, ...)"""
 288         self.log.debug('user %s changed status to %s' % (jid, new_status_type))
 289
 290     def status_message_changed(self, jid, new_status_message):
 291         """Callback for tracking status messages (the free-form status text)"""
 292         self.log.debug('user %s updated text to %s' % (jid, new_status_message))
 293
 294     def broadcast(self, message, only_available=False):
 295         """Broadcast a message to all users 'seen' by this bot.
 296
 297         If the parameter 'only_available' is True, the broadcast
 298         will not go to users whose status is not 'Available'."""
 299         for jid, (show, status) in self.__seen.items():
 300             if not only_available or show is self.AVAILABLE:
 301                 self.send(jid, message)
 302
 303     def callback_presence(self, conn, presence):
 304         self.__lastping = time.time()
 305         jid, type_, show, status = presence.getFrom(), \
 306                 presence.getType(), presence.getShow(), \
 307                 presence.getStatus()
 308
 309         if self.jid.bareMatch(jid) and not self.__acceptownmsgs:
 310             # Ignore our own presence messages
 311             return
 312
 313         if type_ is None:
 314             # Keep track of status message and type changes
 315             old_show, old_status = self.__seen.get(jid, (self.OFFLINE, None))
 316             if old_show != show:
 317                 self.status_type_changed(jid, show)
 318
 319             if old_status != status:
 320                 self.status_message_changed(jid, status)
 321
 322             self.__seen[jid] = (show, status)
 323         elif type_ == self.OFFLINE and jid in self.__seen:
 324             # Notify of user offline status change
 325             del self.__seen[jid]
 326             self.status_type_changed(jid, self.OFFLINE)
 327
 328         try:
 329             subscription = self.roster.getSubscription(unicode(jid.__str__()))
 330         except KeyError, e:
 331             # User not on our roster
 332             subscription = None
 333         except AttributeError, e:
 334             # Recieved presence update before roster built
 335             return
 336
 337         if type_ == 'error':
 338             self.log.error(presence.getError())
 339
 340         self.log.debug('Got presence: %s (type: %s, show: %s, status: %s, subscription: %s)' % (jid, type_, show, status, subscription))
 341
 342         # If subscription is private, disregard anything not from the private domain
 343         if self.__privatedomain and type_ in ('subscribe', 'subscribed', 'unsubscribe'):
 344             if self.__privatedomain == True:
 345                 # Use the bot's domain
 346                 domain = self.jid.getDomain()
 347             else:
 348                 # Use the specified domain
 349                 domain = self.__privatedomain
 350
 351             # Check if the sender is in the private domain
 352             user_domain = jid.getDomain()
 353             if domain != user_domain:
 354                 self.log.info('Ignoring subscribe request: %s does not match private domain (%s)' % (user_domain, domain))
 355                 return
 356
 357         if type_ == 'subscribe':
 358             # Incoming presence subscription request
 359             if subscription in ('to', 'both', 'from'):
 360                 self.roster.Authorize(jid)
 361                 self._send_status()
 362
 363             if subscription not in ('to', 'both'):
 364                 self.roster.Subscribe(jid)
 365
 366             if subscription in (None, 'none'):
 367                 self.send(jid, self.MSG_AUTHORIZE_ME)
 368         elif type_ == 'subscribed':
 369             # Authorize any pending requests for that JID
 370             self.roster.Authorize(jid)
 371         elif type_ == 'unsubscribed':
 372             # Authorization was not granted
 373             self.send(jid, self.MSG_NOT_AUTHORIZED)
 374             self.roster.Unauthorize(jid)
 375
 376     def callback_message(self, conn, mess):
 377         """Messages sent to the bot will arrive here. Command handling + routing is done in this function."""
 378         self.__lastping = time.time()
 379
 380         # Prepare to handle either private chats or group chats
 381         type = mess.getType()
 382         jid = mess.getFrom()
 383         props = mess.getProperties()
 384         text = mess.getBody()
 385         username = self.get_sender_username(mess)
 386
 387         if type not in ("groupchat", "chat"):
 388             self.log.debug("unhandled message type: %s" % type)
 389             return
 390
 391         self.log.debug("*** props = %s" % props)
 392         self.log.debug("*** jid = %s" % jid)
 393         self.log.debug("*** username = %s" % username)
 394         self.log.debug("*** type = %s" % type)
 395         self.log.debug("*** text = %s" % text)
 396
 397         # Ignore messages from before we joined
 398         if xmpp.NS_DELAY in props: return
 399
 400         # Ignore messages from myself
 401         if username == self.__username: return
 402
 403         # If a message format is not supported (eg. encrypted), txt will be None
 404         if not text: return
 405
 406         # Ignore messages from users not seen by this bot
 407         if jid not in self.__seen:
 408             self.log.info('Ignoring message from unseen guest: %s' % jid)
 409             self.log.debug("I've seen: %s" % ["%s" % x for x in self.__seen.keys()])
 410             return
 411
 412         # Remember the last-talked-in thread for replies
 413         self.__threads[jid] = mess.getThread()
 414
 415         if ' ' in text:
 416             command, args = text.split(' ', 1)
 417         else:
 418             command, args = text, ''
 419         cmd = command.lower()
 420         self.log.debug("*** cmd = %s" % cmd)
 421
 422         if self.custom_message_handler is not None:
 423             # Try the custom handler first. It can return None
 424             # if you want JabberBot to fall back to the default.
 425             reply = self.custom_message_handler(mess, text)
 426         else:
 427             reply = None
 428
 429         if reply is None and self.commands.has_key(cmd):
 430             try:
 431                 reply = self.commands[cmd](mess, args)
 432             except Exception, e:
 433                 reply = traceback.format_exc(e)
 434                 self.log.exception('An error happened while processing a message ("%s") from %s: %s"' % (text, jid, reply))
 435         else:
 436             # In private chat, it's okay for the bot to always respond.
 437             # In group chat, the bot should silently ignore commands it
 438             # doesn't understand or aren't handled by unknown_command().
 439             default_reply = 'Unknown command: "%s". Type "help" for available commands.<b>blubb!</b>' % cmd
 440             if type == "groupchat": default_reply = None
 441             reply = self.unknown_command(mess, cmd, args)
 442             if reply is None:
 443                 reply = default_reply
 444         if reply:
 445             self.send_simple_reply(mess, reply)
 446
 447     def unknown_command(self, mess, cmd, args):
 448         """Default handler for unknown commands
 449
 450         Override this method in derived class if you
 451         want to trap some unrecognized commands.  If
 452         'cmd' is handled, you must return some non-false
 453         value, else some helpful text will be sent back
 454         to the sender.
 455         """
 456         return None
 457
 458     def top_of_help_message(self):
 459         """Returns a string that forms the top of the help message
 460
 461         Override this method in derived class if you
 462         want to add additional help text at the
 463         beginning of the help message.
 464         """
 465         return ""
 466
 467     def bottom_of_help_message(self):
 468         """Returns a string that forms the bottom of the help message
 469
 470         Override this method in derived class if you
 471         want to add additional help text at the end
 472         of the help message.
 473         """
 474         return ""
 475
 476     @botcmd
 477     def help(self, mess, args):
 478         """Returns a help string listing available options.
 479
 480         Automatically assigned to the "help" command."""
 481         if not args:
 482             if self.__doc__:
 483                 description = self.__doc__.strip()
 484             else:
 485                 description = 'Available commands:'
 486
 487             usage = '\n'.join(sorted([
 488                 '%s: %s' % (name, (command.__doc__.strip() or '(undocumented)').split('\n', 1)[0])
 489                 for (name, command) in self.commands.iteritems() if name != 'help' and not command._jabberbot_hidden
 490             ]))
 491             usage = usage + '\n\nType help <command name> to get more info about that specific command.'
 492         else:
 493             description = ''
 494             if args in self.commands:
 495                 usage = self.commands[args].__doc__.strip() or 'undocumented'
 496             else:
 497                 usage = 'That command is not defined.'
 498
 499         top = self.top_of_help_message()
 500         bottom = self.bottom_of_help_message()
 501         if top   : top = "%s\n\n" % top
 502         if bottom: bottom = "\n\n%s" % bottom
 503
 504         return '%s%s\n\n%s%s' % (top, description, usage, bottom)
 505
 506     def idle_proc(self):
 507         """This function will be called in the main loop."""
 508         self._idle_ping()
 509
 510     def _idle_ping(self):
 511         """Pings the server, calls on_ping_timeout() on no response.
 512
 513         To enable set self.PING_FREQUENCY to a value higher than zero.
 514         """
 515         if self.PING_FREQUENCY and time.time() - self.__lastping > self.PING_FREQUENCY:
 516             self.__lastping = time.time()
 517             #logging.debug('Pinging the server.')
 518             ping = xmpp.Protocol('iq', typ='get', payload=[xmpp.Node('ping', attrs={'xmlns':'urn:xmpp:ping'})])
 519             try:
 520                 res = self.conn.SendAndWaitForResponse(ping, self.PING_TIMEOUT)
 521                 #logging.debug('Got response: ' + str(res))
 522                 if res is None:
 523                     self.on_ping_timeout()
 524             except IOError, e:
 525                 logging.error('Error pinging the server: %s, treating as ping timeout.' % e)
 526                 self.on_ping_timeout()
 527
 528     def on_ping_timeout(self):
 529         logging.info('Terminating due to PING timeout.')
 530         self.quit()
 531
 532     def shutdown(self):
 533         """This function will be called when we're done serving
 534
 535         Override this method in derived class if you
 536         want to do anything special at shutdown.
 537         """
 538         pass
 539
 540     def serve_forever(self, connect_callback=None, disconnect_callback=None):
 541         """Connects to the server and handles messages."""
 542         conn = self.connect()
 543         if conn:
 544             self.log.info('bot connected. serving forever.')
 545         else:
 546             self.log.warn('could not connect to server - aborting.')
 547             return
 548
 549         if connect_callback:
 550             connect_callback()
 551         self.__lastping = time.time()
 552
 553         while not self.__finished:
 554             try:
 555                 conn.Process(1)
 556                 self.idle_proc()
 557             except KeyboardInterrupt:
 558                 self.log.info('bot stopped by user request. shutting down.')
 559                 break
 560
 561         self.shutdown()
 562
 563         if disconnect_callback:
 564             disconnect_callback()
 565
 566