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