#!/usr/bin/env ruby

####################################
#          EMANAGEMENT
#Ejabberd Management script 
#New features for management and maintenance ejabberd
#@author Aldo Gordillo < agordillos@gmail.com >
#@version 2.1 - 28-2-2012
####################################

require 'logger'

path = "/var/log/ejabberd/scripts.log"
file = File.open(path, File::WRONLY | File::APPEND | File::CREAT)
file.sync = true
$logger = Logger.new(file)
$logger.level = Logger::DEBUG

def getOption(option)
  File.open('/etc/ejabberd/ssconfig.cfg', 'r') do |f1|  
    while line = f1.gets  
      line = line.gsub(/\n/,'')
      if line.match(/^#/)
        #Comments
      elsif line.match(/^#{option}/)
        return line.gsub(/#{option}/,'')
      end  
    end  
  end
  return "Undefined"
end


#Configuration variables
$verbose = (getOption("verbose=")=="true")
$ejabberd_user = getOption("ejabberd_server_user=")
$checkEjabberdctlQuotedString = false


PARAMS_FOR_COMMANDS = {
	'addBuddyToRoster' => 5,
	'removeBuddyFromRoster' => 2,
	'setBidireccionalBuddys' => 6,
	'unsetBidireccionalBuddys' => 4,
	'getRoster' => 1,
	'removeRoster' => 1,
	'removeAllRostersByDomain' => 1,
	'removeAllRosters' => 0,
	'getAllUserJidsWithRosterByDomain' => 1,
	'printAllRostersByDomain' => 1,
	'printAllRosters' => 0,
	'printAllBidirecctionalBuddysByDomain' => 1,
	'checkUserJid' => 1,
	'checkBidirecctionalBuddys' => 2,
	'sendPresence' => 2,
	'setPresence' => 1,
	'unsetPresence' => 1,
	'sendMessageToUser' => 3,
	'getUserResource' => 1,
	'isEjabberdNodeStarted' => 0,
	'broadcast' => 3,
	'broadcastToConnectedUsers' => 3,
        'getConnectedJidsByDomain' => 1,
	'getConnectedJids' => 0,
	'kickUserJid' => 1,
	'createPersistentRoom' => 2,
	'createRoom' => 2,
	'destroyRoom' => 2,
	'destroyAllRoomsByDomain' => 1,
	'destroyAllRooms' => 0,
	'getAllJidsOfRoom' => 2,
	'getAffiliationsOfRoom' => 2,
	'setJidAffiliationOfRoom' => 4,
	'printAllRoomsByDomain' => 1,
	'printAllRooms' => 0,
	'help' => 0,
}

SYNTAX_FOR_COMMANDS = {
	'addBuddyToRoster' => 'addBuddyToRoster userJid buddyJid buddyNick buddyGroup subscription_type',
	'removeBuddyFromRoster' => 'removeBuddyFromRoster userJid buddyJid',
	'setBidireccionalBuddys' => 'setBidireccionalBuddys userAJid userBJid userANick userBNick groupForA groupForB',
	'unsetBidireccionalBuddys' => 'unsetBidireccionalBuddys userJid oldFriendJid oldFriendNick groupForOldFriend',
	'getRoster' => 'getRoster userJid',
	'removeRoster' => 'removeRoster userJid',
	'removeAllRostersByDomain' => 'removeAllRostersByDomain domain',
	'removeAllRosters' => 'removeAllRosters',
	'getAllUserJidsWithRosterByDomain' => 'getAllUserJidsWithRosterByDomain domain',
	'printAllRostersByDomain' => 'printAllRostersByDomain domain',
	'printAllRosters' => 'printAllRosters',
	'printAllBidirecctionalBuddysByDomain' => 'printAllBidirecctionalBuddysByDomain domain',
	'checkUserJid' => 'checkUserJid userJid',
	'checkBidirecctionalBuddys' => 'checkBidirecctionalBuddys userAJid userBJid',
	'sendPresence' => 'sendPresence userJid show',
	'setPresence' => 'setPresence userJid',
	'unsetPresence' => 'unsetPresence userJid',
	'sendMessageToUser' => 'sendMessageToUser fromJid toJid msg',
	'getUserResource' => 'getUserResource userJid',
	'isEjabberdNodeStarted' => 'isEjabberdNodeStarted',
	'broadcast' => 'broadcast admin userJids msg ',
	'broadcastToConnectedUsers' => 'broadcastToConnectedUsers admin userJids msg (userJids value: "all" or jids array [jid1,jid2,...,jidN])',
        'getConnectedJidsByDomain' => 'getConnectedJidsByDomain domain',
	'getConnectedJids' => 'getConnectedJids',
	'kickUserJid' => 'kickUserJid(userJid)',
	'createPersistentRoom' => 'createPersistentRoom roomName domain',
	'createRoom' => 'createRoom roomName domain',
	'destroyRoom' => 'destroyRoom roomName domain',
	'destroyAllRoomsByDomain' => 'destroyAllRoomsByDomain domain',
	'destroyAllRooms' => 'destroyAllRooms',
	'getAllJidsOfRoom' => 'getAllJidsOfRoom roomName domain',
	'getAffiliationsOfRoom' => 'getAffiliationsOfRoom roomName domain',
	'setJidAffiliationOfRoom' => 'setJidAffiliationOfRoom roomName domain jid affiliation (affiliation value: owner/admin/member/outcast/none)',
	'printAllRoomsByDomain' => 'printAllRoomsByDomain domain',
	'printAllRooms' => 'printAllRooms',
	'help' => 'help',
}


#########################
#Debug methods
#########################
def ejabberdLog(text)
	$logger.info "Ejabberd Management Script: " + text
	if $verbose
		#puts "Writing to ejabberdLog: " + "Ejabberd Management Script: " + text
	end
end

def log(msg)
	logWithTitle(msg,nil)
end

def logWithTitle(msg,title)
	puts "------------------------"
	if title
		puts ("<<<" + title + ">>>")
	end
	puts msg
	puts "------------------------"
end


#########################
#Methods to manage rosters from Social Stream Rails App
#########################
def setBidireccionalBuddys(userAJid,userBJid,userANick,userBNick,groupForA,groupForB)
	addBuddyToRoster(userAJid,userBJid,userBNick,groupForB,"both")
	addBuddyToRoster(userBJid,userAJid,userANick,groupForA,"both")
	return "Done"
end

def unsetBidireccionalBuddys(userJid,oldFriendJid,oldFriendNick,groupForOldFriend)
	if checkBidirecctionalBuddys(userJid,oldFriendJid)
		removeBuddyFromRoster(userJid,oldFriendJid)
		removeBuddyFromRoster(oldFriendJid,userJid)
		addBuddyToRoster(userJid,oldFriendJid,oldFriendNick,groupForOldFriend,"to")
		return "Done"
	else
		return userJid + " and " + oldFriendJid + " aren't bidireccional buddys"
	end
end

def addBuddyToRoster(userJid,buddyJid,buddyNick,buddyGroup,subscription_type)
	user = getUsernameFromJid(userJid)
	buddy = getUsernameFromJid(buddyJid)
	userDomain = getDomainFromJid(userJid)
	buddyDomain = getDomainFromJid(buddyJid)
	executeCommand("ejabberdctl add-rosteritem " + user + " " + userDomain + " " + buddy + " " + buddyDomain + " " + buddyNick + " " + buddyGroup + " " + 				subscription_type)
	return "Done"
end

def removeBuddyFromRoster(userJid,buddyJid)
	user = getUsernameFromJid(userJid)
	buddy = getUsernameFromJid(buddyJid)
	userDomain = getDomainFromJid(userJid)
	buddyDomain = getDomainFromJid(buddyJid)
	if checkUserJidInRoster(buddyJid,getRoster(userJid))
		executeCommand("ejabberdctl delete_rosteritem " + user + " " + userDomain + " " + buddy + " " + buddyDomain)
		return "Done"
	else
		return "User " + buddyJid + " not found in " + userJid + " roster."
	end
end


#########################
#Roster Utilities
#########################
def getRoster(userJid)
	if checkUserJid(userJid)
		executeCommand("ejabberdctl get_roster " + getUsernameFromJid(userJid)  + " " + getDomainFromJid(userJid))
	else
		return "Roster not found for user " + userJid
	end
end

def getBuddyJidsFromRoster(roster)
	buddyJids = []
        lines = roster.split("\n")
	lines.each do |line|
		buddyJids << line.split(" ")[0]
	end
	buddyJids
end

def removeRoster(userJid)
	if checkUserJid(userJid)
		executeCommand("ejabberdctl process_rosteritems delete any any " + userJid + " any")
		return "Done"	
	else
		return "Roster not found for userJid " + userJid
	end
end

def removeAllRostersByDomain(domain)
	if(domain=="all")
		executeCommand("ejabberdctl process_rosteritems delete any any any any")
	else
		executeCommand("ejabberdctl process_rosteritems delete any any *@" + domain + " any")
	end
	return "Done";
end

def removeAllRosters()
	return removeAllRostersByDomain("all")
end

def getAllUserJidsWithRosterByDomain(domain)
	if(domain=="all")
		output = executeCommand("ejabberdctl process_rosteritems list any any any any")
	else
		output = executeCommand("ejabberdctl process_rosteritems list any any *@" + domain + " any")
	end

	userJids = []
	lines = output.split("\n");
	items = lines[0].split(" ")[2]

	#test if items string contains a correct number
	if items.to_i.to_s == items
		if items.to_i > 0
			lines.each do |line|
				if line.split(":")[0]=="Matches"
					userJid = line.split(" ")[1]
					#puts "Line:" + line
					unless userJids.include?(userJid)
						userJids << userJid
					end	
				end
			end
		end
	end
	
	return userJids;
end

def getAllUserJidsWithRoster()
	return getAllUserJidsWithRosterByDomain("all")
end

def getAllRostersByDomain(domain)
	rosterList = { }

	userJids = getAllUserJidsWithRosterByDomain(domain)

	userJids.each do |userJid|
		roster = getRoster(userJid)
		rosterList.store(userJid,roster)
	end
	return rosterList
end

def getAllRosters()
	return getAllRostersByDomain("all");
end

def printAllRostersByDomain(domain)
	rosterList = getAllRostersByDomain(domain)
	rosterList.keys.each do |userJid|
		puts "\n"
		puts "-------------------------------------"
		puts userJid + " Roster"
		puts "-------------------------------------"
		puts rosterList[userJid]
		puts "-------------------------------------"
		puts "\n"
	end
	return "Done"
end

def printAllRosters()
	return printAllRostersByDomain("all")
end

def getAllBidirecctionalBuddysByDomain(domain)
	b_buddyJids = []
	userJids = getAllUserJidsWithRosterByDomain(domain)
	nonCheckedUsers = userJids
	userJids.each do |userJid|
		nonCheckedUsers.delete(userJid)
		nonCheckedUsers.each do |checkUserJid|
			if checkBidirecctionalBuddys(userJid,checkUserJid)
				b_buddyJids << [userJid,checkUserJid]
			end
		end
	end
	return b_buddyJids
end

def getAllBidirecctionalBuddys()
	return getAllBidirecctionalBuddysByDomain("all")
end

def printAllBidirecctionalBuddysByDomain(domain)
	puts "This may take a while..."
	b_buddys = getAllBidirecctionalBuddysByDomain(domain)
	b_buddys.each do |contact|
		puts "[" + contact[0] + "," + contact[1] + "]"
	end
	return "Done"
end

def printAllBidirecctionalBuddys()
	return printAllBidirecctionalBuddysByDomain("all")
end

#Check if the user have a roster in his domain
def checkUserJid(userJid)
	domain = getDomainFromJid(userJid)
	return getAllUserJidsWithRosterByDomain(domain).include?(userJid)
end

def checkUserJidInRoster(userJid,roster)
	return getBuddyJidsFromRoster(roster).include?(userJid)
end

def checkBidirecctionalBuddys(userAJid,userBJid)
	rosterA = getRoster(userAJid)
	rosterB = getRoster(userBJid)
	return (checkUserJidInRoster(userAJid,rosterB) and checkUserJidInRoster(userBJid,rosterA))
end


#########################
#Manage stanzas Utilities
#########################
def setPresence(userJid)
	sendPresenceStanzaWithType(userJid,userJid,"available")
end

def unsetPresence(userJid)
	sendPresenceStanzaWithType(userJid,userJid,"unavailable")
end

def sendPresence(username,show)
	sendPresenceWithShow(username,username,show)
end

def sendPresenceWithShow(fromJid,toJid,show)
	resource = getUserResource(fromJid);  
	from_name = getUsernameFromJid(fromJid)
	domain = getDomainFromJid(fromJid)
	pres_stanza = "\\<" + buildQuotedString("presence from=") + "\\\"" + buildQuotedString(fromJid) + "\\\"" + buildQuotedString(" to=") + "\\\"" +
			buildQuotedString(toJid) + "\\\"\\>\\<" + buildQuotedString("show") + "\\>" + buildQuotedString(show) + "\\<" +
			buildQuotedString("/show") + "\\>\\<" + buildQuotedString("/presence") + "\\>"
	executeCommand("ejabberdctl send_stanza_c2s " + from_name + " " + domain + " " + resource + " " + pres_stanza)
	return "Done"
end

def sendPresenceStanzaWithType(fromJid,toJid,presence_type)
	resource = getUserResource(fromJid);
	from_name = getUsernameFromJid(fromJid)
	domain = getDomainFromJid(fromJid)
	pres_stanza = "\\<" + buildQuotedString("presence type=") + "\\\"" + buildQuotedString(presence_type) + "\\\"" + buildQuotedString(" from=") + "\\\"" + 
		buildQuotedString(fromJid) + "\\\"" + buildQuotedString(" to=") + "\\\"" + buildQuotedString(toJid) + "\\\"\\>\\<" + buildQuotedString("/presence") + "\\>" 
	executeCommand("ejabberdctl send_stanza_c2s " + from_name + " " + domain + " " + resource + " " + pres_stanza)
	return "Done"
end

def sendMessageToUser(fromJid,toJid,msg)
	executeCommand("ejabberdctl send_message_chat " + fromJid + " " + toJid + " " + buildQuotedString(msg))
	return "Done"
end

def getUserResource(userJid)
	output = executeCommand("ejabberdctl connected-users")
        lines = output.split("\n");
	lines.each do |line|
		if line.split("/")[0] == userJid
			#puts "Find " + userJid
			resource = line.split("/")[1];
			return resource;	
		end
	end
	return userJid + " no have any active session"
end


#########################
#Utilities
#########################
def getUsernameFromJid(jid)
	return jid.split("@")[0];
end

def getDomainFromJid(jid)
	return jid.split("@")[1];
end

def getElementsFromStringArray(stringArray)
	stringArray=stringArray[1,stringArray.length-2]
	return stringArray.split(",")
end


#Determine how to scape characters for build quoted strings
def checkEjabberdctlQuotedString
	puts "checkForSimpleSlash: " + checkForSimpleSlash.to_s()
	puts "checkForDoubleSlash: " + checkForDoubleSlash.to_s()
end

def checkAndSetEjabberdctlQuotedString
	if checkForSimpleSlash
		$checkForSimpleSlash = true
	end
	if checkForDoubleSlash
		$checkForDoubleSlash = true
	end
	$checkEjabberdctlQuotedString = true	
end

def checkForDoubleSlash
	command = "ejabberdctl send_message_chat example@localhost example@localhost \\'Hello quoted string\\'"
	if execute_as_sudo
		command = "sudo " + command
	end
	
	output = %x[#{command}]
	firstLine = ""
	lines = output.split("\n")
	lines.each do |line|
		if line != ""
			firstLine = line
			break
		end
	end

	if firstLine==""
		return true
	elsif firstLine.split(":")[0]=="Error"
		return false
	else
		#Unknown error		
		return false
	end
end

def checkForSimpleSlash
	command = "ejabberdctl send_message_chat example@localhost example@localhost \'Hello quoted string\'"
	if execute_as_sudo
		command = "sudo " + command
	end
	#puts "Executing " + command
	output = %x[#{command}]
	firstLine = ""
	lines = output.split("\n")
	lines.each do |line|
		if line != ""
			firstLine = line
			break
		end
	end

	if firstLine==""
		return true
	elsif firstLine.split(":")[0]=="Error"
		return false
	else
		#Unknown error		
		return false
	end
end

def buildQuotedString(msg)
	if !$checkEjabberdctlQuotedString
		checkAndSetEjabberdctlQuotedString
	end
	
	if $checkForSimpleSlash
		return "\'" + msg + "\'"
	end
	if $checkForDoubleSlash
		return "\\'" + msg + "\\'"
	end
	return msg
end


#########################
#Connection Info
#########################
def isEjabberdNodeStarted
	output = executeCommand("ejabberdctl status")
	if firstLine = output.split("\n")[0]
	  return ((firstLine.split(":")[1]).strip()=="started")
	end
	return false
end

def getConnectedJidsByDomain(domain)
        jids = []
	if(domain=="all")
		output = executeCommand("ejabberdctl connected-users")
	else
		output = executeCommand("ejabberdctl connected-users | grep @" + domain + "/")
	end
	sessions = output.split("\n")
	sessions.each do |session|  
	  jids << session.split("/")[0]
	end
	return jids
end

def getConnectedJids()
	return getConnectedJidsByDomain("all")
end


#########################
#Advanced features
#########################
def broadcast(admin,userJids,msg)
	getElementsFromStringArray(userJids).each do |userJid|
		sendMessageToUser(admin,userJid,msg)
	end
	return "Done"
end

def broadcastToConnectedUsers(admin,userJids,msg)
	connectedJids = getConnectedJids()
	if(userJids=="all")
		connectedJids.each do |userJid|
			sendMessageToUser(admin,userJid,msg)
		end
	else
		getElementsFromStringArray(userJids).each do |userJid|
			if (connectedJids.include?(userJid))
				sendMessageToUser(admin,userJid,msg)
			end
		end
	end
	return "Done"
end

def kickUserJid(userJid)
	user = getUsernameFromJid(userJid)
	userDomain = getDomainFromJid(userJid)
	resource = getUserResource(userJid);  
	reason = "0"
	executeCommand("ejabberdctl kick_session " + user + " " + userDomain + " " + resource + " " + reason)
	return "Done"
end

def banUserJid(userJid)
	kickUserJid(userJid)
	#Notify Web server to ban user permanently
end



#########################
#MUC Methods (Multi-user-chat)
#########################

$muc_host = "conference"

def printAllRoomsByDomain(domain)
	if(domain=="all")
		return printAllRooms
	end
	output = executeCommand("ejabberdctl muc_online_rooms " + domain)
	return output;
end

def printAllRooms
	output = executeCommand("ejabberdctl muc_online_rooms global")
	return output;
end

def createPersistentRoom(roomName,domain)
	createRoom(roomName,domain)
	setRoomPersistence(roomName,domain,true)
	return "Done"
end

def createRoom(roomName,domain)
	executeCommand("ejabberdctl create_room " + roomName + " " + $muc_host + "." + domain + " " + domain)
	return "Done"
end

def setRoomPersistence(roomName,domain,persistence)
	if (persistence==true)
		executeCommand("ejabberdctl change_room_option " + roomName + " " + $muc_host + "." + domain + " persistent true")
	else
		executeCommand("ejabberdctl change_room_option " + roomName + " " + $muc_host + "." + domain + " persistent false")
	end
end

def destroyAllRoomsByDomain(domain)
	rooms = getAllRoomsByDomain(domain)
	rooms.each do |room|
		destroyRoom(getNameFromRoom(room),getDomainFromRoom(room))
	end
	return "Done"
end

def destroyAllRooms()
	return destroyAllRoomsByDomain("all")
end

def destroyRoom(roomName,domain)
	executeCommand("ejabberdctl destroy_room " + roomName + " " + $muc_host + "." + domain + " " + domain)
	return "Done"
end

def getAllRoomsByDomain(domain)
	if(domain=="all")
		output = executeCommand("ejabberdctl muc_online_rooms global")
	else
		output = executeCommand("ejabberdctl muc_online_rooms " + domain)
	end
	rooms = output.split("\n")
	return rooms
end

def getAllJidsOfRoomObject(room)
	return getAllJidsOfRoom(getNameFromRoom(room),getDomainFromRoom(room))
end

def getAllJidsOfRoom(roomName,domain)
	jids = []
	output = executeCommand("ejabberdctl get_room_occupants " + roomName + " " +  $muc_host + "." + domain)
	lines = output.split("\n")
	lines.each do |line|
		jids << line.split("/")[0]
	end
	return jids
end

def getAffiliationsOfRoomObject(room)
	return getAffiliationsOfRoom(getNameFromRoom(room),getDomainFromRoom(room))
end

def getAffiliationsOfRoom(roomName,domain)
	affiliateds = []
	output = executeCommand("ejabberdctl get_room_affiliations " + roomName + " " +  $muc_host + "." + domain)
	lines = output.split("\n")
	lines.each do |line|
		split=line.split(" ");
		jid = split[0]+"@"+split[1];
		affiliation=split[2]
		affiliateds << [jid,affiliation]
	end
	return affiliateds
end

#affiliation = Owner/Admin/Member/Outcast/None
def setJidAffiliationOfRoom(roomName,domain,jid,affiliation)
	output = executeCommand("ejabberdctl set_room_affiliation " + roomName + " " +  $muc_host + "." + domain + " " + jid + " " + affiliation)
	return "Done"
end

def getDomainFromRoom(room)
	return room.split(".")[1];
end

def getNameFromRoom(room)
	return room.split("@")[0];
end


#########################
#Execution commands methods
#########################
def executeCommand(command)
	#Building...
	command = buildCommand(command)
	
	if $verbose
		#Logging...
		puts "Executing: " + command
		ejabberdLog("Executing (#{command})")
	end

	#Executing...
	output = %x[#{command}]
	return output
end

def buildCommand(command)
	if execute_as_sudo
		command = "sudo -u " + $ejabberd_user + " " + command
	end
	return command
end

def execute_as_sudo
	sudo_users = getOption("users_require_sudo=")

	if sudo_users=="all"
	 return true
	end

	sudo_users_array = sudo_users.split(",")
        current_user = %x["whoami"].split("\n")[0]
	if sudo_users_array.include?(current_user)
		return true
	end
end


#########################
#Support
#########################
def help
	log("Command list")
	SYNTAX_FOR_COMMANDS.values.each do |command|
		puts command
	end
	puts ""
end



#########################
#Main thread
#########################
#log("Init Ejabberd Maintenance script")

begin
	if ARGV[0] and PARAMS_FOR_COMMANDS.keys.include?(ARGV[0])
		if (ARGV.length == (PARAMS_FOR_COMMANDS[ARGV[0]]+1))

			ejabberdLog("Executing (#{ARGV})")

			length = ARGV.length;
			case length
			when 1
			  puts send(ARGV[0])
			when 2
			  puts send(ARGV[0],ARGV[1])
			when 3
			  puts send(ARGV[0],ARGV[1],ARGV[2])
			when 4
			  puts send(ARGV[0],ARGV[1],ARGV[2],ARGV[3])
			when 5
			  puts send(ARGV[0],ARGV[1],ARGV[2],ARGV[3],ARGV[4])
			when 6
			  puts send(ARGV[0],ARGV[1],ARGV[2],ARGV[3],ARGV[4],ARGV[5])
			when 7
			  puts send(ARGV[0],ARGV[1],ARGV[2],ARGV[3],ARGV[4],ARGV[5],ARGV[6])
			else
			  puts send(ARGV[0],ARGV)
			end
			puts ""
		else 
			puts "Error: Params required for command " + (ARGV[0]).to_s() + ": " + (PARAMS_FOR_COMMANDS[ARGV[0]]).to_s()
			puts "Syntax for command " + (ARGV[0]).to_s() + "\n" + (SYNTAX_FOR_COMMANDS[ARGV[0]]).to_s()
			puts "Use 'help' to get information about command syntax"
		end
	else
		if (ARGV.length==0)
			help
		else
			puts "Error: Command not recognized"
			puts "Use 'help' to get information about command syntax"
		end
	end
rescue
	puts "Syntax error"
	puts "Use 'help' to get information about command syntax"
end