jabberbot.py

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