#!/usr/bin/python
# ---------------------------------------------------------------------------- #
# Mario Chamorro | http://www.chamorro.us/ | mario@chamorro.us | 617.895.7038
# ---------------------------------------------------------------------------- #
"""

This chat server script is intended to be run from the command line and
terminated either with control-C or 'kill PID'. The clients should connect
to enter commands using 'telnet localhost PORT' where PORT is configured 
as 'THE_PORT' at the top of the script.

In the interest of scalability, a multi-process and multi-threaded version 
was attempted, but the inter process communication became a bigger task than 
the original puzzle, so the 'select()' based server was used so that global 
list / arrays and dictionary / hashes would be the common data structures.

Robustness was acheived by placing error checking throughout the script.
The primary filter was created by seperating the command validation into 
it's own function ( FX_PARSE ) which would return a standard error message 
for any unknown command. All incoming data is pushed through this function.
Other command handling functions did more specific error checking, such as:

	- checking if the user is already logged in with a different username
	- checking if the username is taken
	- disallowing a client to join a room before logging in
	- disallowing a user to leave a room not joined
	- disallowing a client to message before logging in or joining a room
	- catching a control-C or 'kill' signal to terminate the server

Another aspect to robustness is reusing tested functions rather than
duplicating code and creating another vector for bugs to enter the logic.
In this spirit, FX_MSG_ROOM() calls FX_MSG_USER() for all members of a room
rather than creating more logic to duplicate the same functionality. Other
functions also call each other, such as FX_LOGOUT() and FX_PART(), to ensure
that a user is removed from all rooms before leaving. Likewise, FX_LOGOUT()
is called when a client disconnects abruptly to avoid orphaned nodes in rooms.

There are some inefficiencies like the many regular expressions, integers,
and argument counters for all the commands ( e.g. ARGS[iLOGIN]=nLOGIN ).
These are a compromise between defined constants and elegant command
traversal that might have been more appropriate with a C-like 'switch'
statement, but none is available in Python. Also, time constraints
made a more concise solution difficult to deliver. Certainly if this
was a production server, the extra time would have been justified.

"""
# ---------------------------------------------------------------------------- #

import os , re , string , socket , select , signal

# - USER MODIFIABLE CONSTANTS
THE_HOST = ''; 
THE_PORT = 8888;
MAX_INPUT=64;

# - SOCKET SETUP
SSS = socket.socket ( socket.AF_INET, socket.SOCK_STREAM )
SSS.bind ((THE_HOST, THE_PORT))
SSS.listen(5)

# - COMMUNICATION AND DATA STRUCTURES
QQIN = [SSS];	# - queue incoming
QOUT = []; 	# - queue outgoing
DATA = {}; 	# - data transferred
ADDR = {};	# - sockets / addresses
ROOM = {};	# - chat rooms / users
USER = {};	# - users / addresses

# --- COMMANDS --- #
# LOGIN <USERNAME>
# LOGOUT 
# JOIN  #<ROOM> 
# PART  #<ROOM> 
# MSG   #<ROOM>    <MESSAGE>
# MSG   <USERNAME> <MESSAGE>
# --- COMMANDS --- #

# - COMMAND CONTROL AND EXPRESSIONS 
RX=list(); 
RX.append(re.compile("LOGIN")); RX.append(re.compile("LOGOUT")); 
RX.append(re.compile("JOIN")); RX.append(re.compile("PART")); 
RX.append(re.compile("MSG")); RX.append(re.compile("MSG"));
RX.append(re.compile("#"));

iLOGIN=0; iLOGOUT=1; iJOIN=2; iPART=3; iMSG=4; iMSGU=5; iHASH=6;
nLOGIN=2; nLOGOUT=1; nJOIN=2; nPART=2; nMSG=3; nMSGU=3; # nHASH=1;

ARGS=[]; ARGS.append(nLOGIN); ARGS.append(nLOGOUT); 
ARGS.append(nJOIN); ARGS.append(nPART); 
ARGS.append(nMSG); ARGS.append(nMSGU);

# - GENERAL PURPOSE EXIT FUNCTION
def FX_EXIT():
	print "+-----------------------------------------------------------------------------+"
	print ""
	os._exit(0)

# - SIGNAL HANDLER ( SIGTERM / kill -15 )
def FX_SIGNAL( signal, frame ):
	FX_EXIT()

# - INPUT VALIDATOR
def FX_PARSE( INCOMING ):
	RETURN=["ERROR","INVALID COMMAND\n"]
	INPUT=INCOMING.upper().strip().split()
	# - check for no input at all ( \r\n )
	try:
		len( INPUT[0] )
	except IndexError:
		return RETURN

	# - LOGIN
	if ( RX[iLOGIN].match(INPUT[0]) and len(INPUT) == ARGS[iLOGIN] ) :
		INPUT.insert(0,iLOGIN);
		RETURN=[]
	else:
	# - LOGOUT
		if ( RX[iLOGOUT].match(INPUT[0]) and len(INPUT) == ARGS[iLOGOUT] ) :
			INPUT.insert(0,iLOGOUT);
			RETURN=[]
		else:
	# - JOIN
			if ( RX[iJOIN].match(INPUT[0]) and len(INPUT) == ARGS[iJOIN] and RX[iHASH].match(INPUT[1]) ):
				INPUT.insert(0,iJOIN);
				RETURN=[]
			else:
	# - PART
				if ( RX[iPART].match(INPUT[0]) and len(INPUT) == ARGS[iPART] and RX[iHASH].match(INPUT[1]) ):
					INPUT.insert(0,iPART);
					RETURN=[]
				else:
	# - MSG ROOM
					if ( RX[iMSG].match(INPUT[0]) and len(INPUT) >= ARGS[iMSG] and RX[iHASH].match(INPUT[1]) ):
						INPUT.insert(0,iMSG);
						RETURN=[]
					else:
	# - MSG USER
						if ( RX[iMSG].match(INPUT[0]) and len(INPUT) >= ARGS[iMSG] ) :
							INPUT.insert(0,iMSGU);
							RETURN=[]
						# else: default RETURN error value set initially
	if ( len(RETURN) == 0 ):
		return INPUT
	else:
		return RETURN

# - GIVEN AN (IP,PORT) RETURN USERNAME
def FX_GETUSER( THE_ADDRESS ):
	RETURN=""
	for KKK in USER.keys():
		if ( USER[KKK] == THE_ADDRESS ):
			RETURN=KKK
	return RETURN

# - LOGIN CLIENT AS USERNAME
def FX_LOGIN(NEW_ADDRESS, NEW_USERNAME):
	RETURN="OK\n"
	RESULT=FX_GETUSER( NEW_ADDRESS )
	if ( RESULT ):
		RETURN="ERROR ALREADY LOGGED IN AS " + RESULT + "\n"
	else:
		if ( NEW_USERNAME in USER.keys() ):
			RETURN="ERROR " + NEW_USERNAME + " IS ALREADY TAKEN\n"
		else:
			USER[NEW_USERNAME]=NEW_ADDRESS
	return RETURN

# - LOGOUT USERNAME AND LEAVE ALL ROOMS
def FX_LOGOUT( THE_ADDRESS ):
	RETURN="OK\n"
	RESULT=FX_GETUSER( THE_ADDRESS )
	if ( RESULT ):
		for KKK in ROOM.keys():
			if ( ROOM[KKK].count(RESULT) ):
				FX_PART( RESULT, KKK )
		del USER[RESULT]
	else:
		RETURN="ERROR USER NOT FOUND\n"
	return RETURN

# - JOIN / CREATE ROOM
def FX_JOIN(UUU,RRR):
	RETURN="OK\n"
	if ( len(UUU) < 1 ):
		return "ERROR USER NOT LOGGED IN\n"
		
	if ( ROOM.has_key( RRR ) and ( ROOM[RRR].count(UUU) == 0 ) ) :
		ROOM[RRR].append(UUU)
	else:
		if ( ROOM.has_key( RRR ) and ( ROOM[RRR].count(UUU) > 0 ) ) :
			RETURN="ERROR " + UUU + " ALREADY JOINED " + RRR + "\n"
		else:
			ROOM[RRR]=[UUU]
	return RETURN

# - LEAVE / DEPART A ROOM
def FX_PART(UUU,RRR):
	RETURN="OK\n"
	if ( len(UUU) < 1 ):
		return "ERROR USER NOT LOGGED IN\n"
	if ( ROOM.has_key( RRR ) and ROOM[RRR].count(UUU) > 0 ):
		ROOM[RRR].remove(UUU)
	else:
		RETURN="ERROR " + UUU + " NOT IN " + RRR + "\n"
	return RETURN

# - GIVEN A USERNAME, FIND THE SOCKET AND RELAY MESSAGE
def FX_MSG_USER(UUU, TAG, TXT):
	global ADDR , DATA , QOUT
	RETURN="OK\n"

	if ( USER.has_key(UUU) ):
		OUT_ADDRESS=USER[UUU]
		for KKK in ADDR.keys():
			if ( ADDR[KKK] == OUT_ADDRESS ):
				SSS_USER=KKK
		DATA[SSS_USER] = DATA.get( SSS_USER , '' ) + TAG + " " + TXT + "\n"
		if SSS_USER not in QOUT:
			QOUT.append( SSS_USER )
	else:
		RETURN="ERROR " + UUU + " NOT FOUND\n"

	return RETURN

# - RELAY MESSAGE TO ALL USERS IN A ROOM
def FX_MSG_ROOM(UUU, RRR, TXT):
	RETURN="OK\n"
	if ( ROOM.has_key( RRR ) and ROOM[RRR].count(UUU) > 0 ):
		for USERNAME in ROOM[RRR]:
			RETURN=FX_MSG_USER( USERNAME, "GOTROOMMSG", TXT )
	else:
		RETURN="ERROR ROOM " + RRR + " NOT AVAILABLE\n"
	return RETURN

# --- MAIN PROGRAM ----------------------------------------------------------- #
def FX_MAIN():

	print ""
	print "+-----------------------------------------------------------------------------+"
	print "| Mario Chamorro | http://www.chamorro.us/ | mario@chamorro.us | 617.895.7038 |"
	print "+-----------------------------------------------------------------------------+"

	# - signal handler to terminate server if running as a background process
	signal.signal(signal.SIGTERM, FX_SIGNAL)
	while True:

		# --- INCOMING --- #
		try:	
			SSS_IN, SSS_OUT, SSS_EXCEPTION = select.select( QQIN, QOUT, [] )
		except KeyboardInterrupt:
			FX_EXIT()
		for SSS_EVENT in SSS_IN:
			if SSS_EVENT is SSS:
				# - event is a new connection
				SSS_NEW, SSS_ADDRESS = SSS.accept()
#				print "   Connected: ", SSS_ADDRESS
				QQIN.append( SSS_NEW )
				ADDR[SSS_NEW] = SSS_ADDRESS
			else:
				# - event is either new data or a disconnection
				SSS_INCOMING = SSS_EVENT.recv( MAX_INPUT )
				if SSS_INCOMING:
				# --- MAIN SECTION --------------------------- #
					# - check incoming data
					VALID=FX_PARSE( SSS_INCOMING )
					if ( VALID[0] == "ERROR" ):
						DATA[SSS_EVENT] = DATA.get( SSS_EVENT , '' ) + string.join( VALID )
					else: 
					# - match command, run function, and append to outgoing data
						if ( VALID[0] == iLOGIN ):
							DATA[SSS_EVENT] = DATA.get( SSS_EVENT , '' ) + FX_LOGIN( ADDR[SSS_EVENT], VALID[2] )
						else:
							if ( VALID[0] == iLOGOUT ):
								DATA[SSS_EVENT] = DATA.get( SSS_EVENT , '' ) + FX_LOGOUT( ADDR[ SSS_EVENT ] )
							else:
								if ( VALID[0] == iJOIN ):
									DATA[SSS_EVENT] = DATA.get( SSS_EVENT , '' ) + FX_JOIN( FX_GETUSER( ADDR[SSS_EVENT] ) , VALID[2] )
								else:
									if ( VALID[0] == iPART ):
										DATA[SSS_EVENT] = DATA.get( SSS_EVENT , '' ) + FX_PART( FX_GETUSER( ADDR[SSS_EVENT] ) , VALID[2] )
									else:
										if ( VALID[0] == iMSG ):
											DATA[SSS_EVENT] = DATA.get( SSS_EVENT , '' ) + FX_MSG_ROOM( FX_GETUSER( ADDR[SSS_EVENT] ), VALID[2] , string.join( VALID[2:] ))
										else:
											if ( VALID[0] == iMSGU ):
												DATA[SSS_EVENT] = DATA.get( SSS_EVENT , '' ) + FX_MSG_USER( VALID[2] , "GOTUSERMSG", string.join( VALID[2:] ))
											else:
												DATA[SSS_EVENT] = DATA.get( SSS_EVENT , '' ) + "DEFAULT EVENT\n"

					# - place in outgoing queue
					if SSS_EVENT not in QOUT:
						QOUT.append( SSS_EVENT )

				# --- MAIN SECTION --------------------------- #
				else:
					# - disconnection event
#					print "Disconnected: ", ADDR[SSS_EVENT]
					# - ensure user is cleared out
					FX_LOGOUT( ADDR[SSS_EVENT] )
					del ADDR[SSS_EVENT]
					try:
						QOUT.remove( SSS_EVENT )
					except ValueError:
						pass
					SSS_EVENT.close()
					QQIN.remove( SSS_EVENT )

		# --- OUTGOING --- #
		for SSS_EVENT in SSS_OUT:
			SSS_SEND = DATA.get( SSS_EVENT )
			if SSS_SEND:
				SSS_SEND_NUM = SSS_EVENT.send( SSS_SEND )
				SSS_SEND = SSS_SEND[SSS_SEND_NUM:]
			if SSS_SEND:
				DATA[SSS_EVENT] = SSS_SEND
			else:
				try:
					del DATA[SSS_EVENT]
				except KeyError:
					pass
				QOUT.remove(SSS_EVENT)
	# - THE END
	SSS.close()

# --- MAIN PROGRAM ----------------------------------------------------------- #
FX_MAIN()
# --- MAIN PROGRAM ----------------------------------------------------------- #