Simple XMPP/Jabber bot

Description: xmpp_bot is a minimal, extendable bot which provides easy function integration as well as a basic rights management.
Raw download: xmpp_bot.py / bot.py
Jump to: xmpp_bot.py / bot.py

xmpp_bot.py (Raw)

  1 #!/usr/bin/env python
  2 # -*- coding: utf-8
  3 #
  4 #    xmpp_bot.py - A small bot written for personal fun ;-) It is easily extendable
  5 #                  and very nice to play with in interactive mode. To extend it just
  6 #                  derive it and add new callbacks, gl&hf
  7 #
  8 #    TODO:
  9 #            - Make everything threadsafe
 10 #            - Lots and lots of other stuff ;-)
 11 #                    .    .    .
 12 #    Copyright (C) 2007 Stefan Hacker
 13 #
 14 #    This program is free software; you can redistribute it and/or modify
 15 #    it under the terms of the GNU General Public License as published by
 16 #    the Free Software Foundation; either version 2 of the License, or
 17 #    (at your option) any later version.
 18 #
 19 #    This program is distributed in the hope that it will be useful,
 20 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
 21 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 22 #    GNU General Public License for more details.
 23 #
 24 #    You should have received a copy of the GNU General Public License along
 25 #    with this program; if not, write to the Free Software Foundation, Inc.,
 26 #    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 27 
 28 import os, sys, xmpp, threading, Queue, logging, time
 29 
 30 
 31 log = logging.getLogger(__name__)
 32 
 33 class state(object):
 34     """Possible states of a running xmpp_bot"""
 35     undefined, connect, auth, running, stop, stopped = range(6)
 36 
 37 class user(object):
 38     """
 39     A simple user class
 40     """
 41     def __init__(self, jid, firstname = None, secondname = None, comment = None, rights = None, admin = False):
 42         self.jid          = jid.split('/',1)[0]
 43         self.firstname    = firstname
 44         self.secondname   = secondname
 45         self.comment      = comment
 46         self.rights       = rights or set()
 47         self.admin        = admin
 48         
 49     def __str__(self):
 50         return self.jid
 51     
 52     def __hash__(self):
 53         return hash(self.jid)
 54     
 55     def __eq__(self, other):
 56         if type(other) == str:
 57             return self.jid==other
 58         elif type(other) == instance and self.__class__ is other.__class__:
 59             return self.jid == other.jid
 60         return False
 61     
 62     def whois(self):
 63         """Returns a string with some info about the user"""
 64         ret = 'JID: %s\n' % self.jid
 65         if self.firstname or self.secondname:
 66             ret += 'Name: %s %s' % (self.firstname, self.secondname)
 67         ret += 'Rights: %s' % repr(self.rights)
 68         if self.comment:
 69             ret += 'Comment: %s' % (self.comment)
 70         return ret
 71     
 72     def right_add(self, command):
 73         """
 74         command    - Identifier of the command
 75         
 76         Enables the user to run a certain command
 77         """
 78         return self.rights.add(command)
 79     
 80     def right_revoke(self, command):
 81         """
 82         command    - Identifier of the command
 83         
 84         Revokes the users right to execute given command
 85         """
 86         if command in self.rights:
 87             self.rights.remove(command)
 88             return True
 89         else:
 90             return False
 91         
 92     def is_allowed(self, command):
 93         """
 94         command    - Identifier of the command
 95         
 96         Checks if the user has the right to exeture the
 97         given command
 98         """
 99         return self.admin or command in self.rights
100 
101 class user_manager(object):
102     """
103     A simple user manager class
104     """
105     def __init__(self, users = None):
106         self.users = users or set()
107         
108     def is_allowed(self, us, command):
109         """
110         us      - Username    (str)
111         command - Identifier of the command (str)
112         
113         Checks if a certain user is allowed to run a certain command
114         """
115         us = us.split('/', 1)[0]
116         us = self.user_get(us)
117             
118         if us in self.users:
119             return us.is_allowed(command)
120         return False
121     
122     def user_add(self, us):
123         """
124         us      - User object
125         
126         Adds a user
127         """
128         return self.users.add(us)
129     
130     def user_remove(self, us):
131         """
132         us - User object / Name
133         
134         Removes a user
135         """
136         if us in self.users:
137             self.users.remove(us)
138             return True
139         else:
140             return False
141         
142     def user_get(self, us):
143         """
144         us - Username
145         
146         Returns the object according to the given username
147         or none if the user does not exists.
148         """
149         us = us.split('/',1)[0]
150         for obj in self.users:
151             if us == obj:
152                 return obj
153         return None
154     
155     
156 class xmpp_bot(threading.Thread, user_manager):
157     """
158     A simple bot
159     """
160     def __init__(self, jid, password, ssl=1, sasl=1, server=None, proxy=None, users = None):
161         """Standard constructor"""
162         self.jid           = xmpp.protocol.JID(jid)
163         self.ssl           = ssl
164         self.server        = server
165         self.proxy         = proxy
166         self.password      = password
167         self.sasl          = sasl
168         
169         self._state        = state.undefined
170         self.sema_state    = threading.Semaphore()
171         
172         self._handlers     = {}
173         self.sema_handler  = threading.Semaphore()
174         
175         self._roster       = None
176         self.sema_xmpp     = threading.Semaphore()
177         
178         self._context      = {}
179         
180         self._connection   = xmpp.Client(self.jid.getDomain(), debug=[])        
181         
182         user_manager.__init__(self, users)
183         self._register_default_users()
184         self._register_default_handlers()
185         
186         threading.Thread.__init__(self)
187         
188     def run(self):
189         """Runner of the xmpp_bot"""
190         log.info("Start")
191         #Connect
192         self.sema_xmpp.acquire()
193         log.debug("Try to connect to %s..." % self.jid.getDomain())
194         res = self._connection.connect(self.server, self.proxy, self.ssl)
195         if not res:
196             log.error('Connection failed')
197             self.state = state.stopped
198             self.sema_xmpp.release()
199             return
200         elif res=='ssl':
201             log.debug("Connected [SSL]")
202         else:
203             log.warning("Connected [INSECURE]")
204         
205         #Authenticate
206         log.debug("Auth as '%s'..." % self.jid)
207         res = self._connection.auth(self.jid.getNode(), self.password, self.jid.getResource(), self.sasl)
208         if not res:
209             log.error('Auth failed')
210             self.con.disconnect()
211             self.state = state.stopped
212             self.sema_xmpp.release()
213             return
214         elif res=='sasl':
215             log.debug("Authed [SASL]")
216         else:
217             log.warning("Authed [INSECURE]")
218             
219         #Register handler
220         self._connection.RegisterHandler('message', self._handle_message)
221         self._connection.RegisterDisconnectHandler(self._handle_disconnect)
222         
223         log.debug('Send intial presence')
224         self._connection.sendInitPresence(requestRoster=0)
225         
226         log.debug('Retrieving roster')
227         self._roster = self._connection.getRoster()
228         
229         self.state = state.running
230         log.info('Running')
231         self.sema_xmpp.release()
232         while self.state is state.running:
233             self.sema_xmpp.acquire()
234             self._connection.Process(0.5)
235             self.sema_xmpp.release()
236         
237         self.sema_xmpp.acquire()
238         self.state = state.stop
239         
240         log.debug('Disconnecting')
241         self._connection.disconnect()
242         
243         self.state = state.stopped
244         log.info("Stopped")
245         self.sema_xmpp.release()
246         
247     #State property
248     def _set_state(self, data):
249         self.sema_state.acquire()
250         self._state = data
251         self.sema_state.release()
252         
253     def _get_state(self):
254         self.sema_state.acquire()
255         ret = self._state
256         self.sema_state.release()
257         return ret
258     
259     state = property(_get_state, _set_state, doc='State of the bot')
260     
261     #Normal functions
262     def send_message(self, target, body, type='chat'):
263         """
264         target    - Target jid      (str)
265         body      - message to send (str)
266         type      - Type of message (str)
267             
268         Send a message to a given jid.
269         
270         INFO: Handlers must use _send_message or deadlocks
271         may occur
272         """
273         self.sema_xmpp.acquire()
274         ret = self._send_message(target, body, type)
275         self.sema_xmpp.release()
276         return ret
277     
278     def _send_message(self, target, body, type='chat'):
279         """
280         See send_message. Do not call this function out auf
281         _handle_message and it's handler functions as it is
282         NOT THREAD SAFE
283         """
284         return self._connection.send(xmpp.Message(target, body, type))
285     
286     #Handler management
287     def handler_register(self, command, callback):
288         """
289         command    - Identifier of the command (str)
290         callback   - Callback function
291         
292         Register a new command
293         """
294         log.debug("handler_register called")
295         self.sema_handler.acquire()
296         if not command in self._handlers:
297             log.info("New command '%s' registered!" % command)
298             self._handlers[command] = callback
299             ret = True
300         else:
301             log.warning("Tried to re-register existing command '%s'!" % command)
302             ret = False
303         self.sema_handler.release()
304         log.debug("handler_register return")
305         return ret
306     
307     def handler_unregister(self, command):
308         """
309         command   - Identifier of the command (str)
310         
311         Unregister an existing handler and return its callback.
312         If the handle did not exists None is returned
313         """
314         log.debug("handler_unregister called")
315         self.sema_handler.acquire()
316         if command in self._handlers:
317             log.info("Deregister command '%s'!" % command)
318             ret = self._handlers.pop(command)
319         else:
320             log.warning("Tried to unregister not existing command '%s'!" % command)
321             ret = None
322         self.sema_handler.release()
323         log.debug("handler_unregister return")
324         return ret
325     
326     def handler_get(self, command):
327         """
328         command    - Identifier of the command (str)
329         
330         Return the callback of a command. In case
331         of failure None is returned
332         """
333         self.sema_handler.acquire()
334         if command in self._handles:
335             ret = self._handlers[command]
336         else:
337             ret = None
338         self.sema_handler.release()
339         return ret
340     
341     #Basic handlers
342     def _register_default_handlers(self):
343         """Function in which default handlers get registered"""
344         self.handler_register('help', self._default_handler_help)
345     
346     def _register_default_users(self):
347         """Function in which default users get registered"""
348         self.user_add(user('default', rights = set(['help']), comment = "The default user"))
349     
350     def _handle_message(self, con, msg):
351         """Standard message handling function of the bot"""
352         source = str(msg.getFrom())
353         
354         if source in self._context:
355             handler = self._context[source]
356             param   = msg.getBody()
357             log.info("'%s': Continue in context" % source)
358         else:
359             #Extract the command and the parameters
360             tmp = msg.getBody().split(' ',1)
361             if len(tmp)==2:
362                 (cmd, param) = tmp
363             else:
364                 cmd     = tmp[0]
365                 param   = None
366             
367             self.sema_handler.acquire()
368             if (cmd in self._handlers) and (self.is_allowed(source, cmd)):
369                 handler = self._handlers[cmd]
370                 log.info("'%s': Call '%s' " % (source, cmd))
371             else:
372                 self.sema_handler.release()
373                 log.info("'%s': Unknown command '%s'" % (source, cmd))
374                 self._send_message(source, "Command '%s' is unknown! Try 'help' for more information" % cmd)
375                 return
376             self.sema_handler.release()
377 
378         #Call the corresponding handler
379         start = time.clock()
380         #try:
381         ret = handler(self, source, param)
382         if ret:
383             self._context[source] = ret
384             log.info("New context set")
385         elif source in self._context:
386             del self._context[source]
387             log.info("Returned from context")
388         log.info("Handler returned (Time: %dms)" % (int((time.clock()-start)*1000)))
389 
390     
391     def _handle_disconnect(self):
392         """Standard disconnect handler"""
393         log.debug('Disconnect handler called')
394         self.state = state.stop
395         
396     #Default handlers
397     def _default_handler_help(self, caller, source, param):
398         """List all available commands"""
399         self.sema_handler.acquire()
400         if param:
401             if param in self._handlers and hasattr(self._handlers[param], '_usage'):
402                 help = 'Usage:\n'
403                 help += self._handlers[param]._usage
404             else:
405                 help = "Sorry, no information available for '%s'" % param
406         else:
407             help = 'Help:\n'
408             for cmd, handler in self._handlers.iteritems():
409                 if self.is_allowed(source, cmd):
410                     help += "%s - %s\n" % (cmd, handler.__doc__)        
411             help += "Try help [COMMAND] for further information"
412         self.sema_handler.release()
413         self._send_message(source, help)
414         
415 if __name__== "__main__":
416     import getpass
417     
418     log.setLevel(logging.DEBUG)
419     #create console handler and set level to debug
420     ch = logging.StreamHandler()
421     ch.setLevel(logging.DEBUG)
422     #create formatter
423     formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
424     #add formatter to ch
425     ch.setFormatter(formatter)
426     #add ch to logger
427     log.addHandler(ch)
428     
429     #Gather information
430     username = raw_input("Username: ")
431     server   = raw_input("Server  : ")
432     resource = "bot"
433     port     = 5223
434     #myuser = raw_input("JID/RESOURCE: ")
435     pw   = getpass.getpass()
436     
437     #Start the bot
438     a = xmpp_bot(username+'@'+server+'/'+resource , pw, server =(server, port))
439     a.user_add(user(username+'@'+server, admin=True))
440     a.start()
441     
442     while a.state < state.running:
443         time.sleep(0.5)
444         
445     print "Interactive mode reached:\n"

bot.py (Raw)

  1 #!/usr/bin/env python
  2 # -*- coding: utf-8
  3 #
  4 #    bot.py - Is a playground to test functionality in xmpp_bot.py
  5 #    Copyright (C) 2007 Stefan Hacker
  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 2 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 along
 18 #    with this program; if not, write to the Free Software Foundation, Inc.,
 19 #    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 20 
 21 
 22 import time, getpass, logging, sys, os, random
 23 import xmpp_bot
 24 from xmpp_bot import state, user
 25 
 26 log = logging.getLogger("bot")
 27 
 28 class mybot(xmpp_bot.xmpp_bot):
 29     """My simple bot"""
 30     def __init__(self, jid, password, ssl=1, sasl=1, server=None, proxy=None, users = None):
 31         xmpp_bot.xmpp_bot.__init__(self, jid, password, ssl, sasl, server, proxy, users)
 32         
 33     def _register_default_users(self):
 34         #Define a default user
 35         def_rights=set(['help',
 36                         'time',
 37                         'ping',
 38                         'info',
 39                         'echo',
 40                         'fact',
 41                         'state'])
 42         
 43         self.user_add(user('default', rights=def_rights, comment = 'Default user'))
 44     
 45     def _register_default_handlers(self):
 46         xmpp_bot.xmpp_bot._register_default_handlers(self)
 47         self.handler_register('time', self.timeCB)
 48         self.handler_register('ping', self.pingCB)
 49         self.handler_register('info', self.runtimeCB)
 50         self.handler_register('echo', self.echoCB)
 51         self.handler_register('fact', self.factCB)
 52         self.handler_register('exit', self.exitCB)
 53         self.handler_register('state', self.stateCB)
 54               
 55     def is_allowed(self, us, command):
 56         return xmpp_bot.xmpp_bot.is_allowed(self, 'default', command) or xmpp_bot.xmpp_bot.is_allowed(self, us, command)
 57     
 58     def factCB(self, caller, source, param):
 59         """Print a important fact about life"""
 60         
 61         if not hasattr(self, 'fact_cb_lines'):
 62             log.debug("Initial caching of 'facts.txt'")
 63             try:
 64                 f = open("facts.txt","r")
 65             except IOError:
 66                 log.error("Open file 'facts.txt' failed!")
 67                 self._send_message(source, 'Sorry, facts are out for today ;-)')
 68                 return
 69             else:
 70                 self.fact_cb_lines = f.readlines()
 71                 f.close()
 72         line = self.fact_cb_lines[random.randint(0,len(self.fact_cb_lines)-1)][:-1]
 73         log.debug("Fact: '%s'" % line)
 74         self._send_message(source, line)
 75                 
 76     def echoCB(self, caller, source, param):
 77         """Echos your message"""
 78         if param is not None:
 79             self._send_message(source, param)
 80             log.debug("Echo: '%s'" % param)
 81         
 82     echoCB._usage = "echo [STRING]"
 83         
 84     def timeCB(self, caller, source, param):
 85         """Prints the local time of the bot"""
 86         body = time.asctime(time.localtime())
 87         log.debug("Time: %s" % body)
 88         self._send_message(source, body)
 89     
 90     def pingCB(self, caller, source, param):
 91         """Answers with pong"""
 92         self._send_message(source, 'pong')
 93         
 94     def runtimeCB(self, caller, source, param):
 95         """Prints some runtime info"""
 96         body = "This bot is running under %s with python %s" % (sys.platform,sys.version)
 97         body += ". It uses the xmpp module and is currently connected as %s" % (self.jid) 
 98         log.debug("Runtime: '%s'" % body)
 99         self._send_message(source, body)
100         
101     def exitCB(self, caller, source, param):
102         """Kills the bot"""
103         self._send_message(source, "Goodbye")
104         log.info("User '%s' requested bot shutdown!" % source)
105         self.state = state.stop
106         
107     def stateCB(self, caller, source, param):
108         """A simple test for state handling"""
109         self._send_message(source, "Type 'exit' to leave")
110         log.debug("User '%s' called state testing" % source)
111         return self.stateCB_internal
112     
113     def stateCB_internal(self, caller, source, param):
114         if param=='exit':
115             self._send_message(source, "You left the state test!")
116             log.debug("User %s left state" % source)
117             return None
118         
119         self._send_message(source, "State echo: '%s'" % param)
120         log.debug("SEcho: '%s'" % param)
121         return self.stateCB_internal
122     
123 if __name__ == "__main__":
124     #Enable logging
125     log.setLevel(logging.DEBUG)
126     xmpp_bot.log.setLevel(logging.DEBUG)
127     
128     ch = logging.StreamHandler()
129     ch.setLevel(logging.DEBUG)
130     formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
131     ch.setFormatter(formatter)
132     
133     log.addHandler(ch)
134     xmpp_bot.log.addHandler(ch)
135     #Gather information
136     username = raw_input("Username: ")
137     server   = raw_input("Server  : ")
138     resource = "bot"
139     port     = 5223
140     pw       = getpass.getpass()
141     
142     #Start the bot
143     bot = mybot(username+'@'+server+'/'+resource , pw, server =(server, port))
144     #Add this account as admin
145     bot.user_add(user(username+'@'+server, admin=True))
146     bot.start()
147     
148     while bot.state < state.running:
149         time.sleep(0.5)
150         
151     print "Interactive mode reached:"    

Created on Wed Aug 29 14:52:02 2007 using pygments.