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"
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:"