jabberbot.py

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