brathuhn/brathuhn.py

import xmpp, os, sys, time, traceback, stream_ft
import threading

from config import *

VERSION = "0.2"
MIN_CONNECT_INTERVAL = 300
IGNORE_FIRST = 3

last_connect = 0
commands = dict()
commands_private = dict()
group_listeners = []
private_listeners = []
status_listeners = []
call_ats = dict()
call_ats_id = 0
join_time = time.time()+60*60*24*365 #1 year in the future

users = dict()

STATUS_JOIN, STATUS_ONLINE, STATUS_OFFLINE, STATUS_RENAME, STATUS_XAWAY, STATUS_AWAY, STATUS_UNKNOWN = range(7)
SI, JINGLE = range(2)

def getUsers():
	"""Return a dictionary of every user in the group, the values are status codes."""
	return users

def addCommand(name,help,func):
	"""Add a groupchat command.

	@param name -- The name of the command, if the command keyword 
                   "!name" is used at the start of a message the given
                   handler will be called
	@param help -- Help string displayed on with the !help command
	@param func -- Handler of the form func(msg,args) where msg is the
                   xmpppy representation of the message sent and args is
                   a list of arguments"""
	commands[name]=(help,func)

def addPrivateCommand(name,help,func):
	"""Add a private chat command.

	@param name -- The name of the command, if the command keyword 
                   "!name" is used at the start of a message the given
                   handler will be called
	@param help -- Help string displayed on with the !help command
	@param func -- Handler of the form func(msg,args) where msg is the
                   xmpppy representation of the message sent and args is
                   a list of arguments"""
	commands_private[name]=(help,func)

def addGroupListener(func):
	"""Add a groupchat listener.

	@param func -- Handler of the form func(from,msg) where both arguments
                   are unicode strings"""
	group_listeners.append(func)

def addPrivateListener(func):
	"""Adds a private chat listener.

	@param func -- Handler of the form func(from,msg) where both arguments
                   are unicode strings"""
	private_listeners.append(func)

def addStatusListener(func):
	"""Adds a user status listener.

	@param func -- Handler of the form func(from,msg,status) where first two
                   arguments are unicode strings and status is one of STATUS_*
                   constants"""
	status_listeners.append(func)

def sendMsg(to,body,mtype=None):
	"""Send a message to the server.

    @param to -- Recipient of the message
    @param body -- Message body
    @param mtype -- Type of the message, (e.g. "groupchat" for group 
                    messages)"""
	mymsg = xmpp.protocol.Message(body=body)
	mymsg.setTo(to)
	if mtype!=None:
		mymsg.setType(mtype)
	cl.send(mymsg)

def sendStyledMsg(to,body,mtype=None,style=""):
	"""Send a CSS-styled message to the server.

    @param to -- Recipient of the message
    @param body -- Message body
    @param mtype -- Type of the message, (e.g. "groupchat" for group 
                    messages)
    @param style -- Inline CSS used to format the message"""
	mymsg = xmpp.protocol.Message(body=body)
	mymsg.setTo(to)
	if mtype!=None:
		mymsg.setType(mtype)
	c=mymsg.addChild('html', namespace='http://jabber.org/protocol/xhtml-im')
	c=c.addChild('body',namespace='http://www.w3.org/1999/xhtml')
	c=c.addChild('span',payload=body,attrs={'style':style})
	cl.send(mymsg)

def sendMonoMsg(to,body,mtype=None):
	"""Send a mono-spaced styled message to the server (e.g. ASCII art).

    @param to -- Recipient of the message
    @param body -- Message body
    @param mtype -- Type of the message, (e.g. "groupchat" for group 
                    messages)"""
	mymsg = xmpp.protocol.Message(body=body)
	mymsg.setTo(to)
	mymsg.setType(mtype)
	c=mymsg.addChild('html', namespace='http://jabber.org/protocol/xhtml-im')
	c=c.addChild('body',namespace='http://www.w3.org/1999/xhtml')
	c=c.addChild('span', attrs={'style':'font-family: Courier New, monospace;'})
	c=c.addChild('pre',payload=body,attrs={'style':'background-color: #ffffff;'})
	cl.send(mymsg)

def __help(msg,args):
	sendMsg(ROOM,"Available commands:",'groupchat')
	for cmd in commands:
		sendMsg(ROOM,"    !"+cmd+" - "+commands[cmd][0],'groupchat')
	sendMsg(ROOM,"Available private chat commands:",'groupchat')
	for cmd in commands_private:
		sendMsg(ROOM,"    !"+cmd+" - "+commands_private[cmd][0],'groupchat')

def __help2(msg,args):
	sendMsg(msg.getFrom(),"Available commands:")
	for cmd in commands:
		sendMsg(msg.getFrom(),"    !"+cmd+" - "+commands[cmd][0])
	sendMsg(msg.getFrom(),"Available private chat commands:")
	for cmd in commands_private:
		sendMsg(msg.getFrom(),"    !"+cmd+" - "+commands_private[cmd][0])

def __messageCB(conn,msg):
	if msg.getType() == "groupchat":
		body = unicode(msg.getBody())
		mfrom = msg.getFrom()

		for cmd in commands:
			if body.startswith("!"+cmd) and not room_nick in unicode(mfrom):
				args = body.split()
				commands[cmd][1](msg,args)
		for listener in group_listeners:
			listener( mfrom , msg )
			
	if msg.getType() == "chat":
		if msg.getBody() != None:
			body = unicode(msg.getBody())
			mfrom = msg.getFrom()
			for cmd in commands_private:
				if body.startswith("!"+cmd) and not room_nick in unicode(mfrom):
					args = body.split()
					commands_private[cmd][1](msg,args)
			for listener in private_listeners:
				listener( mfrom, msg )
	raise xmpp.NodeProcessed

def __presenceCB(conn,msg):
	if msg.getType()=='unavailable':
		nick=msg.getNick()
		if nick==None:
			if msg.getFrom() in users:
				del users[msg.getFrom()]
			#Disconnect
			if not msg.getFrom() in users:	
				for listener in status_listeners:
					listener(msg.getFrom(),msg,STATUS_OFFLINE)
		else:
			if msg.getFrom() in users:
				del users[msg.getFrom()]
			#Rename
			if not msg.getFrom() in users:	
				for listener in status_listeners:
					listener(msg.getFrom(),msg,STATUS_RENAME)	
	else:
		show = msg.getShow()
		if show==None:
			if not msg.getFrom() in users:	
				for listener in status_listeners:
					listener(msg.getFrom(),msg,STATUS_JOIN)	

			users[msg.getFrom()] = STATUS_ONLINE
			#Online
		else:
			#print msg.getFrom(),"is now",show
			if not msg.getFrom() in users:	
				for listener in status_listeners:
					listener(msg.getFrom(),msg,STATUS_JOIN)
			if   show=="away":
				users[msg.getFrom()] = STATUS_AWAY
			elif show=="xa":
				users[msg.getFrom()] = STATUS_XAWAY
			else:
				users[msg.getFrom()] = STATUS_UNKNOWN
			# status change

	status = STATUS_OFFLINE
	if msg.getFrom() in users:
		status = users[msg.getFrom()]
	
	for listener in status_listeners:
		listener(msg.getFrom(),msg.getStatus(),status)

def __ignore_first(conn,msg):
	global join_time
	dt = time.time()-join_time
	if dt>IGNORE_FIRST:
		cl.RegisterHandler('message',__messageCB)
		cl.UnregisterHandler('message',__ignore_first)
		__messageCB(conn,msg)


jid=xmpp.protocol.JID(node=NAME,domain=SRV,resource=RESOURCE)
cl=xmpp.Client(jid.getDomain(),debug=[])

def connect_and_join():
	"""Connect to the server and join the group chat room from the config."""
	global room_nick
	global room_to
	global last_connect
	global cl
	global join_time
	last_connect=time.time()
	#connect to server
	con=cl.connect()
	if not con:
		print 'could not connect!'
		raise Exception('could not connect!')
	print 'connected'

	cl.RegisterHandler('message',__ignore_first)
	cl.RegisterHandler('presence',__presenceCB)
	cl.RegisterHandler('iq',stream_ft.ft_listener)

	#authenticate
	auth=cl.auth(NAME,PWD,resource=jid.getResource())
	if not auth:
		print 'could not authenticate!'
		raise Exception('could not authenticate!')
	print 'auth'


	room_to = "%s/%s" % (ROOM, NAME)
	room_nick = ROOM+"/"+NAME
	msg = xmpp.Presence(to=room_to)
	if PWD_ROOM != "":
		msg.addChild('x', namespace='http://jabber.org/protocol/muc').addChild('password',payload=PWD_ROOM)

	#add the help command
	addCommand("help","Prints this help message",__help)
	addPrivateCommand("help","Prints this help message (in private chat)",__help2)

	#join the room
	cl.send(msg)
	join_time = time.time()
	#set status
	cl.send(xmpp.Presence(to=room_to,status=NAME+" --(brathuhn bot framework (v"+VERSION+"), type !help)"))

def loop():
	"""Loop forever and handle the event-queue."""
	global last_connect
	global call_ats
	while True:
		try:
			cl.Process(1)
			to_remove = []
			for cid in call_ats:
				call=call_ats[cid]
				if time.time()>=call[1]:
					call[0](*call[2])
					to_remove.append(cid)
			for cid in to_remove:
				del call_ats[cid]

		except xmpp.protocol.Conflict:
			print "detected another connection, quitting..."
			sys.exit(0)
		except KeyboardInterrupt:
			print "keyboard interrupt, quitting..."
			sys.exit(0)
		except:
			print "-------ERROR OCCURED-------"
			exc_type, exc_value, exc_traceback = sys.exc_info()
			traceback.print_exc(file=sys.stdout)
			if not cl.connected:				
				s=MIN_CONNECT_INTERVAL-(time.time()-last_connect)
				if s>0:
					time.sleep(s)
				connect_and_join()
			print "---------------------------"

def callAt(func,_time,*args):
	"""Call a function at given time; The call is not very exact but it is 
    guaranteed that it will happen after the specified time. The call happens
    in the same thread, so no synchronization is needed. You can append any 
    number of arguments you need. Returns the call id, needed to cancel a call

    @param func -- Function to be called
    @param time -- Message body"""
	global call_ats
	global call_ats_id
	call_ats[call_ats_id]=[func,_time,args]
	call_ats_id+=1
	return call_ats_id-1

def cancelCall(cid):
	"""Remove the call from the list of pending calls

	@param cid -- The id of the call to remove
	"""
	if cid in call_ats:
		del call_ats[cid]

def sendFile(to,path,method=SI):
	"""Send a file using method specified

    @param to     -- Reciever of the file
    @param path   -- Path to the file
    @param method -- Initiation method, currently only stream initiation is supported
	"""

	size = os.path.getsize(path)
	name = os.path.basename(path)
	fd = open(path,'rb')
	if method==SI:
		cim = stream_ft.create_init_message(room_nick,to,name,size,fd)
		cl.send(cim)
	else:
		fd.close()

def sendFileStream(to,size,name,fd,method=SI):
	"""Send a file using method specified

    @param to     -- Reciever of the file
    @param size   -- Number of bytes to send
	@param name   -- The name of the file the user will see
	@param fd     -- A file like object
    @param method -- Initiation method, currently only stream initiation is supported
	"""

	if method==SI:
		cim = stream_ft.create_init_message(room_nick,to,name,size,fd)
		cl.send(cim)
	else:
		fd.close()