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