jabberbot.py

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