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