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