#!/usr/bin/env ruby

#Ejabberd Management script
#New features for management and maintenance ejabberd
#@author Aldo

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
$domain = getOption("server_domain=")
$verbose = (getOption("verbose=")=="true")


PARAMS_FOR_COMMANDS = {
	'addBuddyToRoster' => 5,
	'removeBuddyFromRoster' => 2,
	'setBidireccionalBuddys' => 6,
	'unsetBidireccionalBuddys' => 4,
	'getRoster' => 1,
	'removeRoster' => 1,
	'removeAllRosters' => 0,
	'getBuddysFromRoster' => 1,
	'getAllUsersWithRoster' => 0,
	'getAllRosters' => 0,
	'printAllRosters' => 0,
	'printAllBidirecctionalBuddys' => 0,
	'checkUser' => 1,
	'checkBidirecctionalBuddys' => 2,
	'sendPresence' => 2,
	'setPresence' => 1,
	'unsetPresence' => 1,
	'sendMessageToUser' => 3,
	'getUserResource' => 1,
	'isEjabberdNodeStarted' => 0,
	'checkEjabberdctlQuotedString' => 0,
	'help' => 0,
}

SYNTAX_FOR_COMMANDS = {
	'addBuddyToRoster' => 'addBuddyToRoster userSid buddySid buddyNick buddyGroup subscription_type',
	'removeBuddyFromRoster' => 'removeBuddyFromRoster userSid buddySid',
	'setBidireccionalBuddys' => 'setBidireccionalBuddys userASid userBSid userANick userBNick groupForA groupForB',
	'unsetBidireccionalBuddys' => 'unsetBidireccionalBuddys userSid oldFriendSid oldFriendNick groupForOldFriend',
	'getRoster' => 'getRoster username',
	'removeRoster' => 'removeRoster username',
	'removeAllRosters' => 'removeAllRosters',
	'getBuddysFromRoster' => 'getBuddysFromRoster roster',
	'getAllUsersWithRoster' => 'getAllUsersWithRoster',
	'getAllRosters' => 'getAllRosters',
	'printAllRosters' => 'printAllRosters',
	'printAllBidirecctionalBuddys' => 'printAllBidirecctionalBuddys',
	'checkUser' => 'checkUser user',
	'checkBidirecctionalBuddys' => 'checkBidirecctionalBuddys userASid userBSid',
	'sendPresence' => 'sendPresence username show',
	'setPresence' => 'setPresence username',
	'unsetPresence' => 'unsetPresence username',
	'sendMessageToUser' => 'sendMessageToUser from_name to_name msg',
	'getUserResource' => 'getUserResource username',
	'isEjabberdNodeStarted' => 'isEjabberdNodeStarted',
	'checkEjabberdctlQuotedString' => 'checkEjabberdctlQuotedString',
	'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(userASid,userBSid,userANick,userBNick,groupForA,groupForB)
	addBuddyToRoster(userASid,userBSid,userBNick,groupForB,"both")
	addBuddyToRoster(userBSid,userASid,userANick,groupForA,"both")
	return "Done"
end

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

def addBuddyToRoster(userSid,buddySid,buddyNick,buddyGroup,subscription_type)
	user = userSid.split("@")[0]
	buddy = buddySid.split("@")[0]
	userDomain = userSid.split("@")[1]
	buddyDomain = buddySid.split("@")[1]
	executeCommand("ejabberdctl add-rosteritem " + user + " " + userDomain + " " + buddy + " " + buddyDomain + " " + buddyNick + " " + buddyGroup + " " + 				subscription_type)
	return "Done"
end

def removeBuddyFromRoster(userSid,buddySid)
	user = userSid.split("@")[0]
	buddy = buddySid.split("@")[0]
	userDomain = userSid.split("@")[1]
	buddyDomain = buddySid.split("@")[1]
	if checkUserInRoster(buddy,getRoster(user))
		executeCommand("ejabberdctl delete_rosteritem " + user + " " + userDomain + " " + buddy + " " + buddyDomain)
		return "Done"
	else
		return "User " + buddy + " not found in " + user + " roster."
	end
end


#Roster Utilities

def getRoster(username)
	if checkUser(username)
		executeCommand("ejabberdctl get_roster " + username  + " " + $domain)
	else
		return "Roster not found for user " + username
	end
end

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

def removeRoster(username)
	if checkUser(username)
		user_sid = username + "@" + $domain
		user_roster = getRoster(username);
		user_buddys = getBuddysFromRoster(user_roster);
		user_buddys.each do |buddy|
			buddy_sid = buddy + "@" + $domain
			removeBuddyFromRoster(user_sid,buddy_sid)
		end
		return "Done"	
	else
		return "Roster not found for user " + username
	end
end

def removeAllRosters()
	executeCommand("ejabberdctl process_rosteritems delete any any any any")
	return "Done";
end

def getAllUsersWithRoster()
	output = executeCommand("ejabberdctl process_rosteritems list any any any any")

	users = []

	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"
					user = line.split(" ")[1].split("@")[0]
					#puts i.to_s() + " :" + line
					unless users.include?(user)
						users << user
					end	
				end
			end
		end
	end
	
	return users;
end

def getAllRosters()
	rosterList = { }

	users = getAllUsersWithRoster()

	users.each do |user|
		roster = getRoster(user)
		rosterList.store(user,roster)
	end
	return rosterList
end

def printAllRosters()
	rosterList = getAllRosters()
	rosterList.keys.each do |user|
		puts "\n"
		puts "-------------------------------------"
		puts user + " Roster"
		puts "-------------------------------------"
		puts rosterList[user]
		puts "-------------------------------------"
		puts "\n"
	end
	return "Done"
end

def getAllBidirecctionalBuddys()
	b_buddys = []
	users = getAllUsersWithRoster()
	nonCheckedUsers = users
	users.each do |user|
		nonCheckedUsers.delete(user)
		nonCheckedUsers.each do |checkUser|
			if checkBidirecctionalBuddys(user + "@" + $domain,checkUser + "@" + $domain)
				b_buddys << [user,checkUser]
			end
		end
	end
	return b_buddys
end

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

#Check if the user have a roster
def checkUser(user)
	return getAllUsersWithRoster().include?(user)
end

def checkUserInRoster(user,roster)
	return getBuddysFromRoster(roster).include?(user)
end

def checkBidirecctionalBuddys(userASid,userBSid)
	userA = userASid.split("@")[0]
	userB = userBSid.split("@")[0]
	rosterA = getRoster(userA)
	rosterB = getRoster(userB)
	return (checkUserInRoster(userA,rosterB) and checkUserInRoster(userB,rosterA))
end


#Manage stanzas Utilities

def sendStanzaUserMessage(username,msg)
	resource = getUserResource(username);
	stanza = "\\<\\'message\\'\\>\\<\\'body\\'\\>\\'" + msg + "\\'\\<\\'/body\\'\\>\\<\\'/message\\'\\>"
	executeCommand("ejabberdctl send_stanza_c2s " + username + " " + $domain + " " + resource + " " + stanza)
	return "Done"
end

def setPresence(username)
	sendPresenceStanzaWithType(username,username,"available")
end

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

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

def sendPresenceWithShow(from_name,to_name,show)
	puts from_name
	resource = getUserResource(from_name);
	puts resource
	from_sid = from_name + "@" + $domain;
	to_sid = to_name + "@" + $domain;   
	pres_stanza = "\\<\\'presence from=\\'\\\"\\'" + from_sid + "\\'\\\"\\' to=\\'\\\"\\'" + to_sid + "\\'\\\"\\>\\<\\'show\\'\\>\\'" + show + "\\'\\<\\'/show\\'\\>\\<\\'/presence\\'\\>" 
	executeCommand("ejabberdctl send_stanza_c2s " + from_name + " " + $domain + " " + resource + " " + pres_stanza)
	return "Done"
end

def sendPresenceStanzaWithType(from_name,to_name,presence_type)
	resource = getUserResource(from_name);
	from_sid = from_name + "@" + $domain;
	to_sid = to_name + "@" + $domain;
	pres_stanza = "\\<\\'presence type=\\'\\\"\\'" + presence_type + "\\'\\\"\\' from=\\'\\\"\\'" + from_sid + "\\'\\\"\\' to=\\'\\\"\\'" + to_sid + "\\'\\\"\\>\\<\\'/presence\\'\\>" 
	executeCommand("ejabberdctl send_stanza_c2s " + from_name + " " + $domain + " " + resource + " " + pres_stanza)
	return "Done"
end

def sendMessageToUser(from_name,to_name,msg)
	from_sid = from_name + "@" + $domain;
	to_sid = to_name + "@" + $domain;
	executeCommand("ejabberdctl send_message_chat " + from_sid + " " + to_sid + " " + buildQuotedString(msg))
	return "Done"
end

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


#More utilities

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

def checkEjabberdctlQuotedString
	puts "checkForSimpleSlash: " + checkForSimpleSlash.to_s()
	puts "checkForDoubleSlash: " + checkForDoubleSlash.to_s()
end

def checkForDoubleSlash
	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 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 checkForSimpleSlash
		return "\'" + msg + "\'"
	end
	if checkForDoubleSlash
		return "\\'" + msg + "\\'"
	end
	return msg
end


#Help & Support methods

def executeCommand(command)
	#Building...
	command = buildCommand(command)
	
	if $verbose
		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 " + command
	end
	return command
end

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

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
		puts "Error: Command not recognized"
		puts "Use 'help' to get information about command syntax"
	end
rescue
	puts "Syntax error"
	puts "Use 'help' to get information about command syntax"
end